feat: make code generation runs resumable

.gitignore

@@ -1,5 +1,7 @@
 __pycache__/
 *.pyc
+.pytest_cache/
+.codex_pytest_tmp*/
 .DS_Store
 out
 website/

code_agent/backends.py

@@ -17,6 +17,7 @@ class OpenAICompatibleGenerator:
         temperature: float = 0.7,
         max_tokens: int = 1024,
         open_thinking: bool = False,
+        seed: int | None = None,
         client: Any = None,
     ) -> None:
         if client is None:
@@ -30,15 +31,23 @@ class OpenAICompatibleGenerator:
         self.temperature = temperature
         self.max_tokens = max_tokens
         self.open_thinking = open_thinking
+        self.seed = seed
+        self._request_index = 0
 
     def __call__(self, prompt: str) -> str:
+        request = {
+            "model": self.model,
+            "messages": [{"role": "user", "content": prompt}],
+            "temperature": self.temperature,
+            "max_tokens": self.max_tokens,
+            "stream": False,
+            "extra_body": {"open_thinking": self.open_thinking},
+        }
+        if self.seed is not None:
+            request["seed"] = self.seed + self._request_index
+            self._request_index += 1
         response = self.client.chat.completions.create(
-            model=self.model,
+            **request,
-            messages=[{"role": "user", "content": prompt}],
-            temperature=self.temperature,
-            max_tokens=self.max_tokens,
-            stream=False,
-            extra_body={"open_thinking": self.open_thinking},
         )
         if not response.choices or response.choices[0].message is None:
             raise RuntimeError("Model API returned no completion choice.")

docs/code-agent.md

@@ -37,13 +37,35 @@ Each JSONL record contains a prompt, a required function name, and JSON-serializ
 
 `dataset/code_rl_mini.jsonl` contains four smoke-test tasks. It is deliberately too small for a meaningful training claim; use it to validate the pipeline before preparing a larger train/validation/test split.
 
+For a standard public benchmark, convert the compatible subset of MBPP-sanitized while preserving its official train/validation/test split:
+
+```bash
+python scripts/prepare_mbpp.py --output-dir out/mbpp_sanitized
+```
+
+The converter accepts direct, JSON-serializable function assertions, validates every converted task against the published reference solution, and reports all skipped records. It also emits `train_sft.jsonl` from the official training references, enabling an SFT → execution-reward GRPO ablation without leaking validation or test solutions. Generated benchmark files remain under the ignored `out/` directory.
+
 ## Evaluate saved candidates
 
+Generate multiple candidates from a running OpenAI-compatible model service with deterministic per-request seeds:
+
+```bash
+python scripts/generate_codegen.py \
+  --tasks dataset/code_rl_mini.jsonl \
+  --samples-per-task 10 \
+  --seed 42 \
+  --output out/code_predictions.jsonl
+```
+
+Each completed task is flushed to JSONL immediately. If a long run is interrupted, rerun the same command with `--resume`; task-specific seed ranges keep completed and resumed outputs reproducible.
+
+Then execute and aggregate them:
+
 ```bash
 python scripts/eval_codegen.py \
   --tasks dataset/code_rl_mini.jsonl \
-  --predictions tests/fixtures/code_rl_predictions.jsonl \
+  --predictions out/code_predictions.jsonl \
-  --ks 1 2
+  --ks 1 5 10
 ```
 
 The report separates:
@@ -111,7 +133,7 @@ Use the same held-out tasks and sampling parameters for three comparisons:
 2. inference-time execution feedback and repair;
 3. execution-reward GRPO/CISPO followed by the same inference procedure.
 
-Record pass@1, pass@k, test pass rate, compile rate, timeout rate, average attempts, average generation length, reward mean/std, KL to the reference model, and degenerate-group rate. Model results should only be added to the resume after this protocol has been run on AutoDL and the raw logs are retained.
+Record pass@1, pass@k, test pass rate, compile rate, timeout rate, average attempts, average generation length, reward mean/std, KL to the reference model, and degenerate-group rate. Publish model results only after this protocol has been run and the raw logs and configuration have been retained for reproduction.
 
 ## Attribution
 

requirements.txt

@@ -1,5 +1,6 @@
 datasets==3.6.0
 datasketch==1.6.4
+fastapi==0.115.6
 Flask==3.0.3
 Flask_Cors==4.0.0
 jieba==0.42.1
