lmcache-cli 0.4.5.dev0__py3-none-any.whl

lmcache/v1/cache_interface.py

@@ -0,0 +1,19 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import Optional
+
+# Third Party
+import msgspec
+
+
+class LMCacheModelRequest(
+    msgspec.Struct,
+    array_like=True,  # type: ignore[call-arg]
+    omit_defaults=True,
+):  # type: ignore[call-arg]
+    """
+    User-provided information to control the cache behavior.
+    """
+
+    store_cache: bool = True  # Whether to store the cache
+    ttl: Optional[float] = None  # Time to live

lmcache/v1/check/__init__.py

@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Check mode registry implementation"""
+
+# Standard
+from typing import Callable, Dict, Optional
+import importlib
+import inspect
+import os
+
+# First Party
+from lmcache.logging import init_logger
+
+logger = init_logger(__name__)
+
+
+class CheckModeRegistry:
+    """Registry for dynamically loaded check modes"""
+
+    def __init__(self):
+        self.modes: Dict[str, Callable] = {}
+        self.loaded = False
+
+    def register(self, name: str, func: Callable):
+        """Register a check mode function"""
+        if name in self.modes:
+            raise ValueError(f"Check mode '{name}' already registered")
+        self.modes[name] = func
+
+    def load_modes(self):
+        """Dynamically load all check mode modules"""
+        if self.loaded:
+            return
+
+        # Get current package
+        current_dir = os.path.dirname(__file__)
+
+        # Find all modules with check_mode_ prefix
+        for filename in os.listdir(current_dir):
+            if filename.startswith("check_mode_") and filename.endswith(".py"):
+                module_name = filename[:-3]  # Remove .py
+                try:
+                    module = importlib.import_module(
+                        f".{module_name}", package=__package__
+                    )
+                    # Find and register mode functions
+                    for name, obj in inspect.getmembers(module):
+                        if inspect.isfunction(obj) and hasattr(obj, "is_check_mode"):
+                            self.register(obj.mode_name, obj)
+                except ImportError as e:
+                    logger.error(f"Failed to load check mode module {module_name}: {e}")
+
+        self.loaded = True
+        logger.info(f"Loaded {len(self.modes)} check modes")
+
+    def get_mode(self, name: str) -> Optional[Callable]:
+        """Get registered mode function. Returns None if the mode is not found."""
+        if not self.loaded:
+            self.load_modes()
+        return self.modes.get(name)
+
+
+def check_mode(name: str):
+    """Decorator to mark functions as check modes"""
+
+    def decorator(func):
+        func.is_check_mode = True
+        func.mode_name = name
+        return func
+
+    return decorator
+
+
+# Global registry instance
+registry = CheckModeRegistry()

lmcache/v1/check/check_mode_gen.py

@@ -0,0 +1,86 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Generate mode implementation for key generation"""
+
+# Third Party
+import tqdm
+
+# First Party
+from lmcache.v1.check import check_mode
+from lmcache.v1.check.utils import (
+    _get_default_metadata,
+    create_memory_objects_batch,
+    create_storage_manager_with_config,
+    create_test_key,
+    find_remote_backend,
+    flow_control_check,
+    wait_put_tasks_complete,
+)
+
+
+@check_mode("gen")
+async def run_gen_mode(
+    model: str, num_keys: int, concurrency: int, offset: int = 0, **kwargs
+):
+    """Run key generation mode"""
+    # Create storage manager using common function
+    storage_manager = create_storage_manager_with_config(model)
+    metadata = _get_default_metadata(model)
+
+    try:
+        print("Generate: Passed - Created storage manager with valid config")
+
+        # Find remote backend for flow control
+        remote_backend = find_remote_backend(storage_manager)
+
+        # Create limited number of memory objects for reuse (memory efficiency)
+        batch_size = min(concurrency, 100)  # Limit to 100 for memory efficiency
+        memory_objs = create_memory_objects_batch(storage_manager, metadata, batch_size)
+
+        if not memory_objs:
+            print("Generate: Failed - Could not allocate any memory objects")
+            return
+
+        # Create progress bar
+        progress_bar = tqdm.tqdm(
+            total=num_keys, desc="Generating keys", unit="key", unit_scale=True
+        )
+        sleep_count = 1.0
+        # Process keys in batches of concurrency size
+        for batch_start in range(0, num_keys, concurrency):
+            batch_end = min(batch_start + concurrency, num_keys)
+            batch_keys = []
+            batch_memory_objs = []
+
+            # Create keys and reuse memory objects for this batch
+            for i in range(batch_start, batch_end):
+                key = create_test_key(model, f"gen_{offset + i}")
+                # Reuse memory objects in round-robin fashion
+                memory_obj = memory_objs[i % len(memory_objs)]
+                batch_keys.append(key)
+                batch_memory_objs.append(memory_obj)
+                memory_obj.ref_count_up()
+
+            # Flow control: check if remote backend has too many pending tasks
+            sleep_count = await flow_control_check(
+                remote_backend, concurrency, sleep_count
+            )
+
+            # Use batched_put to store the batch of memory objects
+            storage_manager.batched_put(batch_keys, batch_memory_objs)
+
+            # Update progress bar
+            progress_bar.update(len(batch_keys))
+
+        progress_bar.close()
+        print(f"Generate: Successfully generated {num_keys} keys")
+
+        # Wait for remote backend put_tasks to complete
+        wait_put_tasks_complete(find_remote_backend(storage_manager))
+
+    except Exception as e:
+        print(
+            f"Generate: Failed - Error creating storage manager with valid config: {e}"
+        )
+    finally:
+        if storage_manager:
+            storage_manager.close()

