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 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_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,
        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 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()