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

lmcache/v1/distributed/l2_adapters/mooncake_store_l2_adapter.py

@@ -0,0 +1,135 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Mooncake Store native L2 adapter config and factory.
+"""
+
+# Future
+from __future__ import annotations
+
+# Standard
+from typing import (
+    TYPE_CHECKING,
+    Dict,
+    Optional,
+)
+
+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__)
+
+# Keys consumed only by LMCache (never sent to mooncake).
+_LMCACHE_ONLY_KEYS = {"type", "num_workers", "eviction"}
+
+
+class MooncakeStoreL2AdapterConfig(L2AdapterConfigBase):
+    """Config for an L2 adapter backed by the native
+    C++ Mooncake Store connector.
+
+    ``setup_config`` is a string-to-string dict that is
+    forwarded **as-is** to mooncake's
+    ``RealClient::setup_internal(ConfigDict)``.
+    LMCache does NOT interpret, validate, or fill in
+    defaults for any mooncake keys — that is mooncake's
+    responsibility.
+
+    ``num_workers`` is the only LMCache-specific knob.
+    """
+
+    def __init__(
+        self,
+        setup_config: Dict[str, str],
+        num_workers: int = 4,
+    ):
+        super().__init__()
+        self.setup_config: Dict[str, str] = dict(setup_config)
+        self.num_workers = num_workers
+
+    @classmethod
+    def from_dict(cls, d: dict) -> "MooncakeStoreL2AdapterConfig":
+        num_workers = d.get("num_workers", 4)
+        if not isinstance(num_workers, int) or num_workers <= 0:
+            raise ValueError("num_workers must be a positive integer")
+
+        # Everything except LMCache-only keys is
+        # forwarded to mooncake as str values.
+        setup: Dict[str, str] = {}
+        for k, v in d.items():
+            if k in _LMCACHE_ONLY_KEYS:
+                continue
+            if v is not None:
+                setup[k] = str(v)
+
+        return cls(
+            setup_config=setup,
+            num_workers=num_workers,
+        )
+
+    @classmethod
+    def help(cls) -> str:
+        return (
+            "Mooncake Store L2 adapter config.\n"
+            "All keys except LMCache-only keys are "
+            "forwarded as-is to mooncake's "
+            "setup_internal(ConfigDict).\n"
+            "Refer to mooncake documentation for "
+            "available setup keys.\n"
+            "- num_workers (int): C++ worker threads "
+            "(default 4, >0)"
+        )
+
+
+def _create_mooncake_store_l2_adapter(
+    config: L2AdapterConfigBase,
+    l1_memory_desc: "Optional[L1MemoryDesc]" = None,
+) -> L2AdapterInterface:
+    """Create a NativeConnectorL2Adapter backed by the
+    C++ Mooncake Store connector."""
+    try:
+        # First Party
+        from lmcache.lmcache_mooncake import (
+            LMCacheMooncakeClient,
+        )
+    except ImportError as e:
+        raise RuntimeError(
+            "Mooncake Store L2 adapter requires the "
+            "C++ Mooncake extension. Build with: "
+            "MOONCAKE_INCLUDE_DIR=/path/to/mooncake-"
+            "store/include pip install -e ."
+        ) from e
+
+    # First Party
+    from lmcache.v1.distributed.l2_adapters.native_connector_l2_adapter import (  # noqa: E501
+        NativeConnectorL2Adapter,
+    )
+
+    assert isinstance(config, MooncakeStoreL2AdapterConfig)
+    native_client = LMCacheMooncakeClient(
+        config=config.setup_config,
+        num_workers=config.num_workers,
+    )
+    logger.info(
+        "Created Mooncake Store L2 adapter (workers=%d)",
+        config.num_workers,
+    )
+    return NativeConnectorL2Adapter(native_client)
+
+
+# Self-register config type and adapter factory
+register_l2_adapter_type("mooncake_store", MooncakeStoreL2AdapterConfig)
+register_l2_adapter_factory("mooncake_store", _create_mooncake_store_l2_adapter)

lmcache/v1/distributed/l2_adapters/native_connector_l2_adapter.py

@@ -0,0 +1,468 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+L2 adapter that wraps any pybind-wrapped C++ IStorageConnector (native client).
+
+This bridge lets any native storage connector (Redis, RDMA, Mooncake, etc.)
+serve as an MP-mode L2 adapter.  The same C++ connector implementation is
+also usable in non-MP mode via ConnectorClientBase.
+
+Architecture:
+  - The native client has 1 eventfd + drain_completions() for all operations.
+  - This adapter creates 3 Python eventfds (store, lookup, load) and runs a
+    background demux thread that routes native completions to the right
+    category based on a future_id → op_type mapping.
+  - ObjectKey serialization and MemoryObj buffer extraction happen at the
+    submit call boundary.
+  - Locking is client-side (refcount dict) since remote backends don't have
+    our eviction concept.
+"""
+
+# Future
+from __future__ import annotations
+
+# Standard
+from collections import defaultdict
+import os
+import select
+import threading
+
+# First Party
+from lmcache.logging import init_logger
+from lmcache.native_storage_ops import Bitmap
+from lmcache.v1.distributed.api import ObjectKey
+from lmcache.v1.distributed.l2_adapters.base import (
+    L2AdapterInterface,
+    L2TaskId,
+)
+from lmcache.v1.memory_management import MemoryObj
+
+logger = init_logger(__name__)
+
+
+# Key separator — kept in sync with fs_l2_adapter.py and
+# csrc/storage_backends/fs/connector.cpp. Both ``@`` in ``model_name``
+# and ``@`` in ``cache_salt`` are rejected by ObjectKey.__post_init__
+# so splitting on ``@`` is unambiguous.
+_KEY_SEP = "@"
+
+
+def _object_key_to_string(key: ObjectKey) -> str:
+    """Serialize an ObjectKey to the native-connector wire format.
+
+    Unsalted::
+
+        <model_name>@<kv_rank_hex>@<chunk_hash_hex>
+
+    Salted (trailing ``cache_salt``)::
+
+        <model_name>@<kv_rank_hex>@<chunk_hash_hex>@<cache_salt>
+
+    Keys with ``cache_salt=""`` produce the 3-field shape, which is
+    bit-identical to the format used before ``cache_salt`` existed —
+    so existing un-salted caches remain valid with no migration.
+    """
+    base = (
+        f"{key.model_name}{_KEY_SEP}{key.kv_rank:08x}{_KEY_SEP}{key.chunk_hash.hex()}"
+    )
+    if key.cache_salt:
+        return f"{base}{_KEY_SEP}{key.cache_salt}"
+    return base
+
+
+def _obj_to_memoryview(
+    obj: MemoryObj,
+) -> memoryview:  # type: ignore[type-arg]
+    """
+    Extract a byte-oriented memoryview from a MemoryObj.
+
+    Uses the MemoryObj's byte_array property which returns
+    a ctypes-backed memoryview with itemsize=1, so pybind's
+    buffer_info.size == num_bytes.
+    """
+    return obj.byte_array  # type: ignore[return-value]
+
+
+class NativeConnectorL2Adapter(L2AdapterInterface):
+    """
+    Wraps a pybind-wrapped C++ IStorageConnector to
+    implement L2AdapterInterface.
+
+    The native_client must expose:
+      - event_fd() -> int
+      - submit_batch_get(keys, memoryviews) -> int
+      - submit_batch_set(keys, memoryviews) -> int
+      - submit_batch_exists(keys) -> int
+      - drain_completions()
+          -> list[tuple[int, bool, str, list[bool]|None]]
+      - close()
+    """
+
+    # Operation type tags for the pending-ops map
+    _OP_STORE = "store"
+    _OP_LOOKUP = "lookup"
+    _OP_LOAD = "load"
+    _OP_DELETE = "delete"
+
+    def __init__(self, native_client, max_capacity_gb: float = 0):
+        super().__init__()
+        self._client = native_client
+        self._client_fd: int = int(native_client.event_fd())
+
+        # 3 distinct Python eventfds for the L2 adapter
+        # interface
+        self._store_efd = os.eventfd(0, os.EFD_NONBLOCK | os.EFD_CLOEXEC)
+        self._lookup_efd = os.eventfd(0, os.EFD_NONBLOCK | os.EFD_CLOEXEC)
+        self._load_efd = os.eventfd(0, os.EFD_NONBLOCK | os.EFD_CLOEXEC)
+
+        # Pending ops: native future_id →
+        #   (op_type, task_id, num_keys, keys_for_locking)
+        # keys_for_locking is only set for lookup ops so
+        # we can apply locks
+        self._pending_ops: dict[
+            int,
+            tuple[str, L2TaskId, int, list[ObjectKey] | None],
+        ] = {}
+
+        # Completed results (same pattern as MockL2Adapter)
+        self._completed_stores: dict[L2TaskId, bool] = {}
+        self._completed_lookups: dict[L2TaskId, Bitmap] = {}
+        self._completed_loads: dict[L2TaskId, Bitmap] = {}
+
+        # Client-side lock tracking (refcount per key)
+        self._locked_keys: dict[ObjectKey, int] = defaultdict(int)
+
+        # Delete capability detection
+        self._has_delete = callable(getattr(native_client, "submit_batch_delete", None))
+
+        # Pending delete events for synchronous delete() calls
+        self._pending_delete_events: dict[L2TaskId, threading.Event] = {}
+
+        # Client-side size tracking for get_usage()
+        self._max_capacity_bytes = int(max_capacity_gb * (1024**3))
+        self._current_size_bytes: int = 0
+        self._key_sizes: dict[ObjectKey, int] = {}
+        # Pending store sizes: native future_id -> (keys, per_key_sizes)
+        self._pending_store_sizes: dict[int, tuple[list[ObjectKey], list[int]]] = {}
+
+        # Task ID counter
+        self._next_task_id: L2TaskId = 0
+
+        # Lock for all shared state above
+        self._lock = threading.Lock()
+
+        # Background demux thread
+        self._stop = threading.Event()
+        self._demux_thread = threading.Thread(
+            target=self._demux_loop,
+            daemon=True,
+            name="l2-adapter-demux",
+        )
+        self._demux_thread.start()
+
+    # ---------------------------------------------------------------
+    # Event Fd Interface
+    # ---------------------------------------------------------------
+
+    def get_store_event_fd(self) -> int:
+        return self._store_efd
+
+    def get_lookup_and_lock_event_fd(self) -> int:
+        return self._lookup_efd
+
+    def get_load_event_fd(self) -> int:
+        return self._load_efd
+
+    # ---------------------------------------------------------------
+    # Store Interface
+    # ---------------------------------------------------------------
+
+    def submit_store_task(
+        self,
+        keys: list[ObjectKey],
+        objects: list[MemoryObj],
+    ) -> L2TaskId:
+        key_strings = [_object_key_to_string(k) for k in keys]
+        memviews = [_obj_to_memoryview(obj) for obj in objects]
+        per_key_sizes = [obj.get_size() for obj in objects]
+
+        # Register pending op BEFORE submit to avoid race
+        # with demux thread. The native submit is
+        # non-blocking so holding the lock is brief.
+        with self._lock:
+            task_id = self._get_next_task_id()
+            future_id = int(self._client.submit_batch_set(key_strings, memviews))
+            self._pending_ops[future_id] = (
+                self._OP_STORE,
+                task_id,
+                len(keys),
+                None,
+            )
+            self._pending_store_sizes[future_id] = (list(keys), per_key_sizes)
+
+        return task_id
+
+    def pop_completed_store_tasks(
+        self,
+    ) -> dict[L2TaskId, bool]:
+        with self._lock:
+            completed = self._completed_stores
+            self._completed_stores = {}
+        return completed
+
+    # ---------------------------------------------------------------
+    # Lookup and Lock Interface
+    # ---------------------------------------------------------------
+
+    def submit_lookup_and_lock_task(
+        self,
+        keys: list[ObjectKey],
+    ) -> L2TaskId:
+        key_strings = [_object_key_to_string(k) for k in keys]
+
+        with self._lock:
+            task_id = self._get_next_task_id()
+            future_id = int(self._client.submit_batch_exists(key_strings))
+            self._pending_ops[future_id] = (
+                self._OP_LOOKUP,
+                task_id,
+                len(keys),
+                list(keys),
+            )
+
+        return task_id
+
+    def query_lookup_and_lock_result(self, task_id: L2TaskId) -> Bitmap | None:
+        with self._lock:
+            return self._completed_lookups.pop(task_id, None)
+
+    def submit_unlock(self, keys: list[ObjectKey]) -> None:
+        with self._lock:
+            for key in keys:
+                if key not in self._locked_keys:
+                    continue
+                if self._locked_keys[key] <= 1:
+                    del self._locked_keys[key]
+                else:
+                    self._locked_keys[key] -= 1
+
+    # ---------------------------------------------------------------
+    # Load Interface
+    # ---------------------------------------------------------------
+
+    def submit_load_task(
+        self,
+        keys: list[ObjectKey],
+        objects: list[MemoryObj],
+    ) -> L2TaskId:
+        key_strings = [_object_key_to_string(k) for k in keys]
+        memviews = [_obj_to_memoryview(obj) for obj in objects]
+
+        with self._lock:
+            task_id = self._get_next_task_id()
+            future_id = int(self._client.submit_batch_get(key_strings, memviews))
+            self._pending_ops[future_id] = (
+                self._OP_LOAD,
+                task_id,
+                len(keys),
+                list(keys),
+            )
+
+        return task_id
+
+    def query_load_result(self, task_id: L2TaskId) -> Bitmap | None:
+        with self._lock:
+            return self._completed_loads.pop(task_id, None)
+
+    # ---------------------------------------------------------------
+    # Eviction Interface
+    # ---------------------------------------------------------------
+
+    def delete(self, keys: list[ObjectKey]) -> None:
+        """Delete a batch of keys from the remote backend.
+
+        Submits a batch delete to the native connector and blocks
+        until the demux thread signals completion (up to 30s timeout).
+        Fires ``_notify_keys_deleted`` on success so eviction policy
+        tracking stays in sync.
+
+        No-op if the connector does not expose ``submit_batch_delete``
+        or if the key list is empty.
+        """
+        if not keys or not self._has_delete:
+            return
+
+        key_strings = [_object_key_to_string(k) for k in keys]
+        done_event = threading.Event()
+
+        with self._lock:
+            task_id = self._get_next_task_id()
+            future_id = int(self._client.submit_batch_delete(key_strings))
+            self._pending_ops[future_id] = (
+                self._OP_DELETE,
+                task_id,
+                len(keys),
+                list(keys),
+            )
+            self._pending_delete_events[task_id] = done_event
+
+        # Block until demux thread signals completion
+        if not done_event.wait(timeout=30.0):
+            with self._lock:
+                self._pending_delete_events.pop(task_id, None)
+                # Note: _pending_ops entry may already be consumed
+                # by the demux thread; pop is safe either way.
+                for fid, entry in list(self._pending_ops.items()):
+                    if entry[1] == task_id:
+                        self._pending_ops.pop(fid, None)
+                        break
+            logger.warning(
+                "delete() timed out after 30s for %d keys",
+                len(keys),
+            )
+            return
+
+        self._notify_keys_deleted(keys)
+
+    def get_usage(self) -> tuple[float, float]:
+        if self._max_capacity_bytes <= 0:
+            return (-1.0, -1.0)
+        with self._lock:
+            usage = self._current_size_bytes / self._max_capacity_bytes
+            return (usage, usage)
+
+    # ---------------------------------------------------------------
+    # Cleanup
+    # ---------------------------------------------------------------
+
+    def close(self) -> None:
+        self._stop.set()
+        self._demux_thread.join(timeout=5)
+
+        self._client.close()
+
+        os.close(self._store_efd)
+        os.close(self._lookup_efd)
+        os.close(self._load_efd)
+
+    # ---------------------------------------------------------------
+    # Internal helpers
+    # ---------------------------------------------------------------
+
+    def _get_next_task_id(self) -> L2TaskId:
+        """Increment and return the next task ID.
+        Must be called under _lock."""
+        task_id = self._next_task_id
+        self._next_task_id += 1
+        return task_id
+
+    def _demux_loop(self) -> None:
+        """Background thread that polls the native
+        connector's eventfd, drains completions, and
+        routes them to the correct L2 result category.
+        """
+        poller = select.poll()
+        poller.register(self._client_fd, select.POLLIN)
+
+        while not self._stop.is_set():
+            events = poller.poll(500)
+            if not events:
+                continue
+
+            try:
+                completions = self._client.drain_completions()
+            except Exception:
+                logger.exception("drain_completions failed")
+                continue
+
+            if not completions:
+                continue
+
+            # Collect listener notifications to fire after
+            # releasing the lock.
+            keys_stored: list[ObjectKey] = []
+            keys_accessed: list[ObjectKey] = []
+
+            with self._lock:
+                for (
+                    future_id,
+                    ok,
+                    error,
+                    result_bools,
+                ) in completions:
+                    fid = int(future_id)
+                    entry = self._pending_ops.pop(fid, None)
+                    if entry is None:
+                        logger.warning(
+                            "Received completion for unknown future_id=%d",
+                            fid,
+                        )
+                        continue
+
+                    (
+                        op_type,
+                        task_id,
+                        num_keys,
+                        lookup_keys,
+                    ) = entry
+
+                    if op_type == self._OP_STORE:
+                        self._completed_stores[task_id] = ok
+                        # Update size tracking on success
+                        store_info = self._pending_store_sizes.pop(fid, None)
+                        if ok and store_info is not None:
+                            store_keys, sizes = store_info
+                            for key, size in zip(store_keys, sizes, strict=False):
+                                if key not in self._key_sizes:
+                                    self._key_sizes[key] = size
+                                    self._current_size_bytes += size
+                            keys_stored.extend(store_keys)
+                        os.eventfd_write(self._store_efd, 1)
+
+                    elif op_type == self._OP_LOOKUP:
+                        bitmap = Bitmap(num_keys)
+                        if ok and result_bools is not None:
+                            for i, found in enumerate(result_bools):
+                                if found:
+                                    bitmap.set(i)
+                                    if lookup_keys is not None:
+                                        self._locked_keys[lookup_keys[i]] += 1
+                        self._completed_lookups[task_id] = bitmap
+                        os.eventfd_write(self._lookup_efd, 1)
+
+                    elif op_type == self._OP_LOAD:
+                        bitmap = Bitmap(num_keys)
+                        loaded_keys: list[ObjectKey] = []
+                        if result_bools is not None:
+                            for i, loaded in enumerate(result_bools):
+                                if loaded:
+                                    bitmap.set(i)
+                                    if lookup_keys is not None:
+                                        loaded_keys.append(lookup_keys[i])
+                        elif ok:
+                            # Fallback for connectors that
+                            # do not report per-key results
+                            for i in range(num_keys):
+                                bitmap.set(i)
+                            if lookup_keys is not None:
+                                loaded_keys.extend(lookup_keys)
+                        keys_accessed.extend(loaded_keys)
+                        self._completed_loads[task_id] = bitmap
+                        os.eventfd_write(self._load_efd, 1)
+
+                    elif op_type == self._OP_DELETE:
+                        # Decrement sizes for successfully deleted keys
+                        if result_bools is not None and lookup_keys is not None:
+                            for i, deleted in enumerate(result_bools):
+                                if deleted:
+                                    key = lookup_keys[i]
+                                    size = self._key_sizes.pop(key, 0)
+                                    self._current_size_bytes -= size
+                        evt = self._pending_delete_events.pop(task_id, None)
+                        if evt is not None:
+                            evt.set()
+
+            # Fire listener notifications outside the lock
+            if keys_stored:
+                self._notify_keys_stored(keys_stored)
+            if keys_accessed:
+                self._notify_keys_accessed(keys_accessed)