"""CPU driver tests. Mocked stage execution is not GPU qualification evidence.""" import copy from contextlib import contextmanager import json from pathlib import Path from types import SimpleNamespace import pytest from hessian_recipes import gym_experiment as driver from hessian_recipes import gym_experiment_evaluation as evaluation from hessian_recipes import gym_training_long as training def write_rows(path, rows): path.write_text("".join(json.dumps(row) + "\n" for row in rows)) @pytest.fixture def prepared(tmp_path): baseline, tokenizer = tmp_path / "baseline", tmp_path / "tokenizer" baseline.mkdir() tokenizer.mkdir() (baseline / "model.safetensors").write_bytes(b"fixture baseline, not model weights") (baseline / "config.json").write_text("{}") (tokenizer / "chat_template.jinja").write_text("{{ messages }}") train, dev, mgsm = (tmp_path / name for name in ("train.jsonl", "dev.jsonl", "mgsm.jsonl")) write_rows(train, [{"id": "sft1", "source_id": "gsm8k/train/00000"}]) write_rows( dev, [ { "id": f"dev{i}", "source_id": f"gsm8k/train/{i + 1000:05d}", "prompt": f"Aufgabe {i}", "solution": str(i), } for i in range(128) ], ) write_rows( mgsm, [ { "id": f"mgsm-de-test-{i:03d}", "source_id": f"mgsm/de/test/{i:03d}", "prompt": f"Prüfung {i}", "answer": str(i), "scorer": "final_numeric", "category": "German math", } for i in range(250) ], ) root = tmp_path / "experiment" plan = driver.prepare(baseline, tokenizer, dev, mgsm, root) return SimpleNamespace( root=root, baseline=baseline, plan=plan, path=root / "experiment-plan.json" ) def test_preparation_pins_only_runtime_sources_and_exact_long_schedule(prepared): plan = driver.verify_plan(prepared.path) sources = plan["source_sha256"] assert "src/hessian_recipes/checkpoints.py" in sources assert "scripts/gpu-preflight.py" in sources assert "src/hessian_recipes/german_data.py" not in sources assert all("hub" not in name for name in sources) config = driver._read(plan["training_config"]) assert config["training"]["max_steps"] == 256 and config["training"]["beta"] == 0.1 assert config["training"]["learning_rate"] == 5e-6 assert len(Path(prepared.root / "train.jsonl").read_text().splitlines()) == 256 def test_crash_recovery_conservatively_charges_unclosed_time(prepared, monkeypatch): clock = SimpleNamespace(wall=1000.0, monotonic=10.0) monkeypatch.setattr(driver.time, "time", lambda: clock.wall) monkeypatch.setattr(driver.time, "monotonic", lambda: clock.monotonic) journal = driver._Journal(prepared.path, False) clock.wall += 60 clock.monotonic += 60 journal.save() # SIGKILL: no finally/close. Unknown downtime must not reset the budget. clock.wall += 3600 clock.monotonic += 3600 recovered = driver._Journal(prepared.path, True) assert recovered.state["gpu_wall_seconds"] == 3660 assert recovered.state["attempts"][0]["conservative_unclosed_seconds"] == 3600 recovered.close() clock.wall += 3600 clock.monotonic += 3600 orderly = driver._Journal(prepared.path, True) assert orderly.state["gpu_wall_seconds"] == 3660 orderly.close() def test_exhausted_budget_and_tampered_done_stage_fail_before_subprocess(prepared, monkeypatch): journal = driver._Journal(prepared.path, False) artifact = prepared.root / "artifact.json" artifact.write_text("verified") journal.done("stage", [artifact]) assert journal.complete("stage") journal.state["gpu_wall_seconds"] = 8 * 3600 + 1 monkeypatch.setattr(driver.subprocess, "run", lambda *a, **k: pytest.fail("must not execute")) with pytest.raises(TimeoutError, match="eight-hour"): journal.command("out-of-budget", ["unused"]) journal.close() before = journal.path.read_bytes() artifact.write_text("changed") with pytest.raises(ValueError, match="evidence changed"): driver._Journal(prepared.path, True) assert journal.path.read_bytes() == before def test_stage_commands_have_independent_process_group_deadline(prepared, monkeypatch): journal = driver._Journal(prepared.path, False) calls = [] monkeypatch.setattr(driver.shutil, "which", lambda name: "/usr/bin/timeout") monkeypatch.setattr( driver.subprocess, "run", lambda args, **kwargs: calls.append((args, kwargs)) ) journal.command("fixture", ["-c", "pass"]) args, kwargs = calls[0] assert args[:2] == ["/usr/bin/timeout", "--signal=KILL"] assert 0 < float(args[2][:-1]) <= 8 * 3600 assert kwargs["start_new_session"] and kwargs["check"] journal.close() def test_resume_requires_existing_journal(prepared): with pytest.raises(ValueError, match="without an existing"): driver._Journal(prepared.path, True) def test_qualification_only_returns_the_durable_closed_journal(prepared, monkeypatch): events = [] @contextmanager def service(*args): events.append("start") try: yield "http://127.0.0.1:1", {"explicit_unit_fixture": True} finally: events.append("stop") def command(self, name, args): assert name == "gpu-preflight" driver._write(Path(args[args.index("--output") + 1]), {"explicit_unit_fixture": True}) def qualify(journal, gym_url): proof = journal.root / "resume-qualification/proof.json" driver._write(proof, {"explicit_unit_fixture": True}) journal.done("resume-qualification", [proof]) monkeypatch.setattr(driver, "_service", service) monkeypatch.setattr(driver._Journal, "command", command) monkeypatch.setattr(driver, "_resume_qualification", qualify) returned = driver.run(prepared.path, "not-executed", qualification_only=True) durable = driver._read(prepared.root / "experiment-status.json") assert returned == durable and durable["status"] == "qualification_completed" assert durable["active_attempt"] is None and durable["gpu_wall_seconds"] >= 0 assert durable["attempts"][-1]["status"] == "qualification_completed" assert "resume-qualification" in durable["completed"] and events == ["start", "stop"] def test_completed_resume_revalidates_summary_without_starting_any_service(prepared, monkeypatch): journal = driver._Journal(prepared.path, False) summary = prepared.root / "summary/summary.json" driver._write(summary, {"explicit_unit_fixture": True}) journal.done("summary", [summary]) journal.state["status"] = "completed" journal.close() checked = [] monkeypatch.setattr(driver, "_verify_final_summary", lambda journal: checked.append(True)) monkeypatch.setattr( driver._Journal, "command", lambda *a: pytest.fail("no model stage expected") ) monkeypatch.setattr(driver, "_service", lambda *a: pytest.fail("no service expected")) result = driver.run(prepared.path, "not-executed", resume=True) assert checked == [True] and result["status"] == "completed" assert result == driver._read(journal.path) and result["active_attempt"] is None def test_mismatched_resumed_spec_manifest_fails_before_inference(prepared, monkeypatch): journal = driver._Journal(prepared.path, False) protocol = prepared.plan["evaluation_plan"] entries = evaluation.write_specs(protocol, prepared.baseline, "baseline", "dev") manifest = Path(entries[0]["spec_path"]).parent / "manifest.json" entries[0]["seed"] = 999 driver._write(manifest, entries) monkeypatch.setattr(journal, "command", lambda *a: pytest.fail("must fail before model stage")) with pytest.raises(ValueError, match="suite manifest"): driver._suite(journal, prepared.baseline, "baseline", "dev") journal.close() def test_complete_evaluation_stages_reuse_only_verified_journal_files(prepared, monkeypatch): journal = driver._Journal(prepared.path, False) commands = [] def command(name, args): commands.append(name) output = Path(args[args.index("--output") + 1]) driver._write(output / "report.json", {"explicit_unit_fixture": True}) monkeypatch.setattr(journal, "command", command) monkeypatch.setattr(evaluation.agent, "verify_report", lambda report: None) monkeypatch.setattr(evaluation.numeric, "_verify_report", lambda report: None) monkeypatch.setattr(evaluation, "_load_suite", lambda *a, **k: None) manifest = driver._suite(journal, prepared.baseline, "baseline", "dev") assert len(commands) == 2 driver._suite(journal, prepared.baseline, "baseline", "dev") assert len(commands) == 2 entry = driver._read(manifest)[0] Path(entry["path"]).write_text('{"explicit_unit_fixture":false}') with pytest.raises(ValueError, match="evidence changed"): driver._suite(journal, prepared.baseline, "baseline", "dev") journal.close() @pytest.fixture def ancestry(prepared, monkeypatch): journal = driver._Journal(prepared.path, False) binding, gpu = {"dependencies": {"fixture": "not-GPU"}}, "fixture GPU" driver._write(prepared.root / "training/runtime-binding.json", binding) driver._write(prepared.root / "training/training-manifest.json", {"gpu": gpu}) monkeypatch.setattr( driver, "_verify_resume_proof", lambda journal: {"runtime_binding": binding, "gpu": gpu} ) def validate(output, checkpoint, expected_step=None): path = Path(checkpoint) return { "optimizer_step": expected_step, "checkpoint": str(path), "checkpoint_sha256": driver.checkpoint_sha256(path), "seal_fixture": True, } monkeypatch.setattr(training, "validate_checkpoint", validate) commands = [] def command(name, args): commands.append(name) base, exported, merged = map(Path, args[2:5]) merged.mkdir(parents=True) (merged / "model.safetensors").write_bytes(b"fixture merged " + exported.name.encode()) driver._write( merged.with_name(merged.name + ".merge-manifest.json"), { "status": "completed", "method": "peft_lora_merge", "base_model": {"path": str(base), "sha256": driver.checkpoint_sha256(base)}, "adapter": { "path": str(exported), "sha256": driver.checkpoint_sha256(exported), "base_pointer_used": False, }, "output": {"path": str(merged), "sha256": driver.checkpoint_sha256(merged)}, "dtype": "bfloat16", "device": "cuda", }, ) monkeypatch.setattr(journal, "command", command) monkeypatch.setattr( driver, "_audit_merged", lambda base, merged, adapter: { "status": "passed", "all_finite": True, "changed_target_tensors": 1, "base_sha256": driver.checkpoint_sha256(base), "merged_sha256": driver.checkpoint_sha256(merged), "adapter_sha256": driver.checkpoint_sha256(adapter), }, ) for step in driver.CHECKPOINTS: checkpoint = prepared.root / "training" / f"checkpoint-{step}" checkpoint.mkdir() driver._write(checkpoint / "adapter_config.json", {"peft_type": "LORA"}) (checkpoint / "adapter_model.safetensors").write_bytes(f"fixture adapter {step}".encode()) return journal, commands @pytest.mark.parametrize("tamper", ["export", "merged", "checkpoint"]) def test_checkpoint_export_merge_ancestry_and_stage_reuse(ancestry, tamper): journal, commands = ancestry checkpoint = journal.root / "training/checkpoint-64" merged = driver._export_and_merge(journal, checkpoint, 64) recorded = driver._read(journal.root / "adapters/step64.json") # Actual public validator API uses checkpoint_sha256, never a sha256 alias. assert recorded["checkpoint"]["checkpoint_sha256"] == driver.checkpoint_sha256(checkpoint) assert "sha256" not in recorded["checkpoint"] assert driver._export_and_merge(journal, checkpoint, 64) == merged and commands == ["merge-64"] target = { "export": journal.root / "adapters/step64/adapter_model.safetensors", "merged": merged / "model.safetensors", "checkpoint": checkpoint / "adapter_model.safetensors", }[tamper] target.write_bytes(b"tampered actual artifact") with pytest.raises(ValueError, match="ancestry"): driver._export_and_merge(journal, checkpoint, 64) journal.close() def test_final_summary_must_bind_all_checkpoint_ancestry_and_selection(ancestry, monkeypatch): journal, _ = ancestry for step in driver.CHECKPOINTS: driver._export_and_merge(journal, journal.root / "training" / f"checkpoint-{step}", step) proof = journal.root / "resume-qualification/proof.json" driver._write(proof, {"explicit_unit_fixture": True}) selection = { "selected": { "path": str(journal.root / "merged/step256"), "sha256": driver.checkpoint_sha256(journal.root / "merged/step256"), }, "selected_step": 256, "dev_acceptance": False, } monkeypatch.setattr(evaluation, "_verified_selection", lambda *a: (selection, "selection-hash")) report = journal.root / "fixture-report.json" driver._write(report, {"fixture": True}) summary = { "status": "complete", "release_eligible": False, "plan_sha256": journal.plan["evaluation_plan_sha256"], "selection_sha256": "selection-hash", **selection, "experiment": driver._final_chain(journal), "evidence": {"candidate": [{"path": str(report), "sha256": driver.sha256_file(report)}]}, } target = journal.root / "summary/summary.json" driver._write(target, summary) assert driver._verify_final_summary(journal)["selected_step"] == 256 summary["experiment"]["checkpoint_ancestry"]["128"]["optimizer_step"] = 64 driver._write(target, summary) with pytest.raises(ValueError, match="summary differs"): driver._verify_final_summary(journal) journal.close() @pytest.mark.parametrize("tamper", ["source", "runtime", "step", "config", "evidence"]) def test_resume_qualification_cannot_be_reused_after_identity_change(prepared, monkeypatch, tamper): journal = driver._Journal(prepared.path, False) base_config = driver._read(journal.plan["training_config"]) root = journal.root / "resume-qualification" root.mkdir() write_rows(root / "train.jsonl", driver.generate_tasks("train", 4, seed=43)) binding = {"dependencies": {"fixture": "not actual GPU qualification"}} def validate(output, checkpoint, expected_step=None): assert expected_step == 4 return { "optimizer_step": 4, "checkpoint": str(checkpoint), "checkpoint_sha256": driver.checkpoint_sha256(checkpoint), } monkeypatch.setattr(training, "validate_checkpoint", validate) seals = [] for mode in ("continuous", "resumed"): path = root / mode / "checkpoint-4" path.mkdir(parents=True) (path / "adapter_model.safetensors").write_bytes(b"fixture identical adapter") driver._write(path.parent / "runtime-binding.json", binding) driver._write( path.parent / "training-manifest.json", { "status": "completed", "optimizer_step": 4, "gpu": "fixture", "model": base_config["model"], }, ) (path.parent / "rollout-evidence.jsonl").write_bytes(b"{}\n") config = copy.deepcopy(base_config) config.update( seed=43, output_dir=str(root / mode), dataset={ "path": str(root / "train.jsonl"), "sha256": driver.sha256_file(root / "train.jsonl"), }, ) config["training"]["max_steps"] = 4 driver._write(root / f"{mode}.json", config) seals.append(validate(path.parent, path, 4)) proof = { "status": "passed", "adapter_optimizer_scheduler_equal": True, "exact_rollout_evidence_equal": True, "rng_state_equal": True, "source_sha256": journal.plan["source_sha256"], "experiment_plan_sha256": driver.sha256_file(journal.plan_path), "base_model": base_config["model"], "checkpoints": seals, "runtime_binding": binding, "gpu": "fixture", "evidence_sha256": driver.sha256_file(root / "continuous/rollout-evidence.jsonl"), } driver._write(root / "proof.json", proof) assert driver._verify_resume_proof(journal)["status"] == "passed" if tamper == "source": proof["source_sha256"] = {} driver._write(root / "proof.json", proof) elif tamper == "runtime": driver._write(root / "resumed/runtime-binding.json", {"changed": True}) elif tamper == "step": proof["checkpoints"][1]["optimizer_step"] = 2 driver._write(root / "proof.json", proof) elif tamper == "config": config = driver._read(root / "resumed.json") config["training"]["beta"] = 0.2 driver._write(root / "resumed.json", config) else: (root / "resumed/rollout-evidence.jsonl").write_bytes(b'{"changed": true}\n') with pytest.raises(ValueError, match="qualification"): driver._verify_resume_proof(journal) journal.close() @pytest.mark.parametrize("fault", [None, "unchanged", "non-target", "nonfinite"]) def test_real_saved_bf16_audit_without_model_loading(tmp_path, fault): torch = pytest.importorskip("torch") safetensors = pytest.importorskip("safetensors.torch") base, merged, adapter = (tmp_path / name for name in ("base", "merged", "adapter")) for directory in (base, merged, adapter): directory.mkdir() before = { "layer.weight": torch.ones((2, 2), dtype=torch.bfloat16), "other.weight": torch.ones((2,), dtype=torch.bfloat16), } after = {name: value.clone() for name, value in before.items()} after["layer.weight"][0, 0] = 1.125 if fault == "unchanged": after = before elif fault == "non-target": after["other.weight"][0] = 2 elif fault == "nonfinite": after["layer.weight"][0, 0] = float("nan") safetensors.save_file(before, base / "model.safetensors") safetensors.save_file(after, merged / "model.safetensors") safetensors.save_file( { "base_model.model.layer.lora_A.weight": torch.ones((1, 2)), "base_model.model.layer.lora_B.weight": torch.ones((2, 1)), }, adapter / "adapter_model.safetensors", ) if fault: with pytest.raises(ValueError): driver._audit_merged(base, merged, adapter) else: result = driver._audit_merged(base, merged, adapter) assert result["changed_target_tensors"] == 1 and result["identical_non_target_tensors"] == 1 assert result["max_abs_saved_delta"] == 0.125 and result["all_finite"] @pytest.fixture def comparison_checkpoints(tmp_path): torch = pytest.importorskip("torch") pytest.importorskip("transformers.trainer") import random import numpy as np paths = [tmp_path / mode for mode in ("continuous", "stopped")] rng = {"python": random.getstate(), "numpy": np.random.get_state(), "cpu": torch.get_rng_state()} for path in paths: path.mkdir() for name in ("trainable-state.pt", "optimizer.pt", "scheduler.pt"): torch.save({"value": torch.tensor([0.25], dtype=torch.float64)}, path / name) torch.save(rng, path / "rng_state.pth") return paths @pytest.mark.parametrize("changed", [None, "trainable-state.pt", "optimizer.pt", "scheduler.pt"]) def test_exact_checkpoint_comparison_rejects_any_numerical_drift(comparison_checkpoints, changed): import torch paths = comparison_checkpoints if changed: torch.save({"value": torch.tensor([0.25 + 1e-12], dtype=torch.float64)}, paths[1] / changed) with pytest.raises(RuntimeError, match="equivalence failed"): driver._assert_checkpoint_states_equal(paths, "fixture") else: driver._assert_checkpoint_states_equal(paths, "fixture") @pytest.mark.parametrize("changed", ["trainable-state.pt", "optimizer.pt", "scheduler.pt"]) def test_checkpoint_comparison_rejects_equal_values_with_changed_tensor_dtype( comparison_checkpoints, changed, ): import torch paths = comparison_checkpoints original = torch.tensor([0.25], dtype=torch.float64) altered = original.to(torch.float32) assert torch.equal(original, altered) and original.dtype != altered.dtype torch.save({"value": altered}, paths[1] / changed) with pytest.raises(RuntimeError, match="equivalence failed"): driver._assert_checkpoint_states_equal(paths, "fixture") @pytest.mark.parametrize("changed", ["python", "numpy", "cpu", "numpy_dtype", "cpu_dtype"]) def test_checkpoint_comparison_rejects_nested_rng_state_changes(comparison_checkpoints, changed): import torch import numpy as np from transformers.trainer import safe_globals paths = comparison_checkpoints fixture_allowlist = [] with safe_globals(): rng = torch.load(paths[1] / "rng_state.pth", map_location="cpu", weights_only=True) if changed == "python": version, state, gaussian = rng["python"] rng["python"] = (version, (state[0] ^ 1, *state[1:]), gaussian) elif changed in {"numpy", "numpy_dtype"}: state = list(rng["numpy"]) original = state[1] if changed == "numpy": state[1] = original.copy() state[1][0] ^= np.uint32(1) else: state[1] = original.astype(np.uint64) assert np.array_equal(original, state[1]) and original.dtype != state[1].dtype # Allow this fixture dtype only, so the comparison sees it after safe loading. fixture_allowlist.append(type(state[1].dtype)) rng["numpy"] = tuple(state) elif changed == "cpu": rng["cpu"] = rng["cpu"].clone() rng["cpu"][0] ^= 1 else: original = rng["cpu"] rng["cpu"] = original.to(torch.int16) assert torch.equal(original, rng["cpu"]) and original.dtype != rng["cpu"].dtype torch.save(rng, paths[1] / "rng_state.pth") with torch.serialization.safe_globals(fixture_allowlist): with pytest.raises(RuntimeError, match="equivalence failed for RNG state"): driver._assert_checkpoint_states_equal(paths, "fixture") def test_prefix_drift_stops_before_resumed_gpu_process(prepared, monkeypatch): calls = [] journal = SimpleNamespace(root=prepared.root, plan=prepared.plan, complete=lambda name: False, command=lambda name, arguments: calls.append(name)) def drift(root): raise RuntimeError("independent CUDA update drift") monkeypatch.setattr(driver, "_verify_resume_prefix", drift) with pytest.raises(RuntimeError, match="CUDA update drift"): driver._resume_qualification(journal, "http://unused-fixture") assert calls == ["resume-check-continuous", "resume-check-stopped"]