Spaces:
Running
Running
File size: 6,165 Bytes
e313085 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 | import json
from pathlib import Path
import pytest
from openenv.core.env_server.serialization import serialize_observation
from kernel_env import KernelAction, KernelEnv, KernelObservation
from kernel_env.server.evaluator import MockEvaluator
from kernel_env.server.kernel_env_environment import KernelEnvironment
from kernel_env.server.rubrics import kernel_rubric
from kernel_env.server.tasks import Task, load_tasks
from kernels import GOOD, with_forward
FIXTURES = Path(__file__).parent / "fixtures"
def timed(speedup):
return {**MockEvaluator.DEFAULT, "speedup": speedup}
@pytest.fixture
def tasks():
return load_tasks(str(FIXTURES), levels=(1,))
def test_load_tasks(tasks):
assert list(tasks) == ["level1/19_ReLU"]
assert tasks["level1/19_ReLU"].split in ("train", "eval")
def test_reset_returns_prompt(tasks):
env = KernelEnvironment(MockEvaluator(), tasks, max_turns=2)
obs = env.reset(task_id="level1/19_ReLU")
assert "class Model" in obs.problem and "@triton.jit" in obs.problem
assert obs.turn == 0 and not obs.done
assert env.state.task_id == "level1/19_ReLU"
def test_seeded_task_choice_is_deterministic():
many = {f"level1/{i}_t": Task(f"level1/{i}_t", 1, "") for i in range(50)}
env = KernelEnvironment(MockEvaluator(), many)
picks = {env.reset(seed=7).task_id for _ in range(3)}
assert len(picks) == 1 and many[picks.pop()].split == "train"
assert many[env.reset(seed=1, split="eval").task_id].split == "eval"
def test_episode_rewards_and_done(tasks):
ev = MockEvaluator([timed(2.0), {"stage": "correctness", "correct": False, "feedback": "mismatch"}])
env = KernelEnvironment(ev, tasks, max_turns=3)
env.reset(seed=0, task_id="level1/19_ReLU")
obs = env.step(KernelAction(code=GOOD))
assert obs.stage == "timed" and obs.reward == pytest.approx(1.3) and not obs.done
assert ev.calls[0]["task_id"] == "level1/19_ReLU"
obs = env.step(KernelAction(code=GOOD))
assert obs.stage == "correctness" and obs.reward == 0.0
obs = env.step(KernelAction(code=with_forward("return torch.relu(x)")))
assert obs.stage == "format" and obs.reward == 0.0 and "torch.relu" in obs.feedback
assert obs.done and len(ev.calls) == 2 # rejected code never reaches the GPU
assert env.state.best_speedup == 2.0
assert [h["stage"] for h in env.state.history] == ["timed", "correctness", "format"]
with pytest.raises(RuntimeError):
env.step(KernelAction(code=GOOD))
seeds = [c["seeds"] for c in ev.calls]
assert len(seeds[0]) == 5 and seeds[0] != seeds[1]
def test_suspicious_speedup_flagged(tasks):
env = KernelEnvironment(MockEvaluator([timed(500.0)]), tasks)
env.reset(task_id="level1/19_ReLU")
env.step(KernelAction(code=GOOD))
assert env.state.history[0].get("suspicious")
@pytest.mark.parametrize(
"correct, speedup, reward",
[(False, 10.0, 0.0), (True, 1.0, 0.3), (True, 0.5, 0.1), (True, 4.0, 2.3), (True, 1000.0, 3.3)],
)
def test_rubric(correct, speedup, reward):
obs = KernelObservation(correct=correct, speedup=speedup)
assert kernel_rubric()(None, obs) == pytest.approx(reward)
def test_client_parses_server_payload():
obs = KernelObservation(task_id="t", stage="timed", correct=True, speedup=1.5, turn=1, done=True, reward=0.9)
result = KernelEnv._parse_result(None, serialize_observation(obs))
assert result.observation == obs and result.reward == 0.9 and result.done
def test_eval_log_lines(tasks, tmp_path, monkeypatch):
log = tmp_path / "evals.jsonl"
monkeypatch.setenv("KERNEL_ENV_EVAL_LOG", str(log))
env = KernelEnvironment(MockEvaluator([timed(2.0)]), tasks)
env.reset(seed=0, task_id="level1/19_ReLU")
env.step(KernelAction(code=GOOD))
env.step(KernelAction(code=with_forward("return torch.relu(x)")))
recs = [json.loads(line) for line in log.read_text().splitlines()]
assert [r["event"] for r in recs] == ["reset", "step", "step"]
assert {r["episode_id"] for r in recs} == {env.state.episode_id}
assert recs[1]["stage"] == "timed" and recs[1]["gpu_id"] == "mock" and recs[1]["eval_s"] >= 0
assert recs[2]["stage"] == "format" and recs[2]["queue_wait_s"] == 0.0
assert env.state.history[0]["code_sha"] == recs[1]["code_sha"]
def test_reference_timing_cached_per_task():
ev = MockEvaluator([timed(2.0), timed(3.0), timed(1.0)])
ev.evaluate("ref A", GOOD, "triton", [1])
ev.evaluate("ref A", GOOD, "triton", [2])
ev.evaluate("ref B", GOOD, "triton", [3])
assert ev.calls[0]["ref_cache"] is None
assert ev.calls[1]["ref_cache"]["t_ref_ms"] == MockEvaluator.DEFAULT["t_ref_ms"]
assert ev.calls[2]["ref_cache"] is None # different task
def test_worker_request_carries_task_id(monkeypatch):
from kernel_env.server.evaluator import Evaluator
import kernel_env.server.evaluator as evaluator
def process(args, **kwargs):
# Inspect transport only; do not spawn or execute the candidate.
request = json.loads(Path(args[-2]).read_text())
assert request["task_id"] == "level1/23_Softmax"
Path(args[-1]).write_text(json.dumps({"stage": "correctness", "correct": False}))
class Completed:
returncode = 0
def communicate(self, timeout=None):
return "", ""
return Completed()
monkeypatch.setattr(evaluator.subprocess, "Popen", process)
result = Evaluator(gpus=["0"])._run("0", "reference", GOOD, "triton", [1], task_id="level1/23_Softmax")
assert result["stage"] == "correctness"
@pytest.mark.parametrize(
"stage, correct, speedup, reward",
[("format", False, None, 0.0), ("compile", False, None, 0.0), ("runtime", False, None, 0.02),
("timeout", False, None, 0.02), ("correctness", False, None, 0.05), ("timed", True, 0.5, 0.1)],
)
def test_partial_credit(stage, correct, speedup, reward):
obs = KernelObservation(stage=stage, correct=correct, speedup=speedup)
assert kernel_rubric("speedup", partial=True)(None, obs) == pytest.approx(reward)
assert kernel_rubric("speedup", partial=False)(None, obs) == pytest.approx(reward if correct else 0.0)
|