wehub-resource-sync
1803 lines
66 KiB
Python
1803 lines
66 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||
# Standard
|
||
from collections.abc import Sequence
|
||
from dataclasses import dataclass
|
||
from typing import Any, Callable, NoReturn, Protocol
|
||
import enum
|
||
import os
|
||
import threading
|
||
import uuid
|
||
|
||
# Third Party
|
||
import torch
|
||
import zmq
|
||
|
||
# First Party
|
||
from lmcache import torch_dev
|
||
from lmcache.integration.request_telemetry.factory import RequestTelemetryFactory
|
||
from lmcache.integration.vllm.utils import vllm_layout_hints
|
||
from lmcache.utils import _lmcache_nvtx_annotate, init_logger
|
||
from lmcache.v1.multiprocess.custom_types import (
|
||
BlockAllocationRecord,
|
||
IPCCacheServerKey,
|
||
KVCache,
|
||
)
|
||
from lmcache.v1.multiprocess.group_view import (
|
||
EngineGroupInfo,
|
||
expand_engine_block_ids,
|
||
)
|
||
from lmcache.v1.multiprocess.mq import MessageQueueClient, MessagingFuture
|
||
from lmcache.v1.multiprocess.protocol import RequestType, get_response_class
|
||
from lmcache.v1.multiprocess.transfer_context import (
|
||
EngineDrivenTransferContext,
|
||
TransferContext,
|
||
create_transfer_context,
|
||
)
|
||
from lmcache.v1.periodic_thread import PeriodicThread, ThreadLevel, ThreadRunSummary
|
||
from lmcache.v1.platform import _registry as platform_registry
|
||
|
||
logger = init_logger(__name__)
|
||
|
||
|
||
class ExtraConfigDefault(enum.Enum):
|
||
"""Centralized default values for extra_config keys.
|
||
|
||
Each member's *name* is the key used in the extra_config dict,
|
||
and its *value* is the default.
|
||
"""
|
||
|
||
# Timeout (seconds) for blocking MQ requests: initial
|
||
# chunk-size query, KV cache registration/unregistration,
|
||
# and other synchronous operations.
|
||
mq_timeout = 300.0
|
||
# Interval (seconds) between periodic heartbeat pings
|
||
# to the server.
|
||
heartbeat_interval = 10.0
|
||
# Routing mode for ``create_transfer_context``: ``auto`` keeps the
|
||
# historical CUDA -> lmcache_driven / others -> engine_driven dispatch;
|
||
# ``lmcache_driven`` forces the IPC / SHM zero-copy path where the
|
||
# LMCache server pulls data via device handles;
|
||
# ``engine_driven`` forces the worker-side gather/scatter copy path.
|
||
# Mirrors the ``LMCACHE_MP_TRANSFER_MODE`` env var; this extra_config
|
||
# key wins when both are set.
|
||
mp_transfer_mode = "auto"
|
||
|
||
|
||
# Backward-compatible aliases for the legacy `lmcache_mp_connector_0180`
|
||
# entry point, which still passes these as positional/keyword args.
|
||
DEFAULT_MQ_TIMEOUT: float = ExtraConfigDefault.mq_timeout.value
|
||
DEFAULT_HEARTBEAT_INTERVAL: float = ExtraConfigDefault.heartbeat_interval.value
|
||
|
||
_EXTRA_CONFIG_KEY_PREFIX = "lmcache.mp."
|
||
|
||
# Floor (seconds) of the MP server's worker reap timeout. It only covers the
|
||
# default 10 s heartbeat interval (3 x 10 s); the adapter warns at startup
|
||
# when 3 x heartbeat_interval exceeds it (server timeout must be raised too).
|
||
_SERVER_REAP_TIMEOUT_FLOOR_SECONDS: float = 30.0
|
||
|
||
|
||
def _resolve_extra_config(
|
||
extra_config: dict[str, Any] | None,
|
||
) -> dict[str, Any]:
|
||
"""Resolve extra_config against :class:`ExtraConfigDefault`.
|
||
|
||
Keys in *extra_config* are expected to carry the
|
||
``lmcache.mp.`` prefix (e.g. ``lmcache.mp.mq_timeout``).
|
||
The prefix is stripped before matching against
|
||
:class:`ExtraConfigDefault` members.
|
||
|
||
All ``lmcache.mp.*`` entries are logged. Entries whose
|
||
value differs from the default are additionally marked as
|
||
*overridden*.
|
||
|
||
Args:
|
||
extra_config: User-supplied config dict (may be *None*).
|
||
|
||
Returns:
|
||
A dict keyed by :pyattr:`ExtraConfigDefault` member names.
|
||
"""
|
||
stripped: dict[str, Any] = {}
|
||
if extra_config is not None:
|
||
for k, v in extra_config.items():
|
||
if k.startswith(_EXTRA_CONFIG_KEY_PREFIX):
|
||
short = k[len(_EXTRA_CONFIG_KEY_PREFIX) :]
|
||
stripped[short] = v
|
||
|
||
resolved: dict[str, Any] = {}
|
||
for item in ExtraConfigDefault:
|
||
default = item.value
|
||
raw = stripped.get(item.name)
|
||
value = type(default)(raw) if raw is not None else default
|
||
if value != default:
|
||
logger.info(
|
||
"%s%s = %s (overridden, default: %s)",
|
||
_EXTRA_CONFIG_KEY_PREFIX,
|
||
item.name,
|
||
value,
|
||
default,
|
||
)
|
||
else:
|
||
logger.info(
|
||
"%s%s = %s",
|
||
_EXTRA_CONFIG_KEY_PREFIX,
|
||
item.name,
|
||
value,
|
||
)
|
||
resolved[item.name] = value
|
||
return resolved
|
||
|
||
|
||
class _IpcEvent(Protocol):
|
||
def ipc_handle(self) -> Any: ...
|
||
|
||
def wait(self, stream: Any = None) -> None: ...
|
||
|
||
|
||
def wrap_kv_caches(kv_caches: dict[str, torch.Tensor]) -> KVCache:
|
||
# Emit a per-layer (name, shape, dtype) summary so the operator can
|
||
# verify the exact layer set & tensor geometry being shipped to the
|
||
# LMCache server, then the low-noise count of handles being wrapped.
|
||
kept_summary = [
|
||
(name, tuple(tensor.shape), str(tensor.dtype))
|
||
for name, tensor in kv_caches.items()
|
||
]
|
||
logger.debug(
|
||
"KV cache transfer keeping %d layer(s) (name, shape, dtype):\n%s",
|
||
len(kept_summary),
|
||
"\n".join(
|
||
f" [{i}] {name} shape={shape} dtype={dtype}"
|
||
for i, (name, shape, dtype) in enumerate(kept_summary)
|
||
),
|
||
)
|
||
logger.info("Wrapping %d KV cache tensors for IPC", len(kv_caches))
|
||
# Per-iteration resource management: if wrapping the N-th tensor
|
||
# raises, ``shm_unlink`` whatever earlier iterations already
|
||
# registered with POSIX SHM so the named segments do not outlive
|
||
# the failed batch. CUDA wrappers do not own a named segment and
|
||
# are skipped via the duck-typed ``shm_name`` check.
|
||
wrappers: KVCache = []
|
||
try:
|
||
for tensor in kv_caches.values():
|
||
wrappers.append(wrap_one_kv_cache(tensor))
|
||
except BaseException:
|
||
_release_partial_kv_wrappers(wrappers)
|
||
raise
|
||
return wrappers
|
||
|
||
|
||
def _release_partial_kv_wrappers(wrappers: list[Any]) -> None:
|
||
"""Best-effort unlink of SHM segments owned by partially built wrappers.
|
||
|
||
Used by :func:`wrap_kv_caches` to roll back a half-finished batch
|
||
when a later iteration raises. Only POSIX-SHM-backed wrappers carry
|
||
a ``shm_name`` attribute, so other wrapper kinds (e.g. CUDA-IPC)
|
||
are silently skipped.
|
||
"""
|
||
# First Party
|
||
from lmcache.v1.multiprocess.posix_shm import shm_unlink
|
||
|
||
for w in wrappers:
|
||
name = getattr(w, "shm_name", None)
|
||
if name is None:
|
||
continue
|
||
try:
|
||
shm_unlink(name)
|
||
except Exception: # pragma: no cover - best effort
|
||
logger.debug("shm_unlink failed during rollback", exc_info=True)
|
||
|
||
|
||
def wrap_one_kv_cache(tensor: torch.Tensor) -> Any:
|
||
"""Dispatch by ``tensor.device.type`` via the platform registry.
|
||
|
||
Concrete factories are auto-discovered from
|
||
``DeviceIPCWrapper`` subclasses under ``lmcache.v1.platform``, so
|
||
this call site stays free of if/elif chains and new accelerators
|
||
plug in by shipping a sibling wrapper class.
|
||
"""
|
||
return platform_registry.get_kv_wrapper_factory(tensor.device.type)(tensor)
|
||
|
||
|
||
def send_lmcache_request(
|
||
mq_client: MessageQueueClient,
|
||
request_type: RequestType,
|
||
payloads: list[Any],
|
||
) -> MessagingFuture[Any]:
|
||
"""
|
||
Helper function to send the request to the LMCache multiprocess server
|
||
|
||
Args:
|
||
mq_client: The LMCache multiprocess mode message queue client
|
||
request_type: The request type
|
||
payloads: The request payloads
|
||
|
||
Returns:
|
||
A messaging future for the request
|
||
"""
|
||
|
||
future = mq_client.submit_request(
|
||
request_type, payloads, get_response_class(request_type)
|
||
)
|
||
return future
|
||
|
||
|
||
def get_lmcache_chunk_size(
|
||
mq_client: MessageQueueClient,
|
||
timeout: float = DEFAULT_MQ_TIMEOUT,
|
||
) -> int:
|
||
"""
|
||
Helper function to get the LMCache chunk size from the server
|
||
|
||
Args:
|
||
mq_client: The LMCache multiprocess mode message queue client
|
||
timeout: Timeout in seconds for the blocking request.
|
||
|
||
Returns:
|
||
An integer representing the LMCache chunk size
|
||
"""
|
||
future = send_lmcache_request(mq_client, RequestType.GET_CHUNK_SIZE, [])
|
||
lmcache_tokens_per_chunk = future.result(timeout=timeout)
|
||
return lmcache_tokens_per_chunk
|
||
|
||
|
||
def _raise_server_unreachable(server_url: str, timeout: float) -> NoReturn:
|
||
"""Raise a verbose ConnectionError when the LMCache MP server is
|
||
unreachable.
|
||
|
||
The message intentionally spells out the most common cause (the
|
||
standalone ``lmcache server`` process is not running), the URL that
|
||
was being dialed, and the exact command to start it -- so that users
|
||
landing here via ``vllm serve --kv-offloading-backend lmcache`` are
|
||
not left guessing.
|
||
"""
|
||
hint = (
|
||
"Cannot reach the LMCache MP server at "
|
||
f"'{server_url}' within {timeout}s.\n"
|
||
"This usually means the standalone LMCache server is not "
|
||
"running, or it is listening on a different host/port.\n"
|
||
"To start one locally with the default port (5555):\n"
|
||
" lmcache server --l1-size-gb 20 --eviction-policy LRU\n"
|
||
"To target a different host/port, override via "
|
||
"kv_connector_extra_config (lmcache.mp.host / lmcache.mp.port), "
|
||
"e.g.:\n"
|
||
' --kv-transfer-config \'{"kv_connector":'
|
||
'"LMCacheMPConnector","kv_role":"kv_both",'
|
||
'"kv_connector_extra_config":{"lmcache.mp.host":'
|
||
'"tcp://localhost","lmcache.mp.port":5555}}\'\n'
|
||
"See https://docs.lmcache.ai/mp/quickstart.html for details."
|
||
)
|
||
logger.warning(hint)
|
||
raise ConnectionError(hint) from None
|
||
|
||
|
||
def send_ping(
|
||
mq_client: MessageQueueClient,
|
||
timeout: float,
|
||
instance_id: int | None = None,
|
||
) -> bool:
|
||
"""Send a PING request and return the result.
|
||
|
||
Args:
|
||
mq_client: The message queue client.
|
||
timeout: Seconds to wait for the server's response.
|
||
instance_id: The worker's instance ID so the server can refresh its
|
||
liveness, or None for an untracked prober (scheduler adapter).
|
||
|
||
Returns:
|
||
True if server is healthy, False on timeout or error.
|
||
"""
|
||
try:
|
||
future = send_lmcache_request(mq_client, RequestType.PING, [instance_id])
|
||
return future.result(timeout=timeout)
|
||
except TimeoutError:
|
||
return False
|
||
except Exception:
|
||
logger.debug("Ping failed with exception", exc_info=True)
|
||
return False
|
||
|
||
|
||
@dataclass
|
||
class ParallelStrategy:
|
||
use_mla: bool
|
||
"""Whether to use the MLA."""
|
||
|
||
vllm_world_size: int
|
||
"""Number of workers managed by one vLLM scheduler (TP × PP; excludes DP).
|
||
|
||
Mirrors ``vllm.parallel_config.world_size``.
|
||
"""
|
||
|
||
vllm_worker_id: int
|
||
"""This worker's rank within its scheduler group."""
|
||
|
||
tp_size: int
|
||
"""The tensor parallel size."""
|
||
|
||
pp_size: int
|
||
"""The pipeline parallel size."""
|
||
|
||
n_servers: int
|
||
"""Number of LMCache servers backing this deployment"""
|
||
|
||
@property
|
||
def kv_world_size(self) -> int:
|
||
"""Number of pieces a single token chunk's KV cache is split into
|
||
on the LMCache server storage."""
|
||
if self.use_mla:
|
||
# In this PR we do not support PP + TP + MLA in multi-server mode.
|
||
# A precondition check enforces pp_size == 1, so kv_world_size for
|
||
# MLA can be derived as world_size / tp_size.
|
||
return self.vllm_world_size // self.tp_size
|
||
return self.vllm_world_size // self.n_servers
|
||
|
||
@property
|
||
def kv_worker_id(self) -> int:
|
||
"""Index of the piece of a single token chunk's KV cache
|
||
that the current worker is responsible for,
|
||
in ``[0, kv_world_size)``."""
|
||
if self.use_mla:
|
||
return self.vllm_worker_id // self.tp_size
|
||
return self.vllm_worker_id % (self.vllm_world_size // self.n_servers)
|
||
|
||
@property
|
||
def kv_tp_size(self) -> int:
|
||
"""Tensor-parallel size as seen from a single LMCache server."""
|
||
return self.tp_size // self.n_servers
|
||
|
||
@property
|
||
def is_kv_writer(self) -> bool:
|
||
"""Whether this rank is responsible for storing KV."""
|
||
if not self.use_mla:
|
||
return True
|
||
# MLA: only first rank per node is a writer.
|
||
return self.vllm_worker_id % (self.tp_size // self.n_servers) == 0
|
||
|
||
|
||
def _normalize_adapter_init_args(
|
||
vllm_block_size: int,
|
||
parallel_strategy: ParallelStrategy | int,
|
||
legacy_block_size: int | None,
|
||
mq_timeout: float,
|
||
) -> tuple[int, ParallelStrategy, float]:
|
||
"""Normalize adapter constructor args from old and new vLLM connectors.
|
||
|
||
Args:
|
||
vllm_block_size: The vLLM block size for the current connector API, or
|
||
the legacy KV world size when ``parallel_strategy`` is an int.
|
||
parallel_strategy: The current ``ParallelStrategy`` object, or the
|
||
legacy KV worker id from older vLLM MP connectors.
|
||
legacy_block_size: The legacy vLLM block size passed positionally by
|
||
older vLLM MP connectors.
|
||
mq_timeout: Timeout in seconds for synchronous message queue requests.
|
||
|
||
Returns:
|
||
A tuple of normalized ``(vllm_block_size, parallel_strategy,
|
||
mq_timeout)``.
|
||
|
||
Raises:
|
||
TypeError: If the connector argument shape is not supported.
|
||
"""
|
||
if isinstance(parallel_strategy, ParallelStrategy):
|
||
return vllm_block_size, parallel_strategy, mq_timeout
|
||
if not isinstance(parallel_strategy, int) or legacy_block_size is None:
|
||
raise TypeError(
|
||
"parallel_strategy must be ParallelStrategy, or legacy "
|
||
"(kv_world_size, kv_worker_id, block_size) arguments"
|
||
)
|
||
|
||
kv_world_size = int(vllm_block_size)
|
||
kv_worker_id = int(parallel_strategy)
|
||
strategy = ParallelStrategy(
|
||
use_mla=False,
|
||
vllm_world_size=kv_world_size,
|
||
vllm_worker_id=kv_worker_id,
|
||
tp_size=kv_world_size,
|
||
pp_size=1,
|
||
n_servers=1,
|
||
)
|
||
return int(legacy_block_size), strategy, mq_timeout
|
||
|
||
|
||
class HeartbeatThread(PeriodicThread):
|
||
"""Periodically checks server health via PING.
|
||
|
||
Manages a threading.Event that adapters use to gate operations.
|
||
When unhealthy, the adapter enters degraded mode; if the server
|
||
recovers, the adapter automatically resumes normal operation.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
mq_client: MessageQueueClient,
|
||
health_event: threading.Event,
|
||
interval: float = DEFAULT_HEARTBEAT_INTERVAL,
|
||
instance_id: int | None = None,
|
||
):
|
||
"""
|
||
Args:
|
||
mq_client: The message queue client used to send PING requests.
|
||
health_event: A threading.Event shared with the adapter.
|
||
Set when the server is healthy, cleared when unhealthy.
|
||
Adapters check this event to decide whether to proceed
|
||
with operations or enter degraded mode.
|
||
interval: Seconds between heartbeat pings and ping timeout.
|
||
instance_id: The worker's instance ID sent with each PING so the
|
||
server can refresh its liveness, or None for an untracked
|
||
prober (the scheduler adapter).
|
||
"""
|
||
super().__init__(
|
||
name="lmcache-heartbeat",
|
||
interval=interval,
|
||
level=ThreadLevel.CRITICAL,
|
||
)
|
||
self._mq_client = mq_client
|
||
self._health_event = health_event
|
||
self._interval = interval
|
||
self._instance_id = instance_id
|
||
|
||
# Optional callback invoked on the unhealthy->healthy edge,
|
||
# before the health event is set. See register_recover_callback.
|
||
def noop() -> bool:
|
||
return True
|
||
|
||
self._recover_callback: Callable[[], bool] = noop
|
||
|
||
def register_recover_callback(self, callback: Callable[[], bool]) -> None:
|
||
"""Register a callback fired on the unhealthy->healthy transition.
|
||
|
||
The callback runs **before** the health event is set. It must
|
||
return ``True`` on success (event will be set) or ``False`` on
|
||
failure (event will stay cleared, and the next heartbeat will
|
||
invoke the callback again on the next successful PING).
|
||
|
||
The callback function should NEVER raise exceptions.
|
||
|
||
Intended for setup work that must complete before downstream
|
||
callers observe the recovery — for example, re-registering KV
|
||
caches with a server that just restarted.
|
||
|
||
Should be called before :meth:`start`. Only one callback is
|
||
supported; a second call replaces the first.
|
||
|
||
Args:
|
||
callback: Zero-arg callable returning a success bool.
|
||
"""
|
||
self._recover_callback = callback
|
||
|
||
def _execute(self) -> ThreadRunSummary:
|
||
"""Run one heartbeat cycle: ping, recover callback, event update.
|
||
|
||
A cycle that observes a stop request returns without firing the
|
||
callback or touching the event — a straggler success after
|
||
UNREGISTER must not re-register a ghost context.
|
||
"""
|
||
was_healthy = self._health_event.is_set()
|
||
healthy = send_ping(
|
||
self._mq_client, timeout=self._interval, instance_id=self._instance_id
|
||
)
|
||
|
||
if self.stop_requested:
|
||
return ThreadRunSummary(
|
||
success=True,
|
||
message="stop requested; skipping health update",
|
||
)
|
||
|
||
need_trigger_recover = (
|
||
healthy and not was_healthy and self._recover_callback is not None
|
||
)
|
||
|
||
# Try to call recover callback
|
||
if need_trigger_recover:
|
||
logger.warning(
|
||
"LMCache server is healthy again, triggering recovery callback"
|
||
)
|
||
# If the callback fails, it should not become healthy
|
||
healthy = self._recover_callback()
|
||
|
||
if healthy:
|
||
self._health_event.set()
|
||
if not was_healthy:
|
||
logger.warning(
|
||
"LMCache server is healthy again — resuming normal operation"
|
||
)
|
||
else:
|
||
self._health_event.clear()
|
||
if was_healthy:
|
||
logger.warning("LMCache server is unhealthy — entering degraded mode")
|
||
|
||
return ThreadRunSummary(
|
||
success=True,
|
||
message="healthy" if healthy else "unhealthy",
|
||
)
|
||
|
||
|
||
@dataclass
|
||
class LoadStoreOp:
|
||
token_ids: list[int]
|
||
"""Token IDs for the load/store operation"""
|
||
|
||
block_ids: list[list[int]]
|
||
"""Block IDs for the load/store operation, indexed by engine KV cache
|
||
group (one inner list per engine group). Worker submit paths expand
|
||
this to LMCache KV group order before sending requests to the server.
|
||
"""
|
||
|
||
start: int = 0
|
||
"""Start token index"""
|
||
|
||
end: int = 0
|
||
"""End token index"""
|
||
|
||
skip_first_n_tokens: int = 0
|
||
"""Number of tokens to skip writing at the beginning of the retrieve
|
||
range. Used to avoid overwriting APC-shared GPU blocks during retrieve."""
|
||
|
||
@property
|
||
def flat_block_ids(self) -> list[int]:
|
||
"""Return all block IDs flattened for group-blind error paths.
|
||
|
||
Handles both the normal ``list[list[int]]`` format and the
|
||
IPC-flattened ``list[int]`` format that vLLM v0.19.0 produces when
|
||
``SchedulerOutput`` serializes single-element nested lists across
|
||
process boundaries (e.g. ``[[20, 21]]`` → ``[20, 21]``).
|
||
Returns an empty list when ``block_ids`` is empty.
|
||
"""
|
||
if not self.block_ids:
|
||
return []
|
||
# Defend against IPC serialization flattening [[20, 21, …]] → [20, 21, …]
|
||
if isinstance(self.block_ids[0], int):
|
||
return list(self.block_ids)
|
||
return [
|
||
block_id
|
||
for group_block_ids in self.block_ids
|
||
for block_id in group_block_ids
|
||
]
|
||
|
||
|
||
StoreResult = bool
|
||
RetrieveResult = bool
|
||
LookupResult = int
|
||
|
||
|
||
class LMCacheMPSchedulerAdapter:
|
||
def __init__(
|
||
self,
|
||
server_urls: list[str],
|
||
context: zmq.Context,
|
||
model_name: str,
|
||
vllm_block_size: int,
|
||
parallel_strategy: ParallelStrategy | int,
|
||
legacy_block_size: int | None = None,
|
||
*,
|
||
mq_timeout: float = DEFAULT_MQ_TIMEOUT,
|
||
heartbeat_interval: float = DEFAULT_HEARTBEAT_INTERVAL,
|
||
extra_config: dict[str, Any] | None = None,
|
||
):
|
||
"""
|
||
Args:
|
||
server_urls: The servers URL for the LMCache message queue
|
||
context: The ZMQ context
|
||
model_name: The model name used for LMCache keys
|
||
vllm_block_size: The block size used in vLLM
|
||
parallel_strategy:
|
||
The parallel strategy, which includes `use_mla`,
|
||
`kv_world_size`, `kv_worker_id` and so on. Older vLLM
|
||
connectors pass the KV worker id here.
|
||
legacy_block_size: The vLLM block size passed positionally by
|
||
older vLLM connectors.
|
||
mq_timeout: Timeout in seconds for message queue requests.
|
||
Ignored when ``extra_config`` is provided.
|
||
heartbeat_interval: Interval in seconds between heartbeat pings.
|
||
Ignored when ``extra_config`` is provided.
|
||
extra_config: Optional dict with keys starting with
|
||
``lmcache.mp.`` (e.g., ``lmcache.mp.mq_timeout``). When
|
||
provided, it overrides ``mq_timeout`` / ``heartbeat_interval``.
|
||
"""
|
||
vllm_block_size, parallel_strategy, mq_timeout = _normalize_adapter_init_args(
|
||
vllm_block_size,
|
||
parallel_strategy,
|
||
legacy_block_size,
|
||
mq_timeout,
|
||
)
|
||
assert len(server_urls) >= 1, "At least one server url required"
|
||
self._server_urls: list[str] = list(server_urls)
|
||
self.mq_clients: dict[str, MessageQueueClient] = {
|
||
url: MessageQueueClient(url, context) for url in self._server_urls
|
||
}
|
||
if extra_config is not None:
|
||
cfg = _resolve_extra_config(extra_config)
|
||
mq_timeout = cfg[ExtraConfigDefault.mq_timeout.name]
|
||
heartbeat_interval = cfg[ExtraConfigDefault.heartbeat_interval.name]
|
||
self._mq_timeout = mq_timeout
|
||
|
||
# Lookup state tracking:
|
||
# - _pending_lookups: request_ids submitted but not yet resolved
|
||
# - _finished_lookup_results: cached chunk count keyed by request_id,
|
||
# so that repeated calls to check_lookup_result return the same value
|
||
# even after the server has already popped the job (exactly-once).
|
||
# - _per_server_hits: {request_id: {server_url: hit_chunks}}.
|
||
# Per-server hit counts, used to detect disagreement and free tail locks.
|
||
self._pending_lookups: set[str] = set()
|
||
self._finished_lookup_results: dict[str, int] = {}
|
||
self._per_server_hits: dict[str, dict[str, int]] = {}
|
||
self._lookup_params: dict[str, tuple[list[int], str]] = {}
|
||
|
||
self.model_name = model_name
|
||
self.parallel_strategy = parallel_strategy
|
||
|
||
# Read chunk size from lmcache
|
||
chunk_sizes: dict[str, int] = {}
|
||
for url, client in self.mq_clients.items():
|
||
try:
|
||
chunk_sizes[url] = get_lmcache_chunk_size(
|
||
client, timeout=self._mq_timeout
|
||
)
|
||
except TimeoutError:
|
||
for c in self.mq_clients.values():
|
||
c.close()
|
||
_raise_server_unreachable(url, self._mq_timeout)
|
||
|
||
# All servers must share chunk_size, otherwise the min() aggregation
|
||
# over per-server hits would mix different granularities.
|
||
unique_sizes = set(chunk_sizes.values())
|
||
if len(unique_sizes) != 1:
|
||
raise ValueError(
|
||
f"All LMCache servers must share the same chunk_size, got {chunk_sizes}"
|
||
)
|
||
self.lmcache_tokens_per_chunk = unique_sizes.pop()
|
||
|
||
assert self.lmcache_tokens_per_chunk % vllm_block_size == 0, (
|
||
"LMCache chunk size should be a multiple of vLLM block size"
|
||
)
|
||
self.blocks_in_chunk = self.lmcache_tokens_per_chunk // vllm_block_size
|
||
|
||
# Health state: one Event per server. The adapter is considered healthy
|
||
# only if ALL per-server events are set (any unhealthy server taints
|
||
# the whole adapter, matching the min() semantics used for lookups).
|
||
self._health_events: dict[str, threading.Event] = {}
|
||
for url in self._server_urls:
|
||
ev = threading.Event()
|
||
ev.set()
|
||
self._health_events[url] = ev
|
||
|
||
# Heartbeat thread is created but NOT started yet.
|
||
# It will be lazily started on the first lookup
|
||
# request, by which time vLLM is fully ready.
|
||
self._heartbeat_interval = heartbeat_interval
|
||
self._heartbeats: dict[str, HeartbeatThread] = {}
|
||
self._heartbeat_lock = threading.Lock()
|
||
|
||
@property
|
||
def world_size(self) -> int:
|
||
"""Get the kv world size."""
|
||
return self.parallel_strategy.kv_world_size
|
||
|
||
@property
|
||
def tp_size(self) -> int:
|
||
"""The tensor parallel size."""
|
||
return self.parallel_strategy.kv_tp_size
|
||
|
||
@property
|
||
def is_healthy(self) -> bool:
|
||
"""True iff every backing LMCache server is healthy."""
|
||
return all(ev.is_set() for ev in self._health_events.values())
|
||
|
||
def _ensure_heartbeat_started(self) -> None:
|
||
"""Lazily start the heartbeat thread on first use."""
|
||
if self._heartbeats is not None:
|
||
return
|
||
with self._heartbeat_lock:
|
||
if self._heartbeats is not None:
|
||
return
|
||
for url, client in self.mq_clients.items():
|
||
hb = HeartbeatThread(
|
||
mq_client=client,
|
||
health_event=self._health_events[url],
|
||
interval=self._heartbeat_interval,
|
||
)
|
||
hb.start()
|
||
self._heartbeats[url] = hb
|
||
|
||
@_lmcache_nvtx_annotate
|
||
def maybe_submit_lookup_request(
|
||
self,
|
||
request_id: str,
|
||
token_ids: list[int],
|
||
cache_salt: str = "",
|
||
):
|
||
"""
|
||
Submit a new lookup request to LMCache if there is no ongoing request.
|
||
|
||
Sends a LOOKUP request to the server and blocks until a prefetch
|
||
job ID is returned. The actual prefetch result can then be polled
|
||
via ``check_lookup_result``.
|
||
|
||
Args:
|
||
request_id: The ID of the lookup request. The same ID indicates it's
|
||
from the same request
|
||
token_ids: Token IDs to lookup from LMCache
|
||
cache_salt: Per-user isolation salt. Requests with different
|
||
cache_salt values produce separate cache entries.
|
||
|
||
Returns:
|
||
None
|
||
|
||
Notes:
|
||
This function will have a side-effect: submitting a look up request to
|
||
LMCache, which will essentially 'lock' the KV cache chunks in the LMCache
|
||
for later retrieve operations.
|
||
In the meantime, this function will record the lookup request, and the
|
||
status of the look up request can be checked by `check_lookup_result`.
|
||
"""
|
||
self._ensure_heartbeat_started()
|
||
|
||
if not self.is_healthy:
|
||
return
|
||
|
||
if request_id in self._pending_lookups:
|
||
# Skip if there is already a lookup request
|
||
return
|
||
|
||
aligned_end = (
|
||
len(token_ids) // self.lmcache_tokens_per_chunk
|
||
) * self.lmcache_tokens_per_chunk
|
||
|
||
key = self._create_key(
|
||
token_ids,
|
||
start=0,
|
||
end=aligned_end,
|
||
request_id=request_id,
|
||
cache_salt=cache_salt,
|
||
).no_worker_id_version()
|
||
|
||
futures: dict[str, MessagingFuture[Any]] = {
|
||
url: send_lmcache_request(
|
||
self.mq_clients[url],
|
||
RequestType.LOOKUP,
|
||
[key, self.tp_size],
|
||
)
|
||
for url in self._server_urls
|
||
}
|
||
|
||
# Any one server failure means the whole lookup fails.
|
||
for url, fut in futures.items():
|
||
try:
|
||
fut.result(timeout=self._mq_timeout)
|
||
except TimeoutError:
|
||
logger.warning(
|
||
"LOOKUP to %s timed out after %ss. Marking server as unhealthy.",
|
||
url,
|
||
self._mq_timeout,
|
||
)
|
||
self._health_events[url].clear()
|
||
return
|
||
|
||
self._pending_lookups.add(request_id)
|
||
self._lookup_params[request_id] = (token_ids, cache_salt)
|
||
|
||
def _free_inconsistent_lookup_locks(
|
||
self,
|
||
request_id: str,
|
||
per_server: dict[str, int],
|
||
min_chunks: int,
|
||
) -> None:
|
||
"""Release over-hit tail locks on servers that reported more than min.
|
||
|
||
When servers disagree on hit chunk counts, servers reporting more
|
||
than min_chunks have locked chunks in the range
|
||
[min_chunks * tokens_per_chunk, hit_chunks * tokens_per_chunk)
|
||
that will never be retrieved. This method frees those tail locks.
|
||
|
||
Args:
|
||
request_id: The lookup request ID.
|
||
per_server: Per-server hit chunk counts.
|
||
min_chunks: Minimum hit chunk count across all servers.
|
||
"""
|
||
token_ids_l, cs = self._lookup_params.pop(request_id, (None, None))
|
||
if token_ids_l is not None:
|
||
for url, hit_chunks in per_server.items():
|
||
if hit_chunks <= min_chunks:
|
||
continue
|
||
tail_end = min(
|
||
hit_chunks * self.lmcache_tokens_per_chunk, len(token_ids_l)
|
||
)
|
||
tail_key = self._create_key(
|
||
token_ids=token_ids_l,
|
||
start=min_chunks * self.lmcache_tokens_per_chunk,
|
||
end=tail_end,
|
||
request_id=request_id,
|
||
cache_salt=cs or "",
|
||
).no_worker_id_version()
|
||
send_lmcache_request(
|
||
self.mq_clients[url],
|
||
RequestType.FREE_LOOKUP_LOCKS,
|
||
[tail_key, self.tp_size],
|
||
)
|
||
|
||
@_lmcache_nvtx_annotate
|
||
def check_lookup_result(self, request_id: str) -> int | None:
|
||
"""
|
||
Check the result of a previously submitted lookup request.
|
||
|
||
Sends a QUERY_PREFETCH_STATUS request to the servers and blocks
|
||
until the server responds. Returns the matched token count
|
||
when the prefetch is complete, or None if still in progress.
|
||
|
||
Args:
|
||
request_id: The ID of the lookup request submitted in
|
||
`maybe_submit_lookup_request`
|
||
|
||
Returns:
|
||
An integer representing the total number of tokens matched
|
||
in LMCache (prefix matching), or
|
||
None if the lookup request is not finished yet.
|
||
"""
|
||
if request_id not in self._pending_lookups:
|
||
# No job — either unhealthy at submit time or already cleaned up.
|
||
# Return the cached aggregate if any, otherwise 0.
|
||
return self._finished_lookup_results.get(request_id, 0)
|
||
|
||
if not self.is_healthy:
|
||
# Server went down — give up on this lookup
|
||
return 0
|
||
|
||
if request_id in self._finished_lookup_results:
|
||
# Aggregation already done; return the cached value.
|
||
return self._finished_lookup_results[request_id]
|
||
|
||
# Persistent accumulator for this request. A server present in
|
||
# the dict has already handed over its final hit count and must
|
||
# not be polled again; absence means "not yet observed".
|
||
per_server = self._per_server_hits.setdefault(request_id, {})
|
||
unresolved_urls = [u for u in self._server_urls if u not in per_server]
|
||
|
||
futures: dict[str, MessagingFuture[Any]] = {
|
||
url: send_lmcache_request(
|
||
self.mq_clients[url],
|
||
RequestType.QUERY_PREFETCH_STATUS,
|
||
[request_id],
|
||
)
|
||
for url in unresolved_urls
|
||
}
|
||
|
||
for url, fut in futures.items():
|
||
try:
|
||
r = fut.result(timeout=self._mq_timeout)
|
||
except TimeoutError:
|
||
logger.warning(
|
||
"QUERY_PREFETCH_STATUS to %s timed out. Marking unhealthy.",
|
||
url,
|
||
)
|
||
self._health_events[url].clear()
|
||
return 0
|
||
if r is None:
|
||
continue
|
||
per_server[url] = int(r)
|
||
|
||
if len(per_server) < len(self._server_urls):
|
||
return None
|
||
|
||
min_chunks = min(per_server.values())
|
||
max_chunks = max(per_server.values())
|
||
if min_chunks != max_chunks:
|
||
logger.warning(
|
||
"[req=%s] LMCache hit mismatch across servers: %s → take min=%d",
|
||
request_id,
|
||
dict(per_server),
|
||
min_chunks,
|
||
)
|
||
self._free_inconsistent_lookup_locks(request_id, per_server, min_chunks)
|
||
|
||
token_count = min_chunks * self.lmcache_tokens_per_chunk
|
||
|
||
self._finished_lookup_results[request_id] = token_count
|
||
return token_count
|
||
|
||
def num_blocks_per_chunk(self) -> int:
|
||
"""
|
||
Returns:
|
||
The number of vllm blocks in a LMCache data chunk
|
||
"""
|
||
return self.blocks_in_chunk
|
||
|
||
def cleanup_lookup_result(self, request_id: str) -> None:
|
||
"""
|
||
Clean up lookup state for a finished request to prevent memory leak.
|
||
Args:
|
||
request_id: The ID of the finished request.
|
||
"""
|
||
self._pending_lookups.discard(request_id)
|
||
self._finished_lookup_results.pop(request_id, None)
|
||
self._per_server_hits.pop(request_id, None)
|
||
self._lookup_params.pop(request_id, None)
|
||
|
||
def shutdown(self) -> None:
|
||
"""Shutdown the scheduler adapter and its resources."""
|
||
for client in self.mq_clients.values():
|
||
client.close()
|
||
with self._heartbeat_lock:
|
||
for hb in self._heartbeats.values():
|
||
hb.stop()
|
||
|
||
def free_lookup_locks(
|
||
self,
|
||
token_ids: list[int],
|
||
start: int,
|
||
end: int,
|
||
request_id: str,
|
||
cache_salt: str = "",
|
||
) -> None:
|
||
"""Release read locks acquired during lookup without a full retrieve.
|
||
|
||
Use this when some chunks matched by lookup overlap with blocks that
|
||
vLLM has already computed, so they will never be retrieved. Calling
|
||
this prevents those chunks from holding read locks until TTL expiry.
|
||
|
||
Or use this when a request is cancelled or aborted after lookup but
|
||
before retrieve to avoid holding read locks until TTL expiry.
|
||
|
||
When ``start`` or ``end`` is not aligned to the chunk size, the
|
||
entire chunk containing start boundary is freed but not end boundary.
|
||
It is caller's responsibility to properly align the boundaries.
|
||
|
||
Args:
|
||
token_ids: Token IDs for the key (same as used in lookup).
|
||
start: Start token index.
|
||
end: End token index.
|
||
request_id: The request ID.
|
||
cache_salt: Per-user isolation salt.
|
||
"""
|
||
if not self.is_healthy:
|
||
return
|
||
|
||
# Free [start, end) on every server.
|
||
base_key = self._create_key(
|
||
token_ids,
|
||
start=start,
|
||
end=end,
|
||
request_id=request_id,
|
||
cache_salt=cache_salt,
|
||
).no_worker_id_version()
|
||
for url in self._server_urls:
|
||
send_lmcache_request(
|
||
self.mq_clients[url],
|
||
RequestType.FREE_LOOKUP_LOCKS,
|
||
[base_key, self.tp_size],
|
||
)
|
||
|
||
def end_session(self, request_id: str) -> None:
|
||
"""
|
||
Notify LMCache server to remove the session for a finished request.
|
||
Args:
|
||
request_id: The ID of the finished request.
|
||
"""
|
||
if not self.is_healthy:
|
||
return
|
||
|
||
for url in self._server_urls:
|
||
send_lmcache_request(
|
||
self.mq_clients[url],
|
||
RequestType.END_SESSION,
|
||
[request_id],
|
||
)
|
||
|
||
def report_block_allocations(
|
||
self,
|
||
records: list[BlockAllocationRecord],
|
||
) -> None:
|
||
"""Report vLLM GPU block allocation deltas to LMCache server.
|
||
|
||
Fire-and-forget: does not wait for a response. If the server
|
||
is unhealthy the report is silently dropped.
|
||
|
||
Args:
|
||
records: List of BlockAllocationRecord with per-request
|
||
block and token allocation deltas.
|
||
"""
|
||
if not self.is_healthy or not records:
|
||
return
|
||
|
||
for url in self._server_urls:
|
||
send_lmcache_request(
|
||
self.mq_clients[url],
|
||
RequestType.REPORT_BLOCK_ALLOCATION,
|
||
[os.getpid(), self.model_name, records],
|
||
)
|
||
|
||
# Helper functions
|
||
def _create_key(
|
||
self,
|
||
token_ids: list[int],
|
||
start: int,
|
||
end: int,
|
||
request_id: str,
|
||
cache_salt: str = "",
|
||
) -> IPCCacheServerKey:
|
||
"""Convert token IDs to an IPC cache engine key.
|
||
|
||
Args:
|
||
token_ids: The token IDs.
|
||
start: Start token index.
|
||
end: End token index.
|
||
request_id: The request ID.
|
||
cache_salt: Per-user isolation salt.
|
||
|
||
Returns:
|
||
IPCCacheServerKey: The constructed key.
|
||
"""
|
||
# NOTE: for the scheduler adapter, we don't have a worker id,
|
||
# so we set it to None in the key.
|
||
return IPCCacheServerKey(
|
||
model_name=self.model_name,
|
||
world_size=self.world_size,
|
||
worker_id=None,
|
||
token_ids=tuple(token_ids),
|
||
start=start,
|
||
end=end,
|
||
request_id=request_id,
|
||
cache_salt=cache_salt,
|
||
)
|
||
|
||
|
||
class LMCacheMPWorkerAdapter:
|
||
def __init__(
|
||
self,
|
||
server_url: str,
|
||
context: zmq.Context,
|
||
model_name: str,
|
||
vllm_block_size: int,
|
||
parallel_strategy: ParallelStrategy | int,
|
||
legacy_block_size: int | None = None,
|
||
*,
|
||
mq_timeout: float = DEFAULT_MQ_TIMEOUT,
|
||
heartbeat_interval: float = DEFAULT_HEARTBEAT_INTERVAL,
|
||
extra_config: dict[str, Any] | None = None,
|
||
):
|
||
"""Initialize the worker adapter for current or legacy vLLM callers.
|
||
|
||
Args:
|
||
server_url: The server URL for the LMCache message queue.
|
||
context: The ZMQ context.
|
||
model_name: The model name used for LMCache keys.
|
||
vllm_block_size: The block size used in vLLM, or legacy KV world
|
||
size when ``parallel_strategy`` is an int.
|
||
parallel_strategy: Current ``ParallelStrategy`` metadata, or the
|
||
legacy KV worker id from older vLLM connectors.
|
||
legacy_block_size: The vLLM block size passed positionally by
|
||
older vLLM connectors.
|
||
mq_timeout: Timeout in seconds for message queue requests.
|
||
Ignored when ``extra_config`` is provided.
|
||
heartbeat_interval: Interval in seconds between heartbeat pings.
|
||
Ignored when ``extra_config`` is provided.
|
||
extra_config: Optional dict with keys starting with
|
||
``lmcache.mp.`` (e.g., ``lmcache.mp.mq_timeout``). When
|
||
provided, it overrides ``mq_timeout`` / ``heartbeat_interval``.
|
||
|
||
Raises:
|
||
TypeError: If the connector argument shape is unsupported.
|
||
"""
|
||
vllm_block_size, parallel_strategy, mq_timeout = _normalize_adapter_init_args(
|
||
vllm_block_size,
|
||
parallel_strategy,
|
||
legacy_block_size,
|
||
mq_timeout,
|
||
)
|
||
if extra_config is not None:
|
||
cfg = _resolve_extra_config(extra_config)
|
||
mq_timeout = cfg[ExtraConfigDefault.mq_timeout.name]
|
||
heartbeat_interval = cfg[ExtraConfigDefault.heartbeat_interval.name]
|
||
# Only treat ``mp_transfer_mode`` as an explicit override when
|
||
# the user actually set it in extra_config; otherwise leave it
|
||
# as ``None`` so ``create_transfer_context`` can still consult
|
||
# the ``LMCACHE_MP_TRANSFER_MODE`` env var.
|
||
mp_mode_key = (
|
||
_EXTRA_CONFIG_KEY_PREFIX + ExtraConfigDefault.mp_transfer_mode.name
|
||
)
|
||
if mp_mode_key in extra_config:
|
||
self._mp_transfer_mode = cfg[ExtraConfigDefault.mp_transfer_mode.name]
|
||
else:
|
||
self._mp_transfer_mode = None
|
||
else:
|
||
self._mp_transfer_mode = None
|
||
self.mq_client = MessageQueueClient(server_url, context)
|
||
self._mq_timeout = mq_timeout
|
||
|
||
# Instance id for GPU worker. uuid4-derived (OS entropy) rather
|
||
# than os.getpid() to avoid collision in containerized deployments.
|
||
# Masked to 63 bits to stay signed-int64-safe for any msgpack peer.
|
||
self.instance_id: int = uuid.uuid4().int & ((1 << 63) - 1)
|
||
logger.info(
|
||
"LMCache MP worker adapter created with instance_id=%d",
|
||
self.instance_id,
|
||
)
|
||
|
||
# Registered kv caches from vLLM
|
||
self.kv_caches: dict[str, torch.Tensor] = {}
|
||
self.engine_group_infos: list[EngineGroupInfo] = []
|
||
|
||
# Transport context for transfer operations.
|
||
self.transfer_ctx: TransferContext | None = None
|
||
|
||
# Request futures
|
||
self.store_futures: dict[str, MessagingFuture[StoreResult]] = {}
|
||
# request_id -> (future, block_ids)
|
||
self.retrieve_futures: dict[
|
||
str, tuple[MessagingFuture[RetrieveResult], list[int]]
|
||
] = {}
|
||
# The IPC handle is not enough by itself; CUDA needs the exporting
|
||
# event object to stay alive until the consumer is done with it.
|
||
self.store_events: dict[str, _IpcEvent] = {}
|
||
self.retrieve_events: dict[str, _IpcEvent] = {}
|
||
|
||
# Block IDs that failed due to retrieve timeout
|
||
self.error_block_ids: set[int] = set()
|
||
|
||
# Retrieve request ids dropped by the unhealthy early-return of
|
||
# submit_retrieve_request. get_finished must still report each id
|
||
# exactly once, or async loads hang in WAITING_FOR_REMOTE_KVS.
|
||
self._dropped_retrieves: set[str] = set()
|
||
|
||
# The store requests that have finished execution in LMCache
|
||
self.finished_stores: set[str] = set()
|
||
# The finished request ids that are passed via vLLM and also
|
||
# have corresponding store requests submitted to LMCache before
|
||
self.previously_finished: set[str] = set()
|
||
# Request IDs already returned as finished_sending to the scheduler.
|
||
# Prevents re-reporting the same ID after drain clears tracking sets.
|
||
self._returned_finished: set[str] = set()
|
||
|
||
self.model_name = model_name
|
||
self.parallel_strategy = parallel_strategy
|
||
|
||
# Read chunk size from lmcache
|
||
try:
|
||
lmcache_tokens_per_chunk = get_lmcache_chunk_size(
|
||
self.mq_client, timeout=self._mq_timeout
|
||
)
|
||
except TimeoutError:
|
||
self.mq_client.close()
|
||
_raise_server_unreachable(server_url, self._mq_timeout)
|
||
self.lmcache_tokens_per_chunk = lmcache_tokens_per_chunk
|
||
assert lmcache_tokens_per_chunk % vllm_block_size == 0, (
|
||
"LMCache chunk size should be a multiple of vLLM block size"
|
||
)
|
||
self.blocks_in_chunk = lmcache_tokens_per_chunk // vllm_block_size
|
||
|
||
# Health state (shared with heartbeat thread)
|
||
self._health_event = threading.Event()
|
||
self._health_event.set()
|
||
|
||
# Heartbeat thread is created but NOT started yet.
|
||
# It will be lazily started on the first store or retrieve
|
||
# request, by which time vLLM is fully ready (model loaded,
|
||
# KV caches allocated, warmup & CUDA graph capture done).
|
||
self._heartbeat_interval = heartbeat_interval
|
||
self._heartbeat: HeartbeatThread | None = None
|
||
self._heartbeat_lock = threading.Lock()
|
||
if 3 * heartbeat_interval > _SERVER_REAP_TIMEOUT_FLOOR_SECONDS:
|
||
logger.warning(
|
||
"lmcache.mp.heartbeat_interval is %.1fs, so 3 x "
|
||
"heartbeat_interval (%.1fs) exceeds the MP server's "
|
||
"default worker reap timeout floor (%.1fs). Raise the "
|
||
"server's worker reap timeout to at least 3 x the "
|
||
"heartbeat interval, or the server may reap this "
|
||
"worker between heartbeats.",
|
||
heartbeat_interval,
|
||
3 * heartbeat_interval,
|
||
_SERVER_REAP_TIMEOUT_FLOOR_SECONDS,
|
||
)
|
||
|
||
# request telemetry, used for prefill-decode disagg
|
||
# TODO: pass down the configuration via vLLM connector config
|
||
# instead of env var
|
||
self.request_telemetry = RequestTelemetryFactory.create(
|
||
telemetry_type=os.getenv("LMCACHE_REQUEST_TELEMETRY_TYPE", "noop"),
|
||
config={
|
||
"endpoint": os.getenv(
|
||
"LMCACHE_REQUEST_TELEMETRY_ENDPOINT",
|
||
"http://localhost:5768/api/v1/telemetry",
|
||
),
|
||
},
|
||
)
|
||
|
||
@property
|
||
def is_healthy(self) -> bool:
|
||
"""Whether the LMCache server is healthy.
|
||
|
||
Reflects the most recent heartbeat result. KV cache
|
||
re-registration on the unhealthy->healthy transition is handled
|
||
by the heartbeat thread itself via ``register_recover_callback``,
|
||
so this property only reads the shared event.
|
||
"""
|
||
return self._health_event.is_set()
|
||
|
||
@property
|
||
def world_size(self) -> int:
|
||
"""Get the kv world size."""
|
||
return self.parallel_strategy.kv_world_size
|
||
|
||
@property
|
||
def worker_id(self) -> int:
|
||
"""Get the kv worker id."""
|
||
return self.parallel_strategy.kv_worker_id
|
||
|
||
@property
|
||
def is_kv_writer(self) -> bool:
|
||
"""Whether this worker is responsible for storing KV."""
|
||
return self.parallel_strategy.is_kv_writer
|
||
|
||
def register_kv_caches(
|
||
self,
|
||
kv_caches: dict[str, torch.Tensor],
|
||
engine_group_infos: Sequence[EngineGroupInfo] = (),
|
||
) -> None:
|
||
"""
|
||
Register the kv caches with LMCache server.
|
||
|
||
Args:
|
||
kv_caches: A dict of kv caches to register. The keys are the
|
||
layer names and the values are the corresponding tensors.
|
||
engine_group_infos: LMCache-owned engine KV cache group metadata.
|
||
|
||
Raises:
|
||
ConnectionError: if the server does not respond within
|
||
mq_timeout.
|
||
ValueError: if the LMCache chunk size is not a multiple of an
|
||
engine group's ``tokens_per_block`` (chunk boundaries would
|
||
not align with that group's paged-chunk boundaries).
|
||
"""
|
||
logger.info("Registering kv caches")
|
||
for info in engine_group_infos:
|
||
if (
|
||
info.tokens_per_block > 0
|
||
and self.lmcache_tokens_per_chunk % info.tokens_per_block
|
||
):
|
||
raise ValueError(
|
||
f"LMCache chunk size {self.lmcache_tokens_per_chunk} must be a "
|
||
f"multiple of engine group {info.engine_group_id} "
|
||
f"tokens_per_block {info.tokens_per_block}"
|
||
)
|
||
self.kv_caches = kv_caches
|
||
self.engine_group_infos = list(engine_group_infos)
|
||
self._send_register_kv_caches_request(kv_caches)
|
||
|
||
def _block_ids_per_group(self, op: LoadStoreOp) -> list[list[int]]:
|
||
return expand_engine_block_ids(self.engine_group_infos, op.block_ids)
|
||
|
||
def _send_register_kv_caches_request(
|
||
self,
|
||
kv_caches: dict[str, torch.Tensor],
|
||
) -> None:
|
||
"""Submit a REGISTER_KV_CACHE request and wait for the response.
|
||
|
||
Shared by the public ``register_kv_caches`` entry point and the
|
||
heartbeat recovery path (``_reregister_kv_caches_callback``).
|
||
|
||
Args:
|
||
kv_caches: The KV cache dict to register.
|
||
|
||
Raises:
|
||
ConnectionError: if the server does not respond within
|
||
mq_timeout.
|
||
"""
|
||
self.kv_caches = kv_caches
|
||
transfer_ctx = create_transfer_context(kv_caches, mode=self._mp_transfer_mode)
|
||
layout_hints = vllm_layout_hints()
|
||
self.transfer_ctx = transfer_ctx
|
||
try:
|
||
# Register on the local, not self.transfer_ctx: a concurrent
|
||
# shutdown() may null self.transfer_ctx between publish and this
|
||
# call. The local is always non-None.
|
||
transfer_ctx.register(
|
||
self.instance_id,
|
||
kv_caches,
|
||
self.model_name,
|
||
self.world_size,
|
||
self.blocks_in_chunk,
|
||
self.mq_client,
|
||
self._mq_timeout,
|
||
send_request=send_lmcache_request,
|
||
layout_hints=layout_hints,
|
||
engine_group_infos=self.engine_group_infos,
|
||
)
|
||
except TimeoutError:
|
||
raise ConnectionError(
|
||
"LMCache server did not respond to "
|
||
"register_kv_caches within "
|
||
f"{self._mq_timeout}s. Is the server running?"
|
||
) from None
|
||
|
||
def _ensure_heartbeat_started(self) -> None:
|
||
"""Lazily start the heartbeat thread on first store/retrieve.
|
||
|
||
The heartbeat starts healthy (the event was set at construction). A
|
||
live worker pings every interval, refreshing its server-side
|
||
``last_seen``, so it is never reaped while alive -- no re-registration
|
||
is needed at startup, and the first store/retrieve is not gated. The
|
||
recover callback still re-registers on a genuine unhealthy->healthy
|
||
edge (server restart).
|
||
"""
|
||
if self._heartbeat is not None:
|
||
return
|
||
with self._heartbeat_lock:
|
||
if self._heartbeat is not None:
|
||
return
|
||
heartbeat = HeartbeatThread(
|
||
mq_client=self.mq_client,
|
||
health_event=self._health_event,
|
||
interval=self._heartbeat_interval,
|
||
instance_id=self.instance_id,
|
||
)
|
||
heartbeat.register_recover_callback(self._reregister_kv_caches_callback)
|
||
heartbeat.start()
|
||
self._heartbeat = heartbeat
|
||
|
||
def _heartbeat_stop_requested(self) -> bool:
|
||
"""Whether a created heartbeat thread has a stop requested.
|
||
|
||
Returns:
|
||
``True`` only if a heartbeat exists and its stop was requested.
|
||
"""
|
||
heartbeat = self._heartbeat
|
||
return heartbeat is not None and heartbeat.stop_requested
|
||
|
||
def _reregister_kv_caches_callback(self) -> bool:
|
||
"""Heartbeat recover callback: re-register KV caches after the
|
||
server returns. Runs on the heartbeat thread, before the health
|
||
event is set.
|
||
|
||
Returns:
|
||
``True`` if nothing needs re-registering or registration
|
||
succeeds; ``False`` on failure or a requested heartbeat stop
|
||
(event stays cleared; retried on the next successful PING).
|
||
"""
|
||
if not self.kv_caches:
|
||
# Nothing was registered yet (server flapped before the
|
||
# very first register_kv_caches). Treat as success so the
|
||
# health event can be set.
|
||
return True
|
||
|
||
# Skip the rebuild if a shutdown already requested the heartbeat stop.
|
||
if self._heartbeat_stop_requested():
|
||
logger.info("Heartbeat stop requested; skipping KV cache re-registration")
|
||
return False
|
||
|
||
try:
|
||
self._send_register_kv_caches_request(self.kv_caches)
|
||
except ConnectionError:
|
||
logger.exception(
|
||
"Failed to re-register KV caches after server recovery; "
|
||
"will retry on next heartbeat"
|
||
)
|
||
return False
|
||
except Exception:
|
||
logger.exception(
|
||
"Unexpected error during KV cache re-registration; "
|
||
"will retry on next heartbeat"
|
||
)
|
||
return False
|
||
logger.warning("Finished re-registering KV caches after server recovery")
|
||
return True
|
||
|
||
@_lmcache_nvtx_annotate
|
||
def submit_store_request(
|
||
self,
|
||
request_id: str,
|
||
op: LoadStoreOp,
|
||
event: _IpcEvent,
|
||
cache_salt: str = "",
|
||
):
|
||
"""
|
||
Submit a KV cache store request to LMCache
|
||
|
||
Args:
|
||
request_id: The ID of the request
|
||
op: The LoadStoreOp describing the store operation.
|
||
event: The CUDA event that is recorded after the current
|
||
model inference step
|
||
cache_salt: Per-user isolation salt.
|
||
"""
|
||
self._ensure_heartbeat_started()
|
||
|
||
if not self.is_kv_writer:
|
||
return
|
||
|
||
if not self.is_healthy:
|
||
return
|
||
|
||
assert op.token_ids is not None
|
||
key = self._create_key(
|
||
op.token_ids,
|
||
op.start,
|
||
op.end,
|
||
request_id=request_id,
|
||
cache_salt=cache_salt,
|
||
)
|
||
if self.transfer_ctx is None:
|
||
raise RuntimeError(
|
||
"Transfer context is not initialized. "
|
||
"Call register_kv_caches() before submitting store requests."
|
||
)
|
||
future = self.transfer_ctx.submit_store(
|
||
request_id,
|
||
key,
|
||
self.instance_id,
|
||
self.kv_caches,
|
||
self._block_ids_per_group(op),
|
||
event,
|
||
self.blocks_in_chunk,
|
||
)
|
||
self.store_futures[request_id] = future
|
||
self.store_events[request_id] = event
|
||
|
||
@_lmcache_nvtx_annotate
|
||
def submit_retrieve_request(
|
||
self,
|
||
request_id: str,
|
||
op: LoadStoreOp,
|
||
event: _IpcEvent,
|
||
cache_salt: str = "",
|
||
) -> None:
|
||
"""
|
||
Submit a KV cache retrieve request to LMCache
|
||
|
||
When the server is unhealthy the request is not submitted: blocks
|
||
are flagged via ``error_block_ids`` (vLLM recomputes) and the id is
|
||
recorded so ``get_finished`` still reports it exactly once.
|
||
|
||
Args:
|
||
request_id: The ID of the request
|
||
op: The LoadStoreOp describing the retrieve operation.
|
||
event: The CUDA event that is recorded after the current
|
||
model inference step
|
||
cache_salt: Per-user isolation salt.
|
||
"""
|
||
self._ensure_heartbeat_started()
|
||
|
||
if not self.is_healthy:
|
||
self.error_block_ids.update(op.flat_block_ids)
|
||
self._dropped_retrieves.add(request_id)
|
||
return
|
||
|
||
assert op.token_ids is not None
|
||
key = self._create_key(
|
||
op.token_ids,
|
||
op.start,
|
||
op.end,
|
||
request_id=request_id,
|
||
cache_salt=cache_salt,
|
||
)
|
||
if self.transfer_ctx is None:
|
||
raise RuntimeError(
|
||
"Transfer context is not initialized. "
|
||
"Call register_kv_caches() before submitting retrieve requests."
|
||
)
|
||
future = self.transfer_ctx.submit_retrieve(
|
||
request_id,
|
||
key,
|
||
self.instance_id,
|
||
self.kv_caches,
|
||
self._block_ids_per_group(op),
|
||
event,
|
||
self.blocks_in_chunk,
|
||
skip_first_n_tokens=op.skip_first_n_tokens,
|
||
)
|
||
self.retrieve_futures[request_id] = (future, op.flat_block_ids)
|
||
self.retrieve_events[request_id] = event
|
||
|
||
@_lmcache_nvtx_annotate
|
||
def batched_submit_store_requests(
|
||
self,
|
||
request_ids: list[str],
|
||
ops: list[LoadStoreOp],
|
||
event: _IpcEvent,
|
||
cache_salts: list[str] | None = None,
|
||
):
|
||
"""
|
||
Submit a batched store request to LMCache
|
||
|
||
Args:
|
||
request_ids: The IDs of the requests
|
||
ops: The LoadStoreOps describing the store operations. Should have
|
||
the same length as request_ids
|
||
event: The CUDA event that is recorded after the current
|
||
model inference step
|
||
cache_salts: Per-user isolation salts, one per request. If None,
|
||
all requests use cache_salt="". The list length should be the same as
|
||
request_ids.
|
||
"""
|
||
if cache_salts is None:
|
||
cache_salts = [""] * len(request_ids)
|
||
for request_id, op, salt in zip(request_ids, ops, cache_salts, strict=False):
|
||
self.submit_store_request(request_id, op, event, cache_salt=salt)
|
||
|
||
@_lmcache_nvtx_annotate
|
||
def batched_submit_retrieve_requests(
|
||
self,
|
||
request_ids: list[str],
|
||
ops: list[LoadStoreOp],
|
||
event: _IpcEvent,
|
||
cache_salts: list[str] | None = None,
|
||
):
|
||
"""
|
||
Submit a batched retrieve request to LMCache
|
||
|
||
Args:
|
||
request_ids: The IDs of the requests
|
||
ops: The LoadStoreOps describing the retrieve operations. Should have
|
||
the same length as request_ids
|
||
event: The CUDA event that is recorded after the current
|
||
model inference step
|
||
cache_salts: Per-user isolation salts, one per request. If None,
|
||
all requests use cache_salt="". The list length should be same as
|
||
request_ids.
|
||
"""
|
||
if cache_salts is None:
|
||
cache_salts = [""] * len(request_ids)
|
||
for request_id, op, salt in zip(request_ids, ops, cache_salts, strict=False):
|
||
self.submit_retrieve_request(request_id, op, event, cache_salt=salt)
|
||
|
||
def _process_finished_stores(
|
||
self,
|
||
finished_req_ids_from_lmcache: set[str],
|
||
finished_req_ids_from_engine: set[str],
|
||
) -> set[str]:
|
||
"""Merge LMCache-side and engine-side finished store info."""
|
||
self.finished_stores.update(finished_req_ids_from_lmcache)
|
||
ret_stores = set()
|
||
for req_id in finished_req_ids_from_engine:
|
||
if req_id in self._returned_finished:
|
||
continue
|
||
if req_id in self.finished_stores or req_id in self.store_futures:
|
||
self.previously_finished.add(req_id)
|
||
else:
|
||
ret_stores.add(req_id)
|
||
ret_stores.update(self._update_and_get_finished_store())
|
||
self._returned_finished.update(ret_stores)
|
||
return ret_stores
|
||
|
||
@_lmcache_nvtx_annotate
|
||
def get_finished(
|
||
self, finished_req_ids_from_engine: set[str]
|
||
) -> tuple[set[str] | None, set[str] | None]:
|
||
"""
|
||
Check and get the finished store and retrieve requests.
|
||
|
||
Args:
|
||
finished_req_ids_from_engine: the set of request ids that are
|
||
reported as finished from the vLLM engine side.
|
||
|
||
Returns:
|
||
A tuple of two sets:
|
||
- The first set contains the finished store request ids. The returned
|
||
store request ids MUST be seen before in the
|
||
`finished_req_ids_from_engine`.
|
||
- The second set contains the finished retrieve request ids,
|
||
including retrieves dropped at submit time while unhealthy
|
||
(reported exactly once; blocks already in error_block_ids).
|
||
|
||
Notes:
|
||
When enabling async scheduling in vLLM, the same request ID may appear
|
||
multiple times in `finished_req_ids_from_engine`. The adapter should
|
||
take care of deduplicating the request IDs and only return the request
|
||
IDs that have not been returned before.
|
||
"""
|
||
# If unhealthy, drain all pending futures immediately
|
||
if not self.is_healthy:
|
||
finished_stores = set(self.store_futures.keys())
|
||
finished_retrieves = set()
|
||
for request_id, (
|
||
_r_future,
|
||
r_block_ids,
|
||
) in self.retrieve_futures.items():
|
||
finished_retrieves.add(request_id)
|
||
self.error_block_ids.update(r_block_ids)
|
||
self.store_futures.clear()
|
||
self.retrieve_futures.clear()
|
||
self.store_events.clear()
|
||
self.retrieve_events.clear()
|
||
|
||
# Retrieves dropped at submit time still must be reported,
|
||
# exactly once, or async loads hang in WAITING_FOR_REMOTE_KVS.
|
||
# Swap-drain (not update-then-clear): a concurrent
|
||
# submit_retrieve_request add lands in the old set (reported now)
|
||
# or the fresh set (reported next call), never lost.
|
||
dropped = self._dropped_retrieves
|
||
self._dropped_retrieves = set()
|
||
finished_retrieves.update(dropped)
|
||
|
||
ret_stores = self._process_finished_stores(
|
||
finished_stores, finished_req_ids_from_engine
|
||
)
|
||
# A request may have a pending retrieve AND appear in
|
||
# finished_req_ids_from_engine (it ran without loading KV after
|
||
# the server died). The scheduler processes finished_recving
|
||
# first and deletes the request, so we must not also report it
|
||
# in finished_sending.
|
||
ret_stores -= finished_retrieves
|
||
return ret_stores, finished_retrieves
|
||
|
||
finished_stores = set()
|
||
finished_retrieves = set()
|
||
for request_id, s_future in self.store_futures.items():
|
||
if not s_future.query():
|
||
continue
|
||
|
||
s_result = s_future.result()
|
||
finished_stores.add(request_id)
|
||
|
||
if not s_result:
|
||
logger.error(
|
||
"Something went wrong when processing the "
|
||
"store request for request_id=%s",
|
||
request_id,
|
||
)
|
||
|
||
for request_id, (r_future, _) in self.retrieve_futures.items():
|
||
if not r_future.query():
|
||
continue
|
||
|
||
r_result = r_future.result()
|
||
finished_retrieves.add(request_id)
|
||
|
||
if not r_result:
|
||
logger.error(
|
||
"Something went wrong when processing the "
|
||
"retrieve request for request_id=%s, result=%s",
|
||
request_id,
|
||
r_result,
|
||
)
|
||
|
||
# Remove the finished requests from the tracking dicts
|
||
for request_id in finished_stores:
|
||
self.store_futures.pop(request_id, None)
|
||
self.store_events.pop(request_id, None)
|
||
for request_id in finished_retrieves:
|
||
self.retrieve_futures.pop(request_id, None)
|
||
self.retrieve_events.pop(request_id, None)
|
||
|
||
# Retrieves dropped while unhealthy still must be reported,
|
||
# exactly once, or async loads hang in WAITING_FOR_REMOTE_KVS. No
|
||
# finished_sending dedup is needed (unlike the unhealthy branch): a
|
||
# dropped retrieve's request is parked in WAITING_FOR_REMOTE_KVS until
|
||
# this report, so it cannot also be engine-finished in the same call.
|
||
# Swap-drain so a concurrent submit_retrieve_request add is never lost.
|
||
dropped = self._dropped_retrieves
|
||
self._dropped_retrieves = set()
|
||
finished_retrieves.update(dropped)
|
||
|
||
# Update the internal states
|
||
ret_stores = self._process_finished_stores(
|
||
finished_stores, finished_req_ids_from_engine
|
||
)
|
||
|
||
# the invocation of `get_finished` means that
|
||
# these requests' KV caches are already fully stored.
|
||
# or the requests normally ends without any store.
|
||
if ret_stores:
|
||
self.request_telemetry.on_request_store_finished(
|
||
request_ids_set=ret_stores,
|
||
model_name=self.model_name,
|
||
world_size=self.world_size,
|
||
kv_rank=self.worker_id,
|
||
)
|
||
|
||
return ret_stores, finished_retrieves
|
||
|
||
def num_blocks_per_chunk(self) -> int:
|
||
"""
|
||
Returns:
|
||
The number of vllm blocks in a LMCache data chunk
|
||
"""
|
||
return self.blocks_in_chunk
|
||
|
||
def get_block_ids_with_load_errors(self) -> set[int]:
|
||
"""
|
||
Returns the block IDs that failed due to retrieve timeout,
|
||
then clears the internal set.
|
||
"""
|
||
errors = self.error_block_ids.copy()
|
||
self.error_block_ids.clear()
|
||
return errors
|
||
|
||
def handle_preemptions(self, need_flush_before_forward: bool) -> None:
|
||
"""Handle worker-side preemption hints from connector metadata.
|
||
|
||
When ``need_flush_before_forward`` is true, synchronize deferred
|
||
engine-driven gather work before the next forward pass can overwrite
|
||
paged KV blocks.
|
||
|
||
Args:
|
||
need_flush_before_forward: When True, flush in-flight gather
|
||
operations on the transfer context. When False, this is a no-op.
|
||
"""
|
||
if not need_flush_before_forward:
|
||
return
|
||
if not self.is_healthy or self.transfer_ctx is None:
|
||
return
|
||
self.transfer_ctx.flush_inflight_stores()
|
||
# Force device sync here, compare to preemption, perf panelty is trivial
|
||
torch_dev.synchronize()
|
||
|
||
def shutdown(self) -> None:
|
||
"""
|
||
Shutdown the LMCache MP worker adapter.
|
||
|
||
Stops the heartbeat (if started) before UNREGISTER: no new ping
|
||
on the closing mq_client, and a straggler in-flight cycle cannot
|
||
re-register or flip the health event after unregistration.
|
||
"""
|
||
with self._heartbeat_lock:
|
||
if self._heartbeat is not None:
|
||
self._heartbeat.stop()
|
||
|
||
logger.info("Unregistering kv caches")
|
||
try:
|
||
unregister_type = (
|
||
RequestType.UNREGISTER_KV_CACHE_ENGINE_DRIVEN_CONTEXT
|
||
if isinstance(self.transfer_ctx, EngineDrivenTransferContext)
|
||
else RequestType.UNREGISTER_KV_CACHE
|
||
)
|
||
send_lmcache_request(
|
||
self.mq_client,
|
||
unregister_type,
|
||
[self.instance_id],
|
||
).result(timeout=self._mq_timeout)
|
||
except TimeoutError:
|
||
logger.warning(
|
||
"LMCache server did not respond to unregister within %ss. "
|
||
"Proceeding with shutdown.",
|
||
self._mq_timeout,
|
||
)
|
||
|
||
if self.transfer_ctx is not None:
|
||
self.transfer_ctx.close()
|
||
self.transfer_ctx = None
|
||
|
||
self.mq_client.close()
|
||
self.request_telemetry.close()
|
||
|
||
# Helper functions
|
||
def _update_and_get_finished_store(
|
||
self,
|
||
) -> set[str]:
|
||
"""Converge the internal states about finished stores
|
||
and returns the 'safe finished store request ids' back
|
||
"""
|
||
safe_finished_s = self.finished_stores.intersection(self.previously_finished)
|
||
self.finished_stores.difference_update(self.previously_finished)
|
||
self.previously_finished.difference_update(safe_finished_s)
|
||
|
||
return safe_finished_s
|
||
|
||
def _create_key(
|
||
self,
|
||
token_ids: list[int],
|
||
start: int,
|
||
end: int,
|
||
request_id: str,
|
||
cache_salt: str = "",
|
||
) -> IPCCacheServerKey:
|
||
"""Convert token IDs to an IPC cache engine key.
|
||
|
||
Args:
|
||
token_ids: The token IDs.
|
||
start: Start token index.
|
||
end: End token index.
|
||
request_id: The request ID.
|
||
cache_salt: Per-user isolation salt.
|
||
|
||
Returns:
|
||
IPCCacheServerKey: The constructed key.
|
||
"""
|
||
return IPCCacheServerKey(
|
||
model_name=self.model_name,
|
||
world_size=self.world_size,
|
||
worker_id=self.worker_id,
|
||
token_ids=tuple(token_ids),
|
||
start=start,
|
||
end=end,
|
||
request_id=request_id,
|
||
cache_salt=cache_salt,
|
||
)
|