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

lmcache/v1/check/utils.py

@@ -0,0 +1,571 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Shared utilities for check modes"""
+
+# Standard
+from typing import Optional
+import asyncio
+import hashlib
+import threading
+import time
+
+# Third Party
+import torch
+
+# First Party
+from lmcache.utils import CacheEngineKey
+
+# Import from lmcache with absolute paths
+from lmcache.v1.memory_management import MemoryFormat, MemoryObj
+from lmcache.v1.metadata import LMCacheMetadata
+from lmcache.v1.storage_backend.remote_backend import RemoteBackend
+from lmcache.v1.storage_backend.storage_manager import StorageManager
+
+DTYPE_MAP = {
+    "float32": torch.float32,
+    "float16": torch.float16,
+    "bfloat16": torch.bfloat16,
+}
+
+DEFAULT_KV_DTYPE_STR = "bfloat16"
+DEFAULT_OBJ_SIZE = 1024
+
+
+def _compute_kv_shape(
+    obj_size: int,
+) -> tuple:
+    """Compute a kv_shape that yields the given obj_size.
+
+    The returned shape is in vllm format:
+    ``(num_layers, 2, num_tokens, num_heads, head_size)``.
+
+    The final KV_2LTD tensor has
+    ``2 * num_layers * num_tokens * (num_heads * head_size)``
+    elements, which equals *obj_size*.
+
+    We fix ``num_layers=1, num_heads=1`` and split the
+    remaining factor between ``num_tokens`` and ``head_size``
+    so that ``num_tokens * head_size = obj_size // 2``.
+    """
+    if obj_size % 2 != 0:
+        raise ValueError("obj_size must be even (got %d)" % obj_size)
+    half = obj_size // 2
+    # (num_layers, kv_dim, num_tokens, num_heads, head_size)
+    return (1, 2, half, 1, 1)
+
+
+def parse_kv_dtype(kv_dtype_str: str) -> Optional[torch.dtype]:
+    """Parse a kv_dtype string to a torch.dtype.
+
+    Returns None if the string is not recognized.
+    """
+    return DTYPE_MAP.get(kv_dtype_str)
+
+
+def _get_default_metadata(
+    model: str,
+    kv_dtype: torch.dtype = torch.bfloat16,
+    obj_size: Optional[int] = None,
+) -> LMCacheMetadata:
+    """Get default metadata for testing.
+
+    When *obj_size* is given the ``kv_shape`` is computed so
+    that the resulting KV_2LTD tensor has exactly *obj_size*
+    elements.  Otherwise a small default shape is used.
+    """
+    if obj_size is not None:
+        kv_shape = _compute_kv_shape(obj_size)
+    else:
+        kv_shape = _compute_kv_shape(DEFAULT_OBJ_SIZE)
+    return LMCacheMetadata(
+        model_name=model,
+        world_size=8,
+        local_world_size=8,
+        worker_id=0,
+        local_worker_id=0,
+        kv_dtype=kv_dtype,
+        kv_shape=kv_shape,
+    )
+
+
+def create_test_key(
+    model: str,
+    key_id: str = "test_key",
+    kv_dtype: torch.dtype = torch.bfloat16,
+) -> CacheEngineKey:
+    """Create a test CacheEngineKey."""
+    return CacheEngineKey(
+        model_name=model,
+        world_size=8,
+        worker_id=0,
+        chunk_hash=int(hashlib.sha256(key_id.encode()).hexdigest(), 16),
+        dtype=kv_dtype,
+    )
+
+
+def create_test_memory_obj_for_storage_manager(
+    storage_manager: StorageManager, metadata: LMCacheMetadata
+) -> Optional[MemoryObj]:
+    """Create a test MemoryObj for testing with StorageManager."""
+    # The metadata.kv_shape is in vllm format:
+    # [num_layers, 2, num_tokens, num_heads, head_size]
+    # For KV_2LTD format, we need shape: [2, num_layers, num_tokens, hidden_dim]
+    # where hidden_dim = num_heads * head_size
+
+    vllm_shape = metadata.kv_shape  # [num_layers, 2, num_tokens, num_heads, head_size]
+    num_layers = vllm_shape[0]  # 8
+    kv_dim = vllm_shape[1]  # 2 (K and V)
+    num_tokens = vllm_shape[2]  # 16
+    num_heads = vllm_shape[3]  # 8
+    head_size = vllm_shape[4]  # 16
+
+    # Convert to KV_2LTD format shape: [2, num_layers, num_tokens, hidden_dim]
+    hidden_dim = num_heads * head_size
+    kv_2ltd_shape = torch.Size([kv_dim, num_layers, num_tokens, hidden_dim])
+
+    memory_obj = storage_manager.allocate(
+        kv_2ltd_shape,
+        metadata.kv_dtype,
+        fmt=MemoryFormat.KV_2LTD,
+        eviction=True,
+        busy_loop=False,
+    )
+    return memory_obj
+
+
+def create_storage_manager_with_config(
+    model: str,
+    kv_dtype: torch.dtype = torch.bfloat16,
+    obj_size: Optional[int] = None,
+):
+    """Create storage manager with default configuration"""
+    # First Party
+    from lmcache.integration.vllm.utils import lmcache_get_or_create_config
+    from lmcache.v1.event_manager import EventManager
+
+    config = lmcache_get_or_create_config()
+    metadata = _get_default_metadata(model, kv_dtype=kv_dtype, obj_size=obj_size)
+
+    # Create event manager
+    event_manager = EventManager()
+
+    # Create storage manager
+    storage_manager = StorageManager(
+        config=config,
+        metadata=metadata,
+        event_manager=event_manager,
+    )
+
+    return storage_manager
+
+
+def find_remote_backend(storage_manager: StorageManager) -> Optional[RemoteBackend]:
+    """Find remote backend from storage manager"""
+    for backend_name, backend in storage_manager.storage_backends.items():
+        if isinstance(backend, RemoteBackend):
+            return backend
+    return None
+
+
+def wait_put_tasks_complete(
+    remote_backend: Optional[RemoteBackend], max_wait_time: float = 5.0
+):
+    """Wait for remote backend put tasks to complete"""
+    if remote_backend is None:
+        return
+
+    check_interval = 0.001
+    elapsed_time = 0.0
+
+    while elapsed_time < max_wait_time:
+        if not remote_backend.put_tasks:
+            break
+        time.sleep(check_interval)
+        elapsed_time += check_interval
+
+    # Log warning if timeout
+    remaining_tasks = len(remote_backend.put_tasks)
+    if remaining_tasks > 0:
+        print(
+            f"Warning: {remaining_tasks} remote put tasks still "
+            f"pending after {max_wait_time}s timeout"
+        )
+
+
+def create_memory_objects_batch(
+    storage_manager: StorageManager, metadata: LMCacheMetadata, batch_size: int
+) -> list[MemoryObj]:
+    """Create a batch of memory objects for reuse"""
+    memory_objs = []
+    for i in range(batch_size):
+        memory_obj = create_test_memory_obj_for_storage_manager(
+            storage_manager, metadata
+        )
+        if memory_obj is not None:
+            memory_obj.ref_count_up()
+            memory_objs.append(memory_obj)
+    return memory_objs
+
+
+async def flow_control_check(
+    remote_backend: Optional[RemoteBackend], concurrency: int, sleep_count: float = 1.0
+) -> float:
+    """Check flow control and wait if necessary"""
+    if remote_backend is None:
+        return sleep_count
+
+    high_watermark = 100 * concurrency
+    low_watermark = 10 * concurrency
+    current_tasks = len(remote_backend.put_tasks)
+
+    while current_tasks > high_watermark:
+        current_tasks = len(remote_backend.put_tasks)
+        if current_tasks > high_watermark:
+            # Too many pending tasks, wait before proceeding
+            sleep_sec = 0.1 * sleep_count
+            current_tasks = len(remote_backend.put_tasks)
+            await asyncio.sleep(sleep_sec)
+            current_tasks_after_sleep = len(remote_backend.put_tasks)
+            if current_tasks_after_sleep > low_watermark:
+                sleep_count *= 2.0
+            elif current_tasks_after_sleep == 0:
+                sleep_count /= 2.0
+            continue
+        if current_tasks <= low_watermark:
+            break
+
+    return sleep_count
+
+
+async def run_perf_test_with_timeout(func, args_list, timeout=30.0):
+    """Common performance test framework with timeout handling"""
+    times = []
+    results = []  # Collect results for each operation
+    for i, args in enumerate(args_list):
+        try:
+            start = time.perf_counter()
+            result = await asyncio.wait_for(func(*args), timeout=timeout)
+            end = time.perf_counter()
+            times.append((end - start) * 1000)
+            results.append(result)
+            print(
+                f"  Test {i + 1}/{len(args_list)} completed in "
+                f"{(end - start) * 1000:.2f}ms"
+            )
+        except asyncio.TimeoutError:
+            print(f"  Test {i + 1}/{len(args_list)} timed out after {timeout}s")
+            times.append(timeout * 1000)
+            results.append(None)
+        except Exception as e:
+            print(f"  Test {i + 1}/{len(args_list)} failed: {e}")
+            times.append(0)
+            results.append(None)
+
+    if times:
+        return {
+            "time_stats": {
+                "avg": sum(times) / len(times),
+                "max": max(times),
+                "min": min(times),
+            },
+            "results": results,
+        }
+    else:
+        return {"time_stats": {"avg": 0, "max": 0, "min": 0}, "results": []}
+
+
+def _format_throughput(avg_ms: float, obj_bytes: int) -> str:
+    """Format throughput as a human-readable string."""
+    if avg_ms <= 0 or obj_bytes <= 0:
+        return "N/A"
+    bps = obj_bytes / (avg_ms / 1000.0)
+    if bps >= 1 << 30:
+        return "%.2f GB/s" % (bps / (1 << 30))
+    if bps >= 1 << 20:
+        return "%.2f MB/s" % (bps / (1 << 20))
+    return "%.2f KB/s" % (bps / (1 << 10))
+
+
+def print_performance_results(
+    stats_data,
+    obj_bytes: int = 0,
+    throughput_ops: Optional[set] = None,
+):
+    """Print performance results in a formatted table.
+
+    Args:
+        stats_data: list of (op, stats, results, pass_count).
+        obj_bytes: size of one object in bytes.  When > 0 a
+            throughput column is shown for operations listed
+            in *throughput_ops*.
+        throughput_ops: set of operation name prefixes that
+            should show throughput (e.g. {"STORE", "LOAD"}).
+            Defaults to common data-transfer operations.
+    """
+    if throughput_ops is None:
+        throughput_ops = {
+            "STORE",
+            "LOAD",
+            "PUT",
+            "GET",
+        }
+    show_tp = obj_bytes > 0
+
+    sep_len = 118 if show_tp else 100
+    tp_hdr = " | %s" % "Throughput".center(14) if show_tp else ""
+    print("\nPerformance Results:")
+    print("-" * sep_len)
+    print(
+        f"| {'Operation':<20} | {'Avg (ms)':>12} "
+        f"| {'Max (ms)':>12} "
+        f"| {'Min (ms)':>12} "
+        f"| {'Pass/All':>10} "
+        f"| {'Pass Rate':>10} |" + tp_hdr
+    )
+    print("-" * sep_len)
+    for op, stats, results, pass_count in stats_data:
+        total = len(results)
+        pass_all = f"{pass_count}/{total}"
+        pass_rate = pass_count / total * 100 if total > 0 else 0
+        tp_col = ""
+        if show_tp:
+            is_data_op = any(op.startswith(p) for p in throughput_ops)
+            if is_data_op:
+                tp_col = " | %s" % _format_throughput(
+                    stats["avg"],
+                    obj_bytes,
+                ).center(14)
+            else:
+                tp_col = " | %s" % "-".center(14)
+
+        print(
+            f"| {op:<20} "
+            f"| {stats['avg']:>12.6f} "
+            f"| {stats['max']:>12.6f} "
+            f"| {stats['min']:>12.6f} "
+            f"| {pass_all:>10} "
+            f"| {pass_rate:>9.1f}% |" + tp_col
+        )
+    print("-" * sep_len)
+
+
+def validate_get_results(get_results, exist_keys, exist_memories, num_tests):
+    """Validate GET operation results and return statistics"""
+    content_valid_count = 0
+    for i, result in enumerate(get_results["results"]):
+        if result is None:
+            print(f"  GET for key {exist_keys[i]} returned None result")
+            continue
+        try:
+            if result.tensor is None:
+                print(f"  GET for key {exist_keys[i]} returned None tensor")
+                continue
+
+            if exist_memories[i].tensor is None:
+                print(f"  Original memory object {i} has None tensor")
+                continue
+
+            # Compare data content
+            data_match = torch.equal(result.tensor, exist_memories[i].tensor)
+
+            if data_match:
+                content_valid_count += 1
+            else:
+                print(f"  GET for key {exist_keys[i]} returned incorrect memory object")
+                print("    Data content mismatch detected")
+
+        except Exception as e:
+            print(f"  Data comparison failed for key {exist_keys[i]}: {e}")
+            # Standard
+            import traceback
+
+            traceback.print_exc()
+
+    # Calculate pass rates
+    not_none_count = sum(1 for r in get_results["results"] if r is not None)
+    content_pass_rate = content_valid_count / num_tests * 100
+    print(f"  Validation (not None): {not_none_count}/{num_tests} passed")
+    print(
+        f"  Validation (content correct): {content_valid_count}/{num_tests}"
+        f" passed ({content_pass_rate:.1f}%)"
+    )
+    return content_valid_count, not_none_count
+
+
+async def run_common_test_framework(
+    test_context,
+    model: str,
+    num_tests: int = 5,
+    settle_time: float = 0.0,
+):
+    """
+    Common test framework for both storage manager and remote backend tests.
+
+    Args:
+        test_context: A dictionary containing test-specific functions and objects:
+            - 'create_test_data_func': Function to create test data
+            - 'async_contains_func': Async function for contains operations
+            - 'async_put_func': Async function for put operations
+            - 'async_get_func': Async function for get operations
+            - 'validate_get_func': Function to validate get results
+            - 'test_object': The main test object (storage_manager or backend)
+            - 'extra_args': Extra arguments for test data creation (optional)
+        model: Model name for testing
+        num_tests: Number of tests to run
+    """
+    print("Testing basic operations...")
+
+    # Create test data using the provided function
+    extra_args = test_context.get("extra_args", [])
+    extra_kwargs = {}
+    if "kv_dtype" in test_context:
+        extra_kwargs["kv_dtype"] = test_context["kv_dtype"]
+    if extra_args:
+        non_exist_keys, exist_keys, exist_memories, num_tests = test_context[
+            "create_test_data_func"
+        ](test_context["test_object"], *extra_args, model, num_tests, **extra_kwargs)
+    else:
+        kv_dtype = test_context.get("kv_dtype")
+        obj_size = test_context.get("obj_size")
+        meta_kw = {}
+        if kv_dtype:
+            meta_kw["kv_dtype"] = kv_dtype
+        if obj_size is not None:
+            meta_kw["obj_size"] = obj_size
+        metadata = _get_default_metadata(model, **meta_kw)
+        non_exist_keys, exist_keys, exist_memories, num_tests = test_context[
+            "create_test_data_func"
+        ](test_context["test_object"], metadata, model, num_tests, **extra_kwargs)
+
+    # Phase 1: exists test (key does not exist)
+    print("Phase 1: Testing exists for non-existing keys...")
+
+    exists_non_exist_res = await run_perf_test_with_timeout(
+        test_context["async_contains_func"],
+        [(test_context["test_object"], key) for key in non_exist_keys],
+    )
+    exists_non_exist_stats = exists_non_exist_res["time_stats"]
+    # Validation: All non-existing keys should return False
+    exists_non_exist_pass_count = sum(
+        1 for r in exists_non_exist_res["results"] if r is False
+    )
+    pass_rate = exists_non_exist_pass_count / len(non_exist_keys) * 100
+    print(
+        f"  Validation: {exists_non_exist_pass_count}/{len(non_exist_keys)} "
+        f"passed ({pass_rate:.1f}%)"
+    )
+
+    # Phase 2: put test (create new key)
+    print("Phase 2: Testing put operations...")
+
+    put_res = await run_perf_test_with_timeout(
+        test_context["async_put_func"],
+        [
+            (test_context["test_object"], exist_keys[i], exist_memories[i])
+            for i in range(num_tests)
+        ],
+    )
+    put_stats = put_res["time_stats"]
+    # Validation: All PUT operations should return True
+    put_pass_count = sum(1 for r in put_res["results"] if r is True)
+    pass_rate = put_pass_count / num_tests * 100
+    print(f"  Validation: {put_pass_count}/{num_tests} passed ({pass_rate:.1f}%)")
+
+    if settle_time > 0:
+        print("  Waiting %.1fs for data to settle..." % settle_time)
+        await asyncio.sleep(settle_time)
+
+    # Phase 3: exists test (key exists)
+    print("Phase 3: Testing exists for existing keys...")
+
+    exists_exist_res = await run_perf_test_with_timeout(
+        test_context["async_contains_func"],
+        [(test_context["test_object"], key) for key in exist_keys],
+    )
+    exists_exist_stats = exists_exist_res["time_stats"]
+    # Validation: All existing keys should return True
+    exists_exist_pass_count = sum(1 for r in exists_exist_res["results"] if r is True)
+    pass_rate = exists_exist_pass_count / num_tests * 100
+    print(
+        f"  Validation: {exists_exist_pass_count}/{num_tests} passed ({pass_rate:.1f}%)"
+    )
+
+    # Phase 4: get test (key exists)
+    print("Phase 4: Testing get operations...")
+
+    get_res = await run_perf_test_with_timeout(
+        test_context["async_get_func"],
+        [(test_context["test_object"], key) for key in exist_keys],
+    )
+    get_stats = get_res["time_stats"]
+    # Validation: Check for non-None results and content correctness
+    content_valid_count, not_none_count = test_context["validate_get_func"](
+        get_res, exist_keys, exist_memories, num_tests
+    )
+    # Use content_valid_count as the pass_count for GET operations
+    get_pass_count = content_valid_count
+
+    stats_data = [
+        (
+            "EXISTS (non-exist)",
+            exists_non_exist_stats,
+            exists_non_exist_res["results"],
+            exists_non_exist_pass_count,
+        ),
+        ("PUT", put_stats, put_res["results"], put_pass_count),
+        (
+            "EXISTS (exist)",
+            exists_exist_stats,
+            exists_exist_res["results"],
+            exists_exist_pass_count,
+        ),
+        ("GET", get_stats, get_res["results"], get_pass_count),
+    ]
+
+    # Compute per-object byte size for throughput display
+    tp_obj_size = test_context.get("obj_size") or DEFAULT_OBJ_SIZE
+    tp_kv_dtype = test_context.get("kv_dtype") or torch.bfloat16
+    obj_bytes = tp_obj_size * torch.tensor([], dtype=tp_kv_dtype).element_size()
+
+    # Use common performance results printing
+    print_performance_results(stats_data, obj_bytes=obj_bytes)
+
+
+class EventLoopManager:
+    """Manages a dedicated event loop in a separate thread"""
+
+    def __init__(self):
+        self.loop = None
+        self.thread = None
+        self._loop_started = threading.Event()
+
+    def start(self):
+        """Start the event loop in a separate thread"""
+        if self.thread is not None and self.thread.is_alive():
+            return
+
+        self.loop = asyncio.new_event_loop()
+        self.thread = threading.Thread(target=self._run_loop, daemon=True)
+        self.thread.start()
+        self._loop_started.wait()
+
+    def _run_loop(self):
+        """Run the event loop"""
+        asyncio.set_event_loop(self.loop)
+        self._loop_started.set()
+        try:
+            self.loop.run_forever()
+        except Exception as e:
+            print(f"Event loop error: {e}")
+        finally:
+            self.loop.close()
+
+    def stop(self):
+        """Stop the event loop and thread"""
+        if self.loop and not self.loop.is_closed():
+            self.loop.call_soon_threadsafe(self.loop.stop)
+        if self.thread and self.thread.is_alive():
+            self.thread.join(timeout=5.0)
+
+    def get_loop(self):
+        """Get the event loop"""
+        return self.loop

lmcache/v1/compute/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+

lmcache/v1/compute/attention/abstract.py

@@ -0,0 +1,39 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from typing import TYPE_CHECKING
+import abc
+
+# Third Party
+import torch
+
+if TYPE_CHECKING:
+    # First Party
+    from lmcache.v1.compute.attention.metadata import LMCAttnMetadata
+
+
+class AttentionInterface(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    def forward_contiguous(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        output: torch.Tensor,
+        attn_metadata: "LMCAttnMetadata",
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        Perform forward pass of the attention mechanism.
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def init_attn_metadata(
+        self,
+        input_ids: torch.Tensor,
+        **kwargs,
+    ) -> "LMCAttnMetadata":
+        """
+        Initialize attention metadata.
+        """
+        raise NotImplementedError