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

lmcache/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '0.5.6.dev0'
-__version_tuple__ = version_tuple = (0, 5, 6, 'dev0')
+__version__ = version = '0.5.6rc2.dev0'
+__version_tuple__ = version_tuple = (0, 5, 6, 'rc2', 'dev0')
 
-__commit_id__ = commit_id = 'g05a013b29'
+__commit_id__ = commit_id = 'g9a7b60501'

lmcache/cli/commands/bench/server_bench/client.py

@@ -6,11 +6,9 @@ from __future__ import annotations
 
 # Standard
 from collections.abc import Callable
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any
 import hashlib
-import mmap
-import os
 import time
 
 # First Party
@@ -25,8 +23,8 @@ if TYPE_CHECKING:
     import torch
 
     # First Party
+    from lmcache.v1.multiprocess.transfer_context import TransferContext
     from lmcache.v1.multiprocess.transport.base import RequestClient
-    from lmcache.v1.platform.base.event_ipc import EventIPCBackend
 
 
 @dataclass(frozen=True)
@@ -138,10 +136,8 @@ class WorkerContext:
     """Live resources for one simulated Worker."""
 
     spec: WorkerSpec
-    kv_tensors: "list[torch.Tensor]"
-    ipc_wrappers: list[Any]
-    server_pool: "mmap.mmap | None" = None
-    shm_mappings: list[tuple[int, int]] = field(default_factory=list)
+    kv_caches: "dict[str, torch.Tensor]"
+    transfer_context: "TransferContext"
 
 
 class ServerBenchClient:
@@ -162,14 +158,12 @@ class ServerBenchClient:
         self._log = log
         self._zmq_context: Any | None = None
         self._req_client: "RequestClient | None" = None
-        self._event_backend: "EventIPCBackend | None" = None
         self._workers: list[WorkerContext] = []
-        self._registered_instance_ids: list[int] = []
-        self._shm_names: list[str] = []
         self._started = False
 
         self._chunk_size = 0
         self._block_size = 0
+        self._blocks_in_chunk = 0
         self._num_blocks = 0
         self._num_engine_group_infos = 0
         self._kv_world_size = 0
@@ -321,15 +315,16 @@ class ServerBenchClient:
             RuntimeError: If the client is not started.
             ValueError: If a non-empty range is invalid or not chunk-aligned.
         """
-        req_client = self._require_started()
+        self._require_started()
         if token_count == 0:
             return None
         self._validate_token_range(request, start_token, token_count)
 
         # First Party
         from lmcache.cli.commands.bench.server_bench.helpers import (
+            _TIMEOUT,
             _make_key,
-            _send_store,
+            _wait_for_result,
         )
 
         attempted: list[int] = []
@@ -338,6 +333,8 @@ class ServerBenchClient:
         status = "stored"
         started_at = time.monotonic()
         block_offset = request.block_offset + (start_token // self._block_size)
+        num_blocks = token_count // self._block_size
+        block_ids = list(range(block_offset, block_offset + num_blocks))
         for worker in self._workers:
             if not worker.spec.store_enabled:
                 continue
@@ -350,20 +347,25 @@ class ServerBenchClient:
                 worker_id=worker.spec.kv_worker_id,
                 world_size=self._kv_world_size,
             )
-            worker_status = _send_store(
-                req_client,
+            block_ids_per_group = [
+                block_ids.copy() for _ in range(self._num_engine_group_infos)
+            ]
+            event = worker.transfer_context.create_recorded_event()
+            future = worker.transfer_context.submit_store(
+                request.request_id,
                 key,
-                block_offset=block_offset,
-                block_size=self._block_size,
-                num_engine_group_infos=self._num_engine_group_infos,
-                use_gpu=self._config.is_gpu,
-                use_handle=self._config.uses_handle_transfer,
-                client_tensors=self._data_tensors(worker),
-                chunk_size=self._chunk_size,
-                server_pool=worker.server_pool,
-                instance_id=worker.spec.instance_id,
-                event_backend=self._event_backend,
+                worker.kv_caches,
+                block_ids_per_group,
+                event,
+                self._blocks_in_chunk,
             )
+            reply = _wait_for_result(
+                future, retain_refs=(event,) if event is not None else ()
+            )
+            if reply is _TIMEOUT:
+                worker_status = "timeout"
+            else:
+                worker_status = "stored" if reply else "store_failed"
             if worker_status == "stored":
                 successful.append(worker.spec.rank)
             else:
@@ -412,15 +414,16 @@ class ServerBenchClient:
             RuntimeError: If the client is not started.
             ValueError: If a non-empty range is invalid or not chunk-aligned.
         """