@@ -23,6 +24,7 @@ jinja2==3.1.2
 jsonlines==4.0.0
 trl==0.13.0
 ujson==5.1.0
+uvicorn==0.34.0
 wandb==0.18.3
 streamlit==1.50.0
 einops==0.8.1

scripts/generate_codegen.py

@@ -0,0 +1,108 @@
+"""Generate reproducible multi-sample candidates for execution-based evaluation."""
+
+import argparse
+import json
+import os
+import sys
+from pathlib import Path
+
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+from code_agent import CodeTask, OpenAICompatibleGenerator
+
+
+def read_tasks(path):
+    tasks = []
+    with Path(path).open("r", encoding="utf-8") as handle:
+        for line_number, line in enumerate(handle, start=1):
+            if not line.strip():
+                continue
+            try:
+                tasks.append(CodeTask.from_dict(json.loads(line)))
+            except (KeyError, TypeError, ValueError, json.JSONDecodeError) as exc:
+                raise ValueError(f"Invalid task at {path}:{line_number}: {exc}") from exc
+    return tasks
+
+
+def build_prompt(task):
+    return (
+        "You are solving a Python algorithm task. Return exactly one ```python``` code block.\n"
+        f"Required entry point: {task.entry_point}\n"
+        f"Task:\n{task.prompt}\n"
+        "Do not read files, access the network, spawn processes, or call unsafe builtins."
+    )
+
+
+def generate_predictions(tasks, generator, samples_per_task):
+    if samples_per_task < 1:
+        raise ValueError("samples_per_task must be at least 1")
+    return [
+        {
+            "task_id": task.task_id,
+            "candidates": [generator(build_prompt(task)) for _ in range(samples_per_task)],
+        }
+        for task in tasks
+    ]
+
+
+def read_completed_task_ids(path):
+    completed = set()
+    if not Path(path).exists():
+        return completed
+    with Path(path).open("r", encoding="utf-8") as handle:
+        for line_number, line in enumerate(handle, start=1):
+            if not line.strip():
+                continue
+            try:
+                record = json.loads(line)
+                completed.add(str(record["task_id"]))
+            except (KeyError, json.JSONDecodeError) as exc:
+                raise ValueError(f"Invalid existing prediction at {path}:{line_number}: {exc}") from exc
+    return completed
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Generate code candidates through an OpenAI-compatible API")
+    parser.add_argument("--tasks", required=True, help="CodeTask JSONL file")
+    parser.add_argument("--output", required=True, help="Prediction JSONL output path")
+    parser.add_argument("--samples-per-task", type=int, default=10)
+    parser.add_argument("--base-url", default="http://localhost:8998/v1")
+    parser.add_argument("--api-key", default="not-needed")
+    parser.add_argument("--model", default="minimind")
+    parser.add_argument("--temperature", type=float, default=0.8)
+    parser.add_argument("--max-tokens", type=int, default=512)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--open-thinking", action="store_true")
+    parser.add_argument("--resume", action="store_true", help="Append missing tasks to an existing output file")
+    args = parser.parse_args()
+
+    tasks = read_tasks(args.tasks)
+    output_path = Path(args.output)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    completed = read_completed_task_ids(output_path) if args.resume else set()
+    mode = "a" if args.resume else "w"
+    generated_count = 0
+    with output_path.open(mode, encoding="utf-8") as handle:
+        for task_index, task in enumerate(tasks):
+            if task.task_id in completed:
+                print(f"[{task_index + 1}/{len(tasks)}] {task.task_id}: already complete", flush=True)
+                continue
+            generator = OpenAICompatibleGenerator(
+                base_url=args.base_url,
+                api_key=args.api_key,
+                model=args.model,
+                temperature=args.temperature,
+                max_tokens=args.max_tokens,
+                open_thinking=args.open_thinking,
+                seed=args.seed + task_index * args.samples_per_task,
+            )
+            record = generate_predictions([task], generator, args.samples_per_task)[0]
+            handle.write(json.dumps(record, ensure_ascii=False) + "\n")
+            handle.flush()
+            generated_count += args.samples_per_task
+            print(f"[{task_index + 1}/{len(tasks)}] {task.task_id}: wrote {args.samples_per_task} candidates", flush=True)
+    print(f"Generated {generated_count} new candidates -> {output_path}")
+
+
+if __name__ == "__main__":
+    main()

