lmcache-cli 0.5.3.dev0__py3-none-any.whl → 0.5.3rc2.dev0__py3-none-any.whl

lmcache/__init__.py

@@ -2,6 +2,7 @@
 
 # Standard
 import sys
+import types
 
 # First Party
 from lmcache.logging import init_logger
@@ -9,7 +10,6 @@ from lmcache.logging import init_logger
 # --------------------------
 # Backend instance & Device detection
 # --------------------------
-from lmcache.v1.platform import get_backend
 from lmcache.v1.platform import torch_dev as torch_dev
 from lmcache.v1.platform import torch_device_type as torch_device_type
 
@@ -23,14 +23,34 @@ logger = init_logger(__name__)
 
 __all__ = ["__version__", "torch_dev", "torch_device_type"]
 
-_ops = get_backend(torch_device_type)
-if _ops is not None:
-    # Override lmcache.c_ops with merged module,
-    # in which:
-    #     python_ops_fallback as base,
-    #     use backend implementation if exists
-    sys.modules["lmcache.c_ops"] = _ops
-else:
+
+# --------------------------
+# Backward-compat ``lmcache.c_ops`` shim
+# --------------------------
+def _install_c_ops_shim() -> None:
+    """Register ``lmcache.c_ops`` as the resolved :class:`DeviceOps` instance.
+
+    Resolves the singleton :class:`DeviceOps` instance for the detected
+    device (which calls :meth:`ensure_native` internally), then registers
+    a PEP 562 shim module that forwards attribute access to it.
+    """
+    # First Party
+    from lmcache.v1.platform import resolve_device_ops
+
+    ops = resolve_device_ops(torch_device_type)
+
+    shim = types.ModuleType("lmcache.c_ops")
+    shim.__getattr__ = lambda name: getattr(ops, name)  # type: ignore[method-assign]
+    shim.__dir__ = lambda: dir(ops)  # type: ignore[method-assign]
+    sys.modules["lmcache.c_ops"] = shim
+    globals()["c_ops"] = shim  # parent attr for IMPORT_FROM bytecode
+
+
+try:
+    _install_c_ops_shim()
+except Exception as exc:
     logger.warning(
-        "No compute backend loaded; CLI-only mode (torch/numba not installed)"
+        "No compute backend loaded; CLI-only mode (torch/numba not installed). "
+        "Reason: %s",
+        exc,
     )

lmcache/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '0.5.3.dev0'
-__version_tuple__ = version_tuple = (0, 5, 3, 'dev0')
+__version__ = version = '0.5.3rc2.dev0'
+__version_tuple__ = version_tuple = (0, 5, 3, 'rc2', 'dev0')
 
-__commit_id__ = commit_id = 'gcd2c0d6a6'
+__commit_id__ = commit_id = 'g0427938a4'

lmcache/cli/main.py

@@ -42,3 +42,7 @@ def main() -> None:
     except Exception:  # noqa: BLE001
         logger.exception("Command failed")
         sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

lmcache/integration/sglang/multi_process_adapter.py

@@ -34,6 +34,7 @@ from lmcache.v1.multiprocess.custom_types import (
     IPCCacheServerKey,
     KVCache,
 )
+from lmcache.v1.multiprocess.futures import MessagingFuture
 from lmcache.v1.multiprocess.mq import MessageQueueClient
 from lmcache.v1.multiprocess.protocol import RequestType
 from lmcache.v1.platform.cuda.ipc_wrapper import CudaIPCWrapper
