diff --git a/docs/design/v1/multiprocess/modules/chunk_store.md b/docs/design/v1/multiprocess/modules/chunk_store.md new file mode 100644 index 00000000000..fcc1244207b --- /dev/null +++ b/docs/design/v1/multiprocess/modules/chunk_store.md @@ -0,0 +1,37 @@ +# Chunk-store plugin + +The plugin factory in `chunk_store.py` is loaded through `--server-module` from +the pluggable MP server interface. It requires the existing +`LMCacheDrivenTransferModule` and shares its worker registrations and streams. +It owns no additional device resources. The ordinary store RPC is unchanged. + +`store_with_chunk_events(key, instance_id, block_ids, producer_event_handle)` +returns `(terminal_handle, [(chunk_handle, start, end), ...], success)`. +The typed `RequestClient` method and `ChunkStoreService` protobuf register the +contract on the existing ZMQ/gRPC transport; the dynamically loaded module +provides the handler. No legacy integer operation ID is added. + +The plugin validates the complete range and per-kernel-group block counts before +submitting work. It splits raw block IDs at each group's chunk stride and calls +the existing store once per chunk, preserving the full token prefix, salt and +request metadata. Downsampling, null blocks, storage reservation, observability +and commits remain the ordinary transfer module's responsibility. + +Each returned event follows all object-group copies for its chunk. The terminal +handle is the last submitted chunk's event. Completion permits source-buffer +reuse; publication and cache residency follow ordinary store semantics. A failed +chunk stops further submission and retains any available completion handle. +Earlier successful chunks may already be committed; the batch is not atomic. +An empty, valid range succeeds without device work. + +Event handles use the existing backend's retention contract. In particular, the +default backend retains a bounded ring of exports; arbitrary delayed imports +after that ring turns over are not guaranteed. This plugin does not add event +leases or change backend lifetime behavior. Clients must obey the same handle +usage constraints as the ordinary store path. + +Tests cover actual factory loading, opt-in registration, group geometry, +metadata preservation, rejection before submission, partial failure, both RPC +transports and cross-process GPU source reuse followed by exact KV retrieval. +The implementation issues one ordinary store per chunk, so throughput and CPU +submission overhead need measurement before making performance claims. diff --git a/docs/source/mp/chunk_store.rst b/docs/source/mp/chunk_store.rst new file mode 100644 index 00000000000..f3347a982c8 --- /dev/null +++ b/docs/source/mp/chunk_store.rst @@ -0,0 +1,42 @@ +Chunk-store plugin +================== + +The chunk-store plugin exposes per-chunk source-buffer completion through the +existing MP request transport. Enable it on an LMCache-driven server: + +.. code-block:: bash + + python -m lmcache.v1.multiprocess.server \ + --l1-size-gb 1 --l1-init-size-gb 1 --eviction-policy LRU \ + --server-module '{"module_path":"lmcache.v1.multiprocess.modules.chunk_store"}' + +Add ``--transport grpc`` to use gRPC. The module lives in the LMCache package and +uses the pluggable server-module loader; no separate plugin package is required. + +After registering KV caches through ``register_kv_cache``, an explicit RPC caller +can submit a range of complete chunks: + +.. code-block:: python + + terminal, chunks, success = client.store_with_chunk_events( + key, instance_id, block_ids, producer_event_handle + ).result(timeout=10) + pending = [ + (backend.import_event(handle, device), start, end) + for handle, start, end in chunks + ] + for event, start, end in pending: + backend.synchronize_event(event, device) + # Source blocks for [start, end) can now be reused. + +For nonblocking completion checks, use ``backend.query_event(event)``. Ranges are +absolute and end-exclusive. Keep the producer event alive until submitted work +completes. Returned handles follow the ordinary store's backend lifetime +constraints, including its bounded retention of exported events; this plugin +does not provide a lease for arbitrarily delayed imports. + +``success`` means all chunk stores avoided a fatal error, not that every chunk +was cached. A failure stops later chunks; earlier chunks may already be stored. +The terminal handle covers submitted work and is empty if none was submitted. +Source completion does not guarantee publication has finished. Engine adapters +continue to use ordinary store unless the caller explicitly uses this RPC. diff --git a/docs/source/mp/index.rst b/docs/source/mp/index.rst index 01baa8d299b..79af7d32653 100644 --- a/docs/source/mp/index.rst +++ b/docs/source/mp/index.rst @@ -5,6 +5,7 @@ Overview :hidden: request_transport + chunk_store LMCache multiprocess (MP) mode runs LMCache as a **standalone service** that vLLM instances reach through a configurable ZMQ or gRPC request transport. diff --git a/lmcache/v1/multiprocess/modules/chunk_store.py b/lmcache/v1/multiprocess/modules/chunk_store.py new file mode 100644 index 00000000000..60b6f41134d --- /dev/null +++ b/lmcache/v1/multiprocess/modules/chunk_store.py @@ -0,0 +1,143 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Dynamically loaded chunk-store RPC using the existing transfer module.""" + +# Standard +from dataclasses import replace + +# First Party +from lmcache.v1.multiprocess.custom_types import IPCCacheServerKey +from lmcache.v1.multiprocess.engine_context import MPCacheServerContext +from lmcache.v1.multiprocess.modules.lmcache_driven_transfer import ( + LMCacheDrivenTransferModule, +) +from lmcache.v1.multiprocess.request_handler import HandlerType, request_handler +from lmcache.v1.multiprocess.server_module import ServerModuleBuildContext + + +def build_server_modules(ctx: ServerModuleBuildContext) -> "ChunkStoreModule": + """Build the chunk-store plugin against the server's registered KV contexts. + + Args: + ctx: Plugin factory context supplied by ``--server-module`` loading. + + Returns: + A module exposing the chunk-store RPC on the existing request transport. + + Raises: + ValueError: The server has no LMCache-driven transfer module. + """ + for module in ctx.modules: + if isinstance(module, LMCacheDrivenTransferModule): + return ChunkStoreModule(ctx.server_context, module) + raise ValueError("The chunk-store plugin requires LMCache-driven transfer") + + +class ChunkStoreModule: + """Store chunks independently and report their source-buffer completion. + + The plugin shares registrations, transfer streams, storage and event backend + with the ordinary store path. Each chunk uses that path's commit and error + semantics; a later failure does not roll back earlier chunks. + + Args: + ctx: Shared server context providing chunk geometry. + transfer_module: Existing owner of registered KV caches and transfers. + """ + + def __init__( + self, ctx: MPCacheServerContext, transfer_module: LMCacheDrivenTransferModule + ) -> None: + self._ctx = ctx + self._transfer = transfer_module + + @property + def context(self) -> MPCacheServerContext: + """Return the shared server context.""" + return self._ctx + + def report_status(self) -> dict[str, dict[str, bool]]: + """Report that the chunk-store plugin is loaded.""" + return {"chunk_store": {"is_healthy": True}} + + def close(self) -> None: + """Leave shared registrations and streams to their transfer-module owner.""" + return None + + @request_handler(HandlerType.BLOCKING, requires_client_affinity=True) + def store_with_chunk_events( + self, + key: IPCCacheServerKey, + instance_id: int, + block_ids: list[list[int]], + event_ipc_handle: bytes, + ) -> tuple[bytes, list[tuple[bytes, int, int]], bool]: + """Submit a chunk-aligned range and return one completion per chunk. + + Args: + key: Worker key with the full token prefix and an absolute range + consisting of complete LMCache chunks. + instance_id: Worker ID registered through ``register_kv_cache``. + block_ids: Raw block IDs for the range, in kernel-group order. + event_ipc_handle: Producer event ordering reads of source KV blocks. + + Returns: + Last submitted completion handle, ordered ``(handle, start, end)`` + entries, and whether every chunk store avoided a fatal error. + Ranges are absolute and end-exclusive. Events permit source reuse; + they do not guarantee cache residency or completed publication. + Invalid ranges, block counts and unregistered workers return empty + handles and False before submission. An empty valid range succeeds. + A failed chunk stops submission, preserving any earlier completions + and its own completion handle if device work was submitted. + + Raises: + RuntimeError: The registered backend cannot provide IPC events. + + Notes: + Handles follow the ordinary store's event-backend lifetime contract. + No host/device synchronization is added during submission. + """ + chunk_size = self._ctx.chunk_size + if ( + key.worker_id is None + or key.start < 0 + or key.end < key.start + or key.end > len(key.token_ids) + or key.start % chunk_size + or key.end % chunk_size + ): + return b"", [], False + entry = self._transfer.get_and_touch_context_entry(instance_id) + if entry is None: + return b"", [], False + cache = entry.cache_context + strides = [ + cache.calculate_num_blocks(chunk_size, group) + for group in range(cache.kv_layer_groups_manager.num_kernel_groups) + ] + num_chunks = (key.end - key.start) // chunk_size + if len(block_ids) != len(strides) or any( + len(ids) != num_chunks * stride + for ids, stride in zip(block_ids, strides, strict=True) + ): + return b"", [], False + + terminal = b"" + completions: list[tuple[bytes, int, int]] = [] + for chunk, start in enumerate(range(key.start, key.end, chunk_size)): + end = start + chunk_size + handle, success = self._transfer.store( + replace(key, start=start, end=end), + instance_id, + [ + ids[chunk * stride : (chunk + 1) * stride] + for ids, stride in zip(block_ids, strides, strict=True) + ], + event_ipc_handle, + ) + if handle: + terminal = handle + completions.append((handle, start, end)) + if not success: + return terminal, completions, False + return terminal, completions, True diff --git a/lmcache/v1/multiprocess/transport/base.py b/lmcache/v1/multiprocess/transport/base.py index 2ae7758b1fd..10075dfbebf 100644 --- a/lmcache/v1/multiprocess/transport/base.py +++ b/lmcache/v1/multiprocess/transport/base.py @@ -92,6 +92,31 @@ def store( event_ipc_handle: bytes, ) -> MessagingFuture[tuple[bytes, bool]]: ... + @rpc_method + def store_with_chunk_events( + self, + key: IPCCacheServerKey, + instance_id: int, + block_ids: list[list[int]], + event_ipc_handle: bytes, + ) -> MessagingFuture[tuple[bytes, list[tuple[bytes, int, int]], bool]]: + """Submit complete chunks to the configured chunk-store server plugin. + + Args: + key: Worker key describing an absolute, chunk-aligned token range. + instance_id: Registered worker ID. + block_ids: Raw source blocks for the range, in kernel-group order. + event_ipc_handle: Producer event ordering reads of the source KV. + + Returns: + A future resolving to the terminal handle, per-chunk + ``(handle, start, end)`` completions, and success. Completion permits + source reuse, not guaranteed cache residency. Handles follow the + same backend lifetime contract as ``store``. Requires the server + module ``lmcache.v1.multiprocess.modules.chunk_store`` to be loaded. + """ + ... + @rpc_method def retrieve( self, diff --git a/lmcache/v1/multiprocess/transport/grpc_impl/protos/chunk_store_service.proto b/lmcache/v1/multiprocess/transport/grpc_impl/protos/chunk_store_service.proto new file mode 100644 index 00000000000..c144dc094cd --- /dev/null +++ b/lmcache/v1/multiprocess/transport/grpc_impl/protos/chunk_store_service.proto @@ -0,0 +1,22 @@ +// SPDX-License-Identifier: Apache-2.0 +syntax = "proto3"; + +package lmcache.mp; + +import "lmcache_driven_service.proto"; + +message ChunkCompletion { + bytes event_ipc_handle = 1; + int64 start = 2; + int64 end = 3; +} + +message ChunkStoreResponse { + bytes terminal_event_ipc_handle = 1; + repeated ChunkCompletion chunk_events = 2; + bool success = 3; +} + +service ChunkStoreService { + rpc StoreWithChunkEvents(StoreRequest) returns (ChunkStoreResponse); +} diff --git a/tests/v1/multiprocess/test_chunk_store.py b/tests/v1/multiprocess/test_chunk_store.py new file mode 100644 index 00000000000..af7bca4a142 --- /dev/null +++ b/tests/v1/multiprocess/test_chunk_store.py @@ -0,0 +1,241 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Chunk-store plugin loading, chunk boundaries and transport contracts.""" + +# Standard +from dataclasses import replace +from unittest.mock import MagicMock +import socket + +# Third Party +import pytest + +# First Party +from lmcache.v1.multiprocess.config import MPServerConfig +from lmcache.v1.multiprocess.custom_types import IPCCacheServerKey +from lmcache.v1.multiprocess.modules.chunk_store import ChunkStoreModule +from lmcache.v1.multiprocess.modules.lmcache_driven_transfer import ( + LMCacheDrivenTransferModule, +) +from lmcache.v1.multiprocess.request_handler import iter_request_handlers +from lmcache.v1.multiprocess.server_module import ( + ServerModuleSpec, + load_server_modules, +) +from lmcache.v1.multiprocess.transport.factory import RequestClientFactory +from lmcache.v1.multiprocess.transport.server_factory import create_request_server +from tests.v1.multiprocess.transport_test_utils import ( + REQUEST_TRANSPORTS, + RequestTransport, + request_server_config, + request_server_url, +) + +PLUGIN = "lmcache.v1.multiprocess.modules.chunk_store" + + +@pytest.fixture +def key() -> IPCCacheServerKey: + return IPCCacheServerKey.from_token_ids( + "model", + 1, + 0, + list(range(64)), + start=16, + end=64, + request_id="request", + cache_salt="tenant", + request_configs={"blend": True}, + ) + + +@pytest.fixture +def plugin() -> tuple[ChunkStoreModule, MagicMock]: + ctx = MagicMock(chunk_size=16) + transfer = MagicMock(spec=LMCacheDrivenTransferModule) + cache = transfer.get_and_touch_context_entry.return_value.cache_context + cache.kv_layer_groups_manager.num_kernel_groups = 2 + cache.calculate_num_blocks.side_effect = lambda size, group: (2, 1)[group] + transfer.store.side_effect = [ + (b"chunk-0", True), + (b"chunk-1", True), + (b"chunk-2", True), + ] + modules = load_server_modules( + [ServerModuleSpec(PLUGIN)], + server_context=ctx, + mp_config=MPServerConfig(), + coordinator_config=MagicMock(), + built_modules=[transfer], + ) + assert len(modules) == 1 + module = modules[0] + assert isinstance(module, ChunkStoreModule) + assert module.context is ctx + return module, transfer + + +def test_plugin_requires_lmcache_driven_transfer() -> None: + with pytest.raises(ValueError, match="requires LMCache-driven"): + load_server_modules( + [ServerModuleSpec(PLUGIN)], + server_context=MagicMock(), + mp_config=MPServerConfig(), + coordinator_config=MagicMock(), + built_modules=[], + ) + + +def test_plugin_is_opt_in(plugin: tuple[ChunkStoreModule, MagicMock]) -> None: + module, transfer = plugin + assert ( + load_server_modules( + [], + server_context=module.context, + mp_config=MPServerConfig(), + coordinator_config=MagicMock(), + built_modules=[transfer], + ) + == [] + ) + assert "store_with_chunk_events" not in { + handler.operation + for handler in iter_request_handlers(LMCacheDrivenTransferModule) + } + handlers = list(iter_request_handlers(module)) + assert [handler.operation for handler in handlers] == ["store_with_chunk_events"] + assert module.report_status() == {"chunk_store": {"is_healthy": True}} + module.close() + transfer.close.assert_not_called() + + +def test_chunks_preserve_prefix_metadata_and_group_geometry( + plugin: tuple[ChunkStoreModule, MagicMock], + key: IPCCacheServerKey, +) -> None: + module, transfer = plugin + blocks = [[2, 3, 4, 5, 6, 7], [8, 9, 10]] + assert module.store_with_chunk_events(key, 7, blocks, b"producer") == ( + b"chunk-2", + [(b"chunk-0", 16, 32), (b"chunk-1", 32, 48), (b"chunk-2", 48, 64)], + True, + ) + for chunk, call in enumerate(transfer.store.call_args_list): + chunk_key, instance, chunk_blocks, producer = call.args + assert chunk_key == replace(key, start=16 + chunk * 16, end=32 + chunk * 16) + assert chunk_key.request_id == key.request_id + assert chunk_key.request_configs == key.request_configs + assert instance == 7 and producer == b"producer" + assert chunk_blocks == [ + blocks[0][chunk * 2 : chunk * 2 + 2], + blocks[1][chunk : chunk + 1], + ] + # The ordinary transfer may downsample its input lists in place. + chunk_blocks[0].clear() + assert blocks == [[2, 3, 4, 5, 6, 7], [8, 9, 10]] + + +@pytest.mark.parametrize("start,end", [(-16, 32), (32, 16), (0, 80), (1, 32), (0, 17)]) +def test_invalid_range_submits_nothing( + plugin: tuple[ChunkStoreModule, MagicMock], + key: IPCCacheServerKey, + start: int, + end: int, +) -> None: + module, transfer = plugin + assert module.store_with_chunk_events( + replace(key, start=start, end=end), 7, [[1, 2], [3]], b"producer" + ) == (b"", [], False) + transfer.store.assert_not_called() + + +@pytest.mark.parametrize( + "blocks", [[], [[1] * 6], [[1] * 6, [2] * 2], [[1] * 7, [2] * 3]] +) +def test_invalid_blocks_reject_whole_request( + plugin: tuple[ChunkStoreModule, MagicMock], + key: IPCCacheServerKey, + blocks: list[list[int]], +) -> None: + module, transfer = plugin + assert module.store_with_chunk_events(key, 7, blocks, b"producer") == ( + b"", + [], + False, + ) + transfer.store.assert_not_called() + + +def test_empty_range_is_a_noop( + plugin: tuple[ChunkStoreModule, MagicMock], + key: IPCCacheServerKey, +) -> None: + module, transfer = plugin + assert module.store_with_chunk_events( + replace(key, end=key.start), 7, [[], []], b"producer" + ) == (b"", [], True) + transfer.store.assert_not_called() + + +def test_unregistered_worker_submits_nothing( + plugin: tuple[ChunkStoreModule, MagicMock], + key: IPCCacheServerKey, +) -> None: + module, transfer = plugin + transfer.get_and_touch_context_entry.return_value = None + assert module.store_with_chunk_events(key, 7, [[1] * 6, [2] * 3], b"producer") == ( + b"", + [], + False, + ) + transfer.store.assert_not_called() + + +@pytest.mark.parametrize("failed_handle", [b"failed-chunk", b""]) +def test_failed_chunk_preserves_completion_and_stops_submission( + plugin: tuple[ChunkStoreModule, MagicMock], + key: IPCCacheServerKey, + failed_handle: bytes, +) -> None: + module, transfer = plugin + transfer.store.side_effect = [(b"chunk-0", True), (failed_handle, False)] + expected_chunks = [(b"chunk-0", 16, 32)] + if failed_handle: + expected_chunks.append((failed_handle, 32, 48)) + assert module.store_with_chunk_events(key, 7, [[1] * 6, [2] * 3], b"producer") == ( + failed_handle or b"chunk-0", + expected_chunks, + False, + ) + assert transfer.store.call_count == 2 + + +@pytest.mark.parametrize("transport", REQUEST_TRANSPORTS) +def test_loaded_plugin_rpc_round_trip( + plugin: tuple[ChunkStoreModule, MagicMock], + key: IPCCacheServerKey, + transport: RequestTransport, +) -> None: + module, transfer = plugin + with socket.socket() as listener: + listener.bind(("127.0.0.1", 0)) + port = listener.getsockname()[1] + url = request_server_url(transport, port) + server = create_request_server([module], request_server_config(transport, url)) + server.start() + client = RequestClientFactory.create(url) + try: + assert client.store_with_chunk_events( + key, 7, [[1] * 6, [2] * 3], b"producer" + ).result(10) == ( + b"chunk-2", + [(b"chunk-0", 16, 32), (b"chunk-1", 32, 48), (b"chunk-2", 48, 64)], + True, + ) + assert transfer.store.call_count == 3 + transfer.get_and_touch_context_entry.return_value = None + assert client.store_with_chunk_events( + key, 7, [[1] * 6, [2] * 3], b"producer" + ).result(10) == (b"", [], False) + finally: + client.close() + server.close() diff --git a/tests/v1/multiprocess/test_chunk_store_gpu.py b/tests/v1/multiprocess/test_chunk_store_gpu.py new file mode 100644 index 00000000000..8e863b268d7 --- /dev/null +++ b/tests/v1/multiprocess/test_chunk_store_gpu.py @@ -0,0 +1,210 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Cross-process GPU tests for the dynamically loaded chunk-store plugin.""" + +# Standard +from collections.abc import Iterator +from dataclasses import replace +from multiprocessing.synchronize import Event as ProcessEvent +import multiprocessing as mp +import os +import socket +import time + +# Third Party +import pytest +import torch + +# First Party +from lmcache import torch_dev, torch_device_type +from lmcache.utils import EngineType +from lmcache.v1.distributed.config import ( + EvictionConfig, + L1ManagerConfig, + L1MemoryManagerConfig, + StorageManagerConfig, +) +from lmcache.v1.mp_observability.config import DEFAULT_OBSERVABILITY_CONFIG +from lmcache.v1.mp_observability.errors import LMCacheTimeoutError +from lmcache.v1.multiprocess.config import MPServerConfig +from lmcache.v1.multiprocess.custom_types import IPCCacheServerKey +from lmcache.v1.multiprocess.server import run_cache_server +from lmcache.v1.multiprocess.server_module import ServerModuleSpec +from lmcache.v1.multiprocess.transport.base import RequestClient +from lmcache.v1.multiprocess.transport.factory import RequestClientFactory +from lmcache.v1.platform.base.event_ipc import get_event_ipc_backend +from lmcache.v1.platform.base.ipc_wrapper import DeviceIPCWrapper +from tests.v1.multiprocess.transport_test_utils import ( + REQUEST_TRANSPORTS, + RequestTransport, + request_server_url, +) + +pytestmark = [pytest.mark.cuda, pytest.mark.no_shared_allocator] +if not (torch_device_type == "cuda" and torch_dev.is_available()): + pytest.skip("requires CUDA or ROCm", allow_module_level=True) + +# First Party +from lmcache.v1.platform.devices.cuda.ipc_wrapper import CudaIPCWrapper # noqa: E402 + +CHUNK = 16 +TIMEOUT = 120.0 + + +def _serve(transport: RequestTransport, port: int, stop: ProcessEvent) -> None: + runtime = run_cache_server( + mp_config=MPServerConfig( + transport=transport, + host="127.0.0.1", + port=port, + chunk_size=CHUNK, + separate_object_groups=True, + server_modules=[ + ServerModuleSpec("lmcache.v1.multiprocess.modules.chunk_store") + ], + null_block_id=-1, + ), + storage_manager_config=StorageManagerConfig( + l1_manager_config=L1ManagerConfig( + memory_config=L1MemoryManagerConfig( + size_in_bytes=256 * 1024**2, + init_size_in_bytes=64 * 1024**2, + use_lazy=True, + ) + ), + eviction_config=EvictionConfig(eviction_policy="LRU"), + ), + obs_config=DEFAULT_OBSERVABILITY_CONFIG, + return_engine=True, + start_prometheus_http_server=False, + ) + assert runtime is not None + server, engine = runtime + try: + stop.wait(TIMEOUT * 4) + finally: + server.close() + engine.close() + + +@pytest.fixture(params=REQUEST_TRANSPORTS) +def gpu_client(request: pytest.FixtureRequest) -> Iterator[RequestClient]: + transport = request.param + with socket.socket() as listener: + listener.bind(("127.0.0.1", 0)) + port = listener.getsockname()[1] + ctx = mp.get_context("spawn") + stop = ctx.Event() + process = ctx.Process(target=_serve, args=(transport, port, stop)) + process.start() + client = RequestClientFactory.create(request_server_url(transport, port)) + try: + deadline = time.monotonic() + TIMEOUT + while time.monotonic() < deadline: + assert process.is_alive(), f"server exited: {process.exitcode}" + try: + assert client.get_chunk_size().result(0.5) == CHUNK + break + except LMCacheTimeoutError: + continue + else: + pytest.fail("chunk store server did not become ready") + yield client + finally: + client.close() + stop.set() + process.join(20) + if process.is_alive(): + process.terminate() + process.join(10) + if process.is_alive(): + process.kill() + process.join() + + +@pytest.mark.parametrize("start", [0, 16]) +def test_chunk_completion_allows_source_reuse_and_exact_kv_retrieval( + gpu_client: RequestClient, + start: int, +) -> None: + client = gpu_client + backend = get_event_ipc_backend(0) + tensors = [ + torch.rand((2, 32, 8, 4, 64), device="cuda", dtype=torch.bfloat16) + for _ in range(2) + ] + wrappers: list[DeviceIPCWrapper] = [CudaIPCWrapper(tensor) for tensor in tensors] + instance = os.getpid() + client.register_kv_cache( + instance, wrappers, "chunk-test", 1, EngineType.VLLM, {}, [] + ).result(TIMEOUT) + prefix_blocks = start // 8 + expected = [ + torch.cat((tensor[:, 10 : 10 + prefix_blocks], tensor[:, 1:9]), dim=1) + for tensor in tensors + ] + producer = backend.create_event(0) + backend.record_event(producer, torch_dev.current_stream()) + key = IPCCacheServerKey.from_token_ids( + "chunk-test", + 1, + 0, + list(range(start + 64)), + start=start, + end=start + 64, + request_id=f"chunk-{start}", + ) + producer_handle = backend.export_event(producer, 0) + if start: + prefix = client.store( + replace(key, start=0, end=start), + instance, + [list(range(10, 10 + prefix_blocks))], + producer_handle, + ).to_device_future(device=0, event_backend=backend) + assert prefix.result(TIMEOUT) + terminal, chunks, success = client.store_with_chunk_events( + key, instance, [list(range(1, 9))], producer_handle + ).result(TIMEOUT) + assert success + assert [(left, right) for _, left, right in chunks] == [ + (left, left + CHUNK) for left in range(start, start + 64, CHUNK) + ] + assert terminal == chunks[-1][0] + events = [backend.import_event(handle, 0) for handle, _, _ in chunks] + for chunk, event in enumerate(events): + backend.synchronize_event(event, 0) + assert backend.query_event(event) + # Recycle each chunk's source as soon as its own event completes. + for tensor in tensors: + tensor[:, 1 + chunk * 2 : 3 + chunk * 2].zero_() + + lookup = replace(key, start=0, worker_id=None, request_id=f"lookup-{start}") + deadline = time.monotonic() + TIMEOUT + while time.monotonic() < deadline: + client.lookup(lookup, 1).result(TIMEOUT) + hits = client.wait_prefetch_status(lookup.request_id, TIMEOUT).result(TIMEOUT) + if hits == (start + 64) // CHUNK: + break + # Source completion can precede finish_write publication. Release this + # attempt's read locks before retrying the prefix lookup. + client.free_lookup_locks(lookup, 1).result(TIMEOUT) + client.end_session(lookup.request_id).result(TIMEOUT) + time.sleep(0.01) + else: + pytest.fail("stored chunks were not published before the deadline") + consumed = backend.create_event(0) + backend.record_event(consumed, torch_dev.current_stream()) + restored = client.retrieve( + replace(lookup, worker_id=0), + instance, + [list(range(16, 24 + prefix_blocks))], + backend.export_event(consumed, 0), + 0, + ).to_device_future(device=0, event_backend=backend) + assert restored.result(TIMEOUT) + for tensor, original in zip(tensors, expected, strict=True): + torch.testing.assert_close( + tensor[:, 16 : 24 + prefix_blocks], original, rtol=0, atol=0 + ) + client.end_session(lookup.request_id).result(TIMEOUT) + client.unregister_kv_cache(instance).result(TIMEOUT) diff --git a/tests/v1/multiprocess/test_grpc_transport.py b/tests/v1/multiprocess/test_grpc_transport.py index 064301fb09d..2ec39346fae 100644 --- a/tests/v1/multiprocess/test_grpc_transport.py +++ b/tests/v1/multiprocess/test_grpc_transport.py @@ -30,6 +30,7 @@ RegisterEngineDrivenContextResponse, ) from lmcache.v1.multiprocess.modules.blend import BlendModule +from lmcache.v1.multiprocess.modules.chunk_store import ChunkStoreModule from lmcache.v1.multiprocess.modules.engine_driven_transfer import ( EngineDrivenTransferModule, ) @@ -327,6 +328,7 @@ def test_module_annotations_cover_and_match_generated_grpc_methods() -> None: QStoreModule, BlendModule, ServerModuleRouter, + ChunkStoreModule, ) handlers = { registered.operation: registered.handler