scripts/prepare_mbpp.py

@@ -0,0 +1,149 @@
+"""Convert compatible MBPP-sanitized assertions into MiniMind CodeTask splits."""
+
+import argparse
+import ast
+import json
+import os
+import sys
+from collections import Counter
+from pathlib import Path
+
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+from code_agent import CodeTask, VerificationStatus, verify_code
+
+
+def _literal(node):
+    value = ast.literal_eval(node)
+    json.dumps(value, ensure_ascii=False)
+    return value
+
+
+def parse_assertion(source):
+    statement = ast.parse(source.strip()).body
+    if len(statement) != 1 or not isinstance(statement[0], ast.Assert):
+        raise ValueError("test is not a single assert statement")
+    expression = statement[0].test
+
+    if isinstance(expression, ast.Call):
+        call, expected = expression, True
+    elif isinstance(expression, ast.UnaryOp) and isinstance(expression.op, ast.Not) and isinstance(expression.operand, ast.Call):
+        call, expected = expression.operand, False
+    elif (
+        isinstance(expression, ast.Compare)
+        and len(expression.ops) == 1
+        and isinstance(expression.ops[0], ast.Eq)
+        and len(expression.comparators) == 1
+    ):
+        left, right = expression.left, expression.comparators[0]
+        if isinstance(left, ast.Call):
+            call, expected = left, _literal(right)
+        elif isinstance(right, ast.Call):
+            call, expected = right, _literal(left)
+        else:
+            raise ValueError("equality does not compare a function call")
+    else:
+        raise ValueError("unsupported assertion shape")
+
+    if not isinstance(call.func, ast.Name):
+        raise ValueError("entry point is not a direct function call")
+    args = [_literal(argument) for argument in call.args]
+    kwargs = {}
+    for keyword in call.keywords:
+        if keyword.arg is None:
+            raise ValueError("expanded keyword arguments are unsupported")
+        kwargs[keyword.arg] = _literal(keyword.value)
+    return call.func.id, {"args": args, "kwargs": kwargs, "expected": expected}
+
+
+def convert_record(record):
+    parsed = [parse_assertion(test) for test in record["test_list"]]
+    entry_points = {entry_point for entry_point, _ in parsed}
+    if len(entry_points) != 1:
+        raise ValueError("tests use multiple entry points")
+    payload = {
+        "task_id": f"mbpp-{record['task_id']}",
+        "prompt": record["prompt"].strip(),
+        "entry_point": entry_points.pop(),
+        "tests": [test for _, test in parsed],
+    }
+    task = CodeTask.from_dict(payload)
+    verification = verify_code(record["code"], task)
+    if verification.status != VerificationStatus.PASSED:
+        raise ValueError(f"reference solution failed verifier: {verification.status.value}")
+    return task.to_dict()
+
+
+def build_sft_record(record, task):
+    return {
+        "conversations": [
+            {
+                "role": "user",
+                "content": (
+                    "Solve the Python algorithm task below. Return one complete ```python``` code block.\n"
+                    f"Required entry point: {task['entry_point']}\n"
+                    f"Task:\n{task['prompt']}"
+                ),
+            },
+            {"role": "assistant", "content": f"```python\n{record['code'].strip()}\n```"},
+        ]
+    }
+
+
+def convert_split(records):
+    converted = []
+    sft_records = []
+    skipped = Counter()
+    for record in records:
+        try:
+            task = convert_record(record)
+            converted.append(task)
+            sft_records.append(build_sft_record(record, task))
+        except (KeyError, SyntaxError, TypeError, ValueError) as exc:
+            skipped[str(exc)] += 1
+    return converted, sft_records, skipped
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Prepare verifier-compatible MBPP sanitized splits")
+    parser.add_argument("--dataset", default="google-research-datasets/mbpp")
+    parser.add_argument("--config", default="sanitized")
+    parser.add_argument("--splits", nargs="+", default=["train", "validation", "test"])
+    parser.add_argument("--output-dir", default="out/mbpp_sanitized")
+    args = parser.parse_args()
+
+    from datasets import load_dataset
+
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    summary = {}
+    for split in args.splits:
+        records = load_dataset(args.dataset, args.config, split=split)
+        converted, sft_records, skipped = convert_split(records)
+        output_path = output_dir / f"{split}.jsonl"
+        with output_path.open("w", encoding="utf-8") as handle:
+            for record in converted:
+                handle.write(json.dumps(record, ensure_ascii=False) + "\n")
+        sft_output = None
+        if split == "train":
+            sft_output = output_dir / "train_sft.jsonl"
+            with sft_output.open("w", encoding="utf-8") as handle:
+                for record in sft_records:
+                    handle.write(json.dumps(record, ensure_ascii=False) + "\n")
+        summary[split] = {
+            "source_count": len(records),
+            "converted_count": len(converted),
+            "skipped_count": len(records) - len(converted),
+            "top_skip_reasons": skipped.most_common(10),
+            "output": str(output_path),
+            "sft_output": str(sft_output) if sft_output else None,
+        }
+        print(f"{split}: converted {len(converted)}/{len(records)} -> {output_path}")
+    (output_dir / "summary.json").write_text(
+        json.dumps(summary, ensure_ascii=False, indent=2) + "\n",
+        encoding="utf-8",
+    )
+
+
+if __name__ == "__main__":
+    main()

