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

lmcache/v1/multiprocess/config.py

@@ -0,0 +1,253 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Configuration for the multiprocess (ZMQ) server and HTTP frontend.
+"""
+
+# Standard
+from dataclasses import dataclass, field
+import argparse
+import json
+
+
+@dataclass
+class MPServerConfig:
+    """Configuration for the ZMQ-based multiprocess cache server."""
+
+    host: str = "localhost"
+    """ZMQ server host."""
+
+    port: int = 5555
+    """ZMQ server port."""
+
+    chunk_size: int = 256
+    """Chunk size for KV cache operations."""
+
+    max_workers: int = 1
+    """Base number of worker threads. Sets default for both GPU and CPU pools."""
+
+    max_gpu_workers: int = 1
+    """Worker threads for the GPU affinity pool (STORE/RETRIEVE).
+    Resolved from --max-gpu-workers or --max-workers."""
+
+    max_cpu_workers: int = 1
+    """Worker threads for the normal (CPU) pool (LOOKUP, END_SESSION, etc.).
+    Resolved from --max-cpu-workers or --max-workers."""
+
+    hash_algorithm: str = "blake3"
+    """Hash algorithm for token-based operations (builtin, sha256_cbor, blake3)."""
+
+    engine_type: str = "default"
+    """Cache engine backend type
+    ('default' for MPCacheEngine, 'blend' for BlendEngineV2).
+    """
+
+    runtime_plugin_config: "RuntimePluginConfig" = field(
+        default_factory=lambda: RuntimePluginConfig()
+    )
+    """Runtime plugin configuration (locations + extra config)."""
+
+
+@dataclass
+class RuntimePluginConfig:
+    """Configuration for runtime plugins."""
+
+    locations: list[str] = field(default_factory=list)
+    """Paths to runtime plugin scripts or directories."""
+
+    extra_config: dict = field(default_factory=dict)
+    """Extra key-value config forwarded to runtime plugins
+    via the JSON config blob.
+    Accepts a JSON string on the command line.
+    """
+
+
+DEFAULT_MP_SERVER_CONFIG = MPServerConfig()
+
+
+@dataclass
+class HTTPFrontendConfig:
+    """Configuration for the HTTP frontend (uvicorn/FastAPI)."""
+
+    http_host: str = "0.0.0.0"
+    """HTTP server host."""
+
+    http_port: int = 8080
+    """HTTP server port."""
+
+
+DEFAULT_HTTP_FRONTEND_CONFIG = HTTPFrontendConfig()
+
+
+def add_mp_server_args(
+    parser: argparse.ArgumentParser,
+) -> argparse.ArgumentParser:
+    """
+    Add MP server configuration arguments to an existing parser.
+
+    Args:
+        parser: The argument parser to add arguments to.
+
+    Returns:
+        The same parser with MP server arguments added.
+    """
+    mp_group = parser.add_argument_group(
+        "MP Server", "Configuration for the ZMQ multiprocess cache server"
+    )
+    mp_group.add_argument(
+        "--host",
+        type=str,
+        default="localhost",
+        help="Host to bind the ZMQ server. Default is localhost.",
+    )
+    mp_group.add_argument(
+        "--port",
+        type=int,
+        default=5555,
+        help="Port to bind the ZMQ server. Default is 5555.",
+    )
+    mp_group.add_argument(
+        "--chunk-size",
+        type=int,
+        default=256,
+        help="Chunk size for KV cache operations. Default is 256.",
+    )
+    mp_group.add_argument(
+        "--max-workers",
+        type=int,
+        default=1,
+        help="Base number of worker threads for both GPU and CPU pools. "
+        "Default is 1. Can be overridden per-pool with "
+        "--max-gpu-workers and --max-cpu-workers.",
+    )
+    mp_group.add_argument(
+        "--max-gpu-workers",
+        type=int,
+        default=None,
+        help="Worker threads for the GPU affinity pool (STORE/RETRIEVE). "
+        "Defaults to --max-workers if not specified.",
+    )
+    mp_group.add_argument(
+        "--max-cpu-workers",
+        type=int,
+        default=None,
+        help="Worker threads for the normal CPU pool (LOOKUP, etc.). "
+        "Defaults to --max-workers if not specified.",
+    )
+    mp_group.add_argument(
+        "--hash-algorithm",
+        type=str,
+        default="blake3",
+        help="Hash algorithm for token-based operations "
+        "(builtin, sha256_cbor, blake3). Default is blake3.",
+    )
+    mp_group.add_argument(
+        "--engine-type",
+        type=str,
+        default="default",
+        choices=["default", "blend"],
+        help="Cache engine backend type. 'default' uses MPCacheEngine, "
+        "'blend' uses BlendEngineV2 for cross-request KV reuse. "
+        "Default is 'default'.",
+    )
+    mp_group.add_argument(
+        "--runtime-plugin-locations",
+        type=str,
+        nargs="*",
+        default=[],
+        help="Paths to runtime plugin scripts or "
+        "directories to launch alongside the server.",
+    )
+    mp_group.add_argument(
+        "--runtime-plugin-config",
+        type=str,
+        default="{}",
+        help="JSON string of extra key-value config forwarded to runtime "
+        "plugins via LMCACHE_RUNTIME_PLUGIN_EXTRA_CONFIG. "
+        'Example: \'{"plugin.frontend.heartbeat_url": '
+        '"http://localhost:5000/heartbeat"}\'',
+    )
+    return parser
+
+
+def parse_args_to_mp_server_config(
+    args: argparse.Namespace,
+) -> MPServerConfig:
+    """
+    Convert parsed command line arguments to an MPServerConfig.
+
+    Args:
+        args: Parsed arguments from the argument parser.
+
+    Returns:
+        MPServerConfig: The configuration object.
+    """
+    base = args.max_workers
+    max_gpu = args.max_gpu_workers if args.max_gpu_workers is not None else base
+    max_cpu = args.max_cpu_workers if args.max_cpu_workers is not None else base
+    try:
+        plugin_extra = json.loads(getattr(args, "runtime_plugin_config", None) or "{}")
+    except json.JSONDecodeError as exc:
+        raise ValueError("--runtime-plugin-config is not valid JSON: %s" % exc) from exc
+    return MPServerConfig(
+        host=args.host,
+        port=args.port,
+        chunk_size=args.chunk_size,
+        max_workers=base,
+        max_gpu_workers=max_gpu,
+        max_cpu_workers=max_cpu,
+        hash_algorithm=args.hash_algorithm,
+        engine_type=args.engine_type,
+        runtime_plugin_config=RuntimePluginConfig(
+            locations=(args.runtime_plugin_locations or []),
+            extra_config=plugin_extra,
+        ),
+    )
+
+
+def add_http_frontend_args(
+    parser: argparse.ArgumentParser,
+) -> argparse.ArgumentParser:
+    """
+    Add HTTP frontend configuration arguments to an existing parser.
+
+    Args:
+        parser: The argument parser to add arguments to.
+
+    Returns:
+        The same parser with HTTP frontend arguments added.
+    """
+    http_group = parser.add_argument_group(
+        "HTTP Frontend", "Configuration for the HTTP frontend server"
+    )
+    http_group.add_argument(
+        "--http-host",
+        type=str,
+        default="0.0.0.0",
+        help="Host to bind the HTTP server. Default is 0.0.0.0.",
+    )
+    http_group.add_argument(
+        "--http-port",
+        type=int,
+        default=8080,
+        help="Port to bind the HTTP server. Default is 8080.",
+    )
+    return parser
+
+
+def parse_args_to_http_frontend_config(
+    args: argparse.Namespace,
+) -> HTTPFrontendConfig:
+    """
+    Convert parsed command line arguments to an HTTPFrontendConfig.
+
+    Args:
+        args: Parsed arguments from the argument parser.
+
+    Returns:
+        HTTPFrontendConfig: The configuration object.
+    """
+    return HTTPFrontendConfig(
+        http_host=args.http_host,
+        http_port=args.http_port,
+    )

lmcache/v1/multiprocess/custom_types.py

@@ -0,0 +1,281 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from dataclasses import dataclass, field
+from typing import Any, Callable
+import pickle
+import threading
+
+# Third Party
+import msgspec
+import torch
+
+"""
+Defines the types and the customized encoder/decoders for inter-process
+communications.
+
+Key Types:
+- IPCCacheEngineKey: Token-based cache key
+  - Contains token_ids, start, end, request_id (all required)
+  - Converted to ObjectKey for storage operations via ipc_key_to_object_keys()
+"""
+
+
+class CudaIPCWrapper:
+    _discovered_device_mapping: dict[str, int] = {}
+    _device_mapping_lock = threading.Lock()
+
+    @staticmethod
+    def _get_device_uuid(device_index: int) -> str:
+        """Get the UUID of a GPU device given its index."""
+        return str(torch.cuda.get_device_properties(device_index).uuid)
+
+    @staticmethod
+    def _discover_gpu_devices():
+        """Discover all available GPU devices and map their UUIDs to
+        the physical device ordinals.
+        """
+        if not torch.cuda.is_available():
+            return
+
+        num_devices = torch.cuda.device_count()
+        with CudaIPCWrapper._device_mapping_lock:
+            if CudaIPCWrapper._discovered_device_mapping:
+                return  # Already discovered
+
+            for i in range(num_devices):
+                device_uuid = CudaIPCWrapper._get_device_uuid(i)
+                CudaIPCWrapper._discovered_device_mapping[device_uuid] = i
+
+    @staticmethod
+    def _get_device_index_from_uuid(device_uuid: str) -> int:
+        """Get the physical device ordinal from its UUID."""
+        CudaIPCWrapper._discover_gpu_devices()
+
+        with CudaIPCWrapper._device_mapping_lock:
+            device_index = CudaIPCWrapper._discovered_device_mapping.get(
+                device_uuid, None
+            )
+
+        if device_index is None:
+            raise RuntimeError(
+                f"Device UUID {device_uuid} not found in the discovered devices."
+                "Please make sure the process can see all the GPU devices"
+            )
+        return device_index
+
+    def __init__(self, tensor: torch.Tensor):
+        # First Party
+        from lmcache.v1.gpu_connector.utils import assert_contiguous
+
+        assert_contiguous(tensor)
+
+        storage = tensor.untyped_storage()
+        handle = storage._share_cuda_()
+
+        self.handle = handle
+        self.dtype = tensor.dtype
+        self.shape = tuple(tensor.shape)
+        self.stride = tuple(tensor.stride())
+        self.storage_offset = int(tensor.storage_offset())
+
+        device_index = tensor.device.index
+        self.device_uuid = CudaIPCWrapper._get_device_uuid(device_index)
+
+    def to_tensor(self) -> torch.Tensor:
+        """
+        Note:
+            This function may break if torch cuda is not initialized.
+            We should call `torch.cuda.init()` before using this function.
+        """
+        device_index = CudaIPCWrapper._get_device_index_from_uuid(self.device_uuid)
+
+        storage = torch.UntypedStorage._new_shared_cuda(  # noqa: SLF001
+            device_index, *self.handle[1:]
+        )
+
+        t = torch.empty((), device=f"cuda:{device_index}", dtype=self.dtype)
+        t.set_(storage, self.storage_offset, self.shape, self.stride)
+        return t
+
+    def __eq__(self, other):
+        if not isinstance(other, CudaIPCWrapper):
+            return False
+        return (
+            self.handle == other.handle
+            and self.dtype == other.dtype
+            and self.shape == other.shape
+            and self.stride == other.stride
+            and self.storage_offset == other.storage_offset
+            and self.device_uuid == other.device_uuid
+        )
+
+    @staticmethod
+    def Serialize(obj: "CudaIPCWrapper") -> bytes:
+        return pickle.dumps(obj)
+
+    @staticmethod
+    def Deserialize(data: bytes) -> "CudaIPCWrapper":
+        return pickle.loads(data)
+
+
+@dataclass(order=True, frozen=True)
+class IPCCacheEngineKey:
+    """Cache key for the IPC (multiprocess) protocol.
+
+    This key type is sent by the client over ZMQ (serialized via msgspec).
+
+    The client sends token_ids, start, end, and request_id (all required).
+    The server computes chunk hashes via TokenHasher and converts to
+    ObjectKey for storage operations using ipc_key_to_object_keys().
+
+    The request_id field is for session tracking and is NOT included
+    in equality/hash comparisons (two keys with same content but different
+    request_ids are considered equal for cache purposes).
+    """
+
+    model_name: str
+    world_size: int
+    worker_id: int | None
+
+    token_ids: tuple[int, ...]  # frozen tuple for hashability
+    start: int
+    end: int
+
+    # === Session tracking (not part of cache identity) ===
+    request_id: str = field(compare=False)
+
+    # === Per-user isolation salt (part of cache identity) ===
+    # msgspec encodes dataclasses as maps, so forward wire compatibility
+    # works by field name: an old payload without ``cache_salt`` decodes
+    # on new code using the default "". Placing the field last is a style
+    # choice — all defaulted fields must come after non-defaulted ones.
+    #
+    # Invariant: must not contain ``@``, ``/``, ``\``, or NUL, and
+    # must be <= 128 chars — same rationale as ObjectKey (see
+    # ObjectKey.cache_salt). Validated in __post_init__.
+    cache_salt: str = ""
+
+    # Duplicated from ObjectKey — cannot import ObjectKey here due to
+    # circular dependency (api.py imports IPCCacheEngineKey).
+    _SALT_FORBIDDEN_CHARS = frozenset("@/\\\x00")
+    _SALT_MAX_LEN = 128
+
+    def __post_init__(self) -> None:
+        bad = self._SALT_FORBIDDEN_CHARS & set(self.cache_salt)
+        if bad:
+            raise ValueError(
+                f"cache_salt must not contain {bad!r} (got {self.cache_salt!r})"
+            )
+        if len(self.cache_salt) > self._SALT_MAX_LEN:
+            raise ValueError(
+                f"cache_salt exceeds max length {self._SALT_MAX_LEN} "
+                f"(got {len(self.cache_salt)})"
+            )
+
+    # Helper function for unit tests only
+    @classmethod
+    def from_token_ids(
+        cls,
+        model_name: str,
+        world_size: int,
+        worker_id: int | None,
+        token_ids: list[int],
+        start: int = 0,
+        end: int = 0,
+        request_id: str = "",
+        cache_salt: str = "",
+    ) -> "IPCCacheEngineKey":
+        """Create a key from token ids. Only used by the tests."""
+        return cls(
+            model_name=model_name,
+            world_size=world_size,
+            worker_id=worker_id,
+            token_ids=tuple(token_ids),
+            start=start,
+            end=end,
+            request_id=request_id,
+            cache_salt=cache_salt,
+        )
+
+    def no_worker_id_version(self) -> "IPCCacheEngineKey":
+        """Create a copy with worker_id=None for lookup requests."""
+        return IPCCacheEngineKey(
+            model_name=self.model_name,
+            world_size=self.world_size,
+            worker_id=None,
+            token_ids=self.token_ids,
+            start=self.start,
+            end=self.end,
+            request_id=self.request_id,
+            cache_salt=self.cache_salt,
+        )
+
+
+# Type exports
+KVCache = list[CudaIPCWrapper]
+
+
+@dataclass
+class CustomizedSerdeConfig:
+    serializer: Callable[[Any], bytes]
+    deserializer: Callable[[bytes], Any]
+    code: int
+
+
+_CUSTOMERIZED_SERIALIZERS = {
+    CudaIPCWrapper: CustomizedSerdeConfig(
+        serializer=CudaIPCWrapper.Serialize,
+        deserializer=CudaIPCWrapper.Deserialize,
+        code=1,
+    ),
+}
+
+
+def get_customized_encoder(type: Any) -> msgspec.msgpack.Encoder:
+    # TODO: `type` is not used here
+    def enc_hook(obj: Any) -> Any:
+        for supported_type, cfg in _CUSTOMERIZED_SERIALIZERS.items():
+            if isinstance(obj, supported_type):
+                data = cfg.serializer(obj)
+                return msgspec.msgpack.Ext(cfg.code, data)
+        raise TypeError(f"Unsupported type for serialization: {type(obj)}")
+
+    return msgspec.msgpack.Encoder(enc_hook=enc_hook)
+
+
+def get_customized_decoder(type: Any) -> msgspec.msgpack.Decoder:
+    def ext_hook(code: int, data: bytes) -> Any:
+        for cfg in _CUSTOMERIZED_SERIALIZERS.values():
+            if cfg.code == code:
+                return cfg.deserializer(data)
+        raise TypeError(f"Unsupported ext code for deserialization: {code}")
+
+    return msgspec.msgpack.Decoder(ext_hook=ext_hook, type=type)
+
+
+@dataclass
+class BlockAllocationRecord:
+    """A single per-request GPU block allocation delta from vLLM."""
+
+    req_id: str
+    new_block_ids: list[int]
+    new_token_ids: list[int]
+
+
+@dataclass
+class CBMatchResult:
+    """Result of a sub-sequence match from BlendTokenRangeMatcher.
+
+    Attributes:
+        old_st: Start position in the originally registered (stored) sequence.
+        old_ed: End position in the originally registered (stored) sequence.
+        cur_st: Start position in the query sequence where the match was found.
+        cur_ed: End position in the query sequence where the match was found.
+        hash: Token hash bytes (from registration) used as the storage key.
+    """
+
+    old_st: int
+    old_ed: int
+    cur_st: int
+    cur_ed: int
+    hash: bytes