"""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 accelerate import dispatch_model from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer from obliteratus.abliterate import AbliterationPipeline from obliteratus.config import DatasetConfig, ModelConfig, StrategyConfig, StudyConfig from obliteratus.models.loader import ModelHandle from obliteratus.reporting.report import AblationReport from obliteratus.runner import run_study from tests.fixtures.tiny_offline_model import ( build_tiny_offline_model, build_tiny_offline_moe_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, max_perplexity_increase=1000.0, max_degenerate_fraction=1.0, 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", "baseline", "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_multidirection_pipeline_restores_tiny_model_layer_norms(tmp_path): source = build_tiny_offline_model(tmp_path / "source") output = tmp_path / "output" original = _state_dict(source) layer_weights = { name: tensor.float().norm().item() for name, tensor in original.items() if ".h.0." in name and tensor.ndim >= 2 } pipeline = AbliterationPipeline( model_name=str(source), output_dir=str(output), device="cpu", dtype="float32", method="advanced", n_directions=2, norm_preserve=True, refinement_passes=1, max_seq_length=8, verify_sample_size=1, refusal_max_tokens=1, max_perplexity_increase=1000.0, max_degenerate_fraction=1.0, harmful_prompts=["harmful request", "harmful answer"], harmless_prompts=["harmless request", "harmless answer"], ) pipeline.run() restored = _state_dict(output) assert any( not torch.equal(original[name], restored[name]) for name in layer_weights ) for name, original_norm in layer_weights.items(): assert restored[name].float().norm().item() == pytest.approx( original_norm, rel=1e-5, abs=1e-7, ) metadata = json.loads((output / "abliteration_metadata.json").read_text()) assert metadata["method_config"]["n_directions"] == 2 assert metadata["method_config"]["norm_preserve"] is True def _dispatch_decoder_layers_to_disk(model, offload_dir: Path) -> None: """Offload every decoder layer to disk so its parameters become meta tensors.""" device_map = {"lm_head": "cpu"} for name, _child in model.model.named_children(): if name == "layers": for index in range(len(model.model.layers)): device_map[f"model.layers.{index}"] = "disk" else: device_map[f"model.{name}"] = "cpu" dispatch_model(model, device_map=device_map, offload_dir=str(offload_dir)) def _run_moe_pipeline(source: Path, output: Path, load_model) -> Path: import obliteratus.abliterate as abliterate_module previous = abliterate_module.load_model abliterate_module.load_model = load_model try: pipeline = AbliterationPipeline( model_name=str(source), output_dir=str(output), device="cpu", dtype="float32", method="advanced", max_seq_length=8, verify_sample_size=1, refusal_max_tokens=1, max_perplexity_increase=1000.0, max_degenerate_fraction=1.0, harmful_prompts=["harmful request", "harmful answer"], harmless_prompts=["harmless request", "harmless answer"], ) return pipeline.run() finally: abliterate_module.load_model = previous def _moe_load_model(offload_dir: Path | None, live: dict): def load_model(model_name, task, **_kwargs): config = AutoConfig.from_pretrained(model_name, local_files_only=True) model = AutoModelForCausalLM.from_pretrained( model_name, config=config, dtype=torch.float32, local_files_only=True, ) if offload_dir is not None: _dispatch_decoder_layers_to_disk(model, offload_dir) live["model"] = model return ModelHandle( model=model, tokenizer=AutoTokenizer.from_pretrained(model_name, local_files_only=True), config=config, model_name=model_name, task=task, ) return load_model # transformers emits this advisory whenever ``hf_device_map`` contains cpu/disk # entries. The pipeline materializes the complete state dict itself before # ``save_pretrained`` (see ``_gather_state_dict``), so the advisory is expected # here and unrelated warnings stay fatal. @pytest.mark.filterwarnings( "ignore:Attempting to save a model with offloaded modules.*:UserWarning", ) def test_disk_offloaded_fused_moe_surgery_matches_live_surgery(tmp_path): """Offload must be transparent: checkpoint tensors match live surgery.""" source = build_tiny_offline_moe_model(tmp_path / "source") offload_dir = tmp_path / "offload" offload_dir.mkdir() original = _state_dict(source) live_output = _run_moe_pipeline( source, tmp_path / "live", _moe_load_model(None, {}), ) offloaded_live: dict[str, torch.nn.Module] = {} offloaded_output = _run_moe_pipeline( source, tmp_path / "offloaded", _moe_load_model(offload_dir, offloaded_live), ) assert live_output == tmp_path / "live" assert offloaded_output == tmp_path / "offloaded" model = offloaded_live["model"] expert_parameters = { name: parameter for name, parameter in model.named_parameters() if ".experts." in name } assert expert_parameters, "fixture must expose fused expert parameters" assert all(parameter.dim() == 3 for parameter in expert_parameters.values()) # Surgery must leave the live model offloaded rather than materializing it. assert all(parameter.device.type == "meta" for parameter in expert_parameters.values()) live_state = _state_dict(live_output) offloaded_state = _state_dict(offloaded_output) assert offloaded_state.keys() == live_state.keys() == original.keys() assert all(tensor.device.type != "meta" for tensor in offloaded_state.values()) assert all(torch.isfinite(tensor).all() for tensor in offloaded_state.values()) changed_experts = [ name for name in expert_parameters if not torch.equal(original[name], live_state[name]) ] assert changed_experts, "live surgery did not touch the fused expert tensors" # Bitwise identical on macOS; assert_close with default float32 tolerances # keeps the contract robust to BLAS threading differences on other hosts # while still catching any real divergence (a wrong projection is ~1e-2). for name in live_state: torch.testing.assert_close( offloaded_state[name], live_state[name], msg=lambda detail, name=name: ( f"offloaded surgery diverged from live surgery on {name}: {detail}" ), ) 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", "--max-perplexity-increase", "1000", "--max-degenerate-fraction", "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()