scripts/serve_openai_api.py

@@ -56,6 +56,7 @@ class ChatRequest(BaseModel):
     stream: bool = True
     tools: list = Field(default_factory=list)
     open_thinking: bool = False
+    seed: int | None = None
     chat_template_kwargs: dict = None
     
     def get_open_thinking(self) -> bool:
@@ -102,7 +103,7 @@ def parse_response(text):
     return text.strip(), reasoning_content, tool_calls or None
 
 
-def generate_stream_response(messages, temperature, top_p, max_tokens, tools=None, open_thinking=False):
+def generate_stream_response(messages, temperature, top_p, max_tokens, tools=None, open_thinking=False, seed=None):
     try:
         new_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, tools=tools or None, open_thinking=open_thinking)
         inputs = tokenizer(new_prompt, return_tensors="pt", truncation=True).to(device)
@@ -112,6 +113,7 @@ def generate_stream_response(messages, temperature, top_p, max_tokens, tools=Non
 
         def _generate():
             try:
+                generator = torch.Generator(device=inputs.input_ids.device).manual_seed(seed) if seed is not None else None
                 model.generate(
                     inputs.input_ids,
                     max_new_tokens=max_tokens,
@@ -121,7 +123,8 @@ def generate_stream_response(messages, temperature, top_p, max_tokens, tools=Non
                     attention_mask=inputs.attention_mask,
                     pad_token_id=tokenizer.pad_token_id,
                     eos_token_id=tokenizer.eos_token_id,
-                    streamer=streamer
+                    streamer=streamer,
+                    generator=generator
                 )
             except Exception as e:
                 queue.put({"error": str(e)})
@@ -186,7 +189,8 @@ async def chat_completions(request: ChatRequest):
                     top_p=request.top_p,
                     max_tokens=request.max_tokens,
                     tools=request.tools,
-                    open_thinking=request.get_open_thinking()
+                    open_thinking=request.get_open_thinking(),
+                    seed=request.seed
                 )),
                 media_type="text/event-stream"
             )
@@ -200,6 +204,7 @@ async def chat_completions(request: ChatRequest):
             )
             inputs = tokenizer(new_prompt, return_tensors="pt", truncation=True).to(device)
             with torch.no_grad():
