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

lmcache/integration/request_telemetry/factory.py

@@ -0,0 +1,113 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Factory for creating request telemetry reporters based on configuration.
+"""
+
+# Standard
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any, Optional
+import importlib
+
+if TYPE_CHECKING:
+    # First Party
+    from lmcache.integration.request_telemetry.base import RequestTelemetry
+
+
+class RequestTelemetryFactory:
+    """Factory for creating request telemetry reporters."""
+
+    _registry: dict[str, Callable[[], type["RequestTelemetry"]]] = {}
+    _instances: dict[str, "RequestTelemetry"] = {}
+
+    @classmethod
+    def register(cls, name: str, module_path: str, class_name: str) -> None:
+        """Register a telemetry reporter with lazy-loading module and class name.
+
+        Args:
+            name: The name to register the telemetry reporter under.
+            module_path: The module path to import the class from.
+            class_name: The class name to import from the module.
+
+        Raises:
+            ValueError: If a reporter with the same name is already registered.
+        """
+        if name in cls._registry:
+            raise ValueError(f"Telemetry reporter '{name}' is already registered.")
+
+        def loader() -> type["RequestTelemetry"]:
+            module = importlib.import_module(module_path)
+            return getattr(module, class_name)
+
+        cls._registry[name] = loader
+
+    @classmethod
+    def create(
+        cls,
+        telemetry_type: Optional[str] = None,
+        config: Optional[dict[str, Any]] = None,
+        use_singleton: bool = True,
+    ) -> "RequestTelemetry":
+        """Create a request telemetry reporter based on the specified type.
+
+        Args:
+            telemetry_type: The type of telemetry reporter to create.
+                If None, defaults to "noop".
+                Supported values: "noop", "fastapi".
+            config: Optional dict of configuration options passed to the reporter.
+                Each reporter defines its own expected keys.
+            use_singleton: If True, returns a cached singleton instance.
+                If False, creates a new instance each time.
+                Note: Singleton is keyed by telemetry_type only.
+
+        Returns:
+            A RequestTelemetry instance.
+
+        Raises:
+            ValueError: If an unsupported telemetry type is specified.
+        """
+        if telemetry_type is None:
+            telemetry_type = "noop"
+
+        telemetry_type = telemetry_type.lower()
+
+        if telemetry_type not in cls._registry:
+            raise ValueError(
+                f"Unsupported request telemetry type: {telemetry_type}. "
+                f"Supported types: {list(cls._registry.keys())}"
+            )
+
+        if use_singleton and telemetry_type in cls._instances:
+            return cls._instances[telemetry_type]
+
+        if config is None:
+            config = {}
+
+        telemetry_cls = cls._registry[telemetry_type]()
+        instance = telemetry_cls(config)
+
+        if use_singleton:
+            cls._instances[telemetry_type] = instance
+
+        return instance
+
+    @classmethod
+    def get_registered_types(cls) -> list[str]:
+        """Get a list of all registered telemetry types."""
+        return list(cls._registry.keys())
+
+
+# Register telemetry types here.
+# The registration should not be done in each individual file, as we want to
+# only load the files corresponding to the current telemetry type.
+
+RequestTelemetryFactory.register(
+    "noop",
+    "lmcache.integration.request_telemetry.noop",
+    "NoOpRequestTelemetry",
+)
+
+RequestTelemetryFactory.register(
+    "fastapi",
+    "lmcache.integration.request_telemetry.fastapi",
+    "FastAPIRequestTelemetry",
+)

lmcache/integration/request_telemetry/fastapi.py

@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+FastAPI-based request telemetry reporter.
+"""
+
+# Standard
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any
+from urllib.error import URLError
+from urllib.request import Request, urlopen
+import json
+import threading
+
+# First Party
+from lmcache.integration.request_telemetry.base import RequestTelemetry
+from lmcache.logging import init_logger
+
+logger = init_logger(__name__)
+
+
+class FastAPIRequestTelemetry(RequestTelemetry):
+    """
+    FastAPI-based request telemetry reporter.
+
+    Sends telemetry events to a FastAPI server endpoint asynchronously
+    using a thread pool to avoid blocking the main execution path.
+
+    Config dict keys:
+        endpoint: The FastAPI endpoint URL to send telemetry events to. (required)
+        timeout: Timeout in seconds for HTTP requests. Defaults to 5.0.
+        max_workers: Maximum number of threads for async HTTP requests.
+            Defaults to 2.
+    """
+
+    def __init__(self, config: dict[str, Any]) -> None:
+        endpoint = config.get("endpoint", None)
+        if endpoint is None:
+            raise ValueError(
+                "FastAPIRequestTelemetry requires setting endpoint. "
+                "Please set LMCACHE_REQUEST_TELEMETRY_ENDPOINT envvar."
+            )
+
+        self._endpoint = endpoint
+        self._timeout = config.get("timeout", 5.0)
+        self._executor = ThreadPoolExecutor(
+            max_workers=config.get("max_workers", 2),
+            thread_name_prefix="fastapi-telemetry",
+        )
+        self._lock = threading.Lock()
+        self._closed = False
+
+        logger.info(
+            "LMCache FastAPIRequestTelemetry initialized with endpoint: %s",
+            self._endpoint,
+        )
+
+    def on_request_store_finished(
+        self,
+        request_ids_set: set[str],
+        model_name: str,
+        world_size: int,
+        kv_rank: int,
+    ) -> None:
+        """Send request store finished event to the FastAPI endpoint."""
+        if self._closed:
+            return
+
+        payload = {
+            "event": "request_store_finished",
+            "request_ids_set": list(request_ids_set),
+            "model_name": model_name,
+            "world_size": world_size,
+            "kv_rank": kv_rank,
+        }
+
+        # Submit async HTTP request
+        self._executor.submit(self._send_event, payload)
+
+    def _send_event(self, payload: dict[str, Any]) -> None:
+        """Send an event to the FastAPI endpoint (runs in thread pool)."""
+        try:
+            data = json.dumps(payload).encode("utf-8")
+            request = Request(
+                self._endpoint,
+                data=data,
+                headers={"Content-Type": "application/json"},
+                method="POST",
+            )
+            # TODO: this introduces blocking IO, should use asyncio.
+            with urlopen(request, timeout=self._timeout) as response:
+                if response.status >= 400:
+                    logger.warning(
+                        f"FastAPI telemetry request failed with"
+                        f" status {response.status}"
+                    )
+        except URLError as e:
+            logger.warning(f"FastAPI telemetry request failed: {e}")
+        except Exception as e:
+            logger.warning(f"FastAPI telemetry request failed unexpectedly: {e}")
+
+    def close(self) -> None:
+        """Shutdown the thread pool executor."""
+        with self._lock:
+            if self._closed:
+                return
+            self._closed = True
+
+        self._executor.shutdown(wait=True)
+        logger.info("FastAPIRequestTelemetry closed")

