CalvinBackup/OBLITERATUS
1
0
Fork 0
mirror of https://github.com/elder-plinius/OBLITERATUS.git synced 2026-08-30 06:30:37 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #161 from elder-plinius/fix/158-pipeline-fail-closed

Browse Source 
fix(pipeline): fail closed on invalid output
This commit is contained in:
Joseph Magly
2026-08-24 13:45:28 -04:00
committed by GitHub
parent 7162203da1 d8d0a76231
commit 56442df2d2
7 changed files with 515 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app.py
+178 -37
View File
@@ -71,6 +71,7 @@ from obliteratus.model_load_settings import (
ModelLoadSettings,
resolve_model_load_settings,
)
from obliteratus.persistence_contracts import validate_reloadable_checkpoint
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
# ── ZeroGPU support ─────────────────────────────────────────────────
@@ -121,6 +122,7 @@ def _load_model_to_device(
offload_folder: str | None = None,
low_cpu_mem_usage: bool = False,
token: str | None = None,
local_files_only: bool = False,
) -> AutoModelForCausalLM:
"""Load a causal LM onto the best available device, MPS-safe.
@@ -142,6 +144,8 @@ def _load_model_to_device(
kwargs["low_cpu_mem_usage"] = True
if token is not None:
kwargs["token"] = token
if local_files_only:
kwargs["local_files_only"] = True
if dev.supports_device_map_auto():
kwargs["device_map"] = "auto"
@@ -156,6 +160,87 @@ def _load_model_to_device(
return model
def _resolve_local_checkpoint(checkpoint: str | Path) -> Path:
"""Resolve and fully validate a local checkpoint before Transformers sees it."""
return validate_reloadable_checkpoint(checkpoint)
def _reload_local_checkpoint(
checkpoint: str | Path,
*,
trust_remote_code: bool,
model_kwargs: dict | None = None,
):
"""Load a validated local checkpoint without permitting Hub resolution."""
checkpoint_path = _resolve_local_checkpoint(checkpoint)
model = _load_model_to_device(
checkpoint_path,
trust_remote_code=trust_remote_code,
local_files_only=True,
**(model_kwargs or {}),
)
tokenizer = AutoTokenizer.from_pretrained(
checkpoint_path,
trust_remote_code=trust_remote_code,
local_files_only=True,
)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
return model, tokenizer
def _format_checkpoint_reload_error(error: BaseException) -> str:
"""Return an actionable reload category without misreporting GPU capacity."""
message = str(error) or repr(error)
lowered = message.lower()
if isinstance(error, FileNotFoundError):
return f"Saved checkpoint is missing: {message}"
if "corrupt" in lowered:
return f"Saved checkpoint is corrupt: {message}"
if isinstance(error, ValueError) and "checkpoint" in lowered:
return f"Saved checkpoint is incomplete: {message}"
if dev.is_oom_error(error):
return f"Insufficient GPU memory to reload the completed checkpoint: {message}"
if "repo id" in lowered or "repository" in lowered:
return f"Checkpoint Hub resolution failed: {message}"
return f"Completed checkpoint could not be reloaded: {message}"
def _cancel_pipeline_worker(
worker: threading.Thread,
cancellation: threading.Event,
*,
grace_seconds: float,
) -> bool:
"""Signal cooperative cancellation and report confirmed worker termination."""
cancellation.set()
worker.join(timeout=grace_seconds)
return not worker.is_alive()
def _cleanup_failed_pipeline(pipeline, output_dir: str) -> None:
"""Release failed model state and remove only this UI run's final artifact."""
if pipeline is not None:
try:
pipeline.cleanup_failed_run()
except Exception:
pass
path = Path(output_dir)
try:
resolved = path.resolve()
temp_root = Path("/tmp").resolve()
if (
resolved.parent == temp_root
and resolved.name.startswith("obliterated_")
and resolved.is_dir()
):
import shutil
shutil.rmtree(resolved, ignore_errors=True)
except OSError:
pass
def _checkpoint_load_kwargs(load_settings: ModelLoadSettings) -> dict:
"""Translate resolved settings for direct loading of a saved checkpoint."""
torch_dtype = {
@@ -2072,6 +2157,7 @@ def obliterate(model_choice: str, method_choice: str,
last_yielded = [0]
pipeline_ref = [None]
error_ref = [None]
cancellation = threading.Event()
t_start = time.time()
def _elapsed():
@@ -2128,6 +2214,7 @@ def obliterate(model_choice: str, method_choice: str,
harmless_prompts=harmless_all[:n],
on_stage=on_stage,
on_log=on_log,
cancellation_event=cancellation,
)
pipeline_ref[0] = pipeline
pipeline.run_informed()
@@ -2145,6 +2232,7 @@ def obliterate(model_choice: str, method_choice: str,
harmless_prompts=harmless_all[:n],
on_stage=on_stage,
on_log=on_log,
cancellation_event=cancellation,
# Advanced overrides from UI
n_directions=int(adv_n_directions),
direction_method=adv_direction_method,
@@ -2210,6 +2298,7 @@ def obliterate(model_choice: str, method_choice: str,
# Stream log updates while pipeline runs (max 45 minutes to prevent indefinite hang)
_max_pipeline_secs = 45 * 60
_pipeline_start = time.time()
timed_out = False
status_msg = "**Obliterating\u2026** (0s)"
while worker.is_alive():
status_msg = f"**Obliterating\u2026** ({_elapsed()})"
@@ -2220,13 +2309,47 @@ def obliterate(model_choice: str, method_choice: str,
yield status_msg, "\n".join(log_lines), gr.update(), gr.update(), gr.update(), gr.update()
if time.time() - _pipeline_start > _max_pipeline_secs:
log_lines.append("\nTIMEOUT: Pipeline exceeded 45-minute limit.")
cancellation.set()
timed_out = True
break
time.sleep(0.5)
worker.join(timeout=30)
if timed_out:
terminated = _cancel_pipeline_worker(
worker,
cancellation,
grace_seconds=30,
)
pipeline = pipeline_ref[0]
if terminated:
_cleanup_failed_pipeline(pipeline, save_dir)
detail = "Worker stopped; incomplete output was discarded."
else:
detail = (
"Worker is still stopping cooperatively; no output or session state "
"will be exposed."
)
with _lock:
_state["status"] = "idle"
_state["model"] = None
_state["tokenizer"] = None
log_lines.append(detail)
_state["log"] = log_lines
yield (
"**Timed out:** Pipeline cancelled; no successful artifact was registered.",
"\n".join(log_lines),
get_chat_header(),
gr.update(),
gr.update(),
gr.update(),
)
return
worker.join()
# Handle error
if error_ref[0] is not None:
_cleanup_failed_pipeline(pipeline_ref[0], save_dir)
with _lock:
_state["status"] = "idle"
err_msg = str(error_ref[0]) or repr(error_ref[0])
@@ -2239,6 +2362,7 @@ def obliterate(model_choice: str, method_choice: str,
# Wrapped in try/except to ensure status is never stuck on "obliterating".
try:
pipeline = pipeline_ref[0]
save_dir = str(_resolve_local_checkpoint(save_dir))
can_generate = pipeline._quality_metrics.get("coherence") is not None
# ── Telemetry: log single obliteration to community leaderboard ──
@@ -2345,8 +2469,17 @@ def obliterate(model_choice: str, method_choice: str,
except ImportError:
pass
dtype_label = {
"float16": "FP16",
"bfloat16": "BF16",
"float32": "FP32",
}[load_settings.dtype]
checkpoint_path = _resolve_local_checkpoint(save_dir)
if bnb_available:
log_lines.append("\nModel too large for chat at float16 — reloading in 4-bit...")
log_lines.append(
f"\nModel too large for chat at {dtype_label} — reloading in 4-bit..."
)
last_yielded[0] = len(log_lines)
yield status_msg, "\n".join(log_lines), gr.update(), gr.update(), gr.update(), gr.update()
try:
@@ -2357,17 +2490,11 @@ def obliterate(model_choice: str, method_choice: str,
bnb_4bit_quant_type="nf4",
llm_int8_enable_fp32_cpu_offload=True,
)
model_reloaded = _load_model_to_device(
save_dir,
quantization_config=bnb_cfg,
model_reloaded, tokenizer_reloaded = _reload_local_checkpoint(
checkpoint_path,
trust_remote_code=True,
model_kwargs={"quantization_config": bnb_cfg},
)
tokenizer_reloaded = AutoTokenizer.from_pretrained(
save_dir,
trust_remote_code=True,
)
if tokenizer_reloaded.pad_token is None:
tokenizer_reloaded.pad_token = tokenizer_reloaded.eos_token
# Re-install activation steering hooks on the reloaded model
if steering_meta:
@@ -2382,14 +2509,17 @@ def obliterate(model_choice: str, method_choice: str,
can_generate = True
log_lines.append("Reloaded in 4-bit — chat is ready!")
except Exception as e:
log_lines.append(f"4-bit reload failed: {e}")
log_lines.append(
f"4-bit reload failed: {_format_checkpoint_reload_error(e)}"
)
_clear_gpu()
# -- Attempt 2: CPU offloading (slower but no extra dependencies)
if not can_generate:
import tempfile
log_lines.append(
"\nModel too large for chat at float16 — reloading with CPU offload..."
f"\nModel too large for chat at {dtype_label} — "
"reloading with CPU offload..."
if not bnb_available
else "Falling back to CPU offload..."
)
@@ -2397,22 +2527,18 @@ def obliterate(model_choice: str, method_choice: str,
yield status_msg, "\n".join(log_lines), gr.update(), gr.update(), gr.update(), gr.update()
try:
offload_dir = tempfile.mkdtemp(prefix="obliteratus_offload_")
model_reloaded = _load_model_to_device(
save_dir,
offload_folder=offload_dir,
torch_dtype={
"float16": torch.float16,
"bfloat16": torch.bfloat16,
"float32": torch.float32,
}[load_settings.dtype],
model_reloaded, tokenizer_reloaded = _reload_local_checkpoint(
checkpoint_path,
trust_remote_code=True,
model_kwargs={
"offload_folder": offload_dir,
"torch_dtype": {
"float16": torch.float16,
"bfloat16": torch.bfloat16,
"float32": torch.float32,
}[load_settings.dtype],
},
)
tokenizer_reloaded = AutoTokenizer.from_pretrained(
save_dir,
trust_remote_code=True,
)
if tokenizer_reloaded.pad_token is None:
tokenizer_reloaded.pad_token = tokenizer_reloaded.eos_token
# Re-install activation steering hooks on the reloaded model
if steering_meta:
@@ -2427,8 +2553,13 @@ def obliterate(model_choice: str, method_choice: str,
can_generate = True
log_lines.append("Reloaded with CPU offload — chat is ready (may be slower).")
except Exception as e:
log_lines.append(f"CPU offload reload failed: {e}")
log_lines.append("Chat unavailable. Load the saved model on a larger instance.")
log_lines.append(
f"CPU offload reload failed: {_format_checkpoint_reload_error(e)}"
)
log_lines.append(
"Chat unavailable. The completed checkpoint remains saved; "
"review the reload error above."
)
with _lock:
_state["status"] = "idle"
@@ -2474,6 +2605,7 @@ def obliterate(model_choice: str, method_choice: str,
except Exception as e:
# Ensure status never gets stuck on "obliterating"
_cleanup_failed_pipeline(pipeline_ref[0], save_dir)
with _lock:
_state["status"] = "idle"
err_msg = str(e) or repr(e)
@@ -2564,15 +2696,17 @@ def chat_respond(message: str, history: list[dict], system_prompt: str,
checkpoint = _state.get("output_dir")
if _checkpoint_is_available(checkpoint):
try:
checkpoint = _resolve_local_checkpoint(checkpoint)
is_preset = (_state.get("model_name") or "") in MODELS
load_settings = _settings_from_metadata(_state.get("load_settings"))
model = _load_model_to_device(
checkpoint,
trust_remote_code=is_preset,
local_files_only=True,
**_checkpoint_load_kwargs(load_settings),
)
tokenizer = AutoTokenizer.from_pretrained(
checkpoint, trust_remote_code=is_preset,
checkpoint, trust_remote_code=is_preset, local_files_only=True,
)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
@@ -2584,8 +2718,8 @@ def chat_respond(message: str, history: list[dict], system_prompt: str,
_state["model"] = model
_state["tokenizer"] = tokenizer
_state["status"] = "ready"
except Exception:
yield "Model failed to reload from checkpoint. Try re-obliterating."
except Exception as error:
yield _format_checkpoint_reload_error(error)
return
else:
_clear_stale_model_state()
@@ -2778,14 +2912,16 @@ def load_bench_into_chat(choice: str, progress=gr.Progress()):
if checkpoint and Path(checkpoint).exists():
is_preset = (_state.get("model_name") or "") in MODELS
try:
checkpoint = _resolve_local_checkpoint(checkpoint)
load_settings = _settings_from_metadata(_state.get("load_settings"))
model_loaded = _load_model_to_device(
checkpoint,
trust_remote_code=is_preset,
local_files_only=True,
**_checkpoint_load_kwargs(load_settings),
)
tokenizer_loaded = AutoTokenizer.from_pretrained(
checkpoint, trust_remote_code=is_preset,
checkpoint, trust_remote_code=is_preset, local_files_only=True,
)
if tokenizer_loaded.pad_token is None:
tokenizer_loaded.pad_token = tokenizer_loaded.eos_token
@@ -2799,7 +2935,7 @@ def load_bench_into_chat(choice: str, progress=gr.Progress()):
)
return
except Exception as e:
yield f"**Error:** Could not reload model: {e}", get_chat_header()
yield f"**Error:** {_format_checkpoint_reload_error(e)}", get_chat_header()
return
yield (
"**Error:** Model checkpoint not found. The Space may have restarted — "
@@ -2847,13 +2983,15 @@ def load_bench_into_chat(choice: str, progress=gr.Progress()):
is_preset = cfg["model_choice"] in MODELS
try:
checkpoint_dir = _resolve_local_checkpoint(checkpoint_dir)
model_loaded = _load_model_to_device(
checkpoint_dir,
trust_remote_code=is_preset,
local_files_only=True,
**_checkpoint_load_kwargs(load_settings),
)
tokenizer_loaded = AutoTokenizer.from_pretrained(
checkpoint_dir, trust_remote_code=is_preset,
checkpoint_dir, trust_remote_code=is_preset, local_files_only=True,
)
if tokenizer_loaded.pad_token is None:
tokenizer_loaded.pad_token = tokenizer_loaded.eos_token
@@ -2896,9 +3034,10 @@ def load_bench_into_chat(choice: str, progress=gr.Progress()):
checkpoint_dir,
quantization_config=bnb_cfg,
trust_remote_code=is_preset,
local_files_only=True,
)
tokenizer_loaded = AutoTokenizer.from_pretrained(
checkpoint_dir, trust_remote_code=is_preset,
checkpoint_dir, trust_remote_code=is_preset, local_files_only=True,
)
if tokenizer_loaded.pad_token is None:
tokenizer_loaded.pad_token = tokenizer_loaded.eos_token
@@ -3035,15 +3174,17 @@ def ab_chat_respond(message: str, history_left: list[dict], history_right: list[
model_name = _state.get("model_name") or model_name
if checkpoint and Path(checkpoint).exists():
try:
checkpoint = _resolve_local_checkpoint(checkpoint)
is_preset = (model_name or "") in MODELS
load_settings = _settings_from_metadata(_state.get("load_settings"))
abliterated_model = _load_model_to_device(
checkpoint,
trust_remote_code=is_preset,
local_files_only=True,
**_checkpoint_load_kwargs(load_settings),
)
tokenizer = AutoTokenizer.from_pretrained(
checkpoint, trust_remote_code=is_preset,
checkpoint, trust_remote_code=is_preset, local_files_only=True,
)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
obliteratus/abliterate.py
+79
View File
@@ -29,6 +29,7 @@ import time
import warnings
from dataclasses import dataclass, field
from pathlib import Path
from threading import Event
from typing import Any, Callable, Iterable
import torch
@@ -771,6 +772,23 @@ class StageResult:
details: dict[str, Any] = field(default_factory=dict)
class PipelineFailure(RuntimeError):
"""Terminal pipeline failure with structured stage and metric context."""
def __init__(self, message: str, *, stage: str, metric: str | None = None):
super().__init__(message)
self.stage = stage
self.metric = metric
class PipelineCancelledError(PipelineFailure):
"""Raised at a cooperative cancellation boundary."""
class PipelineValidationError(PipelineFailure):
"""Raised when model-quality validation proves an artifact unsafe."""
def auto_hub_repo_id(model_name: str, *, api=None, org: str | None = None) -> str:
"""Generate a Hub repo ID like ``{namespace}/{short_model}-OBLITERATED``.
@@ -893,6 +911,7 @@ class AbliterationPipeline:
refusal_max_tokens: int | None = None,
on_stage: Callable[[StageResult], None] | None = None,
on_log: Callable[[str], None] | None = None,
cancellation_event: Event | None = None,
):
self.model_name = model_name
self.output_dir = Path(output_dir)
@@ -925,6 +944,8 @@ class AbliterationPipeline:
self.jailbreak_prompts = jailbreak_prompts
self._on_stage = on_stage or (lambda r: None)
self._on_log = on_log or (lambda m: None)
self._cancellation_event = cancellation_event or Event()
self._active_stage = "summon"
self._stage_durations: dict[str, float] = {}
self._excise_modified_count: int | None = None
@@ -1100,9 +1121,12 @@ class AbliterationPipeline:
self._routing_attention_mask: torch.Tensor | None = None
def log(self, msg: str):
self._check_cancelled()
self._on_log(msg)
def _emit(self, key: str, status: str, message: str = "", **details) -> StageResult:
if status != "error":
self._check_cancelled(key)
result = StageResult(stage=key, status=status, message=message, details=details)
if status == "done":
duration = details.get("duration")
@@ -1114,6 +1138,41 @@ class AbliterationPipeline:
self._on_stage(result)
return result
def cancel(self) -> None:
"""Request cooperative cancellation at the next bounded checkpoint."""
self._cancellation_event.set()
def _check_cancelled(self, stage: str | None = None) -> None:
if self._cancellation_event.is_set():
active_stage = stage or self._active_stage
raise PipelineCancelledError(
f"Pipeline cancelled during {active_stage.upper()}.",
stage=active_stage,
metric="cancellation",
)
def _fail_validation(self, metric: str, value: float, reason: str) -> None:
self._quality_metrics[metric] = value
message = f"Validation failed: {reason} ({metric}={value!r})"
self._emit("verify", "error", message, metric=metric, value=value)
raise PipelineValidationError(message, stage="verify", metric=metric)
def cleanup_failed_run(self) -> None:
"""Unload an unsafe/partial in-memory result and release transient state."""
for hook in self._steering_hooks:
hook.remove()
self._steering_hooks.clear()
if self.handle is not None:
self.handle.model = None
self.handle.tokenizer = None
self._harmful_acts.clear()
self._harmless_acts.clear()
self._jailbreak_acts.clear()
self._routing_harmful.clear()
self._routing_harmless.clear()
self._cleanup_offload_dir()
self._free_gpu_memory()
@staticmethod
def _free_gpu_memory():
"""Release unused GPU/accelerator memory between pipeline stages."""
@@ -1329,10 +1388,13 @@ class AbliterationPipeline:
for h in self._steering_hooks:
h.remove()
self._steering_hooks.clear()
self._active_stage = "summon"
self._summon()
self._free_gpu_memory()
self._active_stage = "probe"
self._probe()
self._free_gpu_memory()
self._active_stage = "distill"
self._distill()
# Free raw per-prompt activations now that means/subspaces are extracted
self._harmful_acts.clear()
@@ -1347,10 +1409,13 @@ class AbliterationPipeline:
self._routing_harmless.clear()
self._free_gpu_memory()
self._capture_baseline_kl_logits()
self._active_stage = "excise"
self._excise()
self._free_gpu_memory()
self._active_stage = "verify"
self._verify()
self._free_gpu_memory()
self._active_stage = "rebirth"
return self._rebirth()
# ── Stage 1: SUMMON ─────────────────────────────────────────────────
@@ -6482,6 +6547,11 @@ class AbliterationPipeline:
# All reference texts produced NaN loss — model is completely broken
perplexity = float("inf")
self.log(" Perplexity: inf (model produces NaN outputs — weights may be destroyed)")
self._fail_validation(
"perplexity",
perplexity,
"all reference losses were non-finite; the modified model is unsafe",
)
elif has_nan_loss:
# Some texts produced NaN — compute from valid ones but warn
avg_loss = total_loss / n_tokens
@@ -6498,6 +6568,12 @@ class AbliterationPipeline:
perplexity = float("inf")
self.log(f" Perplexity: {perplexity:.2f}")
self._quality_metrics["perplexity"] = perplexity
if not math.isfinite(perplexity):
self._fail_validation(
"perplexity",
perplexity,
"perplexity is non-finite; generation and export are blocked",
)
# 2. Generation coherence test
test_prompts = [
@@ -7186,7 +7262,10 @@ class AbliterationPipeline:
msg += f" ({disk.free / 1e9:.1f} GB free on {self.output_dir})"
raise type(e)(msg) from e
self._check_cancelled("rebirth")
self.handle.tokenizer.save_pretrained(checkpoint_dir)
self._check_cancelled("rebirth")
(checkpoint_dir / "abliteration_metadata.json").write_text(
metadata_json,
encoding="utf-8",
obliteratus/informed_pipeline.py
+10
View File
@@ -54,6 +54,7 @@ import logging
import time
from dataclasses import dataclass, field
from pathlib import Path
from threading import Event
from typing import Callable
import torch
@@ -182,6 +183,7 @@ class InformedAbliterationPipeline(AbliterationPipeline):
harmless_prompts: list[str] | None = None,
on_stage: Callable[[StageResult], None] | None = None,
on_log: Callable[[str], None] | None = None,
cancellation_event: Event | None = None,
# Base pipeline kwargs forwarded to AbliterationPipeline
push_to_hub: str | None = None,
hub_token: str | None = None,
@@ -216,6 +218,7 @@ class InformedAbliterationPipeline(AbliterationPipeline):
harmless_prompts=harmless_prompts,
on_stage=on_stage,
on_log=on_log,
cancellation_event=cancellation_event,
push_to_hub=push_to_hub,
hub_token=hub_token,
hub_community_org=hub_community_org,
@@ -267,24 +270,31 @@ class InformedAbliterationPipeline(AbliterationPipeline):
t0 = time.time()
# Stage 1: SUMMON
self._active_stage = "summon"
self._summon()
# Stage 2: PROBE
self._active_stage = "probe"
self._probe()
# Stage 3: ANALYZE (new stage — the feedback loop)
self._active_stage = "analyze"
self._analyze()
# Stage 4: DISTILL (informed by analysis)
self._active_stage = "distill"
self._distill_informed()
# Stage 5: EXCISE (informed by analysis)
self._active_stage = "excise"
self._excise_informed()
# Stage 6: VERIFY + Ouroboros compensation loop
self._active_stage = "verify"
self._verify_and_compensate()
# Stage 7: REBIRTH
self._active_stage = "rebirth"
output_path = self._rebirth_informed()
self._report.total_duration = time.time() - t0
obliteratus/persistence_contracts.py
+29 -10
View File
@@ -95,22 +95,23 @@ def _require_nonempty_regular_file(path: Path, artifact_name: str) -> None:
raise ValueError(f"Checkpoint {artifact_name} is missing or empty: {path}")
def validate_local_checkpoint(
def _validate_checkpoint_artifacts(
checkpoint_dir: Path,
expected_metadata_json: str,
) -> None:
"""Validate the minimum artifacts needed for a local Transformers reload."""
expected_metadata_json: str | None,
) -> Path:
"""Validate local checkpoint identity and reload-critical artifacts."""
checkpoint_dir = Path(checkpoint_dir)
if checkpoint_dir.is_symlink() or not checkpoint_dir.is_dir():
raise ValueError(f"Checkpoint staging path is not a directory: {checkpoint_dir}")
raise FileNotFoundError(f"Checkpoint path is not a directory: {checkpoint_dir}")
expected_metadata = json.loads(expected_metadata_json)
actual_metadata = _read_json_object(
checkpoint_dir / "abliteration_metadata.json",
"metadata",
)
if actual_metadata != expected_metadata:
raise ValueError("Checkpoint metadata does not match the prepared transaction")
if expected_metadata_json is not None:
expected_metadata = json.loads(expected_metadata_json)
if actual_metadata != expected_metadata:
raise ValueError("Checkpoint metadata does not match the prepared transaction")
_read_json_object(checkpoint_dir / "config.json", "model config")
_read_json_object(checkpoint_dir / "tokenizer_config.json", "tokenizer config")
@@ -122,7 +123,7 @@ def validate_local_checkpoint(
for weights_path in direct_weights:
if _path_exists(weights_path):
_require_nonempty_regular_file(weights_path, "weights")
return
return checkpoint_dir.resolve()
index_paths = [
checkpoint_dir / "model.safetensors.index.json",
@@ -144,11 +145,29 @@ def validate_local_checkpoint(
checkpoint_dir / shard_name,
"weight shard",
)
return
return checkpoint_dir.resolve()
raise ValueError(f"Checkpoint has no model weights: {checkpoint_dir}")
def validate_local_checkpoint(
checkpoint_dir: Path,
expected_metadata_json: str,
) -> None:
"""Validate a staging checkpoint against its prepared transaction."""
try:
_validate_checkpoint_artifacts(checkpoint_dir, expected_metadata_json)
except FileNotFoundError as error:
raise ValueError(
f"Checkpoint staging path is not a directory: {Path(checkpoint_dir)}",
) from error
def validate_reloadable_checkpoint(checkpoint_dir: Path | str) -> Path:
"""Return an absolute local checkpoint path only when it is reloadable."""
return _validate_checkpoint_artifacts(Path(checkpoint_dir), None)
def _path_exists(path: Path) -> bool:
"""Return whether a path or dangling symlink occupies ``path``."""
return path.exists() or path.is_symlink()
tests/test_abliterate.py
+95
View File
@@ -3,6 +3,7 @@
from __future__ import annotations
import json
import math
from pathlib import Path
from unittest.mock import MagicMock, Mock, patch
@@ -216,6 +217,70 @@ class TestPipelineInit:
assert pipeline.refusal_max_tokens == 128
assert pipeline.handle is None
def test_cancellation_is_terminal_and_cleanup_unloads_model(self):
from threading import Event
from obliteratus.abliterate import PipelineCancelledError
cancellation = Event()
pipeline = AbliterationPipeline(
model_name="test-model",
cancellation_event=cancellation,
)
pipeline.handle = MagicMock()
cancellation.set()
with pytest.raises(PipelineCancelledError, match="cancelled during PROBE"):
pipeline._check_cancelled("probe")
pipeline.cleanup_failed_run()
assert pipeline.handle.model is None
assert pipeline.handle.tokenizer is None
def test_catastrophic_perplexity_aborts_before_generation(self):
from types import SimpleNamespace
from obliteratus.abliterate import PipelineValidationError
class NonFiniteModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.anchor = torch.nn.Parameter(torch.zeros(1))
self.generate_calls = 0
def forward(self, **_kwargs):
return SimpleNamespace(loss=torch.tensor(float("nan")))
def generate(self, **_kwargs):
self.generate_calls += 1
raise AssertionError("generation must not run after catastrophic validation")
class Tokenizer:
def __call__(self, _text, **_kwargs):
return {
"input_ids": torch.tensor([[1, 2, 3]]),
"attention_mask": torch.ones((1, 3), dtype=torch.long),
}
model = NonFiniteModel()
pipeline = AbliterationPipeline(model_name="test-model")
pipeline.handle = SimpleNamespace(model=model, tokenizer=Tokenizer())
stages = []
pipeline._on_stage = stages.append
with pytest.raises(
PipelineValidationError,
match="all reference losses were non-finite",
) as exc_info:
pipeline._verify()
assert exc_info.value.stage == "verify"
assert exc_info.value.metric == "perplexity"
assert math.isinf(pipeline._quality_metrics["perplexity"])
assert model.generate_calls == 0
assert stages[-1].status == "error"
assert stages[-1].details["metric"] == "perplexity"
@pytest.mark.parametrize("invalid", [0, -1, 1.5, True])
def test_refusal_max_tokens_must_be_positive_integer(self, invalid):
with pytest.raises(
@@ -2063,6 +2128,36 @@ class TestExcise:
# ---------------------------------------------------------------------------
class TestRebirth:
def test_rebirth_cancellation_discards_staging_before_promotion(self, handle, tmp_path):
from threading import Event
from obliteratus.abliterate import PipelineCancelledError
cancellation = Event()
output = tmp_path / "output"
pipeline = AbliterationPipeline(
model_name="test-model",
output_dir=str(output),
cancellation_event=cancellation,
)
pipeline.handle = handle
pipeline._strong_layers = [0]
def save_then_cancel(path, **_kwargs):
(Path(path) / "config.json").write_text("{}", encoding="utf-8")
(Path(path) / "model.safetensors").write_bytes(b"partial")
cancellation.set()
handle.model.save_pretrained = MagicMock(side_effect=save_then_cancel)
handle.tokenizer.save_pretrained = MagicMock()
with pytest.raises(PipelineCancelledError, match="cancelled during REBIRTH"):
pipeline._rebirth()
assert not output.exists()
assert list(tmp_path.glob(".output.staging-*")) == []
handle.tokenizer.save_pretrained.assert_not_called()
def test_rebirth_saves_metadata(self, handle, tmp_path):
"""Rebirth should save model and comprehensive metadata JSON."""
pipeline = AbliterationPipeline(
tests/test_app_model_lifecycle.py
+97
View File
@@ -72,3 +72,100 @@ assert app._state["output_dir"] is None
)
assert result.returncode == 0, result.stdout + result.stderr
@pytest.mark.operator_ui
def test_pipeline_cancellation_and_local_checkpoint_reload_contracts(tmp_path):
"""Cover timeout grace outcomes and local-only reload classification."""
script = r'''
import pathlib
import sys
import threading
import time
from unittest.mock import Mock
import app
root = pathlib.Path(sys.argv[1])
checkpoint = root / "completed"
checkpoint.mkdir()
(checkpoint / "abliteration_metadata.json").write_text("{}", encoding="utf-8")
(checkpoint / "config.json").write_text("{}", encoding="utf-8")
(checkpoint / "tokenizer_config.json").write_text("{}", encoding="utf-8")
(checkpoint / "model.safetensors").write_bytes(b"weights")
assert app._resolve_local_checkpoint(checkpoint) == checkpoint.resolve()
calls = []
model = Mock()
model.to.return_value = model
tokenizer = Mock(pad_token=None, eos_token="<eos>")
app.AutoModelForCausalLM.from_pretrained = Mock(
side_effect=lambda source, **kwargs: (calls.append((source, kwargs)), model)[1]
)
app.AutoTokenizer.from_pretrained = Mock(return_value=tokenizer)
app.dev.supports_device_map_auto = lambda: True
app._load_model_to_device(checkpoint, local_files_only=True)
assert calls[0][0] == checkpoint
assert calls[0][1]["local_files_only"] is True
quantization = object()
loaded_model, loaded_tokenizer = app._reload_local_checkpoint(
checkpoint,
trust_remote_code=True,
model_kwargs={"quantization_config": quantization},
)
assert loaded_model is model and loaded_tokenizer is tokenizer
assert calls[-1][1]["quantization_config"] is quantization
assert calls[-1][1]["local_files_only"] is True
assert tokenizer.pad_token == "<eos>"
app._reload_local_checkpoint(
checkpoint,
trust_remote_code=True,
model_kwargs={"offload_folder": str(root / "offload")},
)
assert calls[-1][1]["offload_folder"] == str(root / "offload")
assert calls[-1][1]["local_files_only"] is True
assert app._format_checkpoint_reload_error(FileNotFoundError("gone")).startswith(
"Saved checkpoint is missing"
)
assert app._format_checkpoint_reload_error(ValueError("Checkpoint has no model weights")).startswith(
"Saved checkpoint is incomplete"
)
assert app._format_checkpoint_reload_error(ValueError("Checkpoint metadata is corrupt")).startswith(
"Saved checkpoint is corrupt"
)
assert app._format_checkpoint_reload_error(ValueError("Repo id is invalid")).startswith(
"Checkpoint Hub resolution failed"
)
def run_cooperative(cancel, delay):
cancel.wait()
time.sleep(delay)
cancel = threading.Event()
worker = threading.Thread(target=run_cooperative, args=(cancel, 0.02))
worker.start()
assert app._cancel_pipeline_worker(worker, cancel, grace_seconds=1.0) is True
assert cancel.is_set()
cancel = threading.Event()
release = threading.Event()
worker = threading.Thread(target=lambda: release.wait())
worker.start()
assert app._cancel_pipeline_worker(worker, cancel, grace_seconds=0.01) is False
assert cancel.is_set() and worker.is_alive()
release.set()
worker.join(timeout=1.0)
assert not worker.is_alive()
'''
result = subprocess.run(
[sys.executable, "-c", script, str(tmp_path)],
capture_output=True,
text=True,
timeout=120,
check=False,
)
assert result.returncode == 0, result.stdout + result.stderr
tests/test_persistence_contracts.py
+27
View File
@@ -31,6 +31,33 @@ def _write_valid_local_checkpoint(
(checkpoint_dir / "model.safetensors").write_bytes(b"weights")
def test_reloadable_checkpoint_resolves_absolute_and_relative_paths(tmp_path, monkeypatch):
checkpoint = tmp_path / "completed"
checkpoint.mkdir()
_write_valid_local_checkpoint(checkpoint)
assert persistence.validate_reloadable_checkpoint(checkpoint) == checkpoint.resolve()
monkeypatch.chdir(tmp_path)
assert persistence.validate_reloadable_checkpoint("completed") == checkpoint.resolve()
def test_reloadable_checkpoint_rejects_absent_incomplete_and_corrupt_paths(tmp_path):
with pytest.raises(FileNotFoundError, match="not a directory"):
persistence.validate_reloadable_checkpoint(tmp_path / "absent")
incomplete = tmp_path / "incomplete"
incomplete.mkdir()
(incomplete / "abliteration_metadata.json").write_text("{}", encoding="utf-8")
with pytest.raises(ValueError, match="model config is missing"):
persistence.validate_reloadable_checkpoint(incomplete)
corrupt = tmp_path / "corrupt"
corrupt.mkdir()
(corrupt / "abliteration_metadata.json").write_text("not-json", encoding="utf-8")
with pytest.raises(ValueError, match="metadata is corrupt"):
persistence.validate_reloadable_checkpoint(corrupt)
@pytest.mark.parametrize(
("state_dict", "expected"),
[