+                generator = torch.Generator(device=inputs["input_ids"].device).manual_seed(request.seed) if request.seed is not None else None
                 generated_ids = model.generate(
                     inputs["input_ids"],
                     max_length=inputs["input_ids"].shape[1] + request.max_tokens,
@@ -208,7 +213,8 @@ async def chat_completions(request: ChatRequest):
                     pad_token_id=tokenizer.pad_token_id,
                     eos_token_id=tokenizer.eos_token_id,
                     top_p=request.top_p,
-                    temperature=request.temperature
+                    temperature=request.temperature,
+                    generator=generator
                 )
                 answer = tokenizer.decode(generated_ids[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
             content, reasoning_content, tool_calls = parse_response(answer)

tests/test_code_backend.py

@@ -25,12 +25,17 @@ def test_openai_backend_builds_non_streaming_request():
         temperature=0.2,
         max_tokens=128,
         open_thinking=True,
+        seed=17,
     )
 
     assert generator("solve it").startswith("```python")
     assert completions.kwargs["messages"][0]["content"] == "solve it"
     assert completions.kwargs["stream"] is False
     assert completions.kwargs["extra_body"] == {"open_thinking": True}
+    assert completions.kwargs["seed"] == 17
+
+    generator("solve another")
+    assert completions.kwargs["seed"] == 18
 
 
 def test_openai_backend_rejects_empty_completion():

tests/test_generate_codegen.py

@@ -0,0 +1,37 @@
+from types import SimpleNamespace
+
+import pytest
+
+from scripts.generate_codegen import build_prompt, generate_predictions, read_completed_task_ids
+
+
+def test_generate_predictions_groups_samples_by_task():
+    task = SimpleNamespace(task_id="add", entry_point="add", prompt="Implement add(a, b).")
+    prompts = []
+
+    def generate(prompt):
+        prompts.append(prompt)
+        return f"candidate-{len(prompts)}"
+
+    records = generate_predictions([task], generate, samples_per_task=3)
+
+    assert records == [{"task_id": "add", "candidates": ["candidate-1", "candidate-2", "candidate-3"]}]
+    assert all("Required entry point: add" in prompt for prompt in prompts)
+    assert "exactly one ```python``` code block" in build_prompt(task)
+
+
+def test_generate_predictions_rejects_empty_sample_count():
+    with pytest.raises(ValueError, match="at least 1"):
+        generate_predictions([], lambda prompt: prompt, samples_per_task=0)
+
+
+def test_read_completed_task_ids_supports_resume(tmp_path):
+    output = tmp_path / "predictions.jsonl"
+    output.write_text(
+        '{"task_id":"one","candidates":["a"]}\n'
+        '{"task_id":"two","candidates":["b"]}\n',
+        encoding="utf-8",
+    )
+
+    assert read_completed_task_ids(output) == {"one", "two"}
+    assert read_completed_task_ids(tmp_path / "missing.jsonl") == set()

tests/test_prepare_mbpp.py

@@ -0,0 +1,50 @@
+import pytest
+
+from scripts.prepare_mbpp import build_sft_record, convert_record, convert_split, parse_assertion
+
+
+def test_parse_mbpp_equality_and_boolean_assertions():
+    assert parse_assertion("assert add(1, 2) == 3") == (
+        "add",
+        {"args": [1, 2], "kwargs": {}, "expected": 3},
+    )
+    assert parse_assertion("assert is_even(4)")[-1]["expected"] is True
+    assert parse_assertion("assert not is_even(3)")[-1]["expected"] is False
+
+
+def test_convert_record_validates_reference_solution():
+    record = {
+        "task_id": 1,
+        "prompt": "Write a function that adds two integers.",
+        "code": "def add(a, b):\n    return a + b",
+        "test_list": ["assert add(1, 2) == 3", "assert add(-1, 1) == 0"],
+    }
+
+    converted = convert_record(record)
+
+    assert converted["task_id"] == "mbpp-1"
+    assert converted["entry_point"] == "add"
+    assert len(converted["tests"]) == 2
+
+    sft = build_sft_record(record, converted)
+    assert "Required entry point: add" in sft["conversations"][0]["content"]
+    assert "def add(a, b)" in sft["conversations"][1]["content"]
+
+
+def test_convert_split_keeps_tasks_and_sft_records_aligned():
+    record = {
+        "task_id": 2,
+        "prompt": "Return the input integer.",
+        "code": "def identity(value):\n    return value",
+        "test_list": ["assert identity(7) == 7"],
+    }
+
+    tasks, sft_records, skipped = convert_split([record])
+
+    assert len(tasks) == len(sft_records) == 1
+    assert not skipped
+
+
+def test_parse_mbpp_rejects_nonliteral_expectation():
+    with pytest.raises(ValueError):
+        parse_assertion("assert add(1, 2) == expected")

trainer/train_grpo.py

@@ -168,7 +168,7 @@ def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_mod
             current_aux_loss = aux_loss.item()
             avg_reward_val = rewards.mean().item()
             avg_len_val = completion_mask.sum(dim=1).float().mean().item()
-            kl_ref_val = ((ref_per_token_logps - per_token_logps) * completion_mask).sum().item() / max(completion_mask.sum().item(), 1)
+            kl_ref_val = (per_token_kl * completion_mask).sum().item() / max(completion_mask.sum().item(), 1)
             advantages_mean_val = advantages.mean().item()
             advantages_std_val = advantages.std().item()
             current_lr = optimizer.param_groups[0]['lr']