| |
| """P2: the baseline miniF2F pass@8 run (cap 16384, then the 40960 reference) as ONE supervised, resumable job. |
| |
| .venv/bin/python scripts/p2/launch_baseline.py --run-dir .tmp/p2/<run>/baseline \\ |
| --problems .tmp/p2/problems/minif2f.jsonl # add --print to show the commands only |
| |
| What this process IS: the heavy job. It takes the repo heavy-job flock ONCE (waiting if another job |
| holds it), keeps it for the whole sequence, and runs each phase as a child that *inherits* it |
| (`--inherit-lock <this pid>`; `align.utils.guard.inherit_heavy_job_lock` verifies the claim). |
| Phases, in this order so the pre-registered primary result exists first: |
| |
| gen_c16384 -> ver_c16384 -> gen_c40960 -> ver_c40960 |
| |
| Generation and Lean verification are both heavy; they never overlap, and no other agent's job can slip |
| in between phases, because the lock is not released between them. |
| |
| Properties (each one is a failure this job would otherwise have): |
| * RESUMABLE at every level. Re-running the same command skips phases with a `.done` marker; a phase |
| killed half way resumes from its append-only output (both children skip finished keys). |
| * COMPLETENESS IS CHECKED, not assumed: a generation phase is done only when it holds exactly |
| problems x n records at that cap; a verification phase only when its key set equals the generation's. |
| A child that exits 0 with fewer records is retried, not accepted. |
| * BOUNDED RETRIES that need PROGRESS: a failed phase is retried (max `--max-retries`) only if its |
| output grew since the previous attempt; a crash loop with no progress stops with `stalled`. |
| * NO ORPHANS: every child runs in its own session; after it exits (cleanly or not) its process group |
| is signalled by its exact pgid (vLLM's EngineCore is known to outlive its parent and keep 19 GB of |
| GPU memory), and the next GPU phase waits for GPU memory to drain. Nothing is ever killed by a |
| name pattern. |
| * SIGTERM/SIGINT to this process stops the current child (SIGTERM, grace, SIGKILL of its group), |
| writes `stopped` to the status log and exits 143; the lock goes with the process. |
| * Host rules: preflight (strict) before every phase; scratch and outputs under /apps; vLLM TP=4, |
| `--max-num-seqs 128` (B's measured knee); <= 8 Lean workers. |
| |
| Status: `<run-dir>/supervisor.jsonl` (one event per line) and `<run-dir>/logs/<phase>.log`. |
| """ |
|
|
| 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)) |
|
|
| PY = str(REPO / ".venv" / "bin" / "python") |
| PYR = str(REPO / ".venv-rollout" / "bin" / "python") |
| MODEL = ("/apps/development/rlhf/hf_cache/hub/models--Goedel-LM--Goedel-Prover-V2-8B/snapshots/" |
| "dfd02e6271a58375dfbf3ece0175277cf6b6a89a") |
| KEY = ("problem_id", "sample_idx", "round", "cap_tokens") |
|
|
|
|
| class Stalled(RuntimeError): |
| pass |
|
|
|
|
| def phase_plan(a, parent_pid: int) -> list[dict]: |
| """The ordered phases and their exact argv. Pure (tests pin it).""" |
| d = Path(a.run_dir) |
| plan = [] |
| for cap in a.caps: |
| gen, ver = d / f"gen_minif2f_c{cap}.jsonl", d / f"ver_minif2f_c{cap}.jsonl" |
| plan.append({ |
| "name": f"gen_c{cap}", "kind": "gen", "out": str(gen), "cap": cap, |
| "cmd": [PYR, str(REPO / "extras/vllm_rollout/lean_gen.py"), "--model", a.model, "--problems", a.problems, |
| "--out", str(gen), "--n", str(a.n), "--cap-tokens", str(cap), "--round", "0", "--tp", str(a.tp), |
| "--seed", "0", "--max-num-seqs", str(a.max_num_seqs), "--inherit-lock", str(parent_pid), |
| "--label", f"p2-baseline-gen-c{cap}"], |
| }) |
| plan.append({ |
| "name": f"ver_c{cap}", "kind": "ver", "out": str(ver), "gen": str(gen), "cap": cap, |
| "cmd": [PY, str(REPO / "scripts/p2/verify_generations.py"), "--generations", str(gen), "--problems", a.problems, |
| "--out", str(ver), "--workers", str(a.workers), "--inherit-lock", str(parent_pid)], |
| }) |
| return plan |
|
|
|
|
| def count_keys(path: str, cap: int | None = None) -> set[tuple]: |
| keys = set() |
| p = Path(path) |
| if not p.is_file(): |
| return keys |
| with open(p, "rb") as f: |
| for line in f: |
| if not line.endswith(b"\n"): |
| break |
| r = json.loads(line) |
| if cap is None or r.get("cap_tokens") == cap: |
| keys.add(tuple(r[k] for k in KEY)) |
| return keys |
|
|
|
|
| def phase_complete(ph: dict, n_problems: int, n: int) -> tuple[bool, str]: |
| """Is the phase's OUTPUT complete? gen: problems x n records at this cap. ver: same key set as gen.""" |
| keys = count_keys(ph["out"], ph["cap"]) |
| if ph["kind"] == "gen": |
| want = n_problems * n |
| |
| skipped = len(list(open(ph["out"] + ".skipped.jsonl"))) if os.path.isfile(ph["out"] + ".skipped.jsonl") else 0 |
| return len(keys) + skipped >= want and len(keys) > 0, f"{len(keys)} records + {skipped} skipped of {want}" |
| gkeys = count_keys(ph["gen"], ph["cap"]) |
| return bool(gkeys) and keys == gkeys, f"{len(keys)} verdicts for {len(gkeys)} generations" |
|
|
|
|
| def n_problems_of(path: str) -> int: |
| return sum(1 for l in open(path) if l.strip()) |
|
|
|
|
| def gpu_memory_used_mib() -> list[int]: |
| out = subprocess.run(["nvidia-smi", "--query-gpu=memory.used", "--format=csv,noheader,nounits"], |
| capture_output=True, text=True, timeout=30).stdout.split() |
| return [int(x) for x in out] |
|
|
|
|
| def wait_gpu_drained(limit_mib: int = 1024, timeout_s: float = 180.0) -> None: |
| t0 = time.time() |
| while True: |
| used = gpu_memory_used_mib() |
| if max(used) <= limit_mib: |
| return |
| if time.time() - t0 > timeout_s: |
| raise RuntimeError(f"GPU memory did not drain after the previous phase: {used} MiB (an orphan holds it; " |
| "find it by pid with nvidia-smi and stop it by exact pid)") |
| time.sleep(5) |
|
|
|
|
| SHM_DIR = "/dev/shm" |
| SHM_NAME_RE = r"psm_[0-9a-f]{8}" |
|
|
|
|
| def _paths_in_use(shm_dir: str) -> set[str]: |
| """Every file under `shm_dir` that some live process has open (fd) or mapped (maps).""" |
| held: set[str] = set() |
| prefix = shm_dir.rstrip("/") + "/" |
| for entry in os.listdir("/proc"): |
| if not entry.isdigit(): |
| continue |
| try: |
| for fd in os.listdir(f"/proc/{entry}/fd"): |
| try: |
| t = os.readlink(f"/proc/{entry}/fd/{fd}") |
| except OSError: |
| continue |
| if t.startswith(prefix): |
| held.add(t.split(" (deleted)")[0]) |
| with open(f"/proc/{entry}/maps") as f: |
| for line in f: |
| i = line.find(prefix) |
| if i >= 0: |
| held.add(line[i:].split(" (deleted)")[0].strip()) |
| except (FileNotFoundError, ProcessLookupError, PermissionError): |
| continue |
| return held |
|
|
|
|
| def sweep_orphan_shm(since: float, shm_dir: str = SHM_DIR, only: "set[str] | None" = None) -> tuple[list[str], int]: |
| """Unlink the `psm_*` shared-memory segments a vLLM phase left behind. Returns (names, bytes). |
| |
| THE /dev/shm LEAK (R2 review, measured): every vLLM launch leaves its broadcast ring buffers in /dev/shm |
| (a clean exit: 2 segments, ~0.19 GB; a SIGKILLed one: ~1.4 GB), and nothing removes them -- 46 segments / |
| 6.3 GB were sitting there from P2 alone. tmpfs is RAM on a box with NO swap, and /dev/shm exhaustion is one |
| of the two 2026-08-12 incidents. Only segments that are provably this phase's litter are removed: named |
| exactly like Python's SharedMemory (`psm_` + 8 hex), owned by this uid, modified after the phase began, and |
| held (open or mapped) by NO process. Nothing is matched by pattern against processes. `only` restricts the |
| candidates to the given names (tests use it so that a mutated copy cannot touch other jobs' segments). |
| """ |
| import re |
|
|
| removed, total = [], 0 |
| held = _paths_in_use(shm_dir) |
| for name in sorted(os.listdir(shm_dir)): |
| if only is not None and name not in only: |
| continue |
| if not re.fullmatch(SHM_NAME_RE, name): |
| continue |
| path = os.path.join(shm_dir, name) |
| try: |
| st = os.stat(path) |
| except FileNotFoundError: |
| continue |
| if st.st_uid != os.getuid() or st.st_mtime < since or path in held: |
| continue |
| try: |
| os.unlink(path) |
| except FileNotFoundError: |
| continue |
| removed.append(name) |
| total += st.st_size |
| return removed, total |
|
|
|
|
| class Runner: |
| def __init__(self, a): |
| self.a = a |
| self.run_dir = Path(a.run_dir) |
| self.child: subprocess.Popen | None = None |
| self.stop = False |
|
|
| def log(self, **ev) -> None: |
| ev["t"] = time.strftime("%F %T") |
| with open(self.run_dir / "supervisor.jsonl", "a") as f: |
| f.write(json.dumps(ev) + "\n") |
| print(json.dumps(ev), flush=True) |
|
|
| def _signal(self, signum, _frame) -> None: |
| self.stop = True |
| self.log(event="signal", signum=signum) |
| self.kill_group(self.child, grace=60.0) |
|
|
| @staticmethod |
| def group_alive(pgid: int) -> bool: |
| """Does any NON-ZOMBIE process still belong to process group `pgid`? |
| |
| `os.killpg(pgid, 0)` is the wrong probe: a killed leader stays a zombie until its parent reaps it, |
| and the SIGTERM handler runs in the very thread that is blocked in `Popen.wait()` (which owns the |
| reaping), so the probe saw the dead leader as "alive" for the whole grace period -- every stop took |
| grace + 10 s (70 s) with the heavy-job lock held, whatever the child did.""" |
| for entry in os.listdir("/proc"): |
| if not entry.isdigit(): |
| continue |
| try: |
| with open(f"/proc/{entry}/stat") as f: |
| rest = f.read().rsplit(")", 1)[1].split() |
| except (FileNotFoundError, ProcessLookupError, IndexError): |
| continue |
| if int(rest[2]) == pgid and rest[0] != "Z": |
| return True |
| return False |
|
|
| @staticmethod |
| def kill_group(proc, grace: float) -> None: |
| """SIGTERM then SIGKILL the child's process group, by the exact pgid it was given (its own session).""" |
| if proc is None: |
| return |
| for sig, wait in ((signal.SIGTERM, grace), (signal.SIGKILL, 10.0)): |
| if not Runner.group_alive(proc.pid): |
| return |
| try: |
| os.killpg(proc.pid, sig) |
| except ProcessLookupError: |
| return |
| t0 = time.time() |
| while time.time() - t0 < wait: |
| if not Runner.group_alive(proc.pid): |
| return |
| time.sleep(0.2) |
|
|
| def run_child(self, ph: dict, attempt: int) -> int: |
| log_path = self.run_dir / "logs" / f"{ph['name']}.log" |
| env = dict(os.environ) |
| env.update({"TMPDIR": str(self.run_dir / "scratch"), "HF_HOME": "/apps/development/rlhf/hf_cache", |
| "HF_HUB_DISABLE_XET": "1", "PYTHONUNBUFFERED": "1"}) |
| for k in [k for k in env if k.startswith(("RANK", "LOCAL_RANK", "WORLD_SIZE"))]: |
| del env[k] |
| started_at = time.time() - 1.0 |
| with open(log_path, "a") as lf: |
| lf.write(f"\n=== attempt {attempt} {time.strftime('%F %T')}: {' '.join(ph['cmd'])}\n") |
| lf.flush() |
| self.child = subprocess.Popen(ph["cmd"], stdout=lf, stderr=subprocess.STDOUT, env=env, cwd=str(REPO), |
| start_new_session=True) |
| if self.stop: |
| |
| |
| |
| self.kill_group(self.child, grace=60.0) |
| rc = self.child.wait() |
| |
| self.kill_group(self.child, grace=15.0) |
| self.child = None |
| gone, nbytes = sweep_orphan_shm(started_at) |
| if gone: |
| self.log(event="shm_swept", phase=ph["name"], segments=len(gone), bytes=nbytes) |
| return rc |
|
|
| def run_phase(self, ph: dict, n_problems: int) -> None: |
| marker = self.run_dir / f"{ph['name']}.done" |
| if marker.is_file(): |
| self.log(event="skip_done", phase=ph["name"]) |
| return |
| last = -1 |
| for attempt in range(self.a.max_retries + 1): |
| if self.stop: |
| return |
| from align.utils import guard |
|
|
| guard.preflight_host_check(strict=True) |
| if ph["kind"] == "gen": |
| wait_gpu_drained() |
| t0 = time.time() |
| self.log(event="start", phase=ph["name"], attempt=attempt) |
| rc = self.run_child(ph, attempt) |
| ok, detail = phase_complete(ph, n_problems, self.a.n) |
| size = len(count_keys(ph["out"], ph["cap"])) |
| self.log(event="end", phase=ph["name"], attempt=attempt, rc=rc, complete=ok, detail=detail, |
| wall_s=round(time.time() - t0, 1)) |
| if rc == 0 and ok: |
| marker.write_text(json.dumps({"detail": detail, "t": time.strftime("%F %T")})) |
| return |
| if self.stop: |
| return |
| if size <= last: |
| raise Stalled(f"{ph['name']}: attempt {attempt} made no progress ({detail}); stopping instead of looping") |
| last = size |
| time.sleep(self.a.backoff) |
| raise Stalled(f"{ph['name']}: still incomplete after {self.a.max_retries + 1} attempts ({detail})") |
|
|
| def summary(self) -> dict: |
| out = {} |
| for cap in self.a.caps: |
| vp = self.run_dir / f"ver_minif2f_c{cap}.jsonl" |
| if not vp.is_file(): |
| continue |
| solved, seen, acc = set(), set(), 0 |
| for l in open(vp): |
| r = json.loads(l) |
| seen.add(r["problem_id"]) |
| if r["accepted"] and r["sample_idx"] < self.a.n: |
| solved.add(r["problem_id"]) |
| acc += int(r["accepted"]) |
| out[f"solved@{self.a.n}_cap{cap}"] = f"{len(solved)}/{len(seen)}" |
| out[f"accepted_samples_cap{cap}"] = acc |
| return out |
|
|
| def run(self) -> int: |
| from align.utils import guard |
|
|
| (self.run_dir / "logs").mkdir(parents=True, exist_ok=True) |
| (self.run_dir / "scratch").mkdir(parents=True, exist_ok=True) |
| guard.assert_paths_off_root([str(self.run_dir), str(self.run_dir / "scratch")]) |
| signal.signal(signal.SIGTERM, self._signal) |
| signal.signal(signal.SIGINT, self._signal) |
| t0 = time.time() |
| guard.acquire_heavy_job_lock("p2-baseline-supervisor", blocking=True, lock_path=getattr(self.a, "lock_path", None)) |
| self.log(event="lock_acquired", waited_s=round(time.time() - t0, 1), pid=os.getpid()) |
| n_problems = n_problems_of(self.a.problems) |
| try: |
| for ph in phase_plan(self.a, os.getpid()): |
| self.run_phase(ph, n_problems) |
| if self.stop: |
| self.log(event="stopped") |
| return 143 |
| except Stalled as e: |
| self.log(event="stalled", error=str(e)) |
| return 3 |
| self.log(event="complete", **self.summary()) |
| return 0 |
|
|
|
|
| def main(argv=None) -> int: |
| ap = argparse.ArgumentParser(description=__doc__.split("\n")[0]) |
| ap.add_argument("--run-dir", required=True) |
| ap.add_argument("--problems", required=True, help="miniF2F Problem JSONL (all 244)") |
| ap.add_argument("--model", default=MODEL) |
| ap.add_argument("--caps", type=int, nargs="+", default=[16384, 40960]) |
| ap.add_argument("--n", type=int, default=8) |
| ap.add_argument("--tp", type=int, default=4) |
| ap.add_argument("--max-num-seqs", type=int, default=128) |
| ap.add_argument("--workers", type=int, default=8) |
| ap.add_argument("--max-retries", type=int, default=2) |
| ap.add_argument("--backoff", type=float, default=60.0) |
| ap.add_argument("--print", action="store_true", dest="print_only") |
| a = ap.parse_args(argv) |
| if not 1 <= a.workers <= 8: |
| raise SystemExit("--workers must be in [1, 8] (host rule)") |
| if a.n != 8 or a.caps[0] != 16384: |
| print("note: the registered primary endpoint is n=8 at cap 16384 (docs/P2.md section 1)", file=sys.stderr) |
| for p in (a.run_dir,): |
| if not os.path.abspath(p).startswith("/apps/"): |
| raise SystemExit(f"--run-dir {p} must be on /apps (never / or /tmp)") |
| if a.print_only: |
| for ph in phase_plan(a, os.getpid()): |
| print(" ".join(ph["cmd"])) |
| return 0 |
| return Runner(a).run() |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
