mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-30 14:40:38 +02:00
Merge pull request #188 from jmagly/feat/184-durable-runs
feat: add durable reconnect-safe experiment runs (#184)
This commit is contained in:
@@ -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(),
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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}")
|
||||
|
||||
@@ -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."""
|
||||
|
||||
|
||||
@@ -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"),
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -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())
|
||||
@@ -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")
|
||||
|
||||
@@ -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"
|
||||
Reference in New Issue
Block a user