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

lmcache/v1/distributed/l2_adapters/plugin_l2_adapter.py

@@ -0,0 +1,210 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Plugin L2 adapter -- dynamically loads an external adapter
+class from a user-supplied Python module.
+"""
+
+# Future
+from __future__ import annotations
+
+# Standard
+from typing import TYPE_CHECKING, Any, Optional
+import importlib
+
+if TYPE_CHECKING:
+    # First Party
+    from lmcache.v1.distributed.internal_api import (
+        L1MemoryDesc,
+    )
+
+# First Party
+from lmcache.v1.distributed.l2_adapters.base import L2AdapterInterface as _L2AI
+from lmcache.v1.distributed.l2_adapters.config import (
+    L2AdapterConfigBase,
+    register_l2_adapter_type,
+)
+from lmcache.v1.distributed.l2_adapters.factory import (
+    register_l2_adapter_factory,
+)
+
+# Config class
+
+
+class PluginL2AdapterConfig(L2AdapterConfigBase):
+    """
+    Config for a plugin L2 adapter.
+
+    Dynamically loads an adapter class from a user-supplied
+    Python module at creation time.
+
+    Fields:
+    - module_path: Dotted Python import path of the module
+        containing the adapter class.
+    - class_name: Name of the class inside *module_path*
+        that implements ``L2AdapterInterface``.
+    - adapter_params: Arbitrary dict forwarded to the
+        adapter class constructor.
+    - config_class_name: Optional name of a config class
+        inside *module_path* that subclasses
+        ``L2AdapterConfigBase``.  When set, the factory
+        builds a config instance via ``from_dict()`` and
+        passes it (instead of a raw dict) to the adapter
+        constructor -- matching the built-in convention.
+    """
+
+    def __init__(
+        self,
+        module_path: str,
+        class_name: str,
+        adapter_params: dict[str, Any] | None = None,
+        config_class_name: str | None = None,
+    ):
+        self.module_path = module_path
+        self.class_name = class_name
+        self.adapter_params = adapter_params or {}
+        self.config_class_name = config_class_name
+
+    @classmethod
+    def from_dict(cls, d: dict) -> "PluginL2AdapterConfig":
+        module_path = d.get("module_path")
+        if not isinstance(module_path, str) or not module_path:
+            raise ValueError("module_path must be a non-empty string")
+
+        class_name = d.get("class_name")
+        if not isinstance(class_name, str) or not class_name:
+            raise ValueError("class_name must be a non-empty string")
+
+        adapter_params = d.get("adapter_params", {})
+        if not isinstance(adapter_params, dict):
+            raise ValueError("adapter_params must be a dict")
+
+        config_class_name = d.get("config_class_name")
+        if config_class_name is not None and not isinstance(config_class_name, str):
+            raise ValueError("config_class_name must be a string")
+
+        return cls(
+            module_path=module_path,
+            class_name=class_name,
+            adapter_params=adapter_params,
+            config_class_name=config_class_name,
+        )
+
+    @classmethod
+    def help(cls) -> str:
+        return (
+            "Plugin L2 adapter config fields:\n"
+            "- module_path (str): dotted import path of "
+            "the module containing the adapter class "
+            "(required)\n"
+            "- class_name (str): name of the adapter "
+            "class inside the module (required)\n"
+            "- adapter_params (dict): forwarded to the "
+            "adapter constructor "
+            "(optional, default {})\n"
+            "- config_class_name (str): explicit config "
+            "class name; when omitted the factory "
+            "auto-discovers it (see plugin.md)\n"
+            "\n"
+            "Example JSON (raw dict):\n"
+            '{"type": "plugin", '
+            '"module_path": "my_plugin.l2", '
+            '"class_name": "MyL2Adapter", '
+            '"adapter_params": {"host": "localhost"}}\n'
+            "\n"
+            "Example JSON (with config class):\n"
+            '{"type": "plugin", '
+            '"module_path": "my_plugin.l2", '
+            '"class_name": "MyL2Adapter", '
+            '"config_class_name": "MyL2AdapterConfig", '
+            '"adapter_params": {"host": "localhost"}}'
+        )
+
+
+# Factory function
+
+
+def _create_plugin_adapter(
+    config: L2AdapterConfigBase,
+    l1_memory_desc: "Optional[L1MemoryDesc]" = None,
+) -> _L2AI:
+    """Dynamically load and create a plugin L2 adapter."""
+    assert isinstance(config, PluginL2AdapterConfig)
+
+    try:
+        module = importlib.import_module(config.module_path)
+    except ImportError as e:
+        raise ImportError(
+            "Could not import module '%s': %s" % (config.module_path, e)
+        ) from e
+
+    try:
+        adapter_cls = getattr(module, config.class_name)
+    except AttributeError as e:
+        raise AttributeError(
+            "Module '%s' has no class '%s': %s"
+            % (config.module_path, config.class_name, e)
+        ) from e
+
+    if not (isinstance(adapter_cls, type) and issubclass(adapter_cls, _L2AI)):
+        raise TypeError(
+            "%s.%s is not a subclass of "
+            "L2AdapterInterface" % (config.module_path, config.class_name)
+        )
+
+    cfg_cls = _resolve_config_class(
+        module,
+        config,
+        adapter_cls,
+    )
+    kwargs: dict[str, object] = {}
+    if l1_memory_desc is not None:
+        kwargs["l1_memory_desc"] = l1_memory_desc
+
+    if cfg_cls is not None:
+        return adapter_cls(  # type: ignore[call-arg]
+            cfg_cls.from_dict(config.adapter_params),
+            **kwargs,
+        )
+
+    return adapter_cls(  # type: ignore[call-arg]
+        config.adapter_params,
+        **kwargs,
+    )
+
+
+def _resolve_config_class(
+    module: object,
+    config: PluginL2AdapterConfig,
+    adapter_cls: type,
+) -> type[L2AdapterConfigBase] | None:
+    """Resolve the config class for a plugin adapter.
+
+    Discovery order:
+    1. Explicit ``config.config_class_name`` field.
+    2. Convention: ``config.class_name`` + ``"Config"``.
+    3. ``config_class_name`` attribute on the adapter
+       class itself.
+    4. ``None`` -- fall back to raw dict mode.
+    """
+    candidates: list[str | None] = [
+        config.config_class_name,
+        config.class_name + "Config",
+        getattr(adapter_cls, "config_class_name", None),
+    ]
+    for name in candidates:
+        if not name:
+            continue
+        cls = getattr(module, name, None)
+        if cls is None:
+            continue
+        if isinstance(cls, type) and issubclass(
+            cls,
+            L2AdapterConfigBase,
+        ):
+            return cls
+    return None
+
+
+# Self-register config type and adapter factory
+register_l2_adapter_type("plugin", PluginL2AdapterConfig)
+register_l2_adapter_factory("plugin", _create_plugin_adapter)

lmcache/v1/distributed/l2_adapters/resp_l2_adapter.py

@@ -0,0 +1,176 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+RESP (Redis/Valkey) L2 adapter config and factory.
+
+Backed by the native C++ Redis connector wrapped with
+``NativeConnectorL2Adapter``.
+"""
+
+# Future
+from __future__ import annotations
+
+# Standard
+from typing import TYPE_CHECKING, Optional
+import os
+
+if TYPE_CHECKING:
+    from lmcache.v1.distributed.internal_api import (
+        L1MemoryDesc,
+    )
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.v1.distributed.l2_adapters.base import (
+    L2AdapterInterface,
+)
+from lmcache.v1.distributed.l2_adapters.config import (
+    L2AdapterConfigBase,
+    register_l2_adapter_type,
+)
+from lmcache.v1.distributed.l2_adapters.factory import (
+    register_l2_adapter_factory,
+)
+
+logger = init_logger(__name__)
+
+
+class RESPL2AdapterConfig(L2AdapterConfigBase):
+    """
+    Config for an L2 adapter backed by a native RESP
+    connector (Redis/Valkey).
+
+    Fields:
+    - host: server hostname or IP.
+    - port: server port.
+    - num_workers: C++ worker threads for I/O (default 8).
+    - username: optional auth username.
+    - password: optional auth password.
+    """
+
+    def __init__(
+        self,
+        host: str,
+        port: int,
+        num_workers: int = 8,
+        username: str = "",
+        password: str = "",
+        max_capacity_gb: float = 0,
+    ):
+        super().__init__()
+        self.host = host
+        self.port = port
+        self.num_workers = num_workers
+        self.username = username
+        self.password = password
+        self.max_capacity_gb = max_capacity_gb
+
+    @classmethod
+    def from_dict(cls, d: dict) -> "RESPL2AdapterConfig":
+        host = d.get("host")
+        if not isinstance(host, str) or not host:
+            raise ValueError("host must be a non-empty string")
+
+        port = d.get("port")
+        if not isinstance(port, int) or port <= 0:
+            raise ValueError("port must be a positive integer")
+
+        num_workers = d.get("num_workers", 8)
+        if not isinstance(num_workers, int) or num_workers <= 0:
+            raise ValueError("num_workers must be a positive integer")
+
+        username = d.get("username", "")
+        password = d.get("password", "")
+
+        max_capacity_gb = d.get("max_capacity_gb", 0)
+        if not isinstance(max_capacity_gb, (int, float)) or max_capacity_gb < 0:
+            raise ValueError("max_capacity_gb must be a non-negative number")
+
+        return cls(
+            host=host,
+            port=port,
+            num_workers=num_workers,
+            username=str(username),
+            password=str(password),
+            max_capacity_gb=float(max_capacity_gb),
+        )
+
+    @classmethod
+    def help(cls) -> str:
+        return (
+            "RESP L2 adapter config fields:\n"
+            "- host (str): Redis/Valkey server hostname "
+            "or IP (required)\n"
+            "- port (int): server port (required, >0)\n"
+            "- num_workers (int): C++ worker threads "
+            "for I/O (default 8, >0)\n"
+            "- username (str): auth username "
+            "(default empty)\n"
+            "- password (str): auth password "
+            "(default empty)\n"
+            "- max_capacity_gb (float): max L2 capacity "
+            "in GB for usage tracking / eviction "
+            "(default 0 = disabled)\n\n"
+            "Environment variable defaults (used when "
+            "config value is empty, read at adapter "
+            "creation, not stored in config):\n"
+            "- LMCACHE_RESP_USERNAME: default username\n"
+            "- LMCACHE_RESP_PASSWORD: default password\n"
+            "- LMCACHE_RESP_HOST: default host\n"
+            "- LMCACHE_RESP_PORT: default port"
+        )
+
+
+def _create_resp_l2_adapter(
+    config: L2AdapterConfigBase,
+    l1_memory_desc: "Optional[L1MemoryDesc]" = None,
+) -> L2AdapterInterface:
+    """Create a NativeConnectorL2Adapter backed by the
+    C++ Redis connector."""
+    try:
+        # First Party
+        from lmcache.lmcache_redis import (
+            LMCacheRedisClient,
+        )
+    except ImportError as e:
+        raise RuntimeError(
+            "RESP L2 adapter requires the C++ Redis "
+            "extension. Build with: pip install -e ."
+        ) from e
+
+    # Lazy import to avoid circular dependency
+    # First Party
+    from lmcache.v1.distributed.l2_adapters.native_connector_l2_adapter import (  # noqa: E501
+        NativeConnectorL2Adapter,
+    )
+
+    assert isinstance(config, RESPL2AdapterConfig)
+
+    # Config/CLI args take precedence over environment variables,
+    # which serve as defaults. This keeps secrets out of logged
+    # config while allowing explicit CLI overrides.
+    host = config.host or os.environ.get("LMCACHE_RESP_HOST", "")
+    port = config.port if config.port else int(os.environ.get("LMCACHE_RESP_PORT", "0"))
+    username = config.username or os.environ.get("LMCACHE_RESP_USERNAME", "")
+    password = config.password or os.environ.get("LMCACHE_RESP_PASSWORD", "")
+
+    native_client = LMCacheRedisClient(
+        host,
+        port,
+        config.num_workers,
+        username,
+        password,
+    )
+    logger.info(
+        "Created RESP L2 adapter: %s:%d (workers=%d)",
+        host,
+        port,
+        config.num_workers,
+    )
+    return NativeConnectorL2Adapter(
+        native_client, max_capacity_gb=config.max_capacity_gb
+    )
+
+
+# Self-register config type and adapter factory
+register_l2_adapter_type("resp", RESPL2AdapterConfig)
+register_l2_adapter_factory("resp", _create_resp_l2_adapter)

