"""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.checkpoint_provenance import verify_provenance_record 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 attach_checkpoint_provenance(self, run_id: str, provenance: Any) -> dict[str, Any]: """Attach one canonical artifact identity to the durable run archive.""" try: record = verify_provenance_record(json.loads(provenance.to_json())) except (AttributeError, TypeError, json.JSONDecodeError, ValueError) as error: raise ValueError("checkpoint provenance must be a canonical provenance record") from error artifact_id = record.get("artifact_id") if ( not isinstance(artifact_id, str) or getattr(provenance, "artifact_id", None) != artifact_id ): raise ValueError("checkpoint provenance contract is invalid") manifest = self._load(run_id) existing = manifest.get("artifact_id") if existing is not None and existing != artifact_id: raise ValueError("run archive already names a different artifact identity") relative_path = "checkpoint-provenance.json" path = self._run_dir(run_id) / relative_path _atomic_json(path, record) manifest["artifact_id"] = artifact_id manifest["checkpoint_provenance"] = { "artifact_id": artifact_id, "path": relative_path, "sha256": f"sha256:{_sha256(path)}", } self._save(manifest) self._event(run_id, "checkpoint_provenance_attached", artifact_id=artifact_id) return manifest 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 record_experiment_protocol( self, run_id: str, protocol: dict[str, Any], ) -> dict[str, Any]: """Persist a prompt-identity protocol without storing raw prompts.""" protocol_path = self._run_dir(run_id) / "experiment-protocol.json" _atomic_json(protocol_path, protocol) manifest = self._load(run_id) manifest["experiment_protocol"] = { "path": str(protocol_path), "protocol": protocol.get("protocol"), "manifest_sha256": protocol.get("manifest_sha256"), "counts": protocol.get("counts"), } self._save(manifest) self._event( run_id, "experiment_protocol_resolved", protocol=protocol.get("protocol"), ) 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 begin_evaluation( self, run_id: str, *, partition: str, evaluator: str, ) -> dict[str, Any]: """Reserve one durable saved-checkpoint evaluation attempt. Final-test reservation is intentionally fail-closed: even an interrupted attempt consumes the candidate's single authorized final evaluation. """ if partition not in {"optimizer_tune", "final_test"}: raise ValueError("unsupported evaluation partition") manifest = self.result(run_id) protocol = manifest.get("experiment_protocol") or {} if protocol.get("protocol") != "qwen38-v1": raise ValueError("saved-checkpoint evaluation requires qwen38-v1") evaluations = manifest.setdefault("evaluations", []) if partition == "final_test" and any( item.get("partition") == "final_test" for item in evaluations ): raise ValueError("final-test evaluation was already reserved for this candidate") evaluation_id = f"eval-{uuid.uuid4().hex}" relative_path = f"evaluations/{evaluation_id}.json" record = { "schema_version": SCHEMA_VERSION, "evaluation_id": evaluation_id, "run_id": run_id, "partition": partition, "evaluator": evaluator, "status": "running", "created_at": _now(), "updated_at": _now(), "metrics": None, "failure": None, "log": [], } record_path = self._run_dir(run_id) / relative_path record_path.parent.mkdir(mode=0o750, exist_ok=True) _atomic_json(record_path, record) evaluations.append( { "evaluation_id": evaluation_id, "partition": partition, "evaluator": evaluator, "status": "running", "path": relative_path, } ) self._save(manifest) self._event( run_id, "evaluation_reserved", evaluation_id=evaluation_id, partition=partition, ) return record def finish_evaluation( self, run_id: str, evaluation_id: str, *, metrics: dict[str, Any] | None = None, log: Sequence[str] = (), failure: BaseException | None = None, ) -> dict[str, Any]: """Commit terminal metrics or failure for a reserved evaluation.""" if not re.fullmatch(r"eval-[0-9a-f]{32}", evaluation_id): raise ValueError("invalid evaluation ID") manifest = self._load(run_id) entries = manifest.get("evaluations") or [] entry = next( (item for item in entries if item.get("evaluation_id") == evaluation_id), None, ) if entry is None: raise KeyError(f"unknown evaluation ID: {evaluation_id}") path = self._run_dir(run_id) / str(entry["path"]) record = json.loads(path.read_text(encoding="utf-8")) if record.get("status") != "running": raise ValueError("evaluation is already terminal") record["status"] = "failed" if failure is not None else "succeeded" record["updated_at"] = _now() record["metrics"] = metrics record["log"] = [str(line) for line in log] record["failure"] = ( { "type": type(failure).__name__, "message": _sanitize_message(str(failure) or repr(failure)), } if failure is not None else None ) _atomic_json(path, record) entry["status"] = record["status"] entry["sha256"] = hashlib.sha256(path.read_bytes()).hexdigest() self._save(manifest) self._event( run_id, "evaluation_finished", evaluation_id=evaluation_id, partition=record["partition"], status=record["status"], ) return record 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())