"""A frozen single-GPU NeMo Gym experiment, with DEV-only checkpoint selection. Each model-consuming stage is a subprocess. A journal reuses verified complete stages and archives interrupted evaluation outputs rather than overwriting them. The native trainer owns optimizer/RNG checkpoints and exact rollout boundaries. """ from __future__ import annotations import argparse from contextlib import contextmanager, ExitStack from datetime import datetime, timezone import fcntl import json import math import os from pathlib import Path import shutil import socket import subprocess import sys import time from urllib.error import URLError from urllib.request import urlopen from . import budget_lineage from .agent_rl_tasks import generate_tasks from .agent_rl_training import DEFAULT_LORA, DEFAULT_TRAINING from .evaluation import checkpoint_sha256, sha256_file FORMAT = "hessian-nemo-long-experiment-v1" CHECKPOINTS = (64, 128, 256) def _read(path): from .training import _unique_pairs return json.loads( Path(path).read_text(), object_pairs_hook=_unique_pairs, parse_constant=lambda value: (_ for _ in ()).throw(ValueError(value)), ) def _write(path, value): path = Path(path) path.parent.mkdir(parents=True, exist_ok=True) temporary = path.with_name(path.name + ".writing") with temporary.open("w") as stream: stream.write(json.dumps(value, indent=2, ensure_ascii=False, allow_nan=False) + "\n") stream.flush() os.fsync(stream.fileno()) os.replace(temporary, path) descriptor = os.open(path.parent, os.O_DIRECTORY) try: os.fsync(descriptor) finally: os.close(descriptor) def _sources(): package = Path(__file__).parent names = ( "gym_experiment.py", "budget_lineage.py", "gym_experiment_evaluation.py", "gym_training_long.py", "managed_trial.py", "managed_training.py", "gym_training.py", "gym_client.py", "agent_rl_training.py", "agent_rl_tasks.py", "agent_rl_pilot.py", "agent_rl_evaluation.py", "training.py", "evaluation.py", "trajectories.py", "checkpoints.py", "cli.py", "__main__.py", "__init__.py", ) sources = {f"src/hessian_recipes/{name}": sha256_file(package / name) for name in names} for relative in ("examples/agent_gyms/nemo_gym_server.py", "scripts/gpu-preflight.py", "requirements/german-qwen-managed-cu126.txt"): sources[relative] = sha256_file(package.parent.parent / relative) return sources def prepare(baseline, tokenizer, math_dev, mgsm, output, *, budget_carry=None): from . import gym_experiment_evaluation as evaluation root = Path(output).resolve() root.mkdir(parents=True, exist_ok=True) if (root / "experiment-plan.json").exists() or (root / "training.json").exists(): raise ValueError("Experiment inputs already exist; use a fresh experiment directory") protocol = evaluation.prepare(baseline, tokenizer, math_dev, mgsm, root / "evaluation-plan") train = root / "train.jsonl" rows = generate_tasks("train", 256, seed=20261020) train.write_text("".join(json.dumps(row, ensure_ascii=False) + "\n" for row in rows)) configuration = { "schema_version": 1, "model": protocol["baseline"], "dataset": {"path": str(train), "sha256": sha256_file(train)}, "output_dir": str(root / "training"), "seed": 20261020, "chat_template_kwargs": {"enable_thinking": True}, "training": DEFAULT_TRAINING | {"max_steps": 256, "learning_rate": 5e-6, "beta": 0.1}, "lora": DEFAULT_LORA, } _write(root / "training.json", configuration) plan = { "format": FORMAT, "status": "prepared", "release_eligible": False, "created_at": datetime.now(timezone.utc).isoformat(), "root": str(root), "evaluation_plan": str(root / "evaluation-plan/plan.json"), "evaluation_plan_sha256": sha256_file(root / "evaluation-plan/plan.json"), "training_config": str(root / "training.json"), "training_config_sha256": sha256_file(root / "training.json"), "train_sha256": sha256_file(train), "source_sha256": _sources(), "checkpoint_steps": list(CHECKPOINTS), "checkpoint_interval": 64, "max_gpu_wall_seconds": 8 * 3600, "resume_qualification": {"steps": 4, "stop_step": 2, "seed": 43}, "selection": protocol["selection_rule"], "notice": "One training seed and a narrow scaffolded task family; no general reasoning or release claim.", } if budget_carry is not None: plan["budget_carry"] = budget_lineage.reference(budget_carry) budget_lineage.load(plan) _write(root / "experiment-plan.json", plan) return plan def verify_plan(path): path = Path(path).resolve() plan = _read(path) root = Path(plan["root"]) if ( plan.get("format") != FORMAT or path != root / "experiment-plan.json" or plan["source_sha256"] != _sources() or plan["evaluation_plan_sha256"] != sha256_file(plan["evaluation_plan"]) or plan["training_config_sha256"] != sha256_file(plan["training_config"]) or plan["train_sha256"] != sha256_file(root / "train.jsonl") or plan["checkpoint_steps"] != list(CHECKPOINTS) or plan["checkpoint_interval"] != 64 or plan["max_gpu_wall_seconds"] != 8 * 3600 or plan["resume_qualification"] != {"steps": 4, "stop_step": 2, "seed": 43} ): raise ValueError("Frozen experiment source or inputs changed") from .gym_experiment_evaluation import verify_plan as verify_evaluation_plan protocol, _ = verify_evaluation_plan(plan["evaluation_plan"]) config = _read(plan["training_config"]) expected_training = DEFAULT_TRAINING | {"max_steps": 256, "learning_rate": 5e-6, "beta": 0.1} if ( config["training"] != expected_training or config["model"] != protocol["baseline"] or config["seed"] != 20261020 or config["lora"] != DEFAULT_LORA or config["chat_template_kwargs"] != {"enable_thinking": True} or Path(config["output_dir"]) != root / "training" or config["dataset"] != {"path": str(root / "train.jsonl"), "sha256": plan["train_sha256"]} ): raise ValueError("Frozen training schedule/identity differs from preregistered experiment") budget_lineage.load(plan) return plan @contextmanager def _service(root, gym_python): server = Path(__file__).resolve().parents[2] / "examples/agent_gyms/nemo_gym_server.py" with socket.socket() as sock: sock.bind(("127.0.0.1", 0)) port = sock.getsockname()[1] base = f"http://127.0.0.1:{port}" environment = dict(os.environ) # The Gym CPU service has another Python ABI; never feed it the GPU overlay. environment.pop("PYTHONPATH", None) with (root / f"gym-service-{time.time_ns()}.log").open("w") as log: process = subprocess.Popen( [str(gym_python), str(server), "--port", str(port)], env=environment, stdout=log, stderr=subprocess.STDOUT, ) try: started = time.monotonic() while True: if process.poll() is not None: raise RuntimeError("The owned native Gym service exited during startup") try: with urlopen(base + "/hessian_identity", timeout=2) as response: identity = json.load(response) break except (URLError, TimeoutError): if time.monotonic() - started > 180: raise TimeoutError("Native Gym service startup timed out") time.sleep(0.2) yield base, identity finally: if process.poll() is None: process.terminate() try: process.wait(timeout=15) except subprocess.TimeoutExpired: process.kill() process.wait(timeout=15) class _Journal: def __init__(self, plan_path, resume): self.plan_path = Path(plan_path).resolve() self.plan = verify_plan(self.plan_path) self.root = Path(self.plan["root"]) self.path = self.root / "experiment-status.json" self.carry = budget_lineage.load(self.plan) if self.path.exists() and not resume: raise ValueError("Experiment already started; use --resume with the same frozen plan") if resume and not self.path.exists(): raise ValueError("Cannot resume without an existing experiment journal") fresh = not self.path.exists() self.state = ( _read(self.path) if self.path.exists() else { "status": "running", "plan_sha256": sha256_file(self.plan_path), "release_eligible": False, "completed": {}, "gpu_wall_seconds": 0.0, "attempts": [], "active_attempt": None, "budget_note": "Cumulative driver wall time on GPU workers; an unclosed attempt is conservatively charged through recovery, including unknown downtime.", } ) if self.carry is not None and fresh: self.state["budget_carry"] = self.carry self.state["gpu_wall_seconds"] = self.carry["carry_in_seconds"] self.state["budget_note"] = "Immutable prior-backend carry-in plus all new worker attempts; original absolute deadline retained." if ( self.state["plan_sha256"] != sha256_file(self.plan_path) or self.state["status"] not in {"running", "failed", "completed", "qualification_completed"} or not isinstance(self.state["completed"], dict) or not isinstance(self.state["attempts"], list) or type(self.state["gpu_wall_seconds"]) not in (int, float) or not math.isfinite(self.state["gpu_wall_seconds"]) or self.state["gpu_wall_seconds"] < 0 ): raise ValueError("The existing experiment journal belongs to another plan") for name in list(self.state["completed"]): self.complete(name) # Fail before any stage execution or recovery mutation. now = time.time() if self.carry is not None: budget_lineage.remaining(self.plan, self.state, now) sentinel = self.root / "budget-journal-started.json" if fresh: descriptor = os.open(sentinel, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600) with os.fdopen(descriptor, "w") as stream: json.dump({"plan_sha256": self.state["plan_sha256"], "budget_carry": self.carry, "minimum_worker_seconds": self.state["gpu_wall_seconds"], "minimum_attempt_count": 0}, stream, sort_keys=True) stream.flush() os.fsync(stream.fileno()) descriptor = os.open(self.root, os.O_DIRECTORY | os.O_RDONLY) try: os.fsync(descriptor) finally: os.close(descriptor) else: anchor = _read(sentinel) if (anchor.get("plan_sha256") != self.state["plan_sha256"] or anchor.get("budget_carry") != self.carry or type(anchor.get("minimum_worker_seconds")) not in (int, float) or not math.isfinite(anchor["minimum_worker_seconds"]) or anchor["minimum_worker_seconds"] < self.carry["carry_in_seconds"] or type(anchor.get("minimum_attempt_count")) is not int or anchor["minimum_attempt_count"] < 0 or self.state["gpu_wall_seconds"] < anchor["minimum_worker_seconds"] or len(self.state["attempts"]) < anchor["minimum_attempt_count"]): raise ValueError("Carried journal high-water/creation anchor changed") active = self.state.get("active_attempt") if active is not None: previous = self.state["attempts"][active] unknown = now - previous["accounted_through_epoch"] if not math.isfinite(unknown) or unknown < 0: raise ValueError("Cannot conservatively recover a backwards/invalid wall clock") self.state["gpu_wall_seconds"] += unknown if self.carry is not None: previous["charged_gpu_wall_seconds"] += unknown self.state["gpu_wall_seconds"] = self.carry["carry_in_seconds"] + sum( item["charged_gpu_wall_seconds"] for item in self.state["attempts"]) previous.update(recovered_at_epoch=now, conservative_unclosed_seconds=unknown) self.last_tick = time.monotonic() self.state["active_attempt"] = len(self.state["attempts"]) self.state["attempts"].append( { "started_at_epoch": now, "accounted_through_epoch": now, "hostname": socket.gethostname(), } ) if self.carry is not None: self.state["attempts"][-1]["charged_gpu_wall_seconds"] = 0.0 budget_lineage.remaining(self.plan, self.state, now) self.state.pop("error", None) self.save() def save(self): current = time.monotonic() elapsed = current - self.last_tick if not math.isfinite(elapsed) or elapsed < 0: raise ValueError("Invalid monotonic worker accounting") if self.carry is not None: active = self.state["active_attempt"] if active is not None: self.state["attempts"][active]["charged_gpu_wall_seconds"] += elapsed self.state["gpu_wall_seconds"] = self.carry["carry_in_seconds"] + sum( item["charged_gpu_wall_seconds"] for item in self.state["attempts"]) else: self.state["gpu_wall_seconds"] += elapsed self.last_tick = current active = self.state["active_attempt"] if active is not None: self.state["attempts"][active]["accounted_through_epoch"] = time.time() _write(self.path, self.state) if self.carry is not None: # Journal first: a crash before this second atomic write cannot grant # time. The independent high-water mark rejects journal rollback. _write(self.root / "budget-journal-started.json", { "plan_sha256": self.state["plan_sha256"], "budget_carry": self.carry, "minimum_worker_seconds": self.state["gpu_wall_seconds"], "minimum_attempt_count": len(self.state["attempts"]), }) def close(self): self.save() active = self.state["active_attempt"] if active is not None: self.state["attempts"][active].update( finished_at_epoch=time.time(), status=self.state["status"] ) self.state["active_attempt"] = None _write(self.path, self.state) def command(self, name, arguments): verify_plan(self.plan_path) self.state.update(status="running", stage=name) self.save() remaining = budget_lineage.remaining(self.plan, self.state, time.time()) logs = self.root / "logs" logs.mkdir(exist_ok=True) with (logs / f"{name}-{time.time_ns()}.log").open("w") as log: supervisor = shutil.which("timeout") if supervisor is None: raise RuntimeError("The bounded Linux worker requires GNU timeout") # This timer survives a driver SIGKILL and kills its owned process group. subprocess.run( [supervisor, "--signal=KILL", f"{remaining:.6f}s", sys.executable, *arguments], stdout=log, stderr=subprocess.STDOUT, start_new_session=True, timeout=remaining + 2, check=True, ) self.save() def done(self, name, files): if not files: raise ValueError("A completed stage requires durable evidence files") self.state["completed"][name] = {str(Path(p).resolve()): sha256_file(p) for p in files} self.save() def complete(self, name): files = self.state["completed"].get(name) if files is None: return False if ( not isinstance(files, dict) or not files or any(not Path(p).is_relative_to(self.root) for p in files) ): raise ValueError(f"Invalid completed stage evidence: {name}") if any(not Path(p).is_file() or sha256_file(p) != digest for p, digest in files.items()): raise ValueError(f"Completed stage evidence changed: {name}") return True def _archive(root, path): path = Path(path) if path.exists(): destination = root / "recovery" / f"{path.name}-{time.time_ns()}" destination.parent.mkdir(exist_ok=True) shutil.move(str(path), destination) def _suite(journal, model, label, split, selection=None): from . import gym_experiment_evaluation as evaluation protocol = Path(journal.plan["evaluation_plan"]) manifest = protocol.parent / "specifications" / label / split / "manifest.json" if not manifest.exists(): # An interrupted spec write is not reused as a complete manifest. _archive(journal.root, manifest.parent) evaluation.write_specs(protocol, model, label, split, selection_path=selection) entries = _read(manifest) _validate_suite_manifest(journal, model, label, split, selection, entries) for entry in entries: name = f"{label}-{split}-{entry['kind']}-{entry['seed']}" if journal.complete(name): continue report = Path(entry["path"]) _archive(journal.root, report.parent) module = ( "hessian_recipes.agent_rl_evaluation" if entry["kind"] == "agent" else "hessian_recipes" ) arguments = ["-m", module] if entry["kind"] == "math": arguments.append("eval") arguments += [entry["spec_path"], "--output", str(report.parent)] journal.command(name, arguments) parsed = _read(report) if entry["kind"] == "agent": from .agent_rl_evaluation import verify_report else: from .evaluation import _verify_report as verify_report verify_report(parsed) journal.done(name, [report, entry["spec_path"], manifest]) frozen, digest = evaluation.verify_plan(protocol) evaluation._load_suite( frozen, digest, manifest, split, selection_digest=sha256_file(selection) if selection else None, ) return manifest def _validate_suite_manifest(journal, model, label, split, selection, entries): """Fail before inference if a resumed spec manifest is partial or mismatched.""" from . import gym_experiment_evaluation as evaluation protocol, digest = evaluation.verify_plan(journal.plan["evaluation_plan"]) expected = { (kind, seed) for kind in ("agent", "math") for seed in (protocol["dev_seeds"] if split == "dev" else protocol["test_seeds"][kind]) } if ( not isinstance(entries, list) or len(entries) != len(expected) or {(row["kind"], row["seed"]) for row in entries} != expected ): raise ValueError("Incomplete or mismatched resumed suite manifest") model_hash = checkpoint_sha256(model) selection_hash = sha256_file(selection) if selection else None if split == "test": decision, _ = evaluation._verified_selection(journal.plan["evaluation_plan"], selection) if model_hash not in {protocol["baseline"]["sha256"], decision["selected"]["sha256"]}: raise ValueError("Final suite uses a checkpoint outside frozen DEV selection") for row in entries: kind, seed = row["kind"], row["seed"] name = f"{kind}-{seed}" spec_path = ( Path(journal.plan["evaluation_plan"]).parent / "specifications" / label / split / f"{name}.json" ) report = ( Path(journal.plan["evaluation_plan"]).parent / "evaluations" / label / split / name / "report.json" ) if ( row["split"] != split or row["plan_sha256"] != digest or row["selection_sha256"] != selection_hash or row["spec_path"] != str(spec_path) or row["path"] != str(report) or sha256_file(spec_path) != row["spec_sha256"] ): raise ValueError("Resumed suite specification identity changed") spec = _read(spec_path) frozen_model = spec["model"] if kind == "agent" else spec["adapter"] data = protocol["datasets"][f"{kind}_{split}"] dataset = spec["tasks"] if kind == "agent" else spec["benchmark"] sampling = ( protocol["sampling"] | {"seed": seed} | ({"max_new_tokens": 768} if kind == "math" else {}) ) execution = spec if kind == "agent" else spec["adapter"] mode = ( spec["chat_template_kwargs"] if kind == "agent" else spec["model"]["chat_template_kwargs"] ) if ( frozen_model["path"] != str(Path(model).resolve()) or frozen_model["sha256"] != model_hash or dataset["sha256"] != data["sha256"] or dataset["path"] != data["path"] or spec["sampling"] != sampling or spec["trials"] != 1 or execution["device"] != "cuda" or execution["dtype"] != "bfloat16" or mode != protocol["chat_template_kwargs"] or ( kind == "agent" and ( spec["tokenizer"] != protocol["tokenizer"] or spec["max_completion_length"] != 768 or spec["max_tool_calling_iterations"] != 3 ) ) or ( kind == "math" and ( spec["model"]["revision"] != model_hash or spec["adapter"]["tokenizer_sha256"] != protocol["tokenizer"]["sha256"] or spec["adapter"]["tokenizer_path"] != protocol["tokenizer"]["path"] or spec["model"]["chat_template"] != protocol["chat_template"] or spec["benchmark"]["count"] != data["count"] ) ) ): raise ValueError("Resumed suite does not bind the intended model/data/budget") def _audit_merged(base, merged, adapter): """Compare saved BF16 arrays on CPU; never instantiate or forward a model.""" import re import torch from safetensors import safe_open base, merged, adapter = map(Path, (base, merged, adapter)) with ExitStack() as stack: states = [] for directory in (base, merged): state = {} for file in sorted(directory.glob("*.safetensors")): opened = stack.enter_context(safe_open(file, framework="pt", device="cpu")) for key in opened.keys(): if key in state: raise ValueError("Duplicate saved model tensor key") state[key] = opened states.append(state) if not states[0] or states[0].keys() != states[1].keys(): raise ValueError("Merged/base saved tensor coverage differs") saved_adapter = stack.enter_context( safe_open(adapter / "adapter_model.safetensors", framework="pt", device="cpu") ) targets, sides = set(), {"A": set(), "B": set()} for key in saved_adapter.keys(): match = re.fullmatch(r"base_model\.model\.(.+)\.lora_([AB])\.weight", key) if match is None: raise ValueError("Unsupported saved LoRA tensor identity for BF16 audit") target = match[1] + ".weight" targets.add(target) sides[match[2]].add(target) if not torch.isfinite(saved_adapter.get_tensor(key)).all(): raise ValueError("Nonfinite saved adapter tensor") if not targets or sides["A"] != sides["B"] or not targets <= states[0].keys(): raise ValueError("Saved adapter targets do not match the base model") changed, unchanged_non_targets, maximum = 0, 0, 0.0 for key in states[0]: before, after = [state[key].get_tensor(key) for state in states] if ( before.shape != after.shape or before.dtype != after.dtype or not all(bool(torch.isfinite(tensor).all()) for tensor in (before, after)) ): raise ValueError("Nonfinite or incompatible saved merged tensor") identical = torch.equal(before, after) if key not in targets: if not identical: raise ValueError("A non-LoRA target changed during merging") unchanged_non_targets += 1 else: if after.dtype != torch.bfloat16: raise ValueError("Long experiment requires a saved BF16 merge") changed += not identical if not identical: maximum = max(maximum, float((after.float() - before.float()).abs().max())) if not changed: raise ValueError("All adapter updates vanished in the saved BF16 merge") return { "status": "passed", "method": "Elementwise saved safetensors comparison on CPU; no model forward", "all_finite": True, "model_tensors": len(states[0]), "target_tensors": len(targets), "changed_target_tensors": changed, "unchanged_target_tensors": len(targets) - changed, "identical_non_target_tensors": unchanged_non_targets, "max_abs_saved_delta": maximum, "base_sha256": checkpoint_sha256(base), "merged_sha256": checkpoint_sha256(merged), "adapter_sha256": checkpoint_sha256(adapter), } def _verify_ancestry(journal, checkpoint, step): from .gym_training_long import validate_checkpoint sealed = validate_checkpoint(journal.root / "training", checkpoint, expected_step=step) exported = journal.root / "adapters" / f"step{step}" merged = journal.root / "merged" / f"step{step}" ancestry = _read(exported.with_suffix(".json")) merge_path = merged.with_name(merged.name + ".merge-manifest.json") merge, audit = _read(merge_path), _read(merged.with_suffix(".bf16-audit.json")) base = _read(journal.plan["training_config"])["model"] export_hash, merged_hash = checkpoint_sha256(exported), checkpoint_sha256(merged) if ( ancestry["optimizer_step"] != step or ancestry["checkpoint"] != sealed or ancestry["export_sha256"] != export_hash or ancestry["base"] != base or ancestry["merge_manifest_sha256"] != sha256_file(merge_path) or ancestry["merged_sha256"] != merged_hash or merge["status"] != "completed" or merge["method"] != "peft_lora_merge" or merge["base_model"] != base or merge["adapter"]["path"] != str(exported) or merge["adapter"]["sha256"] != export_hash or merge["adapter"]["base_pointer_used"] is not False or merge["output"] != {"path": str(merged), "sha256": merged_hash} or merge["dtype"] != "bfloat16" or audit["status"] != "passed" or audit["all_finite"] is not True or audit["base_sha256"] != base["sha256"] or audit["merged_sha256"] != merged_hash or audit["adapter_sha256"] != export_hash or audit["changed_target_tensors"] < 1 or ancestry["bf16_audit_sha256"] != sha256_file(merged.with_suffix(".bf16-audit.json")) ): raise ValueError("Checkpoint/export/merge/BF16 ancestry binding mismatch") for name in ("adapter_config.json", "adapter_model.safetensors"): if sha256_file(Path(checkpoint) / name) != sha256_file(exported / name): raise ValueError("Exported adapter differs from sealed native checkpoint") proof = _verify_resume_proof(journal) if ( _read(journal.root / "training/runtime-binding.json") != proof["runtime_binding"] or _read(journal.root / "training/training-manifest.json")["gpu"] != proof["gpu"] ): raise ValueError("Main training runtime differs from qualified native resume runtime") return ancestry def _export_and_merge(journal, checkpoint, step): from .gym_training_long import validate_checkpoint sealed = validate_checkpoint(journal.root / "training", checkpoint, expected_step=step) source = Path(checkpoint) exported = journal.root / "adapters" / f"step{step}" ancestry = exported.with_suffix(".json") merged = journal.root / "merged" / f"step{step}" merge_manifest = merged.with_name(merged.name + ".merge-manifest.json") name = f"merge-{step}" if journal.complete(name): _verify_ancestry(journal, checkpoint, step) return merged _archive(journal.root, exported) _archive(journal.root, ancestry) _archive(journal.root, merged) _archive(journal.root, merge_manifest) _archive(journal.root, merged.with_suffix(".bf16-audit.json")) exported.mkdir(parents=True) for name in ("adapter_config.json", "adapter_model.safetensors"): shutil.copy2(source / name, exported / name) base = _read(journal.plan["training_config"])["model"]["path"] journal.command( f"merge-{step}", [ "-m", "hessian_recipes.checkpoints", base, str(exported), str(merged), "--dtype", "bfloat16", "--device", "cuda", ], ) audit_path = merged.with_suffix(".bf16-audit.json") _write(audit_path, _audit_merged(base, merged, exported)) _write( ancestry, { "optimizer_step": step, "checkpoint": sealed, "export_sha256": checkpoint_sha256(exported), "base": _read(journal.plan["training_config"])["model"], "merge_manifest_sha256": sha256_file(merge_manifest), "merged_sha256": checkpoint_sha256(merged), "bf16_audit_sha256": sha256_file(audit_path), }, ) _verify_ancestry(journal, checkpoint, step) journal.done(f"merge-{step}", [ancestry, merge_manifest, audit_path]) return merged def run(plan_path, gym_python, *, resume=False, qualification_only=False): from . import gym_training_long as training from . import gym_experiment_evaluation as evaluation root = Path(verify_plan(plan_path)["root"]) with (root / ".experiment.lock").open("a") as lock: try: fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB) except BlockingIOError: raise RuntimeError("Another driver owns this experiment") from None journal = _Journal(plan_path, resume) try: if journal.complete("summary"): _verify_final_summary(journal) journal.state.update(status="completed", summary=str(root / "summary/summary.json")) return journal.state preflight = root / f"gpu-preflight-{time.time_ns()}.json" journal.command( "gpu-preflight", [ str(Path(__file__).resolve().parents[2] / "scripts/gpu-preflight.py"), "--output", str(preflight), ], ) journal.done(f"gpu-preflight-attempt-{journal.state['active_attempt']}", [preflight]) with _service(root, gym_python) as (gym_url, identity): _write(root / f"gym-identity-{time.time_ns()}.json", identity) _resume_qualification(journal, gym_url) if qualification_only: journal.state.update(status="qualification_completed") return journal.state protocol = _read(journal.plan["evaluation_plan"]) baseline = Path(protocol["baseline"]["path"]) baseline_dev = _suite(journal, baseline, "baseline", "dev") candidates = {} for step in CHECKPOINTS: checkpoint = root / "training" / f"checkpoint-{step}" if not (checkpoint / training.SEAL).exists(): previous = training.discover_resume_checkpoint(root / "training") args = [ "-m", "hessian_recipes.gym_training_long", journal.plan["training_config"], "--gym-url", gym_url, "--stop-after-step", str(step), ] if previous is not None: args += ["--resume-from-checkpoint", str(previous)] journal.command(f"train-to-{step}", args) training.validate_checkpoint(root / "training", checkpoint, expected_step=step) model = _export_and_merge(journal, checkpoint, step) candidates[str(step)] = str(_suite(journal, model, f"step{step}", "dev")) selection = root / "selection/selection.json" if not journal.complete("selection"): _archive(root, selection.parent) evaluation.select( journal.plan["evaluation_plan"], baseline_dev, candidates, selection.parent ) journal.done("selection", [selection]) evaluation._verified_selection(journal.plan["evaluation_plan"], selection) baseline_test = _suite(journal, baseline, "baseline", "test", selection) decision = _read(selection) candidate_test = _suite( journal, decision["selected"]["path"], "candidate", "test", selection ) summary = root / "summary" if not journal.complete("summary"): _archive(root, summary) result = evaluation.summarize( journal.plan["evaluation_plan"], selection, baseline_test, candidate_test, summary, ) result["experiment"] = _final_chain(journal) _write(summary / "summary.json", result) with (summary / "summary.md").open("a") as stream: stream.write( "\nCheckpoint ancestry and saved BF16 changes are bound in summary.json. " "The actual GPU stop/resume qualification passed before main training.\n" ) journal.done("summary", [summary / "summary.json", summary / "summary.md"]) _verify_final_summary(journal) journal.state.update( status="completed", summary=str(summary / "summary.json"), completed_at=datetime.now(timezone.utc).isoformat(), ) return journal.state except BaseException as error: journal.state.update( status="failed", error={"type": type(error).__name__, "message": str(error)} ) raise finally: journal.close() def _final_chain(journal): proof = _verify_resume_proof(journal) ancestry = { str(step): _verify_ancestry(journal, journal.root / "training" / f"checkpoint-{step}", step) for step in CHECKPOINTS } return { "format": FORMAT, "plan_sha256": sha256_file(journal.plan_path), "training_config_sha256": journal.plan["training_config_sha256"], "train_sha256": journal.plan["train_sha256"], "source_sha256": journal.plan["source_sha256"], "resume_qualification_sha256": sha256_file( journal.root / "resume-qualification/proof.json" ), "qualified_runtime_binding": proof["runtime_binding"], "checkpoint_ancestry": ancestry, "budget_limit_seconds": journal.plan["max_gpu_wall_seconds"], } def _verify_final_summary(journal): from . import gym_experiment_evaluation as evaluation summary = _read(journal.root / "summary/summary.json") selection_path = journal.root / "selection/selection.json" selection, digest = evaluation._verified_selection( journal.plan["evaluation_plan"], selection_path ) chain = _final_chain(journal) selected_model = { "path": str(journal.root / "merged" / f"step{selection['selected_step']}"), "sha256": chain["checkpoint_ancestry"][str(selection["selected_step"])]["merged_sha256"], } if ( summary["status"] != "complete" or summary["release_eligible"] is not False or summary["plan_sha256"] != journal.plan["evaluation_plan_sha256"] or summary["selection_sha256"] != digest or summary["selected"] != selection["selected"] or summary["selected_step"] != selection["selected_step"] or summary["dev_acceptance"] != selection["dev_acceptance"] or summary["experiment"] != chain or selection["selected"] != selected_model ): raise ValueError("Final summary differs from verified selection/checkpoint ancestry") for reports in summary["evidence"].values(): for report in reports: if sha256_file(report["path"]) != report["sha256"]: raise ValueError("Final summary report evidence changed") return summary def _resume_qualification(journal, gym_url): """Require real GPU full-vs-resumed equality before spending the long-run budget.""" proof = journal.root / "resume-qualification/proof.json" if journal.complete("resume-qualification"): return _verify_resume_proof(journal) root = proof.parent if root.exists(): _archive(journal.root, root) root.mkdir(parents=True) data = root / "train.jsonl" data.write_text( "".join( json.dumps(row, ensure_ascii=False) + "\n" for row in generate_tasks("train", 4, seed=43) ) ) config = _read(journal.plan["training_config"]) config.update(seed=43, dataset={"path": str(data), "sha256": sha256_file(data)}) config["training"]["max_steps"] = 4 for mode in ("continuous", "resumed"): config["output_dir"] = str(root / mode) _write(root / f"{mode}.json", config) common = ["--gym-url", gym_url, "--checkpoint-interval", "2"] for name, mode, extra in ( ("continuous", "continuous", []), ("stopped", "resumed", ["--stop-after-step", "2"]), ("resumed", "resumed", ["--resume-from-checkpoint", str(root / "resumed/checkpoint-2")]), ): if name == "resumed": _verify_resume_prefix(root) journal.command( f"resume-check-{name}", [ "-m", "hessian_recipes.gym_training_long", str(root / f"{mode}.json"), *common, *extra, ], ) _compare_resume(journal, root, proof) journal.done("resume-qualification", [proof]) return _verify_resume_proof(journal) def _verify_resume_proof(journal): from .gym_training_long import validate_checkpoint root = journal.root / "resume-qualification" proof = _read(root / "proof.json") config = _read(journal.plan["training_config"]) task_rows = [ json.loads(line) for line in (root / "train.jsonl").read_text().splitlines() if line.strip() ] if task_rows != generate_tasks("train", 4, seed=43): raise ValueError("GPU resume qualification task identity changed") if ( proof["status"] != "passed" or proof["adapter_optimizer_scheduler_equal"] is not True or proof["exact_rollout_evidence_equal"] is not True or proof["rng_state_equal"] is not True or proof["source_sha256"] != journal.plan["source_sha256"] or proof["experiment_plan_sha256"] != sha256_file(journal.plan_path) or proof["base_model"] != config["model"] or len(proof["checkpoints"]) != 2 ): raise ValueError("GPU resume qualification belongs to another experiment or failed") for index, mode in enumerate(("continuous", "resumed")): path = root / mode / "checkpoint-4" actual = validate_checkpoint(path.parent, path, expected_step=4) binding = _read(path.parent / "runtime-binding.json") manifest = _read(path.parent / "training-manifest.json") mode_config = _read(root / f"{mode}.json") if ( actual != proof["checkpoints"][index] or binding != proof["runtime_binding"] or manifest["status"] != "completed" or manifest["optimizer_step"] != 4 or manifest["gpu"] != proof["gpu"] or manifest["model"] != config["model"] or mode_config["training"] != config["training"] | {"max_steps": 4} or mode_config["seed"] != 43 or mode_config["model"] != config["model"] or mode_config["lora"] != config["lora"] or mode_config["chat_template_kwargs"] != config["chat_template_kwargs"] or mode_config["dataset"] != {"path": str(root / "train.jsonl"), "sha256": sha256_file(root / "train.jsonl")} or mode_config["output_dir"] != str(root / mode) or sha256_file(path.parent / "rollout-evidence.jsonl") != proof["evidence_sha256"] ): raise ValueError("GPU resume qualification checkpoint/runtime/configuration changed") return proof def _assert_checkpoint_states_equal(paths, label): import torch import numpy as np from transformers.trainer import safe_globals def equal(a, b): if type(a) is not type(b): return False if torch.is_tensor(a): return a.dtype == b.dtype and torch.equal(a, b) if isinstance(a, np.ndarray): return a.dtype == b.dtype and np.array_equal(a, b) if isinstance(a, dict): return a.keys() == b.keys() and all(equal(a[k], b[k]) for k in a) if isinstance(a, (list, tuple)): return len(a) == len(b) and all(equal(x, y) for x, y in zip(a, b, strict=True)) return a == b for name in ("trainable-state.pt", "optimizer.pt", "scheduler.pt"): states = [torch.load(path / name, map_location="cpu", weights_only=True) for path in paths] if not equal(*states): raise RuntimeError(f"{label} equivalence failed for {name}") with safe_globals(): rng_states = [ torch.load(path / "rng_state.pth", map_location="cpu", weights_only=True) for path in paths ] if not equal(*rng_states): raise RuntimeError(label + " equivalence failed for RNG state") def _verify_resume_prefix(root): """Reject independent numerical drift before spending GPU time on restoration.""" from .gym_training_long import validate_checkpoint paths = [root / mode / "checkpoint-2" for mode in ("continuous", "resumed")] seals = [validate_checkpoint(path.parent, path, expected_step=2) for path in paths] _assert_checkpoint_states_equal(paths, "Independent first-two-step") if seals[0]["evidence_prefix"] != seals[1]["evidence_prefix"]: raise RuntimeError("Independent first-two-step runs changed rollout evidence") _write(root / "prefix-equivalence.json", { "status": "passed", "optimizer_step": 2, "adapter_optimizer_scheduler_rng_equal": True, "exact_rollout_prefix_equal": True, "evidence_prefix": seals[0]["evidence_prefix"], }) def _compare_resume(journal, root, proof): from .gym_training_long import validate_checkpoint paths = [root / mode / "checkpoint-4" for mode in ("continuous", "resumed")] seals = [validate_checkpoint(path.parent, path, expected_step=4) for path in paths] _assert_checkpoint_states_equal(paths, "Actual GPU resume") evidence = [sha256_file(path.parent / "rollout-evidence.jsonl") for path in paths] if len(set(evidence)) != 1: raise RuntimeError("Actual GPU resume changed sampled tokens, contexts, tools or rewards") _write( proof, { "status": "passed", "method": "Actual same-GPU full4 vs stop2/resume4", "adapter_optimizer_scheduler_equal": True, "exact_rollout_evidence_equal": True, "rng_state_equal": True, "evidence_sha256": evidence[0], "checkpoints": seals, "source_sha256": journal.plan["source_sha256"], "experiment_plan_sha256": sha256_file(journal.plan_path), "base_model": _read(journal.plan["training_config"])["model"], "runtime_binding": _read(root / "continuous/runtime-binding.json"), "gpu": _read(root / "continuous/training-manifest.json")["gpu"], "limitation": "Same seed/hardware integration check; not a model-quality result.", }, ) def main(): parser = argparse.ArgumentParser(description=__doc__) commands = parser.add_subparsers(dest="command", required=True) prep = commands.add_parser("prepare") for key in ("baseline", "tokenizer", "math-dev", "mgsm", "output"): prep.add_argument("--" + key, type=Path, required=True) prep.add_argument("--budget-carry", type=Path) execute = commands.add_parser("run") execute.add_argument("--plan", type=Path, required=True) execute.add_argument("--gym-python", type=Path, required=True) execute.add_argument("--resume", action="store_true") execute.add_argument("--qualification-only", action="store_true") args = parser.parse_args() result = ( prepare(args.baseline, args.tokenizer, args.math_dev, args.mgsm, args.output, budget_carry=args.budget_carry) if args.command == "prepare" else run( args.plan, args.gym_python, resume=args.resume, qualification_only=args.qualification_only, ) ) print(json.dumps(result, indent=2)) if __name__ == "__main__": main()