"""Dory coding evaluation. Deploy this APP, then use scripts/spawn_code_eval.py. Only this app is deployed. Generated programs are executed solely in the Modal CPU grading containers, never by the local collection/summary scripts. """ import hashlib import json import os import re import shutil import time from pathlib import Path import modal APP = modal.App("beam-dory-code") weights = modal.Volume.from_name("beam-weights", version=2) ckpt = modal.Volume.from_name("beam-dory-ckpt", version=2) data = modal.Volume.from_name("beam-dory-data", version=2) ROOT = Path("/data/eval/code") VLLM_PR_SHA = "ce9dd0ed884b7cb4014396c1aa715a8b7c9fc4a4" CONFIGS = { "base": ("base", ""), "rE-t050-default": ("merges/rE-t050", ""), "rE-t050-low": ("merges/rE-t050", "low"), "rE-t050-high": ("merges/rE-t050", "high"), "rA-t030-default": ("merges/rA-t030", ""), # coding fix (run rC: rE + Beam TACO solutions with effort replay), blended 50/50 with base "rC-t050-default": ("merges/rC-t050", ""), "rC-t050-low": ("merges/rC-t050", "low"), "rC-t050-high": ("merges/rC-t050", "high"), } INSTRUCTION = ( "\n\nReturn the complete function, including any needed imports and helpers, " "in a single ```python code block." ) DEPS = ( "wget==3.2", "tempdir==0.7.1", "multipledispatch==1.0.0", "appdirs==1.4.4", "numpy>=1.26,<3", "tqdm>=4.66", "termcolor>=2", "fire>=0.6", "tree-sitter==0.25.2", "tree-sitter-python==0.25.0", "rich>=13", "datasets>=2.21", "psutil>=5.9", ) def cache_datasets(): """Run at image build time, so network-disabled grading needs no downloads.""" from evalplus.data import get_human_eval_plus, get_mbpp_plus h, m = get_human_eval_plus(), get_mbpp_plus() assert len(h) == 164 and len(m) == 378 print("Cached EvalPlus datasets:", len(h), len(m), flush=True) # The CUDA/vLLM layers match eval_dory.py, including Dory support from PR #60149. # --no-deps prevents EvalPlus from replacing vLLM's torch/transformers versions. image = ( modal.Image.from_registry("nvidia/cuda:13.0.3-devel-ubuntu24.04", add_python="3.12") .apt_install("git") .env({"VLLM_USE_PRECOMPILED": "1"}) .pip_install(f"vllm @ git+https://github.com/vllm-project/vllm.git@{VLLM_PR_SHA}", "math-verify[antlr4_13_2]") .pip_install("langdetect", "immutabledict", "nltk>=3.9.1") .env({"PYTHONUNBUFFERED": "1", "TOKENIZERS_PARALLELISM": "false", "FLASHINFER_WORKSPACE_BASE": "/data/cache/flashinfer"}) .pip_install("evalplus==0.3.1", extra_options="--no-deps") .pip_install(*DEPS) .run_function(cache_datasets) .env({"OMP_NUM_THREADS": "1", "OPENBLAS_NUM_THREADS": "1", "MKL_NUM_THREADS": "1"}) ) grade_image = ( modal.Image.debian_slim(python_version="3.12") .pip_install("evalplus==0.3.1", extra_options="--no-deps") .pip_install(*DEPS) .env({"PYTHONUNBUFFERED": "1", "OMP_NUM_THREADS": "1", "OPENBLAS_NUM_THREADS": "1", "MKL_NUM_THREADS": "1"}) .run_function(cache_datasets) ) def get_problems(limit=0): from evalplus.data import ( get_human_eval_plus, get_human_eval_plus_hash, get_mbpp_plus, get_mbpp_plus_hash, ) rows = [] hashes = {} for bench, getter, hasher in ( ("humaneval", get_human_eval_plus, get_human_eval_plus_hash), ("mbpp", get_mbpp_plus, get_mbpp_plus_hash), ): problems = sorted(getter().values(), key=lambda r: int(r["task_id"].split("/")[-1])) hashes[bench] = hasher() for problem in problems[:limit or None]: rows.append({"bench": bench, **problem}) return rows, hashes def dump_json(path, obj): path = Path(path) temp = path.with_suffix(path.suffix + ".tmp") temp.write_text(json.dumps(obj, indent=2) + "\n") temp.replace(path) def load_llm(model): from vllm import LLM path = Path("/weights/dory-4b") if model == "base" else Path("/ckpt") / model tok_dir = Path("/tmp/dory-code-tok") tok_dir.mkdir(exist_ok=True) for name in ("tokenizer.json", "tokenizer_config.json", "special_tokens_map.json", "chat_template.jinja"): if (path / name).exists(): shutil.copy(path / name, tok_dir / name) return LLM( model=str(path), tokenizer=str(tok_dir), max_model_len=69632, gpu_memory_utilization=0.92, max_num_seqs=256, enable_prefix_caching=True, hf_overrides={"n_recurrent_loops": 2, "recurrent_kv_cache_mode": "per_loop"}, ) class BudgetExpired(Exception): pass @APP.function( image=image, gpu="H100", cpu=8, memory=65536, volumes={"/weights": weights, "/ckpt": ckpt, "/data": data}, timeout=3000, startup_timeout=600, retries=0, max_containers=5, single_use_containers=True, scaledown_window=2, ) def evaluate(config: str, run: str, limit: int = 0, runtime_s: int = 2700): """One H100 per config. Checkpoint finished problems during continuous batching. runtime_s includes model loading. It never changes the per-sample token cap. An incomplete run is explicitly marked incomplete, not scored as full data. Grading is spawned before returning, so it also survives local disconnects. """ from vllm import SamplingParams import importlib.metadata assert config in CONFIGS assert re.fullmatch(r"[A-Za-z0-9_-]+", run) assert 30 <= runtime_s <= 2700 start = time.monotonic() model, effort = CONFIGS[config] out_dir = ROOT / run / config out_dir.mkdir(parents=True, exist_ok=True) if (out_dir / "manifest.json").exists(): raise FileExistsError(f"Refusing to overwrite {out_dir}") rows, hashes = get_problems(limit) kwargs = {"enable_thinking": True, "medium_effort": False} if effort: kwargs["reasoning_effort"] = effort manifest = { "config": config, "model": model, "reasoning_effort": effort or None, "loops": 2, "run": run, "limit": limit, "n_problems": len(rows), "samples_per_problem": 4, "temperature": 1.0, "top_p": 0.95, "max_tokens": 65536, "max_model_len": 69632, "seed": 20261030, "seed_scheme": "20261030 + problem_index, shared across configurations", "dataset_hashes": hashes, "evalplus_version": importlib.metadata.version("evalplus"), "vllm_version": importlib.metadata.version("vllm"), "vllm_pr_sha": VLLM_PR_SHA, "chat_template_kwargs": kwargs, "instruction": INSTRUCTION, "runtime_cap_s": runtime_s, "started_unix": time.time(), "state": "loading", } dump_json(out_dir / "manifest.json", manifest) with (out_dir / "prompts.jsonl").open("w") as handle: for row in rows: prompt = row["prompt"] + INSTRUCTION handle.write(json.dumps({"task_id": row["task_id"], "bench": row["bench"], "entry_point": row["entry_point"], "prompt": prompt, "sha256": hashlib.sha256(prompt.encode()).hexdigest()}) + "\n") data.commit() llm = load_llm(model) manifest["load_s"] = time.monotonic() - start manifest["state"] = "generating" dump_json(out_dir / "manifest.json", manifest) data.commit() records = [] seen = set() last_commit = time.monotonic() original_step = llm.llm_engine.step # PR-pinned vLLM's _run_engine consumes step() results in request-ID order. # Observe final RequestOutputs without altering scheduling or sampling. This # preserves completed generations even if a runtime budget stops the run. with (out_dir / "generations.jsonl").open("w", buffering=1) as handle: def checkpoint_step(*args, **kw): nonlocal last_commit if time.monotonic() - start > runtime_s: raise BudgetExpired("Runtime budget reached") outputs = original_step(*args, **kw) for output in outputs: if not output.finished or output.request_id in seen: continue row = rows[int(output.request_id)] assert len(output.outputs) == 4 for candidate in output.outputs: record = { "config": config, "bench": row["bench"], "task_id": row["task_id"], "sample_index": candidate.index, "tokens": len(candidate.token_ids), "finish_reason": candidate.finish_reason, "stop_reason": candidate.stop_reason, "thinking_closed": "" in candidate.text, "text": candidate.text, } records.append(record) handle.write(json.dumps(record, ensure_ascii=False) + "\n") seen.add(output.request_id) if time.monotonic() - last_commit >= 30: handle.flush() dump_json(out_dir / "progress.json", { "completed_problems": len(seen), "total_problems": len(rows), "tokens_completed": sum(r["tokens"] for r in records), "elapsed_s": time.monotonic() - start, }) data.commit() last_commit = time.monotonic() print("PROGRESS", config, len(seen), "/", len(rows), flush=True) return outputs llm.llm_engine.step = checkpoint_step try: outputs = llm.chat( [[{"role": "user", "content": r["prompt"] + INSTRUCTION}] for r in rows], [SamplingParams(n=4, temperature=1.0, top_p=0.95, max_tokens=65536, seed=20261030 + i) for i in range(len(rows))], chat_template_kwargs=kwargs, use_tqdm=False, ) assert len(outputs) == len(rows) == len(seen) manifest["state"] = "generated" except BudgetExpired: manifest["state"] = "budget_stopped" manifest.update( elapsed_s=time.monotonic() - start, completed_problems=len(seen), completed_samples=len(records), total_output_tokens=sum(r["tokens"] for r in records), generation_complete=len(seen) == len(rows), ) dump_json(out_dir / "manifest.json", manifest) data.commit() call = grade.spawn(config, run) manifest["grade_call_id"] = call.object_id dump_json(out_dir / "manifest.json", manifest) data.commit() print("GENERATED", json.dumps(manifest), flush=True) return manifest def extract_solution(text, entrypoint): """No code in the thinking section may reach the evaluator.""" from evalplus.sanitize import sanitize if "" not in text: return "" answer = text.rsplit("", 1)[1].strip() match = re.search(r"```(?:python|py)?[ \t]*\r?\n(.*?)(?:```|$)", answer, re.S | re.I) code = match.group(1) if match else answer return sanitize(code, entrypoint=entrypoint) if code.strip() else "" @APP.function( image=grade_image, cpu=4, memory=8192, volumes={"/data": data}, timeout=1800, retries=0, max_containers=5, single_use_containers=True, scaledown_window=2, block_network=True, ) def grade(config: str, run: str): """EvalPlus's actual base/plus checker, with its default timeout policy.""" from concurrent.futures import ProcessPoolExecutor, as_completed from collections import Counter from evalplus.evaluate import check_correctness, get_groundtruth from evalplus.eval import PASS from evalplus.eval._special_oracle import MBPP_OUTPUT_NOT_NONE_TASKS start = time.monotonic() data.reload() out_dir = ROOT / run / config manifest = json.loads((out_dir / "manifest.json").read_text()) rows, hashes = get_problems(manifest["limit"]) assert hashes == manifest["dataset_hashes"] problems = {r["task_id"]: r for r in rows} # Generate reference outputs only in Modal, using the official canonical # functions and special oracles. These are not model-generated programs. groundtruth = {} for bench in ("humaneval", "mbpp"): subset = {k: r for k, r in problems.items() if r["bench"] == bench} # Include the subset IDs in the cache key to keep smoke/full caches apart. cache_hash = hashes[bench] + "_" + hashlib.md5("|".join(subset).encode()).hexdigest() groundtruth.update(get_groundtruth( subset, cache_hash, [] if bench == "humaneval" else MBPP_OUTPUT_NOT_NONE_TASKS, )) records = [json.loads(line) for line in (out_dir / "generations.jsonl").open()] counts = Counter((r["task_id"], r["sample_index"]) for r in records) assert all(n == 1 for n in counts.values()) futures = {} with ProcessPoolExecutor(max_workers=4) as pool, (out_dir / "sanitized.jsonl").open("w") as solutions: for record in records: task = problems[record["task_id"]] solution = extract_solution(record["text"], task["entry_point"]) solutions.write(json.dumps({"task_id": record["task_id"], "sample_index": record["sample_index"], "solution": solution}) + "\n") record["solution_empty"] = not solution.strip() if record["thinking_closed"] and solution.strip(): future = pool.submit( check_correctness, task["bench"], record["sample_index"], task, solution, groundtruth[task["task_id"]], False, True, f"{task['task_id']}:{record['sample_index']}", ) futures[future] = record else: record.update(base_pass=False, plus_pass=False, base_status="format_failure", plus_status="format_failure") for i, future in enumerate(as_completed(futures), 1): record = futures[future] result = future.result() # Infrastructure errors are never turned into model failures. base_status, _ = result["base"] plus_status, _ = result["plus"] record.update(base_status=base_status, plus_status=plus_status, base_pass=base_status == PASS, plus_pass=base_status == plus_status == PASS) if i % 200 == 0: print("GRADED", config, i, "/", len(futures), flush=True) records.sort(key=lambda r: (r["bench"], int(r["task_id"].split("/")[-1]), r["sample_index"])) with (out_dir / "results.jsonl").open("w") as handle: for record in records: handle.write(json.dumps({k: v for k, v in record.items() if k != "text"}) + "\n") summary = { "config": config, "generation_complete": manifest["generation_complete"], "grade_s": time.monotonic() - start, "benches": {}, } for bench in ("humaneval", "mbpp"): rs = [r for r in records if r["bench"] == bench] if not rs: continue per_problem = Counter(r["task_id"] for r in rs) assert set(per_problem.values()) == {4} summary["benches"][bench] = { "problems": len(per_problem), "samples": len(rs), "base_pass_at_1": sum(r["base_pass"] for r in rs) / len(rs), "plus_pass_at_1": sum(r["plus_pass"] for r in rs) / len(rs), "mean_output_tokens": sum(r["tokens"] for r in rs) / len(rs), "truncation_rate": sum(r["finish_reason"] == "length" for r in rs) / len(rs), "unclosed_thinking_rate": sum(not r["thinking_closed"] for r in rs) / len(rs), } dump_json(out_dir / "summary.json", summary) data.commit() print("SUMMARY", json.dumps(summary), flush=True) return summary