| """Generation and verification as ONE heavy job, overlapped. |
| |
| THE IDLE-GPU DEFECT. The P2 queue ran `lean_gen.py` (GPU) and then `verify_generations.py` |
| (CPU, Lean) back to back. Verification of a 1,952-sample miniF2F run took 1.43 h and the pool |
| round's would take ~12 h; the four A10Gs sat at 0% for all of it. Running the verifier as a |
| second job beside the generator is not an option either: /apps/CLAUDE.md rule 2 forbids two heavy |
| jobs at once, and the repo's flock enforces it. |
| |
| So this supervisor takes the heavy-job lock ONCE and runs both phases as children that inherit it |
| (`--inherit-lock`, the mechanism scripts/p2/launch_baseline.py already uses): the generator owns |
| the GPUs, and every `--every` seconds a verification pass checks the COMPLETE lines written so far. |
| A pass works on a snapshot that stops at the last newline, because the generator appends while we |
| read and the strict JSONL reader (p2_common.iter_jsonl) rightly refuses a half-written line. |
| Verification is keyed and resumable, so each pass only verifies what is new, and the final pass |
| after the generator exits covers the tail. One lock, one job, both resources busy. |
| |
| Exit codes: 0 both phases complete; the generator's or verifier's non-zero code otherwise |
| (143 when stopped by a signal); 4 when the generator exited 0 but its manifest does not say |
| `complete` (a partial sample set must never read as a finished stage -- U12-2); 5 when GPU memory |
| did not drain after the generator (an orphan holds it -- U12-3). Re-running the same command |
| resumes both phases. |
| |
| After EVERY generator exit (clean, failed or stopped) the generator's process group is reaped by its |
| exact pgid, orphan `psm_*` /dev/shm segments of this phase are swept, and GPU memory is waited on |
| until drained, using scripts/p2/launch_baseline.py's helpers (vLLM's EngineCore is documented to |
| outlive its parent: 19 GB of GPU memory, 0.19-1.4 GB of /dev/shm). The next GPU stage (sft.py, |
| export, eval generation) therefore starts on a clean host even though this supervisor releases the |
| lock on exit. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
| import signal |
| import subprocess |
| import sys |
| import time |
| from pathlib import Path |
|
|
| REPO = Path(__file__).resolve().parents[2] |
| sys.path.insert(0, str(REPO)) |
|
|
| PYR = str(REPO / ".venv-rollout" / "bin" / "python") |
| PY = str(REPO / ".venv" / "bin" / "python") |
|
|
|
|
| _LB = None |
|
|
|
|
| def launch_baseline(): |
| """scripts/p2/launch_baseline.py, loaded by path (the repo idiom): its kill_group / sweep_orphan_shm / |
| wait_gpu_drained are the reviewed implementation; a second copy would drift.""" |
| global _LB |
| if _LB is None: |
| import importlib.util |
|
|
| spec = importlib.util.spec_from_file_location("_p2_launch_baseline", REPO / "scripts/p2/launch_baseline.py") |
| _LB = importlib.util.module_from_spec(spec) |
| spec.loader.exec_module(_LB) |
| return _LB |
|
|
|
|
| def generation_complete(gen_out: str) -> tuple[bool, str]: |
| """Did the LAST lean_gen session on `gen_out` finish every planned sample? (its manifest says so)""" |
| mp = gen_out + ".manifest.json" |
| if not os.path.exists(mp): |
| return False, "no manifest" |
| with open(mp, encoding="utf-8") as f: |
| sessions = json.load(f).get("sessions", []) |
| if not sessions: |
| return False, "manifest has no sessions" |
| st = sessions[-1].get("status") |
| return st == "complete", f"last lean_gen session status {st!r}" |
|
|
|
|
| def reap_after_gen(proc, started_at: float, events, lb=None, log=None) -> None: |
| """Kill survivors of the generator's process group, sweep its /dev/shm litter, wait for the GPUs to drain.""" |
| lb = lb or launch_baseline() |
| log = log or (lambda **kw: log_event(events, **kw)) |
| lb.Runner.kill_group(proc, 15.0) |
| gone, nbytes = lb.sweep_orphan_shm(started_at) |
| if gone: |
| log(event="shm_swept", segments=len(gone), bytes=nbytes) |
| lb.wait_gpu_drained() |
| log(event="gpu_drained") |
|
|
|
|
| def snapshot_complete_lines(src: Path, dst: Path) -> int: |
| """Copy `src` up to and including its last newline into `dst` (atomically). Returns the |
| number of complete lines. A trailing partial line -- the generator mid-append -- is left out.""" |
| if not src.exists(): |
| return 0 |
| data = src.read_bytes() |
| cut = data.rfind(b"\n") + 1 |
| tmp = dst.with_suffix(dst.suffix + ".tmp") |
| tmp.write_bytes(data[:cut]) |
| os.replace(tmp, dst) |
| return data[:cut].count(b"\n") |
|
|
|
|
| def log_event(path: Path, **kw) -> None: |
| kw["t"] = time.strftime("%Y-%m-%d %H:%M:%S") |
| with open(path, "a", encoding="utf-8") as f: |
| f.write(json.dumps(kw) + "\n") |
|
|
|
|
| def main(argv=None) -> int: |
| ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| ap.add_argument("--run-dir", required=True) |
| ap.add_argument("--gen-args", required=True, help="arguments for lean_gen.py, as ONE quoted string (must include --out)") |
| ap.add_argument("--problems", required=True) |
| ap.add_argument("--verify-out", required=True) |
| ap.add_argument("--workers", type=int, default=8) |
| ap.add_argument("--every", type=float, default=300.0, help="seconds between verification passes") |
| ap.add_argument("--feedback", action="store_true") |
| ap.add_argument("--label", default="p2-gen-verify") |
| a = ap.parse_args(argv) |
|
|
| from align.utils import guard |
|
|
| run = Path(a.run_dir) |
| run.mkdir(parents=True, exist_ok=True) |
| events = run / "pipeline.jsonl" |
| gen_args = a.gen_args.split() |
| gen_out = Path(gen_args[gen_args.index("--out") + 1]) |
| snap = run / (gen_out.stem + ".snapshot.jsonl") |
|
|
| guard.acquire_heavy_job_lock(a.label, blocking=True) |
| guard.preflight_host_check() |
| me = os.getpid() |
| log_event(events, event="lock_acquired", pid=me) |
|
|
| stop = {"sig": None} |
|
|
| def on_signal(sig, _frm): |
| stop["sig"] = sig |
|
|
| signal.signal(signal.SIGTERM, on_signal) |
| signal.signal(signal.SIGINT, on_signal) |
|
|
| gen_log = open(run / "gen.log", "ab") |
| started_at = time.time() - 1.0 |
| gen = subprocess.Popen( |
| [PYR, str(REPO / "extras/vllm_rollout/lean_gen.py"), *gen_args, "--inherit-lock", str(me)], |
| stdout=gen_log, stderr=subprocess.STDOUT, start_new_session=True, |
| ) |
| log_event(events, event="gen_start", pid=gen.pid, args=gen_args) |
|
|
| def verify_pass(source: Path) -> int: |
| cmd = [PY, str(REPO / "scripts/p2/verify_generations.py"), "--generations", str(source), |
| "--problems", a.problems, "--out", a.verify_out, "--workers", str(a.workers), |
| "--inherit-lock", str(me)] |
| if a.feedback: |
| cmd.append("--feedback") |
| with open(run / "verify.log", "ab") as vl: |
| v = subprocess.Popen(cmd, stdout=vl, stderr=subprocess.STDOUT, start_new_session=True) |
| while v.poll() is None: |
| if stop["sig"] is not None: |
| os.killpg(v.pid, signal.SIGTERM) |
| v.wait() |
| break |
| time.sleep(2) |
| return v.returncode |
|
|
| rc_v = 0 |
| last_lines = -1 |
| while gen.poll() is None and stop["sig"] is None: |
| n = snapshot_complete_lines(gen_out, snap) |
| if n > last_lines and n > 0: |
| log_event(events, event="verify_pass", lines=n) |
| rc_v = verify_pass(snap) |
| log_event(events, event="verify_pass_end", lines=n, rc=rc_v) |
| if rc_v != 0 and stop["sig"] is None: |
| |
| |
| log_event(events, event="verify_overlap_disabled", rc=rc_v) |
| last_lines = float("inf") |
| rc_v = 0 |
| continue |
| last_lines = n |
| t_end = time.monotonic() + a.every |
| while time.monotonic() < t_end and gen.poll() is None and stop["sig"] is None: |
| time.sleep(5) |
|
|
| if stop["sig"] is not None: |
| if gen.poll() is None: |
| os.killpg(gen.pid, signal.SIGTERM) |
| try: |
| gen.wait(timeout=120) |
| except subprocess.TimeoutExpired: |
| os.killpg(gen.pid, signal.SIGKILL) |
| gen.wait() |
| try: |
| reap_after_gen(gen, started_at, events) |
| except RuntimeError as e: |
| log_event(events, event="gpu_not_drained", error=str(e)) |
| log_event(events, event="stopped", signal=stop["sig"], verify_rc=rc_v) |
| return 143 |
|
|
| rc_g = gen.wait() |
| log_event(events, event="gen_end", rc=rc_g) |
| try: |
| reap_after_gen(gen, started_at, events) |
| except RuntimeError as e: |
| log_event(events, event="gpu_not_drained", error=str(e)) |
| return rc_g or 5 |
| if rc_g != 0: |
| return rc_g |
| complete, why = generation_complete(gen_out) |
| if not complete: |
| log_event(events, event="gen_incomplete", detail=why) |
| return 4 |
| |
| rc_v = verify_pass(gen_out) |
| log_event(events, event="final_verify_end", rc=rc_v) |
| if snap.exists(): |
| snap.unlink() |
| return 0 if stop["sig"] is None and rc_v == 0 else (143 if stop["sig"] is not None else rc_v) |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
