From 5cfd55363615fc6ea533b775232adf873daa0b4c Mon Sep 17 00:00:00 2001 From: seohui-XCENA Date: Fri, 3 Jul 2026 17:34:57 +0900 Subject: [PATCH 01/28] [MP][Maru] Add L1ManagerInterface protocol and maru L1 config - l1_protocol.py: structural runtime_checkable Protocol mirroring L1Manager's 17-method surface, with per-method listener/lock contract docstrings - config.py: MaruL1Config + maru_config field, __post_init__ DRAM-clamp skip, maru CLI args (--maru-server-url/--maru-pool-size-gb/--maru-instance-id), parse_args_to_config wiring - tests: interface<->L1Manager method-set + signature conformance, maru config parsing Signed-off-by: seohui-XCENA --- lmcache/v1/distributed/config.py | 79 ++++++++++++++ lmcache/v1/distributed/l1_protocol.py | 113 ++++++++++++++++++++ tests/v1/distributed/test_l1_config_maru.py | 49 +++++++++ tests/v1/distributed/test_l1_protocol.py | 98 +++++++++++++++++ 4 files changed, 339 insertions(+) create mode 100644 lmcache/v1/distributed/l1_protocol.py create mode 100644 tests/v1/distributed/test_l1_config_maru.py create mode 100644 tests/v1/distributed/test_l1_protocol.py diff --git a/lmcache/v1/distributed/config.py b/lmcache/v1/distributed/config.py index 4deb8cfbb01..1f5da780531 100644 --- a/lmcache/v1/distributed/config.py +++ b/lmcache/v1/distributed/config.py @@ -128,6 +128,10 @@ class L1MemoryManagerConfig: devdax_size_in_bytes: int = 0 """ Optional Device-DAX overflow size for hybrid DRAM + DAX L1. """ + maru_config: "MaruL1Config | None" = None + """ Optional Maru CXL L1 backend; when set the L1 medium is a shared CXL + pool and the other L1 sizing fields above are ignored. """ + def __post_init__(self): self.init_size_in_bytes = min(self.init_size_in_bytes, self.size_in_bytes) @@ -185,6 +189,36 @@ class GdsL1Config: """Allocation alignment; cuFile/hipFile and O_DIRECT require 4 KiB.""" +@dataclass +class MaruL1Config: + """Config for the Maru CXL-backed L1 backend (mutually exclusive with the + pinned-DRAM / Device-DAX / GDS tiers). Defaults follow the non-MP + ``MaruBackend``; the allocator supplies ``chunk_size_bytes`` (from the KV + layout) and ``auto_connect=False``. + """ + + server_url: str + """MaruServer endpoint, e.g. ``maru://host:port`` or ``tcp://host:port``.""" + + pool_size_bytes: int + """CXL pool size to request from MaruServer, in bytes.""" + + instance_id: str | None = None + """Stable client id for ownership tracking; auto-generated if unset.""" + + timeout_ms: int = 5000 + """RPC timeout in milliseconds.""" + + use_async_rpc: bool = True + """Whether to use the async RPC path to MaruServer.""" + + max_inflight: int = 64 + """Maximum concurrent in-flight RPCs.""" + + eager_map: bool = True + """Eagerly mmap peer regions for cross-instance zero-copy reads.""" + + @dataclass class L1ManagerConfig: """ @@ -421,6 +455,36 @@ def add_storage_manager_args( help="Open the slab file with O_DIRECT (required for the GDS DMA fast " "path on ext4). Default True.", ) + # Maru L1 tier (optional, opt-in via --maru-server-url) + maru_group = parser.add_argument_group( + "Maru L1 tier", + "Optional CXL-backed shared L1 via Maru. Setting --maru-server-url " + "makes the L1 medium a cross-instance CXL pool instead of pinned DRAM; " + "the DRAM L1 settings (--l1-size-gb, --l1-use-lazy, --l1-init-size-gb) " + "are then ignored.", + ) + maru_group.add_argument( + "--maru-server-url", + type=str, + default=None, + help="MaruServer endpoint (maru://host:port or tcp://host:port). " + "Enables the Maru CXL L1 backend when set.", + ) + maru_group.add_argument( + "--maru-pool-size-gb", + type=float, + default=0.0, + help="CXL pool size to request from MaruServer (GB). Required (>0) " + "when --maru-server-url is set.", + ) + maru_group.add_argument( + "--maru-instance-id", + type=str, + default=None, + help="Stable client id reported to MaruServer for ownership tracking. " + "Auto-generated if omitted.", + ) + # L1 Manager Config (TTL settings) ttl_group = parser.add_argument_group( "L1 Manager TTL", "TTL configuration for L1 manager locks" @@ -549,6 +613,19 @@ def parse_args_to_config( Returns: StorageManagerConfig: The configuration object. """ + maru_config: MaruL1Config | None = None + if getattr(args, "maru_server_url", None): + pool_size_gb = getattr(args, "maru_pool_size_gb", 0.0) + if pool_size_gb <= 0: + raise ValueError( + "--maru-pool-size-gb must be > 0 when --maru-server-url is set" + ) + maru_config = MaruL1Config( + server_url=args.maru_server_url, + pool_size_bytes=int(pool_size_gb * (1 << 30)), + instance_id=getattr(args, "maru_instance_id", None), + ) + shm_name = getattr(args, "shm_name", None) if shm_name is None: memory_config = L1MemoryManagerConfig( @@ -556,6 +633,7 @@ def parse_args_to_config( use_lazy=args.l1_use_lazy, init_size_in_bytes=int(args.l1_init_size_gb * (1 << 30)), align_bytes=args.l1_align_bytes, + maru_config=maru_config, devdax_path=args.l1_devdax_path, ) else: @@ -565,6 +643,7 @@ def parse_args_to_config( init_size_in_bytes=int(args.l1_init_size_gb * (1 << 30)), align_bytes=args.l1_align_bytes, shm_name=shm_name, + maru_config=maru_config, devdax_path=args.l1_devdax_path, ) diff --git a/lmcache/v1/distributed/l1_protocol.py b/lmcache/v1/distributed/l1_protocol.py new file mode 100644 index 00000000000..e5920a2f8cf --- /dev/null +++ b/lmcache/v1/distributed/l1_protocol.py @@ -0,0 +1,113 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Structural ``Protocol`` for the L1 control-state manager. + +The seam that lets the stock controllers drive either ``L1Manager`` or the +CXL-backed ``MaruL1Manager``; structural (no inheritance), so ``l1_manager.py`` +stays untouched. The ``Fires ...`` line in each docstring is the listener-event +contract both backends must honor; see ``L1Manager`` for full error semantics. +""" + +# Standard +from typing import Any, Literal, Protocol, runtime_checkable + +# First Party +from lmcache.v1.distributed.api import MemoryLayoutDesc, ObjectKey +from lmcache.v1.distributed.error import L1Error +from lmcache.v1.distributed.internal_api import L1ManagerListener, L1MemoryDesc +from lmcache.v1.distributed.l1_manager import L1ObjectState, L1OperationResult + + +@runtime_checkable +class L1ManagerInterface(Protocol): + """L1 control surface shared by ``L1Manager`` and ``MaruL1Manager``.""" + + def register_listener(self, listener: L1ManagerListener) -> None: + """Register a listener for the ``on_l1_keys_*`` events.""" + ... + + def reserve_read( + self, keys: list[ObjectKey], extra_count: int = 0 + ) -> dict[ObjectKey, L1OperationResult]: + """Reserve read; ``1+extra_count`` holds/key. + Fires ``on_l1_keys_reserved_read``.""" + ... + + def unsafe_read(self, keys: list[ObjectKey]) -> dict[ObjectKey, L1OperationResult]: + """Return read-locked objects without new locks + (between reserve_read and finish_read).""" + ... + + def finish_read( + self, keys: list[ObjectKey], extra_count: int = 0 + ) -> dict[ObjectKey, L1Error]: + """Release ``1+extra_count`` holds/key. + Fires ``on_l1_keys_read_finished`` (+ ``on_l1_keys_deleted_by_manager`` for + temporaries dropped at count 0).""" + ... + + def reserve_write( + self, + keys: list[ObjectKey], + is_temporary: list[bool], + layout_desc: MemoryLayoutDesc, + mode: Literal["new", "update", "all"] = "all", + ) -> dict[ObjectKey, L1OperationResult]: + """Allocate + write-lock buffers; ``is_temporary[i]`` drops key i after read. + Fires ``on_l1_keys_reserved_write``.""" + ... + + def finish_write(self, keys: list[ObjectKey]) -> dict[ObjectKey, L1Error]: + """Release write locks. + Fires ``on_l1_keys_write_finished`` (write-through trigger).""" + ... + + def finish_write_and_reserve_read( + self, keys: list[ObjectKey], extra_count: int = 0 + ) -> dict[ObjectKey, L1OperationResult]: + """Finish write + take ``1+extra_count`` read holds/key (L2->L1 promote). + Fires ``on_l1_keys_finish_write_and_reserve_read`` (NOT write_finished).""" + ... + + def delete(self, keys: list[ObjectKey]) -> dict[ObjectKey, L1Error]: + """Delete unlocked keys (locked keys refused). + Fires ``on_l1_keys_deleted_by_manager`` for keys actually removed.""" + ... + + def touch_keys(self, keys: list[ObjectKey]) -> None: + """Mark keys accessed. + Fires ``on_l1_keys_accessed``.""" + ... + + def clear(self, force: bool = False) -> None: + """Free objects (``force`` frees locked too). + Fires ``on_l1_keys_deleted_by_manager`` for freed keys.""" + ... + + def is_key_evictable(self, key: ObjectKey) -> bool: + """Whether ``key`` exists and is unlocked (lock-free).""" + ... + + def get_memory_usage(self) -> tuple[int, int]: + """Return ``(used_bytes, total_bytes)`` of the L1 medium.""" + ... + + def get_l1_memory_desc(self) -> L1MemoryDesc | None: + """Return the L1 buffer descriptor, or ``None`` if not exposed.""" + ... + + def close(self) -> None: + """Free all objects and release resources.""" + ... + + def report_status(self) -> dict[str, Any]: + """Return a status dict of L1 cache state.""" + ... + + def get_object_state(self, key: ObjectKey) -> L1ObjectState | None: + """Return the internal state of ``key``, or ``None`` if absent.""" + ... + + def memcheck(self) -> bool: + """Run the medium's memory consistency check.""" + ... diff --git a/tests/v1/distributed/test_l1_config_maru.py b/tests/v1/distributed/test_l1_config_maru.py new file mode 100644 index 00000000000..da675a5174e --- /dev/null +++ b/tests/v1/distributed/test_l1_config_maru.py @@ -0,0 +1,49 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for Maru L1 config parsing.""" + +# Third Party +import pytest + +# First Party +from lmcache.v1.distributed.config import MaruL1Config, parse_args + + +def _args(*extra: str, l1_size_gb: str = "1") -> list[str]: + """Minimal required flags (eviction policy + L1 size) plus extras.""" + return [ + "--eviction-policy", + "LRU", + "--l1-size-gb", + l1_size_gb, + "--no-l1-use-lazy", + *extra, + ] + + +def test_no_maru_flags_leaves_maru_config_none(): + cfg = parse_args(_args()) + assert cfg.l1_manager_config.memory_config.maru_config is None + + +def test_maru_flags_build_maru_config(): + cfg = parse_args( + _args( + "--maru-server-url", + "maru://localhost:9000", + "--maru-pool-size-gb", + "8", + "--maru-instance-id", + "node-a", + ) + ) + maru = cfg.l1_manager_config.memory_config.maru_config + assert isinstance(maru, MaruL1Config) + assert maru.server_url == "maru://localhost:9000" + assert maru.pool_size_bytes == 8 * (1 << 30) + assert maru.instance_id == "node-a" + + +def test_maru_without_pool_size_raises(): + with pytest.raises(ValueError): + parse_args(_args("--maru-server-url", "maru://localhost:9000")) diff --git a/tests/v1/distributed/test_l1_protocol.py b/tests/v1/distributed/test_l1_protocol.py new file mode 100644 index 00000000000..4cac18a2b6c --- /dev/null +++ b/tests/v1/distributed/test_l1_protocol.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Structural conformance tests for ``L1ManagerInterface``. + +These bind the Protocol to ``L1Manager`` so that method-set or signature drift +between the two fails CI (the guard the sibling design relies on -- a +``runtime_checkable`` ``issubclass`` alone only checks method *names*). +""" + +# Standard +import inspect + +# Third Party +import pytest + +try: + # First Party + from lmcache.v1.distributed.l1_manager import L1Manager + from lmcache.v1.distributed.l1_protocol import L1ManagerInterface +except ImportError: + pytest.skip( + "L1Manager / L1ManagerInterface unavailable (native ext missing)", + allow_module_level=True, + ) + + +def _interface_methods() -> set[str]: + """Public method names declared on the Protocol.""" + return { + name + for name, val in vars(L1ManagerInterface).items() + if callable(val) and not name.startswith("_") + } + + +def _l1_manager_methods() -> set[str]: + """Public method names on the concrete L1Manager.""" + return { + name + for name in dir(L1Manager) + if not name.startswith("_") and callable(getattr(L1Manager, name)) + } + + +def _unwrap(fn): + """Recover the wrapped function from ``l1_mgr_synchronized``. + + That decorator is a plain closure without ``functools.wraps``, so + ``inspect.signature`` would otherwise see ``(self, *args, **kwargs)``. The + original function is the ``func`` free variable of the wrapper closure. + """ + code = getattr(fn, "__code__", None) + while code is not None and fn.__closure__ and "func" in code.co_freevars: + fn = fn.__closure__[code.co_freevars.index("func")].cell_contents + code = getattr(fn, "__code__", None) + return fn + + +def _params(fn) -> list[tuple[str, object, object]]: + """Parameter (name, kind, default) list excluding ``self``. + + Annotations are excluded on purpose: L1Manager leaves some param/return + annotations off, so only the call shape (names/kinds/defaults) is compared. + """ + return [ + (p.name, p.kind, p.default) + for p in inspect.signature(fn).parameters.values() + if p.name != "self" + ] + + +def test_interface_matches_l1_manager_surface(): + """The Protocol declares exactly L1Manager's public method set.""" + methods = _interface_methods() + assert methods == _l1_manager_methods() + assert len(methods) == 17 # tripwire against wholesale drift + + +def test_signatures_match_l1_manager(): + """Each Protocol method's call shape matches L1Manager's.""" + for name in _interface_methods(): + proto = _params(getattr(L1ManagerInterface, name)) + impl = _params(_unwrap(getattr(L1Manager, name))) + assert proto == impl, name + + +def test_l1_manager_conforms_to_interface(): + """The stock L1Manager satisfies the interface (structural, no instance).""" + assert issubclass(L1Manager, L1ManagerInterface) + + +def test_incomplete_class_does_not_conform(): + """A class missing interface methods is rejected (negative control).""" + + class Partial: + def close(self) -> None: ... + + assert not issubclass(Partial, L1ManagerInterface) From 59eddeef074b8d37721a9d201320732cb5ea1545 Mon Sep 17 00:00:00 2001 From: seohui-XCENA Date: Fri, 3 Jul 2026 18:44:56 +0900 Subject: [PATCH 02/28] [MP][Maru] Add MaruMemoryAllocator (CXL-backed L1 allocator) - thin wrapper over external maru_lmcache.CxlMemoryAdapter, lazy-imported - two-phase init_layout: build MaruHandler + CxlMemoryAdapter on first layout (single-model; layout mismatch rejected) - free/batched_free no-op (page lifecycle owned by MaruServer); abort_alloc discards an allocated-but-unregistered page - MaruL1Config -> maru.MaruConfig mapping; maru-only get_by_location / create_store_handle / handler surface for MaruL1Manager - tests use mocked maru runtime (no CXL required) Signed-off-by: seohui-XCENA --- .../memory_manager/maru_memory_allocator.py | 243 ++++++++++++++++++ .../distributed/test_maru_memory_allocator.py | 160 ++++++++++++ 2 files changed, 403 insertions(+) create mode 100644 lmcache/v1/distributed/memory_manager/maru_memory_allocator.py create mode 100644 tests/v1/distributed/test_maru_memory_allocator.py diff --git a/lmcache/v1/distributed/memory_manager/maru_memory_allocator.py b/lmcache/v1/distributed/memory_manager/maru_memory_allocator.py new file mode 100644 index 00000000000..fab36c4503f --- /dev/null +++ b/lmcache/v1/distributed/memory_manager/maru_memory_allocator.py @@ -0,0 +1,243 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""CXL-backed L1 allocator wrapping the external maru runtime. + +A thin wrapper over ``maru_lmcache.CxlMemoryAdapter``. The maru packages are +imported lazily in :meth:`init_layout` so this module stays importable without +the maru runtime. Two-phase init: the CXL pool is typed by the KV layout, so +the ``MaruHandler`` / ``CxlMemoryAdapter`` are built on the first layout +registration, not at construction. ``free`` is a no-op -- page lifecycle is +owned by MaruServer (pin/unpin/delete), not LMCache refcounts. + +Single-model per instance: the pool is fixed to the first layout; a different +layout is rejected. TODO(maru-multi-model): partition the pool per layout key. +""" + +# Standard +from typing import TYPE_CHECKING, List, Optional, Union + +# Third Party +import torch + +# First Party +from lmcache.integration.vllm.utils import get_size_bytes +from lmcache.logging import init_logger +from lmcache.v1.distributed.config import MaruL1Config +from lmcache.v1.memory_management import ( + MemoryAllocatorInterface, + MemoryFormat, + MemoryObj, +) + +if TYPE_CHECKING: + # Third Party + from maru import MaruHandler + from maru_handler.memory import AllocHandle + from maru_lmcache import CxlMemoryAdapter + +logger = init_logger(__name__) + + +def _to_tcp(url: str) -> str: + """Rewrite a ``maru://`` URL to the ``tcp://`` scheme maru expects.""" + if url.startswith("maru://"): + return "tcp://" + url[len("maru://") :] + return url + + +class MaruMemoryAllocator(MemoryAllocatorInterface): + """L1 allocator over a shared CXL pool (maru). See module docstring.""" + + def __init__(self, config: MaruL1Config) -> None: + self._config = config + self._handler: "MaruHandler | None" = None + self._cxl_adapter: "CxlMemoryAdapter | None" = None + self._shapes: list[torch.Size] | None = None + self._dtypes: list[torch.dtype] | None = None + self._fmt: MemoryFormat | None = None + self._chunk_size_in_tokens: int = 0 + self._single_token_size: int = 0 + + def init_layout( + self, + shapes: list[torch.Size], + dtypes: list[torch.dtype], + fmt: MemoryFormat, + chunk_size_in_tokens: int, + ) -> None: + """Bring up the CXL pool for the KV layout (idempotent per layout). + + Connects the MaruHandler and builds the CxlMemoryAdapter on the first + call. A later call with the same layout is a no-op; a different layout + is rejected (single-model; TODO(maru-multi-model)). Called serially at + startup -- not thread-safe. + """ + if chunk_size_in_tokens <= 0: + raise ValueError("chunk_size_in_tokens must be > 0") + + if self._cxl_adapter is not None: + if ( + self._shapes != shapes + or self._dtypes != dtypes + or self._fmt != fmt + or self._chunk_size_in_tokens != chunk_size_in_tokens + ): + raise ValueError( + "MaruMemoryAllocator is single-model: the KV layout changed " + "on a later init_layout call (TODO(maru-multi-model))." + ) + return + + chunk_bytes = get_size_bytes(shapes, dtypes) + if chunk_bytes <= 0 or chunk_bytes % chunk_size_in_tokens: + raise ValueError( + f"chunk size {chunk_bytes} bytes is not a positive multiple of " + f"{chunk_size_in_tokens} tokens" + ) + + # maru runtime is only needed once a layout is bound; import lazily so + # the module loads on non-maru deployments. + # Third Party + from maru import MaruConfig, MaruHandler + from maru_lmcache import CxlMemoryAdapter + + maru_config = MaruConfig( + server_url=_to_tcp(self._config.server_url), + instance_id=self._config.instance_id, + pool_size=self._config.pool_size_bytes, + chunk_size_bytes=chunk_bytes, + auto_connect=False, + timeout_ms=self._config.timeout_ms, + use_async_rpc=self._config.use_async_rpc, + max_inflight=self._config.max_inflight, + eager_map=self._config.eager_map, + ) + handler = MaruHandler(maru_config) + if not handler.connect(): + raise RuntimeError( + f"failed to connect MaruHandler to {self._config.server_url}" + ) + + self._handler = handler + try: + self._cxl_adapter = CxlMemoryAdapter( + handler=handler, + shapes=shapes, + dtypes=dtypes, + fmt=fmt, + chunk_size=handler.get_chunk_size(), + ) + except Exception: + handler.close() + self._handler = None + raise + self._shapes = shapes + self._dtypes = dtypes + self._fmt = fmt + self._chunk_size_in_tokens = chunk_size_in_tokens + self._single_token_size = chunk_bytes // chunk_size_in_tokens + + @property + def is_initialized(self) -> bool: + """Whether :meth:`init_layout` has bound the CXL pool.""" + return self._cxl_adapter is not None + + def _require_init(self, op: str) -> "CxlMemoryAdapter": + if self._cxl_adapter is None: + raise RuntimeError(f"MaruMemoryAllocator.{op} called before init_layout()") + return self._cxl_adapter + + # MemoryAllocatorInterface + + def allocate( + self, + shapes: Union[torch.Size, list[torch.Size]], + dtypes: Union[torch.dtype, list[torch.dtype]], + fmt: MemoryFormat = MemoryFormat.UNDEFINED, + allocator_type: Optional[str] = None, + ) -> Optional[MemoryObj]: + return self._require_init("allocate").allocate( + shapes, dtypes, fmt, allocator_type + ) + + def batched_allocate( + self, + shapes: Union[torch.Size, list[torch.Size]], + dtypes: Union[torch.dtype, list[torch.dtype]], + batch_size: int, + fmt: MemoryFormat = MemoryFormat.UNDEFINED, + allocator_type: Optional[str] = None, + ) -> Optional[List[MemoryObj]]: + return self._require_init("batched_allocate").batched_allocate( + shapes, dtypes, batch_size, fmt, allocator_type + ) + + def free(self, memory_obj: MemoryObj, allocator_type: Optional[str] = None) -> None: + """No-op: CXL page lifecycle is owned by MaruServer.""" + + def batched_free( + self, + memory_objs: List[MemoryObj], + allocator_type: Optional[str] = None, + update_stats: bool = True, + ) -> None: + """No-op: CXL page lifecycle is owned by MaruServer.""" + + # maru-specific surface (called by MaruL1Manager) + + def get_by_location( + self, + region_id: int, + page_index: int, + actual_size: int, + single_token_size: Optional[int] = None, + ) -> Optional[MemoryObj]: + """Materialize a zero-copy ``MemoryObj`` for a CXL page (no new alloc).""" + adapter = self._require_init("get_by_location") + if single_token_size is None: + single_token_size = self._single_token_size + return adapter.get_by_location( + region_id=region_id, + page_index=page_index, + actual_size=actual_size, + single_token_size=single_token_size, + ) + + def create_store_handle(self, memory_obj: MemoryObj) -> "AllocHandle": + """Build the maru store handle for an allocated page.""" + return self._require_init("create_store_handle").create_store_handle(memory_obj) + + def abort_alloc(self, memory_obj: MemoryObj) -> None: + """Discard an allocated-but-unregistered page (return it to the owner).""" + self._require_init("abort_alloc").free(memory_obj) + + @property + def handler(self) -> "MaruHandler": + """The connected MaruHandler (raises before init_layout()).""" + if self._handler is None: + raise RuntimeError( + "MaruMemoryAllocator.handler accessed before init_layout()" + ) + return self._handler + + @property + def single_token_size(self) -> int: + """Per-token KV byte size (raises before init_layout()).""" + if self._cxl_adapter is None: + raise RuntimeError( + "MaruMemoryAllocator.single_token_size accessed before init_layout()" + ) + return self._single_token_size + + # lifecycle + + def close(self) -> None: + if self._cxl_adapter is not None: + self._cxl_adapter.close() + self._cxl_adapter = None + if self._handler is not None: + self._handler.close() + self._handler = None + + def memcheck(self) -> bool: + return True diff --git a/tests/v1/distributed/test_maru_memory_allocator.py b/tests/v1/distributed/test_maru_memory_allocator.py new file mode 100644 index 00000000000..ab6efee43fd --- /dev/null +++ b/tests/v1/distributed/test_maru_memory_allocator.py @@ -0,0 +1,160 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for MaruMemoryAllocator (mocked maru runtime, no CXL required).""" + +# Standard +from unittest.mock import MagicMock +import sys + +# Third Party +import pytest +import torch + +try: + # First Party + from lmcache.v1.distributed.config import MaruL1Config + from lmcache.v1.distributed.memory_manager.maru_memory_allocator import ( + MaruMemoryAllocator, + _to_tcp, + ) + from lmcache.v1.memory_management import MemoryFormat +except ImportError: + pytest.skip("maru allocator deps unavailable", allow_module_level=True) + + +def _cfg() -> MaruL1Config: + return MaruL1Config(server_url="maru://localhost:9000", pool_size_bytes=1 << 30) + + +# 1 group, shape (4 tokens, 8 feats) fp16 -> 4*8*2 = 64 bytes; 4 tokens/chunk. +_LAYOUT = ([torch.Size([4, 8])], [torch.float16], MemoryFormat.KV_2LTD, 4) + + +@pytest.fixture +def maru_mocks(monkeypatch): + """Inject fake maru / maru_lmcache modules for init_layout's lazy imports.""" + handler = MagicMock() + handler.connect.return_value = True + handler.get_chunk_size.return_value = 64 + adapter = MagicMock() + maru_mod = MagicMock() + maru_mod.MaruHandler.return_value = handler + lmcache_mod = MagicMock() + lmcache_mod.CxlMemoryAdapter.return_value = adapter + monkeypatch.setitem(sys.modules, "maru", maru_mod) + monkeypatch.setitem(sys.modules, "maru_lmcache", lmcache_mod) + return handler, adapter, maru_mod + + +def test_to_tcp(): + assert _to_tcp("maru://h:1") == "tcp://h:1" + assert _to_tcp("tcp://h:1") == "tcp://h:1" + + +def test_methods_before_init_raise(): + alloc = MaruMemoryAllocator(_cfg()) + assert not alloc.is_initialized + with pytest.raises(RuntimeError): + alloc.allocate([torch.Size([4, 8])], [torch.float16]) + with pytest.raises(RuntimeError): + _ = alloc.handler + with pytest.raises(RuntimeError): + _ = alloc.single_token_size + + +def test_init_layout_builds_pool(maru_mocks): + handler, adapter, maru_mod = maru_mocks + alloc = MaruMemoryAllocator(_cfg()) + alloc.init_layout(*_LAYOUT) + + assert alloc.is_initialized + handler.connect.assert_called_once() + assert alloc.handler is handler + assert alloc.single_token_size == 16 # 64 bytes / 4 tokens + + _, kwargs = maru_mod.MaruConfig.call_args + assert kwargs["server_url"] == "tcp://localhost:9000" + assert kwargs["pool_size"] == 1 << 30 + assert kwargs["chunk_size_bytes"] == 64 + assert kwargs["auto_connect"] is False + + +def test_init_layout_same_layout_is_noop(maru_mocks): + handler, _, _ = maru_mocks + alloc = MaruMemoryAllocator(_cfg()) + alloc.init_layout(*_LAYOUT) + alloc.init_layout(*_LAYOUT) + handler.connect.assert_called_once() # not reconnected + + +def test_init_layout_mismatch_raises(maru_mocks): + alloc = MaruMemoryAllocator(_cfg()) + alloc.init_layout(*_LAYOUT) + with pytest.raises(ValueError): + alloc.init_layout( + [torch.Size([8, 8])], [torch.float16], MemoryFormat.KV_2LTD, 4 + ) + + +def test_init_layout_bad_chunk_raises(maru_mocks): + alloc = MaruMemoryAllocator(_cfg()) + with pytest.raises(ValueError): + alloc.init_layout( + [torch.Size([4, 8])], [torch.float16], MemoryFormat.KV_2LTD, 0 + ) + + +def test_init_layout_non_divisible_chunk_raises(maru_mocks): + alloc = MaruMemoryAllocator(_cfg()) + # 64 bytes is not a multiple of 5 tokens. + with pytest.raises(ValueError): + alloc.init_layout( + [torch.Size([4, 8])], [torch.float16], MemoryFormat.KV_2LTD, 5 + ) + + +def test_connect_failure_raises(maru_mocks): + handler, _, _ = maru_mocks + handler.connect.return_value = False + alloc = MaruMemoryAllocator(_cfg()) + with pytest.raises(RuntimeError): + alloc.init_layout(*_LAYOUT) + + +def test_delegation_and_lifecycle(maru_mocks): + _, adapter, _ = maru_mocks + alloc = MaruMemoryAllocator(_cfg()) + alloc.init_layout(*_LAYOUT) + obj = MagicMock() + + adapter.allocate.return_value = obj + assert alloc.allocate([torch.Size([4, 8])], [torch.float16]) is obj + + alloc.get_by_location(1, 2, 64) + adapter.get_by_location.assert_called_once_with( + region_id=1, page_index=2, actual_size=64, single_token_size=16 + ) + + alloc.create_store_handle(obj) + adapter.create_store_handle.assert_called_once_with(obj) + + # abort_alloc returns the page via the adapter's real free + alloc.abort_alloc(obj) + adapter.free.assert_called_once_with(obj) + + # free/batched_free are no-ops (lifecycle owned by MaruServer) + adapter.free.reset_mock() + alloc.free(obj) + alloc.batched_free([obj]) + adapter.free.assert_not_called() + + +def test_close_is_idempotent(maru_mocks): + handler, adapter, _ = maru_mocks + alloc = MaruMemoryAllocator(_cfg()) + alloc.init_layout(*_LAYOUT) + alloc.close() + adapter.close.assert_called_once() + handler.close.assert_called_once() + assert not alloc.is_initialized + alloc.close() # safe to call again From cf88c38ee5798d663afd4c2eeae13826dab17ffe Mon Sep 17 00:00:00 2001 From: seohui-XCENA Date: Fri, 3 Jul 2026 20:07:04 +0900 Subject: [PATCH 03/28] [MP][Maru] Add MaruL1Manager RPC control surface (read/write/delete/lifecycle) - sibling of L1Manager over the maru shared CXL pool: membership/read protection live in the MaruServer directory (pin_count), locally only in-flight staging (_pending_read refcount, _pending_write) - reserve_read: per-key independent pins (1+extra_count via one RPC), rollback on partial pin / retrieve failure / unresolvable page - reserve_write mode=new: local staged check + batch_exists dedup (cross-instance), all-or-nothing OOM; finish_write: batch_store, dup-skip is success, definitive False reclaims the page, unknown server state never recycles - delete: staged keys and pinned keys refuse with KEY_IS_LOCKED (exists() disambiguates the handler's pinned/missing conflation) - clear(force=False) keeps locked staging (stock parity); close drains - PARITY/MARU provenance comments; RPC reply length guards - tests: stateful fake maru runtime with fault-injection knobs, failure-path coverage, and a conformance suite parametrized over both L1Manager (CUDA) and MaruL1Manager Signed-off-by: seohui-XCENA --- lmcache/v1/distributed/maru_l1_manager.py | 733 ++++++++++++++++++ tests/v1/distributed/maru_fakes.py | 220 ++++++ .../test_l1_manager_conformance.py | 168 ++++ tests/v1/distributed/test_maru_l1_manager.py | 402 ++++++++++ 4 files changed, 1523 insertions(+) create mode 100644 lmcache/v1/distributed/maru_l1_manager.py create mode 100644 tests/v1/distributed/maru_fakes.py create mode 100644 tests/v1/distributed/test_l1_manager_conformance.py create mode 100644 tests/v1/distributed/test_maru_l1_manager.py diff --git a/lmcache/v1/distributed/maru_l1_manager.py b/lmcache/v1/distributed/maru_l1_manager.py new file mode 100644 index 00000000000..f6a4112d9c4 --- /dev/null +++ b/lmcache/v1/distributed/maru_l1_manager.py @@ -0,0 +1,733 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""CXL-backed sibling of ``L1Manager`` for MP mode (maru). + +``MaruL1Manager`` implements the ``L1ManagerInterface`` control surface for a +cross-instance shared L1 tier: membership and read protection live in the +MaruServer directory (``pin_count``), not in a local object table, so the +stock ``L1Manager`` state machine cannot be reused. Locally it keeps only +in-flight staging: ``_pending_write`` (reserved-but-unregistered pages) and +``_pending_read`` (pinned reads; the refcount balances N reserves = N pins = +N unpins). + +Provenance convention: ``PARITY(L1Manager.X)`` marks behavior mirrored from +the stock manager (keep in sync with ``l1_manager.py``); ``MARU:`` marks +maru-specific logic. + +Not implemented yet (later commits in this series): listener firing, the +L2->L1 promote transition, and the TTL orphan sweeper. Known gap: a pin whose +RPC reply is lost leaks server-side; reconciliation (per-instance pin ledger) +is a maru-side design item. +""" + +# Standard +from dataclasses import dataclass +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Concatenate, + Literal, + ParamSpec, + TypeVar, +) +import functools +import threading + +# Third Party +import torch + +# First Party +from lmcache.logging import init_logger +from lmcache.v1.distributed.api import MemoryLayoutDesc, ObjectKey +from lmcache.v1.distributed.config import L1ManagerConfig, MaruL1Config +from lmcache.v1.distributed.error import L1Error +from lmcache.v1.distributed.internal_api import L1ManagerListener, L1MemoryDesc +from lmcache.v1.distributed.l1_manager import ( + MAX_READ_LOCK_COUNT, + L1ObjectState, + L1OperationResult, +) +from lmcache.v1.distributed.memory_manager.maru_memory_allocator import ( + MaruMemoryAllocator, +) +from lmcache.v1.memory_management import MemoryFormat, MemoryObj + +if TYPE_CHECKING: + # Third Party + from maru import MaruHandler + from maru_handler.memory import MemoryInfo + +logger = init_logger(__name__) + +P = ParamSpec("P") +R = TypeVar("R") + + +def _maru_l1_synchronized( + func: Callable[Concatenate["MaruL1Manager", P], R], +) -> Callable[Concatenate["MaruL1Manager", P], R]: + """Serialize the method under the manager's non-reentrant lock.""" + + @functools.wraps(func) + def wrapper(self: "MaruL1Manager", *args: P.args, **kwargs: P.kwargs) -> R: + with self._lock: + return func(self, *args, **kwargs) + + return wrapper + + +def object_key_to_string(key: ObjectKey) -> str: + """Encode an ObjectKey as the flat string key MaruServer uses. + + Format ``@@@[@]`` -- + follows the MP L2 adapter key convention (every ObjectKey field encoded). + + Args: + key: The object key to encode. + + Returns: + The flat string key. + """ + base = ( + f"{key.model_name}@{key.kv_rank:08x}" + f"@{key.chunk_hash.hex()}@{key.object_group_id}" + ) + if key.cache_salt: + return f"{base}@{key.cache_salt}" + return base + + +def _clamp_extra_count(extra_count: int) -> int: + # PARITY(L1Manager._validate_extra_count): warn and clamp to [0, MAX-1]. + if extra_count < 0: + logger.warning( + "MaruL1Manager: extra_count=%d is invalid, clamping to 0", extra_count + ) + return 0 + upper = MAX_READ_LOCK_COUNT - 1 + if extra_count > upper: + logger.warning( + "MaruL1Manager: extra_count=%d exceeds limit=%d, clamping", + extra_count, + upper, + ) + return upper + return extra_count + + +@dataclass +class _PendingRead: + """A pinned read staged between reserve_read and the last finish_read.""" + + mem_obj: MemoryObj + refcount: int + is_temporary: bool = False + + +@dataclass +class _PendingWrite: + """A reserved page staged between reserve_write and finish_write.""" + + mem_obj: MemoryObj + is_temporary: bool + + +class MaruL1Manager: + """L1 control manager over the maru shared CXL pool. + + Structurally satisfies ``L1ManagerInterface`` (no nominal base). All + tiering decisions stay with the stock controllers; this class executes + them against the MaruServer directory and the CXL allocator. + """ + + def __init__(self, config: L1ManagerConfig) -> None: + maru_config = config.memory_config.maru_config + if maru_config is None: + raise ValueError("MaruL1Manager requires memory_config.maru_config") + self._config: MaruL1Config = maru_config + self._write_ttl_seconds = config.write_ttl_seconds + self._read_ttl_seconds = config.read_ttl_seconds + self._lock = threading.Lock() + self._allocator = MaruMemoryAllocator(maru_config) + self._pending_read: dict[ObjectKey, _PendingRead] = {} + self._pending_write: dict[ObjectKey, _PendingWrite] = {} + self._registered_listeners: list[L1ManagerListener] = [] + + def register_listener(self, listener: L1ManagerListener) -> None: + """Register a listener for ``on_l1_keys_*`` events. + + NOTE: events are not fired yet; firing lands with the tiering wiring. + + Args: + listener: The listener to register. + """ + # PARITY(L1Manager.register_listener): inline lock, append only. + with self._lock: + self._registered_listeners.append(listener) + + def _safe_unpin(self, handler: "MaruHandler", key_strs: list[str]) -> None: + """Release server pins, logging (not raising) on RPC failure.""" + if not key_strs: + return + try: + results = handler.batch_unpin(key_strs) + except Exception: + logger.exception( + "MaruL1Manager: batch_unpin failed for %d pins", len(key_strs) + ) + return + failed = sum(1 for ok in results if not ok) + if failed: + # A refused unpin means the server held no pin: a balance bug. + logger.warning( + "MaruL1Manager: %d/%d unpins had no pin to release", + failed, + len(key_strs), + ) + + @_maru_l1_synchronized + def reserve_read( + self, keys: list[ObjectKey], extra_count: int = 0 + ) -> dict[ObjectKey, L1OperationResult]: + """Pin keys on MaruServer and stage zero-copy views for reading. + + PARITY(L1Manager.reserve_read): per-key independent results; takes + ``1 + extra_count`` protection units per key. MARU: protection is the + cross-node server ``pin_count``; the local refcount balances the pins + so N finish_read calls release them all. + + Args: + keys: The list of object keys to reserve read access for. + extra_count: Extra protection units on top of the default 1. + + Returns: + A dictionary mapping each key to (L1Error, MemoryObj | None). + + Errors: + KEY_NOT_EXIST: Absent from the directory, unresolvable, or the + pin RPC failed. + KEY_NOT_READABLE: The key is mid-write on this instance. + """ + total = 1 + _clamp_extra_count(extra_count) + ret: dict[ObjectKey, L1OperationResult] = { + k: (L1Error.KEY_NOT_EXIST, None) for k in keys + } + for k in keys: + # PARITY(L1Manager.reserve_read): a mid-write key is not readable + # (distinct from a plain miss). It is unregistered, so the pin + # below misses and this pre-marked result stands. + if k in self._pending_write: + ret[k] = (L1Error.KEY_NOT_READABLE, None) + if not keys: + return ret + handler = self._allocator.handler + key_strs = [object_key_to_string(k) for k in keys] + + # MARU: one RPC takes `total` pins per key (repeat-encoding); a key + # is a hit only if all its pins landed. + pin_list = [ks for ks in key_strs for _ in range(total)] + try: + pin_results = handler.batch_pin(pin_list) + except Exception: + logger.exception("MaruL1Manager: batch_pin failed for %d keys", len(keys)) + return ret + if len(pin_results) != len(pin_list): + # Malformed reply: release whatever was reported taken. + logger.error( + "MaruL1Manager: batch_pin returned %d/%d results; rolling back", + len(pin_results), + len(pin_list), + ) + self._safe_unpin( + handler, + [ks for ks, ok in zip(pin_list, pin_results, strict=False) if ok], + ) + return ret + + hits: list[int] = [] + rollback: list[str] = [] # partial pins to release + for i, ks in enumerate(key_strs): + got = sum(1 for ok in pin_results[i * total : (i + 1) * total] if ok) + if got == total: + hits.append(i) + elif got: + rollback.extend([ks] * got) + + mem_infos: list["MemoryInfo | None"] = [] + if hits: + try: + mem_infos = handler.batch_retrieve([key_strs[i] for i in hits]) + except Exception: + logger.exception( + "MaruL1Manager: batch_retrieve failed for %d keys", len(hits) + ) + for i in hits: + rollback.extend([key_strs[i]] * total) + self._safe_unpin(handler, rollback) + return ret + # Normalize a malformed reply; missing tails roll back below. + mem_infos = list(mem_infos[: len(hits)]) + mem_infos += [None] * (len(hits) - len(mem_infos)) + + for i, mi in zip(hits, mem_infos, strict=False): + mem_obj = ( + self._allocator.get_by_location( + region_id=mi.region_id, + page_index=mi.page_index, + actual_size=len(mi.view), + ) + if mi is not None + else None + ) + if mem_obj is None: + # MARU: pinned but unresolvable (raced delete / pool miss). + rollback.extend([key_strs[i]] * total) + continue + k = keys[i] + staged = self._pending_read.get(k) + if staged is not None: + # Overlapping reserve: same CXL page, one staged object. + staged.refcount += total + ret[k] = (L1Error.SUCCESS, staged.mem_obj) + else: + self._pending_read[k] = _PendingRead(mem_obj=mem_obj, refcount=total) + ret[k] = (L1Error.SUCCESS, mem_obj) + + self._safe_unpin(handler, rollback) + return ret + + @_maru_l1_synchronized + def unsafe_read(self, keys: list[ObjectKey]) -> dict[ObjectKey, L1OperationResult]: + """Return staged read objects without taking new pins. + + Must be called between ``reserve_read`` and ``finish_read``. + + Args: + keys: The list of object keys to read. + + Returns: + A dictionary mapping each key to (L1Error, MemoryObj | None). + + Errors: + KEY_NOT_EXIST: The key has no staged read. + KEY_NOT_READABLE: The key is mid-write on this instance. + """ + ret: dict[ObjectKey, L1OperationResult] = {} + for k in keys: + entry = self._pending_read.get(k) + if entry is not None: + ret[k] = (L1Error.SUCCESS, entry.mem_obj) + elif k in self._pending_write: + ret[k] = (L1Error.KEY_NOT_READABLE, None) + else: + ret[k] = (L1Error.KEY_NOT_EXIST, None) + return ret + + @_maru_l1_synchronized + def finish_read( + self, keys: list[ObjectKey], extra_count: int = 0 + ) -> dict[ObjectKey, L1Error]: + """Release the protection taken by ``reserve_read``. + + Releases ``1 + extra_count`` units per key: the local refcount drops + and the same number of server pins are released; the staged entry is + dropped at refcount zero. + + Args: + keys: The list of object keys to finish read access for. + extra_count: Extra units to release on top of the default 1. + + Returns: + A dictionary mapping each key to an L1Error. + + Errors: + KEY_NOT_EXIST: The key has no staged read. + """ + total = 1 + _clamp_extra_count(extra_count) + ret: dict[ObjectKey, L1Error] = {} + to_unpin: list[str] = [] + for k in keys: + entry = self._pending_read.get(k) + if entry is None: + logger.warning("MaruL1Manager: finish read on unstaged key %s", k) + ret[k] = L1Error.KEY_NOT_EXIST + continue + # MARU: never release more than we hold (over-release would + # corrupt the server pin_count). + released = min(total, entry.refcount) + if released < total: + logger.warning( + "MaruL1Manager: finish read released %d/%d holds for key %s", + released, + total, + k, + ) + entry.refcount -= released + to_unpin.extend([object_key_to_string(k)] * released) + if entry.refcount <= 0: + del self._pending_read[k] + ret[k] = L1Error.SUCCESS + if to_unpin: + self._safe_unpin(self._allocator.handler, to_unpin) + return ret + + @_maru_l1_synchronized + def reserve_write( + self, + keys: list[ObjectKey], + is_temporary: list[bool], + layout_desc: MemoryLayoutDesc, + mode: Literal["new", "update", "all"] = "all", + ) -> dict[ObjectKey, L1OperationResult]: + """Allocate CXL pages and stage them for writing. + + PARITY(L1Manager.reserve_write): in "new" mode existing keys return + KEY_NOT_WRITABLE; allocation is all-or-nothing. MARU: "existing" + covers locally staged keys and directory-registered keys (the latter + check is the cross-instance dedup -- another instance stored it, so + the D2H copy is skipped). In-place update of a registered shared page + is not possible, so only ``mode="new"`` (the only mode MP callers + use) is supported. + + Args: + keys: The list of object keys to reserve write access for. + is_temporary: Per-key flag; temporary objects are dropped after + their read completes. + layout_desc: The memory layout for the allocation. + mode: Reservation mode; must be ``"new"``. + + Returns: + A dictionary mapping each key to (L1Error, MemoryObj | None). + + Raises: + ValueError: If ``mode`` is not ``"new"``. + + Errors: + KEY_NOT_WRITABLE: The key is staged locally or already registered. + OUT_OF_MEMORY: The CXL pool could not fit the batch. + """ + if mode != "new": + raise ValueError(f"MaruL1Manager supports mode='new' only, got {mode!r}") + ret: dict[ObjectKey, L1OperationResult] = {} + candidates: list[tuple[ObjectKey, bool]] = [] + for k, is_temp in zip(keys, is_temporary, strict=False): + if k in self._pending_write or k in self._pending_read: + ret[k] = (L1Error.KEY_NOT_WRITABLE, None) + else: + candidates.append((k, is_temp)) + if not candidates: + return ret + + try: + exists = self._allocator.handler.batch_exists( + [object_key_to_string(k) for k, _ in candidates] + ) + except Exception: + # MARU: existence unknown -- proceed; batch_store dup-skips later. + logger.exception( + "MaruL1Manager: batch_exists failed for %d keys", len(candidates) + ) + exists = [False] * len(candidates) + # Normalize a malformed reply; unknown tails allocate (dup-skip later). + exists = list(exists[: len(candidates)]) + exists += [False] * (len(candidates) - len(exists)) + need_allocate: list[tuple[ObjectKey, bool]] = [] + for (k, is_temp), ex in zip(candidates, exists, strict=False): + if ex: + ret[k] = (L1Error.KEY_NOT_WRITABLE, None) + else: + need_allocate.append((k, is_temp)) + if not need_allocate: + return ret + + objs = self._allocator.batched_allocate( + layout_desc.shapes, layout_desc.dtypes, len(need_allocate) + ) + if objs is None: + # PARITY(L1Manager.reserve_write): allocation failure marks the + # whole batch OUT_OF_MEMORY (batched_allocate is all-or-nothing). + for k, _ in need_allocate: + ret[k] = (L1Error.OUT_OF_MEMORY, None) + return ret + for (k, is_temp), obj in zip(need_allocate, objs, strict=False): + self._pending_write[k] = _PendingWrite(mem_obj=obj, is_temporary=is_temp) + ret[k] = (L1Error.SUCCESS, obj) + return ret + + @_maru_l1_synchronized + def finish_write(self, keys: list[ObjectKey]) -> dict[ObjectKey, L1Error]: + """Register staged pages in the MaruServer directory. + + Args: + keys: The list of object keys to finish write access for. + + Returns: + A dictionary mapping each key to an L1Error. + + Errors: + KEY_NOT_EXIST: The key was never reserved (or already finished). + KEY_IN_WRONG_STATE: Registration failed. + """ + ret: dict[ObjectKey, L1Error] = {} + staged: list[tuple[ObjectKey, _PendingWrite]] = [] + for k in keys: + entry = self._pending_write.pop(k, None) + if entry is None: + logger.warning("MaruL1Manager: finish write on unstaged key %s", k) + ret[k] = L1Error.KEY_NOT_EXIST + else: + staged.append((k, entry)) + if not staged: + return ret + + try: + handles = [ + self._allocator.create_store_handle(e.mem_obj) for _, e in staged + ] + except Exception: + # Pages never reached the server -- safe to reclaim now. + logger.exception( + "MaruL1Manager: create_store_handle failed for %d keys", len(staged) + ) + for k, entry in staged: + self._allocator.abort_alloc(entry.mem_obj) + ret[k] = L1Error.KEY_IN_WRONG_STATE + return ret + + try: + results = self._allocator.handler.batch_store( + [object_key_to_string(k) for k, _ in staged], handles + ) + except Exception: + # MARU: server state unknown -- the pages must NOT be recycled + # (a registered page must never return to the free list). + logger.exception( + "MaruL1Manager: batch_store failed for %d keys", len(staged) + ) + for k, _ in staged: + ret[k] = L1Error.KEY_IN_WRONG_STATE + return ret + + for i, (k, entry) in enumerate(staged): + ok = results[i] if i < len(results) else None + if ok: + # True covers newly-registered and dup-skip (page auto-freed). + ret[k] = L1Error.SUCCESS + elif ok is None: + # Missing reply entry: state unknown -- do not recycle. + ret[k] = L1Error.KEY_IN_WRONG_STATE + else: + # Definitively not registered -- reclaim the page. + self._allocator.abort_alloc(entry.mem_obj) + ret[k] = L1Error.KEY_IN_WRONG_STATE + return ret + + @_maru_l1_synchronized + def finish_write_and_reserve_read( + self, keys: list[ObjectKey], extra_count: int = 0 + ) -> dict[ObjectKey, L1OperationResult]: + """Atomically finish a write and take read holds (L2->L1 promote). + + Raises: + NotImplementedError: The promote transition lands in a later + commit of this series; no MP flow reaches it until the + tiering stack is wired. + """ + raise NotImplementedError( + "MaruL1Manager: finish_write_and_reserve_read is not wired yet" + ) + + @_maru_l1_synchronized + def delete(self, keys: list[ObjectKey]) -> dict[ObjectKey, L1Error]: + """Delete keys from the shared directory. + + PARITY(L1Manager.delete): a key held by any reader or writer refuses + with KEY_IS_LOCKED; the eviction policy keeps it and retries later. + + Args: + keys: The list of object keys to delete. + + Returns: + A dictionary mapping each key to an L1Error. + + Errors: + KEY_NOT_EXIST: The key is not in the directory. + KEY_IS_LOCKED: The key is staged locally, pinned on the server, + or the delete RPC failed (retryable). + """ + ret: dict[ObjectKey, L1Error] = {} + handler = self._allocator.handler + for k in keys: + # MARU: locally staged keys are pinned/write-held by construction. + if k in self._pending_read or k in self._pending_write: + ret[k] = L1Error.KEY_IS_LOCKED + continue + ks = object_key_to_string(k) + try: + if handler.delete(ks): + ret[k] = L1Error.SUCCESS + continue + # MARU: handler.delete conflates pinned and missing; one + # exists() round-trip disambiguates so pinned keys retry. + # TODO(maru): a tri-state delete RPC would remove this. + ret[k] = ( + L1Error.KEY_IS_LOCKED + if handler.exists(ks) + else L1Error.KEY_NOT_EXIST + ) + except Exception: + logger.exception("MaruL1Manager: delete failed for key %s", ks) + ret[k] = L1Error.KEY_IS_LOCKED + return ret + + def touch_keys(self, keys: list[ObjectKey]) -> None: + """Mark keys as accessed. + + NOTE: no-op until listener firing lands with the tiering wiring. + + Args: + keys: The list of object keys touched. + """ + + @_maru_l1_synchronized + def clear(self, force: bool = False) -> None: + """Release this instance's staging (unpin reads, reclaim write pages). + + PARITY(L1Manager.clear): ``force=False`` keeps locked entries -- all + maru staging is locked by construction, so it only logs. MARU: shared + directory data is never deleted (other instances may hold it). + + Args: + force: If True, drain in-flight staging too (unsafe, like stock). + """ + if not force: + if self._pending_read or self._pending_write: + logger.info( + "MaruL1Manager: clear kept %d staged reads / %d staged writes", + len(self._pending_read), + len(self._pending_write), + ) + return + self._drain_staging(force=True) + + def _drain_staging(self, force: bool) -> None: + """Unpin all staged reads (refcount times) and abort staged writes.""" + to_unpin: list[str] = [] + for k, entry in self._pending_read.items(): + to_unpin.extend([object_key_to_string(k)] * entry.refcount) + if force and (self._pending_read or self._pending_write): + logger.warning( + "MaruL1Manager: force-clear drops %d staged reads " + "(%d pins) and %d staged writes", + len(self._pending_read), + len(to_unpin), + len(self._pending_write), + ) + if to_unpin: + self._safe_unpin(self._allocator.handler, to_unpin) + for write_entry in self._pending_write.values(): + self._allocator.abort_alloc(write_entry.mem_obj) + self._pending_read.clear() + self._pending_write.clear() + + def is_key_evictable(self, key: ObjectKey) -> bool: + """Return whether ``key`` has no local hold (lock-free view). + + PARITY(L1Manager.is_key_evictable): deliberately lock-free; delete() + re-checks authoritatively (the server refuses pinned keys). MARU: + directory existence is not checked (remote); delete() drops stale + candidates with KEY_NOT_EXIST. + + Args: + key: The object key to check. + + Returns: + False if the key is staged for read or write locally. + """ + return key not in self._pending_read and key not in self._pending_write + + def get_memory_usage(self) -> tuple[int, int]: + """Return (used, total) bytes of this instance's owned CXL regions. + + MARU: usage is the owned-region allocation pressure -- the basis the + stock eviction watermark expects. Before the pool is up, reports the + configured capacity so watermark math stays sane. + + Returns: + A tuple of (used_bytes, total_bytes). + """ + if not self._allocator.is_initialized: + return 0, self._config.pool_size_bytes + try: + handler = self._allocator.handler + regions = handler.get_stats().get("store_regions") + if not regions: + return 0, self._config.pool_size_bytes + used = regions["total_allocated_pages"] * handler.get_chunk_size() + return used, regions["total_pool_size"] + except Exception: + logger.exception("MaruL1Manager: get_stats failed") + return 0, self._config.pool_size_bytes + + def get_l1_memory_desc(self) -> L1MemoryDesc | None: + """Return None: the shared pool has no single registerable region.""" + return None + + def close(self) -> None: + """Release staged state and tear down the allocator.""" + with self._lock: + self._drain_staging(force=False) + # PARITY(L1Manager.close): backing teardown outside the lock. + self._allocator.close() + + @_maru_l1_synchronized + def report_status(self) -> dict[str, Any]: + """Return a status dict describing the maru L1 state. + + Returns: + A dict with the stock L1 status keys plus ``backend``. + """ + used, total = self.get_memory_usage() + return { + "is_healthy": self._allocator.memcheck(), + "backend": "maru", + "total_object_count": len(self._pending_read) + len(self._pending_write), + "write_locked_count": len(self._pending_write), + "read_locked_count": len(self._pending_read), + "temporary_count": sum( + 1 for e in self._pending_read.values() if e.is_temporary + ) + + sum(1 for e in self._pending_write.values() if e.is_temporary), + "memory_used_bytes": used, + "memory_total_bytes": total, + "memory_usage_ratio": used / total if total > 0 else 0.0, + "write_ttl_seconds": self._write_ttl_seconds, + "read_ttl_seconds": self._read_ttl_seconds, + } + + def get_object_state(self, key: ObjectKey) -> L1ObjectState | None: + """Return None: membership lives in the MaruServer directory.""" + return None + + def memcheck(self) -> bool: + """Delegate to the allocator's consistency check.""" + return self._allocator.memcheck() + + @_maru_l1_synchronized + def register_kv_layout( + self, + shapes: list[torch.Size], + dtypes: list[torch.dtype], + fmt: MemoryFormat, + chunk_size_in_tokens: int, + ) -> None: + """Bind the KV layout, bringing up the CXL pool (idempotent per layout). + + Args: + shapes: Per-group tensor shapes of one chunk. + dtypes: Per-group dtypes of one chunk. + fmt: The KV memory format. + chunk_size_in_tokens: Tokens per chunk. + """ + self._allocator.init_layout(shapes, dtypes, fmt, chunk_size_in_tokens) diff --git a/tests/v1/distributed/maru_fakes.py b/tests/v1/distributed/maru_fakes.py new file mode 100644 index 00000000000..3a1caa79a2c --- /dev/null +++ b/tests/v1/distributed/maru_fakes.py @@ -0,0 +1,220 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Stateful in-memory fakes of the maru runtime for MaruL1Manager tests. + +``FakeMaruHandler`` models the MaruServer directory semantics the manager +relies on (per-key pins, dup-skip store, pinned-delete refusal) so contract +tests exercise real behavior instead of tautological mocks. Assertions should +stay on the manager's observable results plus the pin/free bookkeeping here. +""" + +# Standard +from dataclasses import dataclass +from types import SimpleNamespace +from unittest.mock import MagicMock + +# First Party +from lmcache.v1.distributed.config import ( + L1ManagerConfig, + L1MemoryManagerConfig, + MaruL1Config, +) +from lmcache.v1.distributed.maru_l1_manager import MaruL1Manager + + +@dataclass +class FakeMemoryInfo: + """Stand-in for maru's MemoryInfo (only the fields the manager reads).""" + + view: bytes + region_id: int + page_index: int + + +class FakeMaruHandler: + """Dict-backed MaruServer: directory + per-key pin counts. + + ``fail_*`` knobs make the next matching RPC raise; ``retrieve_none`` + models the real transport-failure mode (batch_retrieve returns all None + instead of raising). + """ + + def __init__(self, chunk_size: int = 64): + self.chunk_size = chunk_size + self.store_map: dict[str, tuple[int, int, int]] = {} # key -> (rid, pid, size) + self.pins: dict[str, int] = {} + self.unpin_log: list[str] = [] + self.fail_store_keys: set[str] = set() # keys batch_store reports False for + self.fail_pin = False + self.fail_retrieve = False + self.retrieve_none = False + self.fail_exists = False + self.fail_store = False + self.fail_delete = False + + def batch_pin(self, keys: list[str]) -> list[bool]: + if self.fail_pin: + raise RuntimeError("rpc fail") + out = [] + for ks in keys: + if ks in self.store_map: + self.pins[ks] = self.pins.get(ks, 0) + 1 + out.append(True) + else: + out.append(False) + return out + + def batch_unpin(self, keys: list[str]) -> list[bool]: + out = [] + for ks in keys: + self.unpin_log.append(ks) + if self.pins.get(ks, 0) > 0: + self.pins[ks] -= 1 + out.append(True) + else: + out.append(False) + return out + + def batch_retrieve(self, keys: list[str]) -> list[FakeMemoryInfo | None]: + if self.fail_retrieve: + raise RuntimeError("rpc fail") + if self.retrieve_none: + return [None] * len(keys) + out: list[FakeMemoryInfo | None] = [] + for ks in keys: + ent = self.store_map.get(ks) + out.append( + FakeMemoryInfo(view=b"x" * ent[2], region_id=ent[0], page_index=ent[1]) + if ent + else None + ) + return out + + def batch_store(self, keys: list[str], handles: list) -> list[bool]: + if self.fail_store: + raise RuntimeError("rpc fail") + out = [] + for ks, h in zip(keys, handles, strict=True): + if ks in self.fail_store_keys: + out.append(False) + elif ks in self.store_map: + out.append(True) # dup-skip is still a success + else: + self.store_map[ks] = (h.region_id, h.page_index, h.size) + out.append(True) + return out + + def batch_exists(self, keys: list[str]) -> list[bool]: + if self.fail_exists: + raise RuntimeError("rpc fail") + return [ks in self.store_map for ks in keys] + + def exists(self, key: str) -> bool: + return key in self.store_map + + def delete(self, key: str) -> bool: + if self.fail_delete: + raise RuntimeError("rpc fail") + if self.pins.get(key, 0) > 0: + return False # pinned: refused + return self.store_map.pop(key, None) is not None + + def get_chunk_size(self) -> int: + return self.chunk_size + + def get_stats(self) -> dict: + return { + "store_regions": { + "total_pool_size": 16 * self.chunk_size, + "total_allocated_pages": len(self.store_map), + } + } + + def close(self) -> None: + pass + + +class FakeCxlAdapter: + """Page-pool stand-in for maru_lmcache.CxlMemoryAdapter.""" + + def __init__(self, chunk_size: int = 64): + self.chunk_size = chunk_size + self.oom = False # set True to make allocation fail + self.resolve_none = False # get_by_location returns None (pool miss) + self.fail_handle = False # create_store_handle raises + self.freed: list[int] = [] # addresses returned via free (abort_alloc) + self._next_pid = 0 + self._pool: dict[tuple[int, int], MagicMock] = {} + + def _page(self, rid: int, pid: int) -> MagicMock: + obj = MagicMock(name=f"cxl-page-{rid}-{pid}") + obj.metadata.address = (rid << 32) | pid + obj.metadata.phy_size = self.chunk_size + return obj + + def batched_allocate(self, shapes, dtypes, batch_size, fmt=None, at=None): + if self.oom: + return None + out = [] + for _ in range(batch_size): + rid, pid = 0, self._next_pid + self._next_pid += 1 + obj = self._page(rid, pid) + self._pool[(rid, pid)] = obj + out.append(obj) + return out + + def create_store_handle(self, memory_obj) -> SimpleNamespace: + if self.fail_handle: + raise RuntimeError("bad handle") + addr = memory_obj.metadata.address + return SimpleNamespace( + region_id=addr >> 32, page_index=addr & 0xFFFFFFFF, size=self.chunk_size + ) + + def get_by_location(self, region_id, page_index, actual_size, single_token_size): + if self.resolve_none: + return None + # A page stored by another instance is materialized on demand + # (mirrors the eager-map view of a peer region). + key = (region_id, page_index) + if key not in self._pool: + self._pool[key] = self._page(region_id, page_index) + return self._pool[key] + + def free(self, memory_obj, allocator_type=None) -> None: + addr = memory_obj.metadata.address + if addr in self.freed: + # Mirror the real allocator's double-free detection. + raise AssertionError(f"double free of page {addr}") + self.freed.append(addr) + + def close(self) -> None: + pass + + +def make_maru_manager( + chunk_size: int = 64, +) -> tuple[MaruL1Manager, FakeMaruHandler, FakeCxlAdapter]: + """Build a MaruL1Manager wired to fresh fakes (post-init_layout state).""" + cfg = L1ManagerConfig( + memory_config=L1MemoryManagerConfig( + size_in_bytes=0, + use_lazy=False, + maru_config=MaruL1Config( + server_url="maru://localhost:5555", + pool_size_bytes=1 << 20, + instance_id="test", + ), + ), + write_ttl_seconds=600, + read_ttl_seconds=300, + ) + manager = MaruL1Manager(cfg) + handler = FakeMaruHandler(chunk_size) + adapter = FakeCxlAdapter(chunk_size) + allocator = manager._allocator + allocator._handler = handler + allocator._cxl_adapter = adapter + allocator._single_token_size = 16 + return manager, handler, adapter diff --git a/tests/v1/distributed/test_l1_manager_conformance.py b/tests/v1/distributed/test_l1_manager_conformance.py new file mode 100644 index 00000000000..475f47f8631 --- /dev/null +++ b/tests/v1/distributed/test_l1_manager_conformance.py @@ -0,0 +1,168 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Shared-contract conformance suite for L1ManagerInterface implementations. + +The same behavioral tests run against the stock ``L1Manager`` and +``MaruL1Manager`` so contract drift between the two backends fails CI. Only +behavior both backends promise is asserted here; backend-specific semantics +live in their own test files. + +CI ownership: the stock param needs CUDA; the maru param is skipped when the +maru package is not installed, so upstream CI never runs (or breaks on) the +maru side. +""" + +# Third Party +import pytest +import torch + +try: + # First Party + from lmcache.v1.distributed.api import MemoryLayoutDesc, ObjectKey + from lmcache.v1.distributed.config import L1ManagerConfig, L1MemoryManagerConfig + from lmcache.v1.distributed.error import L1Error + from lmcache.v1.distributed.l1_manager import L1Manager + + # Local + from .maru_fakes import make_maru_manager +except ImportError: + pytest.skip("L1 manager deps unavailable", allow_module_level=True) + + +def _has_maru() -> bool: + try: + # Third Party + import maru # noqa: F401 + + return True + except ImportError: + return False + + +_LAYOUT = MemoryLayoutDesc(shapes=[torch.Size([100, 2, 512])], dtypes=[torch.bfloat16]) + + +def _key(idx: int) -> ObjectKey: + return ObjectKey( + chunk_hash=idx.to_bytes(4, "big"), model_name="conf-model", kv_rank=1 + ) + + +@pytest.fixture( + params=[ + pytest.param( + "stock", + marks=pytest.mark.skipif( + not torch.cuda.is_available(), reason="CUDA is not available" + ), + ), + # CI-ownership gate (not a dependency: the fakes never import maru): + # the maru param runs only where maru is installed, so upstream CI + # neither runs nor breaks on the maru side. + pytest.param( + "maru", + marks=pytest.mark.skipif( + not _has_maru(), reason="maru runtime not installed" + ), + ), + ] +) +def manager(request): + """Yield one L1ManagerInterface implementation per param.""" + if request.param == "stock": + mgr = L1Manager( + L1ManagerConfig( + memory_config=L1MemoryManagerConfig( + size_in_bytes=128 * 1024 * 1024, + use_lazy=True, + init_size_in_bytes=64 * 1024 * 1024, + align_bytes=0x1000, + ), + write_ttl_seconds=600, + read_ttl_seconds=300, + ) + ) + else: + mgr, _, _ = make_maru_manager() + yield mgr + mgr.close() + + +def _write(mgr, keys): + res = mgr.reserve_write(keys, [False] * len(keys), _LAYOUT, mode="new") + assert all(err == L1Error.SUCCESS for err, _ in res.values()) + fin = mgr.finish_write(keys) + assert all(err == L1Error.SUCCESS for err in fin.values()) + + +def test_reserve_read_missing_key(manager): + k = _key(1) + assert manager.reserve_read([k])[k] == (L1Error.KEY_NOT_EXIST, None) + + +def test_write_read_roundtrip(manager): + """reserve_write -> finish_write -> reserve/unsafe/finish_read.""" + k = _key(2) + _write(manager, [k]) + + err, obj = manager.reserve_read([k])[k] + assert err == L1Error.SUCCESS and obj is not None + assert manager.unsafe_read([k])[k] == (L1Error.SUCCESS, obj) + assert manager.finish_read([k])[k] == L1Error.SUCCESS + + +def test_reserve_write_new_rejects_in_flight_key(manager): + k = _key(3) + first = manager.reserve_write([k], [False], _LAYOUT, mode="new") + assert first[k][0] == L1Error.SUCCESS + second = manager.reserve_write([k], [False], _LAYOUT, mode="new") + assert second[k] == (L1Error.KEY_NOT_WRITABLE, None) + + +def test_unsafe_read_without_reserve_is_not_success(manager): + k = _key(4) + _write(manager, [k]) + assert manager.unsafe_read([k])[k][0] != L1Error.SUCCESS + + +def test_delete_read_held_key_refused_then_retried(manager): + """A read-held key refuses deletion; it succeeds after release.""" + k = _key(5) + _write(manager, [k]) + manager.reserve_read([k]) + + assert manager.delete([k])[k] == L1Error.KEY_IS_LOCKED + manager.finish_read([k]) + assert manager.delete([k])[k] == L1Error.SUCCESS + + +def test_extra_count_balances_across_independent_finishes(manager): + """1+extra holds; that many finish_read calls release them all.""" + k = _key(6) + _write(manager, [k]) + manager.reserve_read([k], extra_count=2) + + assert manager.delete([k])[k] == L1Error.KEY_IS_LOCKED + for _ in range(3): + assert manager.finish_read([k])[k] == L1Error.SUCCESS + assert manager.delete([k])[k] == L1Error.SUCCESS + + +def test_clear_non_force_preserves_read_held_keys(manager): + """clear(force=False) must keep locked entries readable.""" + k = _key(8) + _write(manager, [k]) + manager.reserve_read([k]) + manager.clear() + assert manager.unsafe_read([k])[k][0] == L1Error.SUCCESS + manager.finish_read([k]) + + +def test_is_key_evictable_gates_on_read_hold(manager): + k = _key(7) + _write(manager, [k]) + assert manager.is_key_evictable(k) + manager.reserve_read([k]) + assert not manager.is_key_evictable(k) + manager.finish_read([k]) + assert manager.is_key_evictable(k) diff --git a/tests/v1/distributed/test_maru_l1_manager.py b/tests/v1/distributed/test_maru_l1_manager.py new file mode 100644 index 00000000000..8931f252625 --- /dev/null +++ b/tests/v1/distributed/test_maru_l1_manager.py @@ -0,0 +1,402 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for MaruL1Manager (fake maru runtime, no CXL required).""" + +# Third Party +import pytest +import torch + +try: + # First Party + from lmcache.v1.distributed.api import MemoryLayoutDesc, ObjectKey + from lmcache.v1.distributed.error import L1Error + from lmcache.v1.distributed.l1_protocol import L1ManagerInterface + from lmcache.v1.distributed.maru_l1_manager import ( + MaruL1Manager, + object_key_to_string, + ) + + # Local + from .maru_fakes import make_maru_manager + from .test_l1_protocol import _interface_methods, _params, _unwrap +except ImportError: + pytest.skip("maru manager deps unavailable", allow_module_level=True) + +_LAYOUT = MemoryLayoutDesc(shapes=[torch.Size([4, 8])], dtypes=[torch.float16]) + + +def _key(idx: int, salt: str = "") -> ObjectKey: + return ObjectKey( + chunk_hash=idx.to_bytes(4, "big"), + model_name="test-model", + kv_rank=0xABCD, + cache_salt=salt, + ) + + +def _seed(handler, key: ObjectKey, rid: int = 7, pid: int = 3, size: int = 64): + """Register a key in the fake directory (as if another instance stored it).""" + handler.store_map[object_key_to_string(key)] = (rid, pid, size) + + +def _store(manager, keys: list[ObjectKey]): + """Drive the full write path: reserve + finish.""" + res = manager.reserve_write(keys, [False] * len(keys), _LAYOUT, mode="new") + assert all(err == L1Error.SUCCESS for err, _ in res.values()) + fin = manager.finish_write(keys) + assert all(err == L1Error.SUCCESS for err in fin.values()) + + +# ========================================================================= +# interface conformance (binds MaruL1Manager to the shared Protocol) +# ========================================================================= + + +def test_conforms_to_l1_manager_interface(): + assert issubclass(MaruL1Manager, L1ManagerInterface) + + +def test_signatures_match_interface(): + """Each Protocol method's call shape matches MaruL1Manager's.""" + for name in _interface_methods(): + proto = _params(getattr(L1ManagerInterface, name)) + impl = _params(_unwrap(getattr(MaruL1Manager, name))) + assert proto == impl, name + + +# ========================================================================= +# reserve_read / unsafe_read / finish_read +# ========================================================================= + + +def test_reserve_read_hit_pins_and_stages(): + manager, handler, _ = make_maru_manager() + k = _key(1) + _seed(handler, k) + + res = manager.reserve_read([k]) + err, obj = res[k] + assert err == L1Error.SUCCESS and obj is not None + assert handler.pins[object_key_to_string(k)] == 1 + assert manager.unsafe_read([k])[k] == (L1Error.SUCCESS, obj) + + +def test_reserve_read_miss_leaves_no_pin(): + manager, handler, _ = make_maru_manager() + res = manager.reserve_read([_key(1)]) + assert res[_key(1)] == (L1Error.KEY_NOT_EXIST, None) + assert not handler.pins + + +def test_reserve_read_is_per_key_independent(): + """A miss in the middle must not shadow later hits (unlike prefix-stop).""" + manager, handler, _ = make_maru_manager() + k1, k2, k3 = _key(1), _key(2), _key(3) + _seed(handler, k1, pid=1) + _seed(handler, k3, pid=3) + + res = manager.reserve_read([k1, k2, k3]) + assert res[k1][0] == L1Error.SUCCESS + assert res[k2] == (L1Error.KEY_NOT_EXIST, None) + assert res[k3][0] == L1Error.SUCCESS + assert handler.pins[object_key_to_string(k3)] == 1 + + +def test_extra_count_takes_and_releases_n_pins(): + manager, handler, _ = make_maru_manager() + k = _key(1) + ks = object_key_to_string(k) + _seed(handler, k) + + manager.reserve_read([k], extra_count=2) + assert handler.pins[ks] == 3 + + # Three independent finish_read calls balance the three pins. + for _ in range(3): + assert manager.finish_read([k])[k] == L1Error.SUCCESS + assert handler.pins[ks] == 0 + assert manager.unsafe_read([k])[k] == (L1Error.KEY_NOT_EXIST, None) + + +def test_overlapping_reserve_accumulates_refcount(): + manager, handler, _ = make_maru_manager() + k = _key(1) + ks = object_key_to_string(k) + _seed(handler, k) + + _, obj1 = manager.reserve_read([k])[k] + _, obj2 = manager.reserve_read([k])[k] + assert obj1 is obj2 # one staged object per key + assert handler.pins[ks] == 2 + + manager.finish_read([k]) + assert manager.unsafe_read([k])[k][0] == L1Error.SUCCESS # still staged + manager.finish_read([k]) + assert handler.pins[ks] == 0 + + +def test_finish_read_unstaged_does_not_unpin(): + manager, handler, _ = make_maru_manager() + assert manager.finish_read([_key(1)])[_key(1)] == L1Error.KEY_NOT_EXIST + assert not handler.unpin_log + + +def test_finish_read_never_over_releases(): + manager, handler, _ = make_maru_manager() + k = _key(1) + _seed(handler, k) + manager.reserve_read([k]) # one pin + + assert manager.finish_read([k], extra_count=5)[k] == L1Error.SUCCESS + assert len(handler.unpin_log) == 1 # released min(6, refcount=1) + assert handler.pins[object_key_to_string(k)] == 0 + + +# ========================================================================= +# reserve_write / finish_write +# ========================================================================= + + +def test_write_path_registers_key(): + manager, handler, _ = make_maru_manager() + k = _key(1) + _store(manager, [k]) + assert object_key_to_string(k) in handler.store_map + # Now readable through the directory. + assert manager.reserve_read([k])[k][0] == L1Error.SUCCESS + + +def test_reserve_write_rejects_staged_and_registered_keys(): + manager, handler, _ = make_maru_manager() + staged, registered, fresh = _key(1), _key(2), _key(3) + manager.reserve_write([staged], [False], _LAYOUT, mode="new") + _seed(handler, registered) + + res = manager.reserve_write( + [staged, registered, fresh], [False] * 3, _LAYOUT, mode="new" + ) + assert res[staged] == (L1Error.KEY_NOT_WRITABLE, None) # locally staged + assert res[registered] == (L1Error.KEY_NOT_WRITABLE, None) # cross-instance dedup + assert res[fresh][0] == L1Error.SUCCESS + + +def test_reserve_write_oom_marks_whole_batch(): + manager, _, adapter = make_maru_manager() + adapter.oom = True + res = manager.reserve_write([_key(1), _key(2)], [False] * 2, _LAYOUT, mode="new") + assert all(v == (L1Error.OUT_OF_MEMORY, None) for v in res.values()) + + +def test_reserve_write_rejects_non_new_mode(): + manager, _, _ = make_maru_manager() + with pytest.raises(ValueError): + manager.reserve_write([_key(1)], [False], _LAYOUT, mode="update") + + +def test_finish_write_unstaged_key(): + manager, _, _ = make_maru_manager() + assert manager.finish_write([_key(1)])[_key(1)] == L1Error.KEY_NOT_EXIST + + +def test_finish_write_store_failure_reclaims_page(): + manager, handler, adapter = make_maru_manager() + k = _key(1) + handler.fail_store_keys.add(object_key_to_string(k)) + + res = manager.reserve_write([k], [False], _LAYOUT, mode="new") + addr = res[k][1].metadata.address + assert manager.finish_write([k])[k] == L1Error.KEY_IN_WRONG_STATE + assert addr in adapter.freed # aborted back to the owner free list + + +# ========================================================================= +# delete / clear / misc +# ========================================================================= + + +def test_delete_locally_staged_key_is_locked(): + manager, handler, _ = make_maru_manager() + k = _key(1) + _seed(handler, k) + manager.reserve_read([k]) + + assert manager.delete([k])[k] == L1Error.KEY_IS_LOCKED + assert object_key_to_string(k) in handler.store_map # server untouched + + +def test_delete_remotely_pinned_key_is_locked(): + manager, handler, _ = make_maru_manager() + k = _key(1) + _seed(handler, k) + handler.pins[object_key_to_string(k)] = 1 # pinned by another instance + + assert manager.delete([k])[k] == L1Error.KEY_IS_LOCKED + + +def test_delete_absent_and_present_keys(): + manager, handler, _ = make_maru_manager() + present, absent = _key(1), _key(2) + _seed(handler, present) + + res = manager.delete([present, absent]) + assert res[present] == L1Error.SUCCESS + assert res[absent] == L1Error.KEY_NOT_EXIST + assert object_key_to_string(present) not in handler.store_map + + +def test_clear_non_force_preserves_staging(): + manager, handler, adapter = make_maru_manager() + rk, wk = _key(1), _key(2) + _seed(handler, rk) + manager.reserve_read([rk]) + manager.reserve_write([wk], [False], _LAYOUT, mode="new") + + manager.clear() # force=False keeps locked (staged) entries + assert handler.pins[object_key_to_string(rk)] == 1 + assert not adapter.freed + assert manager.unsafe_read([rk])[rk][0] == L1Error.SUCCESS + + +def test_clear_force_balances_pins_and_reclaims_writes(): + manager, handler, adapter = make_maru_manager() + rk, wk = _key(1), _key(2) + _seed(handler, rk) + manager.reserve_read([rk], extra_count=1) # two pins + manager.reserve_write([wk], [False], _LAYOUT, mode="new") + + manager.clear(force=True) + assert handler.pins[object_key_to_string(rk)] == 0 + assert len(adapter.freed) == 1 # staged write page aborted + assert manager.unsafe_read([rk])[rk] == (L1Error.KEY_NOT_EXIST, None) + + +# ========================================================================= +# failure paths (pin balance + page-reclaim safety under RPC failures) +# ========================================================================= + + +def test_reserve_read_pin_rpc_failure_is_a_miss(): + manager, handler, _ = make_maru_manager() + k = _key(1) + _seed(handler, k) + handler.fail_pin = True + assert manager.reserve_read([k])[k] == (L1Error.KEY_NOT_EXIST, None) + assert not handler.pins + + +def test_reserve_read_retrieve_rpc_failure_rolls_back_pins(): + manager, handler, _ = make_maru_manager() + k = _key(1) + _seed(handler, k) + handler.fail_retrieve = True + assert manager.reserve_read([k])[k] == (L1Error.KEY_NOT_EXIST, None) + assert handler.pins[object_key_to_string(k)] == 0 + + +def test_reserve_read_all_none_retrieve_rolls_back_pins(): + """Real transport failure: batch_retrieve returns [None]*len, no raise.""" + manager, handler, _ = make_maru_manager() + k = _key(1) + _seed(handler, k) + handler.retrieve_none = True + assert manager.reserve_read([k], extra_count=1)[k] == (L1Error.KEY_NOT_EXIST, None) + assert handler.pins[object_key_to_string(k)] == 0 # both pins rolled back + + +def test_reserve_read_pool_miss_rolls_back_pins(): + manager, handler, adapter = make_maru_manager() + k = _key(1) + _seed(handler, k) + adapter.resolve_none = True + assert manager.reserve_read([k])[k] == (L1Error.KEY_NOT_EXIST, None) + assert handler.pins[object_key_to_string(k)] == 0 + + +def test_reserve_read_mid_write_key_not_readable(): + manager, _, _ = make_maru_manager() + k = _key(1) + manager.reserve_write([k], [False], _LAYOUT, mode="new") + assert manager.reserve_read([k])[k] == (L1Error.KEY_NOT_READABLE, None) + assert manager.unsafe_read([k])[k] == (L1Error.KEY_NOT_READABLE, None) + + +def test_finish_write_handle_failure_reclaims_all_pages(): + manager, _, adapter = make_maru_manager() + keys = [_key(1), _key(2)] + adapter.fail_handle = True + manager.reserve_write(keys, [False] * 2, _LAYOUT, mode="new") + res = manager.finish_write(keys) + assert all(err == L1Error.KEY_IN_WRONG_STATE for err in res.values()) + assert len(adapter.freed) == 2 # never reached the server: reclaimed + + +def test_finish_write_store_rpc_failure_never_recycles(): + """Unknown server state: pages must leak rather than be reused.""" + manager, handler, adapter = make_maru_manager() + k = _key(1) + handler.fail_store = True + manager.reserve_write([k], [False], _LAYOUT, mode="new") + assert manager.finish_write([k])[k] == L1Error.KEY_IN_WRONG_STATE + assert not adapter.freed + + +def test_finish_write_dup_skip_is_success_without_abort(): + manager, handler, adapter = make_maru_manager() + k = _key(1) + manager.reserve_write([k], [False], _LAYOUT, mode="new") + _seed(handler, k) # another instance registered it meanwhile + assert manager.finish_write([k])[k] == L1Error.SUCCESS # dup-skip + assert not adapter.freed # no local double-free + + +def test_reserve_write_exists_rpc_failure_proceeds(): + manager, handler, _ = make_maru_manager() + k = _key(1) + handler.fail_exists = True + res = manager.reserve_write([k], [False], _LAYOUT, mode="new") + assert res[k][0] == L1Error.SUCCESS # allocated; dup-skip at store time + + +def test_delete_rpc_failure_reports_locked(): + manager, handler, _ = make_maru_manager() + k = _key(1) + _seed(handler, k) + handler.fail_delete = True + assert manager.delete([k])[k] == L1Error.KEY_IS_LOCKED # retryable + + +def test_get_memory_usage_before_init_reports_pool_size(): + manager, _, _ = make_maru_manager() + manager._allocator._cxl_adapter = None # back to pre-init state + used, total = manager.get_memory_usage() + assert (used, total) == (0, 1 << 20) + + +def test_is_key_evictable_tracks_local_staging(): + manager, handler, _ = make_maru_manager() + k = _key(1) + _seed(handler, k) + assert manager.is_key_evictable(k) + manager.reserve_read([k]) + assert not manager.is_key_evictable(k) + manager.finish_read([k]) + assert manager.is_key_evictable(k) + + +def test_promote_not_implemented_yet(): + manager, _, _ = make_maru_manager() + with pytest.raises(NotImplementedError): + manager.finish_write_and_reserve_read([_key(1)]) + + +def test_report_status_keys(): + manager, handler, _ = make_maru_manager() + k = _key(1) + _seed(handler, k) + manager.reserve_read([k]) + + status = manager.report_status() + assert status["backend"] == "maru" + assert status["read_locked_count"] == 1 + assert status["is_healthy"] is True + assert status["memory_total_bytes"] > 0 From b3e22453b61308e97feaba03736ddbedcfa42ab7 Mon Sep 17 00:00:00 2001 From: seohui-XCENA Date: Mon, 6 Jul 2026 10:47:27 +0900 Subject: [PATCH 04/28] [MP][Maru] Fire L1 listener events and make touch_keys real - reserve_read/reserve_write/finish_write/finish_read/delete/clear now fire the on_l1_keys_* events (feeds the eviction LRU and the store controller), reporting only the keys that actually succeeded - finish_read: temporary reads (local staging, never directory-pinned) are reclaimed via abort_alloc at refcount zero and fire deleted_by_manager; normal reads still unpin -- _drain_staging mirrors the same branch - touch_keys: no-op -> fires on_l1_keys_accessed (unsynchronized, like stock) - event_bus (mp_observability) publish deferred: observability-only, no tiering-control impact - tests: RecordingListener fake; maru firing edge cases (hit-only, store-fail, temporary reclaim, removed-only, clear) + cross-backend conformance lifecycle test parametrized over stock and maru Signed-off-by: seohui-XCENA --- lmcache/v1/distributed/maru_l1_manager.py | 93 ++++++++++-- tests/v1/distributed/maru_fakes.py | 39 ++++++ .../test_l1_manager_conformance.py | 31 +++- tests/v1/distributed/test_maru_l1_manager.py | 132 +++++++++++++++++- 4 files changed, 278 insertions(+), 17 deletions(-) diff --git a/lmcache/v1/distributed/maru_l1_manager.py b/lmcache/v1/distributed/maru_l1_manager.py index f6a4112d9c4..3679776755d 100644 --- a/lmcache/v1/distributed/maru_l1_manager.py +++ b/lmcache/v1/distributed/maru_l1_manager.py @@ -14,10 +14,10 @@ the stock manager (keep in sync with ``l1_manager.py``); ``MARU:`` marks maru-specific logic. -Not implemented yet (later commits in this series): listener firing, the -L2->L1 promote transition, and the TTL orphan sweeper. Known gap: a pin whose -RPC reply is lost leaks server-side; reconciliation (per-instance pin ledger) -is a maru-side design item. +Not implemented yet (later commits in this series): the L2->L1 promote +transition (finish_write_and_reserve_read) and the TTL orphan sweeper. Known +gap: a pin whose RPC reply is lost leaks server-side; reconciliation +(per-instance pin ledger) is a maru-side design item. """ # Standard @@ -157,8 +157,6 @@ def __init__(self, config: L1ManagerConfig) -> None: def register_listener(self, listener: L1ManagerListener) -> None: """Register a listener for ``on_l1_keys_*`` events. - NOTE: events are not fired yet; firing lands with the tiering wiring. - Args: listener: The listener to register. """ @@ -270,6 +268,7 @@ def reserve_read( mem_infos = list(mem_infos[: len(hits)]) mem_infos += [None] * (len(hits) - len(mem_infos)) + successful_keys: list[ObjectKey] = [] for i, mi in zip(hits, mem_infos, strict=False): mem_obj = ( self._allocator.get_by_location( @@ -293,8 +292,13 @@ def reserve_read( else: self._pending_read[k] = _PendingRead(mem_obj=mem_obj, refcount=total) ret[k] = (L1Error.SUCCESS, mem_obj) + successful_keys.append(k) self._safe_unpin(handler, rollback) + # PARITY(L1Manager.reserve_read): notify listeners of the new read + # holds (feeds the eviction LRU and the store controller). + for listener in self._registered_listeners: + listener.on_l1_keys_reserved_read(successful_keys) return ret @_maru_l1_synchronized @@ -347,6 +351,9 @@ def finish_read( total = 1 + _clamp_extra_count(extra_count) ret: dict[ObjectKey, L1Error] = {} to_unpin: list[str] = [] + need_to_free: list[MemoryObj] = [] + need_to_free_keys: list[ObjectKey] = [] + successful_keys: list[ObjectKey] = [] for k in keys: entry = self._pending_read.get(k) if entry is None: @@ -364,12 +371,30 @@ def finish_read( k, ) entry.refcount -= released - to_unpin.extend([object_key_to_string(k)] * released) - if entry.refcount <= 0: - del self._pending_read[k] + # MARU: temporary reads are local staging (never directory-pinned), + # so at refcount zero the page is reclaimed, not unpinned. Normal + # reads release the server pins taken by reserve_read. + if entry.is_temporary: + if entry.refcount <= 0: + need_to_free.append(entry.mem_obj) + need_to_free_keys.append(k) + del self._pending_read[k] + else: + to_unpin.extend([object_key_to_string(k)] * released) + if entry.refcount <= 0: + del self._pending_read[k] ret[k] = L1Error.SUCCESS + successful_keys.append(k) if to_unpin: self._safe_unpin(self._allocator.handler, to_unpin) + for obj in need_to_free: + # MARU: unregistered local page -> discard through the allocator. + self._allocator.abort_alloc(obj) + # PARITY(L1Manager.finish_read): read_finished for every release; + # deleted_by_manager for temporary pages dropped at refcount zero. + for listener in self._registered_listeners: + listener.on_l1_keys_read_finished(successful_keys) + listener.on_l1_keys_deleted_by_manager(need_to_free_keys) return ret @_maru_l1_synchronized @@ -450,9 +475,15 @@ def reserve_write( for k, _ in need_allocate: ret[k] = (L1Error.OUT_OF_MEMORY, None) return ret + successful_keys: list[ObjectKey] = [] for (k, is_temp), obj in zip(need_allocate, objs, strict=False): self._pending_write[k] = _PendingWrite(mem_obj=obj, is_temporary=is_temp) ret[k] = (L1Error.SUCCESS, obj) + successful_keys.append(k) + # PARITY(L1Manager.reserve_write): notify listeners of the new write + # holds (the eviction LRU treats them as unevictable). + for listener in self._registered_listeners: + listener.on_l1_keys_reserved_write(successful_keys) return ret @_maru_l1_synchronized @@ -509,11 +540,13 @@ def finish_write(self, keys: list[ObjectKey]) -> dict[ObjectKey, L1Error]: ret[k] = L1Error.KEY_IN_WRONG_STATE return ret + successful_keys: list[ObjectKey] = [] for i, (k, entry) in enumerate(staged): ok = results[i] if i < len(results) else None if ok: # True covers newly-registered and dup-skip (page auto-freed). ret[k] = L1Error.SUCCESS + successful_keys.append(k) elif ok is None: # Missing reply entry: state unknown -- do not recycle. ret[k] = L1Error.KEY_IN_WRONG_STATE @@ -521,6 +554,11 @@ def finish_write(self, keys: list[ObjectKey]) -> dict[ObjectKey, L1Error]: # Definitively not registered -- reclaim the page. self._allocator.abort_alloc(entry.mem_obj) ret[k] = L1Error.KEY_IN_WRONG_STATE + # PARITY(L1Manager.finish_write): notify listeners of registered pages + # (the store controller stops re-storing them; must NOT be the promote + # event -- that is on_l1_keys_finish_write_and_reserve_read in C5). + for listener in self._registered_listeners: + listener.on_l1_keys_write_finished(successful_keys) return ret @_maru_l1_synchronized @@ -557,6 +595,7 @@ def delete(self, keys: list[ObjectKey]) -> dict[ObjectKey, L1Error]: or the delete RPC failed (retryable). """ ret: dict[ObjectKey, L1Error] = {} + successful_keys: list[ObjectKey] = [] handler = self._allocator.handler for k in keys: # MARU: locally staged keys are pinned/write-held by construction. @@ -567,6 +606,7 @@ def delete(self, keys: list[ObjectKey]) -> dict[ObjectKey, L1Error]: try: if handler.delete(ks): ret[k] = L1Error.SUCCESS + successful_keys.append(k) continue # MARU: handler.delete conflates pinned and missing; one # exists() round-trip disambiguates so pinned keys retry. @@ -579,16 +619,23 @@ def delete(self, keys: list[ObjectKey]) -> dict[ObjectKey, L1Error]: except Exception: logger.exception("MaruL1Manager: delete failed for key %s", ks) ret[k] = L1Error.KEY_IS_LOCKED + # PARITY(L1Manager.delete): notify listeners of the keys actually + # removed from the directory (the eviction LRU drops them). + for listener in self._registered_listeners: + listener.on_l1_keys_deleted_by_manager(successful_keys) return ret def touch_keys(self, keys: list[ObjectKey]) -> None: - """Mark keys as accessed. + """Mark keys as accessed, feeding the eviction LRU recency. - NOTE: no-op until listener firing lands with the tiering wiring. + PARITY(L1Manager.touch_keys): fires ``on_l1_keys_accessed`` without the + manager lock (matching stock); recency lives in the eviction policy. Args: keys: The list of object keys touched. """ + for listener in self._registered_listeners: + listener.on_l1_keys_accessed(keys) @_maru_l1_synchronized def clear(self, force: bool = False) -> None: @@ -609,13 +656,27 @@ def clear(self, force: bool = False) -> None: len(self._pending_write), ) return - self._drain_staging(force=True) + dropped = self._drain_staging(force=True) + # PARITY(L1Manager.clear): a force-clear notifies listeners of the drops. + for listener in self._registered_listeners: + listener.on_l1_keys_deleted_by_manager(dropped) + + def _drain_staging(self, force: bool) -> list[ObjectKey]: + """Unpin staged reads (refcount times) and abort staged writes. + + Args: + force: If True, log the dropped in-flight staging as a warning. - def _drain_staging(self, force: bool) -> None: - """Unpin all staged reads (refcount times) and abort staged writes.""" + Returns: + The keys dropped from staging (staged reads then staged writes). + """ to_unpin: list[str] = [] for k, entry in self._pending_read.items(): - to_unpin.extend([object_key_to_string(k)] * entry.refcount) + # MARU: temporary reads hold no server pin -- reclaim the page. + if entry.is_temporary: + self._allocator.abort_alloc(entry.mem_obj) + else: + to_unpin.extend([object_key_to_string(k)] * entry.refcount) if force and (self._pending_read or self._pending_write): logger.warning( "MaruL1Manager: force-clear drops %d staged reads " @@ -628,8 +689,10 @@ def _drain_staging(self, force: bool) -> None: self._safe_unpin(self._allocator.handler, to_unpin) for write_entry in self._pending_write.values(): self._allocator.abort_alloc(write_entry.mem_obj) + dropped = list(self._pending_read.keys()) + list(self._pending_write.keys()) self._pending_read.clear() self._pending_write.clear() + return dropped def is_key_evictable(self, key: ObjectKey) -> bool: """Return whether ``key`` has no local hold (lock-free view). diff --git a/tests/v1/distributed/maru_fakes.py b/tests/v1/distributed/maru_fakes.py index 3a1caa79a2c..d28641dcfc1 100644 --- a/tests/v1/distributed/maru_fakes.py +++ b/tests/v1/distributed/maru_fakes.py @@ -14,14 +14,53 @@ from unittest.mock import MagicMock # First Party +from lmcache.v1.distributed.api import ObjectKey from lmcache.v1.distributed.config import ( L1ManagerConfig, L1MemoryManagerConfig, MaruL1Config, ) +from lmcache.v1.distributed.internal_api import L1ManagerListener from lmcache.v1.distributed.maru_l1_manager import MaruL1Manager +class RecordingListener(L1ManagerListener): + """Records every ``on_l1_keys_*`` firing as ``(event, keys)``. + + Shared by the maru-specific and the cross-backend conformance suites so + both assert the same listener contract. ``kinds()`` returns the key lists + of each firing of one event, in order. + """ + + def __init__(self) -> None: + self.events: list[tuple[str, list[ObjectKey]]] = [] + + def on_l1_keys_reserved_read(self, keys: list[ObjectKey]) -> None: + self.events.append(("reserved_read", list(keys))) + + def on_l1_keys_read_finished(self, keys: list[ObjectKey]) -> None: + self.events.append(("read_finished", list(keys))) + + def on_l1_keys_reserved_write(self, keys: list[ObjectKey]) -> None: + self.events.append(("reserved_write", list(keys))) + + def on_l1_keys_write_finished(self, keys: list[ObjectKey]) -> None: + self.events.append(("write_finished", list(keys))) + + def on_l1_keys_finish_write_and_reserve_read(self, keys: list[ObjectKey]) -> None: + self.events.append(("finish_write_and_reserve_read", list(keys))) + + def on_l1_keys_deleted_by_manager(self, keys: list[ObjectKey]) -> None: + self.events.append(("deleted_by_manager", list(keys))) + + def on_l1_keys_accessed(self, keys: list[ObjectKey]) -> None: + self.events.append(("accessed", list(keys))) + + def kinds(self, name: str) -> list[list[ObjectKey]]: + """Return the key list of every recorded firing of event ``name``.""" + return [keys for event, keys in self.events if event == name] + + @dataclass class FakeMemoryInfo: """Stand-in for maru's MemoryInfo (only the fields the manager reads).""" diff --git a/tests/v1/distributed/test_l1_manager_conformance.py b/tests/v1/distributed/test_l1_manager_conformance.py index 475f47f8631..8715c1cd31d 100644 --- a/tests/v1/distributed/test_l1_manager_conformance.py +++ b/tests/v1/distributed/test_l1_manager_conformance.py @@ -24,7 +24,7 @@ from lmcache.v1.distributed.l1_manager import L1Manager # Local - from .maru_fakes import make_maru_manager + from .maru_fakes import RecordingListener, make_maru_manager except ImportError: pytest.skip("L1 manager deps unavailable", allow_module_level=True) @@ -166,3 +166,32 @@ def test_is_key_evictable_gates_on_read_hold(manager): assert not manager.is_key_evictable(k) manager.finish_read([k]) assert manager.is_key_evictable(k) + + +def test_listener_fires_across_lifecycle(manager): + """Both backends fire the same on_l1_keys_* events across a lifecycle. + + The eviction LRU and the store controller depend on this contract, so it + must hold identically for stock and maru. + """ + rec = RecordingListener() + manager.register_listener(rec) + k = _key(9) + + manager.reserve_write([k], [False], _LAYOUT, mode="new") + assert [k] in rec.kinds("reserved_write") + + manager.finish_write([k]) + assert [k] in rec.kinds("write_finished") + + manager.reserve_read([k]) + assert [k] in rec.kinds("reserved_read") + + manager.touch_keys([k]) + assert [k] in rec.kinds("accessed") + + manager.finish_read([k]) + assert [k] in rec.kinds("read_finished") + + assert manager.delete([k])[k] == L1Error.SUCCESS + assert [k] in rec.kinds("deleted_by_manager") diff --git a/tests/v1/distributed/test_maru_l1_manager.py b/tests/v1/distributed/test_maru_l1_manager.py index 8931f252625..78b95c929ee 100644 --- a/tests/v1/distributed/test_maru_l1_manager.py +++ b/tests/v1/distributed/test_maru_l1_manager.py @@ -13,11 +13,12 @@ from lmcache.v1.distributed.l1_protocol import L1ManagerInterface from lmcache.v1.distributed.maru_l1_manager import ( MaruL1Manager, + _PendingRead, object_key_to_string, ) # Local - from .maru_fakes import make_maru_manager + from .maru_fakes import RecordingListener, make_maru_manager from .test_l1_protocol import _interface_methods, _params, _unwrap except ImportError: pytest.skip("maru manager deps unavailable", allow_module_level=True) @@ -400,3 +401,132 @@ def test_report_status_keys(): assert status["read_locked_count"] == 1 assert status["is_healthy"] is True assert status["memory_total_bytes"] > 0 + + +# ========================================================================= +# listener firing (C4): feeds the eviction LRU + store controller +# ========================================================================= + + +def test_reserve_read_fires_only_for_hits(): + """A miss must not be reported as a reserved-read hold.""" + manager, handler, _ = make_maru_manager() + hit, miss = _key(1), _key(2) + _seed(handler, hit) + rec = RecordingListener() + manager.register_listener(rec) + + manager.reserve_read([hit, miss]) + assert rec.kinds("reserved_read") == [[hit]] + + +def test_reserve_and_finish_write_fire_with_registered_keys(): + manager, handler, _ = make_maru_manager() + k = _key(1) + rec = RecordingListener() + manager.register_listener(rec) + + manager.reserve_write([k], [False], _LAYOUT, mode="new") + assert rec.kinds("reserved_write") == [[k]] + manager.finish_write([k]) + assert rec.kinds("write_finished") == [[k]] + + +def test_finish_write_store_failure_is_not_fired(): + """A page that never registered must not notify write-finished.""" + manager, handler, _ = make_maru_manager() + k = _key(1) + handler.fail_store_keys.add(object_key_to_string(k)) + rec = RecordingListener() + manager.register_listener(rec) + + manager.reserve_write([k], [False], _LAYOUT, mode="new") + manager.finish_write([k]) + assert rec.kinds("write_finished") == [[]] # fired once, empty + + +def test_finish_read_normal_fires_read_finished_without_delete(): + """A directory-backed read unpins; it is never a manager delete.""" + manager, handler, _ = make_maru_manager() + k = _key(1) + _seed(handler, k) + manager.reserve_read([k]) + rec = RecordingListener() + manager.register_listener(rec) + + manager.finish_read([k]) + assert rec.kinds("read_finished") == [[k]] + assert all(ks == [] for ks in rec.kinds("deleted_by_manager")) + + +def test_finish_read_temporary_frees_page_and_fires_delete(): + """A temporary read (local staging) reclaims its page at refcount zero. + + The promote path that creates temporary reads lands in C5; here the entry + is staged white-box to exercise C4's refcount-zero reclaim branch. + """ + manager, handler, adapter = make_maru_manager() + k = _key(1) + page = adapter.batched_allocate(_LAYOUT.shapes, _LAYOUT.dtypes, 1)[0] + manager._pending_read[k] = _PendingRead(mem_obj=page, refcount=1, is_temporary=True) + rec = RecordingListener() + manager.register_listener(rec) + + assert manager.finish_read([k])[k] == L1Error.SUCCESS + # Temporary pages hold no server pin -- reclaimed, not unpinned. + assert page.metadata.address in adapter.freed + assert not handler.unpin_log + assert rec.kinds("read_finished") == [[k]] + assert rec.kinds("deleted_by_manager") == [[k]] + assert k not in manager._pending_read + + +def test_delete_fires_only_for_removed_keys(): + """Locked / absent keys are not reported as manager deletes.""" + manager, handler, _ = make_maru_manager() + removed, absent, locked = _key(1), _key(2), _key(3) + _seed(handler, removed) + _seed(handler, locked) + handler.pins[object_key_to_string(locked)] = 1 # pinned elsewhere + rec = RecordingListener() + manager.register_listener(rec) + + manager.delete([removed, absent, locked]) + assert rec.kinds("deleted_by_manager") == [[removed]] + + +def test_touch_keys_fires_accessed(): + manager, _, _ = make_maru_manager() + rec = RecordingListener() + manager.register_listener(rec) + + keys = [_key(1), _key(2)] + manager.touch_keys(keys) + assert rec.kinds("accessed") == [keys] + + +def test_clear_force_fires_deleted_for_dropped_staging(): + manager, handler, _ = make_maru_manager() + rk, wk = _key(1), _key(2) + _seed(handler, rk) + manager.reserve_read([rk]) + manager.reserve_write([wk], [False], _LAYOUT, mode="new") + rec = RecordingListener() + manager.register_listener(rec) + + manager.clear(force=True) + fired = rec.kinds("deleted_by_manager") + assert len(fired) == 1 + assert set(fired[0]) == {rk, wk} + + +def test_clear_non_force_fires_nothing(): + manager, handler, _ = make_maru_manager() + rk = _key(1) + _seed(handler, rk) + manager.reserve_read([rk]) + rec = RecordingListener() + manager.register_listener(rec) + + manager.clear() # keeps locked staging -> no drops + assert rec.events == [] From 80c271f7e70579de0d4522e8e6706952e312f719 Mon Sep 17 00:00:00 2001 From: seohui-XCENA Date: Mon, 6 Jul 2026 11:15:47 +0900 Subject: [PATCH 05/28] [MP][Maru] Implement finish_write_and_reserve_read (L2->L1 promote) - promote branches on the staged is_temporary flag (Decision A): temporary (default prefetch) moves the loaded page straight to read staging with no batch_store and no pin -- finish_read reclaims it at refcount zero; retained (prefetch_policy: retain) registers via batch_store then re-resolves the authoritative page with pins so a dup-skip that auto-freed our page still yields the winning shared page - fires on_l1_keys_finish_write_and_reserve_read, never on_l1_keys_write_finished (the latter would make the store controller re-store the key to L2) - extract _pin_retrieve_stage (shared by reserve_read + retained promote) and _store_staged (shared by finish_write + retained promote); reserve_read and finish_write refactored onto them, behavior unchanged - document the load-failure cleanup gap (finish_write->delete on failed keys publishes then removes; caller-side batch-abort is the future fix) - tests: temporary/retained promote, dup-skip re-resolve, extra_count pins, wrong-state/unstaged guards, store-failure, anti re-store event check; cross-backend conformance for retained promote, temporary drop, promote event Signed-off-by: seohui-XCENA --- lmcache/v1/distributed/maru_l1_manager.py | 331 +++++++++++++----- .../test_l1_manager_conformance.py | 38 ++ tests/v1/distributed/test_maru_l1_manager.py | 128 ++++++- 3 files changed, 401 insertions(+), 96 deletions(-) diff --git a/lmcache/v1/distributed/maru_l1_manager.py b/lmcache/v1/distributed/maru_l1_manager.py index 3679776755d..e4d74fa2cda 100644 --- a/lmcache/v1/distributed/maru_l1_manager.py +++ b/lmcache/v1/distributed/maru_l1_manager.py @@ -14,10 +14,15 @@ the stock manager (keep in sync with ``l1_manager.py``); ``MARU:`` marks maru-specific logic. -Not implemented yet (later commits in this series): the L2->L1 promote -transition (finish_write_and_reserve_read) and the TTL orphan sweeper. Known -gap: a pin whose RPC reply is lost leaks server-side; reconciliation -(per-instance pin ledger) is a maru-side design item. +Not implemented yet (later commits in this series): the TTL orphan sweeper. +Known gaps: (1) a pin whose RPC reply is lost leaks server-side; reconciliation +(per-instance pin ledger) is a maru-side design item. (2) the prefetch +controller's load-failure cleanup calls ``finish_write`` then ``delete`` on the +failed keys -- ``finish_write`` publishes the page to the shared directory, so a +peer that pins it in the window before ``delete`` makes ``delete`` refuse +(KEY_IS_LOCKED, which the caller ignores) and the key lingers. Rare (L2-load +failure + concurrent same-key lookup); a caller-side batch-abort is the clean +future fix. """ # Standard @@ -184,41 +189,29 @@ def _safe_unpin(self, handler: "MaruHandler", key_strs: list[str]) -> None: len(key_strs), ) - @_maru_l1_synchronized - def reserve_read( - self, keys: list[ObjectKey], extra_count: int = 0 - ) -> dict[ObjectKey, L1OperationResult]: - """Pin keys on MaruServer and stage zero-copy views for reading. + def _pin_retrieve_stage( + self, + keys: list[ObjectKey], + total: int, + ret: dict[ObjectKey, L1OperationResult], + successful_keys: list[ObjectKey], + ) -> None: + """Pin ``total`` units per key, resolve the page, and stage a read. - PARITY(L1Manager.reserve_read): per-key independent results; takes - ``1 + extra_count`` protection units per key. MARU: protection is the - cross-node server ``pin_count``; the local refcount balances the pins - so N finish_read calls release them all. + Shared by ``reserve_read`` and the retained-promote re-resolve. On a + hit sets ``ret[k] = (SUCCESS, obj)`` and appends ``k`` to + ``successful_keys``; overlapping reads accumulate the refcount. Partial + or unresolvable pins are rolled back; non-hit keys keep whatever value + the caller pre-set in ``ret``. Args: - keys: The list of object keys to reserve read access for. - extra_count: Extra protection units on top of the default 1. - - Returns: - A dictionary mapping each key to (L1Error, MemoryObj | None). - - Errors: - KEY_NOT_EXIST: Absent from the directory, unresolvable, or the - pin RPC failed. - KEY_NOT_READABLE: The key is mid-write on this instance. + keys: Keys to pin and stage (all pinned; misses roll back). + total: Protection units (pins) to take per key. + ret: Result map, mutated in place for hits. + successful_keys: List extended in place with each staged key. """ - total = 1 + _clamp_extra_count(extra_count) - ret: dict[ObjectKey, L1OperationResult] = { - k: (L1Error.KEY_NOT_EXIST, None) for k in keys - } - for k in keys: - # PARITY(L1Manager.reserve_read): a mid-write key is not readable - # (distinct from a plain miss). It is unregistered, so the pin - # below misses and this pre-marked result stands. - if k in self._pending_write: - ret[k] = (L1Error.KEY_NOT_READABLE, None) if not keys: - return ret + return handler = self._allocator.handler key_strs = [object_key_to_string(k) for k in keys] @@ -229,7 +222,7 @@ def reserve_read( pin_results = handler.batch_pin(pin_list) except Exception: logger.exception("MaruL1Manager: batch_pin failed for %d keys", len(keys)) - return ret + return if len(pin_results) != len(pin_list): # Malformed reply: release whatever was reported taken. logger.error( @@ -241,7 +234,7 @@ def reserve_read( handler, [ks for ks, ok in zip(pin_list, pin_results, strict=False) if ok], ) - return ret + return hits: list[int] = [] rollback: list[str] = [] # partial pins to release @@ -263,12 +256,11 @@ def reserve_read( for i in hits: rollback.extend([key_strs[i]] * total) self._safe_unpin(handler, rollback) - return ret + return # Normalize a malformed reply; missing tails roll back below. mem_infos = list(mem_infos[: len(hits)]) mem_infos += [None] * (len(hits) - len(mem_infos)) - successful_keys: list[ObjectKey] = [] for i, mi in zip(hits, mem_infos, strict=False): mem_obj = ( self._allocator.get_by_location( @@ -295,6 +287,108 @@ def reserve_read( successful_keys.append(k) self._safe_unpin(handler, rollback) + + def _store_staged( + self, staged: list[tuple[ObjectKey, _PendingWrite]] + ) -> tuple[list[ObjectKey], dict[ObjectKey, L1Error]]: + """Register staged write pages in the directory and classify the result. + + Shared by ``finish_write`` and the retained-promote path. + + Args: + staged: (key, pending write) pairs whose pages to register. + + Returns: + ``(registered, errors)``: ``registered`` are keys the server stored + or dup-skipped (the directory page is authoritative); ``errors`` + maps every other key to KEY_IN_WRONG_STATE. A handle-build failure + reclaims the pages; a store RPC failure leaves them (unknown server + state must never be recycled). + """ + errors: dict[ObjectKey, L1Error] = {} + if not staged: + return [], errors + try: + handles = [ + self._allocator.create_store_handle(e.mem_obj) for _, e in staged + ] + except Exception: + # Pages never reached the server -- safe to reclaim now. + logger.exception( + "MaruL1Manager: create_store_handle failed for %d keys", len(staged) + ) + for k, entry in staged: + self._allocator.abort_alloc(entry.mem_obj) + errors[k] = L1Error.KEY_IN_WRONG_STATE + return [], errors + + try: + results = self._allocator.handler.batch_store( + [object_key_to_string(k) for k, _ in staged], handles + ) + except Exception: + # MARU: server state unknown -- the pages must NOT be recycled + # (a registered page must never return to the free list). + logger.exception( + "MaruL1Manager: batch_store failed for %d keys", len(staged) + ) + for k, _ in staged: + errors[k] = L1Error.KEY_IN_WRONG_STATE + return [], errors + + registered: list[ObjectKey] = [] + for i, (k, entry) in enumerate(staged): + ok = results[i] if i < len(results) else None + if ok: + # True covers newly-registered and dup-skip (page auto-freed). + registered.append(k) + elif ok is None: + # Missing reply entry: state unknown -- do not recycle. + errors[k] = L1Error.KEY_IN_WRONG_STATE + else: + # Definitively not registered -- reclaim the page. + self._allocator.abort_alloc(entry.mem_obj) + errors[k] = L1Error.KEY_IN_WRONG_STATE + return registered, errors + + @_maru_l1_synchronized + def reserve_read( + self, keys: list[ObjectKey], extra_count: int = 0 + ) -> dict[ObjectKey, L1OperationResult]: + """Pin keys on MaruServer and stage zero-copy views for reading. + + PARITY(L1Manager.reserve_read): per-key independent results; takes + ``1 + extra_count`` protection units per key. MARU: protection is the + cross-node server ``pin_count``; the local refcount balances the pins + so N finish_read calls release them all. + + Args: + keys: The list of object keys to reserve read access for. + extra_count: Extra protection units on top of the default 1. + + Returns: + A dictionary mapping each key to (L1Error, MemoryObj | None). + + Errors: + KEY_NOT_EXIST: Absent from the directory, unresolvable, or the + pin RPC failed. + KEY_NOT_READABLE: The key is mid-write on this instance. + """ + total = 1 + _clamp_extra_count(extra_count) + ret: dict[ObjectKey, L1OperationResult] = { + k: (L1Error.KEY_NOT_EXIST, None) for k in keys + } + for k in keys: + # PARITY(L1Manager.reserve_read): a mid-write key is not readable + # (distinct from a plain miss). It is unregistered, so the pin + # below misses and this pre-marked result stands. + if k in self._pending_write: + ret[k] = (L1Error.KEY_NOT_READABLE, None) + if not keys: + return ret + successful_keys: list[ObjectKey] = [] + # MARU: pin + retrieve + resolve + stage (shared with retained promote). + self._pin_retrieve_stage(keys, total, ret, successful_keys) # PARITY(L1Manager.reserve_read): notify listeners of the new read # holds (feeds the eviction LRU and the store controller). for listener in self._registered_listeners: @@ -509,72 +603,123 @@ def finish_write(self, keys: list[ObjectKey]) -> dict[ObjectKey, L1Error]: ret[k] = L1Error.KEY_NOT_EXIST else: staged.append((k, entry)) - if not staged: - return ret - - try: - handles = [ - self._allocator.create_store_handle(e.mem_obj) for _, e in staged - ] - except Exception: - # Pages never reached the server -- safe to reclaim now. - logger.exception( - "MaruL1Manager: create_store_handle failed for %d keys", len(staged) - ) - for k, entry in staged: - self._allocator.abort_alloc(entry.mem_obj) - ret[k] = L1Error.KEY_IN_WRONG_STATE - return ret - - try: - results = self._allocator.handler.batch_store( - [object_key_to_string(k) for k, _ in staged], handles - ) - except Exception: - # MARU: server state unknown -- the pages must NOT be recycled - # (a registered page must never return to the free list). - logger.exception( - "MaruL1Manager: batch_store failed for %d keys", len(staged) - ) - for k, _ in staged: - ret[k] = L1Error.KEY_IN_WRONG_STATE - return ret - - successful_keys: list[ObjectKey] = [] - for i, (k, entry) in enumerate(staged): - ok = results[i] if i < len(results) else None - if ok: - # True covers newly-registered and dup-skip (page auto-freed). - ret[k] = L1Error.SUCCESS - successful_keys.append(k) - elif ok is None: - # Missing reply entry: state unknown -- do not recycle. - ret[k] = L1Error.KEY_IN_WRONG_STATE - else: - # Definitively not registered -- reclaim the page. - self._allocator.abort_alloc(entry.mem_obj) - ret[k] = L1Error.KEY_IN_WRONG_STATE + registered, errors = self._store_staged(staged) + ret.update(errors) + for k in registered: + ret[k] = L1Error.SUCCESS # PARITY(L1Manager.finish_write): notify listeners of registered pages # (the store controller stops re-storing them; must NOT be the promote - # event -- that is on_l1_keys_finish_write_and_reserve_read in C5). + # event -- that is on_l1_keys_finish_write_and_reserve_read). for listener in self._registered_listeners: - listener.on_l1_keys_write_finished(successful_keys) + listener.on_l1_keys_write_finished(registered) return ret @_maru_l1_synchronized def finish_write_and_reserve_read( self, keys: list[ObjectKey], extra_count: int = 0 ) -> dict[ObjectKey, L1OperationResult]: - """Atomically finish a write and take read holds (L2->L1 promote). + """Finish a write and take read holds in one step (L2->L1 promote). - Raises: - NotImplementedError: The promote transition lands in a later - commit of this series; no MP flow reaches it until the - tiering stack is wired. + Called by the prefetch controller after loading L2 bytes into a + write-reserved page. Branches on the staged ``is_temporary`` flag + (Decision A): + + - temporary (the default prefetch policy): the loaded page is private + staging -- moved straight to read staging without touching the shared + directory; finish_read reclaims it at refcount zero. + - retained (``prefetch_policy: retain``): the page is registered in the + directory (batch_store) and the authoritative page is re-resolved + with pins, so a dup-skip that auto-freed our page still yields the + winning shared page. + + Fires ``on_l1_keys_finish_write_and_reserve_read`` -- never + ``on_l1_keys_write_finished`` (that would make the store controller + re-store the promoted key to L2). + + Args: + keys: Keys to transition from write-staged to read-staged. + extra_count: Extra read holds on top of the default 1 (one per TP + worker for MLA models with TP > 1). + + Returns: + A dictionary mapping each key to (L1Error, MemoryObj | None). + + Errors: + KEY_NOT_EXIST: The key is not write-staged on this instance. + KEY_IN_WRONG_STATE: The key is already read-staged, or registration + or re-resolve failed. """ - raise NotImplementedError( - "MaruL1Manager: finish_write_and_reserve_read is not wired yet" - ) + total = 1 + _clamp_extra_count(extra_count) + ret: dict[ObjectKey, L1OperationResult] = { + k: (L1Error.KEY_NOT_EXIST, None) for k in keys + } + temp_staged: list[tuple[ObjectKey, _PendingWrite]] = [] + retain_staged: list[tuple[ObjectKey, _PendingWrite]] = [] + for k in keys: + entry = self._pending_write.get(k) + if entry is None: + # PARITY(L1Manager): a key not write-held cannot be promoted. + logger.warning("MaruL1Manager: promote on non-write-staged key %s", k) + continue + if k in self._pending_read: + # PARITY(L1Manager): a key already read-held is in wrong state. + ret[k] = (L1Error.KEY_IN_WRONG_STATE, None) + continue + if entry.is_temporary: + temp_staged.append((k, entry)) + else: + retain_staged.append((k, entry)) + + successful_keys: list[ObjectKey] = [] + # MARU temporary promote: private staging -- no batch_store, no pin. + # The loaded local page is authoritative; move it to read staging. + for k, entry in temp_staged: + del self._pending_write[k] + self._pending_read[k] = _PendingRead( + mem_obj=entry.mem_obj, refcount=total, is_temporary=True + ) + ret[k] = (L1Error.SUCCESS, entry.mem_obj) + successful_keys.append(k) + # MARU retained promote: register then re-resolve the authoritative page. + if retain_staged: + self._promote_retained(retain_staged, total, ret, successful_keys) + + for listener in self._registered_listeners: + listener.on_l1_keys_finish_write_and_reserve_read(successful_keys) + return ret + + def _promote_retained( + self, + retain_staged: list[tuple[ObjectKey, _PendingWrite]], + total: int, + ret: dict[ObjectKey, L1OperationResult], + successful_keys: list[ObjectKey], + ) -> None: + """Register retained-promote pages and stage authoritative reads. + + Pops the write staging, registers via ``_store_staged``, then pins and + re-resolves the directory page (a dup-skip auto-freed our own page, so + the pinned+retrieved page is the winning one). Keys that fail to + register or re-resolve are left at KEY_IN_WRONG_STATE. + + Args: + retain_staged: (key, pending write) pairs to register and stage. + total: Read holds (pins) to take per key. + ret: Result map, mutated in place. + successful_keys: List extended in place with each staged key. + """ + for k, _ in retain_staged: + self._pending_write.pop(k, None) + registered, errors = self._store_staged(retain_staged) + for k, err in errors.items(): + ret[k] = (err, None) + if not registered: + return + # A store that cannot be re-resolved to a read view is a wrong-state + # promote (the page is registered but this instance holds no read). + for k in registered: + ret[k] = (L1Error.KEY_IN_WRONG_STATE, None) + self._pin_retrieve_stage(registered, total, ret, successful_keys) @_maru_l1_synchronized def delete(self, keys: list[ObjectKey]) -> dict[ObjectKey, L1Error]: diff --git a/tests/v1/distributed/test_l1_manager_conformance.py b/tests/v1/distributed/test_l1_manager_conformance.py index 8715c1cd31d..e856c814719 100644 --- a/tests/v1/distributed/test_l1_manager_conformance.py +++ b/tests/v1/distributed/test_l1_manager_conformance.py @@ -168,6 +168,44 @@ def test_is_key_evictable_gates_on_read_hold(manager): assert manager.is_key_evictable(k) +def test_retained_promote_transitions_write_to_read(manager): + """Both backends: reserve_write -> promote -> read-held -> release.""" + k = _key(10) + assert manager.reserve_write([k], [False], _LAYOUT, mode="new")[k][0] == ( + L1Error.SUCCESS + ) + assert manager.finish_write_and_reserve_read([k])[k][0] == L1Error.SUCCESS + assert manager.unsafe_read([k])[k][0] == L1Error.SUCCESS + assert manager.delete([k])[k] == L1Error.KEY_IS_LOCKED # read-held + manager.finish_read([k]) + assert manager.delete([k])[k] == L1Error.SUCCESS + + +def test_temporary_promote_dropped_after_read(manager): + """Both backends: a temporary promote is gone after its read finishes.""" + k = _key(11) + assert manager.reserve_write([k], [True], _LAYOUT, mode="new")[k][0] == ( + L1Error.SUCCESS + ) + assert manager.finish_write_and_reserve_read([k])[k][0] == L1Error.SUCCESS + assert manager.unsafe_read([k])[k][0] == L1Error.SUCCESS + assert manager.finish_read([k])[k] == L1Error.SUCCESS + assert manager.unsafe_read([k])[k] == (L1Error.KEY_NOT_EXIST, None) + + +def test_promote_fires_promote_event(manager): + """Both backends fire finish_write_and_reserve_read, never write_finished.""" + k = _key(12) + manager.reserve_write([k], [False], _LAYOUT, mode="new") + rec = RecordingListener() + manager.register_listener(rec) + + manager.finish_write_and_reserve_read([k]) + assert [k] in rec.kinds("finish_write_and_reserve_read") + assert rec.kinds("write_finished") == [] + manager.finish_read([k]) + + def test_listener_fires_across_lifecycle(manager): """Both backends fire the same on_l1_keys_* events across a lifecycle. diff --git a/tests/v1/distributed/test_maru_l1_manager.py b/tests/v1/distributed/test_maru_l1_manager.py index 78b95c929ee..5e29c55c9d5 100644 --- a/tests/v1/distributed/test_maru_l1_manager.py +++ b/tests/v1/distributed/test_maru_l1_manager.py @@ -384,10 +384,132 @@ def test_is_key_evictable_tracks_local_staging(): assert manager.is_key_evictable(k) -def test_promote_not_implemented_yet(): +# ========================================================================= +# finish_write_and_reserve_read (C5): L2->L1 promote +# ========================================================================= + + +def test_temporary_promote_stages_read_without_registering(): + """Default prefetch: private staging -- no batch_store, no server pin.""" + manager, handler, _ = make_maru_manager() + k = _key(1) + page = manager.reserve_write([k], [True], _LAYOUT, mode="new")[k][1] + + err, obj = manager.finish_write_and_reserve_read([k])[k] + assert err == L1Error.SUCCESS and obj is page + assert object_key_to_string(k) not in handler.store_map # never registered + assert not handler.pins # temporary reads hold no pin + assert manager.unsafe_read([k])[k] == (L1Error.SUCCESS, page) + + +def test_temporary_promote_page_reclaimed_after_read(): + manager, handler, adapter = make_maru_manager() + k = _key(1) + page = manager.reserve_write([k], [True], _LAYOUT, mode="new")[k][1] + manager.finish_write_and_reserve_read([k]) + + assert manager.finish_read([k])[k] == L1Error.SUCCESS + assert page.metadata.address in adapter.freed # local page reclaimed + assert not handler.unpin_log # nothing was pinned + assert manager.unsafe_read([k])[k] == (L1Error.KEY_NOT_EXIST, None) + + +def test_retained_promote_registers_and_pins(): + """retain policy: batch_store + authoritative re-resolve with pins.""" + manager, handler, adapter = make_maru_manager() + k = _key(1) + ks = object_key_to_string(k) + manager.reserve_write([k], [False], _LAYOUT, mode="new") + + assert manager.finish_write_and_reserve_read([k])[k][0] == L1Error.SUCCESS + assert ks in handler.store_map # registered in the shared directory + assert handler.pins[ks] == 1 # read hold pinned + + assert manager.finish_read([k])[k] == L1Error.SUCCESS + assert handler.pins[ks] == 0 # unpinned, not freed + assert not adapter.freed + assert ks in handler.store_map # retained page survives the read + + +def test_retained_promote_extra_count_pins_n(): + manager, handler, _ = make_maru_manager() + k = _key(1) + manager.reserve_write([k], [False], _LAYOUT, mode="new") + + manager.finish_write_and_reserve_read([k], extra_count=2) + assert handler.pins[object_key_to_string(k)] == 3 # 1 + extra_count + + +def test_retained_promote_dup_skip_resolves_winner(): + """A peer registered the key first: re-resolve pins the winning page.""" + manager, handler, _ = make_maru_manager() + k = _key(1) + manager.reserve_write([k], [False], _LAYOUT, mode="new") + _seed(handler, k, rid=9, pid=5) # peer registered it meanwhile + + assert manager.finish_write_and_reserve_read([k])[k][0] == L1Error.SUCCESS + assert handler.pins[object_key_to_string(k)] == 1 + + +def test_promote_unstaged_key_is_not_exist(): manager, _, _ = make_maru_manager() - with pytest.raises(NotImplementedError): - manager.finish_write_and_reserve_read([_key(1)]) + assert manager.finish_write_and_reserve_read([_key(1)])[_key(1)] == ( + L1Error.KEY_NOT_EXIST, + None, + ) + + +def test_promote_already_read_staged_is_wrong_state(): + """The both-staged race the read-staged guard defends against.""" + manager, _, adapter = make_maru_manager() + k = _key(1) + manager.reserve_write([k], [False], _LAYOUT, mode="new") + rpage = adapter.batched_allocate(_LAYOUT.shapes, _LAYOUT.dtypes, 1)[0] + manager._pending_read[k] = _PendingRead(mem_obj=rpage, refcount=1) + + assert manager.finish_write_and_reserve_read([k])[k] == ( + L1Error.KEY_IN_WRONG_STATE, + None, + ) + assert k in manager._pending_write # guard returned before popping + + +def test_retained_promote_store_failure_is_wrong_state(): + manager, handler, _ = make_maru_manager() + k = _key(1) + handler.fail_store_keys.add(object_key_to_string(k)) + manager.reserve_write([k], [False], _LAYOUT, mode="new") + + assert manager.finish_write_and_reserve_read([k])[k] == ( + L1Error.KEY_IN_WRONG_STATE, + None, + ) + assert k not in manager._pending_write # write staging drained + + +def test_promote_fires_promote_event_not_write_finished(): + """Anti-#2744: promote must never fire write_finished (would re-store L2).""" + manager, handler, _ = make_maru_manager() + k = _key(1) + manager.reserve_write([k], [False], _LAYOUT, mode="new") + rec = RecordingListener() + manager.register_listener(rec) + + manager.finish_write_and_reserve_read([k]) + assert rec.kinds("finish_write_and_reserve_read") == [[k]] + assert rec.kinds("write_finished") == [] + + +def test_temporary_promote_fires_promote_event(): + manager, _, _ = make_maru_manager() + k = _key(1) + manager.reserve_write([k], [True], _LAYOUT, mode="new") + rec = RecordingListener() + manager.register_listener(rec) + + manager.finish_write_and_reserve_read([k]) + assert rec.kinds("finish_write_and_reserve_read") == [[k]] + assert rec.kinds("write_finished") == [] def test_report_status_keys(): From 7ee386d9171077f7d21c20daa0d0481ecad37655 Mon Sep 17 00:00:00 2001 From: seohui-XCENA Date: Mon, 6 Jul 2026 11:33:14 +0900 Subject: [PATCH 06/28] [MP][Maru] Add TTL sweeper for orphan write/read-pin reclaim - daemon sweeper (started in __init__, stopped in close) scans staging under the manager lock and reclaims entries whose TTL has elapsed: an expired write page is returned to the owner (abort_alloc); an expired read releases its pins (a temporary read reclaims its private page instead) - _PendingRead/_PendingWrite carry a monotonic deadline (default never); set at reserve/promote and refreshed on overlapping reserve (mirrors a stock re-lock extending the TTL). abandonment is a time judgement -- a refcount says how many holds exist, not whether they will ever be released - no listener fires on sweep: a late finish_read/unsafe_read then sees KEY_NOT_EXIST and recomputes (same failure path as a stock TTL expiry), and firing across the daemon thread would be a novel hazard for stock listeners - tests: sweeper lifecycle, expired write/read/temporary reclaim, live staging left intact, overlapping reserve refreshes the deadline, late finish is KEY_NOT_EXIST Signed-off-by: seohui-XCENA --- lmcache/v1/distributed/maru_l1_manager.py | 112 +++++++++++++++++-- tests/v1/distributed/test_maru_l1_manager.py | 89 +++++++++++++++ 2 files changed, 192 insertions(+), 9 deletions(-) diff --git a/lmcache/v1/distributed/maru_l1_manager.py b/lmcache/v1/distributed/maru_l1_manager.py index e4d74fa2cda..2c477cb88d0 100644 --- a/lmcache/v1/distributed/maru_l1_manager.py +++ b/lmcache/v1/distributed/maru_l1_manager.py @@ -14,9 +14,10 @@ the stock manager (keep in sync with ``l1_manager.py``); ``MARU:`` marks maru-specific logic. -Not implemented yet (later commits in this series): the TTL orphan sweeper. -Known gaps: (1) a pin whose RPC reply is lost leaks server-side; reconciliation -(per-instance pin ledger) is a maru-side design item. (2) the prefetch +A background sweeper reclaims staging whose TTL elapses (an abandoned client's +orphan write pages / read pins). Known gaps: (1) a pin whose RPC reply is lost +leaks server-side; reconciliation (per-instance pin ledger) is a maru-side +design item. (2) the prefetch controller's load-failure cleanup calls ``finish_write`` then ``delete`` on the failed keys -- ``finish_write`` publishes the page to the shared directory, so a peer that pins it in the window before ``delete`` makes ``delete`` refuse @@ -38,6 +39,7 @@ ) import functools import threading +import time # Third Party import torch @@ -123,19 +125,30 @@ def _clamp_extra_count(extra_count: int) -> int: @dataclass class _PendingRead: - """A pinned read staged between reserve_read and the last finish_read.""" + """A pinned read staged between reserve_read and the last finish_read. + + ``deadline`` is the monotonic time after which the sweeper treats the read + as orphaned; it defaults to never (the real reserve path sets a finite + value and refreshes it on overlapping reserves). + """ mem_obj: MemoryObj refcount: int is_temporary: bool = False + deadline: float = float("inf") @dataclass class _PendingWrite: - """A reserved page staged between reserve_write and finish_write.""" + """A reserved page staged between reserve_write and finish_write. + + ``deadline`` is the monotonic time after which the sweeper reclaims the + orphaned page (defaults to never; the reserve path sets a finite value). + """ mem_obj: MemoryObj is_temporary: bool + deadline: float = float("inf") class MaruL1Manager: @@ -158,6 +171,16 @@ def __init__(self, config: L1ManagerConfig) -> None: self._pending_read: dict[ObjectKey, _PendingRead] = {} self._pending_write: dict[ObjectKey, _PendingWrite] = {} self._registered_listeners: list[L1ManagerListener] = [] + # MARU: W1+R1 orphan sweeper -- a crashed/abandoned client leaves write + # pages or read pins behind; reclaim them once their TTL elapses. + self._sweep_interval = max( + 1.0, min(self._write_ttl_seconds, self._read_ttl_seconds) / 4 + ) + self._stop_event = threading.Event() + self._sweeper = threading.Thread( + target=self._sweep_loop, name="maru-l1-sweeper", daemon=True + ) + self._sweeper.start() def register_listener(self, listener: L1ManagerListener) -> None: """Register a listener for ``on_l1_keys_*`` events. @@ -276,13 +299,18 @@ def _pin_retrieve_stage( rollback.extend([key_strs[i]] * total) continue k = keys[i] + deadline = time.monotonic() + self._read_ttl_seconds staged = self._pending_read.get(k) if staged is not None: # Overlapping reserve: same CXL page, one staged object. + # Refresh the TTL (mirrors a stock re-lock extending it). staged.refcount += total + staged.deadline = deadline ret[k] = (L1Error.SUCCESS, staged.mem_obj) else: - self._pending_read[k] = _PendingRead(mem_obj=mem_obj, refcount=total) + self._pending_read[k] = _PendingRead( + mem_obj=mem_obj, refcount=total, deadline=deadline + ) ret[k] = (L1Error.SUCCESS, mem_obj) successful_keys.append(k) @@ -351,6 +379,64 @@ def _store_staged( errors[k] = L1Error.KEY_IN_WRONG_STATE return registered, errors + def _sweep_loop(self) -> None: + """Daemon loop: sweep expired staging until ``close`` stops it.""" + while not self._stop_event.wait(self._sweep_interval): + try: + self._sweep_once() + except Exception: + logger.exception("MaruL1Manager: sweep iteration failed") + + def _sweep_once(self) -> None: + """Reclaim staging whose TTL elapsed (orphan write pages / read pins). + + MARU: abandonment can only be judged by time -- a refcount says how + many holds exist, not whether they will ever be released. Mirrors the + stock write_lock/read_lock TTL: an expired write page is returned to + the owner (abort_alloc) and an expired read releases its pins (a + temporary read reclaims its private page instead). No listener fires -- + a late finish_read/unsafe_read then sees KEY_NOT_EXIST and recomputes + (the same failure path as a stock TTL expiry), and firing across the + daemon thread would be a novel hazard for the stock listeners. + """ + now = time.monotonic() + with self._lock: + expired_writes = [ + k for k, e in self._pending_write.items() if e.deadline <= now + ] + expired_reads = [ + k for k, e in self._pending_read.items() if e.deadline <= now + ] + if not expired_writes and not expired_reads: + return + for k in expired_writes: + write_entry = self._pending_write.pop(k) + try: + self._allocator.abort_alloc(write_entry.mem_obj) + except Exception: + logger.exception( + "MaruL1Manager: sweep failed to reclaim write page %s", k + ) + to_unpin: list[str] = [] + for k in expired_reads: + read_entry = self._pending_read.pop(k) + if read_entry.is_temporary: + try: + self._allocator.abort_alloc(read_entry.mem_obj) + except Exception: + logger.exception( + "MaruL1Manager: sweep failed to reclaim read page %s", k + ) + else: + to_unpin.extend([object_key_to_string(k)] * read_entry.refcount) + if to_unpin: + self._safe_unpin(self._allocator.handler, to_unpin) + logger.warning( + "MaruL1Manager: swept %d orphan write(s) / %d orphan read(s)", + len(expired_writes), + len(expired_reads), + ) + @_maru_l1_synchronized def reserve_read( self, keys: list[ObjectKey], extra_count: int = 0 @@ -570,8 +656,11 @@ def reserve_write( ret[k] = (L1Error.OUT_OF_MEMORY, None) return ret successful_keys: list[ObjectKey] = [] + deadline = time.monotonic() + self._write_ttl_seconds for (k, is_temp), obj in zip(need_allocate, objs, strict=False): - self._pending_write[k] = _PendingWrite(mem_obj=obj, is_temporary=is_temp) + self._pending_write[k] = _PendingWrite( + mem_obj=obj, is_temporary=is_temp, deadline=deadline + ) ret[k] = (L1Error.SUCCESS, obj) successful_keys.append(k) # PARITY(L1Manager.reserve_write): notify listeners of the new write @@ -676,7 +765,10 @@ def finish_write_and_reserve_read( for k, entry in temp_staged: del self._pending_write[k] self._pending_read[k] = _PendingRead( - mem_obj=entry.mem_obj, refcount=total, is_temporary=True + mem_obj=entry.mem_obj, + refcount=total, + is_temporary=True, + deadline=time.monotonic() + self._read_ttl_seconds, ) ret[k] = (L1Error.SUCCESS, entry.mem_obj) successful_keys.append(k) @@ -883,7 +975,9 @@ def get_l1_memory_desc(self) -> L1MemoryDesc | None: return None def close(self) -> None: - """Release staged state and tear down the allocator.""" + """Stop the sweeper, release staged state, and tear down the allocator.""" + self._stop_event.set() + self._sweeper.join(timeout=self._sweep_interval + 5.0) with self._lock: self._drain_staging(force=False) # PARITY(L1Manager.close): backing teardown outside the lock. diff --git a/tests/v1/distributed/test_maru_l1_manager.py b/tests/v1/distributed/test_maru_l1_manager.py index 5e29c55c9d5..7236d7542c1 100644 --- a/tests/v1/distributed/test_maru_l1_manager.py +++ b/tests/v1/distributed/test_maru_l1_manager.py @@ -2,6 +2,9 @@ """Tests for MaruL1Manager (fake maru runtime, no CXL required).""" +# Standard +import time + # Third Party import pytest import torch @@ -512,6 +515,92 @@ def test_temporary_promote_fires_promote_event(): assert rec.kinds("write_finished") == [] +# ========================================================================= +# TTL sweeper (C6): reclaim orphan write pages / read pins +# ========================================================================= + + +def test_sweeper_thread_starts_and_stops(): + manager, _, _ = make_maru_manager() + assert manager._sweeper.is_alive() + manager.close() + manager._sweeper.join(timeout=2) + assert not manager._sweeper.is_alive() + + +def test_sweep_reclaims_expired_write(): + manager, _, adapter = make_maru_manager() + k = _key(1) + page = manager.reserve_write([k], [False], _LAYOUT, mode="new")[k][1] + manager._pending_write[k].deadline = time.monotonic() - 1 # force expiry + + manager._sweep_once() + assert k not in manager._pending_write + assert page.metadata.address in adapter.freed # returned to the owner + + +def test_sweep_reclaims_expired_read_unpins(): + manager, handler, adapter = make_maru_manager() + k = _key(1) + ks = object_key_to_string(k) + _seed(handler, k) + manager.reserve_read([k], extra_count=1) # two pins + manager._pending_read[k].deadline = time.monotonic() - 1 + + manager._sweep_once() + assert k not in manager._pending_read + assert handler.pins[ks] == 0 # both pins released + assert not adapter.freed # registered page is not freed + + +def test_sweep_reclaims_expired_temporary_read(): + manager, handler, adapter = make_maru_manager() + k = _key(1) + page = manager.reserve_write([k], [True], _LAYOUT, mode="new")[k][1] + manager.finish_write_and_reserve_read([k]) # temporary read staged + manager._pending_read[k].deadline = time.monotonic() - 1 + + manager._sweep_once() + assert k not in manager._pending_read + assert page.metadata.address in adapter.freed # private page reclaimed + assert not handler.unpin_log # temporary reads hold no pin + + +def test_sweep_leaves_live_staging(): + manager, handler, adapter = make_maru_manager() + k = _key(1) + _seed(handler, k) + manager.reserve_read([k]) # deadline ~ now + read_ttl (300s) + + manager._sweep_once() + assert k in manager._pending_read # not expired + assert handler.pins[object_key_to_string(k)] == 1 + assert not adapter.freed + + +def test_overlapping_reserve_refreshes_read_deadline(): + manager, handler, _ = make_maru_manager() + k = _key(1) + _seed(handler, k) + manager.reserve_read([k]) + manager._pending_read[k].deadline = time.monotonic() - 1 # go stale + + manager.reserve_read([k]) # overlap -> refresh + assert manager._pending_read[k].deadline > time.monotonic() + + +def test_finish_read_after_sweep_is_not_exist(): + """A late finish after a sweep behaves like a stock TTL expiry.""" + manager, handler, _ = make_maru_manager() + k = _key(1) + _seed(handler, k) + manager.reserve_read([k]) + manager._pending_read[k].deadline = time.monotonic() - 1 + manager._sweep_once() + + assert manager.finish_read([k])[k] == L1Error.KEY_NOT_EXIST + + def test_report_status_keys(): manager, handler, _ = make_maru_manager() k = _key(1) From ad2de115b45bdcb1cd5880ead74544d45ba403e4 Mon Sep 17 00:00:00 2001 From: seohui-XCENA Date: Mon, 6 Jul 2026 11:38:29 +0900 Subject: [PATCH 07/28] [MP] Widen storage controller l1_manager param to L1ManagerInterface - store/prefetch/eviction controllers type their l1_manager param as the structural L1ManagerInterface instead of the concrete L1Manager, so either L1Manager or MaruL1Manager can drive them; runtime behavior unchanged - the controllers call only Protocol methods (reserve/finish read+write, finish_write_and_reserve_read, delete, is_key_evictable, get_memory_usage, register_listener), so L1Manager still satisfies the param structurally Signed-off-by: seohui-XCENA --- .../v1/distributed/storage_controllers/eviction_controller.py | 4 ++-- .../v1/distributed/storage_controllers/prefetch_controller.py | 4 ++-- .../v1/distributed/storage_controllers/store_controller.py | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/lmcache/v1/distributed/storage_controllers/eviction_controller.py b/lmcache/v1/distributed/storage_controllers/eviction_controller.py index 56bdb6ac1ce..6bedf910349 100644 --- a/lmcache/v1/distributed/storage_controllers/eviction_controller.py +++ b/lmcache/v1/distributed/storage_controllers/eviction_controller.py @@ -20,7 +20,7 @@ EvictionAction, EvictionDestination, ) -from lmcache.v1.distributed.l1_manager import L1Manager +from lmcache.v1.distributed.l1_protocol import L1ManagerInterface from lmcache.v1.distributed.l2_adapters.base import L2AdapterInterface from lmcache.v1.distributed.storage_controller import StorageControllerInterface from lmcache.v1.mp_observability.event import Event, EventType @@ -96,7 +96,7 @@ class L1EvictionController(EvictionController): def __init__( self, - l1_manager: L1Manager, + l1_manager: L1ManagerInterface, eviction_config: EvictionConfig, ): super().__init__() diff --git a/lmcache/v1/distributed/storage_controllers/prefetch_controller.py b/lmcache/v1/distributed/storage_controllers/prefetch_controller.py index a1bfc010902..03420042c9d 100644 --- a/lmcache/v1/distributed/storage_controllers/prefetch_controller.py +++ b/lmcache/v1/distributed/storage_controllers/prefetch_controller.py @@ -32,7 +32,7 @@ TrimPolicy, ) from lmcache.v1.distributed.error import L1Error -from lmcache.v1.distributed.l1_manager import L1Manager +from lmcache.v1.distributed.l1_protocol import L1ManagerInterface from lmcache.v1.distributed.l2_adapters.base import L2AdapterInterface, L2TaskId from lmcache.v1.distributed.storage_controller import StorageControllerInterface from lmcache.v1.distributed.storage_controllers.adapter_lifecycle import ( @@ -207,7 +207,7 @@ class PrefetchController(StorageControllerInterface): def __init__( self, - l1_manager: L1Manager, + l1_manager: L1ManagerInterface, l2_adapters: list[L2AdapterInterface], adapter_descriptors: list[AdapterDescriptor], policy: PrefetchPolicy, diff --git a/lmcache/v1/distributed/storage_controllers/store_controller.py b/lmcache/v1/distributed/storage_controllers/store_controller.py index 44c20eed17a..dedbf485734 100644 --- a/lmcache/v1/distributed/storage_controllers/store_controller.py +++ b/lmcache/v1/distributed/storage_controllers/store_controller.py @@ -21,7 +21,7 @@ from lmcache.v1.distributed.api import ObjectKey from lmcache.v1.distributed.error import L1Error from lmcache.v1.distributed.internal_api import L1ManagerListener -from lmcache.v1.distributed.l1_manager import L1Manager +from lmcache.v1.distributed.l1_protocol import L1ManagerInterface from lmcache.v1.distributed.l2_adapters.base import L2AdapterInterface, L2TaskId from lmcache.v1.distributed.storage_controller import StorageControllerInterface from lmcache.v1.distributed.storage_controllers.adapter_lifecycle import ( @@ -225,7 +225,7 @@ class StoreController(StorageControllerInterface): def __init__( self, - l1_manager: L1Manager, + l1_manager: L1ManagerInterface, l2_adapters: list[L2AdapterInterface], adapter_descriptors: list[AdapterDescriptor], policy: StorePolicy, From fea97ed96772a018ea1296519e2e8f95b8077964 Mon Sep 17 00:00:00 2001 From: seohui-XCENA Date: Mon, 6 Jul 2026 12:07:50 +0900 Subject: [PATCH 08/28] [MP][Maru] Select MaruL1Manager in StorageManager and add register_kv_layout - StorageManager.__init__ selects the L1 backend: MaruL1Manager when memory_config.maru_config is set, else the stock L1Manager; typed as the shared L1ManagerInterface so the controllers drive either unchanged - new StorageManager.register_kv_layout: maru-gated (isinstance) forward that brings up the CXL pool once the layout is known, rejecting >1 object group (single-model maru limit); a no-op for stock. Engine call site lands later - get_l1_memory_desc is now L1MemoryDesc | None (maru has no single registerable region); the l1_memory_desc property raises if accessed while None (its only consumer, p2p, is rejected at startup for maru) - widen SerdeL2AdapterWrapper l1_manager param to L1ManagerInterface (like the controllers); collect its reserve_write buffers in one pass so the success path types as list[MemoryObj] - tests: maru vs stock selection, register_kv_layout forward / multi-group reject / stock no-op Signed-off-by: seohui-XCENA --- .../distributed/l2_adapters/serde_wrapper.py | 16 ++-- lmcache/v1/distributed/storage_manager.py | 66 +++++++++++++- .../test_distributed_storage_manager.py | 85 +++++++++++++++++++ 3 files changed, 156 insertions(+), 11 deletions(-) diff --git a/lmcache/v1/distributed/l2_adapters/serde_wrapper.py b/lmcache/v1/distributed/l2_adapters/serde_wrapper.py index 5ccd2348afe..208bb30d2b4 100644 --- a/lmcache/v1/distributed/l2_adapters/serde_wrapper.py +++ b/lmcache/v1/distributed/l2_adapters/serde_wrapper.py @@ -41,7 +41,7 @@ from lmcache.v1.distributed.api import KeyListPage, MemoryLayoutDesc, ObjectKey from lmcache.v1.distributed.error import L1Error from lmcache.v1.distributed.internal_api import L2AdapterListener, L2StoreResult -from lmcache.v1.distributed.l1_manager import L1Manager +from lmcache.v1.distributed.l1_protocol import L1ManagerInterface from lmcache.v1.distributed.l2_adapters.base import ( AdapterUsage, L2AdapterInterface, @@ -106,7 +106,7 @@ def __init__( self, inner: L2AdapterInterface, serde: SerdeProcessor, - l1_manager: L1Manager, + l1_manager: L1ManagerInterface, ) -> None: super().__init__() self._inner = inner @@ -604,18 +604,20 @@ def _alloc_temp_buffers( layout_desc=layout, mode="new", ) - # First pass: collect every key whose reserve_write succeeded. - # We must scan the full list (not bail on the first failure) - # so a mixed-success result still releases all reserved keys. + # Scan the full list (not bailing on the first failure) so a mixed + # result still releases every reserved key. A SUCCESS result always + # carries a buffer; the ``is not None`` check keeps that explicit and + # collects the objects in one pass. successful_temp_keys: list[ObjectKey] = [] + temp_objs: list[MemoryObj] = [] for temp_key in temp_keys: r = results.get(temp_key) - if r is not None and r[0] == L1Error.SUCCESS: + if r is not None and r[0] == L1Error.SUCCESS and r[1] is not None: successful_temp_keys.append(temp_key) + temp_objs.append(r[1]) if len(successful_temp_keys) != len(temp_keys): self._release_write_temps(successful_temp_keys) return temp_keys, None - temp_objs = [results[tk][1] for tk in temp_keys] return temp_keys, temp_objs def _release_write_temps(self, temp_keys: list[ObjectKey]) -> None: diff --git a/lmcache/v1/distributed/storage_manager.py b/lmcache/v1/distributed/storage_manager.py index 2154f2ea089..93f5c20a8df 100644 --- a/lmcache/v1/distributed/storage_manager.py +++ b/lmcache/v1/distributed/storage_manager.py @@ -25,6 +25,7 @@ from lmcache.v1.distributed.error import L1Error, strerror from lmcache.v1.distributed.internal_api import L1MemoryDesc, L2AdapterListener from lmcache.v1.distributed.l1_manager import L1Manager +from lmcache.v1.distributed.l1_protocol import L1ManagerInterface from lmcache.v1.distributed.l2_adapters import create_l2_adapter from lmcache.v1.distributed.l2_adapters.base import L2AdapterInterface from lmcache.v1.distributed.l2_adapters.config import L2AdapterConfigBase @@ -33,6 +34,7 @@ L2ReconfigureError, ) from lmcache.v1.distributed.l2_adapters.serde_wrapper import SerdeL2AdapterWrapper +from lmcache.v1.distributed.maru_l1_manager import MaruL1Manager from lmcache.v1.distributed.quota_manager import QuotaManager from lmcache.v1.distributed.serde import create_serde_processor from lmcache.v1.distributed.storage_controllers import ( @@ -49,7 +51,7 @@ AdapterDescriptor, create_store_policy, ) -from lmcache.v1.memory_management import MemoryObj +from lmcache.v1.memory_management import MemoryFormat, MemoryObj from lmcache.v1.mp_observability.errors import LMCacheTimeoutError from lmcache.v1.mp_observability.event import Event, EventType from lmcache.v1.mp_observability.event_bus import get_event_bus @@ -66,7 +68,15 @@ class StorageManager: def __init__(self, config: StorageManagerConfig): - self._l1_manager = L1Manager(config.l1_manager_config) + # Select the L1 backend: the maru shared CXL tier when configured, + # otherwise the stock local-DRAM manager. Both satisfy + # L1ManagerInterface, so the controllers below drive either unchanged. + l1_config = config.l1_manager_config + self._l1_manager: L1ManagerInterface = ( + MaruL1Manager(l1_config) + if l1_config.memory_config.maru_config is not None + else L1Manager(l1_config) + ) self._event_bus = get_event_bus() # L1 eviction controller @@ -80,7 +90,9 @@ def __init__(self, config: StorageManagerConfig): # a ``serde_config``, the adapter is wrapped with # ``SerdeL2AdapterWrapper`` so controllers see a plain L2 adapter # and serde is transparent. - self._l1_memory_desc = self._l1_manager.get_l1_memory_desc() + self._l1_memory_desc: L1MemoryDesc | None = ( + self._l1_manager.get_l1_memory_desc() + ) self._next_adapter_id = 0 # Serializes add_l2_adapter / delete_l2_adapter against each other. self._lifecycle_lock = threading.Lock() @@ -167,6 +179,40 @@ def __init__(self, config: StorageManagerConfig): ) # External APIs for serving engine integration code to call + def register_kv_layout( + self, + layout_desc: MemoryLayoutDesc, + fmt: MemoryFormat, + chunk_size_in_tokens: int, + num_object_groups: int, + ) -> None: + """Bind the KV layout to the L1 backend (maru CXL pool bring-up). + + Stock L1 sizes its pool from config at construction and ignores this; + the maru L1 tier defers pool creation until the layout is known, so the + engine calls this once the KV cache is registered. A no-op for stock. + + Args: + layout_desc: Per-group chunk shapes and dtypes. + fmt: The KV memory format. + chunk_size_in_tokens: Tokens per chunk. + num_object_groups: Number of KV object groups in the layout. + + Raises: + ValueError: The maru L1 tier supports a single object group only + (multi-model support is a maru TODO). + """ + if not isinstance(self._l1_manager, MaruL1Manager): + return + if num_object_groups > 1: + raise ValueError( + "maru L1 supports a single object group only, got " + f"{num_object_groups} (multi-model support is a maru TODO)" + ) + self._l1_manager.register_kv_layout( + layout_desc.shapes, layout_desc.dtypes, fmt, chunk_size_in_tokens + ) + @enable_tracing() def reserve_write( self, @@ -757,7 +803,19 @@ def quota_manager(self) -> QuotaManager: @property def l1_memory_desc(self) -> L1MemoryDesc: - """Descriptor of the L1 memory buffer backing this storage manager.""" + """Descriptor of the L1 memory buffer backing this storage manager. + + Raises: + RuntimeError: The maru L1 tier exposes no single registerable + region (``get_l1_memory_desc`` is None). The only consumer is + p2p, which is rejected at startup for maru, so this never + fires in practice. + """ + if self._l1_memory_desc is None: + raise RuntimeError( + "L1 memory descriptor is unavailable (the maru L1 tier has no " + "single registerable region)" + ) return self._l1_memory_desc def get_l2_usages( diff --git a/tests/v1/distributed/test_distributed_storage_manager.py b/tests/v1/distributed/test_distributed_storage_manager.py index 6617710992c..e921dbe0d84 100644 --- a/tests/v1/distributed/test_distributed_storage_manager.py +++ b/tests/v1/distributed/test_distributed_storage_manager.py @@ -4,6 +4,7 @@ """ # Standard +from unittest.mock import MagicMock import time # Third Party @@ -21,17 +22,21 @@ EvictionConfig, L1ManagerConfig, L1MemoryManagerConfig, + MaruL1Config, StorageManagerConfig, ) from lmcache.v1.distributed.l2_adapters.config import ( L2AdaptersConfig, ) from lmcache.v1.distributed.l2_adapters.mock_l2_adapter import MockL2AdapterConfig +from lmcache.v1.memory_management import MemoryFormat from lmcache.v1.mp_observability.event import Event, EventType from lmcache.v1.mp_observability.event_bus import EventBusConfig, init_event_bus try: # First Party + from lmcache.v1.distributed.l1_manager import L1Manager + from lmcache.v1.distributed.maru_l1_manager import MaruL1Manager from lmcache.v1.distributed.storage_manager import StorageManager except ImportError: # Skip tests if L1Manager cannot be imported @@ -119,6 +124,27 @@ def small_storage_manager_config(small_l1_config): ) +@pytest.fixture +def maru_storage_manager_config(): + """StorageManagerConfig whose L1 tier is maru (shared CXL pool).""" + return StorageManagerConfig( + l1_manager_config=L1ManagerConfig( + memory_config=L1MemoryManagerConfig( + size_in_bytes=0, + use_lazy=False, + maru_config=MaruL1Config( + server_url="maru://localhost:5555", + pool_size_bytes=1 << 20, + instance_id="test-sm", + ), + ), + write_ttl_seconds=600, + read_ttl_seconds=300, + ), + eviction_config=EvictionConfig(eviction_policy="LRU"), + ) + + @pytest.fixture def basic_layout(): """Create a basic MemoryLayoutDesc for testing.""" @@ -902,3 +928,62 @@ def test_sparse_from_l2_loads_all_found( sm.finish_read_prefetched(existing) sm.close() + + +class TestMaruL1Selection: + """C8: StorageManager selects and drives the maru L1 backend.""" + + def test_maru_config_selects_maru_l1_manager(self, maru_storage_manager_config): + sm = StorageManager(maru_storage_manager_config) + try: + assert isinstance(sm._l1_manager, MaruL1Manager) + finally: + sm.close() + + def test_stock_config_selects_stock_l1_manager(self, basic_storage_manager_config): + sm = StorageManager(basic_storage_manager_config) + try: + assert isinstance(sm._l1_manager, L1Manager) + finally: + sm.close() + + def test_register_kv_layout_forwards_to_maru( + self, maru_storage_manager_config, basic_layout + ): + sm = StorageManager(maru_storage_manager_config) + try: + # Spy the forward so the test needs no maru runtime (init_layout). + sm._l1_manager.register_kv_layout = MagicMock() # type: ignore[method-assign] + sm.register_kv_layout( + basic_layout, MemoryFormat.KV_MLA_FMT, 16, num_object_groups=1 + ) + sm._l1_manager.register_kv_layout.assert_called_once_with( + basic_layout.shapes, basic_layout.dtypes, MemoryFormat.KV_MLA_FMT, 16 + ) + finally: + sm.close() + + def test_register_kv_layout_rejects_multi_object_group( + self, maru_storage_manager_config, basic_layout + ): + sm = StorageManager(maru_storage_manager_config) + try: + with pytest.raises(ValueError): + sm.register_kv_layout( + basic_layout, MemoryFormat.KV_MLA_FMT, 16, num_object_groups=2 + ) + finally: + sm.close() + + def test_register_kv_layout_noop_for_stock( + self, basic_storage_manager_config, basic_layout + ): + sm = StorageManager(basic_storage_manager_config) + try: + # Stock sizes its pool at construction: this must be a silent no-op + # (no forward, no raise) even for a would-be-rejected group count. + sm.register_kv_layout( + basic_layout, MemoryFormat.KV_MLA_FMT, 16, num_object_groups=2 + ) + finally: + sm.close() From 8b14ded42b64ce14755f98614a7f734f585e0aab Mon Sep 17 00:00:00 2001 From: seohui-XCENA Date: Mon, 6 Jul 2026 12:52:07 +0900 Subject: [PATCH 09/28] [MP][Maru] Reject unsupported maru combos (gds/devdax/L2-RDMA/skip_l1/p2p/engine) - validate_storage_manager_config: a maru L1 config rejects the other L1 backends (gds-l1-path, l1-devdax-path), store_policy=skip_l1, and L2 adapters that require a single registerable region (nixl / mooncake-rdma); copy-type L2 and the default store policy pass - l1_exposes_single_memory_region returns False for maru (the shared CXL pool has no single registerable region), so the existing p2p startup guard fires; its message now lists maru too - run_http_server rejects maru paired with a non-lmcache_driven transfer mode (engine-driven/auto assume engine-side buffers the shared pool lacks) - tests: maru + gds/devdax/skip_l1/registered-L2 raise; maru + copy-L2 and the default policy pass; l1_exposes False for maru Signed-off-by: seohui-XCENA --- lmcache/v1/distributed/config.py | 32 ++++++- lmcache/v1/multiprocess/http_server.py | 15 ++- tests/v1/distributed/test_l1_config_maru.py | 91 ++++++++++++++++++- tests/v1/distributed/test_l1_single_region.py | 16 ++++ 4 files changed, 149 insertions(+), 5 deletions(-) diff --git a/lmcache/v1/distributed/config.py b/lmcache/v1/distributed/config.py index 1f5da780531..35033f78053 100644 --- a/lmcache/v1/distributed/config.py +++ b/lmcache/v1/distributed/config.py @@ -324,6 +324,33 @@ def validate_storage_manager_config(config: StorageManagerConfig) -> None: ValueError: If mutually exclusive L1 tiers are both configured, or hybrid L1 is paired with incompatible L2 adapters. """ + l1_config = config.l1_manager_config + if l1_config.memory_config.maru_config is not None: + # maru is a standalone shared-CXL L1 tier. It cannot coexist with the + # other L1 backends, and -- exposing no single registerable region -- + # cannot serve L2 adapters or store policies that need one. + if l1_config.gds_l1_config is not None: + raise ValueError("maru L1 cannot be combined with gds-l1-path") + if l1_config.memory_config.devdax_path: + raise ValueError("maru L1 cannot be combined with l1-devdax-path") + if config.store_policy == "skip_l1": + raise ValueError( + "maru L1 does not support store_policy='skip_l1': the shared " + "pool is the store target, not a bypass buffer" + ) + registered_adapters = [ + name + for adapter_config in config.l2_adapter_config.adapters + if (name := _requires_single_l1_memory_region(adapter_config)) is not None + ] + if registered_adapters: + raise ValueError( + "maru L1 has no single registerable memory region, so it " + "cannot be used with L2 adapters that require one: " + f"{', '.join(registered_adapters)}" + ) + return + if ( config.l1_manager_config.gds_l1_config is not None and config.l1_manager_config.memory_config.devdax_path @@ -356,9 +383,12 @@ def l1_exposes_single_memory_region(config: StorageManagerConfig) -> bool: Returns: ``True`` if L1 is a single registerable memory region, ``False`` for - GDS L1 or Device-DAX L1. + GDS L1, Device-DAX L1, or maru L1 (a shared CXL pool with no single + registerable region). """ l1_config = config.l1_manager_config + if l1_config.memory_config.maru_config is not None: + return False if l1_config.gds_l1_config is not None: return False if l1_config.memory_config.devdax_path: diff --git a/lmcache/v1/multiprocess/http_server.py b/lmcache/v1/multiprocess/http_server.py index 43a382d5630..a274cd8163d 100644 --- a/lmcache/v1/multiprocess/http_server.py +++ b/lmcache/v1/multiprocess/http_server.py @@ -236,9 +236,20 @@ def run_http_server( if not l1_exposes_single_memory_region(storage_manager_config): raise ValueError( "P2P requires a single L1 memory region the transfer channel " - "can register; it is incompatible with GDS L1 (--gds-l1-path) " - "and Device-DAX L1 (--l1-devdax-path)." + "can register; it is incompatible with GDS L1 (--gds-l1-path), " + "Device-DAX L1 (--l1-devdax-path), and maru L1 (--maru-*)." ) + # maru drives KV transfer through the LMCache-driven path only; engine- + # driven / auto transfer assumes engine-side buffers the shared CXL pool + # does not provide. + if ( + storage_manager_config.l1_manager_config.memory_config.maru_config is not None + and mp_config.supported_transfer_mode != "lmcache_driven" + ): + raise ValueError( + "maru L1 requires --supported-transfer-mode lmcache_driven (got " + f"{mp_config.supported_transfer_mode!r})." + ) _configs["mp"] = mp_config _configs["storage_manager"] = storage_manager_config _configs["observability"] = obs_config diff --git a/tests/v1/distributed/test_l1_config_maru.py b/tests/v1/distributed/test_l1_config_maru.py index da675a5174e..1bce558ec8a 100644 --- a/tests/v1/distributed/test_l1_config_maru.py +++ b/tests/v1/distributed/test_l1_config_maru.py @@ -1,12 +1,22 @@ # SPDX-License-Identifier: Apache-2.0 -"""Tests for Maru L1 config parsing.""" +"""Tests for Maru L1 config parsing and startup guards.""" # Third Party import pytest # First Party -from lmcache.v1.distributed.config import MaruL1Config, parse_args +from lmcache.v1.distributed.config import ( + EvictionConfig, + GdsL1Config, + L1ManagerConfig, + L1MemoryManagerConfig, + MaruL1Config, + StorageManagerConfig, + parse_args, +) +from lmcache.v1.distributed.l2_adapters.config import L2AdaptersConfig +from lmcache.v1.distributed.l2_adapters.mock_l2_adapter import MockL2AdapterConfig def _args(*extra: str, l1_size_gb: str = "1") -> list[str]: @@ -47,3 +57,80 @@ def test_maru_flags_build_maru_config(): def test_maru_without_pool_size_raises(): with pytest.raises(ValueError): parse_args(_args("--maru-server-url", "maru://localhost:9000")) + + +# --------------------------------------------------------------------------- +# C9: validate_storage_manager_config rejects unsupported maru combinations. +# --------------------------------------------------------------------------- + + +def _maru_memory(**overrides) -> L1MemoryManagerConfig: + return L1MemoryManagerConfig( + size_in_bytes=0, + use_lazy=False, + maru_config=MaruL1Config( + server_url="maru://localhost:5555", + pool_size_bytes=1 << 20, + instance_id="t", + ), + **overrides, + ) + + +def _maru_sm_config( + *, memory=None, gds=None, store_policy="default", adapters=None +) -> StorageManagerConfig: + return StorageManagerConfig( + l1_manager_config=L1ManagerConfig( + memory_config=memory if memory is not None else _maru_memory(), + gds_l1_config=gds, + write_ttl_seconds=600, + read_ttl_seconds=300, + ), + eviction_config=EvictionConfig(eviction_policy="LRU"), + l2_adapter_config=L2AdaptersConfig(adapters=adapters or []), + store_policy=store_policy, + ) + + +def test_maru_plus_copy_l2_is_accepted(): + # A copy-type L2 (mock) needs no registerable region -> allowed. + _maru_sm_config( + adapters=[MockL2AdapterConfig(max_size_gb=1.0, mock_bandwidth_gb=1.0)] + ) + + +def test_maru_rejects_gds_l1(): + with pytest.raises(ValueError, match="gds"): + _maru_sm_config( + gds=GdsL1Config(file_location="/tmp/gds", size_in_bytes=1 << 20) + ) + + +def test_maru_rejects_devdax_l1(): + with pytest.raises(ValueError, match="devdax"): + _maru_sm_config(memory=_maru_memory(devdax_path="/dev/dax0.0", shm_name="")) + + +def test_maru_rejects_skip_l1_store_policy(): + with pytest.raises(ValueError, match="skip_l1"): + _maru_sm_config(store_policy="skip_l1") + + +def test_maru_rejects_registered_l2(monkeypatch): + # Simulate a registered/RDMA adapter via the shared region classifier. + monkeypatch.setattr( + "lmcache.v1.distributed.config._requires_single_l1_memory_region", + lambda adapter_config: "nixl_store", + ) + with pytest.raises(ValueError, match="registerable"): + _maru_sm_config( + adapters=[MockL2AdapterConfig(max_size_gb=1.0, mock_bandwidth_gb=1.0)] + ) + + +# NOTE: the maru transfer-mode guard (maru requires supported_transfer_mode +# == "lmcache_driven") is an inline check in run_http_server, alongside the +# existing p2p startup guards there. Those startup guards are not unit-tested +# (importing run_http_server pulls the full server chain); this one follows +# the same precedent. diff --git a/tests/v1/distributed/test_l1_single_region.py b/tests/v1/distributed/test_l1_single_region.py index 0c7f4a1624b..0a3711ddf86 100644 --- a/tests/v1/distributed/test_l1_single_region.py +++ b/tests/v1/distributed/test_l1_single_region.py @@ -7,6 +7,7 @@ GdsL1Config, L1ManagerConfig, L1MemoryManagerConfig, + MaruL1Config, StorageManagerConfig, l1_exposes_single_memory_region, ) @@ -54,3 +55,18 @@ def test_devdax_l1_is_not_single_region(): ) ) assert l1_exposes_single_memory_region(config) is False + + +def test_maru_l1_is_not_single_region(): + config = _config( + L1MemoryManagerConfig( + size_in_bytes=0, + use_lazy=False, + maru_config=MaruL1Config( + server_url="maru://localhost:5555", + pool_size_bytes=1 << 20, + instance_id="t", + ), + ) + ) + assert l1_exposes_single_memory_region(config) is False From ed25092557e1f2d30e6cc1f7edfb4868cb8af336 Mon Sep 17 00:00:00 2001 From: seohui-XCENA Date: Mon, 6 Jul 2026 13:31:12 +0900 Subject: [PATCH 10/28] [MP][Maru] Wire register_kv_layout into register_kv_cache for maru pool bring-up - after the KV layout is resolved in register_kv_cache, call StorageManager.register_kv_layout with the layout desc, the storage MemoryFormat (KV_MLA_FMT / KV_2LTD via is_mla), chunk size, and object-group count; a no-op for the stock backend, idempotent for maru across instances - on failure (e.g. maru rejecting >1 object group) close the just-built cache context and re-raise so the instance is not left half-registered - runtime test deferred: importing this module needs a c_ops rebuild (built .so lacks execute_object_group_transfer); lands with the C9 G6 test post-rebuild Signed-off-by: seohui-XCENA --- .../modules/lmcache_driven_transfer.py | 26 +++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/lmcache/v1/multiprocess/modules/lmcache_driven_transfer.py b/lmcache/v1/multiprocess/modules/lmcache_driven_transfer.py index a3ccc8b0729..557a351810b 100644 --- a/lmcache/v1/multiprocess/modules/lmcache_driven_transfer.py +++ b/lmcache/v1/multiprocess/modules/lmcache_driven_transfer.py @@ -28,9 +28,9 @@ lmcache_memcpy_async_d2h, lmcache_memcpy_async_h2d, ) -from lmcache.v1.gpu_connector.utils import LayoutHints +from lmcache.v1.gpu_connector.utils import LayoutHints, is_mla from lmcache.v1.lazy_memory_allocator import LazyMemoryAllocator -from lmcache.v1.memory_management import GDSMemoryObject, MemoryObj +from lmcache.v1.memory_management import GDSMemoryObject, MemoryFormat, MemoryObj from lmcache.v1.mp_observability.event import Event, EventType from lmcache.v1.multiprocess.custom_types import ( IPCCacheServerKey, @@ -871,6 +871,28 @@ def register_kv_cache( model_name, world_size, layout_desc, attn_desc ) + # Bring up the maru CXL pool now that the KV layout is known. A no-op + # for the stock L1 backend (StorageManager self-gates on the maru + # config); maru's register_kv_layout is idempotent across instances. + fmt = ( + MemoryFormat.KV_MLA_FMT + if is_mla(cache_context.get_engine_kv_format(0)) + else MemoryFormat.KV_2LTD + ) + try: + self._ctx.storage_manager.register_kv_layout( + layout_desc, + fmt, + self._ctx.chunk_size, + attn_desc.num_object_groups, + ) + except Exception: + # e.g. maru rejects >1 object group -- drop the context we just + # built so the instance is not left half-registered, then surface + # the error on the REGISTER path. + cache_context.close() + raise + with self._lock: self._cache_contexts[instance_id] = ContextEntry( cache_context=cache_context, From 5260f66dd9b31b30498311cb380254ab9fdce42d Mon Sep 17 00:00:00 2001 From: seohui-XCENA Date: Mon, 6 Jul 2026 13:40:17 +0900 Subject: [PATCH 11/28] [MP][Maru] Add deferred startup-hook tests (transfer-mode guard + register_kv_layout) - G6 (C9): run_http_server rejects maru paired with a non-lmcache_driven transfer mode - C10: register_kv_cache forwards to StorageManager.register_kv_layout with the right MemoryFormat (KV_MLA_FMT / KV_2LTD via is_mla); on a rejected layout it closes the cache context and leaves the instance unregistered; stock registers normally - both were blocked by a stale c_ops build (missing execute_object_group_transfer); unblocked after the rebuild Signed-off-by: seohui-XCENA --- tests/v1/distributed/test_l1_config_maru.py | 24 ++++- .../test_maru_register_kv_layout.py | 97 +++++++++++++++++++ 2 files changed, 116 insertions(+), 5 deletions(-) create mode 100644 tests/v1/multiprocess/test_maru_register_kv_layout.py diff --git a/tests/v1/distributed/test_l1_config_maru.py b/tests/v1/distributed/test_l1_config_maru.py index 1bce558ec8a..8e74c835352 100644 --- a/tests/v1/distributed/test_l1_config_maru.py +++ b/tests/v1/distributed/test_l1_config_maru.py @@ -129,8 +129,22 @@ def test_maru_rejects_registered_l2(monkeypatch): ) -# NOTE: the maru transfer-mode guard (maru requires supported_transfer_mode -# == "lmcache_driven") is an inline check in run_http_server, alongside the -# existing p2p startup guards there. Those startup guards are not unit-tested -# (importing run_http_server pulls the full server chain); this one follows -# the same precedent. +def test_maru_rejects_non_lmcache_driven_transfer(): + """maru requires the LMCache-driven transfer path (rejects engine/auto).""" + # First Party + from lmcache.v1.mp_observability.config import ObservabilityConfig + from lmcache.v1.multiprocess.config import ( + CoordinatorConfig, + HTTPFrontendConfig, + MPServerConfig, + ) + from lmcache.v1.multiprocess.http_server import run_http_server + + with pytest.raises(ValueError, match="lmcache_driven"): + run_http_server( + http_config=HTTPFrontendConfig(), + mp_config=MPServerConfig(supported_transfer_mode="engine_driven"), + storage_manager_config=_maru_sm_config(), + obs_config=ObservabilityConfig(), + coordinator_config=CoordinatorConfig(url=""), + ) diff --git a/tests/v1/multiprocess/test_maru_register_kv_layout.py b/tests/v1/multiprocess/test_maru_register_kv_layout.py new file mode 100644 index 00000000000..3e23eeda116 --- /dev/null +++ b/tests/v1/multiprocess/test_maru_register_kv_layout.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the maru register_kv_layout engine hook in register_kv_cache. + +The hook brings up the maru CXL pool once the KV layout is known. These tests +exercise it with a bare transfer module (bypassing the CUDA dispatcher in +__init__) and mocked context creation, mirroring test_worker_liveness.py. +""" + +# Standard +from unittest.mock import MagicMock +import threading + +# Third Party +import pytest + +# First Party +from lmcache.v1.memory_management import MemoryFormat +from lmcache.v1.multiprocess.modules import lmcache_driven_transfer as gpu_mod +from lmcache.v1.multiprocess.modules.lmcache_driven_transfer import ( + LMCacheDrivenTransferModule, +) + + +def _bare_module() -> LMCacheDrivenTransferModule: + """A transfer module with only the state register_kv_cache touches.""" + module = LMCacheDrivenTransferModule.__new__(LMCacheDrivenTransferModule) + module._ctx = MagicMock(name="ctx") + module._cache_contexts = {} + module._lock = threading.Lock() + return module + + +def _register(module: LMCacheDrivenTransferModule) -> None: + module.register_kv_cache(1, MagicMock(), "model", 1, MagicMock(), MagicMock(), []) + + +def test_hook_forwards_kv_2ltd_for_non_mla(monkeypatch): + monkeypatch.setattr( + gpu_mod, "create_cache_context", lambda *a, **kw: MagicMock(num_layers=2) + ) + monkeypatch.setattr(gpu_mod, "get_layout_desc", lambda *a, **kw: MagicMock()) + monkeypatch.setattr(gpu_mod, "is_mla", lambda fmt: False) + module = _bare_module() + + _register(module) + + call = module._ctx.storage_manager.register_kv_layout.call_args + # (layout_desc, fmt, chunk_size, num_object_groups) + assert call.args[1] == MemoryFormat.KV_2LTD + + +def test_hook_forwards_kv_mla_fmt_for_mla(monkeypatch): + monkeypatch.setattr( + gpu_mod, "create_cache_context", lambda *a, **kw: MagicMock(num_layers=2) + ) + monkeypatch.setattr(gpu_mod, "get_layout_desc", lambda *a, **kw: MagicMock()) + monkeypatch.setattr(gpu_mod, "is_mla", lambda fmt: True) + module = _bare_module() + + _register(module) + + call = module._ctx.storage_manager.register_kv_layout.call_args + assert call.args[1] == MemoryFormat.KV_MLA_FMT + + +def test_hook_failure_closes_context_and_does_not_register(monkeypatch): + """A rejected layout (e.g. maru >1 object group) must not half-register.""" + cache_ctx = MagicMock(num_layers=2, name="cache_ctx") + monkeypatch.setattr(gpu_mod, "create_cache_context", lambda *a, **kw: cache_ctx) + monkeypatch.setattr(gpu_mod, "get_layout_desc", lambda *a, **kw: MagicMock()) + monkeypatch.setattr(gpu_mod, "is_mla", lambda fmt: False) + module = _bare_module() + module._ctx.storage_manager.register_kv_layout.side_effect = ValueError( + "maru L1 supports a single object group only" + ) + + with pytest.raises(ValueError, match="single object group"): + _register(module) + + cache_ctx.close.assert_called_once() + assert 1 not in module._cache_contexts # not left half-registered + + +def test_hook_is_noop_style_for_stock_backend(monkeypatch): + """For stock, register_kv_layout is a no-op; the instance still registers.""" + monkeypatch.setattr( + gpu_mod, "create_cache_context", lambda *a, **kw: MagicMock(num_layers=2) + ) + monkeypatch.setattr(gpu_mod, "get_layout_desc", lambda *a, **kw: MagicMock()) + monkeypatch.setattr(gpu_mod, "is_mla", lambda fmt: False) + module = _bare_module() + # Stock StorageManager.register_kv_layout returns None (no-op). + module._ctx.storage_manager.register_kv_layout.return_value = None + + _register(module) + + assert 1 in module._cache_contexts # registration completed normally From 34d452421f241893e277b2565ebf437be15ec644 Mon Sep 17 00:00:00 2001 From: seohui-XCENA Date: Mon, 6 Jul 2026 14:08:55 +0900 Subject: [PATCH 12/28] [MP][Maru] Publish L1 events to mp_observability (event_bus parity) - MaruL1Manager now publishes the same L1 events to the shared event bus as stock L1Manager, alongside the listener notifications: L1_READ_RESERVED, L1_READ_FINISHED (+ L1_KEYS_EVICTED for freed temporaries), L1_WRITE_RESERVED, L1_WRITE_FINISHED, L1_WRITE_FINISHED_AND_READ_RESERVED, and L1_KEYS_EVICTED for delete/clear; touch_keys stays listener-only (matches stock) - add get_event_bus() in __init__ and a _publish helper mirroring stock - closes the observability gap deferred in C4 -- maru L1 ops now surface in the mp_observability dashboards - tests: write/read/delete lifecycle events, promote publishes the promote event (never L1_WRITE_FINISHED), temporary finish_read publishes evicted, touch_keys publishes nothing Signed-off-by: seohui-XCENA --- lmcache/v1/distributed/maru_l1_manager.py | 22 +++++++ tests/v1/distributed/test_maru_l1_manager.py | 68 ++++++++++++++++++++ 2 files changed, 90 insertions(+) diff --git a/lmcache/v1/distributed/maru_l1_manager.py b/lmcache/v1/distributed/maru_l1_manager.py index 2c477cb88d0..1809d3cad6a 100644 --- a/lmcache/v1/distributed/maru_l1_manager.py +++ b/lmcache/v1/distributed/maru_l1_manager.py @@ -59,6 +59,8 @@ MaruMemoryAllocator, ) from lmcache.v1.memory_management import MemoryFormat, MemoryObj +from lmcache.v1.mp_observability.event import Event, EventType +from lmcache.v1.mp_observability.event_bus import get_event_bus if TYPE_CHECKING: # Third Party @@ -171,6 +173,9 @@ def __init__(self, config: L1ManagerConfig) -> None: self._pending_read: dict[ObjectKey, _PendingRead] = {} self._pending_write: dict[ObjectKey, _PendingWrite] = {} self._registered_listeners: list[L1ManagerListener] = [] + # PARITY(L1Manager): observability events go to the shared event bus + # alongside the listener notifications. + self._event_bus = get_event_bus() # MARU: W1+R1 orphan sweeper -- a crashed/abandoned client leaves write # pages or read pins behind; reclaim them once their TTL elapses. self._sweep_interval = max( @@ -192,6 +197,15 @@ def register_listener(self, listener: L1ManagerListener) -> None: with self._lock: self._registered_listeners.append(listener) + def _publish(self, event_type: EventType, keys: list[ObjectKey]) -> None: + """Publish an L1 observability event (PARITY with stock L1Manager). + + Args: + event_type: The L1 event type to publish. + keys: The affected object keys (event metadata). + """ + self._event_bus.publish(Event(event_type=event_type, metadata={"keys": keys})) + def _safe_unpin(self, handler: "MaruHandler", key_strs: list[str]) -> None: """Release server pins, logging (not raising) on RPC failure.""" if not key_strs: @@ -479,6 +493,7 @@ def reserve_read( # holds (feeds the eviction LRU and the store controller). for listener in self._registered_listeners: listener.on_l1_keys_reserved_read(successful_keys) + self._publish(EventType.L1_READ_RESERVED, successful_keys) return ret @_maru_l1_synchronized @@ -575,6 +590,8 @@ def finish_read( for listener in self._registered_listeners: listener.on_l1_keys_read_finished(successful_keys) listener.on_l1_keys_deleted_by_manager(need_to_free_keys) + self._publish(EventType.L1_READ_FINISHED, successful_keys) + self._publish(EventType.L1_KEYS_EVICTED, need_to_free_keys) return ret @_maru_l1_synchronized @@ -667,6 +684,7 @@ def reserve_write( # holds (the eviction LRU treats them as unevictable). for listener in self._registered_listeners: listener.on_l1_keys_reserved_write(successful_keys) + self._publish(EventType.L1_WRITE_RESERVED, successful_keys) return ret @_maru_l1_synchronized @@ -701,6 +719,7 @@ def finish_write(self, keys: list[ObjectKey]) -> dict[ObjectKey, L1Error]: # event -- that is on_l1_keys_finish_write_and_reserve_read). for listener in self._registered_listeners: listener.on_l1_keys_write_finished(registered) + self._publish(EventType.L1_WRITE_FINISHED, registered) return ret @_maru_l1_synchronized @@ -778,6 +797,7 @@ def finish_write_and_reserve_read( for listener in self._registered_listeners: listener.on_l1_keys_finish_write_and_reserve_read(successful_keys) + self._publish(EventType.L1_WRITE_FINISHED_AND_READ_RESERVED, successful_keys) return ret def _promote_retained( @@ -860,6 +880,7 @@ def delete(self, keys: list[ObjectKey]) -> dict[ObjectKey, L1Error]: # removed from the directory (the eviction LRU drops them). for listener in self._registered_listeners: listener.on_l1_keys_deleted_by_manager(successful_keys) + self._publish(EventType.L1_KEYS_EVICTED, successful_keys) return ret def touch_keys(self, keys: list[ObjectKey]) -> None: @@ -897,6 +918,7 @@ def clear(self, force: bool = False) -> None: # PARITY(L1Manager.clear): a force-clear notifies listeners of the drops. for listener in self._registered_listeners: listener.on_l1_keys_deleted_by_manager(dropped) + self._publish(EventType.L1_KEYS_EVICTED, dropped) def _drain_staging(self, force: bool) -> list[ObjectKey]: """Unpin staged reads (refcount times) and abort staged writes. diff --git a/tests/v1/distributed/test_maru_l1_manager.py b/tests/v1/distributed/test_maru_l1_manager.py index 7236d7542c1..72f9a846332 100644 --- a/tests/v1/distributed/test_maru_l1_manager.py +++ b/tests/v1/distributed/test_maru_l1_manager.py @@ -3,6 +3,7 @@ """Tests for MaruL1Manager (fake maru runtime, no CXL required).""" # Standard +from unittest.mock import MagicMock import time # Third Party @@ -19,6 +20,7 @@ _PendingRead, object_key_to_string, ) + from lmcache.v1.mp_observability.event import EventType # Local from .maru_fakes import RecordingListener, make_maru_manager @@ -601,6 +603,72 @@ def test_finish_read_after_sweep_is_not_exist(): assert manager.finish_read([k])[k] == L1Error.KEY_NOT_EXIST +# ========================================================================= +# event_bus observability parity (C11a): mirrors stock L1Manager publishes +# ========================================================================= + + +def _keys_for(bus, event_type): + """Return the ``keys`` metadata of every publish of ``event_type``.""" + return [ + c.args[0].metadata["keys"] + for c in bus.publish.call_args_list + if c.args[0].event_type == event_type + ] + + +def test_event_bus_write_read_delete_lifecycle(): + manager, handler, _ = make_maru_manager() + manager._event_bus = MagicMock() + k = _key(1) + + manager.reserve_write([k], [False], _LAYOUT, mode="new") + manager.finish_write([k]) + manager.reserve_read([k]) + manager.finish_read([k]) + manager.delete([k]) + + assert _keys_for(manager._event_bus, EventType.L1_WRITE_RESERVED) == [[k]] + assert _keys_for(manager._event_bus, EventType.L1_WRITE_FINISHED) == [[k]] + assert _keys_for(manager._event_bus, EventType.L1_READ_RESERVED) == [[k]] + assert _keys_for(manager._event_bus, EventType.L1_READ_FINISHED) == [[k]] + assert [k] in _keys_for(manager._event_bus, EventType.L1_KEYS_EVICTED) # delete + + +def test_event_bus_promote_publishes_promote_event_not_write_finished(): + manager, _, _ = make_maru_manager() + manager._event_bus = MagicMock() + k = _key(1) + manager.reserve_write([k], [False], _LAYOUT, mode="new") + + manager.finish_write_and_reserve_read([k]) + assert _keys_for( + manager._event_bus, EventType.L1_WRITE_FINISHED_AND_READ_RESERVED + ) == [[k]] + # anti re-store at the event-bus level too. + assert _keys_for(manager._event_bus, EventType.L1_WRITE_FINISHED) == [] + + +def test_event_bus_temporary_finish_read_publishes_evicted(): + manager, _, _ = make_maru_manager() + manager._event_bus = MagicMock() + k = _key(1) + manager.reserve_write([k], [True], _LAYOUT, mode="new") + manager.finish_write_and_reserve_read([k]) # temporary read staged + + manager.finish_read([k]) + assert _keys_for(manager._event_bus, EventType.L1_READ_FINISHED) == [[k]] + assert [k] in _keys_for(manager._event_bus, EventType.L1_KEYS_EVICTED) + + +def test_event_bus_touch_keys_does_not_publish(): + manager, _, _ = make_maru_manager() + manager._event_bus = MagicMock() + + manager.touch_keys([_key(1)]) + manager._event_bus.publish.assert_not_called() + + def test_report_status_keys(): manager, handler, _ = make_maru_manager() k = _key(1) From ca1d84fab481cc9860826b969769ceab7b28e878 Mon Sep 17 00:00:00 2001 From: seohui-XCENA Date: Mon, 6 Jul 2026 14:31:23 +0900 Subject: [PATCH 13/28] [MP][Maru] Add maru control-integration tests over the stock tiering stack - a StorageManager harness backed by the maru fakes (fake CXL pool + MaruServer directory) + a mock L2, driving the real Store/Prefetch/Eviction controllers over MaruL1Manager - covers store -> maru directory registration, prefetch of directory-resident keys as a full L1 hit (maru reserve_read), and watermark eviction driving MaruL1Manager.delete on the shared directory - L1<->L2 byte movement (write-through, promote) is stock controller logic that needs real-memory-backed pages and is covered by the stock StorageManager tests, so it is not re-asserted here; extra_count and cross-node PINNED are covered at the manager level (C3/C5) Signed-off-by: seohui-XCENA --- tests/v1/distributed/test_maru_integration.py | 136 ++++++++++++++++++ 1 file changed, 136 insertions(+) create mode 100644 tests/v1/distributed/test_maru_integration.py diff --git a/tests/v1/distributed/test_maru_integration.py b/tests/v1/distributed/test_maru_integration.py new file mode 100644 index 00000000000..1da6bd6c11a --- /dev/null +++ b/tests/v1/distributed/test_maru_integration.py @@ -0,0 +1,136 @@ +# SPDX-License-Identifier: Apache-2.0 +"""C11b: maru control integration through the stock tiering controllers. + +A real StorageManager + StoreController/PrefetchController/EvictionController +drive MaruL1Manager, whose CXL pool + MaruServer directory are the in-memory +fakes from the manager-level tests, plus a mock L2 adapter. These assert maru's +*control* integration (register / read-reserve / evict-delete under the real +controllers). L1<->L2 byte movement is stock controller logic (identical for +any L1 backend) and is covered by the stock StorageManager tests, so it is not +re-asserted here (the fakes back pages with no real memory). +""" + +# Standard +import time + +# Third Party +import pytest +import torch + +try: + # First Party + from lmcache.v1.distributed.api import MemoryLayoutDesc, ObjectKey + from lmcache.v1.distributed.config import ( + EvictionConfig, + L1ManagerConfig, + L1MemoryManagerConfig, + MaruL1Config, + StorageManagerConfig, + ) + from lmcache.v1.distributed.l2_adapters.config import L2AdaptersConfig + from lmcache.v1.distributed.l2_adapters.mock_l2_adapter import MockL2AdapterConfig + from lmcache.v1.distributed.maru_l1_manager import ( + MaruL1Manager, + object_key_to_string, + ) + from lmcache.v1.distributed.storage_manager import StorageManager + + # Local + from .maru_fakes import FakeCxlAdapter, FakeMaruHandler +except ImportError: + pytest.skip("maru integration deps unavailable", allow_module_level=True) + +_LAYOUT = MemoryLayoutDesc(shapes=[torch.Size([4, 8])], dtypes=[torch.float16]) + + +def _key(idx: int) -> ObjectKey: + return ObjectKey(chunk_hash=idx.to_bytes(4, "big"), model_name="m", kv_rank=0) + + +def _wait(predicate, timeout: float = 10.0) -> bool: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return True + time.sleep(0.05) + return False + + +def _maru_sm_with_fakes(chunk_size: int = 64, trigger_watermark: float = 0.8): + """A real StorageManager whose maru L1 tier is backed by the fakes.""" + config = StorageManagerConfig( + l1_manager_config=L1ManagerConfig( + memory_config=L1MemoryManagerConfig( + size_in_bytes=0, + use_lazy=False, + maru_config=MaruL1Config( + server_url="maru://localhost:5555", + pool_size_bytes=1 << 20, + instance_id="t", + ), + ), + write_ttl_seconds=600, + read_ttl_seconds=300, + ), + eviction_config=EvictionConfig( + eviction_policy="LRU", trigger_watermark=trigger_watermark + ), + l2_adapter_config=L2AdaptersConfig( + adapters=[MockL2AdapterConfig(max_size_gb=1.0, mock_bandwidth_gb=1.0)] + ), + ) + sm = StorageManager(config) + assert isinstance(sm._l1_manager, MaruL1Manager) # harness selected maru + handler = FakeMaruHandler(chunk_size) + adapter = FakeCxlAdapter(chunk_size) + alloc = sm._l1_manager._allocator + alloc._handler = handler + alloc._cxl_adapter = adapter + alloc._single_token_size = 16 + return sm, handler, adapter + + +def test_store_registers_in_maru_directory(): + """reserve_write -> finish_write through the full stack registers in maru.""" + sm, handler, _ = _maru_sm_with_fakes() + try: + k = _key(1) + res = sm.reserve_write([k], _LAYOUT, mode="new") + assert res[k] is not None + sm.finish_write([k]) + assert object_key_to_string(k) in handler.store_map + finally: + sm.close() + + +def test_prefetch_hits_l1_resident_keys(): + """Prefetch of directory-resident keys is a full L1 hit (maru reserve_read).""" + sm, _, _ = _maru_sm_with_fakes() + try: + keys = [_key(i) for i in range(3)] + sm.reserve_write(keys, _LAYOUT, mode="new") + sm.finish_write(keys) + + handle = sm.submit_prefetch_task(keys, _LAYOUT) + assert _wait(lambda: sm.query_prefetch_status(handle) is not None) + assert sm.query_prefetch_status(handle).count_leading_ones() == len(keys) + finally: + sm.close() + + +def test_eviction_deletes_from_maru_directory(): + """Watermark eviction drives MaruL1Manager.delete on the shared directory.""" + # Low watermark: the fake pool is 16 pages, so a handful of stored keys + # crosses it and the eviction controller must reclaim some. + sm, handler, _ = _maru_sm_with_fakes(trigger_watermark=0.1) + try: + keys = [_key(i) for i in range(6)] + sm.reserve_write(keys, _LAYOUT, mode="new") + sm.finish_write(keys) + assert len(handler.store_map) == 6 + + # The eviction controller runs on its own thread; it deletes evictable + # keys from the maru directory until usage falls under the watermark. + assert _wait(lambda: len(handler.store_map) < 6) + finally: + sm.close() From 058b1fa1023d1eca65911fe3e0f6adbfc3f04b63 Mon Sep 17 00:00:00 2001 From: seohui-XCENA Date: Mon, 6 Jul 2026 14:40:15 +0900 Subject: [PATCH 14/28] [MP][Maru] Complete the L1 conformance suite (shared error-path contracts) - add cross-backend conformance cases for the shared error paths: reserve_read of a mid-write key is KEY_NOT_READABLE (not a miss), finish_read/finish_write on an unstaged key is KEY_NOT_EXIST, delete of a missing key is KEY_NOT_EXIST - clear(force=True) is intentionally not conformance-tested: stock clears all objects while maru only drops local staging (the shared directory belongs to other instances) -- a legitimate divergence, not drift Signed-off-by: seohui-XCENA --- .../test_l1_manager_conformance.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/v1/distributed/test_l1_manager_conformance.py b/tests/v1/distributed/test_l1_manager_conformance.py index e856c814719..0c7b2d65867 100644 --- a/tests/v1/distributed/test_l1_manager_conformance.py +++ b/tests/v1/distributed/test_l1_manager_conformance.py @@ -206,6 +206,25 @@ def test_promote_fires_promote_event(manager): manager.finish_read([k]) +def test_reserve_read_mid_write_is_not_readable(manager): + """A key reserved-but-not-finished for write is not readable (not a miss).""" + k = _key(13) + manager.reserve_write([k], [False], _LAYOUT, mode="new") + assert manager.reserve_read([k])[k] == (L1Error.KEY_NOT_READABLE, None) + + +def test_finish_on_unstaged_key_is_not_exist(manager): + """finish_read / finish_write on a never-reserved key report KEY_NOT_EXIST.""" + k = _key(14) + assert manager.finish_read([k])[k] == L1Error.KEY_NOT_EXIST + assert manager.finish_write([k])[k] == L1Error.KEY_NOT_EXIST + + +def test_delete_missing_key_is_not_exist(manager): + k = _key(15) + assert manager.delete([k])[k] == L1Error.KEY_NOT_EXIST + + def test_listener_fires_across_lifecycle(manager): """Both backends fire the same on_l1_keys_* events across a lifecycle. From 5d0d1a5f2181cc3f42844f2fd9c793c5fa6c6558 Mon Sep 17 00:00:00 2001 From: seohui-XCENA Date: Mon, 6 Jul 2026 21:42:15 +0900 Subject: [PATCH 15/28] [MP][Maru] Anchor L1 eviction watermark to CXL device free space - get_memory_usage total = owned pool + CXL device free (cxl_pool.free_size from get_stats), so the eviction watermark tracks whole-device fill instead of just this instance's owned pool (which evicts prematurely while the device still has room) - cache the last-known free and reuse it when a get_stats RPC omits cxl_pool (transient timeout / older server) so total never momentarily collapses to the owned pool and fires a spurious eviction Signed-off-by: seohui-XCENA --- lmcache/v1/distributed/maru_l1_manager.py | 41 +++++++++++++++++++---- 1 file changed, 34 insertions(+), 7 deletions(-) diff --git a/lmcache/v1/distributed/maru_l1_manager.py b/lmcache/v1/distributed/maru_l1_manager.py index 1809d3cad6a..3205ac81b93 100644 --- a/lmcache/v1/distributed/maru_l1_manager.py +++ b/lmcache/v1/distributed/maru_l1_manager.py @@ -170,6 +170,11 @@ def __init__(self, config: L1ManagerConfig) -> None: self._read_ttl_seconds = config.read_ttl_seconds self._lock = threading.Lock() self._allocator = MaruMemoryAllocator(maru_config) + # MARU: last-known CXL device free (from get_stats ``cxl_pool``). Reused + # when a get_stats RPC fails to deliver it (e.g. a transient timeout) so + # the eviction watermark's ``total`` does not momentarily collapse to the + # owned pool and trip a spurious eviction. Stays 0 until the first read. + self._last_cxl_free: int = 0 self._pending_read: dict[ObjectKey, _PendingRead] = {} self._pending_write: dict[ObjectKey, _PendingWrite] = {} self._registered_listeners: list[L1ManagerListener] = [] @@ -970,11 +975,24 @@ def is_key_evictable(self, key: ObjectKey) -> bool: return key not in self._pending_read and key not in self._pending_write def get_memory_usage(self) -> tuple[int, int]: - """Return (used, total) bytes of this instance's owned CXL regions. - - MARU: usage is the owned-region allocation pressure -- the basis the - stock eviction watermark expects. Before the pool is up, reports the - configured capacity so watermark math stays sane. + """Return (used, total) bytes for the eviction watermark. + + MARU: ``used`` is this instance's owned-region allocation; ``total`` is + that owned pool **plus the CXL device's free space** + (``cxl_pool.free_size`` from the shared resource manager). Anchoring + ``total`` to "owned pool + free" -- not just the owned pool -- keeps the + stock watermark from tripping while the device still has room: with free + CXL left the pool auto-expands instead of evicting, and only once the + device is full (``free_size`` 0 -> ``total`` collapses to the owned pool) + does eviction engage on this instance's own pages. Before the pool is up, + reports the configured capacity so watermark math stays sane. + + ``free`` is cached across calls (``_last_cxl_free``): a get_stats RPC that + does not deliver ``cxl_pool`` -- a transient timeout, or an older server + that never sends it -- reuses the last-known free instead of dropping to + 0, so a momentary RPC failure does not collapse ``total`` to the owned + pool and fire a spurious eviction. It stays 0 until the first successful + read (older servers therefore keep the prior owned-pool behavior). Returns: A tuple of (used_bytes, total_bytes). @@ -983,11 +1001,20 @@ def get_memory_usage(self) -> tuple[int, int]: return 0, self._config.pool_size_bytes try: handler = self._allocator.handler - regions = handler.get_stats().get("store_regions") + stats = handler.get_stats() + regions = stats.get("store_regions") if not regions: return 0, self._config.pool_size_bytes used = regions["total_allocated_pages"] * handler.get_chunk_size() - return used, regions["total_pool_size"] + own_pool = regions["total_pool_size"] + # A present cxl_pool (incl. free_size 0 when the device is genuinely + # full) updates the cache; a missing one reuses the last-known free. + free = stats.get("cxl_pool", {}).get("free_size") + if free is None: + free = self._last_cxl_free + else: + self._last_cxl_free = free + return used, own_pool + free except Exception: logger.exception("MaruL1Manager: get_stats failed") return 0, self._config.pool_size_bytes From 9fe29549cd8b9767e8eafd6ddab75fff9ef5ea28 Mon Sep 17 00:00:00 2001 From: seohui-XCENA Date: Tue, 7 Jul 2026 13:36:31 +0900 Subject: [PATCH 16/28] [MP][Maru] Fix L1 read-pin accounting and mid-write read handling - Track server pins on _PendingRead.pinned separately from refcount, so a temporary stage that absorbs an overlapping reserve_read's pins releases them; previously those pins leaked and left the page un-evictable on MaruServer. - Exclude mid-write keys from reserve_read's pin/stage step: a key mid-write on this instance stays KEY_NOT_READABLE even when a peer has registered it, instead of being double-staged in both _pending_write and _pending_read (which stranded the in-flight write and failed its promote). - Add regression tests for both paths. Signed-off-by: seohui-XCENA --- lmcache/v1/distributed/maru_l1_manager.py | 75 +++++++++++++------- tests/v1/distributed/test_maru_l1_manager.py | 66 +++++++++++++++++ 2 files changed, 116 insertions(+), 25 deletions(-) diff --git a/lmcache/v1/distributed/maru_l1_manager.py b/lmcache/v1/distributed/maru_l1_manager.py index 3205ac81b93..de978f9d39a 100644 --- a/lmcache/v1/distributed/maru_l1_manager.py +++ b/lmcache/v1/distributed/maru_l1_manager.py @@ -136,6 +136,13 @@ class _PendingRead: mem_obj: MemoryObj refcount: int + # Real MaruServer pins held for this entry (0 <= pinned <= refcount). A + # temporary promote stages a local page with no pin (pinned=0); an + # overlapping reserve_read that pins the directory copy adds to both counts. + # Release paths unpin ``pinned`` -- not ``refcount`` -- so pins absorbed onto + # a temporary entry are never leaked and pin-less holds are never + # over-unpinned. + pinned: int = 0 is_temporary: bool = False deadline: float = float("inf") @@ -322,13 +329,17 @@ def _pin_retrieve_stage( staged = self._pending_read.get(k) if staged is not None: # Overlapping reserve: same CXL page, one staged object. - # Refresh the TTL (mirrors a stock re-lock extending it). + # The pins just taken are real, so track them on ``pinned`` even + # when the existing entry is a temporary (pin-less) stage -- + # otherwise those pins would never be released. Refresh the TTL + # (mirrors a stock re-lock extending it). staged.refcount += total + staged.pinned += total staged.deadline = deadline ret[k] = (L1Error.SUCCESS, staged.mem_obj) else: self._pending_read[k] = _PendingRead( - mem_obj=mem_obj, refcount=total, deadline=deadline + mem_obj=mem_obj, refcount=total, pinned=total, deadline=deadline ) ret[k] = (L1Error.SUCCESS, mem_obj) successful_keys.append(k) @@ -439,6 +450,11 @@ def _sweep_once(self) -> None: to_unpin: list[str] = [] for k in expired_reads: read_entry = self._pending_read.pop(k) + # MARU: release real server pins (pinned); a temporary stage + # (pinned 0) reclaims its private page instead. A temporary that + # absorbed pins does both. + if read_entry.pinned: + to_unpin.extend([object_key_to_string(k)] * read_entry.pinned) if read_entry.is_temporary: try: self._allocator.abort_alloc(read_entry.mem_obj) @@ -446,8 +462,6 @@ def _sweep_once(self) -> None: logger.exception( "MaruL1Manager: sweep failed to reclaim read page %s", k ) - else: - to_unpin.extend([object_key_to_string(k)] * read_entry.refcount) if to_unpin: self._safe_unpin(self._allocator.handler, to_unpin) logger.warning( @@ -483,17 +497,24 @@ def reserve_read( ret: dict[ObjectKey, L1OperationResult] = { k: (L1Error.KEY_NOT_EXIST, None) for k in keys } + # PARITY(L1Manager.reserve_read): a key mid-write on this instance is + # not readable (distinct from a plain miss), mirroring stock's per-key + # write/read lock exclusion. A peer may have already registered the same + # key in the shared directory, so mid-write keys are excluded from the + # pin/stage below -- otherwise the pin would succeed and stage a + # _pending_read entry for a key still in _pending_write (double staging), + # stranding the in-flight write (its promote then returns + # KEY_IN_WRONG_STATE and never pops _pending_write). + readable: list[ObjectKey] = [] for k in keys: - # PARITY(L1Manager.reserve_read): a mid-write key is not readable - # (distinct from a plain miss). It is unregistered, so the pin - # below misses and this pre-marked result stands. if k in self._pending_write: ret[k] = (L1Error.KEY_NOT_READABLE, None) - if not keys: - return ret + else: + readable.append(k) successful_keys: list[ObjectKey] = [] - # MARU: pin + retrieve + resolve + stage (shared with retained promote). - self._pin_retrieve_stage(keys, total, ret, successful_keys) + if readable: + # MARU: pin + retrieve + resolve + stage (shared w/ retained promote). + self._pin_retrieve_stage(readable, total, ret, successful_keys) # PARITY(L1Manager.reserve_read): notify listeners of the new read # holds (feeds the eviction LRU and the store controller). for listener in self._registered_listeners: @@ -571,18 +592,20 @@ def finish_read( k, ) entry.refcount -= released - # MARU: temporary reads are local staging (never directory-pinned), - # so at refcount zero the page is reclaimed, not unpinned. Normal - # reads release the server pins taken by reserve_read. - if entry.is_temporary: - if entry.refcount <= 0: + # MARU: release real server pins up to what we still hold. A pure + # temporary stage has pinned=0 (nothing to unpin); a temporary that + # absorbed an overlapping reserve's pins releases those here. + unpin_now = min(released, entry.pinned) + if unpin_now: + entry.pinned -= unpin_now + to_unpin.extend([object_key_to_string(k)] * unpin_now) + if entry.refcount <= 0: + # MARU: a temporary stage is an unregistered local page -> + # reclaim it through the allocator (a directory read is not). + if entry.is_temporary: need_to_free.append(entry.mem_obj) need_to_free_keys.append(k) - del self._pending_read[k] - else: - to_unpin.extend([object_key_to_string(k)] * released) - if entry.refcount <= 0: - del self._pending_read[k] + del self._pending_read[k] ret[k] = L1Error.SUCCESS successful_keys.append(k) if to_unpin: @@ -791,6 +814,7 @@ def finish_write_and_reserve_read( self._pending_read[k] = _PendingRead( mem_obj=entry.mem_obj, refcount=total, + pinned=0, # local page: not directory-pinned is_temporary=True, deadline=time.monotonic() + self._read_ttl_seconds, ) @@ -926,7 +950,7 @@ def clear(self, force: bool = False) -> None: self._publish(EventType.L1_KEYS_EVICTED, dropped) def _drain_staging(self, force: bool) -> list[ObjectKey]: - """Unpin staged reads (refcount times) and abort staged writes. + """Unpin staged reads (``pinned`` times) and abort staged writes. Args: force: If True, log the dropped in-flight staging as a warning. @@ -936,11 +960,12 @@ def _drain_staging(self, force: bool) -> list[ObjectKey]: """ to_unpin: list[str] = [] for k, entry in self._pending_read.items(): - # MARU: temporary reads hold no server pin -- reclaim the page. + # MARU: release real server pins (pinned); a temporary stage holds a + # private page (pinned may be 0) that is reclaimed instead. + if entry.pinned: + to_unpin.extend([object_key_to_string(k)] * entry.pinned) if entry.is_temporary: self._allocator.abort_alloc(entry.mem_obj) - else: - to_unpin.extend([object_key_to_string(k)] * entry.refcount) if force and (self._pending_read or self._pending_write): logger.warning( "MaruL1Manager: force-clear drops %d staged reads " diff --git a/tests/v1/distributed/test_maru_l1_manager.py b/tests/v1/distributed/test_maru_l1_manager.py index 72f9a846332..5199ad0dc2c 100644 --- a/tests/v1/distributed/test_maru_l1_manager.py +++ b/tests/v1/distributed/test_maru_l1_manager.py @@ -158,6 +158,72 @@ def test_finish_read_never_over_releases(): assert handler.pins[object_key_to_string(k)] == 0 +def test_temporary_read_absorbing_overlapping_pins_releases_them(): + """A peer-registered key read while temporary-staged must not leak pins. + + A temporary promote stages ``k`` locally with no server pin. If a peer has + registered ``k`` in the directory, a later ``reserve_read`` pins the shared + copy and folds those pins into the temporary entry. Draining the reads must + release every absorbed pin -- regression: the temporary release path used to + reclaim the page without unpinning, leaking the pin permanently (the page + became un-evictable on MaruServer). + """ + manager, handler, _ = make_maru_manager() + k = _key(1) + ks = object_key_to_string(k) + + # Temporary promote: staged in _pending_read, no directory pin. + manager.reserve_write([k], [True], _LAYOUT, mode="new") + manager.finish_write_and_reserve_read([k]) + assert handler.pins.get(ks, 0) == 0 # temporary -> no server pin + + # A peer registers k; an overlapping reserve_read pins the shared copy and + # absorbs that pin onto the temporary entry. + _seed(handler, k) + manager.reserve_read([k]) + assert handler.pins[ks] == 1 + + # Drain both holds (temporary hold + absorbed read): every pin released. + manager.finish_read([k]) + manager.finish_read([k]) + assert handler.pins[ks] == 0 # was 1 (leaked) before the fix + assert ks not in manager._pending_read + + +def test_reserve_read_excludes_mid_write_key_no_double_staging(): + """A key mid-write on this instance stays KEY_NOT_READABLE even if a peer + registered it -- it must not be pinned or read-staged. + + A local write reserves ``k`` (staged in _pending_write); before it finishes, + a peer registers ``k`` in the shared directory. A reserve_read that slips in + must refuse ``k``, not serve the peer copy: staging it would leave ``k`` in + both _pending_write and _pending_read (double staging) and strand the write + -- its promote would then return KEY_IN_WRONG_STATE and never pop + _pending_write, orphaning the page until the write-TTL sweeper. + """ + manager, handler, _ = make_maru_manager() + k = _key(1) + ks = object_key_to_string(k) + + # Local write in flight: k staged for write, not yet registered by us. + manager.reserve_write([k], [True], _LAYOUT, mode="new") + # A peer registers the same key in the shared directory mid-write. + _seed(handler, k) + + # A read racing the in-flight write is refused, not served from the peer. + res = manager.reserve_read([k]) + assert res[k] == (L1Error.KEY_NOT_READABLE, None) + assert handler.pins.get(ks, 0) == 0 # no pin taken on the peer copy + assert k not in manager._pending_read # no double staging + assert k in manager._pending_write # write still in flight + + # The write still completes cleanly: promote pops _pending_write and stages + # the read (no orphan, no KEY_IN_WRONG_STATE). + fin = manager.finish_write_and_reserve_read([k]) + assert fin[k][0] == L1Error.SUCCESS + assert k not in manager._pending_write + + # ========================================================================= # reserve_write / finish_write # ========================================================================= From 0b984c1636eb8561679a216faab771d956201767 Mon Sep 17 00:00:00 2001 From: seohui-XCENA Date: Tue, 7 Jul 2026 13:36:42 +0900 Subject: [PATCH 17/28] [MP][Maru] Document MaruMemoryAllocator public methods - Add docstrings to allocate, batched_allocate, close, and memcheck. - Note the RuntimeError raised when called before init_layout(), and the out-of-memory -> None contract that reserve_write's all-or-nothing handling depends on. Signed-off-by: seohui-XCENA --- .../memory_manager/maru_memory_allocator.py | 55 +++++++++++++++++++ 1 file changed, 55 insertions(+) diff --git a/lmcache/v1/distributed/memory_manager/maru_memory_allocator.py b/lmcache/v1/distributed/memory_manager/maru_memory_allocator.py index fab36c4503f..9d8388b0ca0 100644 --- a/lmcache/v1/distributed/memory_manager/maru_memory_allocator.py +++ b/lmcache/v1/distributed/memory_manager/maru_memory_allocator.py @@ -156,6 +156,25 @@ def allocate( fmt: MemoryFormat = MemoryFormat.UNDEFINED, allocator_type: Optional[str] = None, ) -> Optional[MemoryObj]: + """Allocate a single CXL page for one KV chunk. + + Delegates to the maru ``CxlMemoryAdapter`` once the pool is bound. + + Args: + shapes: Tensor shape (or per-object-group shapes) of the KV chunk. + dtypes: Tensor dtype (or per-object-group dtypes) of the KV chunk. + fmt: Memory format tag recorded on the returned object. + allocator_type: Optional allocator selector; unused by the CXL pool + (accepted for interface parity). + + Returns: + A ``MemoryObj`` viewing the allocated CXL page, or ``None`` when the + pool cannot fit the request (out of memory). Callers treat ``None`` + as OUT_OF_MEMORY. + + Raises: + RuntimeError: If called before ``init_layout()`` binds the pool. + """ return self._require_init("allocate").allocate( shapes, dtypes, fmt, allocator_type ) @@ -168,6 +187,29 @@ def batched_allocate( fmt: MemoryFormat = MemoryFormat.UNDEFINED, allocator_type: Optional[str] = None, ) -> Optional[List[MemoryObj]]: + """Allocate ``batch_size`` CXL pages in one call (all-or-nothing). + + Delegates to the maru ``CxlMemoryAdapter`` once the pool is bound. + + Args: + shapes: Tensor shape (or per-object-group shapes) shared by every + page in the batch. + dtypes: Tensor dtype (or per-object-group dtypes) shared by every + page in the batch. + batch_size: Number of pages to allocate. + fmt: Memory format tag recorded on each returned object. + allocator_type: Optional allocator selector; unused by the CXL pool + (accepted for interface parity). + + Returns: + A list of ``batch_size`` ``MemoryObj`` views, or ``None`` when the + whole batch cannot fit (out of memory) -- the allocation is + all-or-nothing, never a partial list. Callers treat ``None`` as + OUT_OF_MEMORY. + + Raises: + RuntimeError: If called before ``init_layout()`` binds the pool. + """ return self._require_init("batched_allocate").batched_allocate( shapes, dtypes, batch_size, fmt, allocator_type ) @@ -232,6 +274,12 @@ def single_token_size(self) -> int: # lifecycle def close(self) -> None: + """Tear down the CXL adapter and disconnect the MaruHandler. + + Idempotent: safe to call when the pool was never brought up. Both the + adapter and handler are cleared, so any later use raises. Shared CXL + pages are not freed here -- their lifecycle is owned by MaruServer. + """ if self._cxl_adapter is not None: self._cxl_adapter.close() self._cxl_adapter = None @@ -240,4 +288,11 @@ def close(self) -> None: self._handler = None def memcheck(self) -> bool: + """Report memory-consistency health. + + Returns: + Always ``True``. The CXL pool's page accounting is owned by + MaruServer, so there is no local invariant to verify (no-op parity + with the stock allocator's memcheck). + """ return True From 74a2d3e68d6d4b69cb58f5aae5b64b73b5e46c00 Mon Sep 17 00:00:00 2001 From: seohui-XCENA Date: Tue, 7 Jul 2026 13:36:50 +0900 Subject: [PATCH 18/28] [MP][Maru] Add L1 design doc and user configuration docs - Add docs/design/v1/distributed/maru_l1.md: the sibling + Protocol seam, the MaruL1Manager API contracts, the pending-read/pending-write state-machine invariants, the TTL sweeper, and the crash-recovery premises. - Document the Maru CXL shared L1 tier in the MP configuration guide and the architecture index; point the deprecated in-process maru page to it. Signed-off-by: seohui-XCENA --- docs/design/v1/distributed/maru_l1.md | 307 ++++++++++++++++++ .../source/kv_cache/storage_backends/maru.rst | 2 +- docs/source/mp/configuration.rst | 71 ++++ docs/source/mp/index.rst | 24 +- 4 files changed, 398 insertions(+), 6 deletions(-) create mode 100644 docs/design/v1/distributed/maru_l1.md diff --git a/docs/design/v1/distributed/maru_l1.md b/docs/design/v1/distributed/maru_l1.md new file mode 100644 index 00000000000..41893ca21b0 --- /dev/null +++ b/docs/design/v1/distributed/maru_l1.md @@ -0,0 +1,307 @@ +# Maru CXL Shared L1 (`MaruL1Manager`) + +Design doc for the MP-mode L1 backend that stores the L1 tier in a +cross-instance **CXL shared pool** instead of local CPU DRAM. Covers the code in +`lmcache/v1/distributed/maru_l1_manager.py`, `l1_protocol.py`, and +`memory_manager/maru_memory_allocator.py`. + +This is a review-oriented summary: it states the **contracts and invariants** the +implementation must satisfy so the code can be checked against them. It is not an +exhaustive walkthrough. + +## 0. One-line summary + +**Maru replaces MP mode's L1 tier storage (CPU DRAM → a cross-instance shared CXL +pool); everything above L1 — L1↔L2 tiering (write-through / discard-evict / +promote-on-miss), the controllers, the L2 adapters, eviction — runs unchanged on +top.** LMCache never uses L1 in isolation, so Maru is not an "L1-only" cache +either. + +The whole integration reduces to one decision: + +> Isolate all Maru logic in a single class, `MaruL1Manager`, and run the existing +> stack unchanged on top of it through the `L1ManagerInterface` Protocol seam. + +## 1. Model — what changes and what does not + +### 1.1 Maru = one cross-instance shared L1 tier + +- **Shared medium**: a physical pool reached over a CXL switch. Multiple LMCache + instances `mmap` the same physical memory → **zero-copy**, no P2P transfer. +- **Shared index**: `MaruServer` holds the `key → (region, offset)` directory, a + per-region `kv_ref_count`, and a cross-node `pin_count`. +- **Access**: a retrieve pins + looks up the key on `MaruServer`, then + materializes a `MemoryObj` pointing at the resident CXL page (no copy). + +### 1.2 Control (LMCache) / Directory (MaruServer) split + +- **LMCache decides**: store / evict / promote / pin are decided by the existing + LMCache controllers. +- **MaruServer is passive**: it executes and records those decisions as a shared + directory + region broker. It holds no policy. +- LMCache's L1 manager therefore needs Maru RPC hooks for store/retrieve/pin/ + delete, and `MaruL1Manager` is the home of those hooks. + +### 1.3 Unchanged (byte-identical to `dev`) + +- `L1Manager`, `L1MemoryManager`, `internal_api.py` — zero Maru code. +- `StoreController` / `PrefetchController` / `L1EvictionController` / the L2 + adapters / `L2EvictionController` — logic unchanged; only the `l1_manager` + parameter type is widened to the Protocol (§3). +- The data plane (`mmap` zero-copy, `metadata.address` bit-pack, `gpu_ops`). + +### 1.4 The three layers of MP L1 — reuse / reimplement / replace + +The recurring confusion is *which layer* "we keep the control logic" refers to. +MP L1 has three layers and the Maru seam cuts each differently: + +| Layer | Components | Stock local state | Role | Maru seam | +|---|---|---|---|---| +| **1 — Controllers / tiering** | `StoreController`, `PrefetchController`, `L1EvictionController`, `L2EvictionController`, L2 adapters, eviction `policy` (LRU) | policy's `OrderedDict` order (fed by listeners) | **decide** write-through / promote / eviction; act only on the L1Manager public surface | **reused** (only the `l1_manager` param type is widened; behavior-neutral) | +| **2 — L1Manager control state machine** | `L1Manager` | `_objects: dict[ObjectKey, L1ObjectState]` with per-key `write_lock`/`read_lock` (TTLLock) | authority for **this instance's** membership + lock state | **reimplemented** as sibling `MaruL1Manager` (MaruServer RPC instead of a local dict) | +| **3 — Memory manager / allocator** | `L1MemoryManager` + `MemoryAllocator` (DRAM / DevDax / GDS) | pinned DRAM/devdax/GDS buffer pool | the medium the KV bytes live in | **replaced** by `MaruMemoryAllocator` (CXL, zero-copy) | + +"Control logic = stock" is a statement about **layer 1**, and it holds. But layer +2 cannot be reused as-is, which is why Maru is a **sibling** (layer-2 +reimplementation), not merely an allocator swap (layer-3 replacement). + +### 1.5 Why an allocator swap (layer-3 only) is not enough + +Keeping stock layer 2 (`L1Manager._objects`) and swapping only the allocator +breaks, because `_objects` is a **per-instance local index**: + +1. **Membership**: if instance B stores key `K`, this instance's local `_objects` + never learns about `K`, so `reserve_read(K)` misses — even though `K` is + physically present in the shared pool. Cross-instance read collapses. +2. **Pin**: stock `read_lock`/`write_lock` are **local** (a C++ atomic TTLLock), + invisible to other instances. Maru must prevent node A from evicting a page + node B is reading, which requires a **cross-node shared refcount** (MaruServer + `pin_count`). A local TTLLock cannot express that. + +GDS/DevDax work as allocator swaps because they are **private per-instance** +tiers, so `_objects` is complete and accurate. Maru is a **shared** tier, so the +authority for membership and pins must be **remote** (MaruServer), and layer 2 is +reimplemented over RPC. (The allocator interface has only +`allocate`/`free`/`get_memory_usage`/…; it has no "does K exist?" or "pin K" +control operation, so a layer-3 swap cannot express Maru's needs.) + +> "Same `MemoryObj`" means the *buffer type* is shared (so the data plane and +> allocator machinery are reusable); it does **not** mean the membership/pin +> control state may stay local. Those are independent. + +### 1.6 Stock local structures → Maru replacements + +Layer 2's authority moves from local to remote (MaruServer); only in-flight state +stays local: + +| Stock (layer 2, local) | Role | Maru replacement | Where | +|---|---|---|---| +| `_objects` dict | membership authority | MaruServer directory `key→region/offset` + `kv_ref_count` | **remote (shared)** | +| `L1ObjectState.memory_obj` | buffer handle | CXL page `MemoryObj` (`get_by_location`, zero-copy view) | local view | +| `write_lock`/`read_lock` (TTLLock) | local eviction/write protection | cross-node pin via MaruServer `pin_count` (delete refuses with PINNED) | **remote (shared)** | +| (no in-flight-write tracking) | reserve_write → finish_write window | `_pending_write` (mem_obj, `is_temporary`, deadline) + write-TTL | local | +| `L1ObjectState.is_temporary` | promote lifetime | `_PendingRead.is_temporary` (propagated from `_pending_write` by promote) | local | +| (no staged-read tracking) | reserve_read → finish_read window | `_pending_read` (`_PendingRead` refcount + `pinned` + `is_temporary`) + read-TTL | local | +| `L1MemoryManager` (DRAM pool) | medium | `MaruMemoryAllocator` (CXL pool) | local | + +**Everything Maru keeps local is in-flight/staged side-channel state; all +authority (membership, pins) is remote.** Layer 1 and the data plane are +untouched. TTL reclaim of the side-channel reuses the stock `write_ttl`/`read_ttl` +config values, but sweeps a monotonic deadline on the pending entries instead of a +TTLLock. + +## 2. The seam — sibling + structural Protocol + +`StorageManager` picks the L1 manager in one line at construction: + +```python +self._l1_manager: L1ManagerInterface = ( + MaruL1Manager(cfg) if maru_config else L1Manager(cfg) +) +``` + +Both implementations satisfy `L1ManagerInterface`, so the stack above does not +know which one it has. + +**Three new files:** + +- `l1_protocol.py` — the structural `typing.Protocol` `L1ManagerInterface` + (`@runtime_checkable`). `L1Manager` and `MaruL1Manager` both satisfy it + **without inheritance**. Each method's docstring states its behavioral contract + (when listeners fire, the PINNED retry, `extra_count` balance). +- `maru_l1_manager.py` — `MaruL1Manager`. Implements the whole RPC control + surface and owns `MaruMemoryAllocator` directly as `self._allocator`. No + separate dispatch layer. +- `memory_manager/maru_memory_allocator.py` — `MaruMemoryAllocator` (same + directory as the devdax/gds allocators). Thin wrapper over the external + `maru_lmcache.CxlMemoryAdapter`; connects the handler and creates the adapter on + the first `init_layout`. `free` is a no-op (CXL page lifecycle is MaruServer's). + +The stock controllers call only the `L1ManagerInterface` surface (no private +access — verified), so widening each controller's `l1_manager` parameter type from +concrete `L1Manager` to `L1ManagerInterface` (runtime-identical) lets the whole +stack run over `MaruL1Manager`. The only remaining Maru-specific concrete +references are the L1-selection branch in `StorageManager` and the maru-only gate +in its `register_kv_layout` wrapper — both legitimately Maru-only. + +## 3. What `MaruL1Manager` implements + +The `L1ManagerInterface` methods, reimplemented over RPC. "Fire" means invoking +the registered listeners' `on_l1_keys_*` callbacks (the eviction LRU and the store +controller consume these). + +| Method | Implemented behavior | +|---|---| +| `reserve_read(keys, extra_count)` | `batch_pin` + `batch_retrieve` + `get_by_location` (zero-copy) → stage in `_pending_read`. Takes `1 + extra_count` server pins per key (MLA + TP>1). A key mid-write on this instance is **excluded** (stays `KEY_NOT_READABLE`, see §4 invariant 1). Fires `on_l1_keys_reserved_read`. | +| `unsafe_read` / `finish_read` | Return the staged object / release `1 + extra_count` holds. Each release drops `refcount` and unpins `min(released, pinned)` server pins (§4 invariant 2). At `refcount == 0` a **temporary** entry frees its private page (`on_l1_keys_deleted_by_manager`); a directory read just drops the entry. | +| `reserve_write` / `finish_write` | Allocate a CXL page + stage in `_pending_write` (`is_temporary` + write-TTL); fire `on_l1_keys_reserved_write`. / `create_store_handle` + `batch_store` (directory register); fire `on_l1_keys_write_finished`. In `mode="new"`, a key already staged locally **or** already in the directory returns `KEY_NOT_WRITABLE` (cross-instance dedup). | +| `finish_write_and_reserve_read` (promote) | Branches on the staged `is_temporary` (§4 invariant 3). **temporary** (the default prefetch policy): no `batch_store`, no pin — the loaded local page moves straight to `_pending_read`; freed at `finish_read == 0`. **retained**: `batch_store`, then re-resolve the authoritative directory page with pins (a dup-skip auto-freed our page). Both fire only `on_l1_keys_finish_write_and_reserve_read` — **never** `on_l1_keys_write_finished` (that would make the store controller re-store the promoted key to L2). | +| `delete(keys)` | Directory delete. A local read-hold → `KEY_IS_LOCKED`; a cross-node pinned key → MaruServer refuses with **PINNED**. Only actually-deleted keys fire `on_l1_keys_deleted_by_manager` (§4 invariant 4). | +| `is_key_evictable(key)` | False if locally pinned (`_pending_read`); lock-free. Cross-node pins are re-checked authoritatively by `delete`. | +| `touch_keys` / `register_listener` | `register_listener` only stores the listener; each method fires its own event. `touch_keys` fires `on_l1_keys_accessed`. | +| `get_memory_usage()` | `handler.get_stats` → `(used, total)`. `used` = owned-pool allocated; `total` = owned pool + CXL device free (`cxl_pool.free_size`), so the eviction watermark tracks whole-device fill, not just the owned pool. The last known free is cached to survive a transient `get_stats` RPC failure. Before init, returns `(0, configured pool size)` to avoid a 0-division in the controller watermark. | +| `get_l1_memory_desc()` | `None` — the shared pool has no single registrable region (§5 L2 limit). | +| `register_kv_layout(...)` | Maru-only: binds the KV layout to `MaruMemoryAllocator` (pool bring-up). `num_object_groups > 1` is rejected by the `StorageManager.register_kv_layout` wrapper (§6). | +| `clear` / `memcheck` / `report_status` / `close` | Delegated under the lock (§4). `clear` empties the local side-channel and releases staged pins but does **not** touch shared data (a shared tier must not be wiped by one node). `close` stops the TTL sweeper. | + +Eviction and LRU stay with the stock stack: `L1EvictionController` decides, and +recency lives in the eviction **policy** (fed by the listener events). +`MaruL1Manager` keeps no local object/LRU dict. + +## 4. Invariants (check the implementation against these) + +These are the contracts a reviewer should verify; the two correctness bugs found +in review were both violations of invariants 1 and 2. + +1. **A key is in at most one of `_pending_write` / `_pending_read`.** This mirrors + stock `L1Manager`'s per-key `write_lock`/`read_lock` mutual exclusion (a + mid-write key is not readable). Enforced in both directions: `reserve_write` + and `finish_write_and_reserve_read` reject a key already staged the other way, + and `reserve_read` **excludes** mid-write keys from the pin/stage step (a peer + may have registered the same key, so the pin would otherwise succeed and double + stage — stranding the in-flight write). + +2. **Pin accounting: `pinned` tracks real server pins separately from + `refcount`** (`0 <= pinned <= refcount`). Release paths unpin `pinned`, not + `refcount`. A pure temporary stage has `pinned == 0` (a private local page, no + server pin); a temporary that absorbs an overlapping `reserve_read`'s pins + records them in `pinned` so they are released, not leaked. `N` reserves take + `N` pins and `N` finishes release `N`; `finish_read` never releases more than + it holds (over-release would corrupt the server `pin_count`). + +3. **temporary vs retained promote.** The default prefetch policy marks every + promote **temporary**: the page is private staging, never registered in the + directory, and discarded after one read (the shared pool is populated only by + store write-through). Only a hot-cache policy marks a promote **retained**, + which registers it in the directory. + +4. **PINNED cross-node retry (no controller change).** MaruServer refuses to + delete any key with `pin_count > 0` (PINNED). `MaruL1Manager.delete` fires + `on_l1_keys_deleted_by_manager` only for keys it **actually** deleted, so a + PINNED-refused key stays in the eviction policy and is retried next cycle. This + relies on the stock LRU contract that a key is removed from the policy only via + the `on_keys_removed` event, never popped by `get_eviction_actions`. + +### Thread safety + +A single non-reentrant `threading.Lock` (`_maru_l1_synchronized`) serializes every +public method that touches the side-channel or allocator. The decorator is typed +with `ParamSpec`/`Concatenate` so wrapped signatures survive the Protocol +conformance check. + +### TTL sweeper + crash-recovery premises + +An entry whose reserve→finish flow is interrupted (an *orphan*) cannot be +distinguished from an in-progress one by inspecting the pending dicts, so **time +(TTL) is the only abandonment signal** (reusing the stock `write_ttl`/`read_ttl` +values). A daemon sweeper thread (scanning under the lock) reclaims both: + +- `_pending_write`: on write-TTL expiry, return the page to the owner's local + free-list (`abort_alloc`) + pop. An unregistered page is invisible to the + server, so client-side reclaim is the only option. +- `_pending_read`: on read-TTL expiry, `batch_unpin × remaining refcount` + pop. + +A late `finish_read`/`unsafe_read` after expiry sees `KEY_NOT_EXIST` → retrieve +returns False → recompute (the same failure path as a stock TTL expiry). Double +unpin is impossible under the serializing lock. + +**Premises for a full client crash** (the MP-server process dies, so the local +sweeper cannot run): + +- read pins are assumed to be released by the **maru side** tracking each client's + pin set and dropping them on disconnect; +- an in-flight write page is assumed to be covered by **region owner-release**. + +These two mechanisms live outside `MaruL1Manager` (on the Maru/MaruServer side) +and are premises of this design, not guarantees provided by this PR. + +## 5. Scope and known limits + +- **Copy-type L2 adapters only.** `get_l1_memory_desc()` returns `None`, so + copy-type adapters (aerospike, dax, fs, fs_native, hfbucket, mock, resp, s3, + plugin, raw_block, mooncake-TCP) work fully, while registered/RDMA-type adapters + (nixl_store, nixl_store_dynamic, mooncake-RDMA) and p2p are **rejected at + startup** (§6 guards; an allowlist, so a new registered adapter fails safe). + Registered-L2 support needs an allocator region accessor + per-region descriptor + and is follow-up work. +- **Single model / single object group.** `MaruMemoryAllocator` fixes the pool to + a single layout on the first `init_layout`, so a different model (different + shapes/dtypes/fmt/chunk) or a hybrid model (`num_object_groups > 1`) fails fast + with a `ValueError`. Multiple instances of the *same* model are fine. Mixed + layouts are follow-up work. +- **Single-device pool bound.** A pool is one contiguous allocation within one DAX + device, so `pool_size_bytes` cannot exceed a single CXL device's capacity. +- **LRU is a per-instance local view.** The policy is fed only by this instance's + listener events, so it does not see other nodes' access recency. A known + approximation for a shared cache. +- **Cross-owner reclaim.** A per-key `delete` returns the page only to the + *owner's* local free-list. Evicting a key another instance owns removes the + directory entry but does not reclaim that owner's page. Region-level return to + the shared pool is an owner-only concern outside this PR's critical path. (This + and the expand-vs-evict watermark interaction are the two multi-instance + consistency issues tracked for a separate team review.) + +## 6. File changes (this PR) + +**New** + +- `l1_protocol.py` — `L1ManagerInterface` + the `L1OperationResult` alias. +- `maru_l1_manager.py` — `MaruL1Manager` (~1050 lines). +- `memory_manager/maru_memory_allocator.py` — `MaruMemoryAllocator`. +- Tests under `tests/v1/distributed/` (fake-maru harness + manager unit + + conformance + config guard + control integration). + +**Modified (all additive / behavior-neutral)** + +- `store_controller.py` / `prefetch_controller.py` / `eviction_controller.py` — + `l1_manager` parameter type widened to `L1ManagerInterface`. +- `storage_manager.py` — the L1-selection branch (`MaruL1Manager` when + `maru_config` is set) + the maru-only `register_kv_layout` wrapper. +- `config.py` — the `maru_config` field, CLI, and validation (§ startup guards). +- `lmcache_driven_transfer.py` — the `register_kv_layout` engine hook (a no-op for + the default backend), wired in `register_kv_cache` and wrapped in try/except so a + failure closes the cache context. + +**Unchanged (byte-identical to `dev`)** + +- `l1_manager.py`, `l1_memory_manager.py`, `internal_api.py`, all controller + logic, all L2 adapters. + +### Startup guards (`validate_storage_manager_config`) + +Unsupported combinations are rejected at startup: maru + gds/devdax (mutual +exclusion), maru + `skip_l1` store policy, maru + registered/RDMA-type L2, maru + +p2p, and maru + engine_driven/auto transfer mode (maru requires +`--supported-transfer-mode lmcache_driven`). Without these guards the failures are +silent or surface as an `AttributeError` deep in the request path. + +## 7. Status + +Implemented on branch `maru-mp-l1`: the sibling `MaruL1Manager` + `l1_protocol.py` ++ `MaruMemoryAllocator`, with `l1_manager.py` / `l1_memory_manager.py` / +`internal_api.py` kept byte-identical to `dev`. The full Maru test suite passes on +GPU (`CUDA_VISIBLE_DEVICES=1`). + +A later phase may extract an `L1StateBackend` seam so Maru becomes a small backend +rather than a full sibling; that is separate follow-up work after this sibling +merges. diff --git a/docs/source/kv_cache/storage_backends/maru.rst b/docs/source/kv_cache/storage_backends/maru.rst index 39a82437077..0a86db84222 100644 --- a/docs/source/kv_cache/storage_backends/maru.rst +++ b/docs/source/kv_cache/storage_backends/maru.rst @@ -3,7 +3,7 @@ Maru .. warning:: - This page documents the behavior of LMCache's in-process mode (deprecated). Please consider using :doc:`LMCache MP mode ` for better feature support and performance. For the MP mode equivalent of this page, see :doc:`/mp/l2_storage/index`. + This page documents the behavior of LMCache's in-process mode (deprecated). Please consider using :doc:`LMCache MP mode ` for better feature support and performance. For the MP mode equivalent, see the *Maru CXL Shared L1* section of :doc:`/mp/configuration`. .. _maru-overview: diff --git a/docs/source/mp/configuration.rst b/docs/source/mp/configuration.rst index 362c38d9f94..ea2cc1f587b 100644 --- a/docs/source/mp/configuration.rst +++ b/docs/source/mp/configuration.rst @@ -298,6 +298,77 @@ flags apply to both; no configuration change is needed to switch vendors. - Open the slab with ``O_DIRECT`` (required for the GDS DMA fast path on ext4). +Maru CXL Shared L1 +------------------ + +Source: ``lmcache/v1/distributed/config.py`` + +Opt-in. Setting ``--maru-server-url`` switches the L1 medium from pinned DRAM to +a **cross-instance shared CXL pool** managed by `Maru +`_. Every LMCache server on the same pool +``mmap``\ s the same physical memory, so an L1 entry produced by one instance is +a zero-copy read for the others. The pinned-DRAM L1 options +(``--l1-size-gb`` / ``--l1-use-lazy`` / ``--l1-init-size-gb``) are then ignored; +the stock L1↔L2 tiering, controllers, and eviction run unchanged on top. + +Requires a running MaruServer and the ``maru`` package (``git clone`` the Maru +repo and run ``./install.sh``). Enable it alongside the lmcache-driven transfer +path: + +.. code-block:: bash + + lmcache server \ + --maru-server-url maru://localhost:5555 \ + --maru-pool-size-gb 4 \ + --supported-transfer-mode lmcache_driven \ + --eviction-policy LRU --max-workers 4 --port 6555 + +.. note:: + + Maru L1 constraints (rejected at startup otherwise): **lmcache-driven** + transfer mode only; mutually exclusive with GDS (``--gds-l1-path``) and + Device-DAX (``--l1-devdax-path``) L1; default store policy only (not + ``skip_l1``); **copy-type L2 adapters only** (``s3`` / ``fs`` / ``raw_block`` + / ``mock`` …; registered/RDMA adapters such as ``nixl`` and P2P are refused); + and **one model / object group per pool** (the pool binds to the first KV + layout it sees). Recency (LRU) is tracked per instance. + +.. list-table:: + :header-rows: 1 + :widths: 30 15 55 + + * - Argument + - Default + - Description + * - ``--maru-server-url`` + - Not set + - MaruServer endpoint (``maru://host:port`` or ``tcp://host:port``). + Setting this enables the Maru CXL shared L1 tier. + * - ``--maru-pool-size-gb`` + - ``0.0`` + - CXL pool size to request from MaruServer, in GB. Required (> 0) when + ``--maru-server-url`` is set. **Must fit within a single CXL device** — + see the note below. + * - ``--maru-instance-id`` + - auto UUID + - Stable client id reported to MaruServer for ownership tracking. + Auto-generated if omitted. + +.. note:: + + **``--maru-pool-size-gb`` is bounded by a single CXL device.** MaruServer + allocates each region as **one contiguous extent within a single DAX + device** — regions do not span devices. This value sizes both the initial + region and every auto-expand step, so it must be ≤ one device's usable + capacity (raw device size minus MaruServer metadata / WAL / alignment + overhead; e.g. a 256 GiB device holds somewhat less). Requesting more than + any device can satisfy fails the MaruServer allocation, and the LMCache + server then **fails to start** (the initial ``MaruHandler.connect()`` returns + an error). With multiple CXL devices, auto-expand can claim additional + regions from other devices, so total capacity is the sum across devices — + but no single region (hence no single ``--maru-pool-size-gb``) may exceed + one device. + L1 Manager TTLs ---------------- diff --git a/docs/source/mp/index.rst b/docs/source/mp/index.rst index 40e74dfbbb9..1ab9f8292dd 100644 --- a/docs/source/mp/index.rst +++ b/docs/source/mp/index.rst @@ -22,9 +22,9 @@ Key Benefits share a single L1 cache, maximizing KV reuse. - **Independent resource scaling** -- Allocate CPU memory for caching independently of GPU memory for inference. -- **Multi-tier storage (L1 + L2)** -- An L1 cache (in CPU DRAM, or an NVMe - slab via GPUDirect Storage) backed by persistent L2 storage via NIXL (GDS, - POSIX, HF3FS, and more). +- **Multi-tier storage (L1 + L2)** -- An L1 cache (in CPU DRAM, an NVMe + slab via GPUDirect Storage, or a cross-instance shared CXL pool via Maru) + backed by persistent L2 storage via NIXL (GDS, POSIX, HF3FS, and more). - **Built-in observability** -- Prometheus metrics and a telemetry event system out of the box. @@ -75,14 +75,18 @@ High-Level Architecture v StorageManager (distributed/storage_manager.py) | - |--- L1Manager (l1_manager.py) + |--- L1Manager (l1_manager.py) [default] | |--- L1MemoryManager (CPU DRAM) or | | GDSL1MemoryManager (NVMe slab via cuFile / hipFile) | |--- TTLLock per object (read/write) + | -- or, when --maru-server-url is set -- + |--- MaruL1Manager (maru_l1_manager.py) [shared CXL L1] + | |--- MaruMemoryAllocator (cross-instance shared CXL pool) + | |--- MaruServer directory (key->region/offset, cross-node pin_count) | |--- StoreController -----> L2 Adapter(s) (async L1->L2 push) |--- PrefetchController ---> L2 Adapter(s) (async L2->L1 load) - |--- EvictionController ----> L1Manager (watermark-triggered eviction) + |--- EvictionController ----> L1Manager / MaruL1Manager (watermark eviction) | v EventBus + OTel providers (observability) @@ -397,6 +401,16 @@ tiers selected at startup (both satisfy ``L1ManagerProtocol``): ROCm; see the *GDS L1 Tier* section of :doc:`configuration` for the vendor-specific requirements. The CPU tier is disabled in this mode. +When ``--maru-server-url`` is set, ``L1Manager`` *itself* is replaced by +``MaruL1Manager`` (``distributed/maru_l1_manager.py``) -- a sibling that +satisfies the manager-level ``L1ManagerInterface``. Unlike the DRAM/GDS tiers +above (allocator swaps *under* ``L1Manager``), Maru makes L1 a cross-instance +*shared* CXL pool: membership and pins live in a remote ``MaruServer`` directory +rather than a local dict, so the control state machine is re-implemented over +RPC. The stock controllers and L2 tiering drive it unchanged through +``L1ManagerInterface``. See the *Maru CXL Shared L1* section of +:doc:`configuration`. + L2 Adapters ~~~~~~~~~~~ From 5d1076e408a1708121b6082032032d386a2b5ed0 Mon Sep 17 00:00:00 2001 From: seohui-XCENA Date: Tue, 7 Jul 2026 15:44:42 +0900 Subject: [PATCH 19/28] [MP][Maru] Add auto_expand knob and anchor the eviction watermark to it - Add MaruL1Config.auto_expand (default True) with a --maru-auto-expand / --no-maru-auto-expand CLI flag, threaded into MaruConfig via the allocator. - Branch get_memory_usage on it: auto_expand on keeps the device-fill watermark (owned pool + CXL device free); off anchors total to the owned pool, so a hard-capped pool evicts before it is exhausted instead of OOMing while the device still has free space the pool cannot grow into. - Extend the fake handler with a cxl_free knob and add a watermark test. Signed-off-by: seohui-XCENA --- lmcache/v1/distributed/config.py | 19 ++++++++++++ lmcache/v1/distributed/maru_l1_manager.py | 30 +++++++++++++------ .../memory_manager/maru_memory_allocator.py | 1 + tests/v1/distributed/maru_fakes.py | 6 +++- tests/v1/distributed/test_maru_l1_manager.py | 21 +++++++++++++ 5 files changed, 67 insertions(+), 10 deletions(-) diff --git a/lmcache/v1/distributed/config.py b/lmcache/v1/distributed/config.py index 35033f78053..2ad7b7fdf8d 100644 --- a/lmcache/v1/distributed/config.py +++ b/lmcache/v1/distributed/config.py @@ -218,6 +218,14 @@ class MaruL1Config: eager_map: bool = True """Eagerly mmap peer regions for cross-instance zero-copy reads.""" + auto_expand: bool = True + """Whether the owned pool auto-expands into free CXL device space when it + fills. When True (default), the pool grows toward the device capacity and + the eviction watermark is anchored to device fill (owned pool + device + free); when False, the pool is hard-capped at ``pool_size_bytes`` and the + watermark is anchored to that pool, so eviction engages before the pool is + exhausted (see ``MaruL1Manager.get_memory_usage``).""" + @dataclass class L1ManagerConfig: @@ -514,6 +522,16 @@ def add_storage_manager_args( help="Stable client id reported to MaruServer for ownership tracking. " "Auto-generated if omitted.", ) + maru_group.add_argument( + "--maru-auto-expand", + action=argparse.BooleanOptionalAction, + default=True, + help="Whether the owned CXL pool auto-expands into free device space " + "when full. Default True: the pool grows toward device capacity and " + "eviction is anchored to device fill. Use --no-maru-auto-expand to " + "hard-cap the pool at --maru-pool-size-gb and evict before it is " + "exhausted.", + ) # L1 Manager Config (TTL settings) ttl_group = parser.add_argument_group( @@ -654,6 +672,7 @@ def parse_args_to_config( server_url=args.maru_server_url, pool_size_bytes=int(pool_size_gb * (1 << 30)), instance_id=getattr(args, "maru_instance_id", None), + auto_expand=getattr(args, "maru_auto_expand", True), ) shm_name = getattr(args, "shm_name", None) diff --git a/lmcache/v1/distributed/maru_l1_manager.py b/lmcache/v1/distributed/maru_l1_manager.py index de978f9d39a..3b99610f376 100644 --- a/lmcache/v1/distributed/maru_l1_manager.py +++ b/lmcache/v1/distributed/maru_l1_manager.py @@ -1002,15 +1002,22 @@ def is_key_evictable(self, key: ObjectKey) -> bool: def get_memory_usage(self) -> tuple[int, int]: """Return (used, total) bytes for the eviction watermark. - MARU: ``used`` is this instance's owned-region allocation; ``total`` is - that owned pool **plus the CXL device's free space** - (``cxl_pool.free_size`` from the shared resource manager). Anchoring - ``total`` to "owned pool + free" -- not just the owned pool -- keeps the - stock watermark from tripping while the device still has room: with free - CXL left the pool auto-expands instead of evicting, and only once the - device is full (``free_size`` 0 -> ``total`` collapses to the owned pool) - does eviction engage on this instance's own pages. Before the pool is up, - reports the configured capacity so watermark math stays sane. + MARU: ``used`` is this instance's owned-region allocation. ``total`` + depends on whether the pool may ``auto_expand``: + + - ``auto_expand`` (default): ``total`` is the owned pool **plus the CXL + device's free space** (``cxl_pool.free_size`` from the shared resource + manager). Anchoring to "owned pool + free" keeps the watermark from + tripping while the device still has room -- the pool auto-expands into + free CXL instead of evicting, and only once the device fills + (``free_size`` 0 -> ``total`` collapses to the owned pool) does + eviction engage on this instance's own pages. + - ``auto_expand`` off: the pool is hard-capped at ``pool_size_bytes``, so + ``total`` is the owned pool alone and eviction engages before it is + exhausted (device free is irrelevant -- the pool cannot grow into it). + + Before the pool is up, reports the configured capacity so watermark math + stays sane. ``free`` is cached across calls (``_last_cxl_free``): a get_stats RPC that does not deliver ``cxl_pool`` -- a transient timeout, or an older server @@ -1032,6 +1039,11 @@ def get_memory_usage(self) -> tuple[int, int]: return 0, self._config.pool_size_bytes used = regions["total_allocated_pages"] * handler.get_chunk_size() own_pool = regions["total_pool_size"] + if not self._config.auto_expand: + # Pool is hard-capped (no expansion): anchor the watermark to the + # owned pool so eviction engages before it is exhausted. Device + # free is irrelevant here -- the pool cannot grow into it. + return used, own_pool # A present cxl_pool (incl. free_size 0 when the device is genuinely # full) updates the cache; a missing one reuses the last-known free. free = stats.get("cxl_pool", {}).get("free_size") diff --git a/lmcache/v1/distributed/memory_manager/maru_memory_allocator.py b/lmcache/v1/distributed/memory_manager/maru_memory_allocator.py index 9d8388b0ca0..12ce63abb75 100644 --- a/lmcache/v1/distributed/memory_manager/maru_memory_allocator.py +++ b/lmcache/v1/distributed/memory_manager/maru_memory_allocator.py @@ -111,6 +111,7 @@ def init_layout( use_async_rpc=self._config.use_async_rpc, max_inflight=self._config.max_inflight, eager_map=self._config.eager_map, + auto_expand=self._config.auto_expand, ) handler = MaruHandler(maru_config) if not handler.connect(): diff --git a/tests/v1/distributed/maru_fakes.py b/tests/v1/distributed/maru_fakes.py index d28641dcfc1..ed46a9c8fc1 100644 --- a/tests/v1/distributed/maru_fakes.py +++ b/tests/v1/distributed/maru_fakes.py @@ -90,6 +90,7 @@ def __init__(self, chunk_size: int = 64): self.fail_exists = False self.fail_store = False self.fail_delete = False + self.cxl_free = 0 # device free bytes reported in get_stats' cxl_pool def batch_pin(self, keys: list[str]) -> list[bool]: if self.fail_pin: @@ -166,7 +167,8 @@ def get_stats(self) -> dict: "store_regions": { "total_pool_size": 16 * self.chunk_size, "total_allocated_pages": len(self.store_map), - } + }, + "cxl_pool": {"free_size": self.cxl_free}, } def close(self) -> None: @@ -234,6 +236,7 @@ def close(self) -> None: def make_maru_manager( chunk_size: int = 64, + auto_expand: bool = True, ) -> tuple[MaruL1Manager, FakeMaruHandler, FakeCxlAdapter]: """Build a MaruL1Manager wired to fresh fakes (post-init_layout state).""" cfg = L1ManagerConfig( @@ -244,6 +247,7 @@ def make_maru_manager( server_url="maru://localhost:5555", pool_size_bytes=1 << 20, instance_id="test", + auto_expand=auto_expand, ), ), write_ttl_seconds=600, diff --git a/tests/v1/distributed/test_maru_l1_manager.py b/tests/v1/distributed/test_maru_l1_manager.py index 5199ad0dc2c..d4910b86663 100644 --- a/tests/v1/distributed/test_maru_l1_manager.py +++ b/tests/v1/distributed/test_maru_l1_manager.py @@ -444,6 +444,27 @@ def test_get_memory_usage_before_init_reports_pool_size(): assert (used, total) == (0, 1 << 20) +def test_get_memory_usage_watermark_tracks_auto_expand(): + """total anchors to device fill when auto_expand is on, to the owned pool + when off -- so a hard-capped pool evicts before it is exhausted instead of + OOMing while the device still has (unusable) free space.""" + chunk = 64 + own_pool = 16 * chunk # FakeMaruHandler.get_stats total_pool_size + free = 100 * chunk + + # auto_expand on (default): total = owned pool + device free. + mgr_on, handler_on, _ = make_maru_manager(chunk_size=chunk) + handler_on.cxl_free = free + _seed(handler_on, _key(1)) # one allocated page + assert mgr_on.get_memory_usage() == (chunk, own_pool + free) + + # auto_expand off: device free is ignored; total is the owned pool alone. + mgr_off, handler_off, _ = make_maru_manager(chunk_size=chunk, auto_expand=False) + handler_off.cxl_free = free + _seed(handler_off, _key(1)) + assert mgr_off.get_memory_usage() == (chunk, own_pool) + + def test_is_key_evictable_tracks_local_staging(): manager, handler, _ = make_maru_manager() k = _key(1) From 8572c18c956d29b32b592acaa4634062f6ea731c Mon Sep 17 00:00:00 2001 From: seohui-XCENA Date: Tue, 7 Jul 2026 17:24:10 +0900 Subject: [PATCH 20/28] [MP][Maru] Revise the maru L1 design doc and fix allocator docstrings - Rewrite docs/design/v1/distributed/maru_l1.md to the l2_adapters house style (Overview + components + operation flow + configuration + limits); document the auto_expand knob and the single-device region bound, and drop the internals-heavy framing. - Correct MaruMemoryAllocator docstrings: page reclamation is driven by LMCache eviction via MaruL1Manager (delete / abort_alloc), not the allocator, so free/batched_free are no-ops. Signed-off-by: seohui-XCENA --- docs/design/v1/distributed/maru_l1.md | 548 +++++++++--------- .../memory_manager/maru_memory_allocator.py | 23 +- 2 files changed, 284 insertions(+), 287 deletions(-) diff --git a/docs/design/v1/distributed/maru_l1.md b/docs/design/v1/distributed/maru_l1.md index 41893ca21b0..e4ecc1eb5cf 100644 --- a/docs/design/v1/distributed/maru_l1.md +++ b/docs/design/v1/distributed/maru_l1.md @@ -1,307 +1,291 @@ -# Maru CXL Shared L1 (`MaruL1Manager`) +# Maru CXL Shared L1 Design -Design doc for the MP-mode L1 backend that stores the L1 tier in a -cross-instance **CXL shared pool** instead of local CPU DRAM. Covers the code in -`lmcache/v1/distributed/maru_l1_manager.py`, `l1_protocol.py`, and -`memory_manager/maru_memory_allocator.py`. +This document describes the `MaruL1Manager` L1 backend for LMCache MP mode: how it +stores the L1 tier in a cross-node CXL shared pool, the components involved, +the control-plane RPC flows, and the state-machine invariants it maintains. -This is a review-oriented summary: it states the **contracts and invariants** the -implementation must satisfy so the code can be checked against them. It is not an -exhaustive walkthrough. - -## 0. One-line summary - -**Maru replaces MP mode's L1 tier storage (CPU DRAM → a cross-instance shared CXL -pool); everything above L1 — L1↔L2 tiering (write-through / discard-evict / -promote-on-miss), the controllers, the L2 adapters, eviction — runs unchanged on -top.** LMCache never uses L1 in isolation, so Maru is not an "L1-only" cache -either. +## Overview -The whole integration reduces to one decision: +Maru replaces MP mode's L1 tier storage — CPU DRAM → a cross-node **CXL shared +pool** reached through the external [maru](https://github.com/xcena-dev/maru) +runtime. Multiple LMCache instances on separate nodes `mmap` the same physical +pool for zero-copy reads. A **single** `MaruServer` — one per shared CXL pool — is +the directory: it maps each key to its `(region, offset)` and holds the cross-node +`pin_count`, and every node sharing the pool connects to it. -> Isolate all Maru logic in a single class, `MaruL1Manager`, and run the existing -> stack unchanged on top of it through the `L1ManagerInterface` Protocol seam. +Control stays with LMCache and only the storage medium changes: the stock L1↔L2 +tiering (write-through, discard-evict, promote-on-miss), the controllers, the L2 +adapters, and eviction all run unchanged on top. `MaruL1Manager` executes those +decisions against the MaruServer directory and the CXL allocator; MaruServer itself +is passive (a directory + region broker with no policy). -## 1. Model — what changes and what does not - -### 1.1 Maru = one cross-instance shared L1 tier - -- **Shared medium**: a physical pool reached over a CXL switch. Multiple LMCache - instances `mmap` the same physical memory → **zero-copy**, no P2P transfer. -- **Shared index**: `MaruServer` holds the `key → (region, offset)` directory, a - per-region `kv_ref_count`, and a cross-node `pin_count`. -- **Access**: a retrieve pins + looks up the key on `MaruServer`, then - materializes a `MemoryObj` pointing at the resident CXL page (no copy). - -### 1.2 Control (LMCache) / Directory (MaruServer) split - -- **LMCache decides**: store / evict / promote / pin are decided by the existing - LMCache controllers. -- **MaruServer is passive**: it executes and records those decisions as a shared - directory + region broker. It holds no policy. -- LMCache's L1 manager therefore needs Maru RPC hooks for store/retrieve/pin/ - delete, and `MaruL1Manager` is the home of those hooks. - -### 1.3 Unchanged (byte-identical to `dev`) +``` +StoreController / PrefetchController / L1EvictionController + │ (L1ManagerInterface only) + ▼ + MaruL1Manager + RPC control surface + _pending_read/_pending_write staging + TTL sweeper + │ + ┌───────────┴───────────┐ + ▼ ▼ + MaruMemoryAllocator MaruHandler (RPC) + (CxlMemoryAdapter) │ + │ ▼ + │ MaruServer (passive directory) + │ key → (region, offset), kv_ref_count, + │ cross-node pin_count + ▼ │ + CXL shared pool ◄───────────┘ + (mmap, zero-copy views) +``` -- `L1Manager`, `L1MemoryManager`, `internal_api.py` — zero Maru code. -- `StoreController` / `PrefetchController` / `L1EvictionController` / the L2 - adapters / `L2EvictionController` — logic unchanged; only the `l1_manager` - parameter type is widened to the Protocol (§3). -- The data plane (`mmap` zero-copy, `metadata.address` bit-pack, `gpu_ops`). +## Goals + +- Store the MP L1 tier in a shared CXL pool so multiple instances hit the same KV + cache with zero-copy reads. +- Keep all tiering control (controllers, L2 cascade, eviction) with the stock + stack; swap only the L1 medium. +- Keep the L1 control-state-machine core (`l1_manager.py`, `l1_memory_manager.py`, + `internal_api.py`) and all controller/L2 logic byte-identical to `dev`. +- Isolate every Maru concern behind one class and one Protocol, so the stack does + not know which L1 backend it runs on. + +## Key Design Choice + +Maru is integrated as a **sibling** L1 manager — `MaruL1Manager` — beside the +existing `L1Manager`, rather than by modifying it. + +Maru is a **shared** tier: LMCache instances on separate nodes read and write the +same CXL pool, so two parts of the L1 control state can no longer be local: + +- **Membership** — whether a key is resident is a property of the shared pool, not + of any one node. A node that never wrote key K must still find it if a peer did. +- **Protection (pins)** — a node must not evict a page another node is still + reading, which needs a reference count visible across nodes. + +Both have to be authoritative in the shared directory (MaruServer) and reached over +RPC; a local dictionary and local locks cannot express them. So Maru needs its own +manager reimplementing the L1 control surface against MaruServer, while +`MaruMemoryAllocator` supplies the CXL medium. Crucially this moves only *where the +L1 state lives*, not the *decisions* — when to evict, pin, write through, or promote +all stay with the stock controllers, which see only the `L1ManagerInterface` and +never learn which manager backs it. + +The sibling form is a deliberate first step. Because Maru is external storage, it +keeps LMCache's core control logic untouched — `L1Manager`, the controllers, and the +tiering stack are all unchanged, and `MaruL1Manager` just satisfies the same +interface alongside them. The trade-off is that it reimplements the entire +L1-manager surface; once the shared-L1 approach is validated and accepted upstream, a +cleaner follow-up is to lift just the state-authority difference into a small +`L1StateBackend` interface inside the existing manager, so a shared tier like Maru +becomes a compact backend instead of a full sibling. + +## Components + +### `MaruL1Manager` (`maru_l1_manager.py`) + +The sibling L1 manager. Reimplements the `L1ManagerInterface` control surface over +MaruServer RPCs and owns: + +- `_pending_write: dict[ObjectKey, _PendingWrite]` — pages reserved between + `reserve_write` and `finish_write` (the CXL page, `is_temporary`, a write-TTL + deadline). +- `_pending_read: dict[ObjectKey, _PendingRead]` — staged reads between + `reserve_read` and the last `finish_read` (the zero-copy `MemoryObj`, a + `refcount`, the real server-pin count `pinned`, `is_temporary`, a read-TTL + deadline). +- The `MaruMemoryAllocator` (as `self._allocator`), the registered listeners, and a + daemon TTL sweeper thread. + +Membership and pins live remotely (MaruServer); these local dicts hold only +in-flight/staged state. A single non-reentrant lock serializes every public method. + +### `L1ManagerInterface` (`l1_protocol.py`) + +A structural `typing.Protocol` (`@runtime_checkable`). Both `L1Manager` and +`MaruL1Manager` satisfy it without inheritance; each method's docstring is the +behavioral contract (listener firing, PINNED retry, `extra_count` balance). +`StorageManager` selects the implementation in one line: -### 1.4 The three layers of MP L1 — reuse / reimplement / replace +```python +self._l1_manager: L1ManagerInterface = ( + MaruL1Manager(cfg) if maru_config else L1Manager(cfg) +) +``` -The recurring confusion is *which layer* "we keep the control logic" refers to. -MP L1 has three layers and the Maru seam cuts each differently: +The stock controllers call only this surface (no private access), so widening their +`l1_manager` parameter type from concrete `L1Manager` to `L1ManagerInterface` +(runtime-identical) lets the whole stack run over either implementation. -| Layer | Components | Stock local state | Role | Maru seam | -|---|---|---|---|---| -| **1 — Controllers / tiering** | `StoreController`, `PrefetchController`, `L1EvictionController`, `L2EvictionController`, L2 adapters, eviction `policy` (LRU) | policy's `OrderedDict` order (fed by listeners) | **decide** write-through / promote / eviction; act only on the L1Manager public surface | **reused** (only the `l1_manager` param type is widened; behavior-neutral) | -| **2 — L1Manager control state machine** | `L1Manager` | `_objects: dict[ObjectKey, L1ObjectState]` with per-key `write_lock`/`read_lock` (TTLLock) | authority for **this instance's** membership + lock state | **reimplemented** as sibling `MaruL1Manager` (MaruServer RPC instead of a local dict) | -| **3 — Memory manager / allocator** | `L1MemoryManager` + `MemoryAllocator` (DRAM / DevDax / GDS) | pinned DRAM/devdax/GDS buffer pool | the medium the KV bytes live in | **replaced** by `MaruMemoryAllocator` (CXL, zero-copy) | +### `MaruMemoryAllocator` (`memory_manager/maru_memory_allocator.py`) -"Control logic = stock" is a statement about **layer 1**, and it holds. But layer -2 cannot be reused as-is, which is why Maru is a **sibling** (layer-2 -reimplementation), not merely an allocator swap (layer-3 replacement). +The CXL medium allocator, alongside the other media allocators — a thin wrapper over +the external `maru_lmcache.CxlMemoryAdapter`. Two-phase init: the pool is typed by +the KV layout, so the `MaruHandler` and `CxlMemoryAdapter` are built on the first +`init_layout` (once at startup, before any store/retrieve). +`allocate`/`batched_allocate` hand out CXL pages and return `None` on OOM. Pages are +reclaimed through LMCache's eviction (executed as the manager's `delete`), not the +allocator, so `free` is a no-op. -### 1.5 Why an allocator swap (layer-3 only) is not enough +## Operation Flow -Keeping stock layer 2 (`L1Manager._objects`) and swapping only the allocator -breaks, because `_objects` is a **per-instance local index**: +All three flows run on the stock controllers; `MaruL1Manager` only implements the +`L1ManagerInterface` calls they make. -1. **Membership**: if instance B stores key `K`, this instance's local `_objects` - never learns about `K`, so `reserve_read(K)` misses — even though `K` is - physically present in the shared pool. Cross-instance read collapses. -2. **Pin**: stock `read_lock`/`write_lock` are **local** (a C++ atomic TTLLock), - invisible to other instances. Maru must prevent node A from evicting a page - node B is reading, which requires a **cross-node shared refcount** (MaruServer - `pin_count`). A local TTLLock cannot express that. +### Store (write-through) -GDS/DevDax work as allocator swaps because they are **private per-instance** -tiers, so `_objects` is complete and accurate. Maru is a **shared** tier, so the -authority for membership and pins must be **remote** (MaruServer), and layer 2 is -reimplemented over RPC. (The allocator interface has only -`allocate`/`free`/`get_memory_usage`/…; it has no "does K exist?" or "pin K" -control operation, so a layer-3 swap cannot express Maru's needs.) +``` +StoreController + ├─ reserve_write(keys) allocate a CXL page, stage in _pending_write + ├─ (D2H copy into the page) + ├─ finish_write(keys) create_store_handle + batch_store (register in directory) + │ → fire on_l1_keys_write_finished + └─ StoreController wakes, unsafe_read()s the zero-copy MemoryObj, copies it to L2 +``` -> "Same `MemoryObj`" means the *buffer type* is shared (so the data plane and -> allocator machinery are reusable); it does **not** mean the membership/pin -> control state may stay local. Those are independent. +A key another instance already registered is dup-skipped by the server and its page +auto-freed. -### 1.6 Stock local structures → Maru replacements +### Retrieve (hit, and miss → promote) -Layer 2's authority moves from local to remote (MaruServer); only in-flight state -stays local: +``` +StorageManager.lookup + ├─ reserve_read(keys) batch_pin + batch_retrieve + get_by_location (zero-copy) + │ stage in _pending_read → fire on_l1_keys_reserved_read + ├─ unsafe_read(keys) return the staged MemoryObj (no new pin), H2D copy to GPU + └─ finish_read(keys) batch_unpin the pins; drop the entry at refcount 0 + +miss → PrefetchController promote: + reserve_write → (load L2 bytes into the page) → finish_write_and_reserve_read + temporary (default): private page, no batch_store/pin, served once then freed + retained (hot-cache): batch_store, then re-resolve the authoritative page +``` -| Stock (layer 2, local) | Role | Maru replacement | Where | -|---|---|---|---| -| `_objects` dict | membership authority | MaruServer directory `key→region/offset` + `kv_ref_count` | **remote (shared)** | -| `L1ObjectState.memory_obj` | buffer handle | CXL page `MemoryObj` (`get_by_location`, zero-copy view) | local view | -| `write_lock`/`read_lock` (TTLLock) | local eviction/write protection | cross-node pin via MaruServer `pin_count` (delete refuses with PINNED) | **remote (shared)** | -| (no in-flight-write tracking) | reserve_write → finish_write window | `_pending_write` (mem_obj, `is_temporary`, deadline) + write-TTL | local | -| `L1ObjectState.is_temporary` | promote lifetime | `_PendingRead.is_temporary` (propagated from `_pending_write` by promote) | local | -| (no staged-read tracking) | reserve_read → finish_read window | `_pending_read` (`_PendingRead` refcount + `pinned` + `is_temporary`) + read-TTL | local | -| `L1MemoryManager` (DRAM pool) | medium | `MaruMemoryAllocator` (CXL pool) | local | +### Eviction (watermark → delete → PINNED retry) -**Everything Maru keeps local is in-flight/staged side-channel state; all -authority (membership, pins) is remote.** Layer 1 and the data plane are -untouched. TTL reclaim of the side-channel reuses the stock `write_ttl`/`read_ttl` -config values, but sweeps a monotonic deadline on the pending entries instead of a -TTLLock. +``` +L1EvictionController (each cycle) + ├─ get_memory_usage() (used, total) over the CXL pool + ├─ if over watermark: the LRU policy picks victims, filtered by is_key_evictable + └─ delete(victims) directory delete; MaruServer refuses PINNED keys (another + node is reading). Only truly-deleted keys fire + on_l1_keys_deleted_by_manager; PINNED keys stay in the + policy and retry next cycle. +``` -## 2. The seam — sibling + structural Protocol +Discard-on-evict is safe because write-through already placed the data in L2. -`StorageManager` picks the L1 manager in one line at construction: +## State-Machine Invariants -```python -self._l1_manager: L1ManagerInterface = ( - MaruL1Manager(cfg) if maru_config else L1Manager(cfg) -) -``` +1. **A key is in at most one of `_pending_write` / `_pending_read`**, mirroring + stock's per-key write/read lock exclusion. `reserve_write` and + `finish_write_and_reserve_read` reject a key already staged the other way, and + `reserve_read` excludes mid-write keys from its pin/stage step — a peer may have + registered the same key, so the pin would otherwise succeed and double-stage, + stranding the in-flight write. -Both implementations satisfy `L1ManagerInterface`, so the stack above does not -know which one it has. - -**Three new files:** - -- `l1_protocol.py` — the structural `typing.Protocol` `L1ManagerInterface` - (`@runtime_checkable`). `L1Manager` and `MaruL1Manager` both satisfy it - **without inheritance**. Each method's docstring states its behavioral contract - (when listeners fire, the PINNED retry, `extra_count` balance). -- `maru_l1_manager.py` — `MaruL1Manager`. Implements the whole RPC control - surface and owns `MaruMemoryAllocator` directly as `self._allocator`. No - separate dispatch layer. -- `memory_manager/maru_memory_allocator.py` — `MaruMemoryAllocator` (same - directory as the devdax/gds allocators). Thin wrapper over the external - `maru_lmcache.CxlMemoryAdapter`; connects the handler and creates the adapter on - the first `init_layout`. `free` is a no-op (CXL page lifecycle is MaruServer's). - -The stock controllers call only the `L1ManagerInterface` surface (no private -access — verified), so widening each controller's `l1_manager` parameter type from -concrete `L1Manager` to `L1ManagerInterface` (runtime-identical) lets the whole -stack run over `MaruL1Manager`. The only remaining Maru-specific concrete -references are the L1-selection branch in `StorageManager` and the maru-only gate -in its `register_kv_layout` wrapper — both legitimately Maru-only. - -## 3. What `MaruL1Manager` implements - -The `L1ManagerInterface` methods, reimplemented over RPC. "Fire" means invoking -the registered listeners' `on_l1_keys_*` callbacks (the eviction LRU and the store -controller consume these). - -| Method | Implemented behavior | -|---|---| -| `reserve_read(keys, extra_count)` | `batch_pin` + `batch_retrieve` + `get_by_location` (zero-copy) → stage in `_pending_read`. Takes `1 + extra_count` server pins per key (MLA + TP>1). A key mid-write on this instance is **excluded** (stays `KEY_NOT_READABLE`, see §4 invariant 1). Fires `on_l1_keys_reserved_read`. | -| `unsafe_read` / `finish_read` | Return the staged object / release `1 + extra_count` holds. Each release drops `refcount` and unpins `min(released, pinned)` server pins (§4 invariant 2). At `refcount == 0` a **temporary** entry frees its private page (`on_l1_keys_deleted_by_manager`); a directory read just drops the entry. | -| `reserve_write` / `finish_write` | Allocate a CXL page + stage in `_pending_write` (`is_temporary` + write-TTL); fire `on_l1_keys_reserved_write`. / `create_store_handle` + `batch_store` (directory register); fire `on_l1_keys_write_finished`. In `mode="new"`, a key already staged locally **or** already in the directory returns `KEY_NOT_WRITABLE` (cross-instance dedup). | -| `finish_write_and_reserve_read` (promote) | Branches on the staged `is_temporary` (§4 invariant 3). **temporary** (the default prefetch policy): no `batch_store`, no pin — the loaded local page moves straight to `_pending_read`; freed at `finish_read == 0`. **retained**: `batch_store`, then re-resolve the authoritative directory page with pins (a dup-skip auto-freed our page). Both fire only `on_l1_keys_finish_write_and_reserve_read` — **never** `on_l1_keys_write_finished` (that would make the store controller re-store the promoted key to L2). | -| `delete(keys)` | Directory delete. A local read-hold → `KEY_IS_LOCKED`; a cross-node pinned key → MaruServer refuses with **PINNED**. Only actually-deleted keys fire `on_l1_keys_deleted_by_manager` (§4 invariant 4). | -| `is_key_evictable(key)` | False if locally pinned (`_pending_read`); lock-free. Cross-node pins are re-checked authoritatively by `delete`. | -| `touch_keys` / `register_listener` | `register_listener` only stores the listener; each method fires its own event. `touch_keys` fires `on_l1_keys_accessed`. | -| `get_memory_usage()` | `handler.get_stats` → `(used, total)`. `used` = owned-pool allocated; `total` = owned pool + CXL device free (`cxl_pool.free_size`), so the eviction watermark tracks whole-device fill, not just the owned pool. The last known free is cached to survive a transient `get_stats` RPC failure. Before init, returns `(0, configured pool size)` to avoid a 0-division in the controller watermark. | -| `get_l1_memory_desc()` | `None` — the shared pool has no single registrable region (§5 L2 limit). | -| `register_kv_layout(...)` | Maru-only: binds the KV layout to `MaruMemoryAllocator` (pool bring-up). `num_object_groups > 1` is rejected by the `StorageManager.register_kv_layout` wrapper (§6). | -| `clear` / `memcheck` / `report_status` / `close` | Delegated under the lock (§4). `clear` empties the local side-channel and releases staged pins but does **not** touch shared data (a shared tier must not be wiped by one node). `close` stops the TTL sweeper. | - -Eviction and LRU stay with the stock stack: `L1EvictionController` decides, and -recency lives in the eviction **policy** (fed by the listener events). -`MaruL1Manager` keeps no local object/LRU dict. - -## 4. Invariants (check the implementation against these) - -These are the contracts a reviewer should verify; the two correctness bugs found -in review were both violations of invariants 1 and 2. - -1. **A key is in at most one of `_pending_write` / `_pending_read`.** This mirrors - stock `L1Manager`'s per-key `write_lock`/`read_lock` mutual exclusion (a - mid-write key is not readable). Enforced in both directions: `reserve_write` - and `finish_write_and_reserve_read` reject a key already staged the other way, - and `reserve_read` **excludes** mid-write keys from the pin/stage step (a peer - may have registered the same key, so the pin would otherwise succeed and double - stage — stranding the in-flight write). - -2. **Pin accounting: `pinned` tracks real server pins separately from - `refcount`** (`0 <= pinned <= refcount`). Release paths unpin `pinned`, not - `refcount`. A pure temporary stage has `pinned == 0` (a private local page, no - server pin); a temporary that absorbs an overlapping `reserve_read`'s pins - records them in `pinned` so they are released, not leaked. `N` reserves take - `N` pins and `N` finishes release `N`; `finish_read` never releases more than - it holds (over-release would corrupt the server `pin_count`). +2. **`pinned` tracks real server pins separately from `refcount`** + (`0 <= pinned <= refcount`); release paths unpin `pinned`, not `refcount`. A pure + temporary stage has `pinned == 0`; a temporary that absorbs an overlapping + `reserve_read`'s pins records them so they are released, not leaked. `finish_read` + never releases more than it holds. 3. **temporary vs retained promote.** The default prefetch policy marks every - promote **temporary**: the page is private staging, never registered in the - directory, and discarded after one read (the shared pool is populated only by - store write-through). Only a hot-cache policy marks a promote **retained**, - which registers it in the directory. - -4. **PINNED cross-node retry (no controller change).** MaruServer refuses to - delete any key with `pin_count > 0` (PINNED). `MaruL1Manager.delete` fires - `on_l1_keys_deleted_by_manager` only for keys it **actually** deleted, so a - PINNED-refused key stays in the eviction policy and is retried next cycle. This - relies on the stock LRU contract that a key is removed from the policy only via - the `on_keys_removed` event, never popped by `get_eviction_actions`. - -### Thread safety - -A single non-reentrant `threading.Lock` (`_maru_l1_synchronized`) serializes every -public method that touches the side-channel or allocator. The decorator is typed -with `ParamSpec`/`Concatenate` so wrapped signatures survive the Protocol -conformance check. - -### TTL sweeper + crash-recovery premises - -An entry whose reserve→finish flow is interrupted (an *orphan*) cannot be -distinguished from an in-progress one by inspecting the pending dicts, so **time -(TTL) is the only abandonment signal** (reusing the stock `write_ttl`/`read_ttl` -values). A daemon sweeper thread (scanning under the lock) reclaims both: - -- `_pending_write`: on write-TTL expiry, return the page to the owner's local - free-list (`abort_alloc`) + pop. An unregistered page is invisible to the - server, so client-side reclaim is the only option. -- `_pending_read`: on read-TTL expiry, `batch_unpin × remaining refcount` + pop. - -A late `finish_read`/`unsafe_read` after expiry sees `KEY_NOT_EXIST` → retrieve -returns False → recompute (the same failure path as a stock TTL expiry). Double -unpin is impossible under the serializing lock. - -**Premises for a full client crash** (the MP-server process dies, so the local -sweeper cannot run): - -- read pins are assumed to be released by the **maru side** tracking each client's - pin set and dropping them on disconnect; -- an in-flight write page is assumed to be covered by **region owner-release**. - -These two mechanisms live outside `MaruL1Manager` (on the Maru/MaruServer side) -and are premises of this design, not guarantees provided by this PR. - -## 5. Scope and known limits - -- **Copy-type L2 adapters only.** `get_l1_memory_desc()` returns `None`, so - copy-type adapters (aerospike, dax, fs, fs_native, hfbucket, mock, resp, s3, - plugin, raw_block, mooncake-TCP) work fully, while registered/RDMA-type adapters - (nixl_store, nixl_store_dynamic, mooncake-RDMA) and p2p are **rejected at - startup** (§6 guards; an allowlist, so a new registered adapter fails safe). - Registered-L2 support needs an allocator region accessor + per-region descriptor - and is follow-up work. -- **Single model / single object group.** `MaruMemoryAllocator` fixes the pool to - a single layout on the first `init_layout`, so a different model (different - shapes/dtypes/fmt/chunk) or a hybrid model (`num_object_groups > 1`) fails fast - with a `ValueError`. Multiple instances of the *same* model are fine. Mixed - layouts are follow-up work. -- **Single-device pool bound.** A pool is one contiguous allocation within one DAX - device, so `pool_size_bytes` cannot exceed a single CXL device's capacity. -- **LRU is a per-instance local view.** The policy is fed only by this instance's - listener events, so it does not see other nodes' access recency. A known - approximation for a shared cache. -- **Cross-owner reclaim.** A per-key `delete` returns the page only to the - *owner's* local free-list. Evicting a key another instance owns removes the - directory entry but does not reclaim that owner's page. Region-level return to - the shared pool is an owner-only concern outside this PR's critical path. (This - and the expand-vs-evict watermark interaction are the two multi-instance - consistency issues tracked for a separate team review.) - -## 6. File changes (this PR) - -**New** - -- `l1_protocol.py` — `L1ManagerInterface` + the `L1OperationResult` alias. -- `maru_l1_manager.py` — `MaruL1Manager` (~1050 lines). -- `memory_manager/maru_memory_allocator.py` — `MaruMemoryAllocator`. -- Tests under `tests/v1/distributed/` (fake-maru harness + manager unit + - conformance + config guard + control integration). - -**Modified (all additive / behavior-neutral)** - -- `store_controller.py` / `prefetch_controller.py` / `eviction_controller.py` — - `l1_manager` parameter type widened to `L1ManagerInterface`. -- `storage_manager.py` — the L1-selection branch (`MaruL1Manager` when - `maru_config` is set) + the maru-only `register_kv_layout` wrapper. -- `config.py` — the `maru_config` field, CLI, and validation (§ startup guards). -- `lmcache_driven_transfer.py` — the `register_kv_layout` engine hook (a no-op for - the default backend), wired in `register_kv_cache` and wrapped in try/except so a - failure closes the cache context. - -**Unchanged (byte-identical to `dev`)** - -- `l1_manager.py`, `l1_memory_manager.py`, `internal_api.py`, all controller - logic, all L2 adapters. - -### Startup guards (`validate_storage_manager_config`) - -Unsupported combinations are rejected at startup: maru + gds/devdax (mutual -exclusion), maru + `skip_l1` store policy, maru + registered/RDMA-type L2, maru + -p2p, and maru + engine_driven/auto transfer mode (maru requires -`--supported-transfer-mode lmcache_driven`). Without these guards the failures are -silent or surface as an `AttributeError` deep in the request path. - -## 7. Status - -Implemented on branch `maru-mp-l1`: the sibling `MaruL1Manager` + `l1_protocol.py` -+ `MaruMemoryAllocator`, with `l1_manager.py` / `l1_memory_manager.py` / -`internal_api.py` kept byte-identical to `dev`. The full Maru test suite passes on -GPU (`CUDA_VISIBLE_DEVICES=1`). - -A later phase may extract an `L1StateBackend` seam so Maru becomes a small backend -rather than a full sibling; that is separate follow-up work after this sibling -merges. + promote temporary — private staging, never registered, discarded after one read + (the shared pool is populated only by store write-through). Only a hot-cache + policy marks a promote retained and registers it in the directory. + +4. **PINNED cross-node retry.** MaruServer refuses to delete a key with + `pin_count > 0`; `delete` fires the deleted-listener only for keys it actually + removed, so a refused key stays in the eviction policy and retries. This relies on + the stock LRU contract that a key leaves the policy only via `on_keys_removed`. + +## Orphan Reclaim and Crash Recovery + +A reserve→finish flow that stalls leaves an orphan the pending dicts cannot +distinguish from an in-progress one, so time is the only abandonment signal (reusing +the stock `write_ttl`/`read_ttl` values). A daemon sweeper reclaims both under the +lock: an expired `_pending_write` page returns to the owner's free-list +(`abort_alloc`); an expired `_pending_read` unpins its remaining refcount. A late +`finish_read`/`unsafe_read` then sees `KEY_NOT_EXIST` and recomputes — the same path +as a stock TTL expiry. + +A full client crash (the MP-server process dies, so the local sweeper cannot run) is +out of scope for `MaruL1Manager` and relies on maru-side mechanisms: read pins are +released when the maru side drops a disconnected client's pin set, and an in-flight +write page is covered by region owner-release. These are premises of this design, +not guarantees provided here. + +## Configuration + +Maru L1 is enabled via CLI flags that build a `MaruL1Config` (mutually exclusive +with the pinned-DRAM / DevDax / GDS tiers): + +- `--maru-server-url` — MaruServer endpoint (`maru://host:port`, rewritten to + `tcp://`). +- `--maru-pool-size-gb` — CXL pool size to request (> 0). A pool is a single region + that must fit within one CXL device's free space (a region cannot span devices); + with `auto_expand` the pool then grows across devices by adding regions. +- `--maru-instance-id` — stable client id for ownership tracking (auto-generated if + unset). +- `--maru-auto-expand` / `--no-maru-auto-expand` — whether the owned pool auto-expands + into free CXL when it fills (default on). + +`MaruL1Config` fields: + +| Field | Default | Purpose | +|---|---|---| +| `server_url` | (required) | MaruServer endpoint. | +| `pool_size_bytes` | (required) | CXL pool size requested from MaruServer. | +| `instance_id` | auto | Stable client id for ownership tracking. | +| `timeout_ms` | 5000 | RPC timeout. | +| `use_async_rpc` | True | Use the async RPC path. | +| `max_inflight` | 64 | Max concurrent in-flight RPCs. | +| `eager_map` | True | Eagerly mmap peer regions for cross-node zero-copy reads. | +| `auto_expand` | True | Grow the owned pool into free CXL when full; off hard-caps it at `pool_size_bytes`. | + +`register_kv_layout` binds the KV layout to the allocator (pool bring-up) on the +first `register_kv_cache`; it is wired only on the `lmcache_driven` transfer path. + +### Startup guards + +`validate_storage_manager_config` rejects unsupported combinations at startup rather +than failing deep in the request path: maru + gds/devdax (mutual exclusion), maru + +`skip_l1` store policy, maru + registered/RDMA-type L2 (only copy-type L2 is +allowlisted, since `get_l1_memory_desc()` is `None`), maru + p2p, and maru + +engine_driven/auto transfer mode (maru requires +`--supported-transfer-mode lmcache_driven`). + +## Capacity and Eviction + +`get_memory_usage()` returns `(used, total)` over the CXL pool: `used` is the +owned-pool allocation and `total` is the owned pool plus the CXL device free +(`cxl_pool.free_size` from `get_stats`), so the eviction watermark tracks +whole-device fill rather than only the owned pool — the pool auto-expands into free +CXL before evicting. With `auto_expand` off, `total` is the owned pool alone (which +is hard-capped at `pool_size_bytes`), so eviction engages before it is exhausted. + +## Current Limits + +- **Copy-type L2 adapters only.** Registered/RDMA-type adapters and p2p are rejected + at startup; registered-L2 support needs an allocator region accessor + per-region + descriptor (follow-up). +- **Single model / single object group.** The pool is fixed to the first layout; a + different or hybrid (`num_object_groups > 1`) layout fails fast. Multiple instances + of the same model are fine. +- **LRU is a per-instance local view** — fed only by this instance's listener + events, so other nodes' access recency is not seen. A known approximation for a + shared cache. +- **Cross-owner reclaim.** A per-key `delete` returns the page only to the owner's + local free-list; evicting a key another instance owns removes the directory entry + but not that owner's page. Region-level return is owner-only, outside this critical + path. + +## References + +- Implementation: `lmcache/v1/distributed/maru_l1_manager.py` +- Protocol: `lmcache/v1/distributed/l1_protocol.py` +- Allocator: `lmcache/v1/distributed/memory_manager/maru_memory_allocator.py` +- Config: `lmcache/v1/distributed/config.py` (`MaruL1Config`) +- User docs: `docs/source/mp/configuration.rst` (Maru CXL Shared L1 section) +- External runtime: [maru](https://github.com/xcena-dev/maru) diff --git a/lmcache/v1/distributed/memory_manager/maru_memory_allocator.py b/lmcache/v1/distributed/memory_manager/maru_memory_allocator.py index 12ce63abb75..375112aba71 100644 --- a/lmcache/v1/distributed/memory_manager/maru_memory_allocator.py +++ b/lmcache/v1/distributed/memory_manager/maru_memory_allocator.py @@ -6,8 +6,10 @@ imported lazily in :meth:`init_layout` so this module stays importable without the maru runtime. Two-phase init: the CXL pool is typed by the KV layout, so the ``MaruHandler`` / ``CxlMemoryAdapter`` are built on the first layout -registration, not at construction. ``free`` is a no-op -- page lifecycle is -owned by MaruServer (pin/unpin/delete), not LMCache refcounts. +registration, not at construction. ``free`` is a no-op -- a CXL page is not +reclaimed through the allocator or by dropping a local ``MemoryObj``; reclamation +is driven by ``MaruL1Manager`` (``delete`` for a registered page, ``abort_alloc`` +for an unregistered one). Single-model per instance: the pool is fixed to the first layout; a different layout is rejected. TODO(maru-multi-model): partition the pool per layout key. @@ -216,7 +218,13 @@ def batched_allocate( ) def free(self, memory_obj: MemoryObj, allocator_type: Optional[str] = None) -> None: - """No-op: CXL page lifecycle is owned by MaruServer.""" + """No-op: a CXL page is not reclaimed through the allocator. + + Reclamation is driven explicitly by ``MaruL1Manager``: ``delete`` for a + registered page (eviction; the shared directory drops the key and the + page returns to its owner), ``abort_alloc`` for a page allocated but + never registered. + """ def batched_free( self, @@ -224,7 +232,11 @@ def batched_free( allocator_type: Optional[str] = None, update_stats: bool = True, ) -> None: - """No-op: CXL page lifecycle is owned by MaruServer.""" + """No-op: CXL pages are not reclaimed through the allocator. + + See :meth:`free`; reclamation is driven by ``MaruL1Manager`` + (``delete`` / ``abort_alloc``), not by dropping the local ``MemoryObj``. + """ # maru-specific surface (called by MaruL1Manager) @@ -279,7 +291,8 @@ def close(self) -> None: Idempotent: safe to call when the pool was never brought up. Both the adapter and handler are cleared, so any later use raises. Shared CXL - pages are not freed here -- their lifecycle is owned by MaruServer. + pages are not freed here -- they belong to the shared pool and are + reclaimed via ``MaruL1Manager.delete``, not by local teardown. """ if self._cxl_adapter is not None: self._cxl_adapter.close() From f8f0cd8afc12a65704b6853e79a3815d6d05517e Mon Sep 17 00:00:00 2001 From: seohui-XCENA Date: Tue, 7 Jul 2026 19:04:36 +0900 Subject: [PATCH 21/28] [MP][Maru] Consolidate and trim the L1 backend test suite - Merge the maru-specific test files into one unit file (test_maru_l1_manager.py): fold in the allocator, config/startup-guard, and control-integration tests, and drop cases already covered by the shared L1Manager conformance suite (notably the notification tests: ~14 -> 4). - Absorb test_l1_protocol.py (structural conformance + shared helpers) into test_l1_manager_conformance.py. - Delete the now-empty test_maru_memory_allocator.py, test_l1_config_maru.py, test_maru_integration.py, and test_l1_protocol.py. Signed-off-by: seohui-XCENA --- tests/v1/distributed/test_l1_config_maru.py | 150 ---- .../test_l1_manager_conformance.py | 102 +++ tests/v1/distributed/test_l1_protocol.py | 98 --- tests/v1/distributed/test_maru_integration.py | 136 --- tests/v1/distributed/test_maru_l1_manager.py | 831 ++++++++---------- .../distributed/test_maru_memory_allocator.py | 160 ---- 6 files changed, 459 insertions(+), 1018 deletions(-) delete mode 100644 tests/v1/distributed/test_l1_config_maru.py delete mode 100644 tests/v1/distributed/test_l1_protocol.py delete mode 100644 tests/v1/distributed/test_maru_integration.py delete mode 100644 tests/v1/distributed/test_maru_memory_allocator.py diff --git a/tests/v1/distributed/test_l1_config_maru.py b/tests/v1/distributed/test_l1_config_maru.py deleted file mode 100644 index 8e74c835352..00000000000 --- a/tests/v1/distributed/test_l1_config_maru.py +++ /dev/null @@ -1,150 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -"""Tests for Maru L1 config parsing and startup guards.""" - -# Third Party -import pytest - -# First Party -from lmcache.v1.distributed.config import ( - EvictionConfig, - GdsL1Config, - L1ManagerConfig, - L1MemoryManagerConfig, - MaruL1Config, - StorageManagerConfig, - parse_args, -) -from lmcache.v1.distributed.l2_adapters.config import L2AdaptersConfig -from lmcache.v1.distributed.l2_adapters.mock_l2_adapter import MockL2AdapterConfig - - -def _args(*extra: str, l1_size_gb: str = "1") -> list[str]: - """Minimal required flags (eviction policy + L1 size) plus extras.""" - return [ - "--eviction-policy", - "LRU", - "--l1-size-gb", - l1_size_gb, - "--no-l1-use-lazy", - *extra, - ] - - -def test_no_maru_flags_leaves_maru_config_none(): - cfg = parse_args(_args()) - assert cfg.l1_manager_config.memory_config.maru_config is None - - -def test_maru_flags_build_maru_config(): - cfg = parse_args( - _args( - "--maru-server-url", - "maru://localhost:9000", - "--maru-pool-size-gb", - "8", - "--maru-instance-id", - "node-a", - ) - ) - maru = cfg.l1_manager_config.memory_config.maru_config - assert isinstance(maru, MaruL1Config) - assert maru.server_url == "maru://localhost:9000" - assert maru.pool_size_bytes == 8 * (1 << 30) - assert maru.instance_id == "node-a" - - -def test_maru_without_pool_size_raises(): - with pytest.raises(ValueError): - parse_args(_args("--maru-server-url", "maru://localhost:9000")) - - -# --------------------------------------------------------------------------- -# C9: validate_storage_manager_config rejects unsupported maru combinations. -# --------------------------------------------------------------------------- - - -def _maru_memory(**overrides) -> L1MemoryManagerConfig: - return L1MemoryManagerConfig( - size_in_bytes=0, - use_lazy=False, - maru_config=MaruL1Config( - server_url="maru://localhost:5555", - pool_size_bytes=1 << 20, - instance_id="t", - ), - **overrides, - ) - - -def _maru_sm_config( - *, memory=None, gds=None, store_policy="default", adapters=None -) -> StorageManagerConfig: - return StorageManagerConfig( - l1_manager_config=L1ManagerConfig( - memory_config=memory if memory is not None else _maru_memory(), - gds_l1_config=gds, - write_ttl_seconds=600, - read_ttl_seconds=300, - ), - eviction_config=EvictionConfig(eviction_policy="LRU"), - l2_adapter_config=L2AdaptersConfig(adapters=adapters or []), - store_policy=store_policy, - ) - - -def test_maru_plus_copy_l2_is_accepted(): - # A copy-type L2 (mock) needs no registerable region -> allowed. - _maru_sm_config( - adapters=[MockL2AdapterConfig(max_size_gb=1.0, mock_bandwidth_gb=1.0)] - ) - - -def test_maru_rejects_gds_l1(): - with pytest.raises(ValueError, match="gds"): - _maru_sm_config( - gds=GdsL1Config(file_location="/tmp/gds", size_in_bytes=1 << 20) - ) - - -def test_maru_rejects_devdax_l1(): - with pytest.raises(ValueError, match="devdax"): - _maru_sm_config(memory=_maru_memory(devdax_path="/dev/dax0.0", shm_name="")) - - -def test_maru_rejects_skip_l1_store_policy(): - with pytest.raises(ValueError, match="skip_l1"): - _maru_sm_config(store_policy="skip_l1") - - -def test_maru_rejects_registered_l2(monkeypatch): - # Simulate a registered/RDMA adapter via the shared region classifier. - monkeypatch.setattr( - "lmcache.v1.distributed.config._requires_single_l1_memory_region", - lambda adapter_config: "nixl_store", - ) - with pytest.raises(ValueError, match="registerable"): - _maru_sm_config( - adapters=[MockL2AdapterConfig(max_size_gb=1.0, mock_bandwidth_gb=1.0)] - ) - - -def test_maru_rejects_non_lmcache_driven_transfer(): - """maru requires the LMCache-driven transfer path (rejects engine/auto).""" - # First Party - from lmcache.v1.mp_observability.config import ObservabilityConfig - from lmcache.v1.multiprocess.config import ( - CoordinatorConfig, - HTTPFrontendConfig, - MPServerConfig, - ) - from lmcache.v1.multiprocess.http_server import run_http_server - - with pytest.raises(ValueError, match="lmcache_driven"): - run_http_server( - http_config=HTTPFrontendConfig(), - mp_config=MPServerConfig(supported_transfer_mode="engine_driven"), - storage_manager_config=_maru_sm_config(), - obs_config=ObservabilityConfig(), - coordinator_config=CoordinatorConfig(url=""), - ) diff --git a/tests/v1/distributed/test_l1_manager_conformance.py b/tests/v1/distributed/test_l1_manager_conformance.py index 0c7b2d65867..f14d0f027c6 100644 --- a/tests/v1/distributed/test_l1_manager_conformance.py +++ b/tests/v1/distributed/test_l1_manager_conformance.py @@ -12,6 +12,9 @@ maru side. """ +# Standard +import inspect + # Third Party import pytest import torch @@ -22,6 +25,8 @@ from lmcache.v1.distributed.config import L1ManagerConfig, L1MemoryManagerConfig from lmcache.v1.distributed.error import L1Error from lmcache.v1.distributed.l1_manager import L1Manager + from lmcache.v1.distributed.l1_protocol import L1ManagerInterface + from lmcache.v1.distributed.maru_l1_manager import MaruL1Manager # Local from .maru_fakes import RecordingListener, make_maru_manager @@ -252,3 +257,100 @@ def test_listener_fires_across_lifecycle(manager): assert manager.delete([k])[k] == L1Error.SUCCESS assert [k] in rec.kinds("deleted_by_manager") + + +# ========================================================================= +# structural conformance: Protocol method-set / signature drift guards +# +# A ``runtime_checkable`` ``issubclass`` only checks method *names*; these bind +# the Protocol's method set and call shapes to both concrete backends so +# signature drift fails CI. (Absorbed from the former protocol structural +# tests and the maru unit file's two interface tests.) +# ========================================================================= + + +def _interface_methods() -> set[str]: + """Public method names declared on the Protocol.""" + return { + name + for name, val in vars(L1ManagerInterface).items() + if callable(val) and not name.startswith("_") + } + + +def _l1_manager_methods() -> set[str]: + """Public method names on the concrete L1Manager.""" + return { + name + for name in dir(L1Manager) + if not name.startswith("_") and callable(getattr(L1Manager, name)) + } + + +def _unwrap(fn): + """Recover the wrapped function from ``l1_mgr_synchronized``. + + That decorator is a plain closure without ``functools.wraps``, so + ``inspect.signature`` would otherwise see ``(self, *args, **kwargs)``. The + original function is the ``func`` free variable of the wrapper closure. + """ + code = getattr(fn, "__code__", None) + while code is not None and fn.__closure__ and "func" in code.co_freevars: + fn = fn.__closure__[code.co_freevars.index("func")].cell_contents + code = getattr(fn, "__code__", None) + return fn + + +def _params(fn) -> list[tuple[str, object, object]]: + """Parameter (name, kind, default) list excluding ``self``. + + Annotations are excluded on purpose: L1Manager leaves some param/return + annotations off, so only the call shape (names/kinds/defaults) is compared. + """ + return [ + (p.name, p.kind, p.default) + for p in inspect.signature(fn).parameters.values() + if p.name != "self" + ] + + +def test_interface_matches_l1_manager_surface(): + """The Protocol declares exactly L1Manager's public method set.""" + methods = _interface_methods() + assert methods == _l1_manager_methods() + assert len(methods) == 17 # tripwire against wholesale drift + + +def test_signatures_match_l1_manager(): + """Each Protocol method's call shape matches L1Manager's.""" + for name in _interface_methods(): + proto = _params(getattr(L1ManagerInterface, name)) + impl = _params(_unwrap(getattr(L1Manager, name))) + assert proto == impl, name + + +def test_l1_manager_conforms_to_interface(): + """The stock L1Manager satisfies the interface (structural, no instance).""" + assert issubclass(L1Manager, L1ManagerInterface) + + +def test_incomplete_class_does_not_conform(): + """A class missing interface methods is rejected (negative control).""" + + class Partial: + def close(self) -> None: ... + + assert not issubclass(Partial, L1ManagerInterface) + + +def test_maru_conforms_to_l1_manager_interface(): + """MaruL1Manager satisfies the interface (structural, binds the sibling).""" + assert issubclass(MaruL1Manager, L1ManagerInterface) + + +def test_maru_signatures_match_interface(): + """Each Protocol method's call shape matches MaruL1Manager's.""" + for name in _interface_methods(): + proto = _params(getattr(L1ManagerInterface, name)) + impl = _params(_unwrap(getattr(MaruL1Manager, name))) + assert proto == impl, name diff --git a/tests/v1/distributed/test_l1_protocol.py b/tests/v1/distributed/test_l1_protocol.py deleted file mode 100644 index 4cac18a2b6c..00000000000 --- a/tests/v1/distributed/test_l1_protocol.py +++ /dev/null @@ -1,98 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -"""Structural conformance tests for ``L1ManagerInterface``. - -These bind the Protocol to ``L1Manager`` so that method-set or signature drift -between the two fails CI (the guard the sibling design relies on -- a -``runtime_checkable`` ``issubclass`` alone only checks method *names*). -""" - -# Standard -import inspect - -# Third Party -import pytest - -try: - # First Party - from lmcache.v1.distributed.l1_manager import L1Manager - from lmcache.v1.distributed.l1_protocol import L1ManagerInterface -except ImportError: - pytest.skip( - "L1Manager / L1ManagerInterface unavailable (native ext missing)", - allow_module_level=True, - ) - - -def _interface_methods() -> set[str]: - """Public method names declared on the Protocol.""" - return { - name - for name, val in vars(L1ManagerInterface).items() - if callable(val) and not name.startswith("_") - } - - -def _l1_manager_methods() -> set[str]: - """Public method names on the concrete L1Manager.""" - return { - name - for name in dir(L1Manager) - if not name.startswith("_") and callable(getattr(L1Manager, name)) - } - - -def _unwrap(fn): - """Recover the wrapped function from ``l1_mgr_synchronized``. - - That decorator is a plain closure without ``functools.wraps``, so - ``inspect.signature`` would otherwise see ``(self, *args, **kwargs)``. The - original function is the ``func`` free variable of the wrapper closure. - """ - code = getattr(fn, "__code__", None) - while code is not None and fn.__closure__ and "func" in code.co_freevars: - fn = fn.__closure__[code.co_freevars.index("func")].cell_contents - code = getattr(fn, "__code__", None) - return fn - - -def _params(fn) -> list[tuple[str, object, object]]: - """Parameter (name, kind, default) list excluding ``self``. - - Annotations are excluded on purpose: L1Manager leaves some param/return - annotations off, so only the call shape (names/kinds/defaults) is compared. - """ - return [ - (p.name, p.kind, p.default) - for p in inspect.signature(fn).parameters.values() - if p.name != "self" - ] - - -def test_interface_matches_l1_manager_surface(): - """The Protocol declares exactly L1Manager's public method set.""" - methods = _interface_methods() - assert methods == _l1_manager_methods() - assert len(methods) == 17 # tripwire against wholesale drift - - -def test_signatures_match_l1_manager(): - """Each Protocol method's call shape matches L1Manager's.""" - for name in _interface_methods(): - proto = _params(getattr(L1ManagerInterface, name)) - impl = _params(_unwrap(getattr(L1Manager, name))) - assert proto == impl, name - - -def test_l1_manager_conforms_to_interface(): - """The stock L1Manager satisfies the interface (structural, no instance).""" - assert issubclass(L1Manager, L1ManagerInterface) - - -def test_incomplete_class_does_not_conform(): - """A class missing interface methods is rejected (negative control).""" - - class Partial: - def close(self) -> None: ... - - assert not issubclass(Partial, L1ManagerInterface) diff --git a/tests/v1/distributed/test_maru_integration.py b/tests/v1/distributed/test_maru_integration.py deleted file mode 100644 index 1da6bd6c11a..00000000000 --- a/tests/v1/distributed/test_maru_integration.py +++ /dev/null @@ -1,136 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""C11b: maru control integration through the stock tiering controllers. - -A real StorageManager + StoreController/PrefetchController/EvictionController -drive MaruL1Manager, whose CXL pool + MaruServer directory are the in-memory -fakes from the manager-level tests, plus a mock L2 adapter. These assert maru's -*control* integration (register / read-reserve / evict-delete under the real -controllers). L1<->L2 byte movement is stock controller logic (identical for -any L1 backend) and is covered by the stock StorageManager tests, so it is not -re-asserted here (the fakes back pages with no real memory). -""" - -# Standard -import time - -# Third Party -import pytest -import torch - -try: - # First Party - from lmcache.v1.distributed.api import MemoryLayoutDesc, ObjectKey - from lmcache.v1.distributed.config import ( - EvictionConfig, - L1ManagerConfig, - L1MemoryManagerConfig, - MaruL1Config, - StorageManagerConfig, - ) - from lmcache.v1.distributed.l2_adapters.config import L2AdaptersConfig - from lmcache.v1.distributed.l2_adapters.mock_l2_adapter import MockL2AdapterConfig - from lmcache.v1.distributed.maru_l1_manager import ( - MaruL1Manager, - object_key_to_string, - ) - from lmcache.v1.distributed.storage_manager import StorageManager - - # Local - from .maru_fakes import FakeCxlAdapter, FakeMaruHandler -except ImportError: - pytest.skip("maru integration deps unavailable", allow_module_level=True) - -_LAYOUT = MemoryLayoutDesc(shapes=[torch.Size([4, 8])], dtypes=[torch.float16]) - - -def _key(idx: int) -> ObjectKey: - return ObjectKey(chunk_hash=idx.to_bytes(4, "big"), model_name="m", kv_rank=0) - - -def _wait(predicate, timeout: float = 10.0) -> bool: - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - if predicate(): - return True - time.sleep(0.05) - return False - - -def _maru_sm_with_fakes(chunk_size: int = 64, trigger_watermark: float = 0.8): - """A real StorageManager whose maru L1 tier is backed by the fakes.""" - config = StorageManagerConfig( - l1_manager_config=L1ManagerConfig( - memory_config=L1MemoryManagerConfig( - size_in_bytes=0, - use_lazy=False, - maru_config=MaruL1Config( - server_url="maru://localhost:5555", - pool_size_bytes=1 << 20, - instance_id="t", - ), - ), - write_ttl_seconds=600, - read_ttl_seconds=300, - ), - eviction_config=EvictionConfig( - eviction_policy="LRU", trigger_watermark=trigger_watermark - ), - l2_adapter_config=L2AdaptersConfig( - adapters=[MockL2AdapterConfig(max_size_gb=1.0, mock_bandwidth_gb=1.0)] - ), - ) - sm = StorageManager(config) - assert isinstance(sm._l1_manager, MaruL1Manager) # harness selected maru - handler = FakeMaruHandler(chunk_size) - adapter = FakeCxlAdapter(chunk_size) - alloc = sm._l1_manager._allocator - alloc._handler = handler - alloc._cxl_adapter = adapter - alloc._single_token_size = 16 - return sm, handler, adapter - - -def test_store_registers_in_maru_directory(): - """reserve_write -> finish_write through the full stack registers in maru.""" - sm, handler, _ = _maru_sm_with_fakes() - try: - k = _key(1) - res = sm.reserve_write([k], _LAYOUT, mode="new") - assert res[k] is not None - sm.finish_write([k]) - assert object_key_to_string(k) in handler.store_map - finally: - sm.close() - - -def test_prefetch_hits_l1_resident_keys(): - """Prefetch of directory-resident keys is a full L1 hit (maru reserve_read).""" - sm, _, _ = _maru_sm_with_fakes() - try: - keys = [_key(i) for i in range(3)] - sm.reserve_write(keys, _LAYOUT, mode="new") - sm.finish_write(keys) - - handle = sm.submit_prefetch_task(keys, _LAYOUT) - assert _wait(lambda: sm.query_prefetch_status(handle) is not None) - assert sm.query_prefetch_status(handle).count_leading_ones() == len(keys) - finally: - sm.close() - - -def test_eviction_deletes_from_maru_directory(): - """Watermark eviction drives MaruL1Manager.delete on the shared directory.""" - # Low watermark: the fake pool is 16 pages, so a handful of stored keys - # crosses it and the eviction controller must reclaim some. - sm, handler, _ = _maru_sm_with_fakes(trigger_watermark=0.1) - try: - keys = [_key(i) for i in range(6)] - sm.reserve_write(keys, _LAYOUT, mode="new") - sm.finish_write(keys) - assert len(handler.store_map) == 6 - - # The eviction controller runs on its own thread; it deletes evictable - # keys from the maru directory until usage falls under the watermark. - assert _wait(lambda: len(handler.store_map) < 6) - finally: - sm.close() diff --git a/tests/v1/distributed/test_maru_l1_manager.py b/tests/v1/distributed/test_maru_l1_manager.py index d4910b86663..ef2d18e702d 100644 --- a/tests/v1/distributed/test_maru_l1_manager.py +++ b/tests/v1/distributed/test_maru_l1_manager.py @@ -1,9 +1,19 @@ # SPDX-License-Identifier: Apache-2.0 -"""Tests for MaruL1Manager (fake maru runtime, no CXL required).""" +"""Maru-specific unit tests for MaruL1Manager, its allocator, and its control. + +The cross-backend L1 contract shared by the stock ``L1Manager`` and +``MaruL1Manager`` (plus the Protocol structural guards) lives in +``test_l1_manager_conformance.py``. This file covers only maru-specific +behavior: CXL-pool allocation, config / startup guards, RPC-failure page/pin +safety, the TTL sweeper, the L2->L1 promote paths, maru-specific notification +nuances, and control integration through the stock tiering stack. Duplicates of +the shared contract are intentionally absent -- see the conformance suite. +""" # Standard from unittest.mock import MagicMock +import sys import time # Third Party @@ -13,18 +23,37 @@ try: # First Party from lmcache.v1.distributed.api import MemoryLayoutDesc, ObjectKey + from lmcache.v1.distributed.config import ( + EvictionConfig, + L1ManagerConfig, + L1MemoryManagerConfig, + MaruL1Config, + StorageManagerConfig, + parse_args, + ) from lmcache.v1.distributed.error import L1Error - from lmcache.v1.distributed.l1_protocol import L1ManagerInterface + from lmcache.v1.distributed.l2_adapters.config import L2AdaptersConfig + from lmcache.v1.distributed.l2_adapters.mock_l2_adapter import MockL2AdapterConfig from lmcache.v1.distributed.maru_l1_manager import ( MaruL1Manager, _PendingRead, object_key_to_string, ) + from lmcache.v1.distributed.memory_manager.maru_memory_allocator import ( + MaruMemoryAllocator, + _to_tcp, + ) + from lmcache.v1.distributed.storage_manager import StorageManager + from lmcache.v1.memory_management import MemoryFormat from lmcache.v1.mp_observability.event import EventType # Local - from .maru_fakes import RecordingListener, make_maru_manager - from .test_l1_protocol import _interface_methods, _params, _unwrap + from .maru_fakes import ( + FakeCxlAdapter, + FakeMaruHandler, + RecordingListener, + make_maru_manager, + ) except ImportError: pytest.skip("maru manager deps unavailable", allow_module_level=True) @@ -45,116 +74,251 @@ def _seed(handler, key: ObjectKey, rid: int = 7, pid: int = 3, size: int = 64): handler.store_map[object_key_to_string(key)] = (rid, pid, size) -def _store(manager, keys: list[ObjectKey]): - """Drive the full write path: reserve + finish.""" - res = manager.reserve_write(keys, [False] * len(keys), _LAYOUT, mode="new") - assert all(err == L1Error.SUCCESS for err, _ in res.values()) - fin = manager.finish_write(keys) - assert all(err == L1Error.SUCCESS for err in fin.values()) - - # ========================================================================= -# interface conformance (binds MaruL1Manager to the shared Protocol) +# allocator (MaruMemoryAllocator): URL normalization, init, delegation +# (absorbed from the former test_maru_memory_allocator.py, mocked runtime) # ========================================================================= -def test_conforms_to_l1_manager_interface(): - assert issubclass(MaruL1Manager, L1ManagerInterface) +def _alloc_cfg() -> MaruL1Config: + return MaruL1Config(server_url="maru://localhost:9000", pool_size_bytes=1 << 30) + + +# 1 group, shape (4 tokens, 8 feats) fp16 -> 4*8*2 = 64 bytes; 4 tokens/chunk. +_ALLOC_LAYOUT = ([torch.Size([4, 8])], [torch.float16], MemoryFormat.KV_2LTD, 4) + + +@pytest.fixture +def maru_mocks(monkeypatch): + """Inject fake maru / maru_lmcache modules for init_layout's lazy imports.""" + handler = MagicMock() + handler.connect.return_value = True + handler.get_chunk_size.return_value = 64 + adapter = MagicMock() + maru_mod = MagicMock() + maru_mod.MaruHandler.return_value = handler + lmcache_mod = MagicMock() + lmcache_mod.CxlMemoryAdapter.return_value = adapter + monkeypatch.setitem(sys.modules, "maru", maru_mod) + monkeypatch.setitem(sys.modules, "maru_lmcache", lmcache_mod) + return handler, adapter, maru_mod + + +def test_to_tcp(): + assert _to_tcp("maru://h:1") == "tcp://h:1" + assert _to_tcp("tcp://h:1") == "tcp://h:1" + + +def test_methods_before_init_raise(): + alloc = MaruMemoryAllocator(_alloc_cfg()) + assert not alloc.is_initialized + with pytest.raises(RuntimeError): + alloc.allocate([torch.Size([4, 8])], [torch.float16]) + with pytest.raises(RuntimeError): + _ = alloc.handler + with pytest.raises(RuntimeError): + _ = alloc.single_token_size + + +def test_init_layout_builds_pool(maru_mocks): + handler, adapter, maru_mod = maru_mocks + alloc = MaruMemoryAllocator(_alloc_cfg()) + alloc.init_layout(*_ALLOC_LAYOUT) + + assert alloc.is_initialized + handler.connect.assert_called_once() + assert alloc.handler is handler + assert alloc.single_token_size == 16 # 64 bytes / 4 tokens + + _, kwargs = maru_mod.MaruConfig.call_args + assert kwargs["server_url"] == "tcp://localhost:9000" + assert kwargs["pool_size"] == 1 << 30 + assert kwargs["chunk_size_bytes"] == 64 + assert kwargs["auto_connect"] is False + + +def test_init_layout_mismatch_raises(maru_mocks): + alloc = MaruMemoryAllocator(_alloc_cfg()) + alloc.init_layout(*_ALLOC_LAYOUT) + with pytest.raises(ValueError): + alloc.init_layout( + [torch.Size([8, 8])], [torch.float16], MemoryFormat.KV_2LTD, 4 + ) + + +def test_connect_failure_raises(maru_mocks): + handler, _, _ = maru_mocks + handler.connect.return_value = False + alloc = MaruMemoryAllocator(_alloc_cfg()) + with pytest.raises(RuntimeError): + alloc.init_layout(*_ALLOC_LAYOUT) + +def test_delegation_and_lifecycle(maru_mocks): + _, adapter, _ = maru_mocks + alloc = MaruMemoryAllocator(_alloc_cfg()) + alloc.init_layout(*_ALLOC_LAYOUT) + obj = MagicMock() + + adapter.allocate.return_value = obj + assert alloc.allocate([torch.Size([4, 8])], [torch.float16]) is obj + + alloc.get_by_location(1, 2, 64) + adapter.get_by_location.assert_called_once_with( + region_id=1, page_index=2, actual_size=64, single_token_size=16 + ) -def test_signatures_match_interface(): - """Each Protocol method's call shape matches MaruL1Manager's.""" - for name in _interface_methods(): - proto = _params(getattr(L1ManagerInterface, name)) - impl = _params(_unwrap(getattr(MaruL1Manager, name))) - assert proto == impl, name + alloc.create_store_handle(obj) + adapter.create_store_handle.assert_called_once_with(obj) + + # abort_alloc returns the page via the adapter's real free + alloc.abort_alloc(obj) + adapter.free.assert_called_once_with(obj) + + # free/batched_free are no-ops (lifecycle owned by MaruServer) + adapter.free.reset_mock() + alloc.free(obj) + alloc.batched_free([obj]) + adapter.free.assert_not_called() # ========================================================================= -# reserve_read / unsafe_read / finish_read +# config / startup guards: flag parsing + rejected backend combinations +# (absorbed from the former test_l1_config_maru.py) # ========================================================================= -def test_reserve_read_hit_pins_and_stages(): - manager, handler, _ = make_maru_manager() - k = _key(1) - _seed(handler, k) +def _args(*extra: str, l1_size_gb: str = "1") -> list[str]: + """Minimal required flags (eviction policy + L1 size) plus extras.""" + return [ + "--eviction-policy", + "LRU", + "--l1-size-gb", + l1_size_gb, + "--no-l1-use-lazy", + *extra, + ] - res = manager.reserve_read([k]) - err, obj = res[k] - assert err == L1Error.SUCCESS and obj is not None - assert handler.pins[object_key_to_string(k)] == 1 - assert manager.unsafe_read([k])[k] == (L1Error.SUCCESS, obj) + +def _maru_memory(**overrides) -> L1MemoryManagerConfig: + return L1MemoryManagerConfig( + size_in_bytes=0, + use_lazy=False, + maru_config=MaruL1Config( + server_url="maru://localhost:5555", + pool_size_bytes=1 << 20, + instance_id="t", + ), + **overrides, + ) -def test_reserve_read_miss_leaves_no_pin(): - manager, handler, _ = make_maru_manager() - res = manager.reserve_read([_key(1)]) - assert res[_key(1)] == (L1Error.KEY_NOT_EXIST, None) - assert not handler.pins +def _maru_sm_config( + *, memory=None, store_policy="default", adapters=None +) -> StorageManagerConfig: + return StorageManagerConfig( + l1_manager_config=L1ManagerConfig( + memory_config=memory if memory is not None else _maru_memory(), + write_ttl_seconds=600, + read_ttl_seconds=300, + ), + eviction_config=EvictionConfig(eviction_policy="LRU"), + l2_adapter_config=L2AdaptersConfig(adapters=adapters or []), + store_policy=store_policy, + ) -def test_reserve_read_is_per_key_independent(): - """A miss in the middle must not shadow later hits (unlike prefix-stop).""" - manager, handler, _ = make_maru_manager() - k1, k2, k3 = _key(1), _key(2), _key(3) - _seed(handler, k1, pid=1) - _seed(handler, k3, pid=3) +def test_maru_flags_build_maru_config(): + cfg = parse_args( + _args( + "--maru-server-url", + "maru://localhost:9000", + "--maru-pool-size-gb", + "8", + "--maru-instance-id", + "node-a", + ) + ) + maru = cfg.l1_manager_config.memory_config.maru_config + assert isinstance(maru, MaruL1Config) + assert maru.server_url == "maru://localhost:9000" + assert maru.pool_size_bytes == 8 * (1 << 30) + assert maru.instance_id == "node-a" - res = manager.reserve_read([k1, k2, k3]) - assert res[k1][0] == L1Error.SUCCESS - assert res[k2] == (L1Error.KEY_NOT_EXIST, None) - assert res[k3][0] == L1Error.SUCCESS - assert handler.pins[object_key_to_string(k3)] == 1 +def test_maru_without_pool_size_raises(): + with pytest.raises(ValueError): + parse_args(_args("--maru-server-url", "maru://localhost:9000")) -def test_extra_count_takes_and_releases_n_pins(): - manager, handler, _ = make_maru_manager() - k = _key(1) - ks = object_key_to_string(k) - _seed(handler, k) - manager.reserve_read([k], extra_count=2) - assert handler.pins[ks] == 3 +def test_maru_rejects_devdax_l1(): + with pytest.raises(ValueError, match="devdax"): + _maru_sm_config(memory=_maru_memory(devdax_path="/dev/dax0.0", shm_name="")) - # Three independent finish_read calls balance the three pins. - for _ in range(3): - assert manager.finish_read([k])[k] == L1Error.SUCCESS - assert handler.pins[ks] == 0 - assert manager.unsafe_read([k])[k] == (L1Error.KEY_NOT_EXIST, None) +def test_maru_rejects_registered_l2(monkeypatch): + # Simulate a registered/RDMA adapter via the shared region classifier. + monkeypatch.setattr( + "lmcache.v1.distributed.config._requires_single_l1_memory_region", + lambda adapter_config: "nixl_store", + ) + with pytest.raises(ValueError, match="registerable"): + _maru_sm_config( + adapters=[MockL2AdapterConfig(max_size_gb=1.0, mock_bandwidth_gb=1.0)] + ) -def test_overlapping_reserve_accumulates_refcount(): - manager, handler, _ = make_maru_manager() - k = _key(1) - ks = object_key_to_string(k) - _seed(handler, k) - _, obj1 = manager.reserve_read([k])[k] - _, obj2 = manager.reserve_read([k])[k] - assert obj1 is obj2 # one staged object per key - assert handler.pins[ks] == 2 +def test_maru_rejects_non_lmcache_driven_transfer(): + """maru requires the LMCache-driven transfer path (rejects engine/auto).""" + # First Party + from lmcache.v1.mp_observability.config import ObservabilityConfig + from lmcache.v1.multiprocess.config import ( + CoordinatorConfig, + HTTPFrontendConfig, + MPServerConfig, + ) + from lmcache.v1.multiprocess.http_server import run_http_server - manager.finish_read([k]) - assert manager.unsafe_read([k])[k][0] == L1Error.SUCCESS # still staged - manager.finish_read([k]) - assert handler.pins[ks] == 0 + with pytest.raises(ValueError, match="lmcache_driven"): + run_http_server( + http_config=HTTPFrontendConfig(), + mp_config=MPServerConfig(supported_transfer_mode="engine_driven"), + storage_manager_config=_maru_sm_config(), + obs_config=ObservabilityConfig(), + coordinator_config=CoordinatorConfig(url=""), + ) -def test_finish_read_unstaged_does_not_unpin(): +# ========================================================================= +# reserve_read / finish_read: pin balance + page-reclaim safety under +# RPC failures, and the two staging regression guards +# ========================================================================= + + +def test_reserve_read_pin_rpc_failure_is_a_miss(): manager, handler, _ = make_maru_manager() - assert manager.finish_read([_key(1)])[_key(1)] == L1Error.KEY_NOT_EXIST - assert not handler.unpin_log + k = _key(1) + _seed(handler, k) + handler.fail_pin = True + assert manager.reserve_read([k])[k] == (L1Error.KEY_NOT_EXIST, None) + assert not handler.pins -def test_finish_read_never_over_releases(): +def test_reserve_read_retrieve_rpc_failure_rolls_back_pins(): manager, handler, _ = make_maru_manager() k = _key(1) _seed(handler, k) - manager.reserve_read([k]) # one pin + handler.fail_retrieve = True + assert manager.reserve_read([k])[k] == (L1Error.KEY_NOT_EXIST, None) + assert handler.pins[object_key_to_string(k)] == 0 + - assert manager.finish_read([k], extra_count=5)[k] == L1Error.SUCCESS - assert len(handler.unpin_log) == 1 # released min(6, refcount=1) +def test_reserve_read_pool_miss_rolls_back_pins(): + manager, handler, adapter = make_maru_manager() + k = _key(1) + _seed(handler, k) + adapter.resolve_none = True + assert manager.reserve_read([k])[k] == (L1Error.KEY_NOT_EXIST, None) assert handler.pins[object_key_to_string(k)] == 0 @@ -225,33 +389,10 @@ def test_reserve_read_excludes_mid_write_key_no_double_staging(): # ========================================================================= -# reserve_write / finish_write +# reserve_write / finish_write: OOM, mode guard, and page-reclaim safety # ========================================================================= -def test_write_path_registers_key(): - manager, handler, _ = make_maru_manager() - k = _key(1) - _store(manager, [k]) - assert object_key_to_string(k) in handler.store_map - # Now readable through the directory. - assert manager.reserve_read([k])[k][0] == L1Error.SUCCESS - - -def test_reserve_write_rejects_staged_and_registered_keys(): - manager, handler, _ = make_maru_manager() - staged, registered, fresh = _key(1), _key(2), _key(3) - manager.reserve_write([staged], [False], _LAYOUT, mode="new") - _seed(handler, registered) - - res = manager.reserve_write( - [staged, registered, fresh], [False] * 3, _LAYOUT, mode="new" - ) - assert res[staged] == (L1Error.KEY_NOT_WRITABLE, None) # locally staged - assert res[registered] == (L1Error.KEY_NOT_WRITABLE, None) # cross-instance dedup - assert res[fresh][0] == L1Error.SUCCESS - - def test_reserve_write_oom_marks_whole_batch(): manager, _, adapter = make_maru_manager() adapter.oom = True @@ -265,133 +406,6 @@ def test_reserve_write_rejects_non_new_mode(): manager.reserve_write([_key(1)], [False], _LAYOUT, mode="update") -def test_finish_write_unstaged_key(): - manager, _, _ = make_maru_manager() - assert manager.finish_write([_key(1)])[_key(1)] == L1Error.KEY_NOT_EXIST - - -def test_finish_write_store_failure_reclaims_page(): - manager, handler, adapter = make_maru_manager() - k = _key(1) - handler.fail_store_keys.add(object_key_to_string(k)) - - res = manager.reserve_write([k], [False], _LAYOUT, mode="new") - addr = res[k][1].metadata.address - assert manager.finish_write([k])[k] == L1Error.KEY_IN_WRONG_STATE - assert addr in adapter.freed # aborted back to the owner free list - - -# ========================================================================= -# delete / clear / misc -# ========================================================================= - - -def test_delete_locally_staged_key_is_locked(): - manager, handler, _ = make_maru_manager() - k = _key(1) - _seed(handler, k) - manager.reserve_read([k]) - - assert manager.delete([k])[k] == L1Error.KEY_IS_LOCKED - assert object_key_to_string(k) in handler.store_map # server untouched - - -def test_delete_remotely_pinned_key_is_locked(): - manager, handler, _ = make_maru_manager() - k = _key(1) - _seed(handler, k) - handler.pins[object_key_to_string(k)] = 1 # pinned by another instance - - assert manager.delete([k])[k] == L1Error.KEY_IS_LOCKED - - -def test_delete_absent_and_present_keys(): - manager, handler, _ = make_maru_manager() - present, absent = _key(1), _key(2) - _seed(handler, present) - - res = manager.delete([present, absent]) - assert res[present] == L1Error.SUCCESS - assert res[absent] == L1Error.KEY_NOT_EXIST - assert object_key_to_string(present) not in handler.store_map - - -def test_clear_non_force_preserves_staging(): - manager, handler, adapter = make_maru_manager() - rk, wk = _key(1), _key(2) - _seed(handler, rk) - manager.reserve_read([rk]) - manager.reserve_write([wk], [False], _LAYOUT, mode="new") - - manager.clear() # force=False keeps locked (staged) entries - assert handler.pins[object_key_to_string(rk)] == 1 - assert not adapter.freed - assert manager.unsafe_read([rk])[rk][0] == L1Error.SUCCESS - - -def test_clear_force_balances_pins_and_reclaims_writes(): - manager, handler, adapter = make_maru_manager() - rk, wk = _key(1), _key(2) - _seed(handler, rk) - manager.reserve_read([rk], extra_count=1) # two pins - manager.reserve_write([wk], [False], _LAYOUT, mode="new") - - manager.clear(force=True) - assert handler.pins[object_key_to_string(rk)] == 0 - assert len(adapter.freed) == 1 # staged write page aborted - assert manager.unsafe_read([rk])[rk] == (L1Error.KEY_NOT_EXIST, None) - - -# ========================================================================= -# failure paths (pin balance + page-reclaim safety under RPC failures) -# ========================================================================= - - -def test_reserve_read_pin_rpc_failure_is_a_miss(): - manager, handler, _ = make_maru_manager() - k = _key(1) - _seed(handler, k) - handler.fail_pin = True - assert manager.reserve_read([k])[k] == (L1Error.KEY_NOT_EXIST, None) - assert not handler.pins - - -def test_reserve_read_retrieve_rpc_failure_rolls_back_pins(): - manager, handler, _ = make_maru_manager() - k = _key(1) - _seed(handler, k) - handler.fail_retrieve = True - assert manager.reserve_read([k])[k] == (L1Error.KEY_NOT_EXIST, None) - assert handler.pins[object_key_to_string(k)] == 0 - - -def test_reserve_read_all_none_retrieve_rolls_back_pins(): - """Real transport failure: batch_retrieve returns [None]*len, no raise.""" - manager, handler, _ = make_maru_manager() - k = _key(1) - _seed(handler, k) - handler.retrieve_none = True - assert manager.reserve_read([k], extra_count=1)[k] == (L1Error.KEY_NOT_EXIST, None) - assert handler.pins[object_key_to_string(k)] == 0 # both pins rolled back - - -def test_reserve_read_pool_miss_rolls_back_pins(): - manager, handler, adapter = make_maru_manager() - k = _key(1) - _seed(handler, k) - adapter.resolve_none = True - assert manager.reserve_read([k])[k] == (L1Error.KEY_NOT_EXIST, None) - assert handler.pins[object_key_to_string(k)] == 0 - - -def test_reserve_read_mid_write_key_not_readable(): - manager, _, _ = make_maru_manager() - k = _key(1) - manager.reserve_write([k], [False], _LAYOUT, mode="new") - assert manager.reserve_read([k])[k] == (L1Error.KEY_NOT_READABLE, None) - assert manager.unsafe_read([k])[k] == (L1Error.KEY_NOT_READABLE, None) - - def test_finish_write_handle_failure_reclaims_all_pages(): manager, _, adapter = make_maru_manager() keys = [_key(1), _key(2)] @@ -421,20 +435,23 @@ def test_finish_write_dup_skip_is_success_without_abort(): assert not adapter.freed # no local double-free -def test_reserve_write_exists_rpc_failure_proceeds(): - manager, handler, _ = make_maru_manager() - k = _key(1) - handler.fail_exists = True - res = manager.reserve_write([k], [False], _LAYOUT, mode="new") - assert res[k][0] == L1Error.SUCCESS # allocated; dup-skip at store time +# ========================================================================= +# delete: cross-node pinned-key refusal +# ========================================================================= -def test_delete_rpc_failure_reports_locked(): +def test_delete_remotely_pinned_key_is_locked(): manager, handler, _ = make_maru_manager() k = _key(1) _seed(handler, k) - handler.fail_delete = True - assert manager.delete([k])[k] == L1Error.KEY_IS_LOCKED # retryable + handler.pins[object_key_to_string(k)] = 1 # pinned by another instance + + assert manager.delete([k])[k] == L1Error.KEY_IS_LOCKED + + +# ========================================================================= +# get_memory_usage (device-fill watermark) / is_key_evictable +# ========================================================================= def test_get_memory_usage_before_init_reports_pool_size(): @@ -477,23 +494,10 @@ def test_is_key_evictable_tracks_local_staging(): # ========================================================================= -# finish_write_and_reserve_read (C5): L2->L1 promote +# finish_write_and_reserve_read (L2->L1 promote): retained vs temporary # ========================================================================= -def test_temporary_promote_stages_read_without_registering(): - """Default prefetch: private staging -- no batch_store, no server pin.""" - manager, handler, _ = make_maru_manager() - k = _key(1) - page = manager.reserve_write([k], [True], _LAYOUT, mode="new")[k][1] - - err, obj = manager.finish_write_and_reserve_read([k])[k] - assert err == L1Error.SUCCESS and obj is page - assert object_key_to_string(k) not in handler.store_map # never registered - assert not handler.pins # temporary reads hold no pin - assert manager.unsafe_read([k])[k] == (L1Error.SUCCESS, page) - - def test_temporary_promote_page_reclaimed_after_read(): manager, handler, adapter = make_maru_manager() k = _key(1) @@ -506,23 +510,6 @@ def test_temporary_promote_page_reclaimed_after_read(): assert manager.unsafe_read([k])[k] == (L1Error.KEY_NOT_EXIST, None) -def test_retained_promote_registers_and_pins(): - """retain policy: batch_store + authoritative re-resolve with pins.""" - manager, handler, adapter = make_maru_manager() - k = _key(1) - ks = object_key_to_string(k) - manager.reserve_write([k], [False], _LAYOUT, mode="new") - - assert manager.finish_write_and_reserve_read([k])[k][0] == L1Error.SUCCESS - assert ks in handler.store_map # registered in the shared directory - assert handler.pins[ks] == 1 # read hold pinned - - assert manager.finish_read([k])[k] == L1Error.SUCCESS - assert handler.pins[ks] == 0 # unpinned, not freed - assert not adapter.freed - assert ks in handler.store_map # retained page survives the read - - def test_retained_promote_extra_count_pins_n(): manager, handler, _ = make_maru_manager() k = _key(1) @@ -543,29 +530,6 @@ def test_retained_promote_dup_skip_resolves_winner(): assert handler.pins[object_key_to_string(k)] == 1 -def test_promote_unstaged_key_is_not_exist(): - manager, _, _ = make_maru_manager() - assert manager.finish_write_and_reserve_read([_key(1)])[_key(1)] == ( - L1Error.KEY_NOT_EXIST, - None, - ) - - -def test_promote_already_read_staged_is_wrong_state(): - """The both-staged race the read-staged guard defends against.""" - manager, _, adapter = make_maru_manager() - k = _key(1) - manager.reserve_write([k], [False], _LAYOUT, mode="new") - rpage = adapter.batched_allocate(_LAYOUT.shapes, _LAYOUT.dtypes, 1)[0] - manager._pending_read[k] = _PendingRead(mem_obj=rpage, refcount=1) - - assert manager.finish_write_and_reserve_read([k])[k] == ( - L1Error.KEY_IN_WRONG_STATE, - None, - ) - assert k in manager._pending_write # guard returned before popping - - def test_retained_promote_store_failure_is_wrong_state(): manager, handler, _ = make_maru_manager() k = _key(1) @@ -579,44 +543,11 @@ def test_retained_promote_store_failure_is_wrong_state(): assert k not in manager._pending_write # write staging drained -def test_promote_fires_promote_event_not_write_finished(): - """Anti-#2744: promote must never fire write_finished (would re-store L2).""" - manager, handler, _ = make_maru_manager() - k = _key(1) - manager.reserve_write([k], [False], _LAYOUT, mode="new") - rec = RecordingListener() - manager.register_listener(rec) - - manager.finish_write_and_reserve_read([k]) - assert rec.kinds("finish_write_and_reserve_read") == [[k]] - assert rec.kinds("write_finished") == [] - - -def test_temporary_promote_fires_promote_event(): - manager, _, _ = make_maru_manager() - k = _key(1) - manager.reserve_write([k], [True], _LAYOUT, mode="new") - rec = RecordingListener() - manager.register_listener(rec) - - manager.finish_write_and_reserve_read([k]) - assert rec.kinds("finish_write_and_reserve_read") == [[k]] - assert rec.kinds("write_finished") == [] - - # ========================================================================= -# TTL sweeper (C6): reclaim orphan write pages / read pins +# TTL sweeper: reclaim orphan write pages / read pins on expiry # ========================================================================= -def test_sweeper_thread_starts_and_stops(): - manager, _, _ = make_maru_manager() - assert manager._sweeper.is_alive() - manager.close() - manager._sweeper.join(timeout=2) - assert not manager._sweeper.is_alive() - - def test_sweep_reclaims_expired_write(): manager, _, adapter = make_maru_manager() k = _key(1) @@ -667,31 +598,9 @@ def test_sweep_leaves_live_staging(): assert not adapter.freed -def test_overlapping_reserve_refreshes_read_deadline(): - manager, handler, _ = make_maru_manager() - k = _key(1) - _seed(handler, k) - manager.reserve_read([k]) - manager._pending_read[k].deadline = time.monotonic() - 1 # go stale - - manager.reserve_read([k]) # overlap -> refresh - assert manager._pending_read[k].deadline > time.monotonic() - - -def test_finish_read_after_sweep_is_not_exist(): - """A late finish after a sweep behaves like a stock TTL expiry.""" - manager, handler, _ = make_maru_manager() - k = _key(1) - _seed(handler, k) - manager.reserve_read([k]) - manager._pending_read[k].deadline = time.monotonic() - 1 - manager._sweep_once() - - assert manager.finish_read([k])[k] == L1Error.KEY_NOT_EXIST - - # ========================================================================= -# event_bus observability parity (C11a): mirrors stock L1Manager publishes +# notifications: maru-specific event_bus / listener nuances +# (the shared listener lifecycle is asserted in the conformance suite) # ========================================================================= @@ -722,82 +631,6 @@ def test_event_bus_write_read_delete_lifecycle(): assert [k] in _keys_for(manager._event_bus, EventType.L1_KEYS_EVICTED) # delete -def test_event_bus_promote_publishes_promote_event_not_write_finished(): - manager, _, _ = make_maru_manager() - manager._event_bus = MagicMock() - k = _key(1) - manager.reserve_write([k], [False], _LAYOUT, mode="new") - - manager.finish_write_and_reserve_read([k]) - assert _keys_for( - manager._event_bus, EventType.L1_WRITE_FINISHED_AND_READ_RESERVED - ) == [[k]] - # anti re-store at the event-bus level too. - assert _keys_for(manager._event_bus, EventType.L1_WRITE_FINISHED) == [] - - -def test_event_bus_temporary_finish_read_publishes_evicted(): - manager, _, _ = make_maru_manager() - manager._event_bus = MagicMock() - k = _key(1) - manager.reserve_write([k], [True], _LAYOUT, mode="new") - manager.finish_write_and_reserve_read([k]) # temporary read staged - - manager.finish_read([k]) - assert _keys_for(manager._event_bus, EventType.L1_READ_FINISHED) == [[k]] - assert [k] in _keys_for(manager._event_bus, EventType.L1_KEYS_EVICTED) - - -def test_event_bus_touch_keys_does_not_publish(): - manager, _, _ = make_maru_manager() - manager._event_bus = MagicMock() - - manager.touch_keys([_key(1)]) - manager._event_bus.publish.assert_not_called() - - -def test_report_status_keys(): - manager, handler, _ = make_maru_manager() - k = _key(1) - _seed(handler, k) - manager.reserve_read([k]) - - status = manager.report_status() - assert status["backend"] == "maru" - assert status["read_locked_count"] == 1 - assert status["is_healthy"] is True - assert status["memory_total_bytes"] > 0 - - -# ========================================================================= -# listener firing (C4): feeds the eviction LRU + store controller -# ========================================================================= - - -def test_reserve_read_fires_only_for_hits(): - """A miss must not be reported as a reserved-read hold.""" - manager, handler, _ = make_maru_manager() - hit, miss = _key(1), _key(2) - _seed(handler, hit) - rec = RecordingListener() - manager.register_listener(rec) - - manager.reserve_read([hit, miss]) - assert rec.kinds("reserved_read") == [[hit]] - - -def test_reserve_and_finish_write_fire_with_registered_keys(): - manager, handler, _ = make_maru_manager() - k = _key(1) - rec = RecordingListener() - manager.register_listener(rec) - - manager.reserve_write([k], [False], _LAYOUT, mode="new") - assert rec.kinds("reserved_write") == [[k]] - manager.finish_write([k]) - assert rec.kinds("write_finished") == [[k]] - - def test_finish_write_store_failure_is_not_fired(): """A page that never registered must not notify write-finished.""" manager, handler, _ = make_maru_manager() @@ -811,20 +644,6 @@ def test_finish_write_store_failure_is_not_fired(): assert rec.kinds("write_finished") == [[]] # fired once, empty -def test_finish_read_normal_fires_read_finished_without_delete(): - """A directory-backed read unpins; it is never a manager delete.""" - manager, handler, _ = make_maru_manager() - k = _key(1) - _seed(handler, k) - manager.reserve_read([k]) - rec = RecordingListener() - manager.register_listener(rec) - - manager.finish_read([k]) - assert rec.kinds("read_finished") == [[k]] - assert all(ks == [] for ks in rec.kinds("deleted_by_manager")) - - def test_finish_read_temporary_frees_page_and_fires_delete(): """A temporary read (local staging) reclaims its page at refcount zero. @@ -861,38 +680,102 @@ def test_delete_fires_only_for_removed_keys(): assert rec.kinds("deleted_by_manager") == [[removed]] -def test_touch_keys_fires_accessed(): - manager, _, _ = make_maru_manager() - rec = RecordingListener() - manager.register_listener(rec) - - keys = [_key(1), _key(2)] - manager.touch_keys(keys) - assert rec.kinds("accessed") == [keys] - - -def test_clear_force_fires_deleted_for_dropped_staging(): - manager, handler, _ = make_maru_manager() - rk, wk = _key(1), _key(2) - _seed(handler, rk) - manager.reserve_read([rk]) - manager.reserve_write([wk], [False], _LAYOUT, mode="new") - rec = RecordingListener() - manager.register_listener(rec) - - manager.clear(force=True) - fired = rec.kinds("deleted_by_manager") - assert len(fired) == 1 - assert set(fired[0]) == {rk, wk} - +# ========================================================================= +# control integration: MaruL1Manager under the real tiering controllers +# (absorbed from the former test_maru_integration.py, C11b) +# +# A real StorageManager + StoreController/PrefetchController/EvictionController +# drive MaruL1Manager, whose CXL pool + MaruServer directory are the in-memory +# fakes. These assert maru's *control* integration (register / read-reserve / +# evict-delete under the real controllers); L1<->L2 byte movement is stock +# controller logic covered by the stock StorageManager tests. +# ========================================================================= -def test_clear_non_force_fires_nothing(): - manager, handler, _ = make_maru_manager() - rk = _key(1) - _seed(handler, rk) - manager.reserve_read([rk]) - rec = RecordingListener() - manager.register_listener(rec) - manager.clear() # keeps locked staging -> no drops - assert rec.events == [] +def _wait(predicate, timeout: float = 10.0) -> bool: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return True + time.sleep(0.05) + return False + + +def _maru_sm_with_fakes(chunk_size: int = 64, trigger_watermark: float = 0.8): + """A real StorageManager whose maru L1 tier is backed by the fakes.""" + config = StorageManagerConfig( + l1_manager_config=L1ManagerConfig( + memory_config=L1MemoryManagerConfig( + size_in_bytes=0, + use_lazy=False, + maru_config=MaruL1Config( + server_url="maru://localhost:5555", + pool_size_bytes=1 << 20, + instance_id="t", + ), + ), + write_ttl_seconds=600, + read_ttl_seconds=300, + ), + eviction_config=EvictionConfig( + eviction_policy="LRU", trigger_watermark=trigger_watermark + ), + l2_adapter_config=L2AdaptersConfig( + adapters=[MockL2AdapterConfig(max_size_gb=1.0, mock_bandwidth_gb=1.0)] + ), + ) + sm = StorageManager(config) + assert isinstance(sm._l1_manager, MaruL1Manager) # harness selected maru + handler = FakeMaruHandler(chunk_size) + adapter = FakeCxlAdapter(chunk_size) + alloc = sm._l1_manager._allocator + alloc._handler = handler + alloc._cxl_adapter = adapter + alloc._single_token_size = 16 + return sm, handler, adapter + + +def test_store_registers_in_maru_directory(): + """reserve_write -> finish_write through the full stack registers in maru.""" + sm, handler, _ = _maru_sm_with_fakes() + try: + k = _key(1) + res = sm.reserve_write([k], _LAYOUT, mode="new") + assert res[k] is not None + sm.finish_write([k]) + assert object_key_to_string(k) in handler.store_map + finally: + sm.close() + + +def test_prefetch_hits_l1_resident_keys(): + """Prefetch of directory-resident keys is a full L1 hit (maru reserve_read).""" + sm, _, _ = _maru_sm_with_fakes() + try: + keys = [_key(i) for i in range(3)] + sm.reserve_write(keys, _LAYOUT, mode="new") + sm.finish_write(keys) + + handle = sm.submit_prefetch_task(keys, _LAYOUT) + assert _wait(lambda: sm.query_prefetch_status(handle) is not None) + assert sm.query_prefetch_status(handle).count_leading_ones() == len(keys) + finally: + sm.close() + + +def test_eviction_deletes_from_maru_directory(): + """Watermark eviction drives MaruL1Manager.delete on the shared directory.""" + # Low watermark: the fake pool is 16 pages, so a handful of stored keys + # crosses it and the eviction controller must reclaim some. + sm, handler, _ = _maru_sm_with_fakes(trigger_watermark=0.1) + try: + keys = [_key(i) for i in range(6)] + sm.reserve_write(keys, _LAYOUT, mode="new") + sm.finish_write(keys) + assert len(handler.store_map) == 6 + + # The eviction controller runs on its own thread; it deletes evictable + # keys from the maru directory until usage falls under the watermark. + assert _wait(lambda: len(handler.store_map) < 6) + finally: + sm.close() diff --git a/tests/v1/distributed/test_maru_memory_allocator.py b/tests/v1/distributed/test_maru_memory_allocator.py deleted file mode 100644 index ab6efee43fd..00000000000 --- a/tests/v1/distributed/test_maru_memory_allocator.py +++ /dev/null @@ -1,160 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -"""Tests for MaruMemoryAllocator (mocked maru runtime, no CXL required).""" - -# Standard -from unittest.mock import MagicMock -import sys - -# Third Party -import pytest -import torch - -try: - # First Party - from lmcache.v1.distributed.config import MaruL1Config - from lmcache.v1.distributed.memory_manager.maru_memory_allocator import ( - MaruMemoryAllocator, - _to_tcp, - ) - from lmcache.v1.memory_management import MemoryFormat -except ImportError: - pytest.skip("maru allocator deps unavailable", allow_module_level=True) - - -def _cfg() -> MaruL1Config: - return MaruL1Config(server_url="maru://localhost:9000", pool_size_bytes=1 << 30) - - -# 1 group, shape (4 tokens, 8 feats) fp16 -> 4*8*2 = 64 bytes; 4 tokens/chunk. -_LAYOUT = ([torch.Size([4, 8])], [torch.float16], MemoryFormat.KV_2LTD, 4) - - -@pytest.fixture -def maru_mocks(monkeypatch): - """Inject fake maru / maru_lmcache modules for init_layout's lazy imports.""" - handler = MagicMock() - handler.connect.return_value = True - handler.get_chunk_size.return_value = 64 - adapter = MagicMock() - maru_mod = MagicMock() - maru_mod.MaruHandler.return_value = handler - lmcache_mod = MagicMock() - lmcache_mod.CxlMemoryAdapter.return_value = adapter - monkeypatch.setitem(sys.modules, "maru", maru_mod) - monkeypatch.setitem(sys.modules, "maru_lmcache", lmcache_mod) - return handler, adapter, maru_mod - - -def test_to_tcp(): - assert _to_tcp("maru://h:1") == "tcp://h:1" - assert _to_tcp("tcp://h:1") == "tcp://h:1" - - -def test_methods_before_init_raise(): - alloc = MaruMemoryAllocator(_cfg()) - assert not alloc.is_initialized - with pytest.raises(RuntimeError): - alloc.allocate([torch.Size([4, 8])], [torch.float16]) - with pytest.raises(RuntimeError): - _ = alloc.handler - with pytest.raises(RuntimeError): - _ = alloc.single_token_size - - -def test_init_layout_builds_pool(maru_mocks): - handler, adapter, maru_mod = maru_mocks - alloc = MaruMemoryAllocator(_cfg()) - alloc.init_layout(*_LAYOUT) - - assert alloc.is_initialized - handler.connect.assert_called_once() - assert alloc.handler is handler - assert alloc.single_token_size == 16 # 64 bytes / 4 tokens - - _, kwargs = maru_mod.MaruConfig.call_args - assert kwargs["server_url"] == "tcp://localhost:9000" - assert kwargs["pool_size"] == 1 << 30 - assert kwargs["chunk_size_bytes"] == 64 - assert kwargs["auto_connect"] is False - - -def test_init_layout_same_layout_is_noop(maru_mocks): - handler, _, _ = maru_mocks - alloc = MaruMemoryAllocator(_cfg()) - alloc.init_layout(*_LAYOUT) - alloc.init_layout(*_LAYOUT) - handler.connect.assert_called_once() # not reconnected - - -def test_init_layout_mismatch_raises(maru_mocks): - alloc = MaruMemoryAllocator(_cfg()) - alloc.init_layout(*_LAYOUT) - with pytest.raises(ValueError): - alloc.init_layout( - [torch.Size([8, 8])], [torch.float16], MemoryFormat.KV_2LTD, 4 - ) - - -def test_init_layout_bad_chunk_raises(maru_mocks): - alloc = MaruMemoryAllocator(_cfg()) - with pytest.raises(ValueError): - alloc.init_layout( - [torch.Size([4, 8])], [torch.float16], MemoryFormat.KV_2LTD, 0 - ) - - -def test_init_layout_non_divisible_chunk_raises(maru_mocks): - alloc = MaruMemoryAllocator(_cfg()) - # 64 bytes is not a multiple of 5 tokens. - with pytest.raises(ValueError): - alloc.init_layout( - [torch.Size([4, 8])], [torch.float16], MemoryFormat.KV_2LTD, 5 - ) - - -def test_connect_failure_raises(maru_mocks): - handler, _, _ = maru_mocks - handler.connect.return_value = False - alloc = MaruMemoryAllocator(_cfg()) - with pytest.raises(RuntimeError): - alloc.init_layout(*_LAYOUT) - - -def test_delegation_and_lifecycle(maru_mocks): - _, adapter, _ = maru_mocks - alloc = MaruMemoryAllocator(_cfg()) - alloc.init_layout(*_LAYOUT) - obj = MagicMock() - - adapter.allocate.return_value = obj - assert alloc.allocate([torch.Size([4, 8])], [torch.float16]) is obj - - alloc.get_by_location(1, 2, 64) - adapter.get_by_location.assert_called_once_with( - region_id=1, page_index=2, actual_size=64, single_token_size=16 - ) - - alloc.create_store_handle(obj) - adapter.create_store_handle.assert_called_once_with(obj) - - # abort_alloc returns the page via the adapter's real free - alloc.abort_alloc(obj) - adapter.free.assert_called_once_with(obj) - - # free/batched_free are no-ops (lifecycle owned by MaruServer) - adapter.free.reset_mock() - alloc.free(obj) - alloc.batched_free([obj]) - adapter.free.assert_not_called() - - -def test_close_is_idempotent(maru_mocks): - handler, adapter, _ = maru_mocks - alloc = MaruMemoryAllocator(_cfg()) - alloc.init_layout(*_LAYOUT) - alloc.close() - adapter.close.assert_called_once() - handler.close.assert_called_once() - assert not alloc.is_initialized - alloc.close() # safe to call again From c2f6cc771ff5c4f25effa433068888f41b98fea3 Mon Sep 17 00:00:00 2001 From: seohui-XCENA Date: Tue, 7 Jul 2026 20:22:28 +0900 Subject: [PATCH 22/28] [MP][Maru] Document the auto_expand flag in the MP configuration guide - Add --maru-auto-expand / --no-maru-auto-expand to the Maru CXL Shared L1 argument table. - Tighten the single-CXL-device pool-size note and align the tier wording with the design doc (cross-node). Signed-off-by: seohui-XCENA --- docs/source/mp/configuration.rst | 26 +++++++++++--------------- 1 file changed, 11 insertions(+), 15 deletions(-) diff --git a/docs/source/mp/configuration.rst b/docs/source/mp/configuration.rst index ea2cc1f587b..737051ff842 100644 --- a/docs/source/mp/configuration.rst +++ b/docs/source/mp/configuration.rst @@ -304,7 +304,7 @@ Maru CXL Shared L1 Source: ``lmcache/v1/distributed/config.py`` Opt-in. Setting ``--maru-server-url`` switches the L1 medium from pinned DRAM to -a **cross-instance shared CXL pool** managed by `Maru +a **cross-node shared CXL pool** managed by `Maru `_. Every LMCache server on the same pool ``mmap``\ s the same physical memory, so an L1 entry produced by one instance is a zero-copy read for the others. The pinned-DRAM L1 options @@ -347,8 +347,11 @@ path: * - ``--maru-pool-size-gb`` - ``0.0`` - CXL pool size to request from MaruServer, in GB. Required (> 0) when - ``--maru-server-url`` is set. **Must fit within a single CXL device** — - see the note below. + ``--maru-server-url`` is set. Bounded by a single CXL device (see note). + * - ``--maru-auto-expand`` / ``--no-maru-auto-expand`` + - ``True`` + - Grow the owned pool into free CXL space when it fills. ``--no-`` hard-caps + the pool at ``--maru-pool-size-gb`` for deterministic eviction. * - ``--maru-instance-id`` - auto UUID - Stable client id reported to MaruServer for ownership tracking. @@ -356,18 +359,11 @@ path: .. note:: - **``--maru-pool-size-gb`` is bounded by a single CXL device.** MaruServer - allocates each region as **one contiguous extent within a single DAX - device** — regions do not span devices. This value sizes both the initial - region and every auto-expand step, so it must be ≤ one device's usable - capacity (raw device size minus MaruServer metadata / WAL / alignment - overhead; e.g. a 256 GiB device holds somewhat less). Requesting more than - any device can satisfy fails the MaruServer allocation, and the LMCache - server then **fails to start** (the initial ``MaruHandler.connect()`` returns - an error). With multiple CXL devices, auto-expand can claim additional - regions from other devices, so total capacity is the sum across devices — - but no single region (hence no single ``--maru-pool-size-gb``) may exceed - one device. + ``--maru-pool-size-gb`` (and each auto-expand step) is one contiguous extent + within a single CXL device — regions do not span devices, so it must be ≤ one + device's usable capacity; a larger request fails MaruServer allocation and the + server does not start. Multiple devices raise total capacity (auto-expand + claims regions from others), but no single region may exceed one device. L1 Manager TTLs ---------------- From 0da8b6be1ecaf65c1baa0d7e32008203d5f6940f Mon Sep 17 00:00:00 2001 From: seohui-XCENA Date: Tue, 7 Jul 2026 20:25:34 +0900 Subject: [PATCH 23/28] [MP][Maru] Keep "watermark-triggered" wording in the architecture diagram Restore the word dropped when MaruL1Manager was added to the EvictionController line, for consistency with the rest of index.rst. Signed-off-by: seohui-XCENA --- docs/source/mp/index.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/mp/index.rst b/docs/source/mp/index.rst index 1ab9f8292dd..d7f42703ddc 100644 --- a/docs/source/mp/index.rst +++ b/docs/source/mp/index.rst @@ -86,7 +86,7 @@ High-Level Architecture | |--- StoreController -----> L2 Adapter(s) (async L1->L2 push) |--- PrefetchController ---> L2 Adapter(s) (async L2->L1 load) - |--- EvictionController ----> L1Manager / MaruL1Manager (watermark eviction) + |--- EvictionController ----> L1Manager / MaruL1Manager (watermark-triggered eviction) | v EventBus + OTel providers (observability) From d29850159dcbf2a9afef881b92708692dcef89b4 Mon Sep 17 00:00:00 2001 From: seohui-XCENA Date: Tue, 7 Jul 2026 20:37:20 +0900 Subject: [PATCH 24/28] [MP][Maru] Fix the Maru L1 CLI example (add required --l1-size-gb) - Add the required --l1-size-gb 0 to the quickstart command (ignored under Maru but the parser requires it). - Note that the example assumes a running MaruServer and that higher-level launchers derive --maru-server-url. Signed-off-by: seohui-XCENA --- docs/source/mp/configuration.rst | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/docs/source/mp/configuration.rst b/docs/source/mp/configuration.rst index 737051ff842..801c5554b26 100644 --- a/docs/source/mp/configuration.rst +++ b/docs/source/mp/configuration.rst @@ -320,9 +320,15 @@ path: lmcache server \ --maru-server-url maru://localhost:5555 \ --maru-pool-size-gb 4 \ + --l1-size-gb 0 \ --supported-transfer-mode lmcache_driven \ --eviction-policy LRU --max-workers 4 --port 6555 +``--l1-size-gb`` is required by the parser but ignored under Maru (pass ``0``, as +the CXL pool replaces the DRAM L1). This assumes a MaruServer is already running +at ``--maru-server-url``; higher-level launchers typically start the MaruServer +and derive this URL for you. + .. note:: Maru L1 constraints (rejected at startup otherwise): **lmcache-driven** From c04788a3810310a5310d048acbf7223cdec2c6bc Mon Sep 17 00:00:00 2001 From: seohui-XCENA Date: Wed, 8 Jul 2026 21:42:47 +0900 Subject: [PATCH 25/28] [MP][Maru] Register the L1 fullness gauges in MaruL1Manager - lmcache_mp.l1_memory_usage_bytes / l1_usage_ratio are registered in L1Manager.__init__, but the maru path builds MaruL1Manager instead of L1Manager (mutually exclusive), so both gauges silently vanished whenever maru is the L1 backend - mirror the stock registration in MaruL1Manager.__init__: same meter and gauge names, singleton dispatch via _gauge_registered/_gauge_target, and a self-contained ratio helper (no reach into the stock private helper) - conftest: stub maru's register_gauge for distributed tests -- OTel honors only the first registration of a gauge name, so a test-built MaruL1Manager could otherwise claim the name and starve the stock gauge assertions in test_l1_l2_state_metrics.py - test: both gauges registered at construction and their callbacks read the live manager, not the zero fallback Signed-off-by: seohui-XCENA --- lmcache/v1/distributed/maru_l1_manager.py | 53 ++++++++++++++++++++ tests/v1/distributed/conftest.py | 26 ++++++++++ tests/v1/distributed/test_maru_l1_manager.py | 38 ++++++++++++++ 3 files changed, 117 insertions(+) diff --git a/lmcache/v1/distributed/maru_l1_manager.py b/lmcache/v1/distributed/maru_l1_manager.py index 3b99610f376..92175e1d2eb 100644 --- a/lmcache/v1/distributed/maru_l1_manager.py +++ b/lmcache/v1/distributed/maru_l1_manager.py @@ -61,6 +61,7 @@ from lmcache.v1.memory_management import MemoryFormat, MemoryObj from lmcache.v1.mp_observability.event import Event, EventType from lmcache.v1.mp_observability.event_bus import get_event_bus +from lmcache.v1.mp_observability.otel_init import register_gauge if TYPE_CHECKING: # Third Party @@ -125,6 +126,22 @@ def _clamp_extra_count(extra_count: int) -> int: return extra_count +def _maru_l1_usage_ratio_or_zero(target: "MaruL1Manager | None") -> float: + """Return ``target.get_memory_usage()`` as a 0.0-1.0 ratio. + + PARITY(L1Manager._l1_usage_ratio_or_zero): duplicated here rather than + imported so maru stays self-contained and does not reach into a private + upstream helper. Returns 0.0 when ``target`` is None or ``total_bytes`` is + zero so the observable-gauge callback never raises during scrape. + """ + if target is None: + return 0.0 + used, total = target.get_memory_usage() + if total <= 0: + return 0.0 + return used / total + + @dataclass class _PendingRead: """A pinned read staged between reserve_read and the last finish_read. @@ -168,6 +185,16 @@ class MaruL1Manager: them against the MaruServer directory and the CXL allocator. """ + # PARITY(L1Manager): singleton dispatch for the L1 fullness gauges. The + # OTel SDK honors only the first registration of a gauge name, so register + # once (``_gauge_registered``) and route the callback to the most recently + # built instance (``_gauge_target``). A real deployment has one + # MaruL1Manager; the indirection just keeps multi-instance tests (which + # share the process-wide meter) reading a live target instead of a stale + # one. + _gauge_registered: bool = False + _gauge_target: "MaruL1Manager | None" = None + def __init__(self, config: L1ManagerConfig) -> None: maru_config = config.memory_config.maru_config if maru_config is None: @@ -199,6 +226,32 @@ def __init__(self, config: L1ManagerConfig) -> None: ) self._sweeper.start() + # PARITY(L1Manager): expose the same L1 fullness gauges. Upstream these + # are registered in L1Manager.__init__, but on the maru path + # StorageManager builds a MaruL1Manager *instead of* an L1Manager + # (they are mutually exclusive), so without this the metric would + # silently vanish whenever maru is the L1 backend. Same meter/gauge + # names as upstream so consumers see one metric regardless of backend. + MaruL1Manager._gauge_target = self + if not MaruL1Manager._gauge_registered: + MaruL1Manager._gauge_registered = True + register_gauge( + "lmcache.l1_manager", + "lmcache_mp.l1_memory_usage_bytes", + "Bytes currently held in L1 cache", + lambda: ( + MaruL1Manager._gauge_target.get_memory_usage()[0] + if MaruL1Manager._gauge_target is not None + else 0 + ), + ) + register_gauge( + "lmcache.l1_manager", + "lmcache_mp.l1_usage_ratio", + "L1 used/total ratio (0.0–1.0)", + lambda: _maru_l1_usage_ratio_or_zero(MaruL1Manager._gauge_target), + ) + def register_listener(self, listener: L1ManagerListener) -> None: """Register a listener for ``on_l1_keys_*`` events. diff --git a/tests/v1/distributed/conftest.py b/tests/v1/distributed/conftest.py index b325073de1c..ff0cc0aac67 100644 --- a/tests/v1/distributed/conftest.py +++ b/tests/v1/distributed/conftest.py @@ -12,6 +12,9 @@ import sys import types +# Third Party +import pytest + if find_spec("lmcache.native_storage_ops") is None: class Bitmap: @@ -89,3 +92,26 @@ class TTLLock: fallback_module_any.Bitmap = Bitmap fallback_module_any.TTLLock = TTLLock sys.modules["lmcache.native_storage_ops"] = fallback_module + + +@pytest.fixture(autouse=True) +def _stub_maru_gauge_registration(monkeypatch: pytest.MonkeyPatch) -> None: + """Keep test-built ``MaruL1Manager``s off the process-global OTel meter. + + OTel honors only the FIRST observable-gauge registration for a given + name; later registrations are silently dropped. Production never builds + both L1 backends in one process (StorageManager picks exactly one), but + tests build both — a ``MaruL1Manager`` constructed before the first + stock ``L1Manager`` would permanently claim + ``lmcache_mp.l1_memory_usage_bytes`` / ``lmcache_mp.l1_usage_ratio`` and + starve the stock gauge assertions in ``test_l1_l2_state_metrics.py``. + Stub the maru-side registration hook so only the stock backend registers + on the shared meter; tests that assert maru's gauge wiring monkeypatch + their own recording fake over this stub. + """ + try: + # First Party + from lmcache.v1.distributed import maru_l1_manager + except ImportError: + return + monkeypatch.setattr(maru_l1_manager, "register_gauge", lambda *args, **kwargs: None) diff --git a/tests/v1/distributed/test_maru_l1_manager.py b/tests/v1/distributed/test_maru_l1_manager.py index ef2d18e702d..01ef3572a19 100644 --- a/tests/v1/distributed/test_maru_l1_manager.py +++ b/tests/v1/distributed/test_maru_l1_manager.py @@ -482,6 +482,44 @@ def test_get_memory_usage_watermark_tracks_auto_expand(): assert mgr_off.get_memory_usage() == (chunk, own_pool) +def test_registers_l1_fullness_gauges(monkeypatch): + """MaruL1Manager must register the L1 fullness gauges itself. + + On the maru path StorageManager builds a MaruL1Manager instead of an + L1Manager (they are mutually exclusive), and the + ``lmcache_mp.l1_memory_usage_bytes`` / ``lmcache_mp.l1_usage_ratio`` gauges + live in L1Manager.__init__ upstream. Without registering them here the + metrics would silently vanish whenever maru is the L1 backend. + """ + # First Party + import lmcache.v1.distributed.maru_l1_manager as maru_mod + + # Reset the process-wide "registered once" guard so this construction is + # observed regardless of managers built earlier in the session. + monkeypatch.setattr(maru_mod.MaruL1Manager, "_gauge_registered", False) + registered: dict[str, object] = {} + monkeypatch.setattr( + maru_mod, + "register_gauge", + lambda meter, name, desc, func: registered.__setitem__(name, func), + ) + + chunk = 64 + manager, handler, _ = make_maru_manager(chunk_size=chunk) + handler.cxl_free = 100 * chunk + _seed(handler, _key(1)) # one allocated page -> used == chunk + + assert "lmcache_mp.l1_memory_usage_bytes" in registered + assert "lmcache_mp.l1_usage_ratio" in registered + + used, total = manager.get_memory_usage() + assert used == chunk + assert total > 0 + # Callbacks read the live manager, not the 0 fallback. + assert registered["lmcache_mp.l1_memory_usage_bytes"]() == used + assert registered["lmcache_mp.l1_usage_ratio"]() == used / total + + def test_is_key_evictable_tracks_local_staging(): manager, handler, _ = make_maru_manager() k = _key(1) From 13885d83d855367b4c1bc4d24ada60c260e286d5 Mon Sep 17 00:00:00 2001 From: seohui-XCENA Date: Thu, 9 Jul 2026 09:18:17 +0000 Subject: [PATCH 26/28] [MP][Maru] Gate the register_kv_layout hook on the maru L1 backend Signed-off-by: seohui-XCENA --- lmcache/v1/distributed/storage_manager.py | 13 ++++++ .../modules/lmcache_driven_transfer.py | 41 ++++++++++--------- .../test_distributed_storage_manager.py | 2 + .../test_lmcache_driven_layout_registry.py | 2 + .../test_maru_register_kv_layout.py | 18 ++++---- 5 files changed, 48 insertions(+), 28 deletions(-) diff --git a/lmcache/v1/distributed/storage_manager.py b/lmcache/v1/distributed/storage_manager.py index 93f5c20a8df..53a0efb4026 100644 --- a/lmcache/v1/distributed/storage_manager.py +++ b/lmcache/v1/distributed/storage_manager.py @@ -179,6 +179,18 @@ def __init__(self, config: StorageManagerConfig): ) # External APIs for serving engine integration code to call + @property + def requires_kv_layout_registration(self) -> bool: + """Whether the engine must call register_kv_layout after KV registration. + + True for the maru L1 tier, which defers pool creation until the KV + layout is known; False for stock L1, which sizes its pool from config + at construction. Callers should skip the register_kv_layout path + entirely (including any layout/format computation feeding it) when + this is False. + """ + return isinstance(self._l1_manager, MaruL1Manager) + def register_kv_layout( self, layout_desc: MemoryLayoutDesc, @@ -202,6 +214,7 @@ def register_kv_layout( ValueError: The maru L1 tier supports a single object group only (multi-model support is a maru TODO). """ + # isinstance (not the property) so mypy narrows for the call below. if not isinstance(self._l1_manager, MaruL1Manager): return if num_object_groups > 1: diff --git a/lmcache/v1/multiprocess/modules/lmcache_driven_transfer.py b/lmcache/v1/multiprocess/modules/lmcache_driven_transfer.py index 557a351810b..eeb36fbb50c 100644 --- a/lmcache/v1/multiprocess/modules/lmcache_driven_transfer.py +++ b/lmcache/v1/multiprocess/modules/lmcache_driven_transfer.py @@ -871,27 +871,28 @@ def register_kv_cache( model_name, world_size, layout_desc, attn_desc ) - # Bring up the maru CXL pool now that the KV layout is known. A no-op - # for the stock L1 backend (StorageManager self-gates on the maru - # config); maru's register_kv_layout is idempotent across instances. - fmt = ( - MemoryFormat.KV_MLA_FMT - if is_mla(cache_context.get_engine_kv_format(0)) - else MemoryFormat.KV_2LTD - ) - try: - self._ctx.storage_manager.register_kv_layout( - layout_desc, - fmt, - self._ctx.chunk_size, - attn_desc.num_object_groups, + # Bring up the maru CXL pool now that the KV layout is known. Skipped + # entirely (including the format probe) for the stock L1 backend; + # maru's register_kv_layout is idempotent across instances. + if self._ctx.storage_manager.requires_kv_layout_registration: + fmt = ( + MemoryFormat.KV_MLA_FMT + if is_mla(cache_context.get_engine_kv_format(0)) + else MemoryFormat.KV_2LTD ) - except Exception: - # e.g. maru rejects >1 object group -- drop the context we just - # built so the instance is not left half-registered, then surface - # the error on the REGISTER path. - cache_context.close() - raise + try: + self._ctx.storage_manager.register_kv_layout( + layout_desc, + fmt, + self._ctx.chunk_size, + attn_desc.num_object_groups, + ) + except Exception: + # e.g. maru rejects >1 object group -- drop the context we + # just built so the instance is not left half-registered, + # then surface the error on the REGISTER path. + cache_context.close() + raise with self._lock: self._cache_contexts[instance_id] = ContextEntry( diff --git a/tests/v1/distributed/test_distributed_storage_manager.py b/tests/v1/distributed/test_distributed_storage_manager.py index e921dbe0d84..1f61c1f3e8c 100644 --- a/tests/v1/distributed/test_distributed_storage_manager.py +++ b/tests/v1/distributed/test_distributed_storage_manager.py @@ -937,6 +937,7 @@ def test_maru_config_selects_maru_l1_manager(self, maru_storage_manager_config): sm = StorageManager(maru_storage_manager_config) try: assert isinstance(sm._l1_manager, MaruL1Manager) + assert sm.requires_kv_layout_registration finally: sm.close() @@ -944,6 +945,7 @@ def test_stock_config_selects_stock_l1_manager(self, basic_storage_manager_confi sm = StorageManager(basic_storage_manager_config) try: assert isinstance(sm._l1_manager, L1Manager) + assert not sm.requires_kv_layout_registration finally: sm.close() diff --git a/tests/v1/multiprocess/test_lmcache_driven_layout_registry.py b/tests/v1/multiprocess/test_lmcache_driven_layout_registry.py index 7b362b4f65d..2a9e307b526 100644 --- a/tests/v1/multiprocess/test_lmcache_driven_layout_registry.py +++ b/tests/v1/multiprocess/test_lmcache_driven_layout_registry.py @@ -86,6 +86,8 @@ def test_unregister_one_shared_gpu_layout_keeps_registry_until_last_instance( ctx = MagicMock() ctx.chunk_size = 16 ctx.layout_desc_registry = LayoutDescRegistry() + # Stock L1: the maru pool bring-up hook must be skipped entirely. + ctx.storage_manager.requires_kv_layout_registration = False def fake_create_cache_context( kv_caches: object, diff --git a/tests/v1/multiprocess/test_maru_register_kv_layout.py b/tests/v1/multiprocess/test_maru_register_kv_layout.py index 3e23eeda116..ddc931ea1db 100644 --- a/tests/v1/multiprocess/test_maru_register_kv_layout.py +++ b/tests/v1/multiprocess/test_maru_register_kv_layout.py @@ -25,6 +25,8 @@ def _bare_module() -> LMCacheDrivenTransferModule: """A transfer module with only the state register_kv_cache touches.""" module = LMCacheDrivenTransferModule.__new__(LMCacheDrivenTransferModule) module._ctx = MagicMock(name="ctx") + # Maru L1 by default; stock tests flip this off. + module._ctx.storage_manager.requires_kv_layout_registration = True module._cache_contexts = {} module._lock = threading.Lock() return module @@ -81,17 +83,17 @@ def test_hook_failure_closes_context_and_does_not_register(monkeypatch): assert 1 not in module._cache_contexts # not left half-registered -def test_hook_is_noop_style_for_stock_backend(monkeypatch): - """For stock, register_kv_layout is a no-op; the instance still registers.""" - monkeypatch.setattr( - gpu_mod, "create_cache_context", lambda *a, **kw: MagicMock(num_layers=2) - ) +def test_hook_skipped_entirely_for_stock_backend(monkeypatch): + """For stock, the hook (format probe included) is skipped; registration + still completes.""" + cache_ctx = MagicMock(num_layers=2, name="cache_ctx") + monkeypatch.setattr(gpu_mod, "create_cache_context", lambda *a, **kw: cache_ctx) monkeypatch.setattr(gpu_mod, "get_layout_desc", lambda *a, **kw: MagicMock()) - monkeypatch.setattr(gpu_mod, "is_mla", lambda fmt: False) module = _bare_module() - # Stock StorageManager.register_kv_layout returns None (no-op). - module._ctx.storage_manager.register_kv_layout.return_value = None + module._ctx.storage_manager.requires_kv_layout_registration = False _register(module) + cache_ctx.get_engine_kv_format.assert_not_called() + module._ctx.storage_manager.register_kv_layout.assert_not_called() assert 1 in module._cache_contexts # registration completed normally From 8947b86f5864ed9b8e7745fa2caaec24136c8c23 Mon Sep 17 00:00:00 2001 From: seohui-XCENA Date: Fri, 10 Jul 2026 07:20:47 +0000 Subject: [PATCH 27/28] [MP][Maru] Skip the maru pool bring-up hook in worker liveness tests Signed-off-by: seohui-XCENA --- tests/v1/multiprocess/test_worker_liveness.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/v1/multiprocess/test_worker_liveness.py b/tests/v1/multiprocess/test_worker_liveness.py index 80947c8ce0b..bc0348baef2 100644 --- a/tests/v1/multiprocess/test_worker_liveness.py +++ b/tests/v1/multiprocess/test_worker_liveness.py @@ -39,6 +39,8 @@ def _bare_gpu_module() -> LMCacheDrivenTransferModule: """ module = LMCacheDrivenTransferModule.__new__(LMCacheDrivenTransferModule) module._ctx = MagicMock(name="ctx") + # Stock L1: skip the maru pool bring-up hook in register_kv_cache. + module._ctx.storage_manager.requires_kv_layout_registration = False module._cache_contexts = {} module._lock = threading.Lock() return module From 092f07035debe31068dce09b0840638c9fbc1daf Mon Sep 17 00:00:00 2001 From: seohui-XCENA Date: Tue, 21 Jul 2026 20:53:42 +0900 Subject: [PATCH 28/28] [MP][Maru] Move engine KV format mapping into MaruL1Manager - Engine hook forwards the raw engine KV format unconditionally (stock no-op) - Map EngineKVFormat to MemoryFormat inside MaruL1Manager.register_kv_layout - Drop requires_kv_layout_registration and the mock pins it forced in tests Signed-off-by: seohui-XCENA --- docs/design/v1/distributed/maru_l1.md | 5 ++ lmcache/v1/distributed/maru_l1_manager.py | 17 +++- lmcache/v1/distributed/storage_manager.py | 41 +++++----- .../modules/lmcache_driven_transfer.py | 41 +++++----- .../test_distributed_storage_manager.py | 23 +++--- tests/v1/distributed/test_maru_l1_manager.py | 27 +++++++ .../test_lmcache_driven_layout_registry.py | 6 +- .../test_maru_register_kv_layout.py | 77 +++++++++---------- tests/v1/multiprocess/test_worker_liveness.py | 2 - 9 files changed, 138 insertions(+), 101 deletions(-) diff --git a/docs/design/v1/distributed/maru_l1.md b/docs/design/v1/distributed/maru_l1.md index e4ecc1eb5cf..ca7a1a31b8c 100644 --- a/docs/design/v1/distributed/maru_l1.md +++ b/docs/design/v1/distributed/maru_l1.md @@ -246,6 +246,11 @@ with the pinned-DRAM / DevDax / GDS tiers): `register_kv_layout` binds the KV layout to the allocator (pool bring-up) on the first `register_kv_cache`; it is wired only on the `lmcache_driven` transfer path. +The engine hook calls `StorageManager.register_kv_layout` unconditionally with the +raw group-0 engine KV format — a silent no-op for the stock L1 backend, so the +hook carries no maru-conditional logic (and never feeds mocked values into native +code under test). The maru backend maps the engine format to its memory format +internally (MLA layouts bind as `KV_MLA_FMT`, everything else as `KV_2LTD`). ### Startup guards diff --git a/lmcache/v1/distributed/maru_l1_manager.py b/lmcache/v1/distributed/maru_l1_manager.py index 3a9fe90e18c..a2380bd651d 100644 --- a/lmcache/v1/distributed/maru_l1_manager.py +++ b/lmcache/v1/distributed/maru_l1_manager.py @@ -58,6 +58,7 @@ from lmcache.v1.distributed.memory_manager.maru_memory_allocator import ( MaruMemoryAllocator, ) +from lmcache.v1.gpu_connector.utils import is_mla from lmcache.v1.memory_management import MemoryFormat, MemoryObj from lmcache.v1.mp_observability.event import Event, EventType from lmcache.v1.mp_observability.event_bus import get_event_bus @@ -68,6 +69,9 @@ from maru import MaruHandler from maru_handler.memory import MemoryInfo + # First Party + import lmcache.c_ops as lmc_ops + logger = init_logger(__name__) P = ParamSpec("P") @@ -1173,15 +1177,24 @@ def register_kv_layout( self, shapes: list[torch.Size], dtypes: list[torch.dtype], - fmt: MemoryFormat, + engine_kv_format: "lmc_ops.EngineKVFormat", chunk_size_in_tokens: int, ) -> None: """Bind the KV layout, bringing up the CXL pool (idempotent per layout). + Maps the engine's KV format to the maru memory format here (MLA + layouts store as KV_MLA_FMT, everything else as KV_2LTD) so engine- + side callers forward the raw format without touching the probe. + Args: shapes: Per-group tensor shapes of one chunk. dtypes: Per-group dtypes of one chunk. - fmt: The KV memory format. + engine_kv_format: The engine's KV format for the layout. chunk_size_in_tokens: Tokens per chunk. """ + fmt = ( + MemoryFormat.KV_MLA_FMT + if is_mla(engine_kv_format) + else MemoryFormat.KV_2LTD + ) self._allocator.init_layout(shapes, dtypes, fmt, chunk_size_in_tokens) diff --git a/lmcache/v1/distributed/storage_manager.py b/lmcache/v1/distributed/storage_manager.py index dc7acf79032..f1eba82b385 100644 --- a/lmcache/v1/distributed/storage_manager.py +++ b/lmcache/v1/distributed/storage_manager.py @@ -5,10 +5,14 @@ # Standard from contextlib import contextmanager -from typing import Iterator, Literal, Optional +from typing import TYPE_CHECKING, Iterator, Literal, Optional import threading import time +if TYPE_CHECKING: + # First Party + import lmcache.c_ops as lmc_ops + # First Party from lmcache.logging import init_logger from lmcache.native_storage_ops import Bitmap, PeriodicEventNotifier @@ -51,7 +55,7 @@ AdapterDescriptor, create_store_policy, ) -from lmcache.v1.memory_management import MemoryFormat, MemoryObj +from lmcache.v1.memory_management import MemoryObj from lmcache.v1.mp_observability.errors import LMCacheTimeoutError from lmcache.v1.mp_observability.event import Event, EventType from lmcache.v1.mp_observability.event_bus import get_event_bus @@ -179,34 +183,25 @@ def __init__(self, config: StorageManagerConfig): ) # External APIs for serving engine integration code to call - @property - def requires_kv_layout_registration(self) -> bool: - """Whether the engine must call register_kv_layout after KV registration. - - True for the maru L1 tier, which defers pool creation until the KV - layout is known; False for stock L1, which sizes its pool from config - at construction. Callers should skip the register_kv_layout path - entirely (including any layout/format computation feeding it) when - this is False. - """ - return isinstance(self._l1_manager, MaruL1Manager) - def register_kv_layout( self, layout_desc: MemoryLayoutDesc, - fmt: MemoryFormat, + engine_kv_format: "lmc_ops.EngineKVFormat", chunk_size_in_tokens: int, num_object_groups: int, ) -> None: """Bind the KV layout to the L1 backend (maru CXL pool bring-up). - Stock L1 sizes its pool from config at construction and ignores this; - the maru L1 tier defers pool creation until the layout is known, so the - engine calls this once the KV cache is registered. A no-op for stock. + Stock L1 sizes its pool from config at construction, so this is a + silent no-op for it (``engine_kv_format`` is never dereferenced); the + maru L1 tier defers pool creation until the layout is known and maps + the engine format to its memory format internally. The engine calls + this unconditionally once the KV cache is registered. Args: layout_desc: Per-group chunk shapes and dtypes. - fmt: The KV memory format. + engine_kv_format: The engine's KV format for object group 0, + forwarded verbatim to the maru backend. chunk_size_in_tokens: Tokens per chunk. num_object_groups: Number of KV object groups in the layout. @@ -214,7 +209,8 @@ def register_kv_layout( ValueError: The maru L1 tier supports a single object group only (multi-model support is a maru TODO). """ - # isinstance (not the property) so mypy narrows for the call below. + # isinstance narrowing: register_kv_layout is maru-only, not part of + # L1ManagerInterface. if not isinstance(self._l1_manager, MaruL1Manager): return if num_object_groups > 1: @@ -223,7 +219,10 @@ def register_kv_layout( f"{num_object_groups} (multi-model support is a maru TODO)" ) self._l1_manager.register_kv_layout( - layout_desc.shapes, layout_desc.dtypes, fmt, chunk_size_in_tokens + layout_desc.shapes, + layout_desc.dtypes, + engine_kv_format, + chunk_size_in_tokens, ) @enable_tracing() diff --git a/lmcache/v1/multiprocess/modules/lmcache_driven_transfer.py b/lmcache/v1/multiprocess/modules/lmcache_driven_transfer.py index 271d4d13014..eaf3481bb91 100644 --- a/lmcache/v1/multiprocess/modules/lmcache_driven_transfer.py +++ b/lmcache/v1/multiprocess/modules/lmcache_driven_transfer.py @@ -28,9 +28,9 @@ lmcache_memcpy_async_d2h, lmcache_memcpy_async_h2d, ) -from lmcache.v1.gpu_connector.utils import LayoutHints, is_mla +from lmcache.v1.gpu_connector.utils import LayoutHints from lmcache.v1.memory_allocators.lazy_memory_allocator import LazyMemoryAllocator -from lmcache.v1.memory_management import GDSMemoryObject, MemoryFormat, MemoryObj +from lmcache.v1.memory_management import GDSMemoryObject, MemoryObj from lmcache.v1.mp_observability.event import Event, EventType from lmcache.v1.multiprocess.custom_types import ( IPCCacheServerKey, @@ -890,28 +890,23 @@ def register_kv_cache( model_name, world_size, layout_desc, attn_desc ) - # Bring up the maru CXL pool now that the KV layout is known. Skipped - # entirely (including the format probe) for the stock L1 backend; - # maru's register_kv_layout is idempotent across instances. - if self._ctx.storage_manager.requires_kv_layout_registration: - fmt = ( - MemoryFormat.KV_MLA_FMT - if is_mla(cache_context.get_engine_kv_format(0)) - else MemoryFormat.KV_2LTD + # Bring up the maru CXL pool now that the KV layout is known; a no-op + # for the stock L1 backend. The raw engine format is forwarded as-is + # (maru maps it to its memory format internally) so this hook carries + # no maru-conditional logic. Idempotent across instances. + try: + self._ctx.storage_manager.register_kv_layout( + layout_desc, + cache_context.get_engine_kv_format(0), + self._ctx.chunk_size, + attn_desc.num_object_groups, ) - try: - self._ctx.storage_manager.register_kv_layout( - layout_desc, - fmt, - self._ctx.chunk_size, - attn_desc.num_object_groups, - ) - except Exception: - # e.g. maru rejects >1 object group -- drop the context we - # just built so the instance is not left half-registered, - # then surface the error on the REGISTER path. - cache_context.close() - raise + except Exception: + # e.g. maru rejects >1 object group -- drop the context we + # just built so the instance is not left half-registered, + # then surface the error on the REGISTER path. + cache_context.close() + raise with self._lock: self._cache_contexts[instance_id] = ContextEntry( diff --git a/tests/v1/distributed/test_distributed_storage_manager.py b/tests/v1/distributed/test_distributed_storage_manager.py index da1207e993b..e801b665b3a 100644 --- a/tests/v1/distributed/test_distributed_storage_manager.py +++ b/tests/v1/distributed/test_distributed_storage_manager.py @@ -29,7 +29,6 @@ L2AdaptersConfig, ) from lmcache.v1.distributed.l2_adapters.mock_l2_adapter import MockL2AdapterConfig -from lmcache.v1.memory_management import MemoryFormat from lmcache.v1.mp_observability.event import Event, EventType from lmcache.v1.mp_observability.event_bus import EventBusConfig, init_event_bus @@ -940,7 +939,6 @@ def test_maru_config_selects_maru_l1_manager(self, maru_storage_manager_config): sm = StorageManager(maru_storage_manager_config) try: assert isinstance(sm._l1_manager, MaruL1Manager) - assert sm.requires_kv_layout_registration finally: sm.close() @@ -948,7 +946,6 @@ def test_stock_config_selects_stock_l1_manager(self, basic_storage_manager_confi sm = StorageManager(basic_storage_manager_config) try: assert isinstance(sm._l1_manager, L1Manager) - assert not sm.requires_kv_layout_registration finally: sm.close() @@ -959,11 +956,10 @@ def test_register_kv_layout_forwards_to_maru( try: # Spy the forward so the test needs no maru runtime (init_layout). sm._l1_manager.register_kv_layout = MagicMock() # type: ignore[method-assign] - sm.register_kv_layout( - basic_layout, MemoryFormat.KV_MLA_FMT, 16, num_object_groups=1 - ) + engine_fmt = MagicMock(name="engine_kv_format") + sm.register_kv_layout(basic_layout, engine_fmt, 16, num_object_groups=1) sm._l1_manager.register_kv_layout.assert_called_once_with( - basic_layout.shapes, basic_layout.dtypes, MemoryFormat.KV_MLA_FMT, 16 + basic_layout.shapes, basic_layout.dtypes, engine_fmt, 16 ) finally: sm.close() @@ -975,7 +971,10 @@ def test_register_kv_layout_rejects_multi_object_group( try: with pytest.raises(ValueError): sm.register_kv_layout( - basic_layout, MemoryFormat.KV_MLA_FMT, 16, num_object_groups=2 + basic_layout, + MagicMock(name="engine_kv_format"), + 16, + num_object_groups=2, ) finally: sm.close() @@ -986,9 +985,13 @@ def test_register_kv_layout_noop_for_stock( sm = StorageManager(basic_storage_manager_config) try: # Stock sizes its pool at construction: this must be a silent no-op - # (no forward, no raise) even for a would-be-rejected group count. + # (no forward, no raise, format never dereferenced) even for a + # would-be-rejected group count. sm.register_kv_layout( - basic_layout, MemoryFormat.KV_MLA_FMT, 16, num_object_groups=2 + basic_layout, + MagicMock(name="engine_kv_format"), + 16, + num_object_groups=2, ) finally: sm.close() diff --git a/tests/v1/distributed/test_maru_l1_manager.py b/tests/v1/distributed/test_maru_l1_manager.py index cc1ec28a994..9c7eca4e21f 100644 --- a/tests/v1/distributed/test_maru_l1_manager.py +++ b/tests/v1/distributed/test_maru_l1_manager.py @@ -46,6 +46,7 @@ from lmcache.v1.distributed.storage_manager import StorageManager from lmcache.v1.memory_management import MemoryFormat from lmcache.v1.mp_observability.event import EventType + import lmcache.c_ops as lmc_ops # Local from .maru_fakes import ( @@ -839,3 +840,29 @@ def test_eviction_deletes_from_maru_directory(): assert _wait(lambda: len(handler.store_map) < 6) finally: sm.close() + + +# ========================================================================= +# register_kv_layout: engine format -> memory format mapping +# ========================================================================= + + +def test_register_kv_layout_maps_engine_format_to_memory_format(): + """MLA engine layouts bind as KV_MLA_FMT; everything else as KV_2LTD. + + The mapping (and the pybind ``is_mla`` probe) lives here in the maru + backend so engine-side callers forward the raw engine format untouched. + """ + manager, _, _ = make_maru_manager() + manager._allocator.init_layout = MagicMock() # type: ignore[method-assign] + shapes, dtypes = [torch.Size([4, 8])], [torch.float16] + + manager.register_kv_layout(shapes, dtypes, lmc_ops.EngineKVFormat.NL_X_NB_BS_HS, 16) + fmt = manager._allocator.init_layout.call_args.args[2] + assert fmt == MemoryFormat.KV_MLA_FMT + + manager.register_kv_layout( + shapes, dtypes, lmc_ops.EngineKVFormat.NB_NL_TWO_BS_NH_HS, 16 + ) + fmt = manager._allocator.init_layout.call_args.args[2] + assert fmt == MemoryFormat.KV_2LTD diff --git a/tests/v1/multiprocess/test_lmcache_driven_layout_registry.py b/tests/v1/multiprocess/test_lmcache_driven_layout_registry.py index 70ad9c65116..aec8a797ada 100644 --- a/tests/v1/multiprocess/test_lmcache_driven_layout_registry.py +++ b/tests/v1/multiprocess/test_lmcache_driven_layout_registry.py @@ -31,6 +31,10 @@ class _FakeGPUContext: num_layers: int = 2 kv_layer_groups_manager: _FakeKVLayerGroupsManager = _FakeKVLayerGroupsManager() + def get_engine_kv_format(self, object_group_id: int) -> object: + """Return a sentinel engine KV format (forwarded, never dereferenced).""" + return object() + def close(self) -> None: """No-op teardown (real GPUCacheContext.close deregisters its GDS buffer).""" @@ -86,8 +90,6 @@ def test_unregister_one_shared_gpu_layout_keeps_registry_until_last_instance( ctx = MagicMock() ctx.chunk_size = 16 ctx.layout_desc_registry = LayoutDescRegistry() - # Stock L1: the maru pool bring-up hook must be skipped entirely. - ctx.storage_manager.requires_kv_layout_registration = False def fake_create_cache_context( kv_caches: object, diff --git a/tests/v1/multiprocess/test_maru_register_kv_layout.py b/tests/v1/multiprocess/test_maru_register_kv_layout.py index ddc931ea1db..1bf5b3d959e 100644 --- a/tests/v1/multiprocess/test_maru_register_kv_layout.py +++ b/tests/v1/multiprocess/test_maru_register_kv_layout.py @@ -1,9 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 -"""Tests for the maru register_kv_layout engine hook in register_kv_cache. - -The hook brings up the maru CXL pool once the KV layout is known. These tests -exercise it with a bare transfer module (bypassing the CUDA dispatcher in -__init__) and mocked context creation, mirroring test_worker_liveness.py. +"""Tests for the KV-layout bring-up hook in register_kv_cache. + +The hook forwards the layout and the raw group-0 engine KV format to +``StorageManager.register_kv_layout`` unconditionally — a silent no-op for the +stock L1 backend, a CXL pool bring-up for maru (which maps the engine format +to its memory format internally). These tests exercise the hook with a bare +transfer module (bypassing the CUDA dispatcher in __init__) and mocked context +creation, mirroring test_worker_liveness.py. The format mapping itself is +covered at the maru level in test_maru_l1_manager.py. """ # Standard @@ -14,7 +18,6 @@ import pytest # First Party -from lmcache.v1.memory_management import MemoryFormat from lmcache.v1.multiprocess.modules import lmcache_driven_transfer as gpu_mod from lmcache.v1.multiprocess.modules.lmcache_driven_transfer import ( LMCacheDrivenTransferModule, @@ -25,8 +28,6 @@ def _bare_module() -> LMCacheDrivenTransferModule: """A transfer module with only the state register_kv_cache touches.""" module = LMCacheDrivenTransferModule.__new__(LMCacheDrivenTransferModule) module._ctx = MagicMock(name="ctx") - # Maru L1 by default; stock tests flip this off. - module._ctx.storage_manager.requires_kv_layout_registration = True module._cache_contexts = {} module._lock = threading.Lock() return module @@ -36,33 +37,25 @@ def _register(module: LMCacheDrivenTransferModule) -> None: module.register_kv_cache(1, MagicMock(), "model", 1, MagicMock(), MagicMock(), []) -def test_hook_forwards_kv_2ltd_for_non_mla(monkeypatch): - monkeypatch.setattr( - gpu_mod, "create_cache_context", lambda *a, **kw: MagicMock(num_layers=2) - ) - monkeypatch.setattr(gpu_mod, "get_layout_desc", lambda *a, **kw: MagicMock()) - monkeypatch.setattr(gpu_mod, "is_mla", lambda fmt: False) - module = _bare_module() - - _register(module) - - call = module._ctx.storage_manager.register_kv_layout.call_args - # (layout_desc, fmt, chunk_size, num_object_groups) - assert call.args[1] == MemoryFormat.KV_2LTD - - -def test_hook_forwards_kv_mla_fmt_for_mla(monkeypatch): - monkeypatch.setattr( - gpu_mod, "create_cache_context", lambda *a, **kw: MagicMock(num_layers=2) - ) - monkeypatch.setattr(gpu_mod, "get_layout_desc", lambda *a, **kw: MagicMock()) - monkeypatch.setattr(gpu_mod, "is_mla", lambda fmt: True) +def test_hook_forwards_layout_and_raw_engine_format(monkeypatch): + """The hook forwards the layout and the group-0 engine format verbatim.""" + engine_fmt = MagicMock(name="engine_kv_format") + cache_ctx = MagicMock(num_layers=2, name="cache_ctx") + cache_ctx.get_engine_kv_format.return_value = engine_fmt + layout = MagicMock(name="layout_desc") + monkeypatch.setattr(gpu_mod, "create_cache_context", lambda *a, **kw: cache_ctx) + monkeypatch.setattr(gpu_mod, "get_layout_desc", lambda *a, **kw: layout) module = _bare_module() _register(module) + cache_ctx.get_engine_kv_format.assert_called_once_with(0) call = module._ctx.storage_manager.register_kv_layout.call_args - assert call.args[1] == MemoryFormat.KV_MLA_FMT + # (layout_desc, engine_kv_format, chunk_size, num_object_groups) + assert call.args[0] is layout + assert call.args[1] is engine_fmt + assert call.args[2] is module._ctx.chunk_size + assert 1 in module._cache_contexts # registration completed normally def test_hook_failure_closes_context_and_does_not_register(monkeypatch): @@ -70,7 +63,6 @@ def test_hook_failure_closes_context_and_does_not_register(monkeypatch): cache_ctx = MagicMock(num_layers=2, name="cache_ctx") monkeypatch.setattr(gpu_mod, "create_cache_context", lambda *a, **kw: cache_ctx) monkeypatch.setattr(gpu_mod, "get_layout_desc", lambda *a, **kw: MagicMock()) - monkeypatch.setattr(gpu_mod, "is_mla", lambda fmt: False) module = _bare_module() module._ctx.storage_manager.register_kv_layout.side_effect = ValueError( "maru L1 supports a single object group only" @@ -83,17 +75,20 @@ def test_hook_failure_closes_context_and_does_not_register(monkeypatch): assert 1 not in module._cache_contexts # not left half-registered -def test_hook_skipped_entirely_for_stock_backend(monkeypatch): - """For stock, the hook (format probe included) is skipped; registration - still completes.""" - cache_ctx = MagicMock(num_layers=2, name="cache_ctx") - monkeypatch.setattr(gpu_mod, "create_cache_context", lambda *a, **kw: cache_ctx) +def test_hook_is_inert_under_a_fully_mocked_context(monkeypatch): + """Upstream-style tests mock the whole ctx; the hook must stay harmless. + + Regression guard for the pybind ``is_mla()`` TypeError class: the hook + must never feed mocked values into native code. Tests that mock the whole + module context (test_worker_liveness.py, test_ipc_memory_reclaim.py, and + future upstream tests alike) rely on this. + """ + monkeypatch.setattr( + gpu_mod, "create_cache_context", lambda *a, **kw: MagicMock(num_layers=2) + ) monkeypatch.setattr(gpu_mod, "get_layout_desc", lambda *a, **kw: MagicMock()) module = _bare_module() - module._ctx.storage_manager.requires_kv_layout_registration = False - _register(module) + _register(module) # must not raise - cache_ctx.get_engine_kv_format.assert_not_called() - module._ctx.storage_manager.register_kv_layout.assert_not_called() - assert 1 in module._cache_contexts # registration completed normally + assert 1 in module._cache_contexts diff --git a/tests/v1/multiprocess/test_worker_liveness.py b/tests/v1/multiprocess/test_worker_liveness.py index bc0348baef2..80947c8ce0b 100644 --- a/tests/v1/multiprocess/test_worker_liveness.py +++ b/tests/v1/multiprocess/test_worker_liveness.py @@ -39,8 +39,6 @@ def _bare_gpu_module() -> LMCacheDrivenTransferModule: """ module = LMCacheDrivenTransferModule.__new__(LMCacheDrivenTransferModule) module._ctx = MagicMock(name="ctx") - # Stock L1: skip the maru pool bring-up hook in register_kv_cache. - module._ctx.storage_manager.requires_kv_layout_registration = False module._cache_contexts = {} module._lock = threading.Lock() return module