@@ -62,6 +63,28 @@ def _wrap_sglang_kv_caches(
     return wrapped
 
 
+def _completed_future(result: bool) -> MessagingFuture[bool]:
+    """Return an already-completed future resolving to ``result``.
+
+    Used by :meth:`LMCacheMPConnector.store_kv_async` for the paths that
+    perform no wire send, so every return value is a pollable future and
+    callers never have to special-case ``None``. ``result`` carries the
+    store outcome for that path: ``False`` when the connector is
+    unhealthy (nothing was stored), ``True`` when there was simply no
+    chunk-aligned range to store (a no-op success).
+
+    Args:
+        result: the success value the returned future resolves to.
+
+    Returns:
+        A ``MessagingFuture`` whose ``result()`` is immediately
+        ``result``.
+    """
+    future: MessagingFuture[bool] = MessagingFuture()
+    future.set_result(result)
+    return future
+
+
 @dataclass
 class _PendingLookup:
     """Per-request_id state retained between ``lookup_kv`` and
@@ -409,7 +432,7 @@ class LMCacheMPConnector:
                 event.ipc_handle(),
                 skip_prefix_n_blocks,
             ],
-        ).to_cuda_future(device=self.device)
+        ).to_device_future(device=self.device)
 
     def retrieve_kv(self, load_metadata: LoadMetadata) -> int:
         """Phase 2 of the two-phase load — fires RETRIEVE only.
@@ -491,6 +514,76 @@ class LMCacheMPConnector:
                     self._pending_lookups[request_id].locks_held = False
         return retrieve_token_num - offset
 
+    def store_kv_async(self, store_metadata: StoreMetadata) -> MessagingFuture[bool]:
+        """Submit a STORE and return its completion future without waiting.
+
+        Fires the STORE request for the chunk-aligned prefix of
+        ``store_metadata`` and returns immediately with a future the
+        caller can poll (``query`` / ``wait``) or block on (``result``)
+        at a later, deferred checkpoint. The future resolves to a
+        ``bool`` success flag once the daemon finishes copying the KV
+        slots GPU → warehouse.
+
+        The KV slots referenced by ``store_metadata`` must remain pinned
+        (not evicted or reused) until the returned future reports done;
+        the caller owns that lifetime. Paths that perform no wire send
+        return an already-completed future so callers never special-case
+        ``None``: an unhealthy connector resolves to ``False`` (nothing
+        was stored), and no chunk-aligned range resolves to ``True`` (a
+        no-op success).
+
+        END_SESSION is owned by ``LMCRadixCache.cache_finished_req``
+        (see :meth:`end_session`); it is not fired here.
+
+        Args:
+            store_metadata: tokens, request id, and KV slot indices for
+                the finished request.
+
+        Returns:
+            A future resolving to ``True`` when the store completes
+            successfully (or there was nothing to store), or ``False``
+            on daemon-side failure or an unhealthy connector.
+        """
+        if not self.is_healthy:
+            return _completed_future(False)
+
+        aligned_end = (len(store_metadata.token_ids) // self._lmcache_chunk_size) * (
+            self._lmcache_chunk_size
+        )
+        if aligned_end == 0:
+            return _completed_future(True)
+
+        request_id = store_metadata.request_id
+        block_ids = self._slot_mapping_to_block_ids(
+            store_metadata.kv_indices[:aligned_end]
+        )
+        event = torch_dev.Event(interprocess=True)
+        event.record(torch_dev.current_stream())
+        future = send_lmcache_request(
+            self.mq_client,
+            RequestType.STORE,
+            [
+                self._create_key(
+                    store_metadata.token_ids,
+                    start=0,
+                    end=aligned_end,
+                    request_id=request_id,
+                ),
+                self.instance_id,
+                # STORE takes per-group block IDs (list[list[int]]); SGLang is
+                # non-hybrid, so wrap the flat list as a single group.
+                [block_ids],
+                event.ipc_handle(),
+            ],
+        ).to_cuda_future(device=self.device)
+        # Keep the exporting CUDA event alive until the caller releases the
+        # future. Since we return without blocking, the local ``event`` would
+        # otherwise be garbage-collected immediately, destroying the underlying
+        # CUDA event before the daemon imports its IPC handle and waits on it.
+        # (Dynamic keepalive attribute; the future type doesn't declare it.)
+        future._export_event = event  # type: ignore[attr-defined]
+        return future
+
     def store_kv(self, store_metadata: StoreMetadata) -> None:
         if not self.is_healthy:
             return
@@ -525,7 +618,7 @@ class LMCacheMPConnector:
                     event.ipc_handle(),
                 ],
             )