lmcache/v1/distributed/memory_manager.py

@@ -0,0 +1,179 @@
+# SPDX-License-Identifier: Apache-2.0
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.v1.distributed.api import MemoryLayoutDesc
+from lmcache.v1.distributed.config import L1MemoryManagerConfig
+from lmcache.v1.distributed.error import L1Error
+from lmcache.v1.distributed.internal_api import L1MemoryDesc
+from lmcache.v1.lazy_memory_allocator import LazyMemoryAllocator
+from lmcache.v1.memory_management import (
+    MemoryAllocatorInterface,
+    MemoryObj,
+    MixedMemoryAllocator,
+)
+
+logger = init_logger(__name__)
+
+
+# HELPER FUNCTIONS
+def create_memory_allocator(config: L1MemoryManagerConfig) -> MemoryAllocatorInterface:
+    """
+    Create a memory allocator based on the provided configuration.
+
+    Args:
+        config (L1MemoryManagerConfig): Configuration for the memory manager.
+
+    Returns:
+        MemoryAllocatorInterface: An instance of a memory allocator.
+    """
+    if config.use_lazy:
+        logger.debug(
+            "use lazy memory allocator, init size is %d bytes, "
+            "final size is %d bytes, align bytes is %d bytes",
+            config.init_size_in_bytes,
+            config.size_in_bytes,
+            config.align_bytes,
+        )
+        return LazyMemoryAllocator(
+            config.init_size_in_bytes, config.size_in_bytes, config.align_bytes
+        )
+    else:
+        logger.debug(
+            "use mixed memory allocator, total size is %d bytes, "
+            "align bytes is %d bytes",
+            config.size_in_bytes,
+            config.align_bytes,
+        )
+        return MixedMemoryAllocator(
+            config.size_in_bytes,
+            align_bytes=config.align_bytes,
+        )
+
+
+# MAIN CLASS
+class L1MemoryManager:
+    """
+    L1MemoryManager manages the allocation and deallocation of L1 memory.
+
+    Observability metrics to emit:
+    1. Memory usage
+    2. Active allocations
+    """
+
+    def __init__(self, config: L1MemoryManagerConfig):
+        self._allocator = create_memory_allocator(config)
+        self._size_in_bytes = config.size_in_bytes
+        self._align_bytes = config.align_bytes
+
+    def allocate(
+        self, layout_desc: MemoryLayoutDesc, count: int
+    ) -> tuple[L1Error, list[MemoryObj]]:
+        """
+        Allocate memory objects based on the provided layout description and count.
+        This function should be thread-safe
+
+        Args:
+            layout_desc (MemoryLayoutDesc): Description of the memory layout.
+            count (int): Number of memory objects to allocate.
+
+        Returns:
+            tuple[L1Error, list[MemoryObj]]: Error code and list of
+            allocated memory objects.
+            Error code will be `L1Error.OUT_OF_MEMORY` if allocation
+            fails; otherwise, it will be `L1Error.SUCCESS`.
+
+        Note:
+            If the allocation fails, the memory object list will be empty.
+        """
+        objects = self._allocator.batched_allocate(
+            layout_desc.shapes, layout_desc.dtypes, count
+        )
+        if objects is None:
+            return L1Error.OUT_OF_MEMORY, []
+        return L1Error.SUCCESS, objects
+
+    def free(self, mem_objs: list[MemoryObj]) -> L1Error:
+        """
+        Free the provided memory objects.
+        This function should be thread-safe.
+
+        Args:
+            mem_objs (list[MemoryObj]): List of memory objects to free.
+
+        Returns:
+            L1Error: Error code indicating the result of the operation.
+            It will be `L1Error.SUCCESS` if the operation succeeds.
+        """
+        self._allocator.batched_free(mem_objs)
+        return L1Error.SUCCESS
+
+    def get_memory_usage(self) -> tuple[int, int]:
+        """
+        Get the current memory usage. This function will mainly be used to support
+        eviction decision.
+
+        Returns:
+            tuple[int, int]: A tuple containing used memory in bytes and total memory
+            in bytes.
+
+        Note:
+            In the future, we may want to make a "callback" based mechanism to
+            trigger eviction when the memory usage reaches a watermark.
+        """
+
+        # HACK: now trying to read this from the address manager in a ad-hoc
+        # manner
+        def get_address_manager(allocator: MemoryAllocatorInterface):
+            if isinstance(allocator, MixedMemoryAllocator) and hasattr(
+                allocator.pin_allocator, "address_manager"
+            ):
+                return allocator.pin_allocator.address_manager
+            elif isinstance(allocator, LazyMemoryAllocator):
+                return allocator.get_address_manager()
+            else:
+                raise NotImplementedError(
+                    "get_memory_usage is not implemented for this allocator type."
+                )
+
+        address_manager = get_address_manager(self._allocator)
+        free_size = address_manager.get_free_size()
+        total_size = address_manager.get_heap_size()
+        used_size = total_size - free_size
+        return used_size, total_size
+
+    def get_l1_memory_desc(self) -> L1MemoryDesc:
+        """
+        Return an L1MemoryDesc describing the underlying memory buffer.
+
+        Returns:
+            L1MemoryDesc: Pointer, size, and alignment of the L1 buffer.
+
+        Raises:
+            NotImplementedError: If the allocator type does not support this operation.
+        """
+        if isinstance(self._allocator, MixedMemoryAllocator):
+            buffer = self._allocator.buffer
+        elif isinstance(self._allocator, LazyMemoryAllocator):
+            # TODO(ApostaC): need to test if the RDMA registration works
+            # before the lazy expansion is finished
+            buffer = self._allocator.get_underlying_buffer()
+        else:
+            raise NotImplementedError(
+                "get_l1_memory_desc is not implemented for this allocator type."
+            )
+        return L1MemoryDesc(
+            ptr=buffer.data_ptr(),
+            size=self._size_in_bytes,
+            align_bytes=self._align_bytes,
+        )
+
+    def close(self) -> None:
+        """
+        Close the memory manager and release all resources.
+        """
+        self._allocator.close()
+
+    # Debugging APIs
+    def memcheck(self):
+        return self._allocator.memcheck()

