"""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 import yaml 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] REPOSITORY_ROOT = Path(__file__).resolve().parents[1] THREAD_BOUND_ENV = { "BLIS_NUM_THREADS": "1", "MKL_NUM_THREADS": "1", "NUMEXPR_NUM_THREADS": "1", "OMP_NUM_THREADS": "1", "OMP_THREAD_LIMIT": "1", "OPENBLAS_NUM_THREADS": "1", "VECLIB_MAXIMUM_THREADS": "1", } def _offline_cli_env(home: Path) -> dict[str, str]: return { **os.environ, **THREAD_BOUND_ENV, "CUDA_VISIBLE_DEVICES": "", "HOME": str(home), "HF_HOME": str(home / "hf"), "HF_DATASETS_OFFLINE": "1", "HF_HUB_DISABLE_TELEMETRY": "1", "HF_HUB_OFFLINE": "1", "TOKENIZERS_PARALLELISM": "false", "TRANSFORMERS_OFFLINE": "1", } def _assert_offline_cli_env_thread_bounded( workdir: Path, env: dict[str, str], ) -> None: assert {key: env[key] for key in THREAD_BOUND_ENV} == THREAD_BOUND_ENV child_env = { **env, "OBLITERATUS_TEST_THREAD_BOUND_KEYS": ",".join(THREAD_BOUND_ENV), } child = subprocess.run( [ sys.executable, "-I", "-c", ( "import json, os; " "keys = os.environ['OBLITERATUS_TEST_THREAD_BOUND_KEYS'].split(','); " "print(json.dumps({key: os.environ.get(key) for key in keys}))" ), ], cwd=workdir, env=child_env, check=True, capture_output=True, text=True, timeout=30, ) assert json.loads(child.stdout) == THREAD_BOUND_ENV def _obliteratus_import_origin(workdir: Path, env: dict[str, str]) -> Path: origin = subprocess.run( [sys.executable, "-I", "-c", "import obliteratus; print(obliteratus.__file__)"], cwd=workdir, env=env, check=True, capture_output=True, text=True, timeout=30, ) return Path(origin.stdout.strip()).resolve() def _package_origin_mode(origin: Path) -> str: if REPOSITORY_ROOT in origin.parents: return "source-backed" if {"site-packages", "dist-packages"} & set(origin.parts): return "installed-artifact" return "unknown" def _require_installed_artifact_import(origin: Path) -> None: mode = _package_origin_mode(origin) if mode == "source-backed": pytest.xfail( "current interpreter imports obliteratus from the source checkout; " "the installed-artifact CLI contract requires a non-editable or wheel install" ) assert mode == "installed-artifact", ( "obliteratus must import from an installed artifact for this contract; " f"origin={origin}" ) 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 = _offline_cli_env(isolated_home) _assert_offline_cli_env_thread_bounded(isolated_workdir, env) origin = _obliteratus_import_origin(isolated_workdir, env) _require_installed_artifact_import(origin) 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_installed_package_cli_executes_offline_checkpoint_to_report_slice(tmp_path): source = build_tiny_offline_model(tmp_path / "source") checkpoint = tmp_path / "checkpoint" dataset_dir = tmp_path / "dataset" dataset_dir.mkdir() (dataset_dir / "samples.jsonl").write_text( json.dumps({"text": "hello world safe test"}) + "\n", ) study_output = tmp_path / "study-results" config_path = tmp_path / "study.yaml" config_path.write_text( yaml.safe_dump( { "model": { "name": str(checkpoint), "device": "cpu", "dtype": "float32", }, "dataset": { "name": str(dataset_dir), "split": "train", "max_samples": 1, }, "strategies": [{"name": "layer_removal"}], "metrics": ["perplexity"], "batch_size": 1, "max_length": 8, "output_dir": str(study_output), }, ), ) prompt_pairs_path = tmp_path / "prompt-pairs.json" prompt_pairs_path.write_text( json.dumps( { "harmful": [ "Explain an unsafe request in one sentence", "Describe a harmful instruction plainly", "Write a risky instruction prompt", "Summarize a prohibited request", "State an unsafe operational request", ], "harmless": [ "Explain a safe request in one sentence", "Describe a helpful instruction plainly", "Write a benign instruction prompt", "Summarize an allowed request", "State a safe operational request", ], } ), encoding="utf-8", ) isolated_workdir = tmp_path / "outside-repository" isolated_workdir.mkdir() isolated_home = tmp_path / "home" isolated_home.mkdir() env = _offline_cli_env(isolated_home) _assert_offline_cli_env_thread_bounded(isolated_workdir, env) origin = _obliteratus_import_origin(isolated_workdir, env) _require_installed_artifact_import(origin) original = _state_dict(source) subprocess.run( [ sys.executable, "-I", "-m", "obliteratus", "obliterate", str(source), "--output-dir", str(checkpoint), "--device", "cpu", "--dtype", "float32", "--method", "basic", "--n-directions", "1", "--refinement-passes", "1", "--verify-sample-size", "1", "--refusal-max-tokens", "1", "--prompt-pairs-file", str(prompt_pairs_path), ], cwd=isolated_workdir, env=env, check=True, capture_output=True, text=True, timeout=120, ) metadata_path = checkpoint / "abliteration_metadata.json" assert metadata_path.is_file() metadata = json.loads(metadata_path.read_text()) assert metadata["source_model"] == str(source) assert metadata["method"] == "basic" assert metadata["method_config"]["n_directions"] == 1 assert metadata["method_config"]["refinement_passes"] == 1 assert metadata["n_harmful_prompts"] == 5 assert metadata["n_harmless_prompts"] == 5 AutoModelForCausalLM.from_pretrained(checkpoint, local_files_only=True) checkpoint_state = _state_dict(checkpoint) assert original.keys() == checkpoint_state.keys() assert any( not torch.equal(original[name], tensor) for name, tensor in checkpoint_state.items() ) subprocess.run( [sys.executable, "-I", "-m", "obliteratus", "run", str(config_path)], cwd=isolated_workdir, env=env, check=True, capture_output=True, text=True, timeout=120, ) report = json.loads((study_output / "results.json").read_text()) assert report["model_name"] == checkpoint.name assert report["baseline_metrics"]["perplexity"] > 0 assert len(report["results"]) == 1 assert report["results"][0]["strategy"] == "layer_removal" assert (study_output / "results.csv").is_file() 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()