lmcache/integration/request_telemetry/noop.py

@@ -0,0 +1,35 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+No-op implementation of request telemetry.
+"""
+
+# Standard
+from typing import Any
+
+# First Party
+from lmcache.integration.request_telemetry.base import RequestTelemetry
+
+
+class NoOpRequestTelemetry(RequestTelemetry):
+    """
+    A no-op implementation of RequestTelemetry.
+
+    This implementation does nothing when events are reported.
+    Use this when telemetry is disabled or not needed.
+    """
+
+    def __init__(self, config: dict[str, Any]) -> None:
+        # NoOp doesn't need any config, but accepts it for interface consistency
+        pass
+
+    def on_request_store_finished(
+        self,
+        request_ids_set: set[str],
+        model_name: str,
+        world_size: int,
+        kv_rank: int,
+    ) -> None:
+        pass
+
+    def close(self) -> None:
+        pass

lmcache/integration/sglang/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+

lmcache/integration/sglang/sglang_adapter.py

@@ -0,0 +1,326 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+from dataclasses import dataclass
+from typing import Any, Iterable, List, Optional
+import uuid
+
+# Third Party
+from sglang.srt.configs.model_config import ModelConfig
+import torch
+import torch.distributed as dist
+
+# First Party
+from lmcache.integration.sglang.utils import ENGINE_NAME, lmcache_get_config
+from lmcache.logging import init_logger
+from lmcache.utils import (
+    CacheStoreEvent,
+    EngineType,
+    mock_up_broadcast_fn,
+    mock_up_broadcast_object_fn,
+)
+from lmcache.v1.cache_engine import LMCacheEngine, LMCacheEngineBuilder
+from lmcache.v1.config import LMCacheEngineConfig
+from lmcache.v1.gpu_connector import CreateGPUConnector
+from lmcache.v1.metadata import LMCacheMetadata
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class StoreMetadata:
+    last_node: Any
+    token_ids: List[int]
+    kv_indices: torch.Tensor
+    offset: int
+
+
+@dataclass
+class LoadMetadata:
+    token_ids: List[int]
+    slot_mapping: torch.Tensor
+    offset: int
+
+
+def init_lmcache_engine(
+    model_config: ModelConfig,
+    tp_size: int,
+    local_rank: int,
+    global_rank: int,
+    kv_dtype: torch.dtype,
+) -> LMCacheEngine:
+    """
+    Initialize LMCache engine for SGLang integration.
+
+    Args:
+        model_config: SGLang model configuration
+        tp_size: Tensor parallel size
+        local_rank: Local GPU device index (for device selection)
+        global_rank: Global tensor parallel rank (for metadata)
+        kv_dtype: Data type for KV cache tensors
+    """
+    if curr_engine := LMCacheEngineBuilder.get(ENGINE_NAME):
+        return curr_engine
+
+    config = lmcache_get_config()
+    assert isinstance(config, LMCacheEngineConfig), (
+        "LMCache v1 configuration is should be passed."
+    )
+
+    # construct kv shape (for mem pool)
+    num_layer = model_config.num_hidden_layers
+    chunk_size = config.chunk_size
+    num_kv_head = model_config.get_num_kv_heads(tp_size)
+    head_dim = model_config.head_dim
+
+    kv_shape = (num_layer, 2, chunk_size, num_kv_head, head_dim)
+
+    # Change current device using local GPU index
+    # Use global rank for metadata (tensor parallel rank)
+    metadata = LMCacheMetadata(
+        model_name=model_config.model_path,
+        world_size=tp_size,
+        local_world_size=tp_size,
+        worker_id=global_rank,
+        local_worker_id=local_rank,
+        kv_dtype=kv_dtype,
+        kv_shape=kv_shape,
+    )
+
+    gpu_connector = CreateGPUConnector(config, metadata, EngineType.SGLANG)
+    engine = LMCacheEngineBuilder.get_or_create(
+        ENGINE_NAME,
+        config,
+        metadata,
+        gpu_connector,
+        mock_up_broadcast_fn,
+        mock_up_broadcast_object_fn,
+    )
+
+    return engine
+
+
+class LMCacheConnector:
+    def __init__(
+        self,
+        sgl_config: ModelConfig,
+        tp_size: int,
+        rank: int,
+        k_pool: List[torch.Tensor],
+        v_pool: List[torch.Tensor],
+    ):
+        if not k_pool:
+            raise ValueError("k_pool cannot be empty during initialization.")
+        kv_dtype = k_pool[0].dtype
+        if k_pool[0].is_cuda and k_pool[0].device.index is not None:
+            local_rank = k_pool[0].device.index
+        else:
+            # Fallback for CPU / odd cases
+            local_rank = rank
+
+        # rank is the global tensor parallel rank (tp_rank) from SGLang
+        # local_rank is the local GPU device index
+        self.lmcache_engine = init_lmcache_engine(
+            sgl_config,
+            tp_size,
+            local_rank,
+            rank,  # global_rank (tp_rank) for metadata
+            kv_dtype,
+        )
+        self.sgl_config = sgl_config
+        self.tp_size = tp_size
+        self.rank = local_rank  # Use local_rank for torch.device() calls
+        self.kvcaches = k_pool + v_pool
+        self.num_layer = sgl_config.num_hidden_layers
+
+        self.lmcache_engine.post_init(kvcaches=self.kvcaches)
+
+    ####################
+    # Worker side APIs
+    ####################
+
+    def load_kv(self, load_metadata: LoadMetadata) -> int:
+        token_ids = torch.tensor(load_metadata.token_ids, dtype=torch.int64).cuda()
+        slot_mapping = load_metadata.slot_mapping.cuda()
+        offset = load_metadata.offset
+
+        assert isinstance(token_ids, torch.Tensor)
+        assert isinstance(slot_mapping, torch.Tensor)
+        assert (len(token_ids) - offset) == len(slot_mapping)
+
+        slot_mapping = slot_mapping.cuda()
+        load_mask = torch.ones_like(token_ids, dtype=torch.bool)
+        load_mask[:offset] = False
+
+        ret_token_mask = self.lmcache_engine.retrieve(
+            token_ids,
+            mask=load_mask,
+            kvcaches=self.kvcaches,
+            slot_mapping=slot_mapping,
+            offset=offset,
+        )
+
+        num_retrieved_tokens = ret_token_mask.sum().item()
+
+        return num_retrieved_tokens
+
+    def store_kv(self, store_metadata: StoreMetadata) -> None:
+        token_ids = torch.tensor(store_metadata.token_ids, dtype=torch.int64).cuda()
+        slot_mapping = store_metadata.kv_indices.to(torch.int64).cuda()
+        offset = store_metadata.offset
+
+        assert isinstance(token_ids, torch.Tensor)
+        assert isinstance(slot_mapping, torch.Tensor)
+        assert len(token_ids) == len(slot_mapping)
+
+        slot_mapping = slot_mapping.cuda()
+        store_mask = torch.ones_like(token_ids, dtype=torch.bool)
+
+        self.lmcache_engine.store(
+            token_ids,
+            mask=store_mask,
+            kvcaches=self.kvcaches,
+            slot_mapping=slot_mapping,
+            offset=offset,
+        )
+
+    def get_kv_events(self) -> Iterable[CacheStoreEvent]:
+        if self.lmcache_engine is not None:
+            return self.lmcache_engine.get_kv_events()
+        return []
+
+    def chunk_size(self):
+        return self.lmcache_engine.config.chunk_size
+
+    def reset(self):
+        self.lmcache_engine.clear()
+
+    def close(self):
+        self.lmcache_engine.close()
+
+
+class LMCacheLayerwiseConnector(LMCacheConnector):
+    def __init__(
+        self,
+        sgl_config: ModelConfig,
+        tp_size: int,
+        rank: int,
+        k_pool: List[torch.Tensor],
+        v_pool: List[torch.Tensor],
+        tp_group: Optional[torch.distributed.ProcessGroup] = None,
+    ):
+        super().__init__(sgl_config, tp_size, rank, k_pool, v_pool)
+        self._lmcache_chunk_size = self.lmcache_engine.config.chunk_size
+        self.layerwise_retrievers: List[Any] = []
+        self.layer_load_layer: List[int] = []
+        self.kvcaches = [k_pool, v_pool]
+        self.tp_group = tp_group
+        self.lookup_id_list: List[str] = []
+
+    @torch.no_grad()
+    def global_min_tokens(
+        self, local_tokens: int, tp_group: dist.ProcessGroup, device: torch.device
+    ):
+        # If tensor parallel size is 1, no need for all_reduce
+        if self.tp_size == 1:
+            return local_tokens
+
+        t = torch.tensor([local_tokens], dtype=torch.int32, device=device)
+        dist.all_reduce(t, op=dist.ReduceOp.MIN, group=tp_group)
+        return int(t.item())
+
+    def load_kv_layerwise(self, layer_id: int) -> None:
+        if len(self.layerwise_retrievers) == 0:
+            return
+
+        indices_to_remove = []
+        for i in range(len(self.layerwise_retrievers)):
+            if self.layer_load_layer[i] == layer_id + 1:
+                next(self.layerwise_retrievers[i])
+                self.layer_load_layer[i] += 1
+                if self.layer_load_layer[i] == self.sgl_config.num_hidden_layers:
+                    indices_to_remove.append(i)
+
+        for i in sorted(indices_to_remove, reverse=True):
+            del self.layerwise_retrievers[i]
+            del self.layer_load_layer[i]
+            self.lmcache_engine.lookup_unpin(self.lookup_id_list[i])
+            del self.lookup_id_list[i]
+
+        return
+
+    def start_load_kv(self, load_metadata: LoadMetadata) -> int:
+        token_ids = torch.tensor(load_metadata.token_ids, dtype=torch.int64).cuda()
+        slot_mapping = load_metadata.slot_mapping.cuda()
+        offset = load_metadata.offset
+
+        assert self.lmcache_engine is not None
+
+        load_mask = torch.ones_like(token_ids, dtype=torch.bool)
+        load_mask[:offset] = False
+
+        lookup_id = str(uuid.uuid4())
+        retrieve_token_num = self.lmcache_engine.lookup(
+            token_ids,
+            lookup_id=lookup_id,
+            pin=True,
+        )
+
+        retrieve_token_num = self.global_min_tokens(
+            retrieve_token_num, self.tp_group, torch.device(f"cuda:{self.rank}")
+        )
+
+        # No new tokens to retrieve from LMCache
+        if retrieve_token_num <= offset:
+            self.lmcache_engine.lookup_unpin(lookup_id)
+            logger.info(
+                f"LMCache retrieve skipped: lookup={retrieve_token_num}, "
+                f"offset={offset}, no new tokens to retrieve"
+            )
+            return 0
+
+        layerwise_retriever = self.lmcache_engine.retrieve_layer(
+            token_ids[:retrieve_token_num],
+            mask=load_mask[:retrieve_token_num],
+            kvcaches=self.kvcaches,
+            slot_mapping=slot_mapping[:retrieve_token_num],
+            sync=False,
+        )
+
+        next(layerwise_retriever)
+        # Load First Layer
+        next(layerwise_retriever)
+
+        self.layerwise_retrievers.append(layerwise_retriever)
+        self.layer_load_layer.append(1)
+
+        self.lookup_id_list.append(lookup_id)
+
+        num_new_tokens = retrieve_token_num - offset
+        logger.info(
+            f"LMCache retrieve started: lookup={retrieve_token_num}, "
+            f"offset={offset}, retrieve {num_new_tokens} new tokens"
+        )
+
+        return num_new_tokens
+
+    def store_kv(self, store_metadata: StoreMetadata) -> None:
+        slot_mapping = store_metadata.kv_indices.to(torch.int64).cuda()
+        token_ids = torch.tensor(store_metadata.token_ids, dtype=torch.int64).cuda()
+        store_mask = torch.ones_like(token_ids, dtype=torch.bool)
+
+        lookup_id = str(uuid.uuid4())
+        self.lmcache_engine.lookup(token_ids, lookup_id=lookup_id, pin=True)
+
+        layerwise_storer = self.lmcache_engine.store_layer(
+            token_ids,
+            mask=store_mask,
+            kvcaches=self.kvcaches,
+            slot_mapping=slot_mapping,
+            offset=store_metadata.offset,
+            sync=False,
+        )
+        next(layerwise_storer)
+        for _ in range(self.sgl_config.num_hidden_layers):
+            next(layerwise_storer)
+
+        self.lmcache_engine.lookup_unpin(lookup_id)

lmcache/integration/sglang/utils.py

@@ -0,0 +1,39 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+import os
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.v1.config import LMCacheEngineConfig
+
+logger = init_logger(__name__)
+ENGINE_NAME = "sglang-instance"
+
+
+def is_false(value: str) -> bool:
+    """Check if the given string value is equivalent to 'false'."""
+    return value.lower() in ("false", "0", "no", "n", "off")
+
+
+def lmcache_get_config() -> LMCacheEngineConfig:
+    """Get the LMCache configuration from the environment variable
+    `LMCACHE_CONFIG_FILE`. If the environment variable is not set, this
+    function will return the default configuration.
+    """
+    logger.info(f"LMCACHE_CONFIG_FILE: {os.getenv('LMCACHE_CONFIG_FILE')}")
+    if "LMCACHE_CONFIG_FILE" not in os.environ:
+        logger.warn(
+            "No LMCache configuration file is set. Trying to read"
+            " configurations from the environment variables."
+        )
+        logger.warn(
+            "You can set the configuration file through "
+            "the environment variable: LMCACHE_CONFIG_FILE"
+        )
+        config = LMCacheEngineConfig.from_env()
+    else:
+        config_file = os.environ["LMCACHE_CONFIG_FILE"]
+        logger.info(f"Loading LMCache config file {config_file}")
+        config = LMCacheEngineConfig.from_file(config_file)
+
+    return config

lmcache/integration/vllm/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0