feat: add durable reconnect-safe experiment runs (#184)

tests/test_gpu_lifecycle.py

@@ -66,6 +66,43 @@ def test_disabled_publisher_is_noop():
     assert publisher.release() is None
 
 
+def test_allocation_phase_can_heartbeat_before_ready(tmp_path):
+    publisher = GpuLifecyclePublisher(tmp_path, run_id="loading-run")
+
+    def acknowledge():
+        while not (tmp_path / "current.json").exists():
+            time.sleep(0.001)
+        request = json.loads((tmp_path / "current.json").read_text())
+        (tmp_path / "ack.json").write_text(json.dumps({
+            "schema_version": 1,
+            "run_id": "loading-run",
+            "request_event_id": request["event_id"],
+            "decision": "grant",
+            "lease_id": "lease-loading",
+            "granted_vram_bytes": 100,
+        }))
+
+    thread = threading.Thread(target=acknowledge)
+    thread.start()
+    publisher.loading("org/model")
+    event = publisher.heartbeat()
+    publisher.release(reason="test_complete")
+    thread.join()
+
+    assert event is not None
+    assert event["event"] == "heartbeat"
+    assert event["run_id"] == "loading-run"
+
+
+def test_environment_uses_stable_experiment_run_id(tmp_path, monkeypatch):
+    from obliteratus.gpu_lifecycle import from_environment
+
+    monkeypatch.setenv("OBLITERATUS_GPU_LIFECYCLE_DIR", str(tmp_path))
+    monkeypatch.setenv("OBLITERATUS_RUN_ID", "run-stable")
+    publisher = from_environment()
+    assert publisher._run_id == "run-stable"
+
+
 def test_runtime_directory_must_exist(tmp_path):
     with pytest.raises(ValueError, match="must already exist"):
         GpuLifecyclePublisher(tmp_path / "missing")

tests/test_run_archive.py

@@ -0,0 +1,329 @@
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+import pytest
+
+from obliteratus.run_archive import RunArchive, _worker
+
+
+class FakeProcess:
+    def __init__(self, pid: int):
+        self.pid = pid
+
+
+def test_launch_persists_manifest_before_worker_and_redacts_secrets(tmp_path, monkeypatch):
+    archive = RunArchive(tmp_path)
+    observed = {}
+
+    def fake_popen(command, **kwargs):
+        run_id = command[command.index("--run-id") + 1]
+        observed["manifest"] = archive._load(run_id)
+        observed["command"] = command
+        observed["kwargs"] = kwargs
+        return FakeProcess(os.getpid())
+
+    monkeypatch.setattr("obliteratus.run_archive._process_start_ticks", lambda _pid: 123)
+    run_id = archive.launch(
+        ["org/model", "--token", "private-value", "--dataset", "builtin"],
+        notes="retain this experiment",
+        popen=fake_popen,
+    )
+
+    assert observed["manifest"]["status"] == "queued"
+    assert observed["manifest"]["phase"] == "manifested"
+    assert observed["manifest"]["arguments"][2] == "[REDACTED]"
+    assert observed["kwargs"]["start_new_session"] is True
+    assert (tmp_path / run_id / "notes.md").read_text() == "retain this experiment\n"
+    assert archive.status(run_id)["phase"] == "worker_started"
+
+
+def test_launch_rejects_archive_escape_and_managed_output_override(tmp_path):
+    archive = RunArchive(tmp_path)
+    with pytest.raises(ValueError, match="model ID"):
+        archive.launch(["--method", "advanced"])
+    with pytest.raises(ValueError, match="managed"):
+        archive.launch(["org/model", "--output-dir", "/tmp/escape"])
+    with pytest.raises(ValueError, match="invalid run ID"):
+        archive.status("../outside")
+
+
+def test_archive_and_run_directory_symlinks_fail_closed(tmp_path):
+    actual = tmp_path / "actual"
+    actual.mkdir()
+    alias = tmp_path / "alias"
+    alias.symlink_to(actual, target_is_directory=True)
+    with pytest.raises(ValueError, match="root must not be a symlink"):
+        RunArchive(alias)
+
+    archive = RunArchive(actual)
+    run_id = "run-" + "c" * 32
+    (actual / run_id).symlink_to(tmp_path, target_is_directory=True)
+    with pytest.raises(ValueError, match="run directory must not be a symlink"):
+        archive.status(run_id)
+
+
+def test_status_recovers_nonterminal_run_when_worker_identity_is_lost(tmp_path, monkeypatch):
+    archive = RunArchive(tmp_path)
+    monkeypatch.setattr("obliteratus.run_archive._process_start_ticks", lambda _pid: 111)
+    run_id = archive.launch(["org/model"], popen=lambda *_args, **_kwargs: FakeProcess(4242))
+    monkeypatch.setattr(archive, "_worker_matches", lambda _manifest: False)
+
+    status = archive.status(run_id)
+
+    assert status["status"] == "failed"
+    assert status["phase"] == "worker_lost"
+    assert status["failure"]["type"] == "WorkerLost"
+
+
+def test_cancel_validates_identity_then_signals_only_worker_group(tmp_path, monkeypatch):
+    archive = RunArchive(tmp_path)
+    monkeypatch.setattr("obliteratus.run_archive._process_start_ticks", lambda _pid: 222)
+    run_id = archive.launch(["org/model"], popen=lambda *_args, **_kwargs: FakeProcess(5252))
+    monkeypatch.setattr(archive, "_worker_matches", lambda _manifest: True)
+    signals = []
+    monkeypatch.setattr("obliteratus.run_archive.os.killpg", lambda pid, sig: signals.append((pid, sig)))
+
+    status = archive.cancel(run_id)
+
+    assert signals and signals[0][0] == 5252
+    assert status["status"] == "cancelling"
+    assert status["phase"] == "cancellation_requested"
+
+
+def test_result_requires_terminal_state(tmp_path, monkeypatch):
+    archive = RunArchive(tmp_path)
+    monkeypatch.setattr("obliteratus.run_archive._process_start_ticks", lambda _pid: 333)
+    run_id = archive.launch(["org/model"], popen=lambda *_args, **_kwargs: FakeProcess(6262))
+    monkeypatch.setattr(archive, "_worker_matches", lambda _manifest: True)
+    with pytest.raises(RuntimeError, match="not complete"):
+        archive.result(run_id)
+
+
+def test_pruning_preserves_notes_logs_manifest_failure_and_hashed_inventory(tmp_path):
+    archive = RunArchive(tmp_path)
+    run_id = "run-" + "a" * 32
+    run_dir = tmp_path / run_id
+    checkpoint = run_dir / "checkpoint"
+    checkpoint.mkdir(parents=True)
+    (checkpoint / "weights.bin").write_bytes(b"weights")
+    (run_dir / "notes.md").write_text("failed candidate\n")
+    (run_dir / "run.log").write_text("failure details\n")
+    manifest = {
+        "schema_version": 1,
+        "run_id": run_id,
+        "status": "failed",
+        "phase": "verify",
+        "failure": {"type": "QualityError", "message": "dominated", "phase": "verify"},
+    }
+    archive._save(manifest)
+
+    result = archive.prune_checkpoint(run_id, reason="dominated run under disk pressure")
+
+    assert not checkpoint.exists()
+    assert (run_dir / "notes.md").is_file()
+    assert (run_dir / "run.log").is_file()
+    assert (run_dir / "manifest.json").is_file()
+    inventory = json.loads((run_dir / "artifact-inventory.json").read_text())
+    assert any(item["path"] == "checkpoint/weights.bin" for item in inventory["artifacts"])
+    assert result["failure"]["message"] == "dominated"
+    assert result["checkpoint_pruned"]["reason"].startswith("dominated")
+
+
+def test_concurrent_run_ids_and_manifests_do_not_collide(tmp_path, monkeypatch):
+    archive = RunArchive(tmp_path)
+    monkeypatch.setattr("obliteratus.run_archive._process_start_ticks", lambda _pid: 444)
+    counter = iter(range(7000, 7020))
+    fake = lambda *_args, **_kwargs: FakeProcess(next(counter))
+
+    run_ids = {archive.launch(["org/model"], popen=fake) for _ in range(20)}
+
+    assert len(run_ids) == 20
+    assert {archive._load(run_id)["run_id"] for run_id in run_ids} == run_ids
+
+
+def test_checkpoint_metrics_are_returned_from_success_metadata(tmp_path):
+    from obliteratus.run_archive import _checkpoint_metrics
+
+    checkpoint = tmp_path / "checkpoint"
+    checkpoint.mkdir()
+    (checkpoint / "abliteration_metadata.json").write_text(
+        json.dumps({"quality_metrics": {"refusal_rate": 0.21, "coherence": 0.8}})
+    )
+    assert _checkpoint_metrics(checkpoint) == {"refusal_rate": 0.21, "coherence": 0.8}
+
+
+def test_list_skips_corrupt_manifests(tmp_path):
+    archive = RunArchive(tmp_path)
+    corrupt = tmp_path / ("run-" + "b" * 32)
+    corrupt.mkdir()
+    (corrupt / "manifest.json").write_text("[]")
+    assert archive.list() == []
+
+
+def test_default_root_honors_service_archive_environment(tmp_path, monkeypatch):
+    monkeypatch.setenv("OBLITERATUS_RUN_ARCHIVE", str(tmp_path / "durable"))
+    assert RunArchive().root == (tmp_path / "durable").resolve()
+
+
+def test_dataset_manifest_records_hash_and_counts_without_prompt_text(tmp_path):
+    archive = RunArchive(tmp_path)
+    run_id = archive.begin(["org/model"])
+    archive.record_dataset(
+        run_id,
+        identifier="builtin",
+        harmful=["sensitive harmful prompt"],
+        harmless=["ordinary harmless prompt"],
+    )
+    raw = (tmp_path / run_id / "manifest.json").read_text()
+    manifest = json.loads(raw)
+    assert "sensitive harmful prompt" not in raw
+    assert manifest["dataset_inputs"][0]["harmful_count"] == 1
+    assert len(manifest["dataset_inputs"][0]["sha256"]) == 64
+
+
+def test_failure_detail_redacts_huggingface_and_bearer_tokens(tmp_path):
+    archive = RunArchive(tmp_path)
+    run_id = archive.begin(["org/model"])
+    result = archive.fail(
+        run_id,
+        RuntimeError("request used hf_abcdefghijklmnopqrstuvwxyz and Bearer abcdefghijklmnop"),
+        phase="download",
+    )
+    assert result["failure"]["message"].count("[REDACTED]") == 2
+    assert "hf_" not in result["failure"]["message"]
+
+
+class FakeLifecycle:
+    def __init__(self):
+        self.events = []
+
+    def loading(self, model):
+        self.events.append(("loading", model))
+
+    def ready(self):
+        self.events.append(("ready", None))
+
+    def resize(self, memory):
+        self.events.append(("resize", memory.reserved_bytes))
+
+    def release(self, *, reason):
+        self.events.append(("release", reason))
+
+
+def _embedded_run(archive):
+    return archive.begin(["org/model"], metadata={"configuration": {"seed": 7}})
+
+
+def test_worker_success_writes_metrics_inventory_and_atomic_marker(tmp_path, monkeypatch):
+    archive = RunArchive(tmp_path)
+    run_id = _embedded_run(archive)
+    lifecycle = FakeLifecycle()
+    monkeypatch.setattr("obliteratus.run_archive.from_environment", lambda: lifecycle)
+
+    class SuccessfulChild:
+        def __init__(self, command, **_kwargs):
+            checkpoint = Path(command[command.index("--output-dir") + 1])
+            checkpoint.mkdir()
+            (checkpoint / "weights.bin").write_bytes(b"model")
+            (checkpoint / "abliteration_metadata.json").write_text(
+                json.dumps({"quality_metrics": {"refusal_rate": 0.2, "coherence": 0.8}})
+            )
+
+        def wait(self, timeout=None):
+            return 0
+
+        def poll(self):
+            return 0
+
+    monkeypatch.setattr("obliteratus.run_archive.subprocess.Popen", SuccessfulChild)
+
+    assert _worker(archive, run_id, ["org/model"]) == 0
+    result = archive.result(run_id)
+    assert result["result"]["metrics"] == {"refusal_rate": 0.2, "coherence": 0.8}
+    assert (tmp_path / run_id / "COMPLETE").is_file()
+    assert lifecycle.events == [
+        ("loading", "org/model"), ("resize", 0), ("ready", None),
+        ("release", "complete")
+    ]
+
+
+def test_worker_failure_preserves_partial_checkpoint_and_failure_detail(tmp_path, monkeypatch):
+    archive = RunArchive(tmp_path)
+    run_id = _embedded_run(archive)
+    lifecycle = FakeLifecycle()
+    monkeypatch.setattr("obliteratus.run_archive.from_environment", lambda: lifecycle)
+
+    class FailedChild:
+        def __init__(self, command, **_kwargs):
+            checkpoint = Path(command[command.index("--output-dir") + 1])
+            checkpoint.mkdir()
+            (checkpoint / "partial.bin").write_bytes(b"partial")
+
+        def wait(self, timeout=None):
+            return 17
+
+        def poll(self):
+            return 17
+
+    monkeypatch.setattr("obliteratus.run_archive.subprocess.Popen", FailedChild)
+
+    assert _worker(archive, run_id, ["org/model"]) == 1
+    result = archive.result(run_id)
+    assert result["status"] == "failed"
+    assert result["failure"]["type"] == "RuntimeError"
+    assert "status 17" in result["failure"]["message"]
+    assert (tmp_path / run_id / "checkpoint" / "partial.bin").is_file()
+    assert lifecycle.events[-1] == ("release", "failed")
+
+
+def test_worker_commits_cancelled_state_after_cooperative_interrupt(tmp_path, monkeypatch):
+    archive = RunArchive(tmp_path)
+    run_id = _embedded_run(archive)
+    lifecycle = FakeLifecycle()
+    monkeypatch.setattr("obliteratus.run_archive.from_environment", lambda: lifecycle)
+
+    class CancelledChild:
+        def __init__(self, _command, **_kwargs):
+            self.terminated = False
+
+        def wait(self, timeout=None):
+            if timeout is not None:
+                return -15
+            manifest = archive._load(run_id)
+            manifest["status"] = "cancelling"
+            archive._save(manifest)
+            raise KeyboardInterrupt("cancelled")
+
+        def poll(self):
+            return None if not self.terminated else -15
+
+        def terminate(self):
+            self.terminated = True
+
+    monkeypatch.setattr("obliteratus.run_archive.subprocess.Popen", CancelledChild)
+
+    assert _worker(archive, run_id, ["org/model"]) == 130
+    result = archive.result(run_id)
+    assert result["status"] == "cancelled"
+    assert result["phase"] == "cancelled"
+    assert lifecycle.events[-1] == ("release", "cancelled")
+
+
+def test_restart_recovery_preserves_partial_save_and_logs(tmp_path, monkeypatch):
+    archive = RunArchive(tmp_path)
+    monkeypatch.setattr("obliteratus.run_archive._process_start_ticks", lambda _pid: 555)
+    run_id = archive.launch(["org/model"], popen=lambda *_args, **_kwargs: FakeProcess(9999))
+    run_dir = tmp_path / run_id
+    (run_dir / "checkpoint").mkdir()
+    (run_dir / "checkpoint" / "partial.bin").write_bytes(b"partial")
+    (run_dir / "run.log").write_text("last durable phase\n")
+    monkeypatch.setattr("obliteratus.run_archive._process_start_ticks", lambda _pid: None)
+
+    recovered = RunArchive(tmp_path).status(run_id)
+
+    assert recovered["status"] == "failed"
+    assert recovered["failure"]["type"] == "WorkerLost"
+    assert (run_dir / "checkpoint" / "partial.bin").is_file()
+    assert (run_dir / "run.log").read_text() == "last durable phase\n"