mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-30 22:50:46 +02:00
330 lines
12 KiB
Python
330 lines
12 KiB
Python
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"
|