wehub-resource-sync
145 lines
5.1 KiB
Python
145 lines
5.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""End-to-end tests for multi-model support in the StoreController.
|
|
|
|
A single ``finish_write`` with mixed-model keys must result in per-model
|
|
``submit_store_task`` calls so each submission sees uniform ``(shape, dtype)``.
|
|
"""
|
|
|
|
# Standard
|
|
from unittest.mock import patch
|
|
import time
|
|
|
|
# Third Party
|
|
import pytest
|
|
import torch
|
|
|
|
# First Party
|
|
from lmcache.v1.distributed.api import MemoryLayoutDesc, ObjectKey
|
|
from lmcache.v1.distributed.config import (
|
|
EvictionConfig,
|
|
L1ManagerConfig,
|
|
L1MemoryManagerConfig,
|
|
StorageManagerConfig,
|
|
)
|
|
from lmcache.v1.distributed.l2_adapters.config import L2AdaptersConfig
|
|
from lmcache.v1.distributed.l2_adapters.mock_l2_adapter import (
|
|
MockL2Adapter,
|
|
MockL2AdapterConfig,
|
|
)
|
|
from lmcache.v1.distributed.storage_manager import StorageManager
|
|
|
|
pytestmark = pytest.mark.skipif(
|
|
not torch.cuda.is_available(), reason="CUDA is not available"
|
|
)
|
|
|
|
|
|
# =============================================================================
|
|
# Helpers
|
|
# =============================================================================
|
|
|
|
|
|
def make_object_key(chunk_id: int, model_name: str) -> ObjectKey:
|
|
return ObjectKey(
|
|
chunk_hash=ObjectKey.IntHash2Bytes(chunk_id),
|
|
model_name=model_name,
|
|
kv_rank=0,
|
|
)
|
|
|
|
|
|
def wait_for_condition(predicate, timeout: float = 10.0, poll_interval: float = 0.05):
|
|
deadline = time.monotonic() + timeout
|
|
while time.monotonic() < deadline:
|
|
if predicate():
|
|
return True
|
|
time.sleep(poll_interval)
|
|
return False
|
|
|
|
|
|
def make_storage_manager(l1_size_mb: int = 256) -> StorageManager:
|
|
cfg = StorageManagerConfig(
|
|
l1_manager_config=L1ManagerConfig(
|
|
memory_config=L1MemoryManagerConfig(
|
|
size_in_bytes=l1_size_mb * 1024 * 1024,
|
|
use_lazy=True,
|
|
init_size_in_bytes=min(l1_size_mb, 64) * 1024 * 1024,
|
|
),
|
|
),
|
|
eviction_config=EvictionConfig(eviction_policy="LRU"),
|
|
l2_adapter_config=L2AdaptersConfig(
|
|
adapters=[MockL2AdapterConfig(max_size_gb=0.1, mock_bandwidth_gb=10.0)],
|
|
),
|
|
)
|
|
return StorageManager(cfg)
|
|
|
|
|
|
# =============================================================================
|
|
# Tests
|
|
# =============================================================================
|
|
|
|
|
|
class TestStoreControllerMultimodel:
|
|
"""Mixed-model keys in one ``finish_write`` must be split per model
|
|
before reaching ``submit_store_task``."""
|
|
|
|
def test_each_submit_store_task_has_uniform_shape(self):
|
|
layout_a = MemoryLayoutDesc(
|
|
shapes=[torch.Size([100, 2, 512])], dtypes=[torch.bfloat16]
|
|
)
|
|
layout_b = MemoryLayoutDesc(
|
|
shapes=[torch.Size([50, 2, 256])], dtypes=[torch.bfloat16]
|
|
)
|
|
keys_a = [make_object_key(i, "model_a") for i in range(3)]
|
|
keys_b = [make_object_key(100 + i, "model_b") for i in range(3)]
|
|
|
|
# Record the (shape, dtype) set observed per submit_store_task call.
|
|
submit_shape_groups = []
|
|
original_submit = MockL2Adapter.submit_store_task
|
|
|
|
def recording_submit(self, keys, objects):
|
|
shape_set = frozenset(
|
|
(tuple(obj.get_shapes()), tuple(obj.get_dtypes())) for obj in objects
|
|
)
|
|
submit_shape_groups.append(shape_set)
|
|
return original_submit(self, keys, objects)
|
|
|
|
with patch.object(MockL2Adapter, "submit_store_task", recording_submit):
|
|
sm = make_storage_manager()
|
|
adapter = sm._l2_adapters[0]
|
|
|
|
ret_a = sm.reserve_write(keys_a, layout_a, mode="new")
|
|
ret_b = sm.reserve_write(keys_b, layout_b, mode="new")
|
|
for i, k in enumerate(keys_a):
|
|
ret_a[k].tensor.fill_(float(i + 1))
|
|
for i, k in enumerate(keys_b):
|
|
ret_b[k].tensor.fill_(float(100 + i))
|
|
|
|
# Single finish_write forces a mixed-model batch in the listener.
|
|
sm.finish_write(keys_a + keys_b)
|
|
|
|
# Wait until every key has landed in L2 — avoids a race on the
|
|
# in_flight counter (which is briefly 0 between pop and submit).
|
|
ok = wait_for_condition(
|
|
lambda: all(
|
|
adapter.debug_has_key(k) # type: ignore[attr-defined]
|
|
for k in keys_a + keys_b
|
|
),
|
|
timeout=10.0,
|
|
)
|
|
assert ok, "Not all keys were stored in L2 within timeout"
|
|
|
|
assert submit_shape_groups, "No submit_store_task calls were recorded"
|
|
for i, shapes in enumerate(submit_shape_groups):
|
|
assert len(shapes) == 1, (
|
|
f"submit_store_task #{i} received keys with mixed "
|
|
f"(shape, dtype): {shapes}. Mixed-model batches must be "
|
|
f"grouped by model before reaching the adapter."
|
|
)
|
|
|
|
distinct_shapes = {s for group in submit_shape_groups for s in group}
|
|
assert len(distinct_shapes) == 2, (
|
|
f"Expected both models' shapes to appear across submits, "
|
|
f"got {len(distinct_shapes)}: {distinct_shapes}"
|
|
)
|
|
|
|
sm.close()
|