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

lmcache/v1/distributed/l2_adapters/config.py

@@ -0,0 +1,385 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Configuration for L2 adapters.
+
+Supports multiple adapter instances (including multiple instances of the same
+adapter type with different configs) via repeatable --l2-adapter <JSON>.
+Each JSON object must include "type" (adapter type name) and type-specific keys.
+"""
+
+# Future
+from __future__ import annotations
+
+# Standard
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, TypeVar
+import argparse
+import json
+
+if TYPE_CHECKING:
+    # First Party
+    from lmcache.v1.distributed.config import EvictionConfig
+
+# First Party
+from lmcache.logging import init_logger
+
+logger = init_logger(__name__)
+
+T = TypeVar("T", bound="L2AdapterConfigBase")
+
+
+@dataclass(frozen=True)
+class PersistConfig:
+    """
+    Configuration for persist on an L2 adapter.
+
+    When enabled, data files are kept on disk at shutdown instead of
+    deleted. Lookup always checks secondary storage (disk) on miss
+    regardless of this setting.
+    """
+
+    persist_enabled: bool = True
+    """ If True, data files are kept on disk at shutdown instead of deleted. """
+
+
+# -----------------------------------------------------------------------------
+# Registry: adapter type name -> config class
+# -----------------------------------------------------------------------------
+
+_L2_ADAPTER_CONFIG_REGISTRY: dict[str, type[L2AdapterConfigBase]] = {}
+
+
+def register_l2_adapter_type(
+    name: str,
+    config_cls: type[L2AdapterConfigBase],
+) -> None:
+    """
+    Register an L2 adapter config class under a type name.
+
+    The type name is used in JSON specs as the "type" field.
+    Each adapter config module should call this at import time.
+
+    Args:
+        name: Adapter type name (e.g. "fs", "mock").
+        config_cls: Config class that can parse from dict
+            via ``from_dict()``.
+    """
+    if name in _L2_ADAPTER_CONFIG_REGISTRY:
+        raise ValueError("L2 adapter type already registered: %s" % repr(name))
+    _L2_ADAPTER_CONFIG_REGISTRY[name] = config_cls
+
+
+def _ensure_config_loaded(name: str) -> None:
+    """Trigger lazy import for *name* if it is not
+    yet in the config registry.
+
+    Raises:
+        ImportError: If the adapter module cannot be
+            imported (missing dependency).
+    """
+    if name in _L2_ADAPTER_CONFIG_REGISTRY:
+        return
+    # Lazy import lives in factory to avoid circular deps
+    # First Party
+    from lmcache.v1.distributed.l2_adapters.factory import (  # noqa: PLC0415
+        ensure_adapter_loaded,
+    )
+
+    ensure_adapter_loaded(name)
+
+
+def get_registered_l2_adapter_types() -> list[str]:
+    """Return all known adapter type names (eager
+    and lazy)."""
+    # First Party
+    from lmcache.v1.distributed.l2_adapters.factory import (  # noqa: PLC0415
+        get_all_registered_names,
+    )
+
+    return get_all_registered_names()
+
+
+def get_type_name_for_config(
+    config: L2AdapterConfigBase,
+) -> str:
+    """
+    Reverse-lookup the registered type name for a config
+    instance.
+
+    Args:
+        config: An L2 adapter config instance.
+
+    Returns:
+        The registered type name (e.g., "mock", "fs").
+
+    Raises:
+        ValueError: If the config's class is not registered.
+    """
+    for name, cls in _L2_ADAPTER_CONFIG_REGISTRY.items():
+        if type(config) is cls:
+            return name
+    raise ValueError("Unregistered L2 adapter config type: %s" % type(config).__name__)
+
+
+# -----------------------------------------------------------------------------
+# Base config class for a single L2 adapter
+# -----------------------------------------------------------------------------
+
+
+class L2AdapterConfigBase(ABC):
+    """
+    Base class for per-adapter configs.
+
+    Each adapter type (e.g. disk, redis) defines a config class that:
+    - Subclasses this base.
+    - Implements from_dict() to parse a dict (from JSON) into an instance.
+    - Is registered via register_l2_adapter_type("type_name", ConfigClass).
+
+    An optional ``"eviction"`` key in the JSON dict enables per-adapter L2
+    eviction. When present it is parsed by ``_parse_eviction_config()`` and
+    stored on ``eviction_config``; the L2EvictionController is then created
+    for this adapter in the StorageManager.
+    """
+
+    #: Populated by ``_parse_eviction_config`` after ``from_dict``; ``None``
+    #: means L2 eviction is disabled for this adapter.
+    eviction_config: EvictionConfig | None = None
+
+    #: Populated by ``_parse_persist_config`` after ``from_dict``.
+    #: Defaults to ``PersistConfig()`` (persist enabled).
+    persist_config: PersistConfig = PersistConfig()
+
+    @staticmethod
+    def _parse_eviction_config(d: dict) -> EvictionConfig | None:
+        """
+        Parse an optional ``"eviction"`` sub-dict from an adapter JSON spec.
+
+        Expected format::
+
+            {
+                "type": "mock",
+                ...
+                "eviction": {
+                    "eviction_policy": "LRU",
+                    "trigger_watermark": 0.8,
+                    "eviction_ratio": 0.2
+                }
+            }
+
+        Returns ``None`` when the key is absent (eviction disabled).
+        """
+        eviction_dict = d.get("eviction")
+        if eviction_dict is None:
+            return None
+
+        # Lazy import to avoid circular dependency:
+        # l2_adapters/config.py <- config.py <- l2_adapters/config.py
+        # First Party
+        from lmcache.v1.distributed.config import EvictionConfig  # noqa: PLC0415
+
+        policy = eviction_dict.get("eviction_policy")
+        if policy not in ("LRU", "noop"):
+            raise ValueError(
+                f"eviction.eviction_policy must be 'LRU' or 'noop', got {policy!r}"
+            )
+        return EvictionConfig(
+            eviction_policy=policy,
+            trigger_watermark=float(eviction_dict.get("trigger_watermark", 0.8)),
+            eviction_ratio=float(eviction_dict.get("eviction_ratio", 0.2)),
+        )
+
+    @staticmethod
+    def _parse_persist_config(d: dict) -> PersistConfig:
+        """
+        Parse optional ``"persist_enabled"`` key from an adapter JSON spec.
+
+        Defaults to ``True``.
+
+        Expected format::
+
+            {
+                "type": "nixl_store_dynamic",
+                ...
+                "persist_enabled": true
+            }
+        """
+        persist_enabled = bool(d.get("persist_enabled", True))
+        return PersistConfig(persist_enabled=persist_enabled)
+
+    @classmethod
+    @abstractmethod
+    def from_dict(cls: type[T], d: dict) -> T:
+        """
+        Build a config instance from a dict (e.g. from parsed JSON).
+
+        The dict will contain the "type" key used for dispatch; the concrete
+        class may ignore it. All other keys are type-specific.
+
+        Args:
+            d: Adapter spec dict (must include type-specific keys).
+
+        Returns:
+            An instance of the config class.
+
+        Raises:
+            ValueError: If required keys are missing or values are invalid.
+        """
+        ...
+
+    @classmethod
+    @abstractmethod
+    def help(cls) -> str:
+        """
+        Return a help string describing the config fields for this adapter type.
+
+        This is used in command-line help to explain the expected JSON format for
+        each adapter type.
+
+        Returns:
+            A help string describing the config fields for this adapter type.
+        """
+        ...
+
+
+# -----------------------------------------------------------------------------
+# Main config: list of adapter configs (order = adapter order)
+# -----------------------------------------------------------------------------
+
+
+@dataclass
+class L2AdaptersConfig:
+    """
+    Main config for L2 adapters.
+
+    Holds an ordered list of adapter configs. Each element corresponds to one
+    L2 adapter instance (e.g. two disk adapters with different paths appear
+    as two entries).
+    """
+
+    adapters: list[L2AdapterConfigBase]
+    """ Ordered list of adapter configs; one per L2 adapter instance. """
+
+
+# -----------------------------------------------------------------------------
+# Command-line: add args and parse to config
+# -----------------------------------------------------------------------------
+
+_L2_ADAPTER_ARG_DEST = "l2_adapter"
+
+
+def add_l2_adapters_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
+    """
+    Add L2 adapter configuration arguments to an existing parser.
+
+    Adds a repeatable --l2-adapter <JSON> argument. Each JSON object specifies
+    one adapter: it must include "type" (registered adapter type name) and
+    type-specific keys. Order of arguments is the order of adapters.
+
+    Args:
+        parser: The argument parser to add arguments to.
+
+    Returns:
+        The same parser with L2 adapter arguments added.
+
+    Example:
+        >>> parser = argparse.ArgumentParser()
+        >>> add_l2_adapters_args(parser)
+        >>> args = parser.parse_args(["--l2-adapter", '{"type":"disk","path":"/data"}'])
+        >>> config = parse_args_to_l2_adapters_config(args)
+    """
+    group = parser.add_argument_group(
+        "L2 Adapters",
+        "L2 adapter instances. Each --l2-adapter is a JSON object with 'type' and "
+        "type-specific keys. Repeat for multiple adapters.",
+    )
+    group.add_argument(
+        "--l2-adapter",
+        dest=_L2_ADAPTER_ARG_DEST,
+        action="append",
+        default=[],
+        type=str,
+        metavar="JSON",
+        help='Adapter spec as JSON with a "type" field and adapter-specific configs'
+        ', e.g. \'{"type":"disk","path":"/data"}\'.'
+        "Repeat for multiple adapters."
+        "Supported adapters are: ["
+        + ", ".join(sorted(get_registered_l2_adapter_types()))
+        + "].",
+    )
+    return parser
+
+
+def parse_args_to_l2_adapters_config(args: argparse.Namespace) -> L2AdaptersConfig:
+    """
+    Build L2AdaptersConfig from parsed command-line arguments.
+
+    Expects args to have the attribute added by add_l2_adapters_args (a list
+    of JSON strings). Each string is parsed; the "type" field selects the
+    config class from the registry, and from_dict() builds the config instance.
+
+    Args:
+        args: Parsed arguments (e.g. from parser.parse_args()).
+
+    Returns:
+        L2AdaptersConfig with one entry per --l2-adapter argument.
+
+    Raises:
+        KeyError: If an adapter "type" is not registered.
+        ValueError: If JSON is invalid or a config class raises from_dict().
+    """
+    raw_list = getattr(args, _L2_ADAPTER_ARG_DEST, None)
+    if raw_list is None:
+        raw_list = []
+
+    adapter_configs: list[L2AdapterConfigBase] = []
+    for i, raw in enumerate(raw_list):
+        try:
+            d = json.loads(raw)
+        except json.JSONDecodeError as e:
+            raise ValueError(f"Invalid JSON for --l2-adapter #{i + 1}: {e}") from e
+
+        if not isinstance(d, dict):
+            raise ValueError(
+                f"--l2-adapter #{i + 1}: expected a JSON object, got {type(d).__name__}"
+            )
+
+        type_name = d.get("type")
+        if type_name is None:
+            raise ValueError("--l2-adapter #%d: missing 'type' field" % (i + 1))
+
+        # Trigger lazy import for this adapter type
+        _ensure_config_loaded(type_name)
+
+        if type_name not in _L2_ADAPTER_CONFIG_REGISTRY:
+            known = ", ".join(sorted(_L2_ADAPTER_CONFIG_REGISTRY)) or "(none)"
+            raise ValueError(
+                f"--l2-adapter #{i + 1}: unknown adapter type "
+                f"{type_name!r}. Known: {known}"
+            )
+
+        config_cls = _L2_ADAPTER_CONFIG_REGISTRY[type_name]
+        try:
+            adapter_cfg = config_cls.from_dict(d)
+            adapter_cfg.eviction_config = L2AdapterConfigBase._parse_eviction_config(d)
+            adapter_cfg.persist_config = L2AdapterConfigBase._parse_persist_config(d)
+            adapter_configs.append(adapter_cfg)
+        except (TypeError, ValueError) as e:
+            logger.error(
+                "Error parsing --l2-adapter #%d (type %r): %s",
+                i + 1,
+                type_name,
+                e,
+            )
+            logger.error(
+                "Adapter config help for %s adapter:\n"
+                "---------------------\n"
+                "%s\n"
+                "---------------------\n\n",
+                type_name,
+                config_cls.help(),
+            )
+            raise ValueError(f"--l2-adapter #{i + 1} ({type_name!r}): {e}") from e
+
+    return L2AdaptersConfig(adapters=adapter_configs)

lmcache/v1/distributed/l2_adapters/factory.py

@@ -0,0 +1,205 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Factory registry for L2 adapters.
+
+Each adapter module self-registers a factory callable via
+``register_l2_adapter_factory`` at import time.  The factory
+signature is
+``(L2AdapterConfigBase, Optional[L1MemoryDesc]) -> L2AdapterInterface``.
+
+Built-in adapters are discovered via ``pkgutil`` but only
+imported on demand -- their modules (and third-party
+dependencies) are loaded only when the adapter type is
+actually requested.
+"""
+
+# Future
+from __future__ import annotations
+
+# Standard
+from typing import TYPE_CHECKING, Callable, Optional
+import importlib
+
+if TYPE_CHECKING:
+    from lmcache.v1.distributed.internal_api import (
+        L1MemoryDesc,
+    )
+    from lmcache.v1.distributed.l2_adapters.base import (
+        L2AdapterInterface,
+    )
+    from lmcache.v1.distributed.l2_adapters.config import (
+        L2AdapterConfigBase,
+    )
+
+# First Party
+from lmcache.logging import init_logger
+
+logger = init_logger(__name__)
+
+# Type alias for factory callables:
+#   (config, l1_memory_desc) -> L2AdapterInterface
+L2AdapterFactory = Callable[
+    ["L2AdapterConfigBase", "Optional[L1MemoryDesc]"],
+    "L2AdapterInterface",
+]
+
+# -----------------------------------------------------------------
+# Registry: adapter type name -> factory callable
+# -----------------------------------------------------------------
+
+_L2_ADAPTER_FACTORY_REGISTRY: dict[str, L2AdapterFactory] = {}
+
+# -----------------------------------------------------------------
+# Pending modules: fully-qualified module paths that have
+# been discovered (e.g. via pkgutil) but not yet imported.
+# When a type name is not found in the eager registry we
+# import these one-by-one until the name appears.
+# -----------------------------------------------------------------
+
+_PENDING_MODULES: list[str] = []
+
+
+def register_l2_adapter_factory(
+    name: str,
+    factory: L2AdapterFactory,
+) -> None:
+    """Register an adapter factory for the given type
+    name.
+
+    Each adapter module should call this at import time
+    **after** its config class has been registered via
+    ``register_l2_adapter_type``.
+
+    Args:
+        name: Adapter type name (must match the name used
+            in ``register_l2_adapter_type``).
+        factory: ``(config, l1_memory_desc)``
+            -> ``L2AdapterInterface``.
+    """
+    if name in _L2_ADAPTER_FACTORY_REGISTRY:
+        raise ValueError("L2 adapter factory already registered: %s" % name)
+    _L2_ADAPTER_FACTORY_REGISTRY[name] = factory
+
+
+def add_pending_module(module_path: str) -> None:
+    """Register a module path for deferred import.
+
+    The module will only be imported when a type name
+    lookup misses the eager registry.  Importing the
+    module triggers its ``register_l2_adapter_type``
+    and ``register_l2_adapter_factory`` calls.
+
+    Args:
+        module_path: Fully-qualified module path
+            (e.g. ``"lmcache.v1.distributed.
+            l2_adapters.fs_l2_adapter"``).
+    """
+    if module_path not in _PENDING_MODULES:
+        _PENDING_MODULES.append(module_path)
+
+
+def ensure_adapter_loaded(name: str) -> None:
+    """Import pending modules until *name* appears in
+    the factory registry (or all pending modules have
+    been tried).
+
+    Raises:
+        ImportError: If a module fails to import due to
+            a missing third-party dependency **and** it
+            was the last candidate.
+    """
+    if name in _L2_ADAPTER_FACTORY_REGISTRY:
+        return
+
+    last_err: ImportError | None = None
+    while _PENDING_MODULES:
+        mod_path = _PENDING_MODULES.pop(0)
+        try:
+            importlib.import_module(mod_path)
+        except ImportError as exc:
+            logger.debug(
+                "Skipping module %s (import failed: %s)",
+                mod_path,
+                exc,
+            )
+            last_err = exc
+            continue
+        # Check if the name is now registered
+        if name in _L2_ADAPTER_FACTORY_REGISTRY:
+            return
+
+    # If we exhausted all pending modules and still
+    # didn't find the name, the last ImportError (if
+    # any) might be the root cause.
+    if last_err is not None and name not in _L2_ADAPTER_FACTORY_REGISTRY:
+        raise last_err
+
+
+def load_all_adapters() -> None:
+    """Force-import all pending adapter modules.
+
+    Useful for CLI help that needs to list every
+    registered type name.
+    """
+    while _PENDING_MODULES:
+        mod_path = _PENDING_MODULES.pop(0)
+        try:
+            importlib.import_module(mod_path)
+        except ImportError:
+            logger.debug(
+                "Skipping module %s during bulk load",
+                mod_path,
+            )
+
+
+def get_all_registered_names() -> list[str]:
+    """Return all known adapter type names.
+
+    Forces import of all pending modules so that the
+    returned list is complete.
+    """
+    load_all_adapters()
+    return sorted(_L2_ADAPTER_FACTORY_REGISTRY)
+
+
+def create_l2_adapter_from_registry(
+    config: "L2AdapterConfigBase",
+    l1_memory_desc: "Optional[L1MemoryDesc]" = None,
+) -> "L2AdapterInterface":
+    """Create an L2 adapter using the factory registry.
+
+    Looks up the type name for *config* via the config
+    registry, then calls the matching factory.
+
+    Args:
+        config: An adapter config instance.
+        l1_memory_desc: Optional L1 memory descriptor,
+            required by adapters that register L1 memory
+            with an external backend (e.g. Nixl).
+
+    Returns:
+        A new ``L2AdapterInterface`` instance.
+
+    Raises:
+        ValueError: If no factory is registered for this
+            config type.
+    """
+    # Import here to avoid circular dependency
+    # First Party
+    from lmcache.v1.distributed.l2_adapters.config import (
+        get_type_name_for_config,
+    )
+
+    name = get_type_name_for_config(config)
+
+    # Trigger lazy import if needed
+    ensure_adapter_loaded(name)
+
+    factory = _L2_ADAPTER_FACTORY_REGISTRY.get(name)
+    if factory is None:
+        raise ValueError(
+            "No adapter factory registered for type "
+            "%s. Make sure the adapter module is "
+            "imported." % name
+        )
+    return factory(config, l1_memory_desc)