wehub-resource-sync
122 lines
3.2 KiB
Python
122 lines
3.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# Standard
|
|
from unittest.mock import patch
|
|
import tempfile
|
|
import time
|
|
|
|
# Third Party
|
|
import pytest
|
|
import torch
|
|
|
|
pytest.importorskip("vllm", reason="EC connector adapter imports vLLM at module top")
|
|
|
|
# First Party
|
|
from lmcache.integration.vllm.vllm_ec_adapter import ( # noqa: E402
|
|
LMCacheECConnectorImpl,
|
|
LMCacheECConnectorMetadata,
|
|
MMMeta,
|
|
)
|
|
|
|
|
|
class _FakeTransferConfig:
|
|
def __init__(self, storage_path: str):
|
|
self._storage_path = storage_path
|
|
|
|
def get_from_extra_config(self, key: str, default=None):
|
|
if key == "shared_storage_path":
|
|
return self._storage_path
|
|
return default
|
|
|
|
|
|
class _FakeVllmConfig:
|
|
def __init__(self, storage_path: str):
|
|
self.ec_transfer_config = _FakeTransferConfig(storage_path)
|
|
self.model_config = _FakeModelConfig()
|
|
self.parallel_config = _FakeParallelConfig()
|
|
self.cache_config = _FakeCacheConfig()
|
|
|
|
|
|
class _FakeModelConfig:
|
|
model = "fake-model"
|
|
served_model_name = "fake-model"
|
|
dtype = torch.float16
|
|
|
|
def get_num_layers(self, parallel_config):
|
|
return 1
|
|
|
|
def get_num_kv_heads(self, parallel_config):
|
|
return 1
|
|
|
|
def get_head_size(self):
|
|
return 1
|
|
|
|
|
|
class _FakeParallelConfig:
|
|
world_size = 1
|
|
rank = 0
|
|
|
|
|
|
class _FakeCacheConfig:
|
|
cache_dtype = "auto"
|
|
|
|
|
|
class _FakeRole:
|
|
name = "WORKER"
|
|
|
|
|
|
class _FakeParent:
|
|
def __init__(self):
|
|
self.is_producer = True
|
|
self._meta = None
|
|
|
|
def _get_connector_metadata(self):
|
|
return self._meta
|
|
|
|
|
|
def test_ec_roundtrip_save_then_load():
|
|
with tempfile.TemporaryDirectory() as td:
|
|
vllm_config = _FakeVllmConfig(td)
|
|
parent = _FakeParent()
|
|
conn = LMCacheECConnectorImpl(
|
|
vllm_config=vllm_config,
|
|
role=_FakeRole(),
|
|
parent=parent,
|
|
)
|
|
|
|
try:
|
|
mm_hash = "hash_abc"
|
|
x = torch.randn(7, 13, dtype=torch.float16)
|
|
encoder_cache = {mm_hash: x}
|
|
|
|
conn.save_caches(encoder_cache, mm_hash)
|
|
|
|
# EC storage is asynchronous across tiers; wait briefly for visibility.
|
|
deadline = time.time() + 2.0
|
|
while not conn.has_cache_item(mm_hash) and time.time() < deadline:
|
|
time.sleep(0.05)
|
|
|
|
assert conn.has_cache_item(mm_hash)
|
|
|
|
# Minimal fake metadata plumbing
|
|
meta = LMCacheECConnectorMetadata()
|
|
meta.add_mm_data(MMMeta.make_meta(mm_hash))
|
|
parent._meta = meta
|
|
|
|
# Pin device to "cpu" for the load path: on CPU-only CI runners
|
|
# vLLM's UnspecifiedPlatform returns an empty device_type, which
|
|
# would make tensor.to(device="") raise. The test exercises the
|
|
# round-trip itself, not vLLM's device discovery.
|
|
encoder_cache2 = {}
|
|
with patch(
|
|
"lmcache.integration.vllm.vllm_ec_adapter.torch_device_type",
|
|
new="cpu",
|
|
):
|
|
conn.start_load_caches(encoder_cache2)
|
|
|
|
assert mm_hash in encoder_cache2
|
|
assert encoder_cache2[mm_hash].shape == x.shape
|
|
assert encoder_cache2[mm_hash].dtype == x.dtype
|
|
finally:
|
|
conn.close()
|