| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
| import tempfile |
| import unittest |
|
|
| import kvcache |
|
|
|
|
| def metadata() -> dict[str, object]: |
| return { |
| "model_revision": "tiny-lab-model@0123456789abcdef", |
| "tokenizer_sha256": "7f" * 32, |
| "rope_theta": 10_000.0, |
| "layers": 4, |
| "kv_heads": 2, |
| "head_dim": 8, |
| "dtype": "f16", |
| "layout": "layer-major-k-then-v", |
| "sequence_start": 0, |
| "sequence_length": 37, |
| } |
|
|
|
|
| def tensor_metadata() -> dict[str, object]: |
| value = metadata() |
| value.update({ |
| "layers": 1, |
| "kv_heads": 1, |
| "head_dim": 2, |
| "sequence_length": 2, |
| "tensor_manifest": { |
| "version": 1, |
| "byte_order": "little", |
| "batch_size": 1, |
| "cache_class": "dynamic", |
| "attention_type": "causal", |
| "tensors": [ |
| { |
| "name": "layer.0.key", |
| "kind": "key", |
| "layer": 0, |
| "shape": [1, 1, 2, 2], |
| "strides": [8, 8, 4, 2], |
| "offset": 0, |
| "nbytes": 8, |
| "dtype": "f16", |
| }, |
| { |
| "name": "layer.0.value", |
| "kind": "value", |
| "layer": 0, |
| "shape": [1, 1, 2, 2], |
| "strides": [8, 8, 4, 2], |
| "offset": 8, |
| "nbytes": 8, |
| "dtype": "f16", |
| }, |
| ], |
| }, |
| }) |
| return value |
|
|
|
|
| class CacheTests(unittest.TestCase): |
| def test_round_trip_tensor_manifest_with_exact_payload_size(self) -> None: |
| with tempfile.TemporaryDirectory() as root: |
| path = Path(root) / "real-tensors.kvc" |
| payload = bytes(range(16)) |
| kvcache.write_cache(path, tensor_metadata(), payload) |
| self.assertEqual((tensor_metadata(), payload), kvcache.read_cache(path, tensor_metadata())) |
| receipt = kvcache.inspect_cache(path) |
| self.assertTrue(receipt["tensor_manifest_verified"]) |
| self.assertEqual(2, receipt["tensor_count"]) |
|
|
| def test_rejects_tensor_manifest_payload_size_mismatch(self) -> None: |
| with tempfile.TemporaryDirectory() as root: |
| path = Path(root) / "wrong-size.kvc" |
| with self.assertRaisesRegex(kvcache.CacheFormatError, "tensor manifest declares 16 payload bytes"): |
| kvcache.write_cache(path, tensor_metadata(), b"too short") |
| self.assertFalse(path.exists()) |
|
|
| def test_rejects_non_contiguous_tensor_offsets(self) -> None: |
| value = tensor_metadata() |
| value["tensor_manifest"]["tensors"][1]["offset"] = 9 |
| with tempfile.TemporaryDirectory() as root: |
| with self.assertRaisesRegex(kvcache.CacheFormatError, "contiguous"): |
| kvcache.write_cache(Path(root) / "gap.kvc", value, bytes(17)) |
|
|
| def test_rejects_tensor_shape_that_disagrees_with_cache_identity(self) -> None: |
| value = tensor_metadata() |
| value["tensor_manifest"]["tensors"][0]["shape"] = [1, 1, 1, 4] |
| with tempfile.TemporaryDirectory() as root: |
| with self.assertRaisesRegex(kvcache.CacheFormatError, "shape must equal"): |
| kvcache.write_cache(Path(root) / "wrong-shape.kvc", value, bytes(16)) |
|
|
| def test_rejects_non_contiguous_tensor_strides(self) -> None: |
| value = tensor_metadata() |
| value["tensor_manifest"]["tensors"][0]["strides"] = [1, 1, 1, 1] |
| with tempfile.TemporaryDirectory() as root: |
| with self.assertRaisesRegex(kvcache.CacheFormatError, "C-contiguous byte strides"): |
| kvcache.write_cache(Path(root) / "wrong-strides.kvc", value, bytes(16)) |
|
|
| def test_legacy_payload_remains_inspectable_but_is_not_tensor_verified(self) -> None: |
| with tempfile.TemporaryDirectory() as root: |
| path = Path(root) / "legacy.kvc" |
| kvcache.write_cache(path, metadata(), b"legacy") |
| receipt = kvcache.inspect_cache(path) |
| self.assertNotIn("tensor_manifest_verified", receipt) |
| self.assertNotIn("tensor_count", receipt) |
|
|
| def test_round_trip_raw_tensor_bytes(self) -> None: |
| with tempfile.TemporaryDirectory() as root: |
| path = Path(root) / "idea.kvc" |
| payload = bytes(range(255)) |
| kvcache.write_cache(path, metadata(), payload) |
| self.assertEqual((metadata(), payload), kvcache.read_cache(path, metadata())) |
|
|
| def test_inspection_receipt(self) -> None: |
| with tempfile.TemporaryDirectory() as root: |
| path = Path(root) / "idea.kvc" |
| kvcache.write_cache(path, metadata(), b"kv") |
| receipt = kvcache.inspect_cache(path) |
| self.assertEqual("KVC1", receipt["format"]) |
| self.assertEqual(2, receipt["payload_bytes"]) |
| json.dumps(receipt) |
|
|
| def test_rejects_wrong_model(self) -> None: |
| with tempfile.TemporaryDirectory() as root: |
| path = Path(root) / "idea.kvc" |
| kvcache.write_cache(path, metadata(), b"kv") |
| wrong = metadata() |
| wrong["model_revision"] = "other@revision" |
| with self.assertRaises(kvcache.CacheCompatibilityError): |
| kvcache.read_cache(path, wrong) |
|
|
| def test_rejects_wrong_position(self) -> None: |
| with tempfile.TemporaryDirectory() as root: |
| path = Path(root) / "idea.kvc" |
| kvcache.write_cache(path, metadata(), b"kv") |
| wrong = metadata() |
| wrong["sequence_start"] = 10 |
| with self.assertRaises(kvcache.CacheCompatibilityError): |
| kvcache.read_cache(path, wrong) |
|
|
| def test_expected_identity_keyword_accepts_a_valid_subset(self) -> None: |
| with tempfile.TemporaryDirectory() as root: |
| path = Path(root) / "idea.kvc" |
| kvcache.write_cache(path, metadata(), b"kv") |
| self.assertEqual( |
| (metadata(), b"kv"), |
| kvcache.read_cache( |
| path, |
| expected_identity={"model_revision": metadata()["model_revision"]}, |
| ), |
| ) |
| with self.assertRaises(kvcache.CacheCompatibilityError): |
| kvcache.read_cache(path, expected_identity={"dtype": "bf16"}) |
|
|
| def test_rejects_corruption(self) -> None: |
| with tempfile.TemporaryDirectory() as root: |
| path = Path(root) / "idea.kvc" |
| kvcache.write_cache(path, metadata(), b"kv") |
| raw = bytearray(path.read_bytes()) |
| raw[-1] ^= 1 |
| path.write_bytes(raw) |
| with self.assertRaises(kvcache.CacheFormatError): |
| kvcache.read_cache(path) |
|
|
| def test_rejects_truncation(self) -> None: |
| with tempfile.TemporaryDirectory() as root: |
| path = Path(root) / "idea.kvc" |
| kvcache.write_cache(path, metadata(), b"kv") |
| path.write_bytes(path.read_bytes()[:-1]) |
| with self.assertRaises(kvcache.CacheFormatError): |
| kvcache.read_cache(path) |
|
|
| def test_rejects_incomplete_identity(self) -> None: |
| value = metadata() |
| value.pop("tokenizer_sha256") |
| with tempfile.TemporaryDirectory() as root: |
| with self.assertRaises(kvcache.CacheFormatError): |
| kvcache.write_cache(Path(root) / "idea.kvc", value, b"kv") |
|
|
| def test_atomic_replacement(self) -> None: |
| with tempfile.TemporaryDirectory() as root: |
| path = Path(root) / "idea.kvc" |
| kvcache.write_cache(path, metadata(), b"first") |
| kvcache.write_cache(path, metadata(), b"second") |
| self.assertEqual(b"second", kvcache.read_cache(path)[1]) |
| self.assertEqual([], list(Path(root).glob("*.tmp"))) |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
