Spaces:
Running
Running
| """Mac-runnable checks. Rollout test needs a mock server: | |
| cd ../kernel_env && KERNEL_ENV_EVALUATOR=mock KERNELBENCH_DIR=<dir with level1/*.py> \ | |
| .venv/bin/uvicorn server.app:app --port 8765 | |
| KERNEL_ENV_URL=http://127.0.0.1:8765 KERNEL_ENV_TASK=level1/19_ReLU .venv/bin/python -m pytest test_train.py | |
| """ | |
| import os | |
| import sys | |
| from collections import defaultdict | |
| from pathlib import Path | |
| from types import SimpleNamespace | |
| import pytest | |
| from baseline_eval import extract_code, run_episodes, run_episodes_pipelined, summarize | |
| from grpo_kernel import KernelArgs, episode_return, feedback_ids, make_rollout_func | |
| sys.path.insert(0, str(Path(__file__).parent.parent / "kernel_env")) | |
| from tests.kernels import GOOD # noqa: E402 (kernel_env's fixture kernel) | |
| URL = os.environ.get("KERNEL_ENV_URL") | |
| TASK = os.environ.get("KERNEL_ENV_TASK", "level1/19_ReLU") | |
| TOKENIZER = "Qwen/Qwen2.5-0.5B-Instruct" | |
| BAD = "class ModelNew:\n pass" # no kernel -> rejected by static checks, reward 0 | |
| def test_extract_code(): | |
| assert extract_code("hi\n```python\na = 1\n```\nbye") == "a = 1" | |
| assert extract_code("```py\nx\n```\n```python\nlast\n```") == "last" | |
| assert extract_code("```\nplain\n```") == "plain" | |
| assert extract_code("```python\nb = 2\n```\n```text\nlog\n```") == "b = 2" # python beats later non-python | |
| assert extract_code("```python\ncut = 1\nmore") == "cut = 1\nmore" # truncated fence | |
| assert extract_code(" import torch\n") == "import torch" # bare source, as the prompt asks | |
| def test_episode_return(): | |
| assert episode_return([0.0, 0.3, 1.3], "discounted", 0.5) == pytest.approx(0.5 * 0.3 + 0.25 * 1.3) | |
| assert episode_return([0.0, 1.3, 0.3], "best", 0.5) == 1.3 | |
| assert episode_return([], "best", 1.0) == 0.0 | |
| def test_summarize(): | |
| eps = [{"turns": [{"stage": "format", "correct": False, "speedup": 0.0}, | |
| {"stage": "timed", "correct": True, "speedup": 2.0}]}, | |
| {"turns": [{"stage": "compile", "correct": False, "speedup": 0.0}]}] | |
| s = summarize(eps) | |
| assert (s["compile_pct"], s["correct_pct"], s["correct_turn1_pct"], s["median_speedup"]) == (50, 50, 0, 2.0) | |
| assert s["stage_hist_furthest"] == {"timed": 1, "compile": 1} | |
| def _tok(): | |
| from transformers import AutoTokenizer | |
| return AutoTokenizer.from_pretrained(TOKENIZER) | |
| def test_feedback_ids_stitch_like_chat_template(): | |
| tok = _tok() | |
| conv = [{"role": "user", "content": "P"}, {"role": "assistant", "content": "A1"}, {"role": "user", "content": "F"}] | |
| want = tok.apply_chat_template(conv, add_generation_prompt=True, tokenize=True, return_dict=False) | |
| p = tok.apply_chat_template(conv[:1], add_generation_prompt=True, tokenize=True, return_dict=False) | |
| got = p + tok.encode("A1") + [tok.eos_token_id] + feedback_ids(tok, "F") | |
| assert got == want | |
| class FakeTrainer: | |
| """Stands in for GRPOTrainer: replies BAD on turn 1 and GOOD (fenced) afterwards.""" | |
| def __init__(self, tok): | |
| self.processing_class, self.tok, self.calls = tok, tok, 0 | |
| self.model = SimpleNamespace(training=True) | |
| self.state = SimpleNamespace(global_step=0) | |
| self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)} | |
| def _generate_single_turn(self, prompt_ids, images, mm, n): | |
| self.calls += 1 | |
| text = BAD if self.calls == 1 else f"Here:\n```python\n{GOOD}\n```" | |
| ids = self.tok.encode(text) + [self.tok.eos_token_id] | |
| return [ids for _ in prompt_ids], None | |
| def test_rollout_func_against_mock_server(tmp_path): | |
| tok, n = _tok(), 3 | |
| trainer = FakeTrainer(tok) | |
| fn = make_rollout_func(KernelArgs(env_url=URL, gamma=0.5), str(tmp_path / "s.jsonl")) | |
| out = fn([TASK] * n, trainer) | |
| assert out["logprobs"] is None and len(out["prompt_ids"]) == n | |
| max_turns = trainer.calls # one generate call per turn, all episodes in lockstep | |
| for p, c, m in zip(out["prompt_ids"], out["completion_ids"], out["env_mask"]): | |
| assert len(c) == len(m) | |
| # stitched token stream == chat template of the same conversation | |
| text = tok.decode(p + c) | |
| assert text.count("<|im_start|>assistant") == max_turns and "stage=format" in text | |
| model_text = tok.decode([t for t, k in zip(c, m) if k]) | |
| assert BAD in model_text and "stage=" not in model_text # feedback tokens are masked out | |
| # mock evaluator: turn 1 rejected (0), later turns speedup 1.0 -> 0.3 each | |
| want = sum(0.5**t * 0.3 for t in range(1, max_turns)) | |
| assert out["episode_return"] == pytest.approx([want] * n) | |
| assert trainer._metrics["train"]["kernel/correct_pct"] == [100.0] | |
| # every episode logged for report.py | |
| import json | |
| from report import render | |
| logged = [json.loads(line) for line in (tmp_path / "rollouts-rank0.jsonl").read_text().splitlines()] | |
| assert len(logged) == n and all(len(ep["completions"]) == len(ep["turns"]) == max_turns for ep in logged) | |
| assert "Turn by turn" in render(logged, "rollouts") | |
| def test_run_episodes_against_mock_server(): | |
| eps = run_episodes(lambda convs: [f"```python\n{GOOD}\n```"] * len(convs), URL, [TASK, TASK]) | |
| assert all(t["correct"] for ep in eps for t in ep["turns"]) | |
| assert summarize(eps)["correct_turn1_pct"] == 100 | |
| def test_pipelined_episodes_against_mock_server(): | |
| import threading | |
| seen_feedback, lock = [], threading.Lock() | |
| def chat(conv): | |
| with lock: | |
| seen_feedback.append(conv[-1]["content"]) | |
| if len(conv) == 1: # first turn: cut off mid-module | |
| return "```python\nimport torch\n# reasoning...", "length", None, 8192 | |
| return f"```python\n{GOOD}\n```", "stop", None, 500 | |
| done = [] | |
| eps = run_episodes_pipelined(chat, URL, [TASK, TASK, "level1/does_not_exist"], on_done=done.append) | |
| assert len(done) == 3 and "error" in eps[2] and not eps[2]["turns"] | |
| for ep in eps[:2]: | |
| assert ep["turns"][0]["truncated"] and ep["turns"][0]["stage"] == "format" | |
| assert all(t["correct"] and not t["truncated"] for t in ep["turns"][1:]) | |
| assert ep["episode_id"] and len(ep["history"]) == len(ep["turns"]) | |
| assert sum("cut off before the module was complete" in f for f in seen_feedback) == 2 | |
| def test_feedback_ids_qwen35_template(): | |
| # Qwen3.5 is not prefix-stable (earlier assistant turns lose their think block); suffix must still be exact | |
| transformers = pytest.importorskip("transformers") | |
| try: | |
| tok = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen3.5-9B") | |
| except OSError: | |
| pytest.skip("Qwen3.5 tokenizer not available offline") | |
| from grpo_kernel import feedback_ids | |
| assert tok.decode(feedback_ids(tok, "fb")) == "\n<|im_start|>user\nfb<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n" | |
| def test_rollout_unwraps_processor(tmp_path): | |
| # VLM checkpoints (Qwen3.5) hand the rollout a Processor whose chat template returns nested lists | |
| tok = _tok() | |
| trainer = FakeTrainer(tok) | |
| trainer.processing_class = SimpleNamespace(tokenizer=tok) | |
| out = make_rollout_func(KernelArgs(env_url=URL), str(tmp_path / "s.jsonl"))([TASK] * 2, trainer) | |
| assert all(isinstance(t, int) for p in out["prompt_ids"] for t in p) | |
| assert all(isinstance(t, int) for c in out["completion_ids"] for t in c) | |