mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-09-21 17:00:50 +02:00
754 lines
28 KiB
Python
754 lines
28 KiB
Python
"""Offline, structure-only checkpoint inspection tests."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import socket
|
|
import subprocess
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
import torch
|
|
from jsonschema import Draft202012Validator, FormatChecker
|
|
from safetensors.torch import save_file
|
|
|
|
import obliteratus.checkpoint_inspection as inspection_module
|
|
from obliteratus.checkpoint_capabilities import (
|
|
AdapterCapability,
|
|
AdapterRegistry,
|
|
ExactDependency,
|
|
)
|
|
from obliteratus.checkpoint_errors import CheckpointContractError
|
|
from obliteratus.checkpoint_inspection import InspectionLimits, inspect_checkpoint
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
DESCRIPTOR_SCHEMA = json.loads(
|
|
(ROOT / "docs/checkpoints/schemas/checkpoint-descriptor-v1.schema.json").read_text()
|
|
)
|
|
|
|
|
|
def _save(path: Path, **tensors: torch.Tensor) -> None:
|
|
save_file(dict(sorted(tensors.items())), path)
|
|
|
|
|
|
def _assert_contract(report) -> dict:
|
|
descriptor = report.to_dict()
|
|
Draft202012Validator(
|
|
DESCRIPTOR_SCHEMA,
|
|
format_checker=FormatChecker(),
|
|
).validate(descriptor)
|
|
assert json.loads(report.to_json()) == descriptor
|
|
return descriptor
|
|
|
|
|
|
def test_direct_hf_safetensors_is_inventory_backed_and_canonical_ready(tmp_path):
|
|
_save(tmp_path / "model.safetensors", weight=torch.arange(6).reshape(2, 3))
|
|
(tmp_path / "config.json").write_text('{"model_type":"tiny"}\n', encoding="utf-8")
|
|
|
|
first = inspect_checkpoint(tmp_path)
|
|
second = inspect_checkpoint(tmp_path)
|
|
descriptor = _assert_contract(first)
|
|
|
|
assert first.primary_format == "hf_safetensors"
|
|
assert first.support_decision == "canonical_hf_ready"
|
|
assert first.descriptor_id == second.descriptor_id
|
|
assert descriptor["classification_confidence"] == "verified"
|
|
assert descriptor["safety"] == {
|
|
"inspection_level": "safe_structure",
|
|
"trust_required": False,
|
|
"inventory_revalidated": True,
|
|
"unsafe_serialization_findings": [],
|
|
"violations": [],
|
|
}
|
|
assert descriptor["state"] == {
|
|
"observed_scopes": ["model_weights"],
|
|
"classification": "weights_only",
|
|
}
|
|
assert descriptor["resource_estimate"]["tensor_count"] == 1
|
|
assert descriptor["resource_estimate"]["logical_bytes"] == 48
|
|
assert not descriptor["blockers"]
|
|
|
|
|
|
def test_indexed_safetensors_validates_safe_weight_map_and_shards(tmp_path):
|
|
_save(tmp_path / "model-00001-of-00002.safetensors", a=torch.tensor([1.0]))
|
|
_save(tmp_path / "model-00002-of-00002.safetensors", b=torch.tensor([2.0, 3.0]))
|
|
(tmp_path / "model.safetensors.index.json").write_text(
|
|
json.dumps(
|
|
{
|
|
"metadata": {"total_size": 12},
|
|
"weight_map": {
|
|
"a": "model-00001-of-00002.safetensors",
|
|
"b": "model-00002-of-00002.safetensors",
|
|
},
|
|
}
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
descriptor = _assert_contract(inspect_checkpoint(tmp_path))
|
|
|
|
assert descriptor["primary_format"] == "hf_safetensors"
|
|
assert descriptor["resource_estimate"]["tensor_count"] == 2
|
|
assert descriptor["resource_estimate"]["shard_count"] == 2
|
|
assert descriptor["support_decision"] == "canonical_hf_ready"
|
|
|
|
|
|
def test_direct_and_indexed_hf_signatures_are_an_ambiguous_collision(tmp_path):
|
|
_save(tmp_path / "model.safetensors", direct=torch.ones(1))
|
|
_save(tmp_path / "model-00001-of-00001.safetensors", indexed=torch.ones(1))
|
|
(tmp_path / "model.safetensors.index.json").write_text(
|
|
json.dumps(
|
|
{
|
|
"weight_map": {
|
|
"indexed": "model-00001-of-00001.safetensors",
|
|
}
|
|
}
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
descriptor = _assert_contract(inspect_checkpoint(tmp_path))
|
|
|
|
assert descriptor["primary_format"] == "ambiguous"
|
|
assert descriptor["support_decision"] == "blocked"
|
|
assert "hf_layout_collision" in descriptor["safety"]["violations"]
|
|
|
|
|
|
@pytest.mark.parametrize("mode", ["missing", "extra"])
|
|
def test_indexed_safetensors_missing_or_extra_shards_are_blocked(tmp_path, mode):
|
|
_save(tmp_path / "model-00001-of-00001.safetensors", a=torch.ones(1))
|
|
referenced = (
|
|
"missing.safetensors"
|
|
if mode == "missing"
|
|
else "model-00001-of-00001.safetensors"
|
|
)
|
|
if mode == "extra":
|
|
_save(tmp_path / "model-00002-of-00002.safetensors", extra=torch.ones(1))
|
|
(tmp_path / "model.safetensors.index.json").write_text(
|
|
json.dumps({"weight_map": {"a": referenced}}),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
descriptor = _assert_contract(inspect_checkpoint(tmp_path))
|
|
|
|
assert descriptor["support_decision"] == "blocked"
|
|
assert "hf_weight_map_shard_mismatch" in descriptor["safety"]["violations"]
|
|
assert descriptor["blockers"][0]["code"] == "DCI_VALIDATION_FAILED"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("header", "payload", "violation"),
|
|
[
|
|
({"weight": {"dtype": "F32", "shape": [2], "data_offsets": [0, 4]}}, b"\0" * 4, "safetensors_range_invalid"),
|
|
({"weight": {"dtype": "UNKNOWN", "shape": [1], "data_offsets": [0, 4]}}, b"\0" * 4, "safetensors_header_invalid"),
|
|
({"weight": {"dtype": "F32", "shape": [1], "data_offsets": [0, 8]}}, b"\0" * 4, "safetensors_range_invalid"),
|
|
({"weight": {"dtype": "F32", "shape": [1], "data_offsets": [4, 8]}}, b"\0" * 8, "safetensors_range_gap"),
|
|
],
|
|
)
|
|
def test_safetensors_dtype_shape_and_range_corruption_is_blocked(
|
|
tmp_path,
|
|
header,
|
|
payload,
|
|
violation,
|
|
):
|
|
encoded = json.dumps(header, separators=(",", ":")).encode()
|
|
(tmp_path / "model.safetensors").write_bytes(
|
|
len(encoded).to_bytes(8, "little") + encoded + payload
|
|
)
|
|
|
|
descriptor = _assert_contract(inspect_checkpoint(tmp_path))
|
|
|
|
assert descriptor["support_decision"] == "blocked"
|
|
assert violation in descriptor["safety"]["violations"]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("files", "expected_format"),
|
|
[
|
|
({"pytorch_model.bin": b"pickle"}, "hf_pytorch_pickle"),
|
|
({".metadata": b"opaque-dcp-metadata", "__0_0.distcp": b"payload"}, "pytorch_dcp"),
|
|
(
|
|
{
|
|
".metadata": b"opaque-dcp-metadata",
|
|
"fsdp_metadata.json": b'{"state_dict_type":"SHARDED_STATE_DICT"}',
|
|
},
|
|
"fsdp_state_dict",
|
|
),
|
|
(
|
|
{
|
|
"metadata.json": b'{"sharded_backend":"torch_dist","version":"1.0"}',
|
|
"common.pt": b"pickle",
|
|
},
|
|
"megatron_torch_dist",
|
|
),
|
|
({"zero_pp_rank_0_mp_rank_00_optim_states.pt": b"pickle"}, "deepspeed_zero"),
|
|
(
|
|
{
|
|
"universal_checkpoint_info.json": b'{"type":"universal"}',
|
|
"zero_pp_rank_0_mp_rank_00_model_states.pt": b"pickle",
|
|
},
|
|
"deepspeed_universal",
|
|
),
|
|
],
|
|
)
|
|
def test_vendor_and_pickle_layouts_are_classified_without_payload_access(
|
|
tmp_path,
|
|
files,
|
|
expected_format,
|
|
):
|
|
for name, payload in files.items():
|
|
(tmp_path / name).write_bytes(payload)
|
|
|
|
descriptor = _assert_contract(inspect_checkpoint(tmp_path))
|
|
|
|
assert descriptor["primary_format"] == expected_format
|
|
assert descriptor["support_decision"] == "trusted_inspection_required"
|
|
assert descriptor["safety"]["trust_required"] is True
|
|
assert descriptor["adapter_resolution"]["status"] == "missing"
|
|
assert {item["code"] for item in descriptor["blockers"]} == {
|
|
"DCI_TRUST_POLICY_REQUIRED"
|
|
}
|
|
|
|
|
|
def test_exact_registered_capability_reports_missing_extra_and_version(tmp_path):
|
|
(tmp_path / ".metadata").write_bytes(b"opaque-dcp-metadata")
|
|
(tmp_path / "__0_0.distcp").write_bytes(b"payload-never-read-as-a-tensor")
|
|
registry = AdapterRegistry(
|
|
(
|
|
AdapterCapability(
|
|
adapter_id="example-dcp",
|
|
adapter_version="0.1.0",
|
|
producer="example-producer",
|
|
producer_version="1.2.3",
|
|
formats=("pytorch_dcp",),
|
|
required_extras=("checkpoint-example",),
|
|
required_dependencies=(
|
|
ExactDependency("obliteratus-test-package-that-does-not-exist", "1.2.3"),
|
|
),
|
|
),
|
|
)
|
|
)
|
|
|
|
descriptor = _assert_contract(
|
|
inspect_checkpoint(tmp_path, adapter_registry=registry)
|
|
)
|
|
|
|
assert descriptor["primary_format"] == "pytorch_dcp"
|
|
assert descriptor["support_decision"] == "trusted_inspection_required"
|
|
assert descriptor["adapter_resolution"]["status"] == "missing"
|
|
assert descriptor["adapter_resolution"]["adapter_id"] == "example-dcp"
|
|
assert "install_extra=obliteratus[checkpoint-example]==0.1.3" in descriptor[
|
|
"adapter_resolution"
|
|
]["reason"]
|
|
assert (
|
|
"required_versions=obliteratus-test-package-that-does-not-exist==1.2.3"
|
|
in descriptor["adapter_resolution"]["reason"]
|
|
)
|
|
assert {item["code"] for item in descriptor["blockers"]} == {
|
|
"DCI_TRUST_POLICY_REQUIRED",
|
|
"DCI_TRUST_RUNTIME_UNAVAILABLE",
|
|
}
|
|
|
|
|
|
def test_ambiguous_registered_capabilities_fail_closed(tmp_path):
|
|
(tmp_path / ".metadata").write_bytes(b"opaque-dcp-metadata")
|
|
dependency = ExactDependency("example-producer", "1.2.3")
|
|
registry = AdapterRegistry(
|
|
tuple(
|
|
AdapterCapability(
|
|
adapter_id=f"example-{index}",
|
|
adapter_version="0.1.0",
|
|
producer="example-producer",
|
|
producer_version="1.2.3",
|
|
formats=("pytorch_dcp",),
|
|
required_extras=("checkpoint-example",),
|
|
required_dependencies=(dependency,),
|
|
)
|
|
for index in range(2)
|
|
)
|
|
)
|
|
|
|
descriptor = _assert_contract(
|
|
inspect_checkpoint(tmp_path, adapter_registry=registry)
|
|
)
|
|
|
|
assert descriptor["adapter_resolution"]["status"] == "ambiguous"
|
|
assert descriptor["support_decision"] == "blocked"
|
|
assert [item["code"] for item in descriptor["blockers"]].count(
|
|
"DCI_UNSUPPORTED_FORMAT_OR_VERSION"
|
|
) == 1
|
|
|
|
|
|
def test_peft_layout_is_safe_safetensors_but_keeps_adapter_identity(tmp_path):
|
|
_save(tmp_path / "adapter_model.safetensors", lora_A=torch.ones(1, 2))
|
|
(tmp_path / "adapter_config.json").write_text(
|
|
'{"base_model_name_or_path":"local/base","peft_type":"LORA"}\n',
|
|
encoding="utf-8",
|
|
)
|
|
|
|
descriptor = _assert_contract(inspect_checkpoint(tmp_path))
|
|
|
|
assert descriptor["primary_format"] == "peft_safetensors"
|
|
assert descriptor["components"][0]["kind"] == "peft_adapter"
|
|
assert descriptor["components"][0]["format"] == "peft_safetensors"
|
|
assert descriptor["state"]["observed_scopes"] == ["adapter_weights"]
|
|
assert descriptor["support_decision"] == "canonical_hf_ready"
|
|
|
|
|
|
def test_mixed_model_and_adapter_components_are_preserved_and_blocked(tmp_path):
|
|
_save(tmp_path / "model.safetensors", weight=torch.ones(2, 2))
|
|
_save(tmp_path / "adapter_model.safetensors", lora_A=torch.ones(1, 2))
|
|
(tmp_path / "adapter_config.json").write_text('{"peft_type":"LORA"}\n')
|
|
|
|
descriptor = _assert_contract(inspect_checkpoint(tmp_path))
|
|
|
|
assert descriptor["primary_format"] == "ambiguous"
|
|
assert [(item["kind"], item["format"]) for item in descriptor["components"]] == [
|
|
("model", "hf_safetensors"),
|
|
("peft_adapter", "peft_safetensors"),
|
|
]
|
|
assert descriptor["state"]["observed_scopes"] == ["adapter_weights", "model_weights"]
|
|
assert descriptor["support_decision"] == "blocked"
|
|
assert descriptor["blockers"][0]["code"] == "DCI_UNSUPPORTED_FORMAT_OR_VERSION"
|
|
|
|
|
|
def test_unknown_layout_is_a_stable_blocked_descriptor(tmp_path):
|
|
(tmp_path / "notes.txt").write_text("not a checkpoint\n")
|
|
|
|
descriptor = _assert_contract(inspect_checkpoint(tmp_path))
|
|
|
|
assert descriptor["primary_format"] == "unknown"
|
|
assert descriptor["classification_confidence"] == "unknown"
|
|
assert descriptor["components"][0]["kind"] == "unknown"
|
|
assert descriptor["support_decision"] == "blocked"
|
|
|
|
|
|
def test_legacy_pickle_payload_is_never_executed(tmp_path):
|
|
marker = tmp_path / "payload-executed"
|
|
command = f"touch {marker}".encode("utf-8")
|
|
malicious = b"cos\nsystem\n(S'" + command + b"'\ntR."
|
|
(tmp_path / "pytorch_model.bin").write_bytes(malicious)
|
|
|
|
descriptor = _assert_contract(inspect_checkpoint(tmp_path))
|
|
|
|
assert descriptor["primary_format"] == "hf_pytorch_pickle"
|
|
assert not marker.exists()
|
|
|
|
|
|
def test_remote_code_declaration_is_inert_and_not_imported(tmp_path, monkeypatch):
|
|
_save(tmp_path / "model.safetensors", weight=torch.ones(1))
|
|
(tmp_path / "config.json").write_text(
|
|
json.dumps({"auto_map": {"AutoModel": "must_not_import.Model"}}),
|
|
encoding="utf-8",
|
|
)
|
|
imported: list[str] = []
|
|
original_import = __import__
|
|
|
|
def guarded_import(name, *args, **kwargs):
|
|
if name.startswith("must_not_import"):
|
|
imported.append(name)
|
|
raise AssertionError("remote code import attempted")
|
|
return original_import(name, *args, **kwargs)
|
|
|
|
monkeypatch.setattr("builtins.__import__", guarded_import)
|
|
|
|
descriptor = _assert_contract(inspect_checkpoint(tmp_path))
|
|
|
|
assert descriptor["primary_format"] == "hf_safetensors"
|
|
assert imported == []
|
|
|
|
|
|
def test_default_inspection_invokes_no_reader_network_process_group_or_plugin(
|
|
tmp_path,
|
|
monkeypatch,
|
|
):
|
|
_save(tmp_path / "model.safetensors", weight=torch.ones(1))
|
|
before = {
|
|
path.name: path.read_bytes()
|
|
for path in tmp_path.iterdir()
|
|
if path.is_file()
|
|
}
|
|
monkeypatch.setattr(torch, "load", lambda *_a, **_k: pytest.fail("torch.load called"))
|
|
monkeypatch.setattr(
|
|
torch.distributed,
|
|
"init_process_group",
|
|
lambda *_a, **_k: pytest.fail("process group initialized"),
|
|
)
|
|
monkeypatch.setattr(
|
|
socket,
|
|
"create_connection",
|
|
lambda *_a, **_k: pytest.fail("network opened"),
|
|
)
|
|
monkeypatch.setattr(
|
|
subprocess,
|
|
"run",
|
|
lambda *_a, **_k: pytest.fail("subprocess started"),
|
|
)
|
|
monkeypatch.setattr(
|
|
"importlib.metadata.entry_points",
|
|
lambda *_a, **_k: pytest.fail("plugin discovery attempted"),
|
|
)
|
|
monkeypatch.setattr(
|
|
"importlib.metadata.version",
|
|
lambda *_a, **_k: pytest.fail("package metadata queried without a capability"),
|
|
)
|
|
|
|
descriptor = _assert_contract(inspect_checkpoint(tmp_path))
|
|
|
|
assert descriptor["support_decision"] == "canonical_hf_ready"
|
|
assert {
|
|
path.name: path.read_bytes()
|
|
for path in tmp_path.iterdir()
|
|
if path.is_file()
|
|
} == before
|
|
|
|
|
|
def test_duplicate_recognized_basenames_fail_closed(tmp_path):
|
|
for directory in (tmp_path / "a", tmp_path / "b"):
|
|
directory.mkdir()
|
|
_save(directory / "model.safetensors", weight=torch.ones(1))
|
|
|
|
with pytest.raises(CheckpointContractError) as caught:
|
|
inspect_checkpoint(tmp_path)
|
|
|
|
assert caught.value.code == "DCI_VALIDATION_FAILED"
|
|
assert caught.value.detail == "duplicate_basename"
|
|
|
|
|
|
def test_symlink_and_special_file_sources_fail_before_classification(tmp_path):
|
|
target = tmp_path / "target.safetensors"
|
|
target.write_bytes(b"not relevant")
|
|
(tmp_path / "model.safetensors").symlink_to(target.name)
|
|
|
|
with pytest.raises(CheckpointContractError) as symlink_error:
|
|
inspect_checkpoint(tmp_path)
|
|
assert symlink_error.value.code == "DCI_SOURCE_BOUNDARY_VIOLATION"
|
|
assert symlink_error.value.detail == "source_symlink"
|
|
|
|
(tmp_path / "model.safetensors").unlink()
|
|
if hasattr(os, "mkfifo"):
|
|
os.mkfifo(tmp_path / "special")
|
|
with pytest.raises(CheckpointContractError) as special_error:
|
|
inspect_checkpoint(tmp_path)
|
|
assert special_error.value.detail == "source_special_file"
|
|
|
|
|
|
def test_resource_limits_apply_before_header_or_json_allocation(tmp_path):
|
|
(tmp_path / "a").write_bytes(b"a")
|
|
(tmp_path / "b").write_bytes(b"b")
|
|
with pytest.raises(CheckpointContractError) as file_error:
|
|
inspect_checkpoint(tmp_path, limits=InspectionLimits(max_files=1))
|
|
assert file_error.value.code == "DCI_RESOURCE_LIMIT"
|
|
assert file_error.value.detail == "max_files"
|
|
|
|
for path in tmp_path.iterdir():
|
|
path.unlink()
|
|
(tmp_path / "model.safetensors").write_bytes((1024).to_bytes(8, "little") + b"{}")
|
|
with pytest.raises(CheckpointContractError) as header_error:
|
|
inspect_checkpoint(
|
|
tmp_path,
|
|
limits=InspectionLimits(max_safetensors_header_bytes=64),
|
|
)
|
|
assert header_error.value.detail == "max_safetensors_header_bytes"
|
|
|
|
|
|
def test_inventory_race_is_detected_and_never_returned_as_success(tmp_path, monkeypatch):
|
|
weights = tmp_path / "model.safetensors"
|
|
_save(weights, weight=torch.ones(1))
|
|
original = inspection_module._hash_regular_file
|
|
changed = False
|
|
|
|
def race(path, expected, limits):
|
|
nonlocal changed
|
|
result = original(path, expected, limits)
|
|
if not changed:
|
|
changed = True
|
|
path.write_bytes(path.read_bytes() + b"changed")
|
|
return result
|
|
|
|
monkeypatch.setattr(inspection_module, "_hash_regular_file", race)
|
|
|
|
with pytest.raises(CheckpointContractError) as caught:
|
|
inspect_checkpoint(tmp_path)
|
|
|
|
assert caught.value.code == "DCI_SOURCE_CHANGED"
|
|
assert caught.value.detail == "source_changed"
|
|
|
|
|
|
def _raw_safetensors(path: Path, header: object, payload: bytes = b"") -> None:
|
|
encoded = json.dumps(header, separators=(",", ":")).encode("utf-8")
|
|
path.write_bytes(len(encoded).to_bytes(8, "little") + encoded + payload)
|
|
|
|
|
|
def test_inspection_limits_and_source_root_types_fail_closed(tmp_path):
|
|
with pytest.raises(ValueError, match="positive integer"):
|
|
InspectionLimits(max_files=0)
|
|
|
|
with pytest.raises(CheckpointContractError) as missing:
|
|
inspect_checkpoint(tmp_path / "missing")
|
|
assert missing.value.detail == "source_missing"
|
|
|
|
target = tmp_path / "target"
|
|
target.mkdir()
|
|
root_link = tmp_path / "root-link"
|
|
root_link.symlink_to(target, target_is_directory=True)
|
|
with pytest.raises(CheckpointContractError) as symlink:
|
|
inspect_checkpoint(root_link)
|
|
assert symlink.value.detail == "source_symlink"
|
|
|
|
if hasattr(os, "mkfifo"):
|
|
fifo = tmp_path / "root-fifo"
|
|
os.mkfifo(fifo)
|
|
with pytest.raises(CheckpointContractError) as special:
|
|
inspect_checkpoint(fifo)
|
|
assert special.value.detail == "source_special_file"
|
|
|
|
|
|
def test_ancestor_symlink_and_open_race_fail_with_stable_source_errors(
|
|
tmp_path,
|
|
monkeypatch,
|
|
):
|
|
actual = tmp_path / "actual"
|
|
source = actual / "checkpoint"
|
|
source.mkdir(parents=True)
|
|
_save(source / "model.safetensors", weight=torch.ones(1))
|
|
alias = tmp_path / "alias"
|
|
alias.symlink_to(actual, target_is_directory=True)
|
|
|
|
with pytest.raises(CheckpointContractError) as symlink:
|
|
inspect_checkpoint(alias / "checkpoint")
|
|
assert symlink.value.code == "DCI_SOURCE_BOUNDARY_VIOLATION"
|
|
assert symlink.value.detail == "source_symlink"
|
|
|
|
weights = source / "model.safetensors"
|
|
real_open = inspection_module.os.open
|
|
|
|
def fail_observed_open(path, flags):
|
|
if Path(path) == weights:
|
|
raise OSError("injected source replacement")
|
|
return real_open(path, flags)
|
|
|
|
monkeypatch.setattr(inspection_module.os, "open", fail_observed_open)
|
|
with pytest.raises(CheckpointContractError) as changed:
|
|
inspect_checkpoint(source)
|
|
assert changed.value.code == "DCI_SOURCE_CHANGED"
|
|
assert changed.value.detail == "source_changed"
|
|
|
|
|
|
def test_single_file_and_directory_byte_limits_are_enforced(tmp_path):
|
|
weights = tmp_path / "single.safetensors"
|
|
_save(weights, weight=torch.ones(1))
|
|
descriptor = _assert_contract(inspect_checkpoint(weights))
|
|
assert descriptor["source_inventory"]["files"][0]["relative_path"] == weights.name
|
|
|
|
nested = tmp_path / "nested"
|
|
nested.mkdir()
|
|
(nested / "file").write_text("x", encoding="utf-8")
|
|
with pytest.raises(CheckpointContractError) as directories:
|
|
inspect_checkpoint(tmp_path, limits=InspectionLimits(max_directories=1))
|
|
assert directories.value.detail == "max_directories"
|
|
|
|
with pytest.raises(CheckpointContractError) as total:
|
|
inspect_checkpoint(weights, limits=InspectionLimits(max_total_bytes=1))
|
|
assert total.value.detail == "max_total_bytes"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("header", "payload", "violation"),
|
|
[
|
|
([], b"", "safetensors_header_invalid"),
|
|
({"": {"dtype": "F32", "shape": [1], "data_offsets": [0, 4]}}, b"\0" * 4, "safetensors_header_invalid"),
|
|
({"x": []}, b"", "safetensors_header_invalid"),
|
|
({"x": {"dtype": "F32", "shape": [1], "data_offsets": [0]}}, b"\0" * 4, "safetensors_header_invalid"),
|
|
({"x": {"dtype": 1, "shape": [1], "data_offsets": [0, 4]}}, b"\0" * 4, "safetensors_header_invalid"),
|
|
({"x": {"dtype": "F32", "shape": [-1], "data_offsets": [0, 0]}}, b"", "safetensors_header_invalid"),
|
|
({"x": {"dtype": "F32", "shape": [1 << 62, 4], "data_offsets": [0, 0]}}, b"", "integer_overflow"),
|
|
({"x": {"dtype": "C128", "shape": [1 << 62], "data_offsets": [0, 0]}}, b"", "integer_overflow"),
|
|
(
|
|
{
|
|
"a": {"dtype": "F32", "shape": [1], "data_offsets": [0, 4]},
|
|
"b": {"dtype": "F32", "shape": [1], "data_offsets": [2, 6]},
|
|
},
|
|
b"\0" * 6,
|
|
"safetensors_range_overlap",
|
|
),
|
|
],
|
|
)
|
|
def test_additional_safetensors_header_corruptions_are_blocked(
|
|
tmp_path,
|
|
header,
|
|
payload,
|
|
violation,
|
|
):
|
|
_raw_safetensors(tmp_path / "model.safetensors", header, payload)
|
|
|
|
descriptor = _assert_contract(inspect_checkpoint(tmp_path))
|
|
|
|
assert violation in descriptor["safety"]["violations"]
|
|
assert descriptor["support_decision"] == "blocked"
|
|
|
|
|
|
def test_truncated_invalid_and_overlarge_safetensors_headers(tmp_path):
|
|
weights = tmp_path / "model.safetensors"
|
|
weights.write_bytes(b"tiny")
|
|
assert "safetensors_truncated" in _assert_contract(inspect_checkpoint(tmp_path))["safety"][
|
|
"violations"
|
|
]
|
|
|
|
weights.write_bytes((20).to_bytes(8, "little") + b"{}")
|
|
assert "safetensors_truncated" in _assert_contract(inspect_checkpoint(tmp_path))["safety"][
|
|
"violations"
|
|
]
|
|
|
|
weights.write_bytes((1).to_bytes(8, "little") + b"{")
|
|
assert "safetensors_header_invalid" in _assert_contract(inspect_checkpoint(tmp_path))[
|
|
"safety"
|
|
]["violations"]
|
|
|
|
_raw_safetensors(
|
|
weights,
|
|
{
|
|
"a": {"dtype": "F32", "shape": [0], "data_offsets": [0, 0]},
|
|
"b": {"dtype": "F32", "shape": [0], "data_offsets": [0, 0]},
|
|
},
|
|
)
|
|
with pytest.raises(CheckpointContractError) as tensor_limit:
|
|
inspect_checkpoint(tmp_path, limits=InspectionLimits(max_tensors=1))
|
|
assert tensor_limit.value.detail == "max_tensors"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("weight_map", "violation"),
|
|
[
|
|
({}, "hf_weight_map_invalid"),
|
|
({"": "model-00001-of-00001.safetensors"}, "hf_weight_map_invalid"),
|
|
({"other": "model-00001-of-00001.safetensors"}, "hf_weight_map_tensor_mismatch"),
|
|
],
|
|
)
|
|
def test_hf_index_shape_and_tensor_membership_are_validated(tmp_path, weight_map, violation):
|
|
_save(tmp_path / "model-00001-of-00001.safetensors", weight=torch.ones(1))
|
|
(tmp_path / "model.safetensors.index.json").write_text(
|
|
json.dumps({"weight_map": weight_map}),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
descriptor = _assert_contract(inspect_checkpoint(tmp_path))
|
|
|
|
assert violation in descriptor["safety"]["violations"]
|
|
|
|
|
|
def test_hf_index_rejects_unmapped_tensor_in_a_referenced_shard(tmp_path):
|
|
_save(
|
|
tmp_path / "model-00001-of-00001.safetensors",
|
|
declared=torch.ones(1),
|
|
undeclared=torch.ones(1),
|
|
)
|
|
(tmp_path / "model.safetensors.index.json").write_text(
|
|
json.dumps(
|
|
{
|
|
"weight_map": {
|
|
"declared": "model-00001-of-00001.safetensors",
|
|
}
|
|
}
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
descriptor = _assert_contract(inspect_checkpoint(tmp_path))
|
|
|
|
assert "hf_weight_map_tensor_mismatch" in descriptor["safety"]["violations"]
|
|
assert descriptor["support_decision"] == "blocked"
|
|
|
|
|
|
def test_duplicate_json_keys_and_aggregate_tensor_limit_fail_closed(tmp_path):
|
|
(tmp_path / ".metadata").write_bytes(b"opaque")
|
|
(tmp_path / "fsdp_metadata.json").write_text(
|
|
'{"state_dict_type":"SHARDED_STATE_DICT","state_dict_type":"FULL_STATE_DICT"}',
|
|
encoding="utf-8",
|
|
)
|
|
descriptor = _assert_contract(inspect_checkpoint(tmp_path))
|
|
assert "json_invalid" in descriptor["safety"]["violations"]
|
|
|
|
for path in tuple(tmp_path.iterdir()):
|
|
path.unlink()
|
|
_save(tmp_path / "model.safetensors", weight=torch.ones(1))
|
|
_save(tmp_path / "adapter_model.safetensors", lora_A=torch.ones(1))
|
|
(tmp_path / "adapter_config.json").write_text("{}", encoding="utf-8")
|
|
with pytest.raises(CheckpointContractError) as tensors:
|
|
inspect_checkpoint(tmp_path, limits=InspectionLimits(max_tensors=1))
|
|
assert tensors.value.code == "DCI_RESOURCE_LIMIT"
|
|
assert tensors.value.detail == "max_tensors"
|
|
|
|
|
|
def test_invalid_bounded_json_is_reported_without_vendor_reader(tmp_path):
|
|
(tmp_path / ".metadata").write_bytes(b"opaque")
|
|
metadata = tmp_path / "fsdp_metadata.json"
|
|
metadata.write_text("[]", encoding="utf-8")
|
|
|
|
descriptor = _assert_contract(inspect_checkpoint(tmp_path))
|
|
assert "json_object_required" in descriptor["safety"]["violations"]
|
|
|
|
metadata.write_text("{", encoding="utf-8")
|
|
descriptor = _assert_contract(inspect_checkpoint(tmp_path))
|
|
assert "json_invalid" in descriptor["safety"]["violations"]
|
|
|
|
metadata.write_text("{}", encoding="utf-8")
|
|
with pytest.raises(CheckpointContractError) as json_limit:
|
|
inspect_checkpoint(tmp_path, limits=InspectionLimits(max_json_bytes=1))
|
|
assert json_limit.value.detail == "max_json_bytes"
|
|
|
|
|
|
def test_final_revalidation_detects_removed_source(tmp_path, monkeypatch):
|
|
weights = tmp_path / "model.safetensors"
|
|
_save(weights, weight=torch.ones(1))
|
|
original = inspection_module._format_components
|
|
|
|
def remove_after_probe(files, limits):
|
|
result = original(files, limits)
|
|
weights.unlink()
|
|
return result
|
|
|
|
monkeypatch.setattr(inspection_module, "_format_components", remove_after_probe)
|
|
with pytest.raises(CheckpointContractError) as caught:
|
|
inspect_checkpoint(tmp_path)
|
|
assert caught.value.code == "DCI_SOURCE_CHANGED"
|
|
|
|
|
|
def test_final_revalidation_detects_nested_inventory_mutation(tmp_path, monkeypatch):
|
|
nested = tmp_path / "nested"
|
|
nested.mkdir()
|
|
_save(nested / "model.safetensors", weight=torch.ones(1))
|
|
original = inspection_module._format_components
|
|
|
|
def add_after_probe(files, limits):
|
|
result = original(files, limits)
|
|
(nested / "late-file").write_text("changed", encoding="utf-8")
|
|
return result
|
|
|
|
monkeypatch.setattr(inspection_module, "_format_components", add_after_probe)
|
|
with pytest.raises(CheckpointContractError) as caught:
|
|
inspect_checkpoint(tmp_path)
|
|
assert caught.value.code == "DCI_SOURCE_CHANGED"
|
|
|
|
|
|
def test_descriptor_read_error_is_mapped_to_stable_source_change(tmp_path, monkeypatch):
|
|
weights = tmp_path / "model.safetensors"
|
|
_save(weights, weight=torch.ones(1))
|
|
real_read = inspection_module.os.read
|
|
calls = 0
|
|
|
|
def fail_after_inventory(descriptor, size):
|
|
nonlocal calls
|
|
calls += 1
|
|
if calls > 2:
|
|
raise OSError("injected read race")
|
|
return real_read(descriptor, size)
|
|
|
|
monkeypatch.setattr(inspection_module.os, "read", fail_after_inventory)
|
|
with pytest.raises(CheckpointContractError) as caught:
|
|
inspect_checkpoint(tmp_path)
|
|
assert caught.value.code == "DCI_SOURCE_CHANGED"
|
|
assert caught.value.detail == "source_changed"
|