mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-17 16:37:30 +02:00
test: add offline integration baseline
This commit is contained in:
Vendored
+14
@@ -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.
|
||||
Vendored
+1
@@ -0,0 +1 @@
|
||||
"""Repository-owned deterministic fixtures for offline integration tests."""
|
||||
Vendored
+80
@@ -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 = {
|
||||
"<pad>": 0,
|
||||
"<eos>": 1,
|
||||
"<unk>": 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="<unk>"))
|
||||
tokenizer_backend.pre_tokenizer = Whitespace()
|
||||
tokenizer = PreTrainedTokenizerFast(
|
||||
tokenizer_object=tokenizer_backend,
|
||||
pad_token="<pad>",
|
||||
eos_token="<eos>",
|
||||
unk_token="<unk>",
|
||||
)
|
||||
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
|
||||
@@ -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()
|
||||
@@ -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
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user