lmcache/v1/check/check_mode_test_l2_adapter.py

@@ -0,0 +1,284 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Test mode implementation for MP mode L2 adapter basic checks"""
+
+# Standard
+import argparse
+import os
+import select
+import time
+
+# Third Party
+import torch
+
+# First Party
+from lmcache.v1.check import check_mode
+from lmcache.v1.check.utils import (
+    DEFAULT_KV_DTYPE_STR,
+    DEFAULT_OBJ_SIZE,
+    parse_kv_dtype,
+    print_performance_results,
+)
+from lmcache.v1.distributed.api import ObjectKey
+from lmcache.v1.distributed.l2_adapters import create_l2_adapter
+from lmcache.v1.distributed.l2_adapters.config import (
+    parse_args_to_l2_adapters_config,
+)
+from lmcache.v1.memory_management import (
+    MemoryFormat,
+    MemoryObjMetadata,
+    TensorMemoryObj,
+)
+
+_POLL_TIMEOUT_MS = 100000
+
+
+def _create_object_key(model: str, key_id: str) -> ObjectKey:
+    """Create a test ObjectKey."""
+    return ObjectKey(
+        chunk_hash=ObjectKey.IntHash2Bytes(hash(key_id) & 0xFFFFFFFF),
+        model_name=model,
+        kv_rank=0,
+    )
+
+
+def _create_memory_obj(
+    fill_value: float = 0.0,
+    obj_size: int = DEFAULT_OBJ_SIZE,
+    dtype: torch.dtype = torch.float32,
+) -> TensorMemoryObj:
+    """Create a test TensorMemoryObj."""
+    raw_data = torch.empty(obj_size, dtype=dtype)
+    raw_data.fill_(fill_value)
+    metadata = MemoryObjMetadata(
+        shape=torch.Size([obj_size]),
+        dtype=dtype,
+        address=0,
+        phy_size=obj_size * raw_data.element_size(),
+        fmt=MemoryFormat.KV_2LTD,
+        ref_count=1,
+    )
+    return TensorMemoryObj(raw_data, metadata, parent_allocator=None)
+
+
+def _wait_event_fd(efd: int, timeout_ms: int = _POLL_TIMEOUT_MS) -> bool:
+    """Wait for an eventfd to be signaled."""
+    poll = select.poll()
+    poll.register(efd, select.POLLIN)
+    events = poll.poll(timeout_ms)
+    if events:
+        try:
+            os.eventfd_read(efd)
+        except BlockingIOError:
+            pass
+        return True
+    return False
+
+
+def _run_store_phase(adapter, keys, objects):
+    """Run store phase and return (stats, success)."""
+    efd = adapter.get_store_event_fd()
+    start = time.perf_counter()
+    task_id = adapter.submit_store_task(keys, objects)
+    if not _wait_event_fd(efd):
+        print("  Store: timed out waiting for eventfd")
+        return None, False
+    completed = adapter.pop_completed_store_tasks()
+    elapsed_ms = (time.perf_counter() - start) * 1000
+    ok = completed.get(task_id, False)
+    return elapsed_ms, ok
+
+
+def _run_lookup_phase(adapter, keys):
+    """Run lookup phase and return (stats, bitmap)."""
+    efd = adapter.get_lookup_and_lock_event_fd()
+    start = time.perf_counter()
+    task_id = adapter.submit_lookup_and_lock_task(keys)
+    if not _wait_event_fd(efd):
+        print("  Lookup: timed out waiting for eventfd")
+        return None, None
+    bitmap = adapter.query_lookup_and_lock_result(task_id)
+    elapsed_ms = (time.perf_counter() - start) * 1000
+    return elapsed_ms, bitmap
+
+
+def _run_load_phase(adapter, keys, buffers):
+    """Run load phase and return (stats, bitmap)."""
+    efd = adapter.get_load_event_fd()
+    start = time.perf_counter()
+    task_id = adapter.submit_load_task(keys, buffers)
+    if not _wait_event_fd(efd):
+        print("  Load: timed out waiting for eventfd")
+        return None, None
+    bitmap = adapter.query_load_result(task_id)
+    elapsed_ms = (time.perf_counter() - start) * 1000
+    return elapsed_ms, bitmap
+
+
+@check_mode("test_l2_adapter")
+async def run_test_mode(model: str, **kwargs):
+    """Run L2 adapter test mode.
+
+    Requires ``l2_adapter`` in *kwargs* (list of JSON
+    strings from ``--l2-adapter``).
+    """
+    l2_adapter_raw = kwargs.get("l2_adapter")
+    if not l2_adapter_raw:
+        print("Error: --l2-adapter is required for test_l2_adapter mode")
+        return
+
+    obj_size = kwargs.get("obj_size") or DEFAULT_OBJ_SIZE
+    kv_dtype_str = kwargs.get("kv_dtype") or DEFAULT_KV_DTYPE_STR
+    kv_dtype = parse_kv_dtype(kv_dtype_str)
+    if kv_dtype is None:
+        print("Error: unsupported --kv-dtype '%s'" % kv_dtype_str)
+        return
+
+    # Build adapter config via the standard parser
+    ns = argparse.Namespace(l2_adapter=l2_adapter_raw)
+    l2_cfg = parse_args_to_l2_adapters_config(ns)
+    if not l2_cfg.adapters:
+        print("Error: no L2 adapter configs parsed")
+        return
+
+    num_tests = kwargs.get("num_keys", 5)
+    settle_time = kwargs.get("settle_time", 0.0)
+
+    for idx, adapter_cfg in enumerate(l2_cfg.adapters):
+        adapter = create_l2_adapter(adapter_cfg)
+        print("=== Testing L2 adapter #%d (%s) ===" % (idx, type(adapter).__name__))
+
+        try:
+            _test_single_adapter(
+                adapter,
+                model,
+                num_tests,
+                obj_size=obj_size,
+                kv_dtype=kv_dtype,
+                settle_time=settle_time,
+            )
+        except Exception as e:
+            print("  Test Failed - Error: %s" % e)
+        finally:
+            adapter.close()
+
+
+def _test_single_adapter(
+    adapter,
+    model,
+    num_tests,
+    obj_size=DEFAULT_OBJ_SIZE,
+    kv_dtype=torch.float32,
+    settle_time=0.0,
+):
+    """Run all test phases against one adapter."""
+    # -- Prepare test data -----------------------------------
+    exist_keys = [_create_object_key(model, "exist_%d" % i) for i in range(num_tests)]
+    non_exist_keys = [
+        _create_object_key(model, "nonexist_%d" % i) for i in range(num_tests)
+    ]
+    store_objs = [
+        _create_memory_obj(float(i + 1), obj_size, kv_dtype) for i in range(num_tests)
+    ]
+
+    # -- Phase 1: lookup non-existing keys -------------------
+    print("Phase 1: Lookup non-existing keys...")
+    lk_absent_ms, lk_bitmap = _run_lookup_phase(adapter, non_exist_keys)
+    if lk_bitmap is None:
+        print("  FAIL: lookup returned None bitmap")
+        ne_pass = 0
+    else:
+        ne_pass = sum(1 for i in range(num_tests) if not lk_bitmap.test(i))
+    print("  Validation: %d/%d correctly absent" % (ne_pass, num_tests))
+    # Unlock the looked-up keys (contract)
+    adapter.submit_unlock(non_exist_keys)
+
+    # -- Phase 2: store existing keys ------------------------
+    print("Phase 2: Store operations (batch of %d)..." % num_tests)
+    st_ms, st_ok = _run_store_phase(adapter, exist_keys, store_objs)
+    store_pass = num_tests if st_ok else 0
+    print("  Batch store %s (%.2fms)" % ("OK" if st_ok else "FAIL", st_ms or 0))
+
+    if settle_time > 0:
+        print("  Waiting %.1fs for data to settle..." % settle_time)
+        time.sleep(settle_time)
+
+    # -- Phase 3: lookup existing keys -----------------------
+    print("Phase 3: Lookup existing keys...")
+    lk_exist_ms, lk_bitmap = _run_lookup_phase(adapter, exist_keys)
+    if lk_bitmap is None:
+        print("  FAIL: lookup returned None bitmap")
+        exist_pass = 0
+    else:
+        exist_pass = sum(1 for i in range(num_tests) if lk_bitmap.test(i))
+    print("  Validation: %d/%d found" % (exist_pass, num_tests))
+
+    # -- Phase 4: load existing keys -------------------------
+    print("Phase 4: Load operations...")
+    load_buffers = [
+        _create_memory_obj(0.0, obj_size, kv_dtype) for _ in range(num_tests)
+    ]
+    ld_ms, ld_bitmap = _run_load_phase(adapter, exist_keys, load_buffers)
+    load_pass = 0
+    content_pass = 0
+    if ld_bitmap is not None:
+        for i in range(num_tests):
+            if ld_bitmap.test(i):
+                load_pass += 1
+                if torch.equal(
+                    load_buffers[i].tensor,
+                    store_objs[i].tensor,
+                ):
+                    content_pass += 1
+                else:
+                    print("  Key %d: data mismatch" % i)
+    print("  Validation (loaded): %d/%d" % (load_pass, num_tests))
+    print("  Validation (content): %d/%d" % (content_pass, num_tests))
+
+    # Unlock after load
+    adapter.submit_unlock(exist_keys)
+
+    # -- Summary ---------------------------------------------
+    total_bytes = obj_size * store_objs[0].tensor.element_size() * num_tests
+    stats_data = [
+        (
+            "LOOKUP (absent)",
+            {
+                "avg": lk_absent_ms or 0,
+                "max": lk_absent_ms or 0,
+                "min": lk_absent_ms or 0,
+            },
+            [False] * num_tests,
+            ne_pass,
+        ),
+        (
+            "STORE",
+            {
+                "avg": st_ms or 0,
+                "max": st_ms or 0,
+                "min": st_ms or 0,
+            },
+            [True] * store_pass + [False] * (num_tests - store_pass),
+            store_pass,
+        ),
+        (
+            "LOOKUP (exist)",
+            {
+                "avg": lk_exist_ms or 0,
+                "max": lk_exist_ms or 0,
+                "min": lk_exist_ms or 0,
+            },
+            [True] * exist_pass + [False] * (num_tests - exist_pass),
+            exist_pass,
+        ),
+        (
+            "LOAD",
+            {
+                "avg": ld_ms or 0,
+                "max": ld_ms or 0,
+                "min": ld_ms or 0,
+            },
+            [True] * content_pass + [False] * (num_tests - content_pass),
+            content_pass,
+        ),
+    ]
+    print_performance_results(stats_data, obj_bytes=total_bytes)

lmcache/v1/check/check_mode_test_remote.py

@@ -0,0 +1,155 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Test mode implementation for basic checks"""
+
+# Standard
+import asyncio
+
+# First Party
+from lmcache.integration.vllm.utils import lmcache_get_or_create_config
+from lmcache.v1.check import check_mode
+
+# Import shared utilities
+from lmcache.v1.check.utils import (
+    DEFAULT_KV_DTYPE_STR,
+    EventLoopManager,
+    _get_default_metadata,
+    create_test_key,
+    parse_kv_dtype,
+    run_common_test_framework,
+    validate_get_results,
+)
+from lmcache.v1.memory_management import MemoryObj
+
+# Import from lmcache with absolute paths
+from lmcache.v1.storage_backend import RemoteBackend
+from lmcache.v1.storage_backend.connector import InstrumentedRemoteConnector
+from lmcache.v1.storage_backend.local_cpu_backend import LocalCPUBackend
+
+
+async def async_contains_backend(backend, key):
+    """Async wrapper for backend contains method"""
+    return backend.contains(key)
+
+
+async def async_get_backend(backend, key):
+    """Async wrapper for backend get_blocking method"""
+    return backend.get_blocking(key)
+
+
+async def async_submit_put_backend(backend, key, memory_obj):
+    """Async wrapper for backend submit_put_task"""
+    future = backend.submit_put_task(key, memory_obj)
+    # Wait for the future to complete with timeout
+    try:
+        await asyncio.wait_for(asyncio.wrap_future(future), timeout=10.0)
+        return True
+    except asyncio.TimeoutError:
+        print(f"Put task timed out for key: {key}")
+        return False
+
+
+def create_test_memory_obj(
+    backend: RemoteBackend, local_cpu_backend: LocalCPUBackend
+) -> MemoryObj:
+    """Create a test MemoryObj for testing."""
+    if backend.connection is None:
+        raise ValueError("Backend connection is None")
+
+    if isinstance(backend.connection, InstrumentedRemoteConnector):
+        connector = backend.connection.getWrappedConnector()
+    else:
+        connector = backend.connection
+
+    return local_cpu_backend.allocate(
+        connector.meta_shapes, connector.meta_dtypes, connector.meta_fmt
+    )
+
+
+def create_test_data_for_backend(
+    backend,
+    local_cpu_backend,
+    model,
+    num_tests,
+    kv_dtype=None,
+):
+    """Create test data for backend based tests"""
+    # Group 1: Non-existing keys
+    kw = {} if kv_dtype is None else {"kv_dtype": kv_dtype}
+    non_exist_keys = [
+        create_test_key(model, f"non_exist_{i}", **kw) for i in range(num_tests)
+    ]
+
+    # Group 2: Existing keys
+    exist_keys = [create_test_key(model, f"exist_{i}", **kw) for i in range(num_tests)]
+    exist_memories = [
+        create_test_memory_obj(backend, local_cpu_backend) for _ in range(num_tests)
+    ]
+
+    return non_exist_keys, exist_keys, exist_memories, num_tests
+
+
+@check_mode("test_remote")
+async def run_test_mode(model: str, **kwargs):
+    """Run connector test mode"""
+    kv_dtype_str = kwargs.get("kv_dtype") or DEFAULT_KV_DTYPE_STR
+    kv_dtype = parse_kv_dtype(kv_dtype_str)
+    if kv_dtype is None:
+        print("Error: unsupported --kv-dtype '%s'" % kv_dtype_str)
+        return
+
+    obj_size = kwargs.get("obj_size")
+
+    config = lmcache_get_or_create_config()
+    metadata = _get_default_metadata(model, kv_dtype=kv_dtype, obj_size=obj_size)
+
+    # Create and start event loop manager
+    loop_manager = EventLoopManager()
+    loop_manager.start()
+
+    local_cpu_backend = LocalCPUBackend(
+        config=config, metadata=metadata, dst_device="cpu"
+    )
+
+    backend = RemoteBackend(
+        config=config,
+        metadata=metadata,
+        loop=loop_manager.get_loop(),
+        local_cpu_backend=local_cpu_backend,
+        dst_device="cpu",
+    )
+
+    try:
+        # Create test context for the common framework
+        test_context = {
+            "create_test_data_func": create_test_data_for_backend,
+            "async_contains_func": async_contains_backend,
+            "async_put_func": async_submit_put_backend,
+            "async_get_func": async_get_backend,
+            "validate_get_func": validate_get_results,
+            "test_object": backend,
+            "extra_args": [
+                local_cpu_backend,
+            ],
+            "kv_dtype": kv_dtype,
+            "obj_size": obj_size,
+        }
+
+        # Run the common test framework
+        num_tests = kwargs.get("num_keys", 5)
+        settle_time = kwargs.get("settle_time", 0.0)
+        await run_common_test_framework(
+            test_context, model, num_tests=num_tests, settle_time=settle_time
+        )
+
+    except Exception as e:
+        print(f"Test Failed - Error: {e}")
+    finally:
+        # Clean up
+        try:
+            if backend:
+                backend.close()
+        except Exception as e:
+            print(f"Error closing backend: {e}")
+
+        # Stop the event loop
+        loop_manager.stop()