From e8a36d796c5fdc4f073317dcc0a98e81b27ec695 Mon Sep 17 00:00:00 2001 From: Joseph Magly <1159087+jmagly@users.noreply.github.com> Date: Fri, 28 Aug 2026 23:47:20 -0400 Subject: [PATCH] feat: add durable reconnect-safe experiment runs (#184) --- app.py | 182 +++++++-- ci/test-risk-map.json | 2 + docs/deployment/shared-gpu-host.md | 25 ++ obliteratus/__init__.py | 4 + obliteratus/cli.py | 54 +++ obliteratus/gpu_lifecycle.py | 7 +- obliteratus/run_archive.py | 609 +++++++++++++++++++++++++++++ tests/test_gpu_lifecycle.py | 37 ++ tests/test_run_archive.py | 329 ++++++++++++++++ 9 files changed, 1204 insertions(+), 45 deletions(-) create mode 100644 obliteratus/run_archive.py create mode 100644 tests/test_run_archive.py diff --git a/app.py b/app.py index f8e9f08..0a737cc 100644 --- a/app.py +++ b/app.py @@ -18,6 +18,7 @@ ZeroGPU Support: from __future__ import annotations import gc +import hashlib import json as _json import os import re @@ -379,48 +380,55 @@ def _persist_session_meta(output_dir: str, label: str, meta: dict) -> None: def _recover_sessions_from_disk() -> None: - """Scan /tmp for obliterated checkpoints and repopulate _session_models. + """Scan temporary and durable checkpoints and repopulate session models. Called on startup and when a stale dropdown value is detected. Skips directories that are already registered. """ global _last_obliterated_label, _obliterate_counter found_any = False + candidates = [] for pattern in ("obliterated_*", "obliterated", "bench_*", "obliteratus_tourney/r*"): - for p in Path("/tmp").glob(pattern): - if not p.is_dir(): - continue - meta_file = p / _SESSION_META_FILE - if not meta_file.exists(): - continue + candidates.extend(Path("/tmp").glob(pattern)) + try: + from obliteratus.run_archive import default_archive_root + candidates.extend(default_archive_root().glob("run-*/checkpoint")) + except (OSError, ValueError): + pass + for p in candidates: + if not p.is_dir(): + continue + meta_file = p / _SESSION_META_FILE + if not meta_file.exists(): + continue + try: + data = _json.loads(meta_file.read_text()) + except Exception: + continue + label = data.get("label", p.name) + if label in _session_models: + continue # already registered + _session_models[label] = { + "model_id": data.get("model_id", ""), + "model_choice": data.get("model_choice", data.get("model_id", "")), + "method": data.get("method", "unknown"), + "dataset_key": data.get("dataset_key", ""), + "prompt_volume": data.get("prompt_volume", 0), + "output_dir": str(p), + "source": data.get("source", "recovered"), + "load_settings": data.get("load_settings", {}), + } + found_any = True + # Track the latest for auto-select + _last_obliterated_label = label + # Keep counter above any existing numbered dirs + if p.name.startswith("obliterated_"): try: - data = _json.loads(meta_file.read_text()) - except Exception: + idx = int(p.name.split("_", 1)[1]) + if idx >= _obliterate_counter: + _obliterate_counter = idx + 1 + except (ValueError, IndexError): continue - label = data.get("label", p.name) - if label in _session_models: - continue # already registered - _session_models[label] = { - "model_id": data.get("model_id", ""), - "model_choice": data.get("model_choice", data.get("model_id", "")), - "method": data.get("method", "unknown"), - "dataset_key": data.get("dataset_key", ""), - "prompt_volume": data.get("prompt_volume", 0), - "output_dir": str(p), - "source": data.get("source", "recovered"), - "load_settings": data.get("load_settings", {}), - } - found_any = True - # Track the latest for auto-select - _last_obliterated_label = label - # Keep counter above any existing numbered dirs - if p.name.startswith("obliterated_"): - try: - idx = int(p.name.split("_", 1)[1]) - if idx >= _obliterate_counter: - _obliterate_counter = idx + 1 - except (ValueError, IndexError): - pass # If we recovered sessions but _state has no output_dir, set it to the # most recent checkpoint so chat_respond can reload from disk. if found_any and not _state.get("output_dir"): @@ -2236,11 +2244,52 @@ def obliterate(model_choice: str, method_choice: str, _state["model_name"] = model_choice _state["method"] = method + from obliteratus.run_archive import RunArchive + + run_config = { + key: value + for key, value in locals().copy().items() + if key.startswith("adv_") + or key in { + "model_choice", "model_id", "method_choice", "method", + "prompt_volume_choice", "prompt_volume", "dataset_source_choice", + "quantization_choice", "dtype_choice", + } + } + run_config["load_settings"] = load_settings.metadata() + run_config["custom_prompt_counts"] = { + "harmful_chars": len(custom_harmful or ""), + "harmless_chars": len(custom_harmless or ""), + } + run_config["custom_prompt_sha256"] = { + "harmful": hashlib.sha256((custom_harmful or "").encode()).hexdigest(), + "harmless": hashlib.sha256((custom_harmless or "").encode()).hexdigest(), + } + try: + run_archive = RunArchive() + run_id = run_archive.begin( + [model_id], + notes="Launched from the OBLITERATUS operator UI.", + metadata={"source": "operator_ui", "configuration": run_config}, + ) + save_dir = str(run_archive._run_dir(run_id) / "checkpoint") + except Exception as exc: + with _lock: + _state["status"] = "idle" + detail = f"{type(exc).__name__}: {exc}" + yield ( + f"**Run manifest failed:** {detail}", + f"RUN MANIFEST FAILED before GPU allocation: {detail}", + get_chat_header(), gr.update(), gr.update(), gr.update(), + ) + return + # Admission must be granted before cleanup can touch the CUDA runtime or # the worker can enter any model-loading path. try: _gpu_lifecycle.loading(model_id) except AdmissionError as exc: + run_archive.fail(run_id, exc, phase="admission") with _lock: _state["status"] = "idle" _state["model"] = None @@ -2250,7 +2299,9 @@ def obliterate(model_choice: str, method_choice: str, f"Diagnostic: {exc.diagnostic_message()}", ] yield ( - f"**GPU admission failed:** {exc.user_message()}", + f"**GPU admission failed ({type(exc).__name__}, phase admission):** " + f"{exc.user_message()} (run `{run_id}`; " + f"log `{run_archive._run_dir(run_id) / 'run.log'}`)", "\n".join(_state["log"]), get_chat_header(), gr.update(), @@ -2260,11 +2311,6 @@ def obliterate(model_choice: str, method_choice: str, return _clear_gpu(release_lifecycle=False) - with _lock: - global _obliterate_counter - _obliterate_counter += 1 - save_dir = f"/tmp/obliterated_{_obliterate_counter}" - log_lines = [] last_yielded = [0] pipeline_ref = [None] @@ -2278,6 +2324,7 @@ def obliterate(model_choice: str, method_choice: str, def on_log(msg): log_lines.append(msg) + run_archive.append_log(run_id, msg) def on_stage(result): stage_key = result.stage @@ -2285,6 +2332,12 @@ def obliterate(model_choice: str, method_choice: str, "excise": "\u2702\ufe0f", "verify": "\u2705", "rebirth": "\u2b50"}.get(stage_key, "\u25b6") if result.status == "running": log_lines.append(f"\n{icon} {stage_key.upper()} \u2014 {result.message}") + run_archive.append_log(run_id, f"{stage_key.upper()} — {result.message}") + run_archive.mark_running(run_id, phase=stage_key) + elif stage_key == "summon" and result.status == "done": + memory = measure_torch_memory(torch) + _gpu_lifecycle.resize(memory) + _gpu_lifecycle.ready(memory) stage_order = {"summon": 0, "probe": 1, "distill": 2, "excise": 3, "verify": 4, "rebirth": 5} idx = stage_order.get(stage_key, 0) @@ -2311,6 +2364,12 @@ def obliterate(model_choice: str, method_choice: str, n = min(prompt_volume, len(harmful_all), len(harmless_all)) else: n = min(len(harmful_all), len(harmless_all)) + run_archive.record_dataset( + run_id, + identifier="custom" if use_custom else dataset_key, + harmful=harmful_all[:n], + harmless=harmless_all[:n], + ) if method == "informed": # Use the analysis-guided InformedAbliterationPipeline @@ -2395,6 +2454,8 @@ def obliterate(model_choice: str, method_choice: str, source_info = DATASET_SOURCES.get(dataset_key) source_label = source_info.label if source_info else dataset_key log_lines.append(f"Target: {model_id}") + log_lines.append(f"Run ID: {run_id}") + log_lines.append(f"Durable log: {run_archive._run_dir(run_id) / 'run.log'}") log_lines.append(f"Method: {method}") if _adaptive_info: log_lines.append(_adaptive_info) @@ -2403,6 +2464,8 @@ def obliterate(model_choice: str, method_choice: str, log_lines.append(f"Prompt volume: {vol_label} pairs") log_lines.append(f"Resolved load mode: {load_settings.summary}") log_lines.append("") + for header_line in log_lines: + run_archive.append_log(run_id, header_line) worker = threading.Thread(target=run_pipeline, daemon=True) worker.start() @@ -2447,9 +2510,15 @@ def obliterate(model_choice: str, method_choice: str, _state["model"] = None _state["tokenizer"] = None log_lines.append(detail) + run_archive.fail( + run_id, + TimeoutError("pipeline exceeded the 45-minute limit"), + phase="timeout", + ) _state["log"] = log_lines yield ( - "**Timed out:** Pipeline cancelled; no successful artifact was registered.", + f"**Timed out (TimeoutError, phase timeout):** Pipeline cancelled (run `{run_id}`; " + f"log `{run_archive._run_dir(run_id) / 'run.log'}`).", "\n".join(log_lines), get_chat_header(), gr.update(), @@ -2467,9 +2536,15 @@ def obliterate(model_choice: str, method_choice: str, with _lock: _state["status"] = "idle" err_msg = str(error_ref[0]) or repr(error_ref[0]) + run_archive.fail(run_id, error_ref[0], phase="pipeline") log_lines.append(f"\nERROR: {err_msg}") _state["log"] = log_lines - yield f"**Error:** {err_msg}", "\n".join(log_lines), get_chat_header(), gr.update(), gr.update(), gr.update() + yield ( + f"**Error ({type(error_ref[0]).__name__}, phase pipeline):** {err_msg} " + f"(run `{run_id}`; " + f"log `{run_archive._run_dir(run_id) / 'run.log'}`)", + "\n".join(log_lines), get_chat_header(), gr.update(), gr.update(), gr.update(), + ) return # Success — keep model in memory for chat. @@ -2480,6 +2555,18 @@ def obliterate(model_choice: str, method_choice: str, _gpu_lifecycle.resize(memory) _gpu_lifecycle.ready(memory) save_dir = str(_resolve_local_checkpoint(save_dir)) + model_config = getattr(pipeline.handle.model, "config", None) + tokenizer_kwargs = getattr(pipeline.handle.tokenizer, "init_kwargs", {}) or {} + run_archive.record_revisions( + run_id, + model_revision=getattr(model_config, "_commit_hash", None), + tokenizer_revision=tokenizer_kwargs.get("_commit_hash"), + ) + run_archive.complete( + run_id, + checkpoint=save_dir, + metrics=dict(pipeline._quality_metrics), + ) can_generate = pipeline._quality_metrics.get("coherence") is not None # ── Telemetry: log single obliteration to community leaderboard ── @@ -2724,13 +2811,22 @@ def obliterate(model_choice: str, method_choice: str, except Exception as e: # Ensure status never gets stuck on "obliterating" _gpu_lifecycle.release(reason="post_load_failed") - _cleanup_failed_pipeline(pipeline_ref[0], save_dir) + completed = run_archive.status(run_id).get("status") == "succeeded" + if not completed: + _cleanup_failed_pipeline(pipeline_ref[0], save_dir) + run_archive.fail(run_id, e, phase="post_pipeline") with _lock: _state["status"] = "idle" err_msg = str(e) or repr(e) log_lines.append(f"\nERROR (post-pipeline): {err_msg}") + run_archive.append_log(run_id, f"ERROR (post-pipeline): {err_msg}") _state["log"] = log_lines - yield f"**Error:** {err_msg}", "\n".join(log_lines), get_chat_header(), gr.update(), gr.update(), gr.update() + yield ( + f"**Error ({type(e).__name__}, phase post-pipeline):** {err_msg} " + f"(run `{run_id}`; " + f"log `{run_archive._run_dir(run_id) / 'run.log'}`)", + "\n".join(log_lines), get_chat_header(), gr.update(), gr.update(), gr.update(), + ) # --------------------------------------------------------------------------- diff --git a/ci/test-risk-map.json b/ci/test-risk-map.json index 6f13763..f433720 100644 --- a/ci/test-risk-map.json +++ b/ci/test-risk-map.json @@ -184,6 +184,7 @@ "obliteratus/interactive.py", "obliteratus/local_ui.py", "obliteratus/gpu_lifecycle.py", + "obliteratus/run_archive.py", "obliteratus/models_client.py", "obliteratus/remote.py", "obliteratus/remote_contracts.py", @@ -197,6 +198,7 @@ "tests/test_interactive_contracts.py", "tests/test_local_ui_contracts.py", "tests/test_local_ui_portability.py", + "tests/test_run_archive.py", "tests/test_bestiary_contracts.py", "tests/test_models_client_contracts.py", "tests/test_remote_boundaries.py", diff --git a/docs/deployment/shared-gpu-host.md b/docs/deployment/shared-gpu-host.md index 0202db2..39df30a 100644 --- a/docs/deployment/shared-gpu-host.md +++ b/docs/deployment/shared-gpu-host.md @@ -112,9 +112,34 @@ RuntimeDirectoryMode=0770 Environment=OBLITERATUS_GPU_LIFECYCLE_DIR=/run/obliteratus-gpu-lifecycle Environment=OBLITERATUS_GPU_HEARTBEAT_SECONDS=15 Environment=OBLITERATUS_GPU_ADMISSION_TIMEOUT_SECONDS=120 +Environment=OBLITERATUS_RUN_ARCHIVE=/srv/obliteratus/run-archive ExecStart=/srv/obliteratus/current/.venv/bin/obliteratus ui --host 127.0.0.1 ``` +Create the archive before starting the service and make it writable only by the +service account (for example, owner `obliteratus`, mode `0750`). Every UI and +headless run receives a stable `run-*` ID and writes `manifest.json`, +`events.jsonl`, `run.log`, `notes.md`, checkpoint hashes, and an atomic +`COMPLETE` marker below this root. The manifest is committed before GPU +admission, so a disconnect or process restart remains diagnosable. + +The same contract is available without the UI: + +```bash +obliteratus runs --archive-root /srv/obliteratus/run-archive launch -- \ + Qwen/Qwen3.8-27B --method advanced --dtype bfloat16 +obliteratus runs --archive-root /srv/obliteratus/run-archive status RUN_ID +obliteratus runs --archive-root /srv/obliteratus/run-archive cancel RUN_ID +obliteratus runs --archive-root /srv/obliteratus/run-archive result RUN_ID +``` + +Checkpoint retention is deliberately separate from evidence retention. Under +verified disk pressure, an operator may run `runs prune-checkpoint RUN_ID +--reason TEXT`; this hashes and inventories the checkpoint first, then removes +only the model payload. Notes, configuration, logs, metrics, failure details, +events, and hashes remain in the archive. Do not automate this command without +a policy that ranks and selects dominated runs. + The runtime directory must be writable by the application group and readable by the root supervisor; `ack.json` should be root-owned and group-readable. The supervisor watches `current.json`, completes acquire/prepare, atomically writes diff --git a/obliteratus/__init__.py b/obliteratus/__init__.py index 7707d30..6248a92 100644 --- a/obliteratus/__init__.py +++ b/obliteratus/__init__.py @@ -22,6 +22,7 @@ __all__ = [ "Watchtower", "get_watchtower", "AutoObliterator", + "RunArchive", ] @@ -80,4 +81,7 @@ def __getattr__(name): if name == "AutoObliterator": from obliteratus.auto_obliterate import AutoObliterator return AutoObliterator + if name == "RunArchive": + from obliteratus.run_archive import RunArchive + return RunArchive raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/obliteratus/cli.py b/obliteratus/cli.py index 2905b66..9c896e0 100644 --- a/obliteratus/cli.py +++ b/obliteratus/cli.py @@ -560,6 +560,31 @@ def main(argv: list[str] | None = None): help="VRAM per GPU in GB (default: 80 for A100-80GB)", ) + # --- durable headless runs --- + runs_parser = subparsers.add_parser( + "runs", help="Launch and manage reconnect-safe durable obliteration runs", + ) + runs_parser.add_argument( + "--archive-root", default=None, + help="Durable run root (default: OBLITERATUS_RUN_ARCHIVE or the user state directory)", + ) + runs_subparsers = runs_parser.add_subparsers(dest="runs_command", required=True) + runs_launch = runs_subparsers.add_parser("launch", help="Launch a detached obliteration run") + runs_launch.add_argument("--notes", default="", help="Operator notes retained with the run") + runs_launch.add_argument( + "arguments", nargs=argparse.REMAINDER, + help="Arguments after --, beginning with the model ID", + ) + for command_name in ("status", "cancel", "result"): + command_parser = runs_subparsers.add_parser(command_name) + command_parser.add_argument("run_id") + runs_subparsers.add_parser("list") + runs_prune = runs_subparsers.add_parser( + "prune-checkpoint", help="Delete model payload while retaining evidence", + ) + runs_prune.add_argument("run_id") + runs_prune.add_argument("--reason", required=True) + args = parser.parse_args(argv) if getattr(args, "prompt_pairs_file", None): @@ -582,6 +607,8 @@ def main(argv: list[str] | None = None): if args.command == "gpu-calc": _cmd_gpu_calc(args) return + elif args.command == "runs": + _cmd_runs(args) elif args.command == "run": if getattr(args, "remote", None): _cmd_remote_run(args) @@ -660,6 +687,33 @@ def main(argv: list[str] | None = None): _cmd_abliterate(args) +def _cmd_runs(args): + """Dispatch the stable durable-run control surface.""" + + import json + + from obliteratus.run_archive import RunArchive + + archive = RunArchive(args.archive_root) + if args.runs_command == "launch": + arguments = list(args.arguments) + if arguments[:1] == ["--"]: + arguments = arguments[1:] + run_id = archive.launch(arguments, notes=args.notes) + payload = archive.status(run_id) + elif args.runs_command == "status": + payload = archive.status(args.run_id) + elif args.runs_command == "cancel": + payload = archive.cancel(args.run_id) + elif args.runs_command == "result": + payload = archive.result(args.run_id) + elif args.runs_command == "prune-checkpoint": + payload = archive.prune_checkpoint(args.run_id, reason=args.reason) + else: + payload = archive.list() + print(json.dumps(payload, indent=2, sort_keys=True)) + + def _cmd_self_improve(args): """Run one recursive hard-negative OBLITERATUS iteration.""" diff --git a/obliteratus/gpu_lifecycle.py b/obliteratus/gpu_lifecycle.py index 6b0ee91..43fd96b 100644 --- a/obliteratus/gpu_lifecycle.py +++ b/obliteratus/gpu_lifecycle.py @@ -128,7 +128,9 @@ class GpuLifecyclePublisher: self._phase = "admission_granted" self._publish("admission_granted") self._phase = "allocation_started" - return self._publish("allocation_started") + event = self._publish("allocation_started") + self._start_heartbeat() + return event def resize(self, memory: MemoryUsage) -> dict | None: with self._lock: @@ -150,7 +152,7 @@ class GpuLifecyclePublisher: def heartbeat(self) -> dict | None: with self._lock: - if self._phase != "ready": + if self._phase in {"released", "intent_published"}: return None return self._publish("heartbeat") @@ -336,6 +338,7 @@ def from_environment() -> GpuLifecyclePublisher: os.environ.get("OBLITERATUS_GPU_LIFECYCLE_DIR"), heartbeat_seconds=heartbeat_seconds, admission_timeout_seconds=admission_timeout_seconds, + run_id=os.environ.get("OBLITERATUS_RUN_ID"), ) diff --git a/obliteratus/run_archive.py b/obliteratus/run_archive.py new file mode 100644 index 0000000..5ee0a45 --- /dev/null +++ b/obliteratus/run_archive.py @@ -0,0 +1,609 @@ +"""Durable, reconnect-safe lifecycle for headless obliteration runs.""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib.metadata +import json +import os +import platform +import re +import shutil +import signal +import subprocess +import sys +import uuid +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Sequence + +from obliteratus.gpu_lifecycle import AdmissionError, MemoryUsage, from_environment + + +SCHEMA_VERSION = 1 +TERMINAL_STATES = frozenset({"succeeded", "failed", "cancelled"}) +RUN_ID_PATTERN = re.compile(r"^run-[0-9a-f]{32}$") +_SECRET_OPTION = re.compile(r"(?:token|secret|password|credential|api[-_]?key)", re.I) +_SECRET_VALUE = re.compile(r"(?i)(?:hf_[a-z0-9]{12,}|bearer\s+[a-z0-9._~+/-]{12,})") + + +def _now() -> str: + return datetime.now(timezone.utc).isoformat() + + +def default_archive_root() -> Path: + """Return a durable per-user root; service deployments should override it.""" + + configured = os.environ.get("OBLITERATUS_RUN_ARCHIVE") + if configured: + return Path(configured).expanduser() + state_home = os.environ.get("XDG_STATE_HOME") + base = Path(state_home).expanduser() if state_home else Path.home() / ".local" / "state" + return base / "obliteratus" / "runs" + + +def _atomic_json(path: Path, value: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp") + payload = json.dumps(value, indent=2, sort_keys=True) + "\n" + with temporary.open("x", encoding="utf-8") as stream: + stream.write(payload) + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary, path) + try: + directory_fd = os.open(path.parent, os.O_RDONLY) + try: + os.fsync(directory_fd) + finally: + os.close(directory_fd) + except OSError: + # Directory fsync is unavailable on some supported platforms. + pass + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _redact_arguments(arguments: Sequence[str]) -> list[str]: + result: list[str] = [] + redact_next = False + for argument in arguments: + if redact_next: + result.append("[REDACTED]") + redact_next = False + continue + if argument.startswith("--") and "=" in argument: + option, value = argument.split("=", 1) + result.append(f"{option}=[REDACTED]" if _SECRET_OPTION.search(option) else argument) + continue + result.append(argument) + if argument.startswith("--") and _SECRET_OPTION.search(argument): + redact_next = True + return result + + +def _sanitize_message(message: str) -> str: + return _SECRET_VALUE.sub("[REDACTED]", message) + + +def _process_start_ticks(pid: int) -> int | None: + try: + fields = Path(f"/proc/{pid}/stat").read_text(encoding="utf-8").split() + return int(fields[21]) + except (OSError, ValueError, IndexError): + return None + + +def _software_inventory() -> dict[str, Any]: + packages: dict[str, str] = {} + for name in ("obliteratus", "torch", "transformers", "accelerate", "bitsandbytes"): + try: + packages[name] = importlib.metadata.version(name) + except importlib.metadata.PackageNotFoundError: + continue + inventory: dict[str, Any] = { + "python": platform.python_version(), + "platform": platform.platform(), + "packages": packages, + } + try: + import torch + + inventory["cuda_runtime"] = torch.version.cuda + inventory["cuda_available"] = torch.cuda.is_available() + inventory["cuda_devices"] = ( + [torch.cuda.get_device_name(index) for index in range(torch.cuda.device_count())] + if torch.cuda.is_available() + else [] + ) + except (ImportError, RuntimeError): + inventory["cuda_available"] = False + inventory["cuda_devices"] = [] + return inventory + + +class RunArchive: + """Persistent launch/status/cancel/result API for obliteration experiments.""" + + def __init__(self, root: str | os.PathLike[str] | None = None): + configured = Path(root) if root is not None else default_archive_root() + configured = configured.expanduser() + if configured.is_symlink(): + raise ValueError("run archive root must not be a symlink") + self.root = configured.resolve() + self.root.mkdir(parents=True, exist_ok=True) + + def _run_dir(self, run_id: str) -> Path: + if not RUN_ID_PATTERN.fullmatch(run_id): + raise ValueError("invalid run ID") + path = self.root / run_id + if path.is_symlink(): + raise ValueError("run directory must not be a symlink") + return path + + def _manifest_path(self, run_id: str) -> Path: + return self._run_dir(run_id) / "manifest.json" + + def _load(self, run_id: str) -> dict[str, Any]: + path = self._manifest_path(run_id) + try: + value = json.loads(path.read_text(encoding="utf-8")) + except FileNotFoundError as exc: + raise KeyError(f"unknown run ID: {run_id}") from exc + if not isinstance(value, dict) or value.get("run_id") != run_id: + raise ValueError(f"invalid manifest for {run_id}") + return value + + def _save(self, manifest: dict[str, Any]) -> None: + manifest["updated_at"] = _now() + _atomic_json(self._manifest_path(str(manifest["run_id"])), manifest) + + def _event(self, run_id: str, event: str, **fields: Any) -> None: + record = {"schema_version": SCHEMA_VERSION, "at": _now(), "event": event, **fields} + path = self._run_dir(run_id) / "events.jsonl" + with path.open("a", encoding="utf-8") as stream: + stream.write(json.dumps(record, sort_keys=True) + "\n") + stream.flush() + os.fsync(stream.fileno()) + + def launch( + self, + arguments: Sequence[str], + *, + notes: str = "", + metadata: dict[str, Any] | None = None, + popen: Any = subprocess.Popen, + ) -> str: + """Launch ``obliteratus obliterate`` in a detached, durable worker.""" + + args = [str(value) for value in arguments] + run_id = self.begin(args, notes=notes, metadata=metadata) + run_dir = self._run_dir(run_id) + command = [ + sys.executable, + "-m", + "obliteratus.run_archive", + "worker", + "--archive-root", + str(self.root), + "--run-id", + run_id, + "--", + *args, + ] + process = popen( + command, + cwd=str(run_dir), + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + start_new_session=True, + ) + manifest = self._load(run_id) + manifest["worker"] = { + "pid": int(process.pid), + "uid": os.getuid() if hasattr(os, "getuid") else None, + "start_ticks": _process_start_ticks(int(process.pid)), + } + if manifest.get("status") not in TERMINAL_STATES: + manifest["phase"] = "worker_started" + self._save(manifest) + self._event(run_id, "worker_started", pid=int(process.pid)) + return run_id + + def begin( + self, + arguments: Sequence[str], + *, + notes: str = "", + metadata: dict[str, Any] | None = None, + ) -> str: + """Persist a new immutable run identity before any resource allocation.""" + + args = [str(value) for value in arguments] + if not args or args[0].startswith("-"): + raise ValueError("obliteration arguments must start with a model ID or local path") + if any(arg == "--output-dir" or arg.startswith("--output-dir=") for arg in args): + raise ValueError("output directory is managed by the run archive") + run_id = f"run-{uuid.uuid4().hex}" + run_dir = self._run_dir(run_id) + run_dir.mkdir(mode=0o750) + (run_dir / "notes.md").write_text(notes.rstrip() + ("\n" if notes else ""), encoding="utf-8") + created = _now() + manifest: dict[str, Any] = { + "schema_version": SCHEMA_VERSION, + "run_id": run_id, + "status": "queued", + "phase": "manifested", + "created_at": created, + "updated_at": created, + "arguments": _redact_arguments(args), + "model": args[0], + "model_revision": _option_value(args, "--revision"), + "tokenizer_revision": _option_value(args, "--tokenizer-revision"), + "seed": _option_value(args, "--seed"), + "dataset_inputs": _dataset_inputs(args), + "metadata": metadata or {}, + "software": _software_inventory(), + "paths": { + "run_dir": str(run_dir), + "checkpoint": str(run_dir / "checkpoint"), + "log": str(run_dir / "run.log"), + "notes": str(run_dir / "notes.md"), + }, + "failure": None, + "result": None, + "worker": { + "pid": os.getpid(), + "uid": os.getuid() if hasattr(os, "getuid") else None, + "start_ticks": _process_start_ticks(os.getpid()), + "mode": "manifest_owner", + }, + } + self._save(manifest) + self._event(run_id, "manifested") + return run_id + + def mark_running(self, run_id: str, *, phase: str) -> dict[str, Any]: + manifest = self._load(run_id) + manifest["status"] = "running" + manifest["phase"] = phase + self._save(manifest) + self._event(run_id, phase) + return manifest + + def append_log(self, run_id: str, message: str) -> None: + path = self._run_dir(run_id) / "run.log" + with path.open("a", encoding="utf-8") as stream: + stream.write(str(message).rstrip("\n") + "\n") + stream.flush() + + def record_dataset( + self, + run_id: str, + *, + identifier: str, + harmful: Sequence[str], + harmless: Sequence[str], + ) -> dict[str, Any]: + """Record prompt provenance by count and content hash, never raw prompts.""" + + encoded = json.dumps( + {"harmful": list(harmful), "harmless": list(harmless)}, + ensure_ascii=False, + separators=(",", ":"), + ).encode("utf-8") + manifest = self._load(run_id) + manifest["dataset_inputs"] = [ + { + "identifier": identifier, + "harmful_count": len(harmful), + "harmless_count": len(harmless), + "sha256": hashlib.sha256(encoded).hexdigest(), + } + ] + self._save(manifest) + self._event(run_id, "dataset_resolved", identifier=identifier) + return manifest + + def record_revisions( + self, + run_id: str, + *, + model_revision: str | None, + tokenizer_revision: str | None, + ) -> dict[str, Any]: + manifest = self._load(run_id) + manifest["model_revision"] = model_revision + manifest["tokenizer_revision"] = tokenizer_revision + self._save(manifest) + self._event(run_id, "revisions_resolved") + return manifest + + def fail(self, run_id: str, error: BaseException, *, phase: str) -> dict[str, Any]: + inventory_path = self._run_dir(run_id) / "artifact-inventory.json" + _atomic_json(inventory_path, {"artifacts": self._inventory(run_id)}) + manifest = self._load(run_id) + manifest["status"] = "failed" + manifest["phase"] = phase + manifest["failure"] = { + "type": type(error).__name__, + "message": _sanitize_message(str(error) or repr(error)), + "phase": phase, + } + manifest["failure_inventory"] = str(inventory_path) + self._save(manifest) + self._event(run_id, "failed", failure=manifest["failure"]) + return manifest + + def complete( + self, + run_id: str, + *, + checkpoint: str | os.PathLike[str], + metrics: dict[str, Any] | None = None, + ) -> dict[str, Any]: + checkpoint_path = Path(checkpoint).resolve() + expected = (self._run_dir(run_id) / "checkpoint").resolve() + if checkpoint_path != expected or not checkpoint_path.is_dir(): + raise ValueError("completed checkpoint must be the managed run checkpoint") + inventory_path = self._run_dir(run_id) / "artifact-inventory.json" + _atomic_json(inventory_path, {"artifacts": self._inventory(run_id)}) + manifest = self._load(run_id) + manifest["status"] = "succeeded" + manifest["phase"] = "complete" + manifest["result"] = { + "checkpoint": str(checkpoint_path), + "inventory": str(inventory_path), + "metrics": metrics or _checkpoint_metrics(checkpoint_path), + } + self._save(manifest) + (self._run_dir(run_id) / "COMPLETE").write_text(_now() + "\n", encoding="utf-8") + self._event(run_id, "complete") + return manifest + + def status(self, run_id: str) -> dict[str, Any]: + """Return the latest durable status and recover a vanished worker.""" + + manifest = self._load(run_id) + if manifest.get("status") not in TERMINAL_STATES and manifest.get("worker"): + if not self._worker_matches(manifest): + manifest["status"] = "failed" + manifest["phase"] = "worker_lost" + manifest["failure"] = { + "type": "WorkerLost", + "message": "worker process vanished before a terminal result was committed", + "phase": "worker_lost", + } + self._save(manifest) + self._event(run_id, "failed", reason="worker_lost") + return manifest + + def result(self, run_id: str) -> dict[str, Any]: + manifest = self.status(run_id) + if manifest.get("status") not in TERMINAL_STATES: + raise RuntimeError(f"run {run_id} is not complete") + return manifest + + def cancel(self, run_id: str) -> dict[str, Any]: + manifest = self.status(run_id) + if manifest.get("status") in TERMINAL_STATES: + return manifest + if not self._worker_matches(manifest): + return self.status(run_id) + pid = int(manifest["worker"]["pid"]) + try: + os.killpg(pid, signal.SIGTERM) + except (ProcessLookupError, PermissionError) as exc: + raise RuntimeError(f"cannot cancel worker for {run_id}: {exc}") from exc + manifest["status"] = "cancelling" + manifest["phase"] = "cancellation_requested" + self._save(manifest) + self._event(run_id, "cancellation_requested") + return manifest + + def list(self) -> list[dict[str, Any]]: + runs = [] + for path in sorted(self.root.glob("run-*/manifest.json"), reverse=True): + try: + runs.append(self.status(path.parent.name)) + except (KeyError, ValueError): + continue + return runs + + def prune_checkpoint(self, run_id: str, *, reason: str) -> dict[str, Any]: + """Remove only model payload after durable evidence and hashes exist.""" + + if not reason.strip(): + raise ValueError("checkpoint pruning requires an operator reason") + manifest = self.result(run_id) + checkpoint = self._run_dir(run_id) / "checkpoint" + inventory_path = self._run_dir(run_id) / "artifact-inventory.json" + if not inventory_path.is_file(): + inventory = self._inventory(run_id) + _atomic_json(inventory_path, {"artifacts": inventory}) + if checkpoint.exists(): + if checkpoint.is_symlink() or not checkpoint.is_dir(): + raise ValueError("checkpoint target is not a safe directory") + shutil.rmtree(checkpoint) + manifest["checkpoint_pruned"] = {"at": _now(), "reason": reason.strip()} + self._save(manifest) + self._event(run_id, "checkpoint_pruned", reason=reason.strip()) + return manifest + + def _worker_matches(self, manifest: dict[str, Any]) -> bool: + worker = manifest.get("worker") or {} + pid = worker.get("pid") + if not isinstance(pid, int) or pid <= 1: + return False + expected_uid = worker.get("uid") + if expected_uid is not None and hasattr(os, "getuid") and expected_uid != os.getuid(): + return False + start_ticks = _process_start_ticks(pid) + return start_ticks is not None and start_ticks == worker.get("start_ticks") + + def _inventory(self, run_id: str) -> list[dict[str, Any]]: + run_dir = self._run_dir(run_id) + records = [] + candidates = [run_dir / "notes.md", *(run_dir / "checkpoint").rglob("*")] + for path in sorted(candidates): + if not path.is_file(): + continue + records.append( + { + "path": path.relative_to(run_dir).as_posix(), + "bytes": path.stat().st_size, + "sha256": _sha256(path), + } + ) + return records + + +def _option_value(arguments: Sequence[str], option: str) -> str | None: + for index, value in enumerate(arguments): + if value == option and index + 1 < len(arguments): + return arguments[index + 1] + if value.startswith(f"{option}="): + return value.split("=", 1)[1] + return None + + +def _dataset_inputs(arguments: Sequence[str]) -> list[dict[str, Any]]: + inputs: list[dict[str, Any]] = [] + dataset = _option_value(arguments, "--dataset") or "builtin" + inputs.append({"option": "--dataset", "identifier": dataset, "sha256": None}) + for option in ("--prompt-pairs-file", "--residue-file"): + for index, value in enumerate(arguments): + candidate = None + if value == option and index + 1 < len(arguments): + candidate = arguments[index + 1] + elif value.startswith(f"{option}="): + candidate = value.split("=", 1)[1] + if candidate: + path = Path(candidate).expanduser() + inputs.append( + { + "option": option, + "path": str(path.resolve()), + "sha256": _sha256(path) if path.is_file() else None, + } + ) + return inputs + + +def _worker(archive: RunArchive, run_id: str, arguments: Sequence[str]) -> int: + manifest = archive._load(run_id) + log_path = archive._run_dir(run_id) / "run.log" + checkpoint = archive._run_dir(run_id) / "checkpoint" + os.environ["OBLITERATUS_RUN_ID"] = run_id + lifecycle = from_environment() + child: subprocess.Popen[Any] | None = None + previous_sigterm = signal.getsignal(signal.SIGTERM) + + def _cancel_worker(_signum: int, _frame: Any) -> None: + raise KeyboardInterrupt("cancellation requested") + + signal.signal(signal.SIGTERM, _cancel_worker) + try: + manifest["status"] = "running" + manifest["phase"] = "admission" + archive._save(manifest) + archive._event(run_id, "loading") + lifecycle.loading(str(arguments[0])) + manifest["phase"] = "pipeline" + archive._save(manifest) + archive._event(run_id, "admitted") + command = [ + sys.executable, + "-m", + "obliteratus", + "obliterate", + *arguments, + "--output-dir", + str(checkpoint), + ] + with log_path.open("a", encoding="utf-8") as log: + child = subprocess.Popen(command, stdin=subprocess.DEVNULL, stdout=log, stderr=subprocess.STDOUT) + return_code = child.wait() + if return_code != 0: + raise RuntimeError(f"obliteration process exited with status {return_code}") + archive.complete(run_id, checkpoint=checkpoint) + lifecycle.resize(MemoryUsage()) + lifecycle.ready() + lifecycle.release(reason="complete") + signal.signal(signal.SIGTERM, previous_sigterm) + return 0 + except AdmissionError as exc: + failure = { + "type": type(exc).__name__, + "message": _sanitize_message(exc.user_message()), + "phase": "admission", + } + except BaseException as exc: + phase = "cancelled" if isinstance(exc, KeyboardInterrupt) else "pipeline" + failure = { + "type": type(exc).__name__, + "message": _sanitize_message(str(exc) or repr(exc)), + "phase": phase, + } + if child is not None and child.poll() is None: + child.terminate() + try: + child.wait(timeout=30) + except subprocess.TimeoutExpired: + child.kill() + child.wait(timeout=10) + inventory_path = archive._run_dir(run_id) / "artifact-inventory.json" + _atomic_json(inventory_path, {"artifacts": archive._inventory(run_id)}) + manifest = archive._load(run_id) + cancelling = manifest.get("status") == "cancelling" + manifest["status"] = "cancelled" if cancelling else "failed" + manifest["phase"] = "cancelled" if cancelling else failure["phase"] + manifest["failure"] = failure + manifest["failure_inventory"] = str(inventory_path) + archive._save(manifest) + archive._event(run_id, manifest["status"], failure=failure) + lifecycle.release(reason=manifest["status"]) + signal.signal(signal.SIGTERM, previous_sigterm) + return 130 if cancelling else 1 + + +def _checkpoint_metrics(checkpoint: Path) -> dict[str, Any]: + metadata_path = checkpoint / "abliteration_metadata.json" + try: + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + return {} + metrics = metadata.get("quality_metrics", metadata.get("metrics", {})) + return metrics if isinstance(metrics, dict) else {} + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + subparsers = parser.add_subparsers(dest="command", required=True) + worker = subparsers.add_parser("worker") + worker.add_argument("--archive-root", required=True) + worker.add_argument("--run-id", required=True) + worker.add_argument("arguments", nargs=argparse.REMAINDER) + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + args = _parser().parse_args(argv) + arguments = list(args.arguments) + if arguments[:1] == ["--"]: + arguments = arguments[1:] + if not arguments: + raise SystemExit("worker requires obliteration arguments after --") + return _worker(RunArchive(args.archive_root), args.run_id, arguments) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_gpu_lifecycle.py b/tests/test_gpu_lifecycle.py index 2216012..db3f6f1 100644 --- a/tests/test_gpu_lifecycle.py +++ b/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") diff --git a/tests/test_run_archive.py b/tests/test_run_archive.py new file mode 100644 index 0000000..f737985 --- /dev/null +++ b/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"