240 lines
10 KiB
Python
240 lines
10 KiB
Python
import json
|
|
import os
|
|
import random
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest import mock
|
|
|
|
import torch
|
|
from safetensors.torch import save_file
|
|
|
|
from toolkit.trigger_binding_artifacts import (
|
|
ARTIFACT_TYPES,
|
|
ArtifactValidationError,
|
|
artifact_reference,
|
|
build_checkpoint_manifest,
|
|
canonical_json_dumps,
|
|
config_fingerprint,
|
|
decode_rng_state,
|
|
encode_rng_state,
|
|
load_artifact,
|
|
load_checkpoint_manifest,
|
|
phase_fingerprint,
|
|
save_artifact,
|
|
save_checkpoint_manifest,
|
|
sha256_bytes,
|
|
tensor_sha256,
|
|
source_fingerprint,
|
|
)
|
|
|
|
|
|
class TriggerBindingArtifactsTest(unittest.TestCase):
|
|
def setUp(self):
|
|
self.phase = {"name": "a1", "step": 12}
|
|
self.source = {"model": "ideogram-4", "revision": "abc123"}
|
|
self.config = {"rank": 1, "dtype": "bf16"}
|
|
self.tensors = {
|
|
"adapter.down.weight": torch.arange(6, dtype=torch.float32).reshape(2, 3),
|
|
"adapter.up.weight": torch.ones(3, 2, dtype=torch.bfloat16),
|
|
}
|
|
|
|
def _save(self, directory, artifact_type="te_adapter", tensors=None):
|
|
path = Path(directory) / f"{artifact_type}.safetensors"
|
|
manifest = save_artifact(
|
|
path,
|
|
artifact_type,
|
|
tensors or self.tensors,
|
|
phase=self.phase,
|
|
source=self.source,
|
|
config=self.config,
|
|
extra={"trigger": "<literal>"},
|
|
)
|
|
return path, manifest
|
|
|
|
def test_all_artifact_types_round_trip_with_metadata_and_hashes(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
for artifact_type in sorted(ARTIFACT_TYPES):
|
|
with self.subTest(artifact_type=artifact_type):
|
|
path, manifest = self._save(temp_dir, artifact_type)
|
|
loaded, loaded_manifest = load_artifact(
|
|
path,
|
|
expected_type=artifact_type,
|
|
expected_keys=self.tensors.keys(),
|
|
expected_shapes={key: tensor.shape for key, tensor in self.tensors.items()},
|
|
expected_phase_fingerprint=phase_fingerprint(self.phase),
|
|
expected_source_fingerprint=source_fingerprint(self.source),
|
|
expected_config_fingerprint=config_fingerprint(self.config),
|
|
)
|
|
self.assertEqual(manifest, loaded_manifest)
|
|
for key, tensor in self.tensors.items():
|
|
self.assertTrue(torch.equal(loaded[key], tensor))
|
|
self.assertEqual(
|
|
len(loaded_manifest["tensors"][key]["sha256"]),
|
|
64,
|
|
)
|
|
|
|
def test_tensor_hash_supports_bfloat16_scalars_empty_and_noncontiguous_tensors(self):
|
|
scalar = torch.tensor(1.25, dtype=torch.bfloat16)
|
|
vector = scalar.reshape(1)
|
|
self.assertEqual(tensor_sha256(scalar), tensor_sha256(vector))
|
|
|
|
empty = torch.empty(0, dtype=torch.bfloat16)
|
|
self.assertEqual(tensor_sha256(empty), sha256_bytes(b""))
|
|
|
|
base = torch.arange(12, dtype=torch.float32).reshape(3, 4)
|
|
noncontiguous = base.transpose(0, 1)
|
|
self.assertFalse(noncontiguous.is_contiguous())
|
|
self.assertEqual(
|
|
tensor_sha256(noncontiguous),
|
|
tensor_sha256(noncontiguous.contiguous()),
|
|
)
|
|
|
|
def test_artifact_round_trip_supports_bfloat16_scalar_tensor(self):
|
|
tensors = {
|
|
"adapter.scale": torch.tensor(1.0, dtype=torch.bfloat16),
|
|
"adapter.weight": torch.ones(2, 2, dtype=torch.bfloat16),
|
|
}
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
path, manifest = self._save(temp_dir, tensors=tensors)
|
|
loaded, loaded_manifest = load_artifact(path, expected_type="te_adapter")
|
|
self.assertEqual(manifest, loaded_manifest)
|
|
self.assertEqual(loaded["adapter.scale"].shape, torch.Size([]))
|
|
self.assertTrue(torch.equal(loaded["adapter.scale"], tensors["adapter.scale"]))
|
|
|
|
def test_fingerprints_are_canonical_and_order_independent(self):
|
|
self.assertEqual(config_fingerprint({"a": 1, "b": 2}), config_fingerprint({"b": 2, "a": 1}))
|
|
self.assertNotEqual(config_fingerprint({"a": 1}), config_fingerprint({"a": 2}))
|
|
|
|
def test_load_fails_closed_on_expected_key_shape_and_fingerprint_mismatch(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
path, _ = self._save(temp_dir)
|
|
cases = [
|
|
{"expected_keys": ["adapter.down.weight"]},
|
|
{"expected_shapes": {key: [99] for key in self.tensors}},
|
|
{"expected_phase_fingerprint": "0" * 64},
|
|
{"expected_source_fingerprint": "1" * 64},
|
|
{"expected_config_fingerprint": "2" * 64},
|
|
{"expected_file_sha256": "3" * 64},
|
|
{"expected_type": "embedding"},
|
|
]
|
|
for kwargs in cases:
|
|
with self.subTest(kwargs=kwargs), self.assertRaises(ArtifactValidationError):
|
|
load_artifact(path, **kwargs)
|
|
|
|
def test_load_rejects_unmanaged_safetensors_without_required_metadata(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
path = Path(temp_dir) / "bare.safetensors"
|
|
save_file({"weight": torch.ones(1)}, str(path))
|
|
with self.assertRaises(ArtifactValidationError):
|
|
load_artifact(path)
|
|
|
|
def test_load_rejects_tampered_tensor_manifest_hash(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
path, manifest = self._save(temp_dir)
|
|
tampered = json.loads(json.dumps(manifest))
|
|
tampered["tensors"]["adapter.down.weight"]["sha256"] = "0" * 64
|
|
manifest_json = canonical_json_dumps(tampered)
|
|
metadata = {
|
|
"trigger_binding.schema": tampered["schema"],
|
|
"trigger_binding.schema_version": str(tampered["schema_version"]),
|
|
"trigger_binding.artifact_type": tampered["artifact_type"],
|
|
"trigger_binding.manifest": manifest_json,
|
|
"trigger_binding.manifest_sha256": sha256_bytes(manifest_json.encode("utf-8")),
|
|
}
|
|
save_file(self.tensors, str(path), metadata=metadata)
|
|
with self.assertRaisesRegex(ArtifactValidationError, "SHA-256 mismatch"):
|
|
load_artifact(path)
|
|
|
|
def test_atomic_save_preserves_previous_destination_on_replace_failure(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
path = Path(temp_dir) / "te_adapter.safetensors"
|
|
path.write_bytes(b"previous-good-file")
|
|
with mock.patch("toolkit.trigger_binding_artifacts.os.replace", side_effect=OSError("replace failed")):
|
|
with self.assertRaises(OSError):
|
|
save_artifact(
|
|
path,
|
|
"te_adapter",
|
|
self.tensors,
|
|
phase=self.phase,
|
|
source=self.source,
|
|
config=self.config,
|
|
)
|
|
self.assertEqual(path.read_bytes(), b"previous-good-file")
|
|
leftovers = [item for item in os.listdir(temp_dir) if item != path.name]
|
|
self.assertEqual(leftovers, [])
|
|
|
|
def test_rng_state_json_round_trip_supports_python_numpy_and_torch(self):
|
|
state = {
|
|
"python": random.Random(7).getstate(),
|
|
"torch": torch.get_rng_state(),
|
|
"cuda": [torch.arange(8, dtype=torch.uint8)],
|
|
}
|
|
try:
|
|
import numpy as np
|
|
|
|
state["numpy"] = np.random.RandomState(11).get_state()
|
|
except ImportError:
|
|
np = None
|
|
|
|
encoded = encode_rng_state(state)
|
|
json.dumps(encoded, allow_nan=False)
|
|
decoded = decode_rng_state(encoded)
|
|
self.assertEqual(decoded["python"], state["python"])
|
|
self.assertTrue(torch.equal(decoded["torch"], state["torch"]))
|
|
self.assertTrue(torch.equal(decoded["cuda"][0], state["cuda"][0]))
|
|
if np is not None:
|
|
self.assertEqual(decoded["numpy"][0], state["numpy"][0])
|
|
self.assertTrue(np.array_equal(decoded["numpy"][1], state["numpy"][1]))
|
|
self.assertEqual(decoded["numpy"][2:], state["numpy"][2:])
|
|
|
|
def test_checkpoint_manifest_round_trip_and_artifact_verification(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
artifact_path, _ = self._save(temp_dir, "embedding")
|
|
reference = artifact_reference(artifact_path, relative_to=temp_dir)
|
|
manifest = build_checkpoint_manifest(
|
|
phase=self.phase,
|
|
source=self.source,
|
|
config=self.config,
|
|
step=42,
|
|
artifacts={"embedding": reference},
|
|
rng_state={"python": random.Random(3).getstate(), "torch": torch.get_rng_state()},
|
|
extra={"optimizer": "adamw"},
|
|
)
|
|
manifest_path = Path(temp_dir) / "checkpoint_manifest.json"
|
|
save_checkpoint_manifest(manifest_path, manifest)
|
|
loaded = load_checkpoint_manifest(
|
|
manifest_path,
|
|
expected_phase_fingerprint=phase_fingerprint(self.phase),
|
|
expected_source_fingerprint=source_fingerprint(self.source),
|
|
expected_config_fingerprint=config_fingerprint(self.config),
|
|
)
|
|
self.assertEqual(loaded, manifest)
|
|
decoded_rng = decode_rng_state(loaded["rng_state"])
|
|
self.assertTrue(torch.equal(decoded_rng["torch"], torch.get_rng_state()))
|
|
|
|
with open(artifact_path, "ab") as handle:
|
|
handle.write(b"tamper")
|
|
with self.assertRaises(ArtifactValidationError):
|
|
load_checkpoint_manifest(manifest_path)
|
|
|
|
def test_checkpoint_manifest_rejects_unknown_schema_keys(self):
|
|
manifest = build_checkpoint_manifest(
|
|
phase=self.phase,
|
|
source=self.source,
|
|
config=self.config,
|
|
step=0,
|
|
artifacts={},
|
|
rng_state=None,
|
|
)
|
|
manifest["unexpected"] = True
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
path = Path(temp_dir) / "checkpoint.json"
|
|
path.write_text(json.dumps(manifest), encoding="utf-8")
|
|
with self.assertRaises(ArtifactValidationError):
|
|
load_checkpoint_manifest(path, verify_artifacts=False)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|