OBLITERATUS/tests/test_offline_integration.py

"""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: same checkpoint as live surgery, byte for byte."""
    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()