diff --git a/app.py b/app.py index 52778c4..a60a1bc 100644 --- a/app.py +++ b/app.py @@ -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 diff --git a/obliteratus/abliterate.py b/obliteratus/abliterate.py index 2d93e22..fce7fcb 100644 --- a/obliteratus/abliterate.py +++ b/obliteratus/abliterate.py @@ -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", diff --git a/obliteratus/informed_pipeline.py b/obliteratus/informed_pipeline.py index 24d7f1f..562c79f 100644 --- a/obliteratus/informed_pipeline.py +++ b/obliteratus/informed_pipeline.py @@ -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 diff --git a/obliteratus/persistence_contracts.py b/obliteratus/persistence_contracts.py index 9ff91e4..6ca7401 100644 --- a/obliteratus/persistence_contracts.py +++ b/obliteratus/persistence_contracts.py @@ -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() diff --git a/tests/test_abliterate.py b/tests/test_abliterate.py index fc62109..ca3c51f 100644 --- a/tests/test_abliterate.py +++ b/tests/test_abliterate.py @@ -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( diff --git a/tests/test_app_model_lifecycle.py b/tests/test_app_model_lifecycle.py index 15e99cc..1b8cbc0 100644 --- a/tests/test_app_model_lifecycle.py +++ b/tests/test_app_model_lifecycle.py @@ -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="") +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 == "" + +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 diff --git a/tests/test_persistence_contracts.py b/tests/test_persistence_contracts.py index e04bfc3..5d86586 100644 --- a/tests/test_persistence_contracts.py +++ b/tests/test_persistence_contracts.py @@ -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"), [