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

lmcache/v1/distributed/api.py

@@ -0,0 +1,248 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Defines the data structures that will be used by the
+distributed storage manager public functions
+
+Could be implemented by native code in the future
+"""
+
+# Standard
+from dataclasses import dataclass
+
+# Third Party
+import torch
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.v1.multiprocess.custom_types import IPCCacheEngineKey
+
+logger = init_logger(__name__)
+
+
+@dataclass(frozen=True)
+class ObjectKey:
+    """
+    The unique identifier for an object in the distributed storage manager
+    """
+
+    chunk_hash: bytes
+    """ Content hash of this particular chunk """
+
+    model_name: str
+    """ Name of the model this chunk belongs to.
+
+    Invariant: must not contain ``@``. The L2 adapters use ``@`` as the
+    field separator in serialized keys/filenames and rely on this
+    invariant for unambiguous parsing. HuggingFace model IDs use
+    alphanumerics + ``/-_.`` so this rejects nothing that appears in
+    practice.
+    """
+
+    kv_rank: int
+    """ The rank that uniquely identifies the slice of the KV cache """
+
+    cache_salt: str = ""
+    """ Per-user isolation salt. Same content from different users with
+    different cache_salt values produces different ObjectKeys, giving
+    strict per-user cache isolation. Defaults to empty string, in which
+    case serialized keys and filenames match the pre-cache_salt shape
+    (no trailing salt field) — no migration is needed for un-salted
+    deployments.
+
+    Invariant: must not contain ``@``, ``/``, ``\\``, or NUL. The L2
+    adapters use ``@`` as the field separator; ``/`` and ``\\`` are
+    filesystem path separators (FS adapter embeds the salt into
+    filenames); NUL terminates C strings (C++ connector). Max length
+    128 to stay well within ``NAME_MAX`` (255) after the model, rank,
+    hash, and extension are added.
+    """
+
+    _SALT_FORBIDDEN_CHARS = frozenset("@/\\\x00")
+    _SALT_MAX_LEN = 128
+
+    def __post_init__(self) -> None:
+        if "@" in self.model_name:
+            raise ValueError(
+                f"model_name must not contain '@' (got {self.model_name!r})"
+            )
+        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)})"
+            )
+
+    @staticmethod
+    def IntHash2Bytes(chunk_hash: int) -> bytes:
+        # NOTE: this is only used by tests
+        return chunk_hash.to_bytes(4, byteorder="big")
+
+    @staticmethod
+    def Bytes2IntHash(chunk_hash: bytes) -> int:
+        # NOTE: this is only used by tests
+        return int.from_bytes(chunk_hash, byteorder="big") & ((1 << 64) - 1)
+
+    @staticmethod
+    def ComputeKVRank(
+        world_size: int,
+        global_rank: int,
+        local_world_size: int,
+        local_rank: int,
+    ) -> int:
+        """
+        Compute the kv_rank from world_size and worker_id
+
+        Args:
+            world_size (int): The total number of workers (include TP + PP)
+            global_rank (int): The global worker id (from 0 to world_size - 1)
+            local_world_size (int): The local world size (for local node),
+                should NOT be greater than 8
+            local_rank (int): The local world rank (for local node)
+
+        Returns:
+            The special KV rank (bitmap) used by the objectkey
+
+        Example:
+            In the case of TP=4, PP=2, the TP worker 1 on node 1 has:
+            - world_size = 8
+            - global_rank = 5
+            - local_world_size = 4
+            - local_rank = 1
+
+            The output KV rank is the bitmap:
+            +--head--+
+            |00000000|
+            |00000000|
+            |00000000|
+            |00000000| layers
+            |00001100|
+            |00001100|
+            |00001100|
+            |00001100|
+            +--------+
+        """
+        # TODO(ApostaC): in the long run, we want to have the above bitmap based
+        # representation for asymmetric parallelism (e.g., sharing across different
+        # TP/PP settings).
+        # For now, let's have a simple implementation that just
+        # differentiate between different parallel setups
+
+        # For each number, we use 8-bit, and pack them together
+        return (
+            (world_size << 24)
+            | (global_rank << 16)
+            | (local_world_size << 8)
+            | local_rank
+        )
+
+
+@dataclass(frozen=True)
+class MemoryLayoutDesc:
+    """
+    Describes the layout of a memory object
+    """
+
+    shapes: list[torch.Size]
+    dtypes: list[torch.dtype]
+
+    def __post_init__(self):
+        if len(self.shapes) != len(self.dtypes):
+            raise ValueError(
+                "MemoryLayoutDesc: shapes and dtype must have the same length"
+            )
+
+
+@dataclass(frozen=True)
+class PrefetchHandle:
+    """Opaque handle returned by ``StorageManager.submit_prefetch_task``.
+
+    Carries the bookkeeping needed to later query lookup / prefetch status
+    without exposing controller internals.
+    """
+
+    prefetch_request_id: int
+    """Opaque ID for tracking L2 prefetch in the controller.
+    -1 if no L2 request was submitted."""
+
+    external_request_id: str
+    """Request ID from the caller for end-to-end tracing."""
+
+    l1_prefix_hit_count: int
+    """Number of leading keys already in L1 at submission time."""
+
+    total_requested_keys: int
+    """Total number of keys originally requested."""
+
+    submit_time: float
+    """Monotonic timestamp when the prefetch task was submitted."""
+
+
+def ipc_key_to_object_keys(
+    ipc_key: IPCCacheEngineKey,
+    chunk_hashes: list[bytes],
+) -> list[ObjectKey]:
+    """
+    Convert a single IPCCacheEngineKey and its chunk hashes to a list of ObjectKey.
+
+    When the ipc_key's worker_id is None, each chunk hash is exploded into
+    multiple ObjectKeys (one per worker in world_size).
+
+    ``cache_salt`` is read directly from ``ipc_key`` so the produced
+    ObjectKeys are per-user isolated whenever the sender set a non-empty
+    salt. There is intentionally no separate ``cache_salt`` parameter —
+    duplicating the source of truth would risk silent isolation bugs
+    where a caller passes ``ipc_key`` but forgets the salt.
+
+    Args:
+        ipc_key: The IPC key providing model_name, world_size, worker_id,
+            and cache_salt.
+        chunk_hashes: List of chunk hash bytes, one per chunk.
+
+    Returns:
+        list[ObjectKey]: The converted list of ObjectKey.
+    """
+    cache_salt = ipc_key.cache_salt
+    storage_keys = []
+    for chunk_hash in chunk_hashes:
+        if ipc_key.worker_id is None:
+            # For look up request, we want to expand to all workers
+            for worker_id in range(ipc_key.world_size):
+                # TODO (ApostaC): include local world size/rank info
+                # in the future once it's in IPCCacheEngineKey
+                kv_rank = ObjectKey.ComputeKVRank(
+                    world_size=ipc_key.world_size,
+                    global_rank=worker_id,
+                    local_world_size=ipc_key.world_size,
+                    local_rank=worker_id,
+                )
+
+                storage_keys.append(
+                    ObjectKey(
+                        chunk_hash=chunk_hash,
+                        model_name=ipc_key.model_name,
+                        kv_rank=kv_rank,
+                        cache_salt=cache_salt,
+                    )
+                )
+        else:
+            kv_rank = ObjectKey.ComputeKVRank(
+                world_size=ipc_key.world_size,
+                global_rank=ipc_key.worker_id,
+                local_world_size=ipc_key.world_size,
+                local_rank=ipc_key.worker_id,
+            )
+
+            storage_keys.append(
+                ObjectKey(
+                    chunk_hash=chunk_hash,
+                    model_name=ipc_key.model_name,
+                    kv_rank=kv_rank,
+                    cache_salt=cache_salt,
+                )
+            )
+
+    return storage_keys

lmcache/v1/distributed/config.py

@@ -0,0 +1,321 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Configuration for distributed storage manager
+"""
+
+# Standard
+from dataclasses import dataclass, field
+from typing import Literal
+import argparse
+
+# First Party
+from lmcache.v1.distributed.l2_adapters.config import (
+    L2AdaptersConfig,
+    add_l2_adapters_args,
+    parse_args_to_l2_adapters_config,
+)
+
+
+@dataclass
+class L1MemoryManagerConfig:
+    """
+    The configuration for L1 memory manager.
+    """
+
+    size_in_bytes: int
+    """ The size of L1 memory in bytes. """
+
+    use_lazy: bool
+    """ Whether to use lazy initialization for L1 memory. """
+
+    init_size_in_bytes: int = field(default=20 << 30)
+    """ The initial size when using lazy allocation. Default is 20GB. """
+
+    align_bytes: int = field(default=0x1000)
+    """ The alignment size in bytes. Default is 4KB. """
+
+    def __post_init__(self):
+        self.init_size_in_bytes = min(self.init_size_in_bytes, self.size_in_bytes)
+
+
+@dataclass
+class L1ManagerConfig:
+    """
+    Special config for the L1 Object/Key manager
+    """
+
+    memory_config: L1MemoryManagerConfig
+    """ The memory manager configuration for L1 cache. """
+
+    write_ttl_seconds: int = field(default=600)
+    """ Time to live for each object's write lock. Default is 600s (10 minutes). """
+
+    read_ttl_seconds: int = field(default=300)
+    """ Time to live for each object's read lock. Default is 300s (5 minutes). """
+
+
+@dataclass
+class EvictionConfig:
+    """
+    The configuration for eviction policies (L1 and optionally L2).
+    """
+
+    eviction_policy: Literal["LRU", "noop"]
+    """ The eviction policy to use. """
+
+    trigger_watermark: float = field(default=0.8)
+    """ The memory usage watermark to trigger eviction (0.0 to 1.0). """
+
+    eviction_ratio: float = field(default=0.2)
+    """ The fraction of *allocated* memory to evict when triggered (0.0 to 1.0). """
+
+
+@dataclass
+class StorageManagerConfig:
+    """
+    The configuration for the distributed storage manager.
+    """
+
+    l1_manager_config: L1ManagerConfig
+    """ The configuration for the L1 manager. """
+
+    eviction_config: EvictionConfig
+    """ The configuration for eviction policies. """
+
+    l2_adapter_config: L2AdaptersConfig = field(
+        default_factory=lambda: L2AdaptersConfig([])
+    )
+    """ The configuration for L2 adapters. """
+
+    store_policy: str = "default"
+    """ The L2 store policy name. """
+
+    prefetch_policy: str = "default"
+    """ The L2 prefetch policy name. """
+
+    prefetch_max_in_flight: int = 8
+    """ Maximum number of concurrent prefetch requests. """
+
+
+def add_storage_manager_args(
+    parser: argparse.ArgumentParser,
+) -> argparse.ArgumentParser:
+    """
+    Add storage manager configuration arguments to an existing parser.
+
+    This function allows other modules to integrate storage manager arguments
+    into their own argument parsers. Arguments are organized into groups to
+    avoid naming conflicts with other modules.
+
+    Args:
+        parser: The argument parser to add arguments to.
+
+    Returns:
+        argparse.ArgumentParser: The same parser with storage manager
+            arguments added.
+
+    Example:
+        >>> # In another module that needs its own arguments
+        >>> parser = argparse.ArgumentParser(description="My Application")
+        >>> parser.add_argument("--my-arg", type=str)
+        >>> add_storage_manager_args(parser)
+        >>> args = parser.parse_args()
+        >>> config = parse_args_to_config(args)
+    """
+    # L1 Memory Manager Config
+    memory_group = parser.add_argument_group(
+        "L1 Memory Manager", "Configuration for L1 memory manager"
+    )
+    memory_group.add_argument(
+        "--l1-size-gb",
+        type=float,
+        required=True,
+        help="The size of L1 memory in GB.",
+    )
+    memory_group.add_argument(
+        "--l1-use-lazy",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="Whether to use lazy loading for L1 memory. (Default is True)",
+    )
+    memory_group.add_argument(
+        "--l1-init-size-gb",
+        type=int,
+        default=20,
+        help="The initial size (GB) when using lazy allocation. Default is 20.",
+    )
+    memory_group.add_argument(
+        "--l1-align-bytes",
+        type=int,
+        default=4096,
+        help="The alignment size in bytes. Default is 4KB (4096 bytes).",
+    )
+
+    # L1 Manager Config (TTL settings)
+    ttl_group = parser.add_argument_group(
+        "L1 Manager TTL", "TTL configuration for L1 manager locks"
+    )
+    ttl_group.add_argument(
+        "--l1-write-ttl-seconds",
+        type=int,
+        default=600,
+        help="Time to live for each object's write lock. Default is 600s.",
+    )
+    ttl_group.add_argument(
+        "--l1-read-ttl-seconds",
+        type=int,
+        default=300,
+        help="Time to live for each object's read lock. Default is 300s.",
+    )
+
+    # Eviction Config
+    eviction_group = parser.add_argument_group(
+        "Eviction Policy", "Configuration for eviction policies"
+    )
+    eviction_group.add_argument(
+        "--eviction-policy",
+        type=str,
+        choices=["LRU", "noop"],
+        required=True,
+        help="The eviction policy to use ('LRU' or 'noop').",
+    )
+    eviction_group.add_argument(
+        "--eviction-trigger-watermark",
+        type=float,
+        default=0.8,
+        help="The memory usage watermark to trigger eviction (0.0 to 1.0). "
+        "Default is 0.8.",
+    )
+    eviction_group.add_argument(
+        "--eviction-ratio",
+        type=float,
+        default=0.2,
+        help="The fraction of memory to evict when triggered (0.0 to 1.0). "
+        "Default is 0.2.",
+    )
+
+    # L2 Policies
+    # Import here to break circular dependency:
+    # config.py <-> storage_controllers (via eviction_controller)
+    # Safe because config.py is fully initialized by the time this
+    # function is called.
+    # First Party
+    from lmcache.v1.distributed.storage_controllers.prefetch_policy import (
+        get_registered_prefetch_policies,
+    )
+    from lmcache.v1.distributed.storage_controllers.store_policy import (
+        get_registered_store_policies,
+    )
+    import lmcache.v1.distributed.storage_controllers  # noqa: F401
+
+    policy_group = parser.add_argument_group(
+        "L2 Policies", "Store and prefetch policy selection for L2 adapters"
+    )
+    policy_group.add_argument(
+        "--l2-store-policy",
+        type=str,
+        choices=get_registered_store_policies(),
+        default="default",
+        help="L2 store policy. Determines which adapters receive each key "
+        "and whether keys are deleted from L1 after L2 store. "
+        "Default is 'default' (store all keys to all adapters, keep L1).",
+    )
+    policy_group.add_argument(
+        "--l2-prefetch-policy",
+        type=str,
+        choices=get_registered_prefetch_policies(),
+        default="default",
+        help="L2 prefetch policy. Determines which adapter loads each key "
+        "when multiple adapters have it. "
+        "Default is 'default' (pick the first adapter by index).",
+    )
+    policy_group.add_argument(
+        "--l2-prefetch-max-in-flight",
+        type=int,
+        default=8,
+        help="Maximum number of concurrent prefetch requests. Default is 8.",
+    )
+
+    # Adapter config
+    add_l2_adapters_args(parser)
+    return parser
+
+
+def get_arg_parser() -> argparse.ArgumentParser:
+    """
+    Get a standalone argument parser for storage manager configuration.
+
+    This creates a new parser with only storage manager arguments.
+    For integrating with other modules' parsers, use add_storage_manager_args()
+    instead.
+
+    Returns:
+        argparse.ArgumentParser: The argument parser with all storage manager
+            configuration options.
+    """
+    parser = argparse.ArgumentParser(
+        description="Distributed Storage Manager Configuration"
+    )
+    return add_storage_manager_args(parser)
+
+
+def parse_args_to_config(
+    args: argparse.Namespace,
+) -> StorageManagerConfig:
+    """
+    Convert parsed command line arguments to a StorageManagerConfig.
+
+    Args:
+        args: Parsed arguments from the argument parser.
+
+    Returns:
+        StorageManagerConfig: The configuration object.
+    """
+    memory_config = L1MemoryManagerConfig(
+        size_in_bytes=int(args.l1_size_gb * (1 << 30)),
+        use_lazy=args.l1_use_lazy,
+        init_size_in_bytes=int(args.l1_init_size_gb * (1 << 30)),
+        align_bytes=args.l1_align_bytes,
+    )
+
+    l1_manager_config = L1ManagerConfig(
+        memory_config=memory_config,
+        write_ttl_seconds=args.l1_write_ttl_seconds,
+        read_ttl_seconds=args.l1_read_ttl_seconds,
+    )
+
+    eviction_config = EvictionConfig(
+        eviction_policy=args.eviction_policy,
+        trigger_watermark=args.eviction_trigger_watermark,
+        eviction_ratio=args.eviction_ratio,
+    )
+
+    l2_adapter_config = parse_args_to_l2_adapters_config(args)
+
+    return StorageManagerConfig(
+        l1_manager_config=l1_manager_config,
+        eviction_config=eviction_config,
+        l2_adapter_config=l2_adapter_config,
+        store_policy=args.l2_store_policy,
+        prefetch_policy=args.l2_prefetch_policy,
+        prefetch_max_in_flight=args.l2_prefetch_max_in_flight,
+    )
+
+
+def parse_args(args: list[str] | None = None) -> StorageManagerConfig:
+    """
+    Parse command line arguments and return a StorageManagerConfig.
+
+    This is a convenience function that combines get_arg_parser() and
+    parse_args_to_config().
+
+    Args:
+        args: Optional list of arguments to parse. If None, uses sys.argv.
+
+    Returns:
+        StorageManagerConfig: The configuration object.
+    """
+    parser = get_arg_parser()
+    parsed_args = parser.parse_args(args)
+    return parse_args_to_config(parsed_args)

