Files
OBLITERATUS/obliteratus/run_archive.py
T

610 lines
23 KiB
Python

"""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())