From 8bbb5f2926b823677fcbb23562d24d90154bdfa2 Mon Sep 17 00:00:00 2001 From: Joseph Magly <1159087+jmagly@users.noreply.github.com> Date: Fri, 14 Aug 2026 13:12:41 -0400 Subject: [PATCH] test: add offline integration baseline --- .github/workflows/ci.yml | 2 +- CONTRIBUTING.md | 4 +- README.md | 15 +- obliteratus/abliterate.py | 237 +++++++++++++++--------- obliteratus/auto_obliterate.py | 35 +++- pyproject.toml | 2 +- tests/fixtures/README.md | 14 ++ tests/fixtures/__init__.py | 1 + tests/fixtures/tiny_offline_model.py | 80 ++++++++ tests/test_auto_obliterate.py | 265 +++++++++++++++++++++++++++ tests/test_checkpoint_atomicity.py | 214 +++++++++++++++++++++ tests/test_offline_integration.py | 183 ++++++++++++++++++ 12 files changed, 948 insertions(+), 104 deletions(-) create mode 100644 tests/fixtures/README.md create mode 100644 tests/fixtures/__init__.py create mode 100644 tests/fixtures/tiny_offline_model.py create mode 100644 tests/test_auto_obliterate.py create mode 100644 tests/test_checkpoint_atomicity.py create mode 100644 tests/test_offline_integration.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 1b9b8b0..561dcbb 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -282,7 +282,7 @@ jobs: run: >- "$TEST_ENV/bin/python" scripts/check_coverage_thresholds.py "test-results/coverage-py${{ matrix.python-version }}.json" - --min-line 55 + --min-line 60 --min-branch 42 --min-file obliteratus/device.py=70 --min-file obliteratus/models/loader.py=70 diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 2a873d4..c604c4a 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -32,7 +32,7 @@ python -m obliteratus --help ``` All tests must pass before submitting a PR. Tests are designed to run on CPU without downloading models. -The mandatory gate currently requires at least 55% repository statement coverage, +The mandatory gate currently requires at least 60% repository statement coverage, 42% branch coverage, 90% coverage of changed executable lines, and 70% statement coverage in the device, loader, architecture-profile, CLI, and simulated MLX boundary modules. New changes should raise these floors rather than consume the @@ -108,7 +108,7 @@ obliteratus/ models/ # Model loading utilities reporting/ # Report generation strategies/ # Ablation strategies (layer, head, FFN, embedding) -tests/ # 41 test files +tests/ # 44 test files paper/ # LaTeX paper examples/ # YAML config examples ``` diff --git a/README.md b/README.md index b0a9176..fde8bbc 100644 --- a/README.md +++ b/README.md @@ -651,7 +651,7 @@ obliteratus run examples/preset_quick.yaml | Analysis-informed abliteration | Yes (closed-loop feedback) | N/A | N/A | N/A | N/A | N/A | | Auto parameter optimization | Analysis-guided | N/A | Bayesian (Optuna) | N/A | N/A | N/A | | Model compatibility | Any HuggingFace model | ~50 architectures | 16/16 tested | TransformerLens only | HuggingFace | TransformerLens | -| Test suite | 970+ tests | Community | Unknown | None | Minimal | Moderate | +| Test suite | 1,001 tests | Community | Unknown | None | Minimal | Moderate | ## Community-powered research — every run advances the science @@ -745,7 +745,7 @@ If you use OBLITERATUS in your research, please cite: author = {{OBLITERATUS Contributors}}, year = {2026}, url = {https://github.com/elder-plinius/OBLITERATUS}, - note = {15 analysis modules, 970+ tests} + note = {15 analysis modules, 1,001 tests} } ``` @@ -756,10 +756,13 @@ pip install -e ".[dev]" pytest ``` -The mandatory CPU suite currently contains 975 tests across 41 test files, -covering the CLI, model/device/quantization/MLX boundaries, all analysis modules, -the abliteration pipeline, architecture detection, visualization sanitization, -community contributions, edge cases, and evaluation metrics. +The mandatory CPU suite currently contains 1,001 tests across 44 test files, +including a repository-owned synthetic model that exercises the offline pipeline, +installed-wheel CLI, study runner, transactional checkpoint recovery, and resumable +auto-obliteration state. The suite also covers model/device/quantization/MLX +boundaries, all analysis modules, architecture detection, visualization sanitization, +community contributions, edge cases, and evaluation metrics. CI enforces at least +60% statement coverage, 42% branch coverage, and 90% changed-line coverage. ## License diff --git a/obliteratus/abliterate.py b/obliteratus/abliterate.py index e6bebf8..8cb8249 100644 --- a/obliteratus/abliterate.py +++ b/obliteratus/abliterate.py @@ -24,8 +24,12 @@ import json import logging import math import os +import shutil +import tempfile import time +import uuid import warnings +from contextlib import contextmanager from dataclasses import dataclass, field from pathlib import Path from typing import Any, Callable, Iterable @@ -56,6 +60,66 @@ logger = logging.getLogger(__name__) # compounding norm drift across many layers/directions. _MAX_NORM_RATIO = 1.10 + +def _remove_checkpoint_path(path: Path) -> None: + """Remove a staging/backup path without following directory symlinks.""" + if path.is_symlink() or path.is_file(): + path.unlink(missing_ok=True) + elif path.exists(): + shutil.rmtree(path, ignore_errors=True) + + +@contextmanager +def _atomic_checkpoint_directory(destination: Path): + """Yield a sibling staging directory and atomically promote it on success. + + An existing checkpoint is moved to a uniquely named backup immediately + before promotion. If promotion fails, that backup is restored. Exceptions + while writing only remove the staging directory and leave the destination + untouched. + """ + destination = Path(destination) + destination.parent.mkdir(parents=True, exist_ok=True) + staging = Path( + tempfile.mkdtemp( + prefix=f".{destination.name or 'checkpoint'}.staging-", + dir=destination.parent, + ), + ) + backup: Path | None = None + try: + yield staging + if destination.exists() or destination.is_symlink(): + backup = destination.with_name( + f".{destination.name}.backup-{uuid.uuid4().hex}", + ) + os.replace(destination, backup) + try: + os.replace(staging, destination) + except Exception as promotion_error: + if backup is not None and (backup.exists() or backup.is_symlink()): + try: + os.replace(backup, destination) + except Exception as restore_error: + raise RuntimeError( + "Checkpoint promotion and rollback both failed; " + f"recover the previous checkpoint from {backup}", + ) from restore_error + raise promotion_error + if backup is not None: + try: + _remove_checkpoint_path(backup) + except OSError as cleanup_error: + logger.warning( + "Checkpoint promoted, but previous-checkpoint backup %s " + "could not be removed: %s", + backup, + cleanup_error, + ) + except Exception: + _remove_checkpoint_path(staging) + raise + # ── Abliteration method presets ─────────────────────────────────────────── METHODS = { @@ -6597,16 +6661,18 @@ class AbliterationPipeline: Only safe AFTER the state_dict has been gathered into memory — disk-offloaded weights live in this directory and would be lost. """ - import shutil as _shutil - offload_dir = getattr(self.handle, "_offload_dir", None) - if offload_dir and Path(offload_dir).exists(): - size_mb = sum( - f.stat().st_size for f in Path(offload_dir).rglob("*") if f.is_file() - ) / (1024 ** 2) - if size_mb > 0: - _shutil.rmtree(offload_dir, ignore_errors=True) + owns_offload_dir = getattr(self.handle, "_owns_offload_dir", False) + if owns_offload_dir and offload_dir: + offload_path = Path(offload_dir) + if offload_path.exists(): + size_mb = sum( + f.stat().st_size for f in offload_path.rglob("*") if f.is_file() + ) / (1024 ** 2) + shutil.rmtree(offload_path, ignore_errors=True) self.log(f"Cleaned up offload dir ({size_mb:.0f} MiB reclaimed)") + self.handle._offload_dir = None + self.handle._owns_offload_dir = False def _gather_state_dict(self) -> dict: """Gather a complete state dict, materializing any offloaded tensors. @@ -6660,9 +6726,57 @@ class AbliterationPipeline: return state_dict + def _write_local_checkpoint( + self, + checkpoint_dir: Path, + metadata: dict, + state_dict: dict, + ) -> None: + """Write every local checkpoint component into an isolated directory.""" + model = self.handle.model + + # Strip native quantization metadata (e.g. Mxfp4) so save_pretrained + # treats this as a plain float model after EXCISE dequantizes weights. + if hasattr(model, "hf_quantizer") and model.hf_quantizer is not None: + self.log("Stripping native quantization config (weights are now float16)") + model.hf_quantizer.remove_quantization_config(model) + + # Avoid unsupported reverse conversions when saving a new HF-native artifact. + if hasattr(model, "_weight_conversions"): + del model._weight_conversions + + try: + model.save_pretrained( + checkpoint_dir, + state_dict=state_dict, + max_shard_size="2GB", + save_original_format=False, + ) + except Exception as e: + msg = str(e) + if not msg: + msg = repr(e) + if hasattr(e, "errno") and e.errno is not None: + import errno as errno_mod + msg = f"{errno_mod.errorcode.get(e.errno, f'errno {e.errno}')}: {os.strerror(e.errno)}" + if e.errno == 28: # ENOSPC + disk = shutil.disk_usage(checkpoint_dir) + msg += f" ({disk.free / 1e9:.1f} GB free on {self.output_dir})" + raise type(e)(msg) from e + + self.handle.tokenizer.save_pretrained(checkpoint_dir) + (checkpoint_dir / "abliteration_metadata.json").write_text( + json.dumps(metadata, indent=2), + encoding="utf-8", + ) + + if self._lora_adapters: + from obliteratus.lora_ablation import save_lora_adapters + adapter_path = save_lora_adapters(self._lora_adapters, checkpoint_dir) + self.log(f"Saved LoRA adapters to {adapter_path}") + def _rebirth(self) -> Path: - """Save the abliterated model with comprehensive metadata.""" - import shutil + """Atomically save the abliterated model with comprehensive metadata.""" dest = self.push_to_hub or str(self.output_dir) self._emit("rebirth", "running", f"Saving to {dest}...") @@ -6679,94 +6793,35 @@ class AbliterationPipeline: param_bytes = sum(v.numel() * v.element_size() for v in state_dict.values()) self.log(f"State dict: {len(state_dict)} tensors, {param_bytes / 1e9:.1f} GB") - # 3. Save model + tokenizer + metadata - # NOTE: offload dir cleanup is deferred until AFTER save_pretrained - # completes, because accelerate's dispatch hooks may still access - # the offload dir during serialization (even when state_dict is - # explicitly provided). - self.output_dir.mkdir(parents=True, exist_ok=True) + # 3. Save every component to a sibling staging directory, then promote + # the complete checkpoint atomically. The offload dir remains intact + # until promotion so a failed write can be retried safely. self.log(f"Saving model to {self.output_dir}/") - - # Check disk space with the actual state dict size. - try: - disk = shutil.disk_usage(self.output_dir) - # Need ~1.1x the raw param bytes for safetensors overhead + metadata - needed = int(param_bytes * 1.1) - if disk.free < needed: - raise OSError( - f"Insufficient disk space: " - f"{disk.free / 1e9:.1f} GB free, need ~{param_bytes / 1e9:.1f} GB. " - f"Try a different --output-dir on a larger filesystem." + with _atomic_checkpoint_directory(self.output_dir) as checkpoint_dir: + try: + disk = shutil.disk_usage(checkpoint_dir) + needed = int(param_bytes * 1.1) + if disk.free < needed: + raise OSError( + f"Insufficient disk space: " + f"{disk.free / 1e9:.1f} GB free, need ~{param_bytes / 1e9:.1f} GB. " + f"Try a different --output-dir on a larger filesystem." + ) + self.log( + f"Disk space: {disk.free / 1e9:.1f} GB free, " + f"need ~{param_bytes / 1e9:.1f} GB", ) - self.log(f"Disk space: {disk.free / 1e9:.1f} GB free, need ~{param_bytes / 1e9:.1f} GB") - except OSError: - raise - except Exception: - pass # Non-critical — don't block save on stat failure + except OSError: + raise + except Exception: + pass # Non-critical — don't block save on stat failure + self._write_local_checkpoint(checkpoint_dir, metadata, state_dict) - # Strip native quantization metadata (e.g. Mxfp4) so save_pretrained - # treats this as a plain float model. After EXCISE the weights are - # dequantized float16 — the original quantization format is gone, and - # save_pretrained's quantizer hook would crash trying to access - # format-specific internals (Triton storage layout, etc.). - model = self.handle.model - if hasattr(model, "hf_quantizer") and model.hf_quantizer is not None: - self.log("Stripping native quantization config (weights are now float16)") - model.hf_quantizer.remove_quantization_config(model) - - # Clear _weight_conversions unconditionally. For natively-quantized - # models (e.g. MXFP4) the list includes Mxfp4Deserialize whose - # reverse_op is not implemented — revert_weight_conversion() would - # raise NotImplementedError. hf_quantizer may already be None even - # when these conversions are present, so we can't gate on it. - if hasattr(model, "_weight_conversions"): - del model._weight_conversions - - # Use 2 GB shards to reduce peak memory during serialization (default - # is 5 GB which can cause OOM when GPU tensors are copied to CPU). - # - # save_original_format=False: the abliterated model is a new artifact - # and doesn't need the original checkpoint's key naming convention. - # HF-native format loads correctly with from_pretrained. This also - # avoids revert_weight_conversion() which can fail for quantizer ops. - try: - model.save_pretrained( - self.output_dir, - state_dict=state_dict, - max_shard_size="2GB", - save_original_format=False, - ) - except Exception as e: - msg = str(e) - if not msg: - msg = repr(e) - if hasattr(e, "errno") and e.errno is not None: - import errno as errno_mod - msg = f"{errno_mod.errorcode.get(e.errno, f'errno {e.errno}')}: {os.strerror(e.errno)}" - if e.errno == 28: # ENOSPC - disk = shutil.disk_usage(self.output_dir) - msg += f" ({disk.free / 1e9:.1f} GB free on {self.output_dir})" - raise type(e)(msg) from e - - # Free the state dict to reclaim memory before tokenizer save + # Free the state dict and temporary offload only after promotion. del state_dict self._free_gpu_memory() - - # NOW it's safe to clean up the offload dir — save_pretrained is done. self._cleanup_offload_dir() - self.handle.tokenizer.save_pretrained(self.output_dir) - - (self.output_dir / "abliteration_metadata.json").write_text( - json.dumps(metadata, indent=2) - ) - - # Save LoRA adapters if they exist (reversible ablation mode) - if self._lora_adapters: - from obliteratus.lora_ablation import save_lora_adapters - adapter_path = save_lora_adapters(self._lora_adapters, self.output_dir) - self.log(f"Saved LoRA adapters to {adapter_path}") - # 5. Optionally push the saved directory to the Hub. if self.push_to_hub: from huggingface_hub import HfApi diff --git a/obliteratus/auto_obliterate.py b/obliteratus/auto_obliterate.py index 26ce96e..9522ecc 100644 --- a/obliteratus/auto_obliterate.py +++ b/obliteratus/auto_obliterate.py @@ -149,25 +149,54 @@ class AutoObliterator: try: if self._state_file.exists(): data = json.loads(self._state_file.read_text(encoding="utf-8")) - self._result = AutoObliterateResult.from_dict(data) + if not isinstance(data, dict): + raise ValueError("state root must be a JSON object") + loaded = AutoObliterateResult.from_dict(data) + if loaded.model_id != self.model_id: + raise ValueError( + f"state belongs to {loaded.model_id!r}, not {self.model_id!r}", + ) + if not all(isinstance(item, IterationResult) for item in loaded.iterations): + raise ValueError("state contains an invalid iteration entry") + self._result = loaded self._resume_from = len(self._result.iterations) logger.info( "AutoObliterate: resuming from iteration %d", self._resume_from, ) except Exception as e: - logger.warning("AutoObliterate: failed to load state: %s", e) + timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S%fZ") + quarantine = self._state_file.with_name( + f"{self._state_file.name}.corrupt-{timestamp}", + ) + try: + os.replace(self._state_file, quarantine) + logger.warning( + "AutoObliterate: invalid state quarantined at %s: %s", + quarantine, + e, + ) + except Exception as quarantine_error: + logger.warning( + "AutoObliterate: invalid state at %s could not be quarantined: %s " + "(validation error: %s)", + self._state_file, + quarantine_error, + e, + ) def _save_state(self): """Persist current state to disk.""" + tmp = self._state_file.with_suffix(".tmp") try: Path(self.output_base).mkdir(parents=True, exist_ok=True) data = self._result.to_dict() - tmp = self._state_file.with_suffix(".tmp") tmp.write_text(json.dumps(data, indent=2, default=str), encoding="utf-8") tmp.replace(self._state_file) except Exception as e: logger.warning("AutoObliterate: failed to save state: %s", e) + finally: + tmp.unlink(missing_ok=True) # ── Quick benchmark ─────────────────────────────────────────────── diff --git a/pyproject.toml b/pyproject.toml index 5cdf2db..71d6f4e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -87,7 +87,7 @@ addopts = [ "--strict-markers", "--cov=obliteratus", "--cov-report=term-missing", - "--cov-fail-under=55", + "--cov-fail-under=60", ] filterwarnings = [ "error", diff --git a/tests/fixtures/README.md b/tests/fixtures/README.md new file mode 100644 index 0000000..8cd904b --- /dev/null +++ b/tests/fixtures/README.md @@ -0,0 +1,14 @@ +# Offline model fixture + +`tiny_offline_model.py` creates a one-layer, 4,480-parameter GPT-2 causal +language model and an eleven-token WordLevel tokenizer at test time. The model +is initialized from a fixed seed (`20260814`); it is not trained and contains +no downloaded weights or dataset content. + +The generated artifact is repository-owned test data under the project's +AGPL-3.0-only license. Its purpose is software integration testing only. It +does not support research, safety, capability, or model-quality claims. + +Tests must build the fixture inside pytest's temporary directory and load it +with Hugging Face offline mode enabled. Do not replace it with a Hub model or +depend on a pre-populated cache. diff --git a/tests/fixtures/__init__.py b/tests/fixtures/__init__.py new file mode 100644 index 0000000..c3393ad --- /dev/null +++ b/tests/fixtures/__init__.py @@ -0,0 +1 @@ +"""Repository-owned deterministic fixtures for offline integration tests.""" diff --git a/tests/fixtures/tiny_offline_model.py b/tests/fixtures/tiny_offline_model.py new file mode 100644 index 0000000..c954485 --- /dev/null +++ b/tests/fixtures/tiny_offline_model.py @@ -0,0 +1,80 @@ +"""Build a deterministic, synthetic Hugging Face causal language model.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import torch +from tokenizers import Tokenizer +from tokenizers.models import WordLevel +from tokenizers.pre_tokenizers import Whitespace +from transformers import GPT2Config, GPT2LMHeadModel, PreTrainedTokenizerFast + + +FIXTURE_SEED = 20260814 +FIXTURE_VOCAB = { + "": 0, + "": 1, + "": 2, + "harmful": 3, + "harmless": 4, + "request": 5, + "answer": 6, + "hello": 7, + "world": 8, + "safe": 9, + "test": 10, +} + + +def build_tiny_offline_model(destination: Path) -> Path: + """Create a tiny random-init GPT-2 model without downloads or caches.""" + destination = Path(destination) + destination.mkdir(parents=True, exist_ok=False) + + torch.manual_seed(FIXTURE_SEED) + tokenizer_backend = Tokenizer(WordLevel(FIXTURE_VOCAB, unk_token="")) + tokenizer_backend.pre_tokenizer = Whitespace() + tokenizer = PreTrainedTokenizerFast( + tokenizer_object=tokenizer_backend, + pad_token="", + eos_token="", + unk_token="", + ) + config = GPT2Config( + vocab_size=len(FIXTURE_VOCAB), + n_positions=128, + n_ctx=128, + n_embd=16, + n_layer=1, + n_head=2, + n_inner=32, + bos_token_id=1, + eos_token_id=1, + pad_token_id=0, + ) + model = GPT2LMHeadModel(config) + model.save_pretrained(destination, safe_serialization=True) + tokenizer.save_pretrained(destination) + + manifest = { + "fixture": "tiny-offline-gpt2", + "provenance": "generated locally from configuration with random initialization", + "training_data": None, + "third_party_weights": None, + "license": "AGPL-3.0-only (part of the OBLITERATUS test suite)", + "seed": FIXTURE_SEED, + "architecture": { + "model_type": "gpt2", + "layers": 1, + "hidden_size": 16, + "attention_heads": 2, + "vocabulary_size": len(FIXTURE_VOCAB), + }, + } + (destination / "fixture-provenance.json").write_text( + json.dumps(manifest, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + return destination diff --git a/tests/test_auto_obliterate.py b/tests/test_auto_obliterate.py new file mode 100644 index 0000000..d82b658 --- /dev/null +++ b/tests/test_auto_obliterate.py @@ -0,0 +1,265 @@ +"""Behavior tests for resumable auto-obliteration orchestration.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +import obliteratus.abliterate +from obliteratus.auto_obliterate import ( + AutoObliterateResult, + AutoObliterator, + IterationResult, +) + + +pytestmark = pytest.mark.cpu + + +def _finish(generator): + yielded = [] + while True: + try: + yielded.append(next(generator)) + except StopIteration as completed: + return yielded, completed.value + + +class _SuccessfulPipeline: + created = [] + + def __init__(self, **kwargs): + self.kwargs = kwargs + self._quality_metrics = { + "perplexity": 4.25, + "coherence": 0.9, + "refusal_rate": 0.04, + "kl_divergence": 0.02, + } + self._strong_layers = [0, 1] + self._expert_directions = {0: {0: object(), 1: object()}} + self.handle = type("Handle", (), {"model": object(), "tokenizer": object()})() + self.created.append(self) + + def run(self): + Path(self.kwargs["output_dir"]).mkdir(parents=True) + self.kwargs["on_log"]("pipeline ran") + return Path(self.kwargs["output_dir"]) + + +def test_result_round_trip_ignores_unknown_forward_compatible_fields(): + original = AutoObliterateResult( + model_id="local/model", + iterations=[ + IterationResult( + iteration=1, + method="aggressive", + prompt_volume=8, + categories_targeted=["test"], + ), + ], + success=True, + ) + encoded = original.to_dict() + encoded["future_field"] = "ignored" + encoded["iterations"][0]["future_field"] = "ignored" + + restored = AutoObliterateResult.from_dict(encoded) + + assert restored == original + assert isinstance(restored.iterations[0], IterationResult) + + +def test_valid_state_resumes_at_next_iteration(tmp_path): + output = tmp_path / "state" + output.mkdir() + state = AutoObliterateResult( + model_id="local/model", + iterations=[IterationResult(1, "aggressive", 8, output_dir="first")], + ) + (output / "auto_state.json").write_text(json.dumps(state.to_dict())) + + auto = AutoObliterator("local/model", output_base=str(output), max_iterations=2) + + assert auto._resume_from == 1 + assert auto._result.iterations[0].method == "aggressive" + + +@pytest.mark.parametrize( + "invalid_state", + [ + "not json", + "[]", + '{"model_id": "different/model"}', + '{"model_id": "local/model", "iterations": [5]}', + ], +) +def test_invalid_state_is_quarantined_with_actionable_warning( + tmp_path, + caplog, + invalid_state, +): + output = tmp_path / "state" + output.mkdir() + state_file = output / "auto_state.json" + state_file.write_text(invalid_state) + + auto = AutoObliterator("local/model", output_base=str(output)) + + quarantined = list(output.glob("auto_state.json.corrupt-*")) + assert auto._resume_from == 0 + assert not state_file.exists() + assert len(quarantined) == 1 + assert quarantined[0].read_text() == invalid_state + assert "quarantined at" in caplog.text + + +def test_invalid_state_is_retained_when_quarantine_fails(tmp_path, monkeypatch, caplog): + output = tmp_path / "state" + output.mkdir() + state_file = output / "auto_state.json" + state_file.write_text("not json", encoding="utf-8") + + def fail_quarantine(_source, _target): + raise OSError("read-only filesystem") + + monkeypatch.setattr("obliteratus.auto_obliterate.os.replace", fail_quarantine) + auto = AutoObliterator("local/model", output_base=str(output)) + + assert auto._resume_from == 0 + assert state_file.read_text() == "not json" + assert "could not be quarantined" in caplog.text + assert "read-only filesystem" in caplog.text + + +def test_interrupted_state_replace_preserves_previous_checkpoint( + tmp_path, + monkeypatch, + caplog, +): + output = tmp_path / "state" + output.mkdir() + state_file = output / "auto_state.json" + state_file.write_text('{"sentinel": true}', encoding="utf-8") + auto = AutoObliterator.__new__(AutoObliterator) + auto.model_id = "local/model" + auto.output_base = str(output) + auto._state_file = state_file + auto._result = AutoObliterateResult(model_id="local/model") + + def fail_replace(_self, _target): + raise OSError("simulated interrupted replace") + + monkeypatch.setattr(Path, "replace", fail_replace) + auto._save_state() + + assert json.loads(state_file.read_text()) == {"sentinel": True} + assert not (output / "auto_state.tmp").exists() + assert "simulated interrupted replace" in caplog.text + + +def test_auto_loop_runs_pipeline_persists_metrics_and_stops_at_target( + tmp_path, + monkeypatch, +): + _SuccessfulPipeline.created.clear() + monkeypatch.setattr( + obliteratus.abliterate, + "AbliterationPipeline", + _SuccessfulPipeline, + ) + monkeypatch.setattr( + AutoObliterator, + "_get_expanded_prompts", + staticmethod(lambda _iteration: (["harmful"] * 3, ["harmless"] * 3)), + ) + auto = AutoObliterator( + "local/model", + output_base=str(tmp_path / "run"), + max_iterations=3, + target_refusal_rate=0.05, + trust_remote_code=False, + ) + + yielded, result = _finish(auto.run()) + + assert result.success is True + assert result.final_refusal_rate == 0.04 + assert len(result.iterations) == 1 + assert result.iterations[0].strong_layers == 2 + assert result.iterations[0].ega_expert_dirs == 2 + assert result.final_output_dir.endswith("iter_1") + assert len(yielded) == 5 + assert yielded[-1][0] == "✅ Complete" + created = _SuccessfulPipeline.created[0] + assert created.kwargs["model_name"] == "local/model" + assert created.kwargs["method"] == "aggressive" + assert created.kwargs["harmful_prompts"] == ["harmful"] * 3 + saved = json.loads((tmp_path / "run" / "auto_state.json").read_text()) + assert saved["success"] is True + + +def test_auto_loop_records_failures_and_completes_without_success( + tmp_path, + monkeypatch, +): + class FailingPipeline: + def __init__(self, **_kwargs): + pass + + def run(self): + raise RuntimeError("simulated pipeline failure") + + monkeypatch.setattr(obliteratus.abliterate, "AbliterationPipeline", FailingPipeline) + monkeypatch.setattr( + AutoObliterator, + "_get_expanded_prompts", + staticmethod(lambda _iteration: (["harmful"], ["harmless"])), + ) + auto = AutoObliterator( + "local/model", + output_base=str(tmp_path / "run"), + max_iterations=1, + ) + + yielded, result = _finish(auto.run()) + + assert result.success is False + assert result.final_output_dir == "" + assert result.iterations[0].error == "simulated pipeline failure" + assert yielded[-1][0] == "⚠️ Complete (target not met)" + assert "simulated pipeline" in yielded[-1][2] + + +def test_prompt_expansion_and_benchmark_fallbacks(monkeypatch): + monkeypatch.delenv("OPENROUTER_API_KEY", raising=False) + assert AutoObliterator._quick_benchmark_claude("missing", "model") == { + "method": "skipped", + "reason": "no OPENROUTER_API_KEY", + } + pipeline = type( + "Pipeline", + (), + {"_quality_metrics": {"refusal_rate": 0.2, "coherence": 0.8}}, + )() + assert AutoObliterator._quick_benchmark_heuristic(pipeline) == { + "refusal_rate": 0.2, + "perplexity": None, + "coherence": 0.8, + "kl_divergence": None, + "method": "heuristic", + } + + +def test_reset_clears_persisted_state(tmp_path): + auto = AutoObliterator("local/model", output_base=str(tmp_path / "run")) + auto._result.iterations.append(IterationResult(1, "aggressive", 1)) + auto._save_state() + + auto.reset() + + assert auto._resume_from == 0 + assert auto._result.iterations == [] + assert not auto._state_file.exists() diff --git a/tests/test_checkpoint_atomicity.py b/tests/test_checkpoint_atomicity.py new file mode 100644 index 0000000..9868e95 --- /dev/null +++ b/tests/test_checkpoint_atomicity.py @@ -0,0 +1,214 @@ +"""Failure and recovery tests for transactional checkpoint writes.""" + +from __future__ import annotations + +import os +from pathlib import Path +from unittest.mock import MagicMock + +import pytest +import torch + +import obliteratus.abliterate as abliterate +from obliteratus.abliterate import AbliterationPipeline, _atomic_checkpoint_directory + + +pytestmark = pytest.mark.cpu + + +def _temporary_artifacts(parent: Path, name: str) -> list[Path]: + return [ + *parent.glob(f".{name}.staging-*"), + *parent.glob(f".{name}.backup-*"), + ] + + +def test_atomic_checkpoint_replaces_existing_destination(tmp_path): + destination = tmp_path / "checkpoint" + destination.mkdir() + (destination / "old.txt").write_text("old", encoding="utf-8") + + with _atomic_checkpoint_directory(destination) as staging: + (staging / "new.txt").write_text("new", encoding="utf-8") + + assert not (destination / "old.txt").exists() + assert (destination / "new.txt").read_text() == "new" + assert _temporary_artifacts(tmp_path, destination.name) == [] + + +def test_atomic_checkpoint_replaces_an_invalid_file_destination(tmp_path): + destination = tmp_path / "checkpoint" + destination.write_text("not a checkpoint", encoding="utf-8") + + with _atomic_checkpoint_directory(destination) as staging: + (staging / "config.json").write_text("{}", encoding="utf-8") + + assert destination.is_dir() + assert (destination / "config.json").read_text() == "{}" + assert _temporary_artifacts(tmp_path, destination.name) == [] + + +def test_atomic_checkpoint_warns_if_obsolete_backup_cannot_be_removed( + tmp_path, + monkeypatch, + caplog, +): + destination = tmp_path / "checkpoint" + destination.write_text("old", encoding="utf-8") + real_remove = abliterate._remove_checkpoint_path + + def fail_backup_cleanup(path): + if ".backup-" in path.name: + raise PermissionError("simulated cleanup denial") + return real_remove(path) + + monkeypatch.setattr(abliterate, "_remove_checkpoint_path", fail_backup_cleanup) + with _atomic_checkpoint_directory(destination) as staging: + (staging / "config.json").write_text("{}", encoding="utf-8") + + assert (destination / "config.json").is_file() + assert "could not be removed" in caplog.text + backups = list(tmp_path.glob(".checkpoint.backup-*")) + assert len(backups) == 1 + assert backups[0].read_text() == "old" + + +def test_atomic_checkpoint_preserves_destination_on_write_failure(tmp_path): + destination = tmp_path / "checkpoint" + destination.mkdir() + sentinel = destination / "sentinel.txt" + sentinel.write_text("preserve me", encoding="utf-8") + + with pytest.raises(OSError, match="simulated write failure"): + with _atomic_checkpoint_directory(destination) as staging: + (staging / "partial.bin").write_bytes(b"partial") + raise OSError("simulated write failure") + + assert sentinel.read_text() == "preserve me" + assert _temporary_artifacts(tmp_path, destination.name) == [] + + +def test_atomic_checkpoint_restores_destination_on_promotion_failure( + tmp_path, + monkeypatch, +): + destination = tmp_path / "checkpoint" + destination.mkdir() + sentinel = destination / "sentinel.txt" + sentinel.write_text("preserve me", encoding="utf-8") + real_replace = os.replace + + def fail_staging_promotion(source, target): + if ".staging-" in Path(source).name: + raise OSError("simulated promotion failure") + return real_replace(source, target) + + monkeypatch.setattr(abliterate.os, "replace", fail_staging_promotion) + with pytest.raises(OSError, match="simulated promotion failure"): + with _atomic_checkpoint_directory(destination) as staging: + (staging / "new.txt").write_text("new", encoding="utf-8") + + assert sentinel.read_text() == "preserve me" + assert _temporary_artifacts(tmp_path, destination.name) == [] + + +def test_atomic_checkpoint_reports_recoverable_backup_when_rollback_fails( + tmp_path, + monkeypatch, +): + destination = tmp_path / "checkpoint" + destination.mkdir() + (destination / "sentinel.txt").write_text("recover me", encoding="utf-8") + real_replace = os.replace + replacement_calls = 0 + + def fail_promotion_and_restore(source, target): + nonlocal replacement_calls + replacement_calls += 1 + if replacement_calls >= 2: + raise OSError("simulated replace failure") + return real_replace(source, target) + + monkeypatch.setattr(abliterate.os, "replace", fail_promotion_and_restore) + with pytest.raises(RuntimeError, match="recover the previous checkpoint from"): + with _atomic_checkpoint_directory(destination) as staging: + (staging / "new.txt").write_text("new", encoding="utf-8") + + backups = list(tmp_path.glob(".checkpoint.backup-*")) + assert len(backups) == 1 + assert (backups[0] / "sentinel.txt").read_text() == "recover me" + + +def test_rebirth_failure_preserves_checkpoint_and_owned_offload(tmp_path): + destination = tmp_path / "checkpoint" + destination.mkdir() + sentinel = destination / "sentinel.txt" + sentinel.write_text("preserve me", encoding="utf-8") + offload = tmp_path / "owned-offload" + offload.mkdir() + (offload / "weight.bin").write_bytes(b"still needed") + + pipeline = AbliterationPipeline( + model_name="test-model", + output_dir=str(destination), + method="basic", + ) + pipeline._on_log = lambda _message: None + pipeline._on_stage = lambda _event: None + pipeline.handle = MagicMock() + pipeline.handle.model.state_dict.return_value = {"weight": torch.ones(1)} + pipeline.handle.model.save_pretrained.side_effect = OSError("disk vanished") + pipeline.handle._offload_dir = str(offload) + pipeline.handle._owns_offload_dir = True + + with pytest.raises(OSError, match="disk vanished"): + pipeline._rebirth() + + assert sentinel.read_text() == "preserve me" + assert (offload / "weight.bin").read_bytes() == b"still needed" + assert _temporary_artifacts(tmp_path, destination.name) == [] + + +def test_cleanup_only_removes_pipeline_owned_offload_directory(tmp_path): + caller_owned = tmp_path / "caller-owned" + caller_owned.mkdir() + (caller_owned / "weight.bin").write_bytes(b"owned by caller") + pipeline = AbliterationPipeline(model_name="test-model", method="basic") + pipeline._on_log = lambda _message: None + pipeline.handle = MagicMock() + pipeline.handle._offload_dir = str(caller_owned) + pipeline.handle._owns_offload_dir = False + + pipeline._cleanup_offload_dir() + + assert (caller_owned / "weight.bin").read_bytes() == b"owned by caller" + + +def test_cleanup_removes_and_clears_pipeline_owned_offload_directory(tmp_path): + owned = tmp_path / "pipeline-owned" + owned.mkdir() + (owned / "weight.bin").write_bytes(b"temporary") + pipeline = AbliterationPipeline(model_name="test-model", method="basic") + pipeline._on_log = lambda _message: None + pipeline.handle = MagicMock() + pipeline.handle._offload_dir = str(owned) + pipeline.handle._owns_offload_dir = True + + pipeline._cleanup_offload_dir() + + assert not owned.exists() + assert pipeline.handle._offload_dir is None + assert pipeline.handle._owns_offload_dir is False + + +def test_cleanup_clears_stale_owned_offload_reference(tmp_path): + pipeline = AbliterationPipeline(model_name="test-model", method="basic") + pipeline._on_log = lambda _message: None + pipeline.handle = MagicMock() + pipeline.handle._offload_dir = str(tmp_path / "already-gone") + pipeline.handle._owns_offload_dir = True + + pipeline._cleanup_offload_dir() + + assert pipeline.handle._offload_dir is None + assert pipeline.handle._owns_offload_dir is False diff --git a/tests/test_offline_integration.py b/tests/test_offline_integration.py new file mode 100644 index 0000000..90e3723 --- /dev/null +++ b/tests/test_offline_integration.py @@ -0,0 +1,183 @@ +"""Offline integration coverage across real model, pipeline, and CLI boundaries.""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +from pathlib import Path + +import pytest +import torch +from datasets import Dataset +from transformers import AutoModelForCausalLM, AutoTokenizer + +from obliteratus.abliterate import AbliterationPipeline +from obliteratus.config import DatasetConfig, ModelConfig, StrategyConfig, StudyConfig +from obliteratus.reporting.report import AblationReport +from obliteratus.runner import run_study +from tests.fixtures.tiny_offline_model import build_tiny_offline_model + + +pytestmark = [pytest.mark.cpu, pytest.mark.integration] + + +def _state_dict(path: Path) -> dict[str, torch.Tensor]: + return AutoModelForCausalLM.from_pretrained( + path, + local_files_only=True, + ).state_dict() + + +def test_fixture_is_deterministic_and_documents_provenance(tmp_path): + first = build_tiny_offline_model(tmp_path / "first") + second = build_tiny_offline_model(tmp_path / "second") + + first_state = _state_dict(first) + second_state = _state_dict(second) + assert first_state.keys() == second_state.keys() + assert all(torch.equal(first_state[key], second_state[key]) for key in first_state) + + first_manifest = json.loads((first / "fixture-provenance.json").read_text()) + second_manifest = json.loads((second / "fixture-provenance.json").read_text()) + assert first_manifest == second_manifest + assert first_manifest["training_data"] is None + assert first_manifest["third_party_weights"] is None + + +def test_full_pipeline_saves_and_reloads_a_real_offline_model(tmp_path): + source = build_tiny_offline_model(tmp_path / "source") + output = tmp_path / "output" + events = [] + original = _state_dict(source) + + pipeline = AbliterationPipeline( + model_name=str(source), + output_dir=str(output), + device="cpu", + dtype="float32", + method="basic", + n_directions=1, + max_seq_length=8, + verify_sample_size=1, + harmful_prompts=["harmful request"], + harmless_prompts=["harmless request"], + on_stage=events.append, + ) + result = pipeline.run() + + assert result == output + assert [(event.stage, event.status) for event in events] == [ + (stage, status) + for stage in ("summon", "probe", "distill", "excise", "verify", "rebirth") + for status in ("running", "done") + ] + assert (output / "abliteration_metadata.json").is_file() + assert not list(tmp_path.glob(".output.staging-*")) + assert not list(tmp_path.glob(".output.backup-*")) + + reloaded = AutoModelForCausalLM.from_pretrained(output, local_files_only=True) + tokenizer = AutoTokenizer.from_pretrained(output, local_files_only=True) + batch = tokenizer("hello world", return_tensors="pt") + with torch.no_grad(): + logits = reloaded(**batch).logits + assert logits.shape == (1, 2, len(tokenizer)) + assert torch.isfinite(logits).all() + assert any( + not torch.equal(original[name], tensor) + for name, tensor in reloaded.state_dict().items() + ) + assert set(pipeline._stage_durations) == { + "summon", + "probe", + "distill", + "excise", + "verify", + "rebirth", + } + assert all(duration >= 0 for duration in pipeline._stage_durations.values()) + + +def test_installed_wheel_cli_loads_local_model_without_repository_imports(tmp_path): + source = build_tiny_offline_model(tmp_path / "source") + isolated_workdir = tmp_path / "outside-repository" + isolated_workdir.mkdir() + isolated_home = tmp_path / "home" + isolated_home.mkdir() + env = { + **os.environ, + "HOME": str(isolated_home), + "HF_HOME": str(isolated_home / "hf"), + "HF_DATASETS_OFFLINE": "1", + "HF_HUB_DISABLE_TELEMETRY": "1", + "HF_HUB_OFFLINE": "1", + "TRANSFORMERS_OFFLINE": "1", + } + + origin = subprocess.run( + [sys.executable, "-I", "-c", "import obliteratus; print(obliteratus.__file__)"], + cwd=isolated_workdir, + env=env, + check=True, + capture_output=True, + text=True, + timeout=30, + ) + assert "site-packages" in origin.stdout + + result = subprocess.run( + [ + sys.executable, + "-I", + "-m", + "obliteratus", + "info", + str(source), + "--device", + "cpu", + "--dtype", + "float32", + ], + cwd=isolated_workdir, + env=env, + check=True, + capture_output=True, + text=True, + timeout=30, + ) + assert "architecture:" in result.stdout.lower() + assert "gpt2" in result.stdout.lower() + + +def test_study_runner_evaluates_ablates_restores_and_reports( + tmp_path, + monkeypatch, +): + source = build_tiny_offline_model(tmp_path / "source") + output = tmp_path / "study-results" + dataset = Dataset.from_dict({"text": ["hello world safe test"]}) + monkeypatch.setattr("obliteratus.runner.load_dataset", lambda **_kwargs: dataset) + monkeypatch.setattr(AblationReport, "plot_impact", lambda *_args, **_kwargs: None) + monkeypatch.setattr(AblationReport, "plot_heatmap", lambda *_args, **_kwargs: None) + config = StudyConfig( + model=ModelConfig(name=str(source), device="cpu", dtype="float32"), + dataset=DatasetConfig(name="synthetic/offline", max_samples=1), + strategies=[StrategyConfig(name="layer_removal")], + metrics=["perplexity"], + batch_size=1, + max_length=8, + output_dir=str(output), + ) + + report = run_study(config) + + assert report.model_name == str(source) + assert report.baseline_metrics["perplexity"] > 0 + assert len(report.results) == 1 + assert report.results[0].strategy == "layer_removal" + assert report.results[0].component == "layer_0" + assert report.results[0].metrics["perplexity"] > 0 + saved = json.loads((output / "results.json").read_text()) + assert saved["baseline_metrics"] == report.baseline_metrics + assert (output / "results.csv").is_file()