lmcache/v1/distributed/error.py

@@ -0,0 +1,64 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Definition of errors for class APIs.
+"""
+
+# Standard
+import enum
+
+
+class L1Error(enum.Enum):
+    """Errors for L1Manager class APIs."""
+
+    SUCCESS = enum.auto()
+    """ Operation succeeded. """
+
+    KEY_NOT_EXIST = enum.auto()
+    """ The specified key does not exist. """
+
+    KEY_NOT_READABLE = enum.auto()
+    """ The specified key exists but cannot be read """
+
+    KEY_NOT_WRITABLE = enum.auto()
+    """ The specified key exists but cannot be written """
+
+    KEY_IN_WRONG_STATE = enum.auto()
+    """ The specified key is in the wrong state for the operation. """
+
+    KEY_IS_LOCKED = enum.auto()
+    """ The specified key is locked and cannot perform the operation. """
+
+    OUT_OF_MEMORY = enum.auto()
+    """ Not enough memory to complete the operation. """
+
+
+ErrorType = L1Error
+
+
+def strerror(error: ErrorType) -> str:
+    """Convert error code to human-readable string.
+
+    Args:
+        error (ErrorType): The error code.
+
+    Returns:
+        str: The human-readable string.
+    """
+    if isinstance(error, L1Error):
+        if error == L1Error.SUCCESS:
+            return "Operation succeeded."
+        elif error == L1Error.KEY_NOT_EXIST:
+            return "The specified key does not exist."
+        elif error == L1Error.KEY_NOT_READABLE:
+            return "The specified key exists but cannot be read."
+        elif error == L1Error.KEY_NOT_WRITABLE:
+            return "The specified key exists but cannot be written."
+        elif error == L1Error.KEY_IN_WRONG_STATE:
+            return "The specified key is in the wrong state for the operation."
+        elif error == L1Error.KEY_IS_LOCKED:
+            return "The specified key is locked and cannot perform the operation."
+        elif error == L1Error.OUT_OF_MEMORY:
+            return "Not enough memory to complete the operation."
+
+    return "Unknown error."