lmcache/v1/distributed/storage_controller.py

@@ -0,0 +1,39 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Interface for storage controllers
+
+Storage controllers are separate modules/threads that sees the L1 Manager
+and can operate on it.
+"""
+
+# Standard
+from abc import ABC, abstractmethod
+
+
+class StorageControllerInterface(ABC):
+    @abstractmethod
+    def start(self):
+        """
+        Start the storage controller.
+        This function should be implemented by subclasses to start
+        any necessary threads or processes.
+        """
+        pass
+
+    @abstractmethod
+    def stop(self):
+        """
+        Stop the storage controller.
+        This function should be implemented by subclasses to stop
+        any running threads or processes.
+        """
+        pass
+
+    @abstractmethod
+    def report_status(self) -> dict:
+        """
+        Return a status dict for this controller.
+
+        Must include at least ``is_healthy: bool``.
+        """
+        pass

lmcache/v1/distributed/storage_controllers/__init__.py

@@ -0,0 +1,43 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Storage controllers with automatic module discovery.
+
+All modules under this package that call ``register_store_policy`` or
+``register_prefetch_policy`` at import time are discovered automatically.
+To add a new policy, simply create a new module in this directory -- **no
+changes to any existing file are needed**.
+"""
+
+# Standard
+import importlib
+import pkgutil
+
+# First Party
+from lmcache.v1.distributed.storage_controllers.eviction_controller import (
+    EvictionController,
+    L1EvictionController,
+    L2AdapterEvictionState,
+    L2EvictionController,
+)
+from lmcache.v1.distributed.storage_controllers.prefetch_controller import (
+    PrefetchController,
+)
+from lmcache.v1.distributed.storage_controllers.store_controller import (
+    StoreController,
+)
+
+# Auto-discover and import every module in this package so
+# that each policy's self-registration code runs.
+_PACKAGE_PATH = __path__  # type: ignore[name-defined]
+_PACKAGE_NAME = __name__
+for _finder, _mod_name, _is_pkg in pkgutil.iter_modules(_PACKAGE_PATH):
+    importlib.import_module(f".{_mod_name}", _PACKAGE_NAME)
+
+__all__ = [
+    "EvictionController",
+    "L1EvictionController",
+    "L2EvictionController",
+    "L2AdapterEvictionState",
+    "PrefetchController",
+    "StoreController",
+]