test: enforce transactional persistence contracts

tests/test_checkpoint_atomicity.py

@@ -9,7 +9,7 @@ from unittest.mock import MagicMock
 import pytest
 import torch
 
-import obliteratus.abliterate as abliterate
+import obliteratus.persistence_contracts as persistence
 from obliteratus.abliterate import AbliterationPipeline, _atomic_checkpoint_directory
 
 
@@ -55,14 +55,14 @@ def test_atomic_checkpoint_warns_if_obsolete_backup_cannot_be_removed(
 ):
     destination = tmp_path / "checkpoint"
     destination.write_text("old", encoding="utf-8")
-    real_remove = abliterate._remove_checkpoint_path
+    real_remove = persistence._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)
+    monkeypatch.setattr(persistence, "_remove_checkpoint_path", fail_backup_cleanup)
     with _atomic_checkpoint_directory(destination) as staging:
         (staging / "config.json").write_text("{}", encoding="utf-8")
 
@@ -103,7 +103,7 @@ def test_atomic_checkpoint_restores_destination_on_promotion_failure(
             raise OSError("simulated promotion failure")
         return real_replace(source, target)
 
-    monkeypatch.setattr(abliterate.os, "replace", fail_staging_promotion)
+    monkeypatch.setattr(persistence.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")
@@ -129,7 +129,7 @@ def test_atomic_checkpoint_reports_recoverable_backup_when_rollback_fails(
             raise OSError("simulated replace failure")
         return real_replace(source, target)
 
-    monkeypatch.setattr(abliterate.os, "replace", fail_promotion_and_restore)
+    monkeypatch.setattr(persistence.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")

tests/test_persistence_contracts.py

@@ -0,0 +1,215 @@
+"""Boundary contracts for deterministic and transactional checkpoint state."""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+import obliteratus.persistence_contracts as persistence
+
+
+pytestmark = pytest.mark.cpu
+
+@pytest.mark.parametrize(
+    ("state_dict", "expected"),
+    [
+        ({}, 0),
+        ({"float": torch.ones(3, dtype=torch.float32)}, 12),
+        (
+            {
+                "half": torch.ones(5, dtype=torch.float16),
+                "index": torch.ones(2, dtype=torch.int64),
+            },
+            26,
+        ),
+    ],
+)
+def test_state_dict_size_bytes_is_exact(state_dict, expected):
+    assert persistence.state_dict_size_bytes(state_dict) == expected
+
+
+@pytest.mark.parametrize(
+    ("payload", "required"),
+    [(0, 0), (1, 2), (9, 10), (10, 11), (101, 112)],
+)
+def test_required_checkpoint_bytes_rounds_headroom_up(payload, required):
+    assert persistence.required_checkpoint_bytes(payload) == required
+
+
+def test_required_checkpoint_bytes_rejects_negative_payload():
+    with pytest.raises(ValueError) as exc_info:
+        persistence.required_checkpoint_bytes(-1)
+
+    assert str(exc_info.value) == "checkpoint payload size cannot be negative"
+
+
+def test_checkpoint_capacity_accepts_exact_boundary_and_rejects_one_byte_less():
+    persistence.ensure_checkpoint_capacity(free_bytes=11, payload_bytes=10)
+
+    with pytest.raises(
+        OSError,
+        match=r"Insufficient disk space: 0\.0 GB free, need ~0\.0 GB",
+    ):
+        persistence.ensure_checkpoint_capacity(free_bytes=10, payload_bytes=10)
+
+
+def test_checkpoint_capacity_error_reports_decimal_gigabytes_exactly():
+    payload_bytes = 1_000_000_000_000_000_000
+
+    with pytest.raises(OSError) as exc_info:
+        persistence.ensure_checkpoint_capacity(
+            free_bytes=payload_bytes,
+            payload_bytes=payload_bytes,
+        )
+
+    assert str(exc_info.value) == (
+        "Insufficient disk space: 1000000000.0 GB free, "
+        "need ~1000000000.0 GB. "
+        "Try a different --output-dir on a larger filesystem."
+    )
+
+
+def test_metadata_serialization_is_stable_strict_json():
+    first = persistence.serialize_checkpoint_metadata({"z": 1, "a": [True, None]})
+    second = persistence.serialize_checkpoint_metadata({"a": [True, None], "z": 1})
+
+    assert first == second
+    assert json.loads(first) == {"a": [True, None], "z": 1}
+    assert first == (
+        '{\n  "a": [\n    true,\n    null\n  ],\n  "z": 1\n}'
+    )
+
+
+def test_metadata_serializer_passes_strict_format_options(monkeypatch):
+    real_dumps = json.dumps
+    observed = {}
+
+    def record_dumps(metadata, **kwargs):
+        observed.update(kwargs)
+        return real_dumps(metadata, **kwargs)
+
+    monkeypatch.setattr(persistence.json, "dumps", record_dumps)
+
+    assert json.loads(persistence.serialize_checkpoint_metadata({"schema": 1})) == {
+        "schema": 1,
+    }
+    assert observed == {"indent": 2, "sort_keys": True, "allow_nan": False}
+
+
+@pytest.mark.parametrize("value", [object(), float("nan"), float("inf")])
+def test_metadata_serialization_rejects_nonportable_values(value):
+    with pytest.raises((TypeError, ValueError)):
+        persistence.serialize_checkpoint_metadata({"invalid": value})
+
+
+def test_remove_checkpoint_path_handles_file_directory_symlink_and_missing(tmp_path):
+    file_path = tmp_path / "file"
+    file_path.write_text("data", encoding="utf-8")
+    directory = tmp_path / "directory"
+    directory.mkdir()
+    (directory / "nested").write_text("data", encoding="utf-8")
+    target = tmp_path / "target"
+    target.mkdir()
+    (target / "sentinel").write_text("preserve", encoding="utf-8")
+    symlink = tmp_path / "link"
+    symlink.symlink_to(target, target_is_directory=True)
+
+    for path in (file_path, directory, symlink, tmp_path / "missing"):
+        persistence._remove_checkpoint_path(path)
+
+    assert not file_path.exists()
+    assert not directory.exists()
+    assert not symlink.exists()
+    assert (target / "sentinel").read_text(encoding="utf-8") == "preserve"
+
+
+def test_remove_checkpoint_path_requests_race_safe_file_unlink():
+    path = MagicMock()
+    path.is_symlink.return_value = True
+
+    persistence._remove_checkpoint_path(path)
+
+    path.unlink.assert_called_once_with(missing_ok=True)
+    path.is_file.assert_not_called()
+
+
+def test_remove_checkpoint_path_requests_best_effort_directory_cleanup(monkeypatch):
+    path = MagicMock()
+    path.is_symlink.return_value = False
+    path.is_file.return_value = False
+    path.exists.return_value = True
+    remove_tree = MagicMock()
+    monkeypatch.setattr(persistence.shutil, "rmtree", remove_tree)
+
+    persistence._remove_checkpoint_path(path)
+
+    remove_tree.assert_called_once_with(path, ignore_errors=True)
+
+
+def test_atomic_checkpoint_creates_parent_and_promotes_new_destination(tmp_path):
+    destination = tmp_path / "nested" / "checkpoint"
+
+    with persistence.atomic_checkpoint_directory(destination) as staging:
+        assert staging.parent == destination.parent
+        assert staging.name.startswith(".checkpoint.staging-")
+        (staging / "model.bin").write_bytes(b"complete")
+
+    assert (destination / "model.bin").read_bytes() == b"complete"
+    assert list(destination.parent.glob(".checkpoint.*-*")) == []
+
+
+def test_atomic_checkpoint_cleans_staging_when_new_destination_promotion_fails(
+    tmp_path,
+    monkeypatch,
+):
+    destination = tmp_path / "checkpoint"
+
+    def fail_promotion(source, target):
+        raise OSError(f"cannot promote {Path(source).name} to {Path(target).name}")
+
+    monkeypatch.setattr(persistence.os, "replace", fail_promotion)
+    with pytest.raises(OSError, match="cannot promote"):
+        with persistence.atomic_checkpoint_directory(destination) as staging:
+            (staging / "model.bin").write_bytes(b"partial")
+
+    assert not destination.exists()
+    assert list(tmp_path.glob(".checkpoint.staging-*")) == []
+
+
+def test_atomic_checkpoint_replaces_symlink_without_touching_target(tmp_path):
+    target = tmp_path / "target"
+    target.mkdir()
+    sentinel = target / "sentinel"
+    sentinel.write_text("preserve", encoding="utf-8")
+    destination = tmp_path / "checkpoint"
+    destination.symlink_to(target, target_is_directory=True)
+
+    with persistence.atomic_checkpoint_directory(destination) as staging:
+        (staging / "model.bin").write_bytes(b"replacement")
+
+    assert not destination.is_symlink()
+    assert (destination / "model.bin").read_bytes() == b"replacement"
+    assert sentinel.read_text(encoding="utf-8") == "preserve"
+
+
+def test_atomic_checkpoint_uses_os_replace_not_copy(tmp_path, monkeypatch):
+    destination = tmp_path / "checkpoint"
+    calls = []
+    real_replace = os.replace
+
+    def record_replace(source, target):
+        calls.append((Path(source), Path(target)))
+        return real_replace(source, target)
+
+    monkeypatch.setattr(persistence.os, "replace", record_replace)
+    with persistence.atomic_checkpoint_directory(destination) as staging:
+        (staging / "model.bin").write_bytes(b"saved")
+
+    assert len(calls) == 1
+    assert calls[0][1] == destination
+    assert ".staging-" in calls[0][0].name

tests/test_persistence_pipeline.py

@@ -0,0 +1,231 @@
+"""Integration contracts between checkpoint helpers and the model pipeline."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+import obliteratus.abliterate as abliterate
+from obliteratus.abliterate import AbliterationPipeline
+
+
+pytestmark = pytest.mark.cpu
+
+
+def _pipeline(tmp_path: Path) -> AbliterationPipeline:
+    pipeline = AbliterationPipeline(
+        model_name="test-model",
+        output_dir=str(tmp_path / "checkpoint"),
+        method="basic",
+    )
+    pipeline._on_log = lambda _message: None
+    pipeline._on_stage = lambda _event: None
+    pipeline.handle = MagicMock()
+    return pipeline
+
+
+def test_rebirth_rejects_invalid_metadata_before_gathering_state(tmp_path):
+    pipeline = _pipeline(tmp_path)
+    destination = pipeline.output_dir
+    destination.mkdir()
+    sentinel = destination / "sentinel"
+    sentinel.write_text("old", encoding="utf-8")
+    pipeline._build_metadata = MagicMock(return_value={"invalid": object()})
+    pipeline._gather_state_dict = MagicMock()
+
+    with pytest.raises(TypeError):
+        pipeline._rebirth()
+
+    pipeline._gather_state_dict.assert_not_called()
+    assert sentinel.read_text(encoding="utf-8") == "old"
+
+
+def test_rebirth_insufficient_capacity_preserves_checkpoint_and_runtime_state(
+    tmp_path,
+    monkeypatch,
+):
+    pipeline = _pipeline(tmp_path)
+    destination = pipeline.output_dir
+    destination.mkdir()
+    sentinel = destination / "sentinel"
+    sentinel.write_text("old", encoding="utf-8")
+    pipeline._build_metadata = MagicMock(return_value={"schema": 1})
+    pipeline._gather_state_dict = MagicMock(
+        return_value={"weight": torch.ones(4, dtype=torch.float32)},
+    )
+    pipeline._write_local_checkpoint = MagicMock()
+    pipeline._free_gpu_memory = MagicMock()
+    pipeline._cleanup_offload_dir = MagicMock()
+    monkeypatch.setattr(
+        abliterate.shutil,
+        "disk_usage",
+        lambda _path: SimpleNamespace(free=16),
+    )
+
+    with pytest.raises(OSError, match="Insufficient disk space"):
+        pipeline._rebirth()
+
+    pipeline._write_local_checkpoint.assert_not_called()
+    pipeline._free_gpu_memory.assert_not_called()
+    pipeline._cleanup_offload_dir.assert_not_called()
+    assert sentinel.read_text(encoding="utf-8") == "old"
+    assert list(tmp_path.glob(".checkpoint.staging-*")) == []
+
+
+def test_rebirth_ignores_non_os_disk_probe_failure_and_promotes_checkpoint(
+    tmp_path,
+    monkeypatch,
+):
+    pipeline = _pipeline(tmp_path)
+    pipeline._build_metadata = MagicMock(return_value={"schema": 1})
+    state_dict = {"weight": torch.ones(1)}
+    pipeline._gather_state_dict = MagicMock(return_value=state_dict)
+
+    def write_checkpoint(checkpoint_dir, metadata_json, received_state_dict):
+        assert json.loads(metadata_json) == {"schema": 1}
+        assert received_state_dict is state_dict
+        (checkpoint_dir / "model.bin").write_bytes(b"saved")
+
+    pipeline._write_local_checkpoint = MagicMock(side_effect=write_checkpoint)
+    pipeline._free_gpu_memory = MagicMock()
+    pipeline._cleanup_offload_dir = MagicMock()
+
+    def fail_probe(_path):
+        raise RuntimeError("filesystem probe unavailable")
+
+    monkeypatch.setattr(abliterate.shutil, "disk_usage", fail_probe)
+
+    assert pipeline._rebirth() == pipeline.output_dir
+    assert (pipeline.output_dir / "model.bin").read_bytes() == b"saved"
+    pipeline._free_gpu_memory.assert_called_once_with()
+    pipeline._cleanup_offload_dir.assert_called_once_with()
+
+
+def test_rebirth_pushes_only_after_local_checkpoint_promotion(tmp_path, monkeypatch):
+    import huggingface_hub
+
+    pipeline = _pipeline(tmp_path)
+    pipeline.push_to_hub = "auto"
+    pipeline.hub_token = "test-token"
+    pipeline.hub_community_org = "test-org"
+    pipeline._build_metadata = MagicMock(return_value={"schema": 1})
+    pipeline._gather_state_dict = MagicMock(return_value={"weight": torch.ones(1)})
+    pipeline._free_gpu_memory = MagicMock()
+    pipeline._cleanup_offload_dir = MagicMock()
+
+    def write_checkpoint(checkpoint_dir, _metadata_json, _state_dict):
+        (checkpoint_dir / "model.bin").write_bytes(b"saved")
+
+    pipeline._write_local_checkpoint = MagicMock(side_effect=write_checkpoint)
+    api = MagicMock()
+    api_factory = MagicMock(return_value=api)
+    monkeypatch.setattr(huggingface_hub, "HfApi", api_factory)
+    auto_name = MagicMock(return_value="test-org/test-model-OBLITERATED")
+    monkeypatch.setattr(abliterate, "auto_hub_repo_id", auto_name)
+
+    assert pipeline._rebirth() == pipeline.output_dir
+
+    assert (pipeline.output_dir / "model.bin").read_bytes() == b"saved"
+    api_factory.assert_called_once_with(token="test-token")
+    auto_name.assert_called_once_with(
+        "test-model",
+        api=api,
+        org="test-org",
+    )
+    api.create_repo.assert_called_once_with(
+        "test-org/test-model-OBLITERATED",
+        exist_ok=True,
+    )
+    api.upload_folder.assert_called_once_with(
+        folder_path=str(pipeline.output_dir),
+        repo_id="test-org/test-model-OBLITERATED",
+        commit_message="OBLITERATUS: abliterated test-model (basic)",
+    )
+
+
+def test_rebirth_uses_fallback_token_for_explicit_hub_destination(
+    tmp_path,
+    monkeypatch,
+):
+    import huggingface_hub
+
+    pipeline = _pipeline(tmp_path)
+    pipeline.push_to_hub = "test-org/explicit-repo"
+    pipeline._build_metadata = MagicMock(return_value={"schema": 1})
+    pipeline._gather_state_dict = MagicMock(return_value={"weight": torch.ones(1)})
+    pipeline._write_local_checkpoint = MagicMock(
+        side_effect=lambda path, *_args: (path / "model.bin").write_bytes(b"saved"),
+    )
+    pipeline._free_gpu_memory = MagicMock()
+    pipeline._cleanup_offload_dir = MagicMock()
+    api = MagicMock()
+    api_factory = MagicMock(return_value=api)
+    monkeypatch.setattr(huggingface_hub, "HfApi", api_factory)
+    monkeypatch.setenv("HF_TOKEN", "fallback-token")
+
+    pipeline._rebirth()
+
+    api_factory.assert_called_once_with(token="fallback-token")
+    api.create_repo.assert_called_once_with(
+        "test-org/explicit-repo",
+        exist_ok=True,
+    )
+
+
+def test_write_local_checkpoint_strips_runtime_only_state_and_writes_metadata(
+    tmp_path,
+    monkeypatch,
+):
+    import obliteratus.lora_ablation as lora_ablation
+
+    class Quantizer:
+        def __init__(self):
+            self.models = []
+
+        def remove_quantization_config(self, model):
+            self.models.append(model)
+
+    class Model:
+        def __init__(self):
+            self.hf_quantizer = Quantizer()
+            self._weight_conversions = {"legacy": "conversion"}
+            self.saved = None
+
+        def save_pretrained(self, path, **kwargs):
+            self.saved = (path, kwargs)
+
+    pipeline = _pipeline(tmp_path)
+    model = Model()
+    tokenizer = MagicMock()
+    pipeline.handle = SimpleNamespace(model=model, tokenizer=tokenizer)
+    pipeline._lora_adapters = {"layer": (torch.ones(1), torch.ones(1))}
+    checkpoint_dir = tmp_path / "staging"
+    adapter_path = checkpoint_dir / "lora"
+    save_adapters = MagicMock(return_value=adapter_path)
+    monkeypatch.setattr(lora_ablation, "save_lora_adapters", save_adapters)
+    state_dict = {"weight": torch.ones(1)}
+    metadata_json = '{"schema": 1}'
+    checkpoint_dir.mkdir()
+
+    pipeline._write_local_checkpoint(checkpoint_dir, metadata_json, state_dict)
+
+    assert model.hf_quantizer.models == [model]
+    assert not hasattr(model, "_weight_conversions")
+    assert model.saved == (
+        checkpoint_dir,
+        {
+            "state_dict": state_dict,
+            "max_shard_size": "2GB",
+            "save_original_format": False,
+        },
+    )
+    tokenizer.save_pretrained.assert_called_once_with(checkpoint_dir)
+    save_adapters.assert_called_once_with(pipeline._lora_adapters, checkpoint_dir)
+    assert (checkpoint_dir / "abliteration_metadata.json").read_text(
+        encoding="utf-8",
+    ) == metadata_json

tests/test_quality_policy.py

@@ -18,7 +18,9 @@ def test_mutation_campaign_preloads_native_modules_before_covered_line_discovery
 
     assert "mutate_only_covered_lines = true" in mutmut_config
     assert '"obliteratus/runtime_contracts.py"' in mutmut_config
+    assert '"obliteratus/persistence_contracts.py"' in mutmut_config
     assert '"tests/test_runtime_contracts.py"' in mutmut_config
+    assert '"tests/test_persistence_contracts.py"' in mutmut_config
     assert "import torch, yaml; from mutmut.__main__ import cli; cli()" in workflow