From 1240bf2adbcac0a1c9dee79c25c581ba1240d073 Mon Sep 17 00:00:00 2001 From: Joseph Magly <1159087+jmagly@users.noreply.github.com> Date: Fri, 14 Aug 2026 23:11:38 -0400 Subject: [PATCH] test: enforce transactional persistence contracts --- ci/test-quality-policy.json | 1 + ci/test-risk-map.json | 18 ++- obliteratus/abliterate.py | 86 ++-------- obliteratus/persistence_contracts.py | 115 +++++++++++++ pyproject.toml | 2 + tests/test_checkpoint_atomicity.py | 10 +- tests/test_persistence_contracts.py | 215 +++++++++++++++++++++++++ tests/test_persistence_pipeline.py | 231 +++++++++++++++++++++++++++ tests/test_quality_policy.py | 2 + 9 files changed, 600 insertions(+), 80 deletions(-) create mode 100644 obliteratus/persistence_contracts.py create mode 100644 tests/test_persistence_contracts.py create mode 100644 tests/test_persistence_pipeline.py diff --git a/ci/test-quality-policy.json b/ci/test-quality-policy.json index 6def675..7984660 100644 --- a/ci/test-quality-policy.json +++ b/ci/test-quality-policy.json @@ -11,6 +11,7 @@ }, "critical_cpu_paths": [ "obliteratus/runtime_contracts.py", + "obliteratus/persistence_contracts.py", "obliteratus/device.py", "obliteratus/models/loader.py", "obliteratus/architecture_profiles.py", diff --git a/ci/test-risk-map.json b/ci/test-risk-map.json index 3a1d82e..06f6949 100644 --- a/ci/test-risk-map.json +++ b/ci/test-risk-map.json @@ -35,6 +35,7 @@ "contract_types": ["model-mutation", "model-runtime", "orchestration", "persistence"], "paths": [ "obliteratus/abliterate.py", + "obliteratus/persistence_contracts.py", "obliteratus/auto_obliterate.py", "obliteratus/bayesian_optimizer.py", "obliteratus/informed_pipeline.py", @@ -50,7 +51,9 @@ "tests/test_informed_pipeline.py", "tests/test_offline_integration.py", "tests/test_runner_boundaries.py", - "tests/test_checkpoint_atomicity.py" + "tests/test_checkpoint_atomicity.py", + "tests/test_persistence_contracts.py", + "tests/test_persistence_pipeline.py" ] }, { @@ -315,6 +318,17 @@ "required_tests": ["tests/test_mlx_backend_boundaries.py"], "conditional_gates": ["mlx-runtime"] }, + { + "path": "obliteratus/persistence_contracts.py", + "risk_class": "cpu-contract", + "risk": "checkpoint sizing, metadata serialization, atomic promotion, rollback, and cleanup", + "required_tests": [ + "tests/test_checkpoint_atomicity.py", + "tests/test_persistence_contracts.py", + "tests/test_persistence_pipeline.py" + ], + "conditional_gates": [] + }, { "path": "obliteratus/abliterate.py", "risk_class": "mixed-runtime", @@ -323,6 +337,8 @@ "tests/test_abliterate.py", "tests/test_abliterate_extended.py", "tests/test_checkpoint_atomicity.py", + "tests/test_persistence_contracts.py", + "tests/test_persistence_pipeline.py", "tests/test_offline_integration.py" ], "conditional_gates": ["model-download-runtime"] diff --git a/obliteratus/abliterate.py b/obliteratus/abliterate.py index 1bfce95..060a4d3 100644 --- a/obliteratus/abliterate.py +++ b/obliteratus/abliterate.py @@ -20,16 +20,12 @@ Novel contributions (OBLITERATUS): from __future__ import annotations -import json import logging import math import os import shutil -import tempfile import time -import uuid import warnings -from contextlib import contextmanager from dataclasses import dataclass, field from pathlib import Path from typing import Any, Callable, Iterable @@ -45,6 +41,12 @@ from obliteratus import device as dev # noqa: E402 — must import before CUDA dev.configure_cuda_alloc() from obliteratus.models.loader import ModelHandle, load_model # noqa: E402 +from obliteratus.persistence_contracts import ( # noqa: E402 + atomic_checkpoint_directory as _atomic_checkpoint_directory, + ensure_checkpoint_capacity, + serialize_checkpoint_metadata, + state_dict_size_bytes, +) from obliteratus.strategies.utils import ( # noqa: E402 get_attention_module, get_ffn_module, @@ -61,65 +63,6 @@ logger = logging.getLogger(__name__) _MAX_NORM_RATIO = 1.10 -def _remove_checkpoint_path(path: Path) -> None: - """Remove a staging/backup path without following directory symlinks.""" - if path.is_symlink() or path.is_file(): - path.unlink(missing_ok=True) - elif path.exists(): - shutil.rmtree(path, ignore_errors=True) - - -@contextmanager -def _atomic_checkpoint_directory(destination: Path): - """Yield a sibling staging directory and atomically promote it on success. - - An existing checkpoint is moved to a uniquely named backup immediately - before promotion. If promotion fails, that backup is restored. Exceptions - while writing only remove the staging directory and leave the destination - untouched. - """ - destination = Path(destination) - destination.parent.mkdir(parents=True, exist_ok=True) - staging = Path( - tempfile.mkdtemp( - prefix=f".{destination.name or 'checkpoint'}.staging-", - dir=destination.parent, - ), - ) - backup: Path | None = None - try: - yield staging - if destination.exists() or destination.is_symlink(): - backup = destination.with_name( - f".{destination.name}.backup-{uuid.uuid4().hex}", - ) - os.replace(destination, backup) - try: - os.replace(staging, destination) - except Exception as promotion_error: - if backup is not None and (backup.exists() or backup.is_symlink()): - try: - os.replace(backup, destination) - except Exception as restore_error: - raise RuntimeError( - "Checkpoint promotion and rollback both failed; " - f"recover the previous checkpoint from {backup}", - ) from restore_error - raise promotion_error - if backup is not None: - try: - _remove_checkpoint_path(backup) - except OSError as cleanup_error: - logger.warning( - "Checkpoint promoted, but previous-checkpoint backup %s " - "could not be removed: %s", - backup, - cleanup_error, - ) - except Exception: - _remove_checkpoint_path(staging) - raise - # ── Abliteration method presets ─────────────────────────────────────────── METHODS = { @@ -6737,7 +6680,7 @@ class AbliterationPipeline: def _write_local_checkpoint( self, checkpoint_dir: Path, - metadata: dict, + metadata_json: str, state_dict: dict, ) -> None: """Write every local checkpoint component into an isolated directory.""" @@ -6774,7 +6717,7 @@ class AbliterationPipeline: self.handle.tokenizer.save_pretrained(checkpoint_dir) (checkpoint_dir / "abliteration_metadata.json").write_text( - json.dumps(metadata, indent=2), + metadata_json, encoding="utf-8", ) @@ -6791,6 +6734,7 @@ class AbliterationPipeline: t0 = time.time() metadata = self._build_metadata() + metadata_json = serialize_checkpoint_metadata(metadata) # 1. Gather state dict FIRST (while offload dir still exists, so we # can read any disk-offloaded weights). @@ -6798,7 +6742,7 @@ class AbliterationPipeline: state_dict = self._gather_state_dict() # 2. Estimate serialized size from the gathered state dict. - param_bytes = sum(v.numel() * v.element_size() for v in state_dict.values()) + param_bytes = state_dict_size_bytes(state_dict) self.log(f"State dict: {len(state_dict)} tensors, {param_bytes / 1e9:.1f} GB") # 3. Save every component to a sibling staging directory, then promote @@ -6808,13 +6752,7 @@ class AbliterationPipeline: with _atomic_checkpoint_directory(self.output_dir) as checkpoint_dir: try: disk = shutil.disk_usage(checkpoint_dir) - needed = int(param_bytes * 1.1) - if disk.free < needed: - raise OSError( - f"Insufficient disk space: " - f"{disk.free / 1e9:.1f} GB free, need ~{param_bytes / 1e9:.1f} GB. " - f"Try a different --output-dir on a larger filesystem." - ) + ensure_checkpoint_capacity(disk.free, param_bytes) self.log( f"Disk space: {disk.free / 1e9:.1f} GB free, " f"need ~{param_bytes / 1e9:.1f} GB", @@ -6823,7 +6761,7 @@ class AbliterationPipeline: raise except Exception: pass # Non-critical — don't block save on stat failure - self._write_local_checkpoint(checkpoint_dir, metadata, state_dict) + self._write_local_checkpoint(checkpoint_dir, metadata_json, state_dict) # Free the state dict and temporary offload only after promotion. del state_dict diff --git a/obliteratus/persistence_contracts.py b/obliteratus/persistence_contracts.py new file mode 100644 index 0000000..9f8adc3 --- /dev/null +++ b/obliteratus/persistence_contracts.py @@ -0,0 +1,115 @@ +"""Deterministic and transactional contracts for local model checkpoints.""" + +from __future__ import annotations + +import json +import logging +import os +import shutil +import tempfile +import uuid +from collections.abc import Mapping +from contextlib import contextmanager +from pathlib import Path +from typing import Any, Iterator, Protocol + + +logger = logging.getLogger(__name__) + + +class SizedTensor(Protocol): + """Structural subset used to estimate a serialized state dictionary.""" + + def numel(self) -> int: ... + + def element_size(self) -> int: ... + + +def state_dict_size_bytes(state_dict: Mapping[str, SizedTensor]) -> int: + """Return the exact unsharded tensor payload size for a state dictionary.""" + return sum(value.numel() * value.element_size() for value in state_dict.values()) + + +def required_checkpoint_bytes(payload_bytes: int) -> int: + """Return the checkpoint capacity requirement with ten percent headroom.""" + if payload_bytes < 0: + raise ValueError("checkpoint payload size cannot be negative") + return (payload_bytes * 11 + 9) // 10 + + +def ensure_checkpoint_capacity( + free_bytes: int, + payload_bytes: int, +) -> None: + """Reject a checkpoint write that cannot satisfy the headroom policy.""" + if free_bytes < required_checkpoint_bytes(payload_bytes): + raise OSError( + f"Insufficient disk space: {free_bytes / 1e9:.1f} GB free, " + f"need ~{payload_bytes / 1e9:.1f} GB. " + f"Try a different --output-dir on a larger filesystem.", + ) + + +def serialize_checkpoint_metadata(metadata: Mapping[str, Any]) -> str: + """Serialize checkpoint metadata before any model artifact is written.""" + return json.dumps(metadata, indent=2, sort_keys=True, allow_nan=False) + + +def _remove_checkpoint_path(path: Path) -> None: + """Remove a staging/backup path without following directory symlinks.""" + if path.is_symlink() or path.is_file(): + path.unlink(missing_ok=True) + elif path.exists(): + shutil.rmtree(path, ignore_errors=True) + + +@contextmanager +def atomic_checkpoint_directory(destination: Path) -> Iterator[Path]: + """Yield a sibling staging directory and atomically promote it on success. + + An existing checkpoint is moved to a uniquely named backup immediately + before promotion. If promotion fails, that backup is restored. Exceptions + while writing only remove the staging directory and leave the destination + untouched. + """ + destination = Path(destination) + destination.parent.mkdir(parents=True, exist_ok=True) + staging = Path( + tempfile.mkdtemp( + prefix=f".{destination.name or 'checkpoint'}.staging-", + dir=destination.parent, + ), + ) + backup: Path | None = None + try: + yield staging + if destination.exists() or destination.is_symlink(): + backup = destination.with_name( + f".{destination.name}.backup-{uuid.uuid4().hex}", + ) + os.replace(destination, backup) + try: + os.replace(staging, destination) + except Exception as promotion_error: + if backup is not None and (backup.exists() or backup.is_symlink()): + try: + os.replace(backup, destination) + except Exception as restore_error: + raise RuntimeError( + "Checkpoint promotion and rollback both failed; " + f"recover the previous checkpoint from {backup}", + ) from restore_error + raise promotion_error + if backup is not None: + try: + _remove_checkpoint_path(backup) + except OSError as cleanup_error: + logger.warning( + "Checkpoint promoted, but previous-checkpoint backup %s " + "could not be removed: %s", + backup, + cleanup_error, + ) + except Exception: + _remove_checkpoint_path(staging) + raise diff --git a/pyproject.toml b/pyproject.toml index 3c78897..073ef90 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -119,6 +119,7 @@ only_mutate = [ "obliteratus/config.py", "obliteratus/analysis/numerical_contracts.py", "obliteratus/runtime_contracts.py", + "obliteratus/persistence_contracts.py", "scripts/check_coverage_thresholds.py", ] pytest_add_cli_args = ["--no-cov", "-q"] @@ -127,6 +128,7 @@ pytest_add_cli_args_test_selection = [ "tests/test_config_properties.py", "tests/test_coverage_thresholds.py", "tests/test_numerical_contracts.py", + "tests/test_persistence_contracts.py", "tests/test_runtime_contracts.py", ] mutate_only_covered_lines = true diff --git a/tests/test_checkpoint_atomicity.py b/tests/test_checkpoint_atomicity.py index 9868e95..93bb79f 100644 --- a/tests/test_checkpoint_atomicity.py +++ b/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") diff --git a/tests/test_persistence_contracts.py b/tests/test_persistence_contracts.py new file mode 100644 index 0000000..0725146 --- /dev/null +++ b/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 diff --git a/tests/test_persistence_pipeline.py b/tests/test_persistence_pipeline.py new file mode 100644 index 0000000..cdc29f3 --- /dev/null +++ b/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 diff --git a/tests/test_quality_policy.py b/tests/test_quality_policy.py index 9c44557..588ee46 100644 --- a/tests/test_quality_policy.py +++ b/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