Spaces:
Running
Running
| """Submission evaluation: static checks in-process, then worker.py in a subprocess on a leased GPU.""" | |
| import hashlib | |
| import json | |
| import os | |
| import queue | |
| import signal | |
| import subprocess | |
| import sys | |
| import tempfile | |
| import threading | |
| import time | |
| from functools import lru_cache | |
| from pathlib import Path | |
| from typing import Dict, List, Optional, Sequence | |
| from .roofline import describe, roofline | |
| from .static_checks import check_source | |
| WORKER = Path(__file__).with_name("worker.py") | |
| # t_new below this fraction of the roofline floor is physically impossible -> rejected as a timing hack | |
| ROOFLINE_REJECT = float(os.environ.get("KERNEL_ENV_ROOFLINE_REJECT", "0.8")) | |
| def apply_roofline(res: Dict) -> Dict: | |
| timed = res.get("stage") == "timed" | |
| r = roofline(res.get("flops"), res.get("bytes"), res.get("dtype"), res.get("gpu"), | |
| res.get("t_ref_ms"), res.get("t_new_ms") if timed else None) | |
| if not r: | |
| return res | |
| res = {**res, **r} | |
| if timed and res["t_new_ms"] < ROOFLINE_REJECT * r["t_min_ms"]: | |
| res.update( | |
| stage="correctness", correct=False, roofline_violation=True, | |
| feedback=f"Rejected: measured {res['t_new_ms']:.4f} ms is below the hardware floor of {r['t_min_ms']:.4f} ms " | |
| f"for this much work on {res['gpu']}. No real kernel can do that; outputs were reused, work was " | |
| f"skipped, or launches were not synchronized.", | |
| ) | |
| elif timed: | |
| res["feedback"] = f"{res.get('feedback', '')}\n{describe(r, res['gpu'])}".strip() | |
| return res | |
| REF_KEYS = ("t_ref_ms", "flops", "bytes", "dtype") | |
| class Evaluator: | |
| """Thread-safe; one instance is shared by all sessions so the GPU pool is global. | |
| workers_per_gpu > 1 lets evals share a GPU for compile + correctness (CPU-heavy JIT overlaps | |
| other work); benchmarks still run alone on the GPU via the worker's exclusive flock. The | |
| reference timing + roofline counts are measured once per task and reused (assumes one GPU type | |
| per evaluator). | |
| """ | |
| def __init__( | |
| self, | |
| gpus: Optional[Sequence[str]] = None, | |
| timeout_s: Optional[float] = None, | |
| atol: float = 1e-2, | |
| rtol: float = 1e-2, | |
| warmup_ms: int = 25, | |
| rep_ms: int = 100, | |
| workers_per_gpu: Optional[int] = None, | |
| ): | |
| self.workers_per_gpu = workers_per_gpu or int(os.environ.get("KERNEL_ENV_WORKERS_PER_GPU", "1")) | |
| self._gpus: "queue.Queue[str]" = queue.Queue() | |
| for g in gpus if gpus is not None else os.environ.get("KERNEL_ENV_GPUS", "0").split(","): | |
| for _ in range(self.workers_per_gpu): | |
| self._gpus.put(g.strip()) | |
| self._ref_cache: Dict[str, Dict] = {} | |
| self._ref_lock = threading.Lock() | |
| # a hung kernel holds its GPU slot (and, in lockstep GRPO rollouts, the whole step) until this fires | |
| self.timeout_s = timeout_s or float(os.environ.get("KERNEL_ENV_EVAL_TIMEOUT", "300")) | |
| self.atol, self.rtol = atol, rtol | |
| self.warmup_ms, self.rep_ms = warmup_ms, rep_ms | |
| def evaluate(self, ref_src: str, code: str, language: str, seeds: List[int], task_id: str = "") -> Dict: | |
| """Result dict also carries gpu_id, queue_wait_s (waiting for a free GPU) and eval_s.""" | |
| t0 = time.monotonic() | |
| errors = check_source(code, language) | |
| if errors: | |
| return {"stage": "format", "feedback": "Rejected before running:\n- " + "\n- ".join(errors[:10]), | |
| "gpu_id": None, "queue_wait_s": 0.0, "eval_s": time.monotonic() - t0} | |
| key = hashlib.sha1(ref_src.encode()).hexdigest() | |
| with self._ref_lock: | |
| cached = self._ref_cache.get(key) | |
| gpu = self._gpus.get() | |
| t1 = time.monotonic() | |
| try: | |
| res = self._run(gpu, ref_src, code, language, seeds, cached, task_id=task_id) | |
| if self.workers_per_gpu > 1 and "OutOfMemoryError" in res.get("feedback", ""): | |
| res = self._run(gpu, ref_src, code, language, seeds, cached, exclusive=True, task_id=task_id) | |
| finally: | |
| self._gpus.put(gpu) | |
| if cached is None and res.get("t_ref_ms") is not None: | |
| with self._ref_lock: | |
| self._ref_cache.setdefault(key, {k: res[k] for k in REF_KEYS if res.get(k) is not None}) | |
| return apply_roofline({**res, "gpu_id": gpu, "queue_wait_s": t1 - t0, "eval_s": time.monotonic() - t1}) | |
| def _run(self, gpu: str, ref_src: str, code: str, language: str, seeds: List[int], | |
| ref_cache: Optional[Dict] = None, exclusive: bool = False, task_id: str = "") -> Dict: | |
| with tempfile.TemporaryDirectory(prefix="kernel_env_") as tmp: | |
| req, out, code_path = Path(tmp, "request.json"), Path(tmp, "result.json"), Path(tmp, "candidate.py") | |
| code_path.write_text(code) | |
| req.write_text(json.dumps({ | |
| "ref_src": ref_src, "code_path": str(code_path), "language": language, "seeds": seeds, "task_id": task_id, | |
| "atol": self.atol, "rtol": self.rtol, "warmup_ms": self.warmup_ms, "rep_ms": self.rep_ms, | |
| "ref_cache": ref_cache, "exclusive": exclusive, | |
| "gpu_lock": f"{tempfile.gettempdir()}/kernel_env_gpu{gpu}.lock" if self.workers_per_gpu > 1 else None, | |
| })) | |
| env = { | |
| **os.environ, | |
| "CUDA_VISIBLE_DEVICES": gpu, | |
| "HOME": tmp, | |
| "TRITON_CACHE_DIR": f"{tmp}/triton", | |
| "TORCH_EXTENSIONS_DIR": f"{tmp}/torch_ext", | |
| } | |
| # own process group so a timeout also kills compiler children (nvcc, ptxas) | |
| proc = subprocess.Popen( | |
| [sys.executable, "-I", str(WORKER), str(req), str(out)], | |
| cwd=tmp, env=env, start_new_session=True, | |
| stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True, | |
| ) | |
| try: | |
| _, stderr = proc.communicate(timeout=self.timeout_s) | |
| except subprocess.TimeoutExpired: | |
| os.killpg(proc.pid, signal.SIGKILL) | |
| proc.communicate() | |
| return {"stage": "timeout", "feedback": f"Evaluation exceeded {self.timeout_s:.0f}s (hang or very slow kernel)"} | |
| if not out.exists(): | |
| return {"stage": "runtime", "feedback": f"Worker crashed (exit {proc.returncode}):\n{stderr[-2000:]}"} | |
| return json.loads(out.read_text()) | |
| class MockEvaluator(Evaluator): | |
| """No-GPU stand-in for Mac dev/tests: real static checks, scripted (or default) results.""" | |
| DEFAULT = {"stage": "timed", "correct": True, "max_abs_err": 0.0, "t_ref_ms": 1.0, "t_new_ms": 1.0, | |
| "speedup": 1.0, "feedback": "[mock] Correct. speedup 1.00x"} | |
| def __init__(self, results: Optional[List[Dict]] = None): | |
| super().__init__(gpus=["mock"]) | |
| self.results = list(results or []) | |
| self.calls: List[Dict] = [] | |
| def _run(self, gpu, ref_src, code, language, seeds, ref_cache=None, exclusive=False, task_id=""): | |
| self.calls.append({"code": code, "language": language, "seeds": seeds, "ref_cache": ref_cache, "task_id": task_id}) | |
| return self.results.pop(0) if self.results else dict(self.DEFAULT) | |
| def default_evaluator() -> Evaluator: | |
| """Process-wide evaluator. KERNEL_ENV_EVALUATOR=mock for machines without a CUDA GPU.""" | |
| return MockEvaluator() if os.environ.get("KERNEL_ENV_EVALUATOR") == "mock" else Evaluator() | |