| """ |
| ENGRAM Protocol — Serializer Tests |
| Tests for .eng safetensors serialize/deserialize round-trip (D7). |
| """ |
|
|
| from __future__ import annotations |
|
|
| from pathlib import Path |
|
|
| import pytest |
| import torch |
| from safetensors.torch import load_file |
|
|
| from kvcos.core.serializer import EngramSerializer, SerializationError |
| from kvcos.core.types import CompressionMethod |
| from tests.conftest import make_synthetic_kv |
| from kvcos.core.cache_spec import LLAMA_3_1_8B |
|
|
|
|
| class TestSerializeRoundTrip: |
| """Serialize → deserialize preserves shape, dtype, metadata.""" |
|
|
| def test_round_trip_shape(self, tmp_data_dir: Path) -> None: |
| keys, values = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=256) |
| s = EngramSerializer() |
| eng = tmp_data_dir / "test.eng" |
|
|
| s.serialize( |
| keys=keys, values=values, |
| agent_id="test-agent", task_description="unit test", |
| model_id=LLAMA_3_1_8B["model_id"], output_path=eng, |
| compression=CompressionMethod.FP16, |
| ) |
| k_out, v_out, meta = s.deserialize(eng) |
|
|
| assert k_out.shape == keys.shape |
| assert v_out.shape == values.shape |
|
|
| def test_metadata_fields(self, tmp_data_dir: Path) -> None: |
| keys, values = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=64) |
| s = EngramSerializer() |
| eng = tmp_data_dir / "meta.eng" |
|
|
| s.serialize( |
| keys=keys, values=values, |
| agent_id="agent-42", task_description="metadata check", |
| model_id=LLAMA_3_1_8B["model_id"], output_path=eng, |
| compression=CompressionMethod.Q8_0, |
| ) |
| _, _, meta = s.deserialize(eng) |
|
|
| assert meta["agent_id"] == "agent-42" |
| assert meta["task_description"] == "metadata check" |
| assert meta["compression"] == "q8_0" |
| assert meta["n_layers"] == "32" |
| assert meta["model_family"] == "llama" |
|
|
| def test_safetensors_loadable(self, tmp_data_dir: Path) -> None: |
| """D7: File must be valid safetensors.""" |
| keys, values = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=64) |
| s = EngramSerializer() |
| eng = tmp_data_dir / "valid.eng" |
|
|
| s.serialize( |
| keys=keys, values=values, |
| agent_id="test", task_description="safetensors check", |
| model_id=LLAMA_3_1_8B["model_id"], output_path=eng, |
| compression=CompressionMethod.FP16, |
| ) |
| tensors = load_file(str(eng)) |
| assert "layer_0_keys" in tensors |
| assert "layer_0_values" in tensors |
|
|
| def test_result_dict(self, tmp_data_dir: Path) -> None: |
| keys, values = make_synthetic_kv(LLAMA_3_1_8B, ctx_len=64) |
| s = EngramSerializer() |
| eng = tmp_data_dir / "result.eng" |
|
|
| result = s.serialize( |
| keys=keys, values=values, |
| agent_id="test", task_description="result check", |
| model_id=LLAMA_3_1_8B["model_id"], output_path=eng, |
| ) |
| assert "cache_id" in result |
| assert result["size_bytes"] > 0 |
| assert result["n_layers"] == 32 |
|
|
|
|
| class TestSerializerErrors: |
| """Edge cases and error handling.""" |
|
|
| def test_shape_mismatch_raises(self, tmp_data_dir: Path) -> None: |
| keys = torch.randn(32, 8, 64, 128, dtype=torch.float16) |
| values = torch.randn(32, 8, 32, 128, dtype=torch.float16) |
| s = EngramSerializer() |
|
|
| with pytest.raises(SerializationError, match="mismatch"): |
| s.serialize( |
| keys=keys, values=values, |
| agent_id="t", task_description="t", |
| model_id="test", output_path=tmp_data_dir / "bad.eng", |
| ) |
|
|
| def test_3d_tensor_raises(self, tmp_data_dir: Path) -> None: |
| keys = torch.randn(8, 64, 128, dtype=torch.float16) |
| s = EngramSerializer() |
|
|
| with pytest.raises(SerializationError, match="4D"): |
| s.serialize( |
| keys=keys, values=keys, |
| agent_id="t", task_description="t", |
| model_id="test", output_path=tmp_data_dir / "bad.eng", |
| ) |
|
|
| def test_missing_file_raises(self, tmp_data_dir: Path) -> None: |
| s = EngramSerializer() |
| with pytest.raises(SerializationError, match="not found"): |
| s.deserialize(tmp_data_dir / "nonexistent.eng") |
|
|
