Blame for code/modal_app/code_eval.py

"""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": "</think>" 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 "</think>" not in text:
        return ""
    answer = text.rsplit("</think>", 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