diff --git a/src/memos/memories/activation/kv.py b/src/memos/memories/activation/kv.py index 1981b958f..90b35996f 100644 --- a/src/memos/memories/activation/kv.py +++ b/src/memos/memories/activation/kv.py @@ -1,18 +1,148 @@ +import hashlib +import json import os import pickle from datetime import datetime +from typing import Any from transformers import DynamicCache from memos.configs.memory import KVCacheMemoryConfig from memos.dependency import require_python_package from memos.llms.factory import LLMFactory +from memos.log import get_logger from memos.memories.activation.base import BaseActMemory from memos.memories.activation.item import KVCacheItem from memos.memories.textual.item import TextualMemoryItem +logger = get_logger(__name__) + + +# --------------------------------------------------------------------------- +# Producer fingerprint helpers (see openspec change 2300-... for the contract) +# --------------------------------------------------------------------------- + +# The strict subset of fingerprint fields we compare on load. A field mismatches +# only if both sides are non-None and unequal — unknown fields never fail. +_STRICT_FIELDS: tuple[str, ...] = ( + "model_name_or_path", + "config_fingerprint", + "tokenizer_fingerprint", + "torch_dtype", + "quantization", +) + + +def _compute_producer_fingerprint(llm: Any) -> dict[str, Any]: + """Return an identity fingerprint for the LLM that produced a KV cache. + + Every field is best-effort — a missing attribute becomes ``None`` rather + than an exception. The loader treats ``None`` as unknown, not as mismatch, + so a partial fingerprint stays useful. + """ + fp: dict[str, Any] = { + "model_name_or_path": None, + "backend": None, + "config_fingerprint": None, + "tokenizer_fingerprint": None, + "torch_dtype": None, + "quantization": None, + "architectures": None, + "num_hidden_layers": None, + "num_kv_heads": None, + "head_dim": None, + "transformers_version": None, + } + + llm_config = getattr(llm, "config", None) + if llm_config is not None: + fp["model_name_or_path"] = getattr(llm_config, "model_name_or_path", None) + fp["backend"] = type(llm_config).__name__ + + model = getattr(llm, "model", None) + if model is not None: + try: + model_config = getattr(model, "config", None) + if model_config is not None: + to_diff_dict = getattr(model_config, "to_diff_dict", None) + if callable(to_diff_dict): + diff = to_diff_dict() + fp["config_fingerprint"] = hashlib.sha256( + json.dumps(diff, sort_keys=True, default=str).encode() + ).hexdigest() + quant = getattr(model_config, "quantization_config", None) + fp["quantization"] = str(quant) if quant is not None else None + arch = getattr(model_config, "architectures", None) + if arch: + fp["architectures"] = list(arch) + fp["num_hidden_layers"] = getattr(model_config, "num_hidden_layers", None) + _num_kv = getattr(model_config, "num_key_value_heads", None) + fp["num_kv_heads"] = ( + _num_kv + if _num_kv is not None + else getattr(model_config, "num_attention_heads", None) + ) + head_dim = getattr(model_config, "head_dim", None) + if head_dim is None: + hidden = getattr(model_config, "hidden_size", None) + n_heads = getattr(model_config, "num_attention_heads", None) + if hidden and n_heads: + head_dim = hidden // n_heads + fp["head_dim"] = head_dim + dtype = getattr(model, "dtype", None) + if dtype is not None: + fp["torch_dtype"] = str(dtype) + except Exception: + logger.warning( + "Failed to introspect HF model config for KV cache fingerprint", + exc_info=True, + ) + + tokenizer = getattr(llm, "tokenizer", None) + if tokenizer is not None: + try: + backend_tokenizer = getattr(tokenizer, "backend_tokenizer", None) + if backend_tokenizer is not None and hasattr(backend_tokenizer, "to_str"): + fp["tokenizer_fingerprint"] = hashlib.sha256( + backend_tokenizer.to_str().encode() + ).hexdigest() + except Exception: + logger.warning( + "Failed to compute tokenizer fingerprint for KV cache", + exc_info=True, + ) + + try: + import transformers + + fp["transformers_version"] = transformers.__version__ + except Exception: + logger.debug("Could not determine transformers version for fingerprint", exc_info=True) + + return fp + + +def _fingerprint_mismatch_reasons(saved: dict[str, Any] | None, live: dict[str, Any]) -> list[str]: + """Return a list of human-readable mismatch reasons. + + Empty list means "safe to install" (either every field matched or one side + was unknown). Only strict fields participate in the decision. + """ + if not isinstance(saved, dict): + return [] + reasons: list[str] = [] + for field in _STRICT_FIELDS: + s_val = saved.get(field) + l_val = live.get(field) + if s_val is None or l_val is None: + continue + if s_val != l_val: + reasons.append(f"{field}: saved={s_val!r} live={l_val!r}") + return reasons + + class KVCacheMemory(BaseActMemory): """ Key-Value Cache Memory for activation memories. @@ -29,6 +159,10 @@ def __init__(self, config: KVCacheMemoryConfig) -> None: self.llm = LLMFactory.from_config(config.extractor_llm) self.kv_cache_memories: dict[str, KVCacheItem] = {} + def _producer_fingerprint(self) -> dict[str, Any]: + """Fingerprint of the LLM currently wired into this memory.""" + return _compute_producer_fingerprint(self.llm) + def extract(self, text: str) -> KVCacheItem: """Extract memory based on the text. @@ -46,7 +180,11 @@ def extract(self, text: str) -> KVCacheItem: # Create a KVCacheItem with the extracted cache cache_item = KVCacheItem( memory=kv_cache, - metadata={"source_text": text, "extracted_at": datetime.now().isoformat()}, + metadata={ + "source_text": text, + "extracted_at": datetime.now().isoformat(), + "producer": self._producer_fingerprint(), + }, ) return cache_item @@ -134,7 +272,53 @@ def from_textual_memory(self, mem: TextualMemoryItem) -> KVCacheItem: """ # Build KV cache from the textual memory content kv_cache = self.llm.build_kv_cache(mem.memory) - return KVCacheItem(memory=kv_cache, metadata=mem.metadata.model_dump()) + metadata = mem.metadata.model_dump() + metadata["producer"] = self._producer_fingerprint() + return KVCacheItem(memory=kv_cache, metadata=metadata) + + def _verify_and_filter(self, memories: dict[str, KVCacheItem]) -> dict[str, KVCacheItem]: + """Compare each item's saved producer fingerprint against the live LLM. + + Drops items whose fingerprint disagrees on any strict field. Items + without a fingerprint are kept with a warning (backward compatibility + with pre-2.0.30 caches). + """ + if not memories: + return {} + + # Live fingerprint may itself fail to compute for an unusual backend. + # In that case skip verification with a warning — do not regress + # loading reliability. + try: + live_fp = self._producer_fingerprint() + except Exception: + logger.warning( + "Failed to compute live producer fingerprint; loading KV cache without verification", + exc_info=True, + ) + return memories + + verified: dict[str, KVCacheItem] = {} + for item_id, item in memories.items(): + item_metadata = getattr(item, "metadata", None) or {} + saved_fp = item_metadata.get("producer") if isinstance(item_metadata, dict) else None + if saved_fp is None: + logger.warning( + "KV cache item %s has no producer fingerprint (pre-2.0.30 cache); loading unchecked", + item_id, + ) + verified[item_id] = item + continue + reasons = _fingerprint_mismatch_reasons(saved_fp, live_fp) + if reasons: + logger.error( + "KV cache item %s dropped: producer fingerprint mismatch (%s)", + item_id, + "; ".join(reasons), + ) + continue + verified[item_id] = item + return verified def load(self, dir: str) -> None: """Load memories from os.path.join(dir, self.config.memory_filename) @@ -163,21 +347,30 @@ def load(self, dir: str) -> None: memories = data["kv_cache_memories"] if isinstance(memories, list): # Convert list to dict format - self.kv_cache_memories = {item.id: item for item in memories} + candidate = {item.id: item for item in memories} else: - self.kv_cache_memories = memories + candidate = memories + self.kv_cache_memories = self._verify_and_filter(candidate) else: # Reset to empty if no memories in data self.kv_cache_memories = {} elif isinstance(data, list): # Backward compatibility: convert list to dict - self.kv_cache_memories = {item.id: item for item in data} + candidate = {item.id: item for item in data} + self.kv_cache_memories = self._verify_and_filter(candidate) else: # Reset to empty if data format is unexpected self.kv_cache_memories = {} - except (EOFError, pickle.UnpicklingError, Exception): - # If loading fails, start with empty memories + except Exception: + # Corrupt or incompatible cache — log the reason so the failure is + # distinguishable from an empty cache in production. Loader stays + # resilient by resetting to an empty dict. + logger.warning( + "Failed to load KV cache from %s; resetting to empty", + file_path, + exc_info=True, + ) self.kv_cache_memories = {} def dump(self, dir: str) -> None: diff --git a/src/memos/memories/activation/vllmkv.py b/src/memos/memories/activation/vllmkv.py index 4b74115f4..216bb8ff8 100644 --- a/src/memos/memories/activation/vllmkv.py +++ b/src/memos/memories/activation/vllmkv.py @@ -2,15 +2,24 @@ import pickle from datetime import datetime +from typing import Any from memos.configs.memory import KVCacheMemoryConfig from memos.dependency import require_python_package from memos.llms.factory import LLMFactory +from memos.log import get_logger from memos.memories.activation.base import BaseActMemory from memos.memories.activation.item import VLLMKVCacheItem +from memos.memories.activation.kv import ( + _compute_producer_fingerprint, + _fingerprint_mismatch_reasons, +) from memos.memories.textual.item import TextualMemoryItem +logger = get_logger(__name__) + + class VLLMKVCacheMemory(BaseActMemory): """ VLLM Key-Value Cache Memory for activation memories. @@ -28,6 +37,10 @@ def __init__(self, config: KVCacheMemoryConfig) -> None: self.llm = LLMFactory.from_config(config.extractor_llm) self.kv_cache_memories: dict[str, VLLMKVCacheItem] = {} + def _producer_fingerprint(self) -> dict[str, Any]: + """Fingerprint of the LLM currently wired into this memory.""" + return _compute_producer_fingerprint(self.llm) + def extract(self, text: str) -> VLLMKVCacheItem: """Extract memory based on the text. @@ -47,7 +60,11 @@ def extract(self, text: str) -> VLLMKVCacheItem: # Create a VLLMKVCacheItem with the extracted prompt cache_item = VLLMKVCacheItem( memory=prompt, - metadata={"source_text": text, "extracted_at": datetime.now().isoformat()}, + metadata={ + "source_text": text, + "extracted_at": datetime.now().isoformat(), + "producer": self._producer_fingerprint(), + }, ) return cache_item @@ -140,7 +157,51 @@ def from_textual_memory(self, mem: TextualMemoryItem) -> VLLMKVCacheItem: """ # Build vLLM KV cache from the textual memory content prompt = self.llm.build_vllm_kv_cache(mem.memory) - return VLLMKVCacheItem(memory=prompt, metadata=mem.metadata.model_dump()) + metadata = mem.metadata.model_dump() + metadata["producer"] = self._producer_fingerprint() + return VLLMKVCacheItem(memory=prompt, metadata=metadata) + + def _verify_and_filter( + self, memories: dict[str, VLLMKVCacheItem] + ) -> dict[str, VLLMKVCacheItem]: + """Compare each item's saved producer fingerprint against the live LLM. + + Drops items with an incompatible fingerprint. Items without a + fingerprint are kept with a warning (backward compatibility with + pre-2.0.30 caches). + """ + if not memories: + return {} + try: + live_fp = self._producer_fingerprint() + except Exception: + logger.warning( + "Failed to compute live producer fingerprint; loading vLLM KV cache without verification", + exc_info=True, + ) + return memories + + verified: dict[str, VLLMKVCacheItem] = {} + for item_id, item in memories.items(): + item_metadata = getattr(item, "metadata", None) or {} + saved_fp = item_metadata.get("producer") if isinstance(item_metadata, dict) else None + if saved_fp is None: + logger.warning( + "vLLM KV cache item %s has no producer fingerprint (pre-2.0.30 cache); loading unchecked", + item_id, + ) + verified[item_id] = item + continue + reasons = _fingerprint_mismatch_reasons(saved_fp, live_fp) + if reasons: + logger.warning( + "vLLM KV cache item %s dropped: producer fingerprint mismatch (%s)", + item_id, + "; ".join(reasons), + ) + continue + verified[item_id] = item + return verified def load(self, dir: str) -> None: """Load memories from os.path.join(dir, self.config.memory_filename) @@ -169,21 +230,30 @@ def load(self, dir: str) -> None: memories = data["kv_cache_memories"] if isinstance(memories, list): # Convert list to dict format - self.kv_cache_memories = {item.id: item for item in memories} + candidate = {item.id: item for item in memories} else: - self.kv_cache_memories = memories + candidate = memories + self.kv_cache_memories = self._verify_and_filter(candidate) else: # Reset to empty if no memories in data self.kv_cache_memories = {} elif isinstance(data, list): # Backward compatibility: convert list to dict - self.kv_cache_memories = {item.id: item for item in data} + candidate = {item.id: item for item in data} + self.kv_cache_memories = self._verify_and_filter(candidate) else: # Reset to empty if data format is unexpected self.kv_cache_memories = {} - except (EOFError, pickle.UnpicklingError, Exception): - # If loading fails, start with empty memories + except Exception: + # Corrupt or incompatible cache — log the failure so it is + # distinguishable from an empty file in production. Loader + # stays resilient by resetting to an empty dict. + logger.warning( + "Failed to load vLLM KV cache from %s; resetting to empty", + file_path, + exc_info=True, + ) self.kv_cache_memories = {} def dump(self, dir: str) -> None: diff --git a/tests/memories/activation/test_kv.py b/tests/memories/activation/test_kv.py index 6490d687f..599a6ae80 100644 --- a/tests/memories/activation/test_kv.py +++ b/tests/memories/activation/test_kv.py @@ -1,3 +1,6 @@ +import logging +import pickle + from unittest.mock import MagicMock import pytest @@ -10,77 +13,286 @@ from memos.memories.activation.kv import KVCacheMemory +# ------------------------------------------------------------------ fixtures -- + + @pytest.fixture def dummy_config(): - # Minimal config mock for KVCacheMemory config = MagicMock(spec=KVCacheMemoryConfig) config.extractor_llm = MagicMock() config.memory_filename = "test_kv_cache.pkl" return config +def _make_fake_llm( + *, + model_name="test/writer", + config_fp="cfg-A", + tokenizer_fp="tok-A", + torch_dtype="torch.float32", + quantization=None, + architectures=("LlamaForCausalLM",), + num_hidden_layers=2, + num_kv_heads=4, + head_dim=8, +): + """Build a MagicMock LLM whose _compute_producer_fingerprint output is + controllable by the test. + + We take advantage of the fingerprint helper reading .config / + .model.config / .tokenizer.backend_tokenizer — the helper is deliberately + defensive so a MagicMock with the right shape suffices. + """ + llm = MagicMock() + llm.config.model_name_or_path = model_name + llm.config.__class__.__name__ = "HFLLMConfig" + + # model.config.to_diff_dict() → drives config_fingerprint + llm.model.config.to_diff_dict.return_value = {"marker": config_fp} + llm.model.config.quantization_config = quantization + llm.model.config.architectures = list(architectures) + llm.model.config.num_hidden_layers = num_hidden_layers + llm.model.config.num_key_value_heads = num_kv_heads + llm.model.config.num_attention_heads = num_kv_heads + llm.model.config.head_dim = head_dim + llm.model.dtype = torch_dtype + + # tokenizer.backend_tokenizer.to_str() → drives tokenizer_fingerprint + llm.tokenizer.backend_tokenizer.to_str.return_value = tokenizer_fp + + # build_kv_cache should return a real DynamicCache so pickle round-trips + llm.build_kv_cache = MagicMock(return_value=DynamicCache()) + return llm + + @pytest.fixture -def kv_memory(dummy_config): - # Patch LLMFactory to avoid real LLM calls - with pytest.MonkeyPatch.context() as m: - from memos.llms import factory - - m.setattr( - factory.LLMFactory, - "from_config", - lambda cfg: MagicMock(build_kv_cache=lambda x: DynamicCache()), - ) - yield KVCacheMemory(dummy_config) +def kv_memory_factory(dummy_config, monkeypatch): + """Build KVCacheMemory instances with an injected fake LLM per test.""" + from memos.llms import factory as llm_factory + + def _factory(llm=None): + llm = llm or _make_fake_llm() + monkeypatch.setattr(llm_factory.LLMFactory, "from_config", lambda cfg: llm) + return KVCacheMemory(dummy_config), llm + + return _factory -def make_filled_cache(): - # Create a DynamicCache with at least one dummy tensor layer +# ------------------------------------------------------------------ helpers -- + + +def _make_populated_cache(): + """Return a DynamicCache with one layer, using the new-API .update path.""" cache = DynamicCache() - cache.key_cache.append(torch.zeros(1, 2, 3)) - cache.value_cache.append(torch.zeros(1, 2, 3)) + keys = torch.zeros(1, 2, 3, 4) + values = torch.zeros(1, 2, 3, 4) + cache.update(keys, values, 0) return cache -def test_extract_and_add_and_get(kv_memory): - # Test extract, add, and get functionality - item = kv_memory.extract("hello world") +def _dump_and_load(kv, tmpdir): + kv.dump(str(tmpdir)) + kv.kv_cache_memories = {} + kv.load(str(tmpdir)) + return kv + + +# -------------------------------------------------------------------- tests -- + + +class TestProducerFingerprintCapture: + def test_extract_attaches_producer_metadata(self, kv_memory_factory): + kv, _ = kv_memory_factory() + item = kv.extract("hello world") + assert "producer" in item.metadata + producer = item.metadata["producer"] + # required minimums + assert producer["model_name_or_path"] == "test/writer" + assert producer["backend"] == "HFLLMConfig" + assert producer["config_fingerprint"] is not None + assert producer["tokenizer_fingerprint"] is not None + assert producer["torch_dtype"] == "torch.float32" + + def test_from_textual_memory_attaches_producer(self, kv_memory_factory): + kv, _ = kv_memory_factory() + + class DummyTextual: + memory = "foo" + metadata = MagicMock(model_dump=lambda: {"bar": 1}) + + item = kv.from_textual_memory(DummyTextual()) + assert item.metadata["bar"] == 1 + assert "producer" in item.metadata + assert item.metadata["producer"]["model_name_or_path"] == "test/writer" + + def test_extract_survives_broken_fingerprint_source(self, kv_memory_factory): + # A backend that cannot expose model.config still produces an item — + # extract() must not raise even when introspection fails partially. + llm = _make_fake_llm() + llm.model.config.to_diff_dict.side_effect = RuntimeError("boom") + kv, _ = kv_memory_factory(llm=llm) + item = kv.extract("robust") + assert isinstance(item, KVCacheItem) + # producer still present, with at least model_name_or_path + assert item.metadata["producer"]["model_name_or_path"] == "test/writer" + # config_fingerprint is None because introspection failed + assert item.metadata["producer"].get("config_fingerprint") is None + + +class TestProducerMismatchOnLoad: + def test_match_installs_item_silently(self, kv_memory_factory, tmp_path, caplog): + kv, _ = kv_memory_factory() + item = kv.extract("prompt") + kv.add([item]) + with caplog.at_level(logging.WARNING, logger="memos"): + _dump_and_load(kv, tmp_path) + # item still present + assert item.id in kv.kv_cache_memories + # nothing angry logged + errors = [r for r in caplog.records if r.levelno >= logging.ERROR] + assert not errors + + def test_config_mismatch_drops_item_and_logs_error(self, kv_memory_factory, tmp_path, caplog): + writer = _make_fake_llm(config_fp="writer-cfg") + kv, _ = kv_memory_factory(llm=writer) + item = kv.extract("prompt") + kv.add([item]) + kv.dump(str(tmp_path)) + + # Simulate a fresh process with a different model loaded. + reader = _make_fake_llm(config_fp="reader-cfg") + kv2, _ = kv_memory_factory(llm=reader) + with caplog.at_level(logging.ERROR, logger="memos"): + kv2.load(str(tmp_path)) + assert kv2.kv_cache_memories == {} + error_lines = [r.getMessage() for r in caplog.records if r.levelno >= logging.ERROR] + assert any("config_fingerprint" in msg for msg in error_lines) + assert any(item.id in msg for msg in error_lines) + + def test_tokenizer_mismatch_drops_item(self, kv_memory_factory, tmp_path, caplog): + writer = _make_fake_llm(tokenizer_fp="writer-tok") + kv, _ = kv_memory_factory(llm=writer) + item = kv.extract("prompt") + kv.add([item]) + kv.dump(str(tmp_path)) + + reader = _make_fake_llm(tokenizer_fp="reader-tok") + kv2, _ = kv_memory_factory(llm=reader) + with caplog.at_level(logging.ERROR, logger="memos"): + kv2.load(str(tmp_path)) + assert kv2.kv_cache_memories == {} + + def test_dtype_mismatch_drops_item(self, kv_memory_factory, tmp_path, caplog): + writer = _make_fake_llm(torch_dtype="torch.float32") + kv, _ = kv_memory_factory(llm=writer) + item = kv.extract("prompt") + kv.add([item]) + kv.dump(str(tmp_path)) + + reader = _make_fake_llm(torch_dtype="torch.bfloat16") + kv2, _ = kv_memory_factory(llm=reader) + with caplog.at_level(logging.ERROR, logger="memos"): + kv2.load(str(tmp_path)) + assert kv2.kv_cache_memories == {} + + def test_missing_producer_key_installs_with_warning(self, kv_memory_factory, tmp_path, caplog): + # Simulate a pre-fix cache: manually write an item without a + # producer key. Loader must warn but keep the item (deprecation + # window semantics). + kv, _ = kv_memory_factory() + legacy_item = KVCacheItem( + memory=DynamicCache(), + metadata={"source_text": "legacy", "extracted_at": "old"}, + ) + payload = {"kv_cache_memories": {legacy_item.id: legacy_item}} + with open(tmp_path / kv.config.memory_filename, "wb") as f: + pickle.dump(payload, f, protocol=pickle.HIGHEST_PROTOCOL) + + with caplog.at_level(logging.WARNING, logger="memos"): + kv.load(str(tmp_path)) + assert legacy_item.id in kv.kv_cache_memories + assert any( + "no producer" in r.getMessage().lower() + for r in caplog.records + if r.levelno == logging.WARNING + ) + + def test_mixed_batch_only_good_items_installed(self, kv_memory_factory, tmp_path, caplog): + writer = _make_fake_llm(config_fp="writer-cfg") + kv, _ = kv_memory_factory(llm=writer) + good = kv.extract("stays") + bad = kv.extract("drops") + # Corrupt one producer to force a mismatch on load + bad.metadata["producer"]["config_fingerprint"] = "mismatch" + kv.add([good, bad]) + kv.dump(str(tmp_path)) + + reader = _make_fake_llm(config_fp="writer-cfg") + kv2, _ = kv_memory_factory(llm=reader) + with caplog.at_level(logging.ERROR, logger="memos"): + kv2.load(str(tmp_path)) + assert good.id in kv2.kv_cache_memories + assert bad.id not in kv2.kv_cache_memories + + +class TestCorruptCacheLogging: + def test_corrupt_pickle_logs_warning_and_resets(self, kv_memory_factory, tmp_path, caplog): + kv, _ = kv_memory_factory() + file_path = tmp_path / kv.config.memory_filename + # Write a garbage byte stream that will fail to unpickle + file_path.write_bytes(b"\x00\x01not-a-pickle-file") + + with caplog.at_level(logging.WARNING, logger="memos"): + kv.load(str(tmp_path)) + assert kv.kv_cache_memories == {} + assert any( + "kv cache" in r.getMessage().lower() or "load" in r.getMessage().lower() + for r in caplog.records + if r.levelno == logging.WARNING + ) + + +# ------------------------------------------ smoke tests for previously-passing -- + + +def test_extract_and_add_and_get(kv_memory_factory): + kv, _ = kv_memory_factory() + item = kv.extract("hello world") assert isinstance(item, KVCacheItem) assert isinstance(item.memory, DynamicCache) - kv_memory.add([item]) - got = kv_memory.get(item.id) - assert got is item + kv.add([item]) + assert kv.get(item.id) is item -def test_get_cache_merge(kv_memory): - # Test merging multiple KVCacheItems into a single DynamicCache - item1 = KVCacheItem(memory=make_filled_cache()) - item2 = KVCacheItem(memory=make_filled_cache()) - kv_memory.add([item1, item2]) - merged = kv_memory.get_cache([item1.id, item2.id]) +def test_get_cache_merge(kv_memory_factory): + kv, _ = kv_memory_factory() + item1 = KVCacheItem(memory=_make_populated_cache()) + item2 = KVCacheItem(memory=_make_populated_cache()) + kv.add([item1, item2]) + merged = kv.get_cache([item1.id, item2.id]) assert isinstance(merged, DynamicCache) - # Check the number of layers in merged key/value cache - assert len(merged.key_cache) == 1 - assert len(merged.value_cache) == 1 + assert len(merged.layers) == 1 + +def test_delete_and_get_all(kv_memory_factory): + kv, _ = kv_memory_factory() + item = KVCacheItem(memory=_make_populated_cache()) + kv.add([item]) + assert item in kv.get_all() + kv.delete([item.id]) + assert kv.get(item.id) is None + kv.add([item]) + kv.delete_all() + assert kv.get_all() == [] -def test_delete_and_get_all(kv_memory): - # Test delete and get_all functionality - item = KVCacheItem(memory=make_filled_cache()) - kv_memory.add([item]) - assert item in kv_memory.get_all() - kv_memory.delete([item.id]) - assert kv_memory.get(item.id) is None - kv_memory.add([item]) - kv_memory.delete_all() - assert kv_memory.get_all() == [] +def test_from_textual_memory(kv_memory_factory): + kv, _ = kv_memory_factory() -def test_from_textual_memory(kv_memory): - # Test conversion from textual memory to KVCacheItem class DummyTextualMemory: memory = "foo" metadata = MagicMock(model_dump=lambda: {"bar": 1}) - item = kv_memory.from_textual_memory(DummyTextualMemory()) + item = kv.from_textual_memory(DummyTextualMemory()) assert isinstance(item, KVCacheItem) assert item.metadata["bar"] == 1 diff --git a/tests/memories/activation/test_vllmkv.py b/tests/memories/activation/test_vllmkv.py new file mode 100644 index 000000000..63393a3dc --- /dev/null +++ b/tests/memories/activation/test_vllmkv.py @@ -0,0 +1,111 @@ +import hashlib +import logging +import os +import pickle + +from unittest.mock import MagicMock + +import pytest + +from memos.configs.memory import KVCacheMemoryConfig +from memos.memories.activation.item import VLLMKVCacheItem +from memos.memories.activation.vllmkv import VLLMKVCacheMemory + + +@pytest.fixture +def dummy_config(): + config = MagicMock(spec=KVCacheMemoryConfig) + config.extractor_llm = MagicMock() + config.memory_filename = "test_vllm_kv_cache.pkl" + return config + + +def _make_fake_vllm(model_name="test/vllm-writer", tokenizer_fp="tok-A"): + llm = MagicMock() + llm.config.model_name_or_path = model_name + llm.config.__class__.__name__ = "VLLMLLMConfig" + # vLLM backend often has no local model, but may have a tokenizer via HF. + llm.model = None + llm.tokenizer.backend_tokenizer.to_str.return_value = tokenizer_fp + llm.build_vllm_kv_cache = MagicMock(side_effect=lambda t: f"prompt::{t}") + return llm + + +@pytest.fixture +def vllm_memory_factory(dummy_config, monkeypatch): + from memos.llms import factory as llm_factory + + def _factory(llm=None): + llm = llm or _make_fake_vllm() + monkeypatch.setattr(llm_factory.LLMFactory, "from_config", lambda cfg: llm) + return VLLMKVCacheMemory(dummy_config), llm + + return _factory + + +class TestVLLMProducerFingerprint: + def test_extract_attaches_producer(self, vllm_memory_factory): + kv, _ = vllm_memory_factory() + item = kv.extract("hello") + assert "producer" in item.metadata + producer = item.metadata["producer"] + assert producer["model_name_or_path"] == "test/vllm-writer" + assert producer["backend"] == "VLLMLLMConfig" + # vLLM has no HF model → config_fingerprint is None; that's fine + assert producer["config_fingerprint"] is None + # Fingerprint must be the exact SHA-256 of the tokenizer string, not + # just any truthy value — locks in that the right field is being + # populated with the right hash function. + expected = hashlib.sha256(b"tok-A").hexdigest() + assert producer["tokenizer_fingerprint"] == expected + + def test_tokenizer_mismatch_drops_item(self, vllm_memory_factory, tmp_path, caplog): + writer = _make_fake_vllm(tokenizer_fp="writer-tok") + kv, _ = vllm_memory_factory(llm=writer) + item = kv.extract("prompt") + kv.add([item]) + kv.dump(str(tmp_path)) + + reader = _make_fake_vllm(tokenizer_fp="reader-tok") + kv2, _ = vllm_memory_factory(llm=reader) + with caplog.at_level(logging.WARNING, logger="memos"): + kv2.load(str(tmp_path)) + assert kv2.kv_cache_memories == {} + # Lock in that the drop was actually logged — otherwise a silent + # code path that empties the dict for a different reason would + # still make this test pass. + assert any( + r.levelno == logging.WARNING and "fingerprint mismatch" in r.getMessage().lower() + for r in caplog.records + ) + + def test_missing_producer_installs_with_warning(self, vllm_memory_factory, tmp_path, caplog): + kv, _ = vllm_memory_factory() + legacy = VLLMKVCacheItem( + memory="legacy-prompt", + metadata={"source_text": "old", "extracted_at": "old"}, + ) + payload = {"kv_cache_memories": {legacy.id: legacy}} + with open(os.path.join(tmp_path, kv.config.memory_filename), "wb") as f: + pickle.dump(payload, f, protocol=pickle.HIGHEST_PROTOCOL) + + with caplog.at_level(logging.WARNING, logger="memos"): + kv.load(str(tmp_path)) + assert legacy.id in kv.kv_cache_memories + assert any( + "no producer" in r.getMessage().lower() + for r in caplog.records + if r.levelno == logging.WARNING + ) + + def test_corrupt_pickle_logs_warning(self, vllm_memory_factory, tmp_path, caplog): + kv, _ = vllm_memory_factory() + file_path = tmp_path / kv.config.memory_filename + file_path.write_bytes(b"\x00garbage") + with caplog.at_level(logging.WARNING, logger="memos"): + kv.load(str(tmp_path)) + assert kv.kv_cache_memories == {} + assert any( + r.levelno == logging.WARNING and "load" in r.getMessage().lower() + for r in caplog.records + )