-            .to_cuda_future(device=self.device)
+            .to_device_future(device=self.device)
             .result(timeout=self._mq_timeout)
         )
         # END_SESSION is owned by ``LMCRadixCache.cache_finished_req`` so

lmcache/integration/sglang/sglang_adapter.py

@@ -286,8 +286,10 @@ class LMCacheLayerwiseConnector(LMCacheConnector):
         if retrieve_token_num <= offset:
             self.lmcache_engine.lookup_unpin(lookup_id)
             logger.info(
-                f"LMCache retrieve skipped: lookup={retrieve_token_num}, "
-                f"offset={offset}, no new tokens to retrieve"
+                "LMCache retrieve skipped: lookup=%d, "
+                "offset=%d, no new tokens to retrieve",
+                retrieve_token_num,
+                offset,
             )
             return 0
 
@@ -310,8 +312,10 @@ class LMCacheLayerwiseConnector(LMCacheConnector):
 
         num_new_tokens = retrieve_token_num - offset
         logger.info(
-            f"LMCache retrieve started: lookup={retrieve_token_num}, "
-            f"offset={offset}, retrieve {num_new_tokens} new tokens"
+            "LMCache retrieve started: lookup=%d, offset=%d, retrieve %d new tokens",
+            retrieve_token_num,
+            offset,
+            num_new_tokens,
         )
 
         return num_new_tokens
@@ -324,9 +328,10 @@ class LMCacheLayerwiseConnector(LMCacheConnector):
         store_mask = torch.ones_like(token_ids, dtype=torch.bool)
 
         logger.info(
-            f"LMCache store_kv started: tokens={len(token_ids)}, "
-            f"num_layers={self.sgl_config.num_hidden_layers}, "
-            f"offset={store_metadata.offset}"
+            "LMCache store_kv started: tokens=%d, num_layers=%d, offset=%d",
+            len(token_ids),
+            self.sgl_config.num_hidden_layers,
+            store_metadata.offset,
         )
 
         lookup_id = str(uuid.uuid4())
@@ -357,20 +362,22 @@ class LMCacheLayerwiseConnector(LMCacheConnector):
                     next(layerwise_storer)
                 except StopIteration:
                     logger.error(
-                        (
-                            f"store_layer generator stopped at layer {layer_idx}/"
-                            f"{self.sgl_config.num_hidden_layers}"
-                        )
+                        "store_layer generator stopped at layer %d/%d",
+                        layer_idx,
+                        self.sgl_config.num_hidden_layers,
                     )
                     break
 
             self.lmcache_engine.lookup_unpin(lookup_id)
-            logger.info(f"LMCache store_kv completed: stored {len(token_ids)} tokens")
+            logger.info("LMCache store_kv completed: stored %d tokens", len(token_ids))
         except Exception as e:
             logger.error(
-                f"LMCache store_kv failed: {type(e).__name__}: {e}", exc_info=True
+                "LMCache store_kv failed: %s: %s",
+                type(e).__name__,
+                e,
+                exc_info=True,
             )
             try:
                 self.lmcache_engine.lookup_unpin(lookup_id)
             except Exception as unpin_err:
-                logger.error(f"Failed to unpin lookup: {unpin_err}", exc_info=True)
+                logger.error("Failed to unpin lookup: %s", unpin_err, exc_info=True)

lmcache/integration/sglang/utils.py

