couchbase-lm-cache 0.2.0__py3-none-any.whl

couchbase_lm_cache/__init__.py

@@ -0,0 +1,15 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Couchbase remote storage connector plugin for LMCache."""
+
+from couchbase_lm_cache.config import CouchbaseLMCacheConfig
+from couchbase_lm_cache.connector import CouchbaseConnector
+from couchbase_lm_cache.store import CouchbaseKVStore
+
+__version__ = "0.2.0"
+
+__all__ = [
+    "CouchbaseConnector",
+    "CouchbaseKVStore",
+    "CouchbaseLMCacheConfig",
+    "__version__",
+]

couchbase_lm_cache/config.py

@@ -0,0 +1,166 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Configuration helpers for the LMCache Couchbase connector."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from os import environ
+from typing import Any
+
+_PLUGIN_PREFIX = "remote_storage_plugin"
+_ENV_PREFIX = "LMCACHE_COUCHBASE_"
+# Payloads above this size are split across several documents. Couchbase Server
+# caps documents at 20 MiB (one 256-token chunk of an 8B model is ~36 MiB), and
+# ~1 MiB documents read about twice as fast as 16 MiB ones.
+DEFAULT_SHARD_SIZE_BYTES = 1024 * 1024
+DEFAULT_MAX_WORKERS = 32
+# The SDK default (2.5 s) is tuned for small documents; a 16 MiB shard read
+# from disk on a full-eviction bucket can take longer.
+DEFAULT_KV_TIMEOUT_SECONDS = 10.0
+
+
+def _as_int(value: Any, default: int | None = None) -> int | None:
+    if value is None or value == "":
+        return default
+    return int(value)
+
+
+def _as_float(value: Any, default: float) -> float:
+    if value is None or value == "":
+        return default
+    return float(value)
+
+
+@dataclass(frozen=True)
+class CouchbaseLMCacheConfig:
+    """Connection and behavior settings for Couchbase-backed LMCache storage.
+
+    Values can come from LMCache ``extra_config`` keys named
+    ``remote_storage_plugin.<plugin_name>.<field>`` or from environment variables
+    with the ``LMCACHE_COUCHBASE_`` prefix.
+    """
+
+    connection_string: str
+    username: str
+    password: str
+    bucket_name: str
+    scope_name: str = "_default"
+    collection_name: str = "_default"
+    key_prefix: str = "lmcache:"
+    ttl_seconds: int | None = None
+    kv_timeout_seconds: float = DEFAULT_KV_TIMEOUT_SECONDS
+    connect_timeout_seconds: float = 10.0
+    max_workers: int = DEFAULT_MAX_WORKERS
+    shard_size_bytes: int = DEFAULT_SHARD_SIZE_BYTES
+
+    @classmethod
+    def from_env(cls) -> CouchbaseLMCacheConfig:
+        """Build configuration from environment variables.
+
+        Raises:
+            ValueError: if required variables are absent.
+        """
+
+        return cls._build(lambda field: None)
+
+    @classmethod
+    def from_lmcache_extra_config(
+        cls,
+        extra_config: dict[str, Any] | None,
+        plugin_name: str = "couchbase",
+    ) -> CouchbaseLMCacheConfig:
+        """Build configuration from LMCache ``extra_config`` plus env fallback.
+
+        Args:
+            extra_config: LMCache ``LMCacheEngineConfig.extra_config`` mapping.
+            plugin_name: Remote storage plugin name, for example ``couchbase`` or
+                ``couchbase.primary``.
+
+        Raises:
+            ValueError: if required connection fields are absent.
+        """
+
+        extra_config = extra_config or {}
+
+        def get(field: str) -> Any:
+            specific = f"{_PLUGIN_PREFIX}.{plugin_name}.{field}"
+            if specific in extra_config:
+                return extra_config[specific]
+            plugin_type = plugin_name.split(".", 1)[0]
+            return extra_config.get(f"{_PLUGIN_PREFIX}.{plugin_type}.{field}")
+
+        return cls._build(get)
+
+    @classmethod
+    def from_dict(cls, values: dict[str, Any] | None) -> CouchbaseLMCacheConfig:
+        """Build configuration from plain field names plus env fallback.
+
+        Used by the LMCache multiprocess-mode L2 adapter, whose
+        ``adapter_params`` use the same field names as ``extra_config``
+        without the ``remote_storage_plugin.<name>.`` prefix.
+
+        Raises:
+            ValueError: if required connection fields are absent.
+        """
+
+        values = values or {}
+        return cls._build(values.get)
+
+    @classmethod
+    def _build(cls, lookup: Callable[[str], Any]) -> CouchbaseLMCacheConfig:
+        def get(field: str, env_name: str | None = None, default: Any = None) -> Any:
+            value = lookup(field)
+            if value is not None:
+                return value
+            return environ.get(_ENV_PREFIX + (env_name or field).upper(), default)
+
+        return cls(
+            connection_string=_required_value(
+                get("connection_string", "CONNSTR"), "connection_string"
+            ),
+            username=_required_value(get("username"), "username"),
+            password=_required_value(get("password"), "password"),
+            bucket_name=_required_value(get("bucket_name", "BUCKET"), "bucket_name"),
+            scope_name=str(get("scope_name", "SCOPE", "_default")),
+            collection_name=str(get("collection_name", "COLLECTION", "_default")),
+            key_prefix=str(get("key_prefix", "KEY_PREFIX", "lmcache:")),
+            ttl_seconds=_as_int(get("ttl_seconds", "TTL_SECONDS")),
+            kv_timeout_seconds=_as_float(
+                get("kv_timeout_seconds", "KV_TIMEOUT_SECONDS"),
+                DEFAULT_KV_TIMEOUT_SECONDS,
+            ),
+            connect_timeout_seconds=_as_float(
+                get("connect_timeout_seconds", "CONNECT_TIMEOUT_SECONDS"), 10.0
+            ),
+            max_workers=_as_int(get("max_workers", "MAX_WORKERS"), DEFAULT_MAX_WORKERS)
+            or DEFAULT_MAX_WORKERS,
+            shard_size_bytes=_as_int(
+                get("shard_size_bytes", "SHARD_SIZE_BYTES"), DEFAULT_SHARD_SIZE_BYTES
+            )
+            or DEFAULT_SHARD_SIZE_BYTES,
+        )
+
+    def redacted(self) -> dict[str, Any]:
+        """Return a safe dictionary for logging or support output."""
+
+        return {
+            "connection_string": self.connection_string,
+            "username": self.username,
+            "password": "***" if self.password else "",
+            "bucket_name": self.bucket_name,
+            "scope_name": self.scope_name,
+            "collection_name": self.collection_name,
+            "key_prefix": self.key_prefix,
+            "ttl_seconds": self.ttl_seconds,
+            "kv_timeout_seconds": self.kv_timeout_seconds,
+            "connect_timeout_seconds": self.connect_timeout_seconds,
+            "max_workers": self.max_workers,
+            "shard_size_bytes": self.shard_size_bytes,
+        }
+
+
+def _required_value(value: Any, name: str) -> str:
+    if value is None or str(value) == "":
+        raise ValueError(f"Missing required Couchbase LMCache setting: {name}")
+    return str(value)

couchbase_lm_cache/connector.py

@@ -0,0 +1,319 @@
+# SPDX-License-Identifier: Apache-2.0
+"""LMCache RemoteConnector implementation backed by Couchbase KV."""
+
+from __future__ import annotations
+
+import asyncio
+import builtins
+import logging
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any
+
+from couchbase_lm_cache.config import CouchbaseLMCacheConfig
+from couchbase_lm_cache.store import CouchbaseKVStore
+
+try:  # pragma: no cover - exercised when installed alongside LMCache
+    from lmcache.logging import init_logger
+    from lmcache.utils import CacheEngineKey
+    from lmcache.v1.memory_management import MemoryObj
+    from lmcache.v1.protocol import RemoteMetadata
+    from lmcache.v1.storage_backend.connector.base_connector import RemoteConnector
+
+    _HAS_LMCACHE = True
+except Exception:  # pragma: no cover - unit tests run without heavy LMCache deps
+    CacheEngineKey = Any  # type: ignore[misc,assignment]
+    MemoryObj = Any  # type: ignore[misc,assignment]
+    RemoteMetadata = None  # type: ignore[misc,assignment]
+
+    class RemoteConnector:  # type: ignore[no-redef]
+        """Fallback base so unit tests can run without LMCache installed."""
+
+        save_chunk_meta = False
+        remote_metadata_bytes = 0
+
+        def reshape_partial_chunk(self, memory_obj: Any, bytes_read: int) -> Any:
+            del bytes_read
+            return memory_obj
+
+    _HAS_LMCACHE = False
+    init_logger = logging.getLogger
+
+logger = init_logger(__name__)
+
+
+def _byte_view(memory_obj: Any) -> memoryview:
+    """Return a flat unsigned-byte view over a ``MemoryObj`` buffer."""
+
+    buffer = memory_obj.byte_array
+    view = buffer if isinstance(buffer, memoryview) else memoryview(buffer)
+    if view.format != "B":
+        view = view.cast("B")
+    return view
+
+
+def _release(memory_obj: Any) -> None:
+    release = getattr(memory_obj, "ref_count_down", None)
+    if release is not None:
+        release()
+
+
+class CouchbaseConnector(RemoteConnector):
+    """LMCache remote storage connector that stores chunks in Couchbase KV.
+
+    This class is loaded by LMCache with ``remote_storage_plugins`` and implements
+    the ``RemoteConnector`` methods expected by ``RemoteBackend``.
+
+    Each LMCache chunk is one binary Couchbase document. When LMCache's
+    ``save_chunk_meta`` is enabled (the LMCache default), the document starts with
+    LMCache's serialized ``RemoteMetadata`` header followed by the chunk bytes, so
+    header and payload are always written and expired atomically.
+    """
+
+    def __init__(
+        self,
+        loop: asyncio.AbstractEventLoop,
+        local_cpu_backend: Any,
+        config: Any,
+        plugin_name: str | None = None,
+    ) -> None:
+        """Initialize the connector.
+
+        Args:
+            loop: LMCache event loop.
+            local_cpu_backend: LMCache local CPU backend used to allocate buffers
+                for retrieved objects.
+            config: LMCache engine config containing ``extra_config``.
+            plugin_name: Name from ``remote_storage_plugins``; defaults to
+                ``couchbase``.
+        """
+
+        if _HAS_LMCACHE:
+            super().__init__(local_cpu_backend.config, local_cpu_backend.metadata)
+        self.loop = loop
+        self.local_cpu_backend = local_cpu_backend
+        self.config = config
+        configured_plugins = getattr(config, "remote_storage_plugins", None) or []
+        inferred_plugin_name = (
+            configured_plugins[0] if len(configured_plugins) == 1 else "couchbase"
+        )
+        self.plugin_name = plugin_name or inferred_plugin_name
+        extra_config = getattr(config, "extra_config", None)
+        self.cfg = CouchbaseLMCacheConfig.from_lmcache_extra_config(
+            extra_config, self.plugin_name
+        )
+        self.store = CouchbaseKVStore(self.cfg)
+        self.executor = ThreadPoolExecutor(
+            max_workers=self.cfg.max_workers, thread_name_prefix="couchbase-lmcache"
+        )
+        logger.info(
+            "Couchbase LMCache connector ready: %s (save_chunk_meta=%s)",
+            self.cfg.redacted(),
+            self.save_chunk_meta,
+        )
+
+    async def _run_blocking(self, func: Any, *args: Any) -> Any:
+        return await self.loop.run_in_executor(self.executor, func, *args)
+
+    def _encode(self, memory_obj: MemoryObj) -> bytes:
+        payload = _byte_view(memory_obj)
+        if not self.save_chunk_meta:
+            return bytes(payload)
+        header = RemoteMetadata(
+            len(payload),
+            memory_obj.get_shapes(),
+            memory_obj.get_dtypes(),
+            memory_obj.get_memory_format(),
+        ).serialize()
+        return header + bytes(payload)
+
+    def _decode(self, key: CacheEngineKey, document: bytes) -> MemoryObj | None:
+        if self.save_chunk_meta:
+            header_size = self.remote_metadata_bytes
+            if len(document) < header_size:
+                raise ValueError(
+                    f"Couchbase document for {key} has {len(document)} bytes, "
+                    f"shorter than the {header_size}-byte LMCache metadata header"
+                )
+            metadata = RemoteMetadata.deserialize(document[:header_size])
+            payload = memoryview(document)[header_size:]
+            if len(payload) != metadata.length:
+                raise ValueError(
+                    f"Couchbase document for {key} has {len(payload)} payload "
+                    f"bytes but its metadata header says {metadata.length}"
+                )
+            memory_obj = self.local_cpu_backend.allocate(
+                metadata.shapes, metadata.dtypes, metadata.fmt
+            )
+        else:
+            payload = memoryview(document)
+            memory_obj = self.local_cpu_backend.allocate(
+                self.meta_shapes, self.meta_dtypes, self.meta_fmt
+            )
+        if memory_obj is None:
+            logger.warning("LMCache local CPU backend could not allocate MemoryObj")
+            return None
+        try:
+            target = _byte_view(memory_obj)
+            if len(payload) > len(target):
+                raise ValueError(
+                    f"Couchbase payload for {key} has {len(payload)} bytes but "
+                    f"allocated LMCache buffer has {len(target)} bytes"
+                )
+            target[: len(payload)] = payload
+            if self.save_chunk_meta:
+                return memory_obj
+            return self.reshape_partial_chunk(memory_obj, len(payload))
+        except Exception:
+            _release(memory_obj)
+            raise
+
+    async def exists(self, key: CacheEngineKey) -> bool:
+        """Return whether Couchbase contains the LMCache key."""
+
+        return bool(await self._run_blocking(self.store.exists, key))
+
+    def exists_sync(self, key: CacheEngineKey) -> bool:
+        """Synchronous existence check used by LMCache's optional fast path."""
+
+        return self.store.exists(key)
+
+    async def get(self, key: CacheEngineKey) -> MemoryObj | None:
+        """Fetch a chunk from Couchbase into an LMCache ``MemoryObj``."""
+
+        document = await self._run_blocking(self.store.get_bytes, key)
+        if document is None:
+            return None
+        return self._decode(key, document)
+
+    async def put(self, key: CacheEngineKey, memory_obj: MemoryObj) -> None:
+        """Persist a chunk in Couchbase."""
+
+        await self._run_blocking(self.store.put_bytes, key, self._encode(memory_obj))
+
+    def remove_sync(self, key: CacheEngineKey) -> bool:
+        """Remove a chunk; used by LMCache ``RemoteBackend.remove``."""
+
+        return self.store.remove(key)
+
+    async def list(self) -> builtins.list[str]:
+        """List is intentionally unsupported for the remote connector contract."""
+
+        return []
+
+    def support_ping(self) -> bool:
+        """Couchbase connectivity can be health-checked."""
+
+        return True
+
+    async def ping(self) -> int:
+        """Return 0 when the Couchbase KV service answers, 1 otherwise."""
+
+        return 0 if await self._run_blocking(self.store.ping) else 1
+
+    async def close(self) -> None:
+        """Close Couchbase and executor resources."""
+
+        await self._run_blocking(self.store.close)
+        self.executor.shutdown(wait=True, cancel_futures=False)
+
+    def support_batched_put(self) -> bool:
+        """Couchbase writes are batched through concurrent executor jobs."""
+
+        return True
+
+    async def batched_put(
+        self, keys: builtins.list[CacheEngineKey], memory_objs: builtins.list[MemoryObj]
+    ) -> None:
+        """Persist multiple chunks concurrently."""
+
+        documents = [self._encode(memory_obj) for memory_obj in memory_objs]
+        await asyncio.gather(
+            *(
+                self._run_blocking(self.store.put_bytes, key, document)
+                for key, document in zip(keys, documents, strict=True)
+            )
+        )
+
+    def support_batched_get(self) -> bool:
+        """Couchbase reads are batched through concurrent executor jobs."""
+
+        return True
+
+    async def batched_get(
+        self, keys: builtins.list[CacheEngineKey]
+    ) -> builtins.list[MemoryObj | None]:
+        """Fetch multiple chunks concurrently."""
+
+        return list(await asyncio.gather(*(self.get(key) for key in keys)))
+
+    def support_batched_contains(self) -> bool:
+        """Enable LMCache prefix-hit counting over multiple keys."""
+
+        return True
+
+    def batched_contains(self, keys: builtins.list[CacheEngineKey]) -> int:
+        """Return the number of consecutive existing keys from the start."""
+
+        count = 0
+        for key in keys:
+            if not self.store.exists(key):
+                return count
+            count += 1
+        return count
+
+    def support_batched_async_contains(self) -> bool:
+        """Enable async prefix-hit counting."""
+
+        return True
+
+    async def batched_async_contains(
+        self,
+        lookup_id: str,
+        keys: builtins.list[CacheEngineKey],
+        pin: bool = False,
+    ) -> int:
+        """Return the async number of consecutive existing keys from the start."""
+
+        del lookup_id, pin
+        results = await asyncio.gather(
+            *(self._run_blocking(self.store.exists, key) for key in keys)
+        )
+        count = 0
+        for exists in results:
+            if not exists:
+                return count
+            count += 1
+        return count
+
+    def support_batched_get_non_blocking(self) -> bool:
+        """Enable LMCache prefetch reads."""
+
+        return True
+
+    async def batched_get_non_blocking(
+        self,
+        lookup_id: str,
+        keys: builtins.list[CacheEngineKey],
+    ) -> builtins.list[MemoryObj]:
+        """Fetch the consecutive prefix of chunks for prefetch paths.
+
+        LMCache maps the returned objects positionally onto ``keys``, so results
+        stop at the first miss or error and any later objects are released.
+        """
+
+        del lookup_id
+        results = await asyncio.gather(
+            *(self.get(key) for key in keys), return_exceptions=True
+        )
+        prefix: builtins.list[MemoryObj] = []
+        broken = False
+        for result in results:
+            if result is None or isinstance(result, BaseException):
+                if isinstance(result, BaseException):
+                    logger.warning("Couchbase prefetch failed: %s", result)
+                broken = True
+            elif broken:
+                _release(result)
+            else:
+                prefix.append(result)
+        return prefix