wehub-resource-sync
336 lines
13 KiB
Python
336 lines
13 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Tests for the unified touch feature in end_session.
|
|
|
|
This module tests:
|
|
1. L1EvictionPolicy.on_l1_keys_accessed bridges to policy.on_keys_touched
|
|
2. L1EvictionPolicy.on_l1_keys_read_finished also triggers touch
|
|
3. L1Manager.touch_keys dispatches to all registered listeners
|
|
4. MPCacheServer.end_session performs unified touch with correct keys
|
|
5. StorageManager.touch_l1_keys delegates to L1Manager
|
|
"""
|
|
|
|
# Standard
|
|
from unittest.mock import MagicMock
|
|
|
|
# Third Party
|
|
import pytest
|
|
|
|
# First Party
|
|
from lmcache.v1.distributed.api import ObjectKey, ipc_key_to_object_keys
|
|
from lmcache.v1.distributed.eviction import L1EvictionPolicy
|
|
from lmcache.v1.distributed.eviction_policy import LRUEvictionPolicy
|
|
from lmcache.v1.multiprocess.custom_types import IPCCacheServerKey
|
|
from lmcache.v1.multiprocess.session import SessionManager
|
|
from lmcache.v1.multiprocess.token_hasher import TokenHasher
|
|
|
|
# =============================================================================
|
|
# Helper Functions
|
|
# =============================================================================
|
|
|
|
|
|
def make_key(chunk_hash: int, model: str = "model", kv_rank: int = 0) -> ObjectKey:
|
|
"""Create an ObjectKey for testing."""
|
|
hash_bytes = ObjectKey.IntHash2Bytes(chunk_hash)
|
|
return ObjectKey(chunk_hash=hash_bytes, model_name=model, kv_rank=kv_rank)
|
|
|
|
|
|
def make_ipc_key(
|
|
token_ids: list[int],
|
|
chunk_size: int = 4,
|
|
model_name: str = "model",
|
|
world_size: int = 1,
|
|
worker_id: int | None = None,
|
|
request_id: str = "req-1",
|
|
) -> IPCCacheServerKey:
|
|
"""Create an IPCCacheServerKey for testing."""
|
|
return IPCCacheServerKey.from_token_ids(
|
|
model_name=model_name,
|
|
world_size=world_size,
|
|
worker_id=worker_id,
|
|
token_ids=token_ids,
|
|
start=0,
|
|
end=len(token_ids),
|
|
request_id=request_id,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def hasher() -> TokenHasher:
|
|
"""TokenHasher with small chunk_size for testing."""
|
|
return TokenHasher(chunk_size=4, hash_algorithm="blake3")
|
|
|
|
|
|
# =============================================================================
|
|
# L1EvictionPolicy Bridge Tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestL1EvictionPolicyAccessedBridge:
|
|
"""Tests for L1EvictionPolicy.on_l1_keys_accessed bridging."""
|
|
|
|
def test_on_l1_keys_accessed_calls_policy_on_keys_touched(self):
|
|
"""on_l1_keys_accessed should delegate to policy.on_keys_touched."""
|
|
mock_policy = MagicMock()
|
|
listener = L1EvictionPolicy(mock_policy)
|
|
keys = [make_key(1), make_key(2)]
|
|
|
|
listener.on_l1_keys_accessed(keys)
|
|
|
|
mock_policy.on_keys_touched.assert_called_once_with(keys)
|
|
|
|
def test_on_l1_keys_read_finished_also_touches(self):
|
|
"""on_l1_keys_read_finished should call on_keys_touched."""
|
|
mock_policy = MagicMock()
|
|
listener = L1EvictionPolicy(mock_policy)
|
|
keys = [make_key(1), make_key(2)]
|
|
|
|
listener.on_l1_keys_read_finished(keys)
|
|
|
|
mock_policy.on_keys_touched.assert_called_once_with(keys)
|
|
|
|
def test_on_l1_keys_accessed_with_lru_policy(self):
|
|
"""on_l1_keys_accessed should update LRU order via real LRU policy."""
|
|
lru = LRUEvictionPolicy()
|
|
listener = L1EvictionPolicy(lru)
|
|
|
|
# Create keys in order: 1, 2, 3
|
|
keys = [make_key(1), make_key(2), make_key(3)]
|
|
lru.on_keys_created(keys)
|
|
|
|
# Touch key 1 via the listener bridge
|
|
listener.on_l1_keys_accessed([make_key(1)])
|
|
|
|
# Eviction order should now be: 3, 2, 1 (1 is most recent)
|
|
candidates = lru.get_eviction_candidates(3)
|
|
assert ObjectKey.Bytes2IntHash(candidates[0].chunk_hash) == 3
|
|
assert ObjectKey.Bytes2IntHash(candidates[1].chunk_hash) == 2
|
|
assert ObjectKey.Bytes2IntHash(candidates[2].chunk_hash) == 1
|
|
|
|
def test_on_l1_keys_read_finished_updates_lru(self):
|
|
"""on_l1_keys_read_finished should update LRU order."""
|
|
lru = LRUEvictionPolicy()
|
|
listener = L1EvictionPolicy(lru)
|
|
|
|
# Create keys in order: 1, 2, 3
|
|
keys = [make_key(1), make_key(2), make_key(3)]
|
|
lru.on_keys_created(keys)
|
|
|
|
# Call on_l1_keys_read_finished on key 1 - should touch it
|
|
listener.on_l1_keys_read_finished([make_key(1)])
|
|
|
|
# Key 1 is now most recently touched, eviction order: 2, 3, 1
|
|
candidates = lru.get_eviction_candidates(3)
|
|
assert ObjectKey.Bytes2IntHash(candidates[0].chunk_hash) == 3
|
|
assert ObjectKey.Bytes2IntHash(candidates[1].chunk_hash) == 2
|
|
assert ObjectKey.Bytes2IntHash(candidates[2].chunk_hash) == 1
|
|
|
|
|
|
class TestL1EvictionPolicyUnifiedTouchOrder:
|
|
"""Tests for the unified touch producing correct eviction order.
|
|
|
|
The key benefit of unified touch is that all chunks of a request
|
|
are touched in order, so eviction proceeds from the last chunk
|
|
backwards (most recent chunk evicted last).
|
|
"""
|
|
|
|
def test_unified_touch_produces_reverse_eviction_order(self):
|
|
"""Touching keys in order [1,2,3] should evict 1 first (LRU)."""
|
|
lru = LRUEvictionPolicy()
|
|
listener = L1EvictionPolicy(lru)
|
|
|
|
keys = [make_key(i) for i in range(1, 6)]
|
|
lru.on_keys_created(keys)
|
|
|
|
# Unified touch: touch all keys in order (simulating end_session)
|
|
listener.on_l1_keys_accessed(keys)
|
|
|
|
# Eviction should be in reverse order: 1 first (least recently touched)
|
|
# because on_keys_touched processes them sequentially,
|
|
# so key 5 is the most recently touched
|
|
candidates = lru.get_eviction_candidates(5)
|
|
evicted_ids = [ObjectKey.Bytes2IntHash(c.chunk_hash) for c in candidates]
|
|
# LRU evicts from the back of the ordered dict (least recent first)
|
|
# After touching [1,2,3,4,5] in order, 1 is least recent, 5 is most recent
|
|
assert evicted_ids == [5, 4, 3, 2, 1]
|
|
|
|
def test_per_read_touch_vs_unified_touch_order(self):
|
|
"""Compare per-read touch (old behavior) vs unified touch (new behavior).
|
|
|
|
Old behavior: each retrieve touches its keys immediately.
|
|
New behavior: all keys are touched together at end_session.
|
|
|
|
With unified touch, the eviction order reflects the full request's
|
|
chunk ordering, not the order of individual retrieve calls.
|
|
"""
|
|
# Scenario: request has chunks [1,2,3,4,5]
|
|
# Retrieve touches [1,2,3] (prefix hit), Store creates [4,5]
|
|
lru = LRUEvictionPolicy()
|
|
|
|
keys = [make_key(i) for i in range(1, 6)]
|
|
lru.on_keys_created(keys)
|
|
|
|
# Unified touch at end_session: touch all in order
|
|
lru.on_keys_touched(keys)
|
|
|
|
candidates = lru.get_eviction_candidates(5)
|
|
evicted_ids = [ObjectKey.Bytes2IntHash(c.chunk_hash) for c in candidates]
|
|
# Last chunk (5) should be evicted last (most recently touched)
|
|
assert evicted_ids[-1] == 1
|
|
|
|
|
|
# =============================================================================
|
|
# End Session Touch Integration Tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestEndSessionTouchKeys:
|
|
"""Tests for the end_session unified touch key generation logic.
|
|
|
|
These tests verify that end_session correctly:
|
|
1. Retrieves the session with accumulated hashes
|
|
2. Generates the correct ObjectKeys using lookup_ipc_key
|
|
3. Handles edge cases (no session, no lookup key, empty hashes)
|
|
"""
|
|
|
|
def test_end_session_generates_correct_keys(self, hasher: TokenHasher):
|
|
"""end_session should generate ObjectKeys from session hashes and lookup key."""
|
|
mgr = SessionManager(hasher, ttl=600, cleanup_interval=None)
|
|
session = mgr.get_or_create("req-1")
|
|
|
|
tokens = list(range(12)) # 3 chunks of 4
|
|
session.set_tokens(tokens)
|
|
session.get_hashes(0, 12)
|
|
|
|
ipc_key = make_ipc_key(tokens, chunk_size=4, worker_id=None, request_id="req-1")
|
|
session.lookup_ipc_key = ipc_key
|
|
|
|
# Simulate end_session logic
|
|
removed = mgr.remove("req-1")
|
|
assert removed is not None
|
|
assert removed.lookup_ipc_key is not None
|
|
|
|
chunk_hashes = [TokenHasher.hash_to_bytes(h) for h in removed.get_hashes(0)]
|
|
obj_keys = ipc_key_to_object_keys(removed.lookup_ipc_key, chunk_hashes, [0])[0]
|
|
|
|
# With world_size=1 and worker_id=None, should have 3 keys
|
|
assert len(obj_keys) == 3
|
|
|
|
def test_end_session_expands_keys_for_world_size(self, hasher: TokenHasher):
|
|
"""end_session should expand keys for all workers when world_size > 1."""
|
|
mgr = SessionManager(hasher, ttl=600, cleanup_interval=None)
|
|
session = mgr.get_or_create("req-1")
|
|
|
|
tokens = list(range(8)) # 2 chunks of 4
|
|
session.set_tokens(tokens)
|
|
session.get_hashes(0, 8)
|
|
|
|
ipc_key = make_ipc_key(
|
|
tokens,
|
|
chunk_size=4,
|
|
world_size=2,
|
|
worker_id=None,
|
|
request_id="req-1",
|
|
)
|
|
session.lookup_ipc_key = ipc_key
|
|
|
|
removed = mgr.remove("req-1")
|
|
assert removed is not None
|
|
assert removed.lookup_ipc_key is not None
|
|
|
|
chunk_hashes = [TokenHasher.hash_to_bytes(h) for h in removed.get_hashes(0)]
|
|
obj_keys = ipc_key_to_object_keys(removed.lookup_ipc_key, chunk_hashes, [0])[0]
|
|
|
|
# 2 chunks * 2 workers = 4 keys
|
|
assert len(obj_keys) == 4
|
|
|
|
def test_end_session_no_session_skips_touch(self, hasher: TokenHasher):
|
|
"""end_session should skip touch when session doesn't exist."""
|
|
mgr = SessionManager(hasher, ttl=600, cleanup_interval=None)
|
|
removed = mgr.remove("nonexistent")
|
|
assert removed is None
|
|
# No error should occur
|
|
|
|
def test_end_session_no_lookup_key_skips_touch(self, hasher: TokenHasher):
|
|
"""end_session should skip touch when session has no lookup_ipc_key."""
|
|
mgr = SessionManager(hasher, ttl=600, cleanup_interval=None)
|
|
session = mgr.get_or_create("req-1")
|
|
session.set_tokens(list(range(8)))
|
|
session.get_hashes(0, 8)
|
|
# Don't set lookup_ipc_key
|
|
|
|
removed = mgr.remove("req-1")
|
|
assert removed is not None
|
|
assert removed.lookup_ipc_key is None
|
|
# No error should occur
|
|
|
|
def test_end_session_empty_hashes_produces_no_keys(self, hasher: TokenHasher):
|
|
"""end_session should produce no keys when session has no computed hashes."""
|
|
mgr = SessionManager(hasher, ttl=600, cleanup_interval=None)
|
|
session = mgr.get_or_create("req-1")
|
|
|
|
ipc_key = make_ipc_key(
|
|
list(range(8)), chunk_size=4, worker_id=None, request_id="req-1"
|
|
)
|
|
session.lookup_ipc_key = ipc_key
|
|
# Don't compute any hashes
|
|
|
|
removed = mgr.remove("req-1")
|
|
assert removed is not None
|
|
assert removed.lookup_ipc_key is not None
|
|
|
|
chunk_hashes = [TokenHasher.hash_to_bytes(h) for h in removed.get_hashes(0)]
|
|
obj_keys = ipc_key_to_object_keys(removed.lookup_ipc_key, chunk_hashes, [0])[0]
|
|
assert len(obj_keys) == 0
|
|
|
|
def test_end_session_hashes_cover_retrieve_and_store(self, hasher: TokenHasher):
|
|
"""Session hashes should cover both retrieve and store ranges.
|
|
|
|
Simulates the typical flow:
|
|
1. lookup computes all hashes (but via TokenHasher, not Session)
|
|
2. retrieve calls session.get_hashes(0, hit_end)
|
|
3. store calls session.get_hashes(hit_end, total)
|
|
4. end_session uses session.get_hashes(0) to get all
|
|
"""
|
|
mgr = SessionManager(hasher, ttl=600, cleanup_interval=None)
|
|
session = mgr.get_or_create("req-1")
|
|
|
|
tokens = list(range(20)) # 5 chunks of 4
|
|
session.set_tokens(tokens)
|
|
|
|
# Simulate retrieve: first 3 chunks
|
|
retrieve_hashes = session.get_hashes(0, 12)
|
|
assert len(retrieve_hashes) == 3
|
|
|
|
# Simulate store: last 2 chunks
|
|
store_hashes = session.get_hashes(12, 20)
|
|
assert len(store_hashes) == 2
|
|
|
|
# All hashes should cover all 5 chunks
|
|
all_hashes = session.get_hashes(0)
|
|
assert len(all_hashes) == 5
|
|
assert all_hashes == retrieve_hashes + store_hashes
|
|
|
|
|
|
# =============================================================================
|
|
# StoreListener No-op Tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestStoreListenerAccessedNoop:
|
|
"""Tests that StoreListener.on_l1_keys_accessed is a no-op."""
|
|
|
|
def test_store_listener_accessed_is_noop(self):
|
|
"""StoreListener.on_l1_keys_accessed should not raise."""
|
|
# Import here to avoid import errors if store_controller has
|
|
# heavy dependencies
|
|
# First Party
|
|
from lmcache.v1.distributed.storage_controllers.store_controller import (
|
|
StoreListener,
|
|
)
|
|
|
|
# StoreListener requires an event_fd; use a mock
|
|
listener = MagicMock(spec=StoreListener)
|
|
# Call the real method on the class
|
|
StoreListener.on_l1_keys_accessed(listener, [make_key(1)])
|
|
# Should not raise - that's the test
|