#!/usr/bin/env python3 """Resume the ELF quality matrix on independent single-GPU workers (no timing).""" from __future__ import annotations import argparse from concurrent.futures import ThreadPoolExecutor, wait from dataclasses import asdict, dataclass from datetime import datetime, timezone import json import os from pathlib import Path import queue import signal import subprocess import sys import threading import time from elf_common import ROOT, SOURCE_COMMIT, default_sampling, manifest, read_jsonl, sha256, write_json from inspect_elf_progress import inspect_run, OWT_MAIN_STEPS EOS_POLICY = "released_ids_or_raw_plus_eos_then_source_cap_v2" TASKS = ("owt", "owt-prefix", "wmt14", "xsum") @dataclass(frozen=True) class Cell: task: str split: str steps: int count: int input_path: str | None = None input_sha256: str | None = None @property def key(self): return f"{self.task}/{self.split}/steps_{self.steps}/seed_42/compile_0" def build_matrix(root, tasks): cells, blocked = [], [] for task in tasks: if task == "owt": cells.extend(Cell(task, "unconditional", s, 1024) for s in OWT_MAIN_STEPS) continue splits = (("c64-text-t5", OWT_MAIN_STEPS, 1024),) if task == "owt-prefix" else ( ("validation", (1, 2, 4, 8, 16, 32, 64), 3000 if task == "wmt14" else 11332), ("test", (32, 64), 3003 if task == "wmt14" else 11334), ) for split, steps, expected_count in splits: path = root / "data/elf" / ("owt-prefix.jsonl" if task == "owt-prefix" else f"{task}-{split}.jsonl") try: meta = json.loads(path.with_suffix(".manifest.json").read_text()) digest = sha256(path) rows = read_jsonl(path) if (meta["task"] != task or meta["sha256"] != digest or meta["count"] != expected_count or len(rows) != expected_count or len({r["id"] for r in rows}) != expected_count): raise ValueError("input task/count/hash/IDs mismatch") if task == "owt-prefix": if meta.get("protocol") != "c64_text_t5_v1": raise ValueError("prefix protocol mismatch") source = root / "data/conditional/owt-c64/prompts.jsonl" if meta.get("source_sha256") != sha256(source): raise ValueError("prefix source changed; prepare prefix again") elif meta.get("split") != split: raise ValueError("input split mismatch") except (OSError, KeyError, TypeError, ValueError) as error: blocked.append({"task": task, "split": split, "status": "blocked_input", "reason": str(error)}) continue cells.extend(Cell(task, split, s, expected_count, str(path), digest) for s in steps) return sorted(cells, key=lambda c: (c.steps, c.task, c.split)), blocked def matches_request(request, cell, asset_hash, assets): checkpoint_task = "owt" if cell.task == "owt-prefix" else cell.task record = assets[checkpoint_task] expected = { "task": cell.task, "steps": cell.steps, "seed": 42, "compile": False, "batch_size": 8, "source_commit": SOURCE_COMMIT, "assets_manifest_sha256": asset_hash, "checkpoint_task": checkpoint_task, "checkpoint_selection": "ema_params1", "checkpoint_sha256": record["sha256"][record["checkpoint"]], "sampling": default_sampling(cell.task, cell.steps), "model_canvas_t5_tokens": 1088 if cell.task == "xsum" else 128 if cell.task == "wmt14" else 1024, "protocol": "c64_text_t5_v1" if cell.task == "owt-prefix" else "elf_native_v1", } if any(request.get(key) != value for key, value in expected.items()): return False if cell.task in ("wmt14", "xsum") and request.get("condition_tokenization_policy") != EOS_POLICY: return False return (request.get("input") or {}).get("sha256") == cell.input_sha256 def reusable_result(root, cell, asset_hash, assets): candidates, rejected = [], [] for path in sorted((root / "results").glob(f"elf*/quality/{cell.key}")): try: request = json.loads((path / "request.json").read_text()) if not matches_request(request, cell, asset_hash, assets): continue row = inspect_run(path) if row["status"] not in ("quality_valid", "quality_invalid", "needs_evaluation"): rejected.append({"run": str(path), "reason": row.get("reason", row["status"])}) continue gen = json.loads((path / "generation.json").read_text()) keys = ("sampling", "input", "assets_manifest_sha256", "checkpoint_sha256", "source_commit", "condition_tokenization_policy", "batch_size", "protocol") if any(gen.get(k) != request.get(k) for k in keys): raise ValueError("request/generation provenance differs") diversity = cell.task == "owt-prefix" total = cell.count + (min(256, cell.count) * 4 if diversity else 0) if gen["prompt_count"] != cell.count or gen["sample_count"] != total or bool(gen.get("diversity")) != diversity: continue # e.g. 1000-row sanity is not a 1024-row primary run. if row["status"] == "quality_invalid": reason = row.get("reason") or "" if not reason.startswith("At least one completion has fewer than two GPT-2 tokens;"): rejected.append({"run": str(path), "reason": "unrecognized invalid-quality reason"}) continue candidates.append(row) except (OSError, ValueError, KeyError, TypeError) as error: rejected.append({"run": str(path), "reason": str(error)}) # A valid score and a recorded invalid score for one exact cell conflict; # never silently choose the better run. completed = [r for r in candidates if r["status"] != "needs_evaluation"] if len({r["status"] for r in completed}) > 1: raise ValueError(f"Conflicting completed results for {cell.key}; inspect runs manually") return (completed or candidates or [None])[0], rejected def fresh_path(output_root, cell): for attempt in range(10000): base = output_root if attempt == 0 else output_root.with_name(output_root.name + f"-retry{attempt}") path = base / "quality" / cell.key if not path.exists(): return path raise RuntimeError("Too many existing retry directories") def commands(root, python, cell, path, generate): result = [] if generate: cmd = [python, "-B", str(root / "scripts/run_elf.py"), "generate", "--root", str(root), "--task", cell.task, "--steps", str(cell.steps), "--seed", "42", "--batch-size", "8", "--num-samples", str(cell.count), "--output", str(path)] if cell.input_path: cmd += ["--input", cell.input_path] if cell.task == "owt-prefix": cmd += ["--diversity"] result.append(("generate", cmd)) result.append(("evaluate", [python, "-B", str(root / "scripts/evaluate_elf.py"), "--root", str(root), "--run", str(path), "--batch-size", "8"])) return result def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--root", type=Path, default=ROOT) parser.add_argument("--output-root", type=Path) parser.add_argument("--tasks", nargs="+", choices=TASKS, default=list(TASKS)) parser.add_argument("--gpus", nargs="+", default=["0", "1"]) parser.add_argument("--plan", action="store_true", help="Read-only preview; no GPU queries, writes or inference") args = parser.parse_args() root = args.root.resolve() output_root = (args.output_root or root / "results/elf-gpu01").resolve() if output_root.parent != root / "results" or not output_root.name.startswith("elf"): parser.error("output-root must be an elf* directory directly under ROOT/results") if len(set(args.tasks)) != len(args.tasks): parser.error("duplicate tasks") assets = manifest(root)["assets"] asset_hash = sha256(root / "data/elf/assets.json") cells, blocked = build_matrix(root, args.tasks) jobs, reports = [], [] for cell in cells: existing, rejected = reusable_result(root, cell, asset_hash, assets) action = ("skip" if existing["status"] != "needs_evaluation" else "evaluate") if existing else "generate+evaluate" path = Path(existing["run"]) if existing else fresh_path(output_root, cell) row = {"cell": asdict(cell), "action": action, "run": str(path), "rejected_candidates": rejected, "status": existing["status"] if action == "skip" else "pending"} reports.append(row) print(f"{action:18} {cell.key} -> {path}", flush=True) if action != "skip": jobs.append((cell, row)) for row in blocked: print(f"BLOCKED {row['task']}/{row['split']}: {row['reason']}", flush=True) print(f"Plan: {len(cells)} cells, {len(jobs)} queued, {len(blocked)} blocked input groups; no timing.", flush=True) if args.plan: return 2 if blocked else 0 if sys.platform != "linux": parser.error("GPU execution is server-only; use --plan for a read-only preview") if not jobs: print("Nothing to run. Existing invalid-quality results remain invalid, not successful PPL points.") return 2 if blocked else 0 import fcntl from contextlib import ExitStack with ExitStack() as stack: gpu_records = [] lock_dir = root / "results/elf-queue-locks" lock_dir.mkdir(parents=True, exist_ok=True) for gpu in args.gpus: line = subprocess.check_output(["nvidia-smi", "-i", gpu, "--query-gpu=uuid,name", "--format=csv,noheader"], text=True).strip() uuid, name = [s.strip() for s in line.split(",", 1)] if not uuid.startswith("GPU-") or any(r["uuid"] == uuid for r in gpu_records): raise ValueError("GPU selections must resolve to distinct physical GPUs") handle = stack.enter_context((lock_dir / f"{uuid}.lock").open("a")) fcntl.flock(handle, fcntl.LOCK_EX | fcntl.LOCK_NB) gpu_records.append({"physical_selection": gpu, "uuid": uuid, "name": name}) stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S") + f"-{os.getpid()}" log_dir = output_root / "logs" / stamp log_dir.mkdir(parents=True) report_path = log_dir / "queue.json" report = {"schema": "elf-quality-queue-v1", "gpus": gpu_records, "runs": reports, "blocked": blocked, "assets_manifest_sha256": asset_hash, "timing": False} mutex, stop = threading.Lock(), threading.Event() processes = {} work = queue.Queue() for job in jobs: work.put(job) def publish(): write_json(report_path, report) def worker(gpu): env = dict(os.environ, CUDA_VISIBLE_DEVICES=gpu["uuid"], PYTHONDONTWRITEBYTECODE="1", HF_HUB_OFFLINE="1", TRANSFORMERS_OFFLINE="1", WANDB_MODE="disabled") while not stop.is_set(): try: cell, row = work.get_nowait() except queue.Empty: break path = Path(row["run"]) with mutex: row.update(status="running", gpu=gpu["uuid"]) publish() print(f"START GPU {gpu['physical_selection']} {cell.key}", flush=True) try: for stage, command in commands(root, sys.executable, cell, path, row["action"] == "generate+evaluate"): if stop.is_set(): raise RuntimeError("queue interrupted") log = log_dir / (cell.key.replace("/", "__") + f"__{stage}.log") with log.open("w") as stream: stream.write(json.dumps(command) + "\n") stream.flush() with mutex: if stop.is_set(): raise RuntimeError("queue interrupted") proc = subprocess.Popen(command, cwd=root, env=env, stdout=stream, stderr=subprocess.STDOUT, start_new_session=True) processes[gpu["uuid"]] = proc code = proc.wait() with mutex: processes.pop(gpu["uuid"], None) with mutex: row[stage + "_log"] = str(log) if code != 0: checked, _ = reusable_result(root, cell, asset_hash, assets) if stage == "evaluate" and checked and checked["run"] == str(path) and checked["status"] == "quality_invalid": break raise RuntimeError(f"{stage} exit={code}; see {log}") checked, _ = reusable_result(root, cell, asset_hash, assets) if not checked or checked["run"] != str(path) or checked["status"] not in ("quality_valid", "quality_invalid"): raise RuntimeError("finished process did not publish a verified quality result") with mutex: row.update(status=checked["status"], reason=checked.get("reason")) except Exception as error: with mutex: row.update(status="failed", reason=str(error)) finally: with mutex: publish() print(f"DONE GPU {gpu['physical_selection']} {cell.key}: {row['status']}", flush=True) work.task_done() publish() print(f"Logs and live queue report: {log_dir}", flush=True) pool = ThreadPoolExecutor(max_workers=len(gpu_records)) futures = [pool.submit(worker, gpu) for gpu in gpu_records] interrupted = False try: pending = futures while pending: _, pending = wait(pending, timeout=30) if pending: with mutex: running = [r['cell'] for r in reports if r['status'] == 'running'] print(f"RUNNING {running}; remaining queued={work.qsize()}", flush=True) for future in futures: future.result() except KeyboardInterrupt: interrupted = True stop.set() with mutex: active = list(processes.values()) for proc in active: try: os.killpg(proc.pid, signal.SIGTERM) except ProcessLookupError: pass for proc in active: try: proc.wait(timeout=15) except subprocess.TimeoutExpired: try: os.killpg(proc.pid, signal.SIGKILL) except ProcessLookupError: pass finally: pool.shutdown(wait=True) with mutex: publish() counts = {s: sum(r["status"] == s for r in reports) for s in sorted({r["status"] for r in reports})} print(f"Final statuses: {counts}; blocked input groups={len(blocked)}; report={report_path}", flush=True) return 130 if interrupted else 2 if blocked or any(r["status"] in ("failed", "pending", "running") for r in reports) else 0 if __name__ == "__main__": raise SystemExit(main())