-        req_client = self._require_started()
+        self._require_started()
         if token_count == 0:
             return None
         self._validate_token_range(request, start_token, token_count)
 
         # First Party
         from lmcache.cli.commands.bench.server_bench.helpers import (
+            _TIMEOUT,
             _make_key,
-            _send_retrieve,
+            _wait_for_result,
         )
 
         attempted: list[int] = []
@@ -429,7 +432,8 @@ class ServerBenchClient:
         status = "retrieved"
         started_at = time.monotonic()
         block_offset = request.block_offset + (start_token // self._block_size)
-        hit_chunks = token_count // self._chunk_size
+        num_blocks = token_count // self._block_size
+        block_ids = list(range(block_offset, block_offset + num_blocks))
         for worker in self._workers:
             if not worker.spec.retrieve_enabled:
                 continue
@@ -442,21 +446,26 @@ class ServerBenchClient:
                 worker_id=worker.spec.kv_worker_id,
                 world_size=self._kv_world_size,
             )
-            worker_status = _send_retrieve(
-                req_client,
+            block_ids_per_group = [
+                block_ids.copy() for _ in range(self._num_engine_group_infos)
+            ]
+            event = worker.transfer_context.create_recorded_event()
+            future = worker.transfer_context.submit_retrieve(
+                request.request_id,
                 key,
-                self._chunk_size,
-                hit_chunks,
-                block_offset=block_offset,
-                block_size=self._block_size,
-                num_engine_group_infos=self._num_engine_group_infos,
-                use_gpu=self._config.is_gpu,
-                use_handle=self._config.uses_handle_transfer,
-                client_tensors=self._data_tensors(worker),
-                server_pool=worker.server_pool,
-                instance_id=worker.spec.instance_id,
-                event_backend=self._event_backend,
+                worker.kv_caches,
+                block_ids_per_group,
+                event,
+                self._blocks_in_chunk,
+                skip_first_n_tokens=0,
+            )
+            reply = _wait_for_result(
+                future, retain_refs=(event,) if event is not None else ()
             )
+            if reply is _TIMEOUT:
+                worker_status = "timeout"
+            else:
+                worker_status = "retrieved" if reply else "retrieve_failed"
             if worker_status == "retrieved":
                 successful.append(worker.spec.rank)
             else:
@@ -519,7 +528,7 @@ class ServerBenchClient:
         for worker in self._workers:
             if worker.spec.retrieve_enabled:
                 _zero_fill_client_blocks(
-                    worker.kv_tensors,
+                    list(worker.kv_caches.values()),
                     block_offset,
                     num_blocks,
                 )
@@ -578,7 +587,7 @@ class ServerBenchClient:
                     continue
                 parts.extend(
                     _compute_client_checksums(
-                        worker.kv_tensors,
+                        list(worker.kv_caches.values()),
                         block_offset,
                         num_blocks,
                         self._block_size,
@@ -613,49 +622,41 @@ class ServerBenchClient:
         """Idempotently release resources, including after partial startup."""
         # First Party
         from lmcache.cli.commands.bench.server_bench.helpers import (
-            _send_unregister_kv_cache,
+            _DEFAULT_RPC_TIMEOUT_S,
         )
 
         if self._req_client is not None:
-            for instance_id in self._registered_instance_ids:
+            for worker in self._workers:
+                instance_id = worker.spec.instance_id
                 try:
-                    ok = _send_unregister_kv_cache(
-                        self._req_client,
-                        instance_id=instance_id,
-                        use_handle=self._config.uses_handle_transfer,
-                    )
+                    worker.transfer_context.flush_inflight_stores()
+                except Exception as exc:
                     self._log(
-                        "[iid %d] UNREGISTER_KV_CACHE: %s"
-                        % (instance_id, "OK" if ok else "FAIL")
+                        "  [warning] TransferContext flush failed for iid %d: %s"
+                        % (instance_id, exc)
                     )
+                try:
+                    future = worker.transfer_context.unregister()
+                    if future is not None:
+                        future.result(timeout=_DEFAULT_RPC_TIMEOUT_S)
+                        self._log("[iid %d] UNREGISTER_KV_CACHE: OK" % instance_id)
                 except Exception as exc:
                     self._log(
                         "  [warning] UNREGISTER_KV_CACHE failed for "
                         "iid %d: %s" % (instance_id, exc)
                     )
-
         for worker in self._workers:
-            if worker.server_pool is None:
-                continue
             try:
-                worker.server_pool.close()
-            except (BufferError, ValueError):
-                pass
+                worker.transfer_context.close()
+            except Exception as exc:
+                self._log(
+                    "  [warning] TransferContext close failed for iid %d: %s"
+                    % (worker.spec.instance_id, exc)
+                )
 
-        # Drop buffer owners before unmapping SHM.
+        # Drop tensors after contexts release their local transfer resources.
         for worker in self._workers:
-            worker.kv_tensors.clear()
-            worker.ipc_wrappers.clear()
-        if any(worker.shm_mappings for worker in self._workers):
-            # First Party
-            from lmcache.v1.platform.cpu.shm import shm_munmap
-
-            for worker in self._workers:
-                for address, size in worker.shm_mappings:
-                    try:
-                        shm_munmap(address, size)
-                    except OSError:
-                        pass
+            worker.kv_caches.clear()
 
         if self._req_client is not None:
             try:
@@ -672,20 +673,7 @@ class ServerBenchClient:
             finally:
                 self._zmq_context = None
 
-        if self._shm_names:
-            # First Party
-            from lmcache.v1.platform.cpu.shm import shm_unlink
-
-            for name in self._shm_names:
-                try:
-                    shm_unlink(name)
-                except OSError:
-                    pass
-
-        self._registered_instance_ids.clear()
         self._workers.clear()
-        self._shm_names.clear()
-        self._event_backend = None
         self._started = False
 
     def _initialize(self) -> None:
@@ -697,22 +685,26 @@ class ServerBenchClient:
 
         # First Party
         from lmcache.cli.commands.bench.server_bench.helpers import (
+            _DEFAULT_RPC_TIMEOUT_S,
             _INSTANCE_ID_BASE,
+            _MODEL_NAME,
             DTYPE_MAP,
-            _allocate_cpu_shm_kv_cache,
-            _allocate_gpu_kv_cache,
+            _allocate_kv_cache,
             _get_chunk_size,
             _is_mla_kv_size,
-            _send_register_kv_cache,
-            shm_open_pool_as_mmap,
         )
+        from lmcache.utils import EngineType
+        from lmcache.v1.gpu_connector.utils import LayoutHints
         from lmcache.v1.kv_layer_groups import (
             format_kvcache_shape_spec,
             parse_kvcache_shape_spec,
         )
         from lmcache.v1.multiprocess.group_view import EngineGroupInfo
+        from lmcache.v1.multiprocess.transfer_context import (
+            EngineDrivenTransferContext,
+            create_transfer_context,
+        )
         from lmcache.v1.multiprocess.transport.factory import RequestClientFactory
-        from lmcache.v1.platform.base.event_ipc import get_event_ipc_backend
 
         config = self._config
         use_gpu = config.is_gpu
@@ -724,17 +716,13 @@ class ServerBenchClient:
                 "  [info] --transfer-mode=lmcache_driven on cpu mode: "
                 "using REGISTER_KV_CACHE + STORE/RETRIEVE over POSIX SHM"
             )
-        if use_gpu and use_handle:
-            device = torch_dev.current_device()
-            event_backend = get_event_ipc_backend(device)
-            event_backend.check_event_support(device)
-            self._event_backend = event_backend
 
         self._log(
             "Connecting to LMCache MP Server at %s (mode=%s) ..."
             % (config.rpc_url, config.mode)
         )
         self._zmq_context = zmq.Context()
+        self._zmq_context.setsockopt(zmq.LINGER, 0)
         self._req_client = RequestClientFactory.create(
             config.rpc_url,
             context=self._zmq_context,
@@ -768,6 +756,7 @@ class ServerBenchClient:
                 "Server chunk_size %d must be a multiple of block_size %d"
                 % (self._chunk_size, self._block_size)
             )
+        self._blocks_in_chunk = self._chunk_size // self._block_size
         if spec_nb and spec_nb != config.num_blocks:
             self._log(
                 "  [info] spec nb=%d overrides --num-blocks=%d"
@@ -800,15 +789,7 @@ class ServerBenchClient:
         else:
             dtype_string = "mixed"
 
-        layout_hints = {
-            "num_layers": num_layers,
-            "num_heads": num_heads_display,
-            "head_size": head_size_display,
-            "num_blocks": self._num_blocks,
-            "block_size": self._block_size,
-            "dtype": dtype_string,
-            "kv_size": kv_size_display,
-        }
+        layout_hints: LayoutHints = {"kv_layout": "NHD"}
         engine_group_infos = [
             EngineGroupInfo(
                 engine_group_id=group_index,
@@ -853,46 +834,36 @@ class ServerBenchClient:
             )
         self._kv_world_size = 1 if use_mla else config.tp_size
 
-        if use_gpu:
-            # First Party
-            from lmcache.v1.platform.kv_wrap import wrap_one_kv_cache
-        else:
-            # First Party
-            from lmcache.v1.platform.cpu.shm import CpuShmTensorWrapper
-
         for rank in range(config.tp_size):
             instance_id = _INSTANCE_ID_BASE + rank
             kv_worker_id = 0 if use_mla else rank
-            if use_gpu:
-                allocated = _allocate_gpu_kv_cache(groups=layer_groups)
-                self._log(
-                    "[rank %d] Allocated %d GPU tensors on %s"
-                    % (rank, len(allocated), allocated[0].device)
+            tensors = _allocate_kv_cache(
+                groups=layer_groups,
+                device=None if use_gpu else "cpu",
+            )
+            kv_caches = {
+                "layer.%d" % layer_index: tensor
+                for layer_index, tensor in enumerate(tensors)
+            }
+            self._log(
+                "[rank %d] Allocated %d KV tensors on %s"
+                % (rank, len(tensors), tensors[0].device)
+            )
+            # CPU bench tensors can coexist with process-global accelerator
+            # support. Keep this workaround local until factory dispatch is
+            # made device-aware.
+            transfer_context: TransferContext
+            if not use_gpu and config.transfer_mode in ("auto", "engine_driven"):
+                transfer_context = EngineDrivenTransferContext(
+                    instance_id, self._req_client
                 )
-                kv_wrappers = [wrap_one_kv_cache(tensor) for tensor in allocated]
-                client_kv_tensors = allocated
-                rank_shm_mappings: list[tuple[int, int]] = []
             else:
-                shm_prefix = CpuShmTensorWrapper.SHM_NAME_PREFIX + "%s_r%d" % (
-                    os.getpid(),
-                    rank,
-                )
-                (
-                    cpu_tensors,
-                    cpu_wrappers,
-                    rank_shm_names,
-                    rank_shm_mappings,
-                ) = _allocate_cpu_shm_kv_cache(
-                    groups=layer_groups,
-                    shm_prefix=shm_prefix,
-                )
-                self._shm_names.extend(rank_shm_names)
-                self._log(
-                    "[rank %d] Allocated %d CPU SHM tensors (prefix=%s)"
-                    % (rank, len(cpu_tensors), shm_prefix)
+                transfer_context = create_transfer_context(
+                    kv_caches,
+                    mode=config.transfer_mode,
+                    instance_id=instance_id,
+                    req_client=self._req_client,
                 )
-                kv_wrappers = list(cpu_wrappers)
-                client_kv_tensors = cpu_tensors
 
             worker = WorkerContext(
                 spec=WorkerSpec(
@@ -903,41 +874,37 @@ class ServerBenchClient:
                     store_enabled=(rank == 0) if use_mla else True,
                     retrieve_enabled=True,
                 ),
-                kv_tensors=client_kv_tensors,
-                ipc_wrappers=kv_wrappers,
-                shm_mappings=rank_shm_mappings,
+                kv_caches=kv_caches,
+                transfer_context=transfer_context,
             )
             self._workers.append(worker)
 
-            register_result = _send_register_kv_cache(
-                self._req_client,
-                instance_id=instance_id,
-                world_size=self._kv_world_size,
-                layout_hints=layout_hints,
-                kv_caches=kv_wrappers if use_handle else None,
-                use_gpu=use_gpu,
-                use_handle=use_handle,
-                engine_group_infos=engine_group_infos,
-                num_physical_slots=self._chunk_size,
-            )
-            self._log(
-                "[rank %d] REGISTER_KV_CACHE: %s"
-                % (rank, "OK" if register_result else "FAIL")
-            )
-            if not register_result:
+            if (
+                isinstance(transfer_context, EngineDrivenTransferContext)
+                and len(layer_groups) > 1
+            ):
+                raise ValueError(
+                    "engine_driven supports only one KV group; "
+                    "use --transfer-mode lmcache_driven for multiple groups"
+                )
+
+            try:
+                transfer_context.register(
+                    kv_caches,
+                    _MODEL_NAME,
+                    self._kv_world_size,
+                    self._blocks_in_chunk,
+                    _DEFAULT_RPC_TIMEOUT_S,
+                    layout_hints=layout_hints,
+                    engine_group_infos=engine_group_infos,
+                    engine_type=EngineType.VLLM,
+                )
+            except TimeoutError:
                 raise RuntimeError(
                     "REGISTER_KV_CACHE failed for rank %d (instance_id=%d)"
                     % (rank, instance_id)
-                )
-            self._registered_instance_ids.append(instance_id)
-
-            rank_server_pool: mmap.mmap | None = None
-            if not use_handle and not isinstance(register_result, bool):
-                shm_name = getattr(register_result, "shm_name", "")
-                pool_size = getattr(register_result, "pool_size", 0)
-                if shm_name and pool_size > 0:
-                    rank_server_pool = shm_open_pool_as_mmap(shm_name, pool_size)
-            worker.server_pool = rank_server_pool
+                ) from None
+            self._log("[rank %d] REGISTER_KV_CACHE: OK" % rank)
 
         self._log("")
 
@@ -960,9 +927,3 @@ class ServerBenchClient:
             raise ValueError("token range exceeds request")
         if start_token % self._chunk_size or token_count % self._chunk_size:
             raise ValueError("token range must be chunk-aligned")
-
-    def _data_tensors(self, worker: WorkerContext) -> "list[torch.Tensor] | None":
-        """Return tensors only for the engine-driven transfer path."""
-        if self._config.uses_handle_transfer:
-            return None
-        return worker.kv_tensors

lmcache/cli/commands/bench/server_bench/command.py

@@ -103,8 +103,8 @@ def add_server_arguments(parser: argparse.ArgumentParser) -> None:
         default="gpu",
         help=(
             "Run mode (default: gpu). In cpu mode the client allocates "
-            "POSIX-SHM-backed KV cache tensors and the server maps the "
-            "same physical pages."
+            "regular CPU KV tensors; the selected TransferContext owns "
+            "any SHM migration or staging."
         ),
     )
     parser.add_argument(