Files
OBLITERATUS/tests/test_run_archive.py
T

638 lines
23 KiB
Python

from __future__ import annotations
import hashlib
import json
import os
from pathlib import Path
from types import SimpleNamespace
import pytest
from obliteratus.checkpoint_provenance import (
ArtifactIdentity,
ToolIdentity,
build_provenance,
)
from obliteratus.run_archive import (
RunArchive,
_atomic_json,
_checkpoint_metrics,
_dataset_inputs,
_option_value,
_process_start_ticks,
_redact_arguments,
_worker,
default_archive_root,
main,
)
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_experiment_protocol_is_durable_without_raw_prompts(tmp_path):
archive = RunArchive(tmp_path / "runs")
run_id = archive.begin(["model"])
protocol = {
"protocol": "qwen38-v1",
"manifest_sha256": "a" * 64,
"counts": {"train": 500, "tune": 142, "test": 200},
"pair_ids": {"train": ["b" * 64], "tune": [], "test": []},
}
manifest = archive.record_experiment_protocol(run_id, protocol)
assert manifest["experiment_protocol"]["protocol"] == "qwen38-v1"
path = tmp_path / "runs" / run_id / "experiment-protocol.json"
assert json.loads(path.read_text(encoding="utf-8")) == protocol
def _completed_protocol_run(archive: RunArchive) -> str:
run_id = archive.begin(["org/model"])
archive.record_experiment_protocol(
run_id,
{
"protocol": "qwen38-v1",
"manifest_sha256": "a" * 64,
"counts": {"train": 500, "tune": 142, "test": 200},
"pair_ids": {"train": [], "tune": [], "test": []},
},
)
checkpoint = archive._run_dir(run_id) / "checkpoint"
checkpoint.mkdir()
(checkpoint / "weights.bin").write_bytes(b"weights")
archive.complete(run_id, checkpoint=checkpoint, metrics={})
return run_id
def test_saved_checkpoint_evaluation_is_durable_and_hashed(tmp_path):
archive = RunArchive(tmp_path)
run_id = _completed_protocol_run(archive)
reservation = archive.begin_evaluation(
run_id,
partition="optimizer_tune",
evaluator="test-v1",
)
record = archive.finish_evaluation(
run_id,
reservation["evaluation_id"],
metrics={"refusal_rate": 0.1, "coherence": 0.9},
log=["verified"],
)
assert record["status"] == "succeeded"
manifest = archive.result(run_id)
entry = manifest["evaluations"][0]
path = archive._run_dir(run_id) / entry["path"]
assert path.is_file()
assert entry["sha256"] == hashlib.sha256(path.read_bytes()).hexdigest()
def test_final_evaluation_reservation_is_single_use_even_if_interrupted(tmp_path):
archive = RunArchive(tmp_path)
run_id = _completed_protocol_run(archive)
archive.begin_evaluation(run_id, partition="final_test", evaluator="test-v1")
with pytest.raises(ValueError, match="already reserved"):
archive.begin_evaluation(run_id, partition="final_test", evaluator="test-v1")
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"
def test_archive_helpers_cover_fallbacks_redaction_and_atomic_directory_sync(
tmp_path,
monkeypatch,
):
monkeypatch.delenv("OBLITERATUS_RUN_ARCHIVE", raising=False)
monkeypatch.setenv("XDG_STATE_HOME", str(tmp_path / "state"))
assert default_archive_root() == tmp_path / "state" / "obliteratus" / "runs"
assert _redact_arguments(["model", "--api-key=value", "--revision", "safe"]) == [
"model",
"--api-key=[REDACTED]",
"--revision",
"safe",
]
assert _process_start_ticks(2**31 - 1) is None
target = tmp_path / "record.json"
monkeypatch.setattr(
"obliteratus.run_archive.os.open",
lambda *_args, **_kwargs: (_ for _ in ()).throw(OSError("unsupported")),
)
_atomic_json(target, {"ok": True})
assert json.loads(target.read_text(encoding="utf-8")) == {"ok": True}
def test_manifest_operations_record_running_log_and_revisions(tmp_path):
archive = RunArchive(tmp_path)
run_id = archive.begin(["org/model"])
running = archive.mark_running(run_id, phase="loading")
archive.append_log(run_id, "first line\n")
revised = archive.record_revisions(
run_id,
model_revision="model-sha",
tokenizer_revision="tokenizer-sha",
)
assert running["status"] == "running"
assert revised["model_revision"] == "model-sha"
assert revised["tokenizer_revision"] == "tokenizer-sha"
assert (archive._run_dir(run_id) / "run.log").read_text() == "first line\n"
def test_checkpoint_provenance_attachment_rejects_invalid_and_conflicting_records(
tmp_path,
):
archive = RunArchive(tmp_path)
run_id = archive.begin(["org/model"])
with pytest.raises(ValueError, match="canonical provenance"):
archive.attach_checkpoint_provenance(run_id, object())
invalid = SimpleNamespace(to_json=lambda: json.dumps({"artifact_id": "artifact-a"}))
with pytest.raises(ValueError, match="canonical provenance"):
archive.attach_checkpoint_provenance(run_id, invalid)
def provenance(output_digest: str):
return build_provenance(
sources=(
ArtifactIdentity(
"generated",
"run-archive-test",
"v1",
"sha256:" + "a" * 64,
),
),
converter=ToolIdentity("test", "1", "c" * 40),
obliteratus_commit="c" * 40,
configuration_digest=None,
tokenizer=None,
base_model=None,
command=("checkpoint", "attach"),
environment={"python": "test", "platform": "cpu", "packages": {}},
source_topology={},
lineage=(),
input_digests=("sha256:" + "a" * 64,),
output_digests=(output_digest,),
transformations=(),
observed_scopes=("model_weights",),
lost_state=(),
)
first = provenance("sha256:" + "b" * 64)
attached = archive.attach_checkpoint_provenance(run_id, first)
assert attached["artifact_id"] == first.artifact_id
assert attached["checkpoint_provenance"]["sha256"].startswith("sha256:")
second = provenance("sha256:" + "d" * 64)
with pytest.raises(ValueError, match="different artifact identity"):
archive.attach_checkpoint_provenance(run_id, second)
def test_evaluation_validation_and_failed_terminal_record(tmp_path):
archive = RunArchive(tmp_path)
run_id = archive.begin(["org/model"])
with pytest.raises(ValueError, match="unsupported evaluation partition"):
archive.begin_evaluation(run_id, partition="training", evaluator="test-v1")
checkpoint = archive._run_dir(run_id) / "checkpoint"
checkpoint.mkdir()
(checkpoint / "weights.bin").write_bytes(b"weights")
archive.complete(run_id, checkpoint=checkpoint, metrics={})
with pytest.raises(ValueError, match="requires qwen38-v1"):
archive.begin_evaluation(run_id, partition="optimizer_tune", evaluator="test-v1")
archive.record_experiment_protocol(
run_id,
{"protocol": "qwen38-v1", "manifest_sha256": "a" * 64, "counts": {}},
)
reservation = archive.begin_evaluation(
run_id,
partition="optimizer_tune",
evaluator="test-v1",
)
failed = archive.finish_evaluation(
run_id,
reservation["evaluation_id"],
failure=RuntimeError("Bearer abcdefghijklmnop"),
log=["failed safely"],
)
assert failed["status"] == "failed"
assert failed["failure"] == {"type": "RuntimeError", "message": "[REDACTED]"}
with pytest.raises(ValueError, match="invalid evaluation ID"):
archive.finish_evaluation(run_id, "invalid")
with pytest.raises(KeyError, match="unknown evaluation ID"):
archive.finish_evaluation(run_id, "eval-" + "e" * 32)
with pytest.raises(ValueError, match="already terminal"):
archive.finish_evaluation(run_id, reservation["evaluation_id"])
def test_complete_cancel_prune_and_worker_identity_fail_closed(tmp_path, monkeypatch):
archive = RunArchive(tmp_path)
run_id = archive.begin(["org/model"])
with pytest.raises(ValueError, match="managed run checkpoint"):
archive.complete(run_id, checkpoint=tmp_path)
with pytest.raises(ValueError, match="operator reason"):
archive.prune_checkpoint(run_id, reason=" ")
queued = archive._load(run_id)
assert archive._worker_matches({"worker": {"pid": 1}}) is False
queued["worker"] = {"pid": 9001, "uid": -1, "start_ticks": 10}
assert archive._worker_matches(queued) is False
queued["worker"] = {"pid": 9001, "uid": os.getuid(), "start_ticks": 10}
monkeypatch.setattr("obliteratus.run_archive._process_start_ticks", lambda _pid: 10)
assert archive._worker_matches(queued) is True
archive.fail(run_id, RuntimeError("failed"), phase="pipeline")
assert archive.cancel(run_id)["status"] == "failed"
live_run = archive.begin(["org/model"])
live = archive._load(live_run)
live["worker"] = {"pid": 9002, "uid": os.getuid(), "start_ticks": 11}
archive._save(live)
monkeypatch.setattr("obliteratus.run_archive._process_start_ticks", lambda _pid: 11)
monkeypatch.setattr(
"obliteratus.run_archive.os.killpg",
lambda *_args: (_ for _ in ()).throw(PermissionError("denied")),
)
with pytest.raises(RuntimeError, match="cannot cancel worker"):
archive.cancel(live_run)
def test_option_dataset_metric_and_main_edge_contracts(tmp_path, monkeypatch):
payload = tmp_path / "pairs.json"
payload.write_bytes(b"pairs")
arguments = [
"model",
"--dataset=custom",
"--prompt-pairs-file",
str(payload),
"--residue-file=missing.json",
]
assert _option_value(arguments, "--dataset") == "custom"
assert _option_value(arguments, "--absent") is None
inputs = _dataset_inputs(arguments)
assert inputs[0]["identifier"] == "custom"
assert inputs[1]["sha256"] == hashlib.sha256(b"pairs").hexdigest()
assert inputs[2]["sha256"] is None
checkpoint = tmp_path / "checkpoint"
checkpoint.mkdir()
assert _checkpoint_metrics(checkpoint) == {}
(checkpoint / "abliteration_metadata.json").write_text(
json.dumps({"metrics": ["invalid"]}),
encoding="utf-8",
)
assert _checkpoint_metrics(checkpoint) == {}
with pytest.raises(SystemExit, match="requires obliteration arguments"):
main(["worker", "--archive-root", str(tmp_path), "--run-id", "run-" + "f" * 32])
observed = []
monkeypatch.setattr(
"obliteratus.run_archive._worker",
lambda archive, run_id, args: observed.append((archive.root, run_id, args)) or 7,
)
run_id = "run-" + "f" * 32
assert (
main(
[
"worker",
"--archive-root",
str(tmp_path),
"--run-id",
run_id,
"--",
"org/model",
]
)
== 7
)
assert observed == [(tmp_path.resolve(), run_id, ["org/model"])]