@@ -23,7 +23,7 @@ def lmcache_get_config(config_file: str = "") -> LMCacheEngineConfig:
         A validated ``LMCacheEngineConfig``.
     """
     if config_file:
-        logger.info(f"Loading LMCache config file {config_file}")
+        logger.info("Loading LMCache config file %s", config_file)
         config = LMCacheEngineConfig.from_file(config_file)
     else:
         config = LMCacheEngineConfig.from_defaults()

lmcache/integration/vllm/experimental/__init__.py

@@ -0,0 +1,12 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Connector-side experimental tensor-transfer features and their dispatcher."""
+
+# First Party
+from lmcache.integration.vllm.experimental.dispatcher import (
+    Dispatcher,
+    FeatureContext,
+    dispatch,
+    init_dispatcher,
+)
+
+__all__ = ["Dispatcher", "FeatureContext", "dispatch", "init_dispatcher"]

lmcache/integration/vllm/experimental/dispatcher.py

@@ -0,0 +1,213 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Fan-out of connector/adapter lifecycle hooks to active experimental features.
+
+The connector and worker adapter know only this dispatcher; the dispatcher
+knows the features; each feature knows the SDK. Adding a tensor type never
+touches the connector or the adapter -- see ``registry.py``.
+"""
+
+# Future
+from __future__ import annotations
+
+# Standard
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Callable
+
+# First Party
+from lmcache.sdk.cache_kind import LMCacheSDKCacheKind
+from lmcache.sdk.qringbuffer import QRingBufferAdapter, QRingBufferCapture
+from lmcache.utils import init_logger
+from lmcache.v1.multiprocess.modules.experimental import TRANSFER_QUERY
+
+if TYPE_CHECKING:
+    # Third Party
+    from vllm.config import VllmConfig
+    from vllm.v1.kv_cache_interface import KVCacheConfig
+    import torch
+
+    # First Party
+    from lmcache.integration.vllm.lmcache_mp_connector import (
+        LMCacheMPConnectorMetadata,
+    )
+    from lmcache.integration.vllm.vllm_multi_process_adapter import (
+        LMCacheMPWorkerAdapter,
+        _IpcEvent,
+    )
+
+logger = init_logger(__name__)
+
+
+class QTensorFeature:
+    """Captures per-layer query tensors into a paged-Q ring and stores them."""
+
+    def __init__(self, ctx: FeatureContext) -> None:
+        self._q_ring_adapter = QRingBufferAdapter(
+            adapter=ctx.worker_adapter,
+            q_model_name=LMCacheSDKCacheKind.QUERY.server_model_name(
+                ctx.worker_adapter.model_name
+            ),
+            send_lmcache_request=ctx.send_lmcache_request,
+        )
+        self._capture = QRingBufferCapture(ctx.worker_adapter, self._q_ring_adapter)
+
+    def register(
+        self,
+        kv_caches: dict[str, torch.Tensor],
+        kv_cache_config: KVCacheConfig | None,
+        vllm_config: VllmConfig,
+    ) -> None:
+        self._capture.setup_q_ring(kv_caches, kv_cache_config, vllm_config)
+
+    def save_kv_layer(
+        self,
+        layer_name: str,
+        metadata: LMCacheMPConnectorMetadata,
+        **layer_io: Any,
+    ) -> None:
+        self._capture.save_q_layer(layer_name, metadata, **layer_io)
+
+    def wait_for_save(self, event: _IpcEvent | None) -> None:
+        self._capture.batched_submit_qstore_requests(event=event)
+
+    def reclaim(self) -> None:
+        self._q_ring_adapter.reclaim_finished_q_stores()
+
+    def reregister(self) -> bool:
+        try:
+            self._q_ring_adapter.reregister_q_ring()
+        except ConnectionError:
+            logger.exception(
+                "Failed to re-register experimental tensors after server "
+                "recovery; will retry on next heartbeat"
+            )
+            return False
+        except Exception:
+            logger.exception(
+                "Unexpected error during experimental tensor "
+                "re-registration; will retry on next heartbeat"
+            )
+            return False
+        logger.warning(
+            "Finished re-registering experimental tensors after server recovery"
+        )
+        return True
+
+    def shutdown(self) -> None:
+        self._q_ring_adapter.shutdown_q_ring()
+
+
+@dataclass
+class FeatureContext:
+    """
+    Collection of arguments needed for each feature's initialization.
+    """
+
+    worker_adapter: LMCacheMPWorkerAdapter
+    send_lmcache_request: Callable[..., Any]
+
+
+FeatureFactory = Callable[[FeatureContext], "QTensorFeature"]
+
+FEATURE_REGISTRY: dict[str, FeatureFactory] = {
+    TRANSFER_QUERY: QTensorFeature,
+}
+
+
+def init_dispatcher(
+    ctx: FeatureContext,
+    requested: set[str],
+) -> Dispatcher:
+    """Build the dispatcher for intermediate tensors both the server and
+    connector agree to transfer.
+
+    Args:
+        ctx: The feature context.
+        requested: The set of experimental features from vllm_config.
+
+    Returns:
+        A dispatcher for all experimental features built in the server.
+
+    Raises:
+        ValueError: If a requested capability is not advertised by the server.
+    """
+    experimental = ctx.worker_adapter.experimental
+    features: list[QTensorFeature] = []
+    for name, factory in FEATURE_REGISTRY.items():
+        if name not in requested:
+            continue
+        if name not in experimental:
+            raise ValueError(f"Connector enables {name} but server does not.")
+        features.append(factory(ctx))
+    return Dispatcher(features)
+
+
+def dispatch(
+    dispatcher: Dispatcher | None,
+    fn_name: str,
+    **kwargs: Any,
+) -> None:
+    """Dispatch the given arguments to the dispatcher.
+
+    Args:
+        dispatcher: The dispatcher to dispatch to.
+        **kwargs: The arguments to dispatch.
+    """
+    if dispatcher is None:
+        return
+    getattr(dispatcher, fn_name)(**kwargs)
+
+
+class Dispatcher:
+    """Holds the experimental features built into the server and connector, and
+    fans out lifecycle hooks to all of them. The connector and adapter doesn't
+    need to distinguish the types of intermediate tensors. Dispatcher holds the
+    TensorFeature which talks to the SDK.
+
+    Args:
+        features: The active features, in fan-out order.
+    """
+
+    def __init__(self, features: list[QTensorFeature]) -> None:
+        self._features = features
+
+    def register(
+        self,
+        kv_caches: dict[str, torch.Tensor],
+        kv_cache_config: KVCacheConfig | None,
+        vllm_config: VllmConfig,
+    ) -> None:
+        """Fan out ``register`` (connector ``register_kv_caches``)."""
+        for feature in self._features:
+            feature.register(kv_caches, kv_cache_config, vllm_config)
+
+    def save_kv_layer(
+        self,
+        layer_name: str,
+        metadata: LMCacheMPConnectorMetadata,
+        **layer_io: Any,
+    ) -> None:
+        """Fan out ``save_kv_layer`` (connector ``save_kv_layer``; hot path)."""
+        for feature in self._features:
+            feature.save_kv_layer(layer_name, metadata, **layer_io)
+
+    def wait_for_save(self, event: _IpcEvent | None) -> None:
+        """Fan out ``wait_for_save`` (connector ``wait_for_save``)."""
+        for feature in self._features:
+            feature.wait_for_save(event)
+
+    def reclaim(self) -> None:
+        """Fan out ``reclaim`` (adapter ``get_finished``)."""
+        for feature in self._features:
+            feature.reclaim()
+
+    def reregister(self) -> bool:
+        """Fan out ``reregister`` (adapter heartbeat recovery)."""
+        for feature in self._features:
+            if not feature.reregister():
+                return False
+        return True
+
+    def shutdown(self) -> None:
+        """Fan out ``shutdown`` (adapter ``shutdown``)."""
+        for feature in self._features:
+            feature.shutdown()