Spaces:
Running
Running
| # Copyright (c) Meta Platforms, Inc. and affiliates. | |
| # All rights reserved. | |
| # | |
| # This source code is licensed under the BSD-style license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| """ | |
| Kernel Env Environment Implementation. | |
| reset() picks a KernelBench task and returns the prompt; each step() evaluates one ModelNew | |
| submission (static checks -> compile -> correctness on hidden seeds -> timing) and returns | |
| feedback the agent can use to refine. Episode ends after `max_turns` submissions. | |
| """ | |
| import asyncio | |
| import hashlib | |
| import json | |
| import logging | |
| import os | |
| import random | |
| import threading | |
| import time | |
| from typing import Any, Dict, Optional | |
| from uuid import uuid4 | |
| from openenv.core.env_server.interfaces import Environment | |
| try: | |
| from ..models import KernelAction, KernelObservation, KernelState | |
| from .evaluator import Evaluator, default_evaluator | |
| from .rubrics import kernel_rubric | |
| from .tasks import Task, build_prompt, load_tasks | |
| except ImportError: | |
| from models import KernelAction, KernelObservation, KernelState | |
| from server.evaluator import Evaluator, default_evaluator | |
| from server.rubrics import kernel_rubric | |
| from server.tasks import Task, build_prompt, load_tasks | |
| log = logging.getLogger(__name__) | |
| N_TRIALS = 5 | |
| SUSPICIOUS_SPEEDUP = 50.0 # flagged in state.history for manual review of possible reward hacks | |
| RESULT_FIELDS = ("stage", "feedback", "correct", "max_abs_err", "speedup", "t_ref_ms", "t_new_ms", | |
| "t_min_ms", "bound", "arithmetic_intensity", "roofline_frac", "ref_roofline_frac") | |
| _log_lock = threading.Lock() | |
| def emit(event: str, **fields: Any) -> None: | |
| """One JSON line per event on stdout (Modal / docker logs); also appended to $KERNEL_ENV_EVAL_LOG if set.""" | |
| line = json.dumps({"event": event, "ts": round(time.time(), 3), **fields}, default=str) | |
| print(line, flush=True) | |
| path = os.environ.get("KERNEL_ENV_EVAL_LOG") | |
| if path: | |
| with _log_lock, open(path, "a") as f: | |
| f.write(line + "\n") | |
| class KernelEnvironment(Environment[KernelAction, KernelObservation, KernelState]): | |
| """Multi-turn GPU kernel optimization over KernelBench tasks.""" | |
| # Episode state is per instance; GPU access is serialized by the shared Evaluator. | |
| SUPPORTS_CONCURRENT_SESSIONS: bool = True | |
| def __init__( | |
| self, | |
| evaluator: Optional[Evaluator] = None, | |
| tasks: Optional[Dict[str, Task]] = None, | |
| max_turns: Optional[int] = None, | |
| ): | |
| super().__init__(rubric=kernel_rubric()) | |
| self._evaluator = evaluator or default_evaluator() | |
| self._tasks = tasks if tasks is not None else load_tasks() | |
| self.max_turns = max_turns or int(os.environ.get("KERNEL_ENV_MAX_TURNS", "4")) | |
| self._state = KernelState(episode_id=str(uuid4())) | |
| self._task: Optional[Task] = None | |
| self._rng = random.Random() | |
| def reset( | |
| self, | |
| seed: Optional[int] = None, | |
| episode_id: Optional[str] = None, | |
| task_id: Optional[str] = None, | |
| split: str = "train", | |
| language: str = "triton", | |
| **kwargs: Any, | |
| ) -> KernelObservation: | |
| """Start an episode on `task_id`, or a seeded random task from `split` ("train"/"eval").""" | |
| self._reset_rubric() | |
| self._rng = random.Random(seed) | |
| if task_id is None: | |
| pool = sorted(t for t, task in self._tasks.items() if task.split == split) | |
| if not pool: | |
| raise ValueError(f"no tasks in split {split!r}") | |
| task_id = self._rng.choice(pool) | |
| if task_id not in self._tasks: | |
| raise ValueError(f"unknown task_id {task_id!r}") | |
| self._task = self._tasks[task_id] | |
| self._state = KernelState( | |
| episode_id=episode_id or str(uuid4()), task_id=task_id, level=self._task.level | |
| ) | |
| emit("reset", episode_id=self._state.episode_id, task_id=task_id, split=split, seed=seed) | |
| return KernelObservation( | |
| task_id=task_id, | |
| problem=build_prompt(self._task, language, self.max_turns), | |
| max_turns=self.max_turns, | |
| done=False, | |
| reward=0.0, | |
| ) | |
| def step(self, action: KernelAction, timeout_s: Optional[float] = None, **kwargs: Any) -> KernelObservation: # type: ignore[override] | |
| if self._task is None: | |
| raise RuntimeError("call reset() before step()") | |
| if self._state.step_count >= self.max_turns: | |
| raise RuntimeError("episode is done; call reset()") | |
| self._state.step_count += 1 | |
| # hidden per-step input seeds: outputs can't be memorized across turns | |
| seeds = [self._rng.randrange(2**31) for _ in range(N_TRIALS)] | |
| res = self._evaluator.evaluate(self._task.src, action.code, action.language, seeds, task_id=self._task.task_id) | |
| obs = KernelObservation( | |
| task_id=self._task.task_id, | |
| turn=self._state.step_count, | |
| max_turns=self.max_turns, | |
| done=self._state.step_count >= self.max_turns, | |
| **{k: res[k] for k in RESULT_FIELDS if res.get(k) is not None}, | |
| ) | |
| obs.reward = self._apply_rubric(action, obs) | |
| entry = { | |
| "turn": obs.turn, "stage": obs.stage, "speedup": obs.speedup, "reward": obs.reward, | |
| "roofline_frac": obs.roofline_frac, "roofline_violation": bool(res.get("roofline_violation")), | |
| "gpu_id": res.get("gpu_id"), "queue_wait_s": res.get("queue_wait_s"), "eval_s": res.get("eval_s"), | |
| "code_sha": hashlib.sha1(action.code.encode()).hexdigest()[:12], | |
| "phases": res.get("phases"), # worker seconds per phase: inputs, JIT, correctness, bench, lock wait | |
| } | |
| if obs.correct and (obs.speedup or 0) > SUSPICIOUS_SPEEDUP: | |
| entry["suspicious"] = True | |
| log.warning("suspicious speedup %.1fx on %s (episode %s)", obs.speedup, obs.task_id, self._state.episode_id) | |
| self._state.history.append(entry) | |
| emit( | |
| "step", episode_id=self._state.episode_id, task_id=obs.task_id, **entry, | |
| t_ref_ms=obs.t_ref_ms, t_new_ms=obs.t_new_ms, t_min_ms=obs.t_min_ms, bound=obs.bound, | |
| ref_roofline_frac=obs.ref_roofline_frac, max_abs_err=obs.max_abs_err, | |
| gpu=res.get("gpu"), feedback_head=obs.feedback[:200], | |
| ) | |
| if obs.correct: | |
| self._state.best_speedup = max(self._state.best_speedup, obs.speedup or 0.0) | |
| return obs | |
| async def step_async(self, action: KernelAction, timeout_s: Optional[float] = None, **kwargs: Any) -> KernelObservation: # type: ignore[override] | |
| # evaluation blocks for seconds; keep the server event loop free for other sessions | |
| return await asyncio.to_thread(self.step, action, timeout_s, **kwargs) | |
| def state(self) -> KernelState: | |
| return self._state | |