diff --git a/docs/design/v1/distributed/maru_l1.md b/docs/design/v1/distributed/maru_l1.md new file mode 100644 index 00000000000..ca7a1a31b8c --- /dev/null +++ b/docs/design/v1/distributed/maru_l1.md @@ -0,0 +1,296 @@ +# Maru CXL Shared L1 Design + +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. + +## Overview + +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. + +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). + +``` +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) +``` + +## 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: + +```python +self._l1_manager: L1ManagerInterface = ( + MaruL1Manager(cfg) if maru_config else L1Manager(cfg) +) +``` + +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. + +### `MaruMemoryAllocator` (`memory_manager/maru_memory_allocator.py`) + +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. + +## Operation Flow + +All three flows run on the stock controllers; `MaruL1Manager` only implements the +`L1ManagerInterface` calls they make. + +### Store (write-through) + +``` +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 +``` + +A key another instance already registered is dup-skipped by the server and its page +auto-freed. + +### Retrieve (hit, and miss → promote) + +``` +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 +``` + +### Eviction (watermark → delete → PINNED retry) + +``` +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. +``` + +Discard-on-evict is safe because write-through already placed the data in L2. + +## State-Machine Invariants + +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. + +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 — 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. +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 + +`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/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 5e0b54fa006..1e1526f6dff 100644 --- a/docs/source/mp/configuration.rst +++ b/docs/source/mp/configuration.rst @@ -298,6 +298,79 @@ 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-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 +(``--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 \ + --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** + 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. 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. + Auto-generated if omitted. + +.. note:: + + ``--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 ---------------- diff --git a/docs/source/mp/index.rst b/docs/source/mp/index.rst index 40e74dfbbb9..d7f42703ddc 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-triggered 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 ~~~~~~~~~~~ diff --git a/lmcache/v1/distributed/config.py b/lmcache/v1/distributed/config.py index de36535a57d..467bbe2bc91 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,44 @@ 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.""" + + 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: """ @@ -296,6 +338,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 @@ -328,9 +397,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: @@ -427,6 +499,46 @@ 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.", + ) + 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( "L1 Manager TTL", "TTL configuration for L1 manager locks" @@ -555,6 +667,20 @@ 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), + auto_expand=getattr(args, "maru_auto_expand", True), + ) + shm_name = getattr(args, "shm_name", None) if shm_name is None: memory_config = L1MemoryManagerConfig( @@ -562,6 +688,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: @@ -571,6 +698,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..9926314e197 --- /dev/null +++ b/lmcache/v1/distributed/l1_protocol.py @@ -0,0 +1,118 @@ +# 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], force: bool = False + ) -> dict[ObjectKey, L1Error]: + """Delete unlocked keys (locked keys refused). + ``force`` requests deletion of locked keys too; backends that cannot + honor it (e.g. a shared pool with cross-process pins) may still refuse + with ``KEY_IS_LOCKED``. + 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/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/maru_l1_manager.py b/lmcache/v1/distributed/maru_l1_manager.py new file mode 100644 index 00000000000..a2380bd651d --- /dev/null +++ b/lmcache/v1/distributed/maru_l1_manager.py @@ -0,0 +1,1200 @@ +# 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. + +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 +(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 +from dataclasses import dataclass +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Concatenate, + Literal, + ParamSpec, + TypeVar, +) +import functools +import threading +import time + +# 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.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 +from lmcache.v1.mp_observability.otel_init import register_gauge + +if TYPE_CHECKING: + # Third Party + 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") +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 + + +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. + + ``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 + # 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") + + +@dataclass +class _PendingWrite: + """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: + """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. + """ + + # 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: + 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) + # 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] = [] + # 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( + 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() + + # 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. + + Args: + listener: The listener to register. + """ + # PARITY(L1Manager.register_listener): inline lock, append only. + 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: + 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), + ) + + 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. + + 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: 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. + """ + if not keys: + return + 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 + 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 + + 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 + # 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] + 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. + # 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, pinned=total, deadline=deadline + ) + ret[k] = (L1Error.SUCCESS, mem_obj) + 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 + + 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) + # 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) + except Exception: + logger.exception( + "MaruL1Manager: sweep failed to reclaim read page %s", k + ) + 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 + ) -> 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 + } + # 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: + if k in self._pending_write: + ret[k] = (L1Error.KEY_NOT_READABLE, None) + else: + readable.append(k) + successful_keys: list[ObjectKey] = [] + 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: + listener.on_l1_keys_reserved_read(successful_keys) + self._publish(EventType.L1_READ_RESERVED, successful_keys) + 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] = [] + 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: + 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 + # 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] + 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) + self._publish(EventType.L1_READ_FINISHED, successful_keys) + self._publish(EventType.L1_KEYS_EVICTED, need_to_free_keys) + 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 + 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, deadline=deadline + ) + 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) + self._publish(EventType.L1_WRITE_RESERVED, successful_keys) + 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)) + 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). + 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 + def finish_write_and_reserve_read( + self, keys: list[ObjectKey], extra_count: int = 0 + ) -> dict[ObjectKey, L1OperationResult]: + """Finish a write and take read holds in one step (L2->L1 promote). + + 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. + """ + 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, + pinned=0, # local page: not directory-pinned + is_temporary=True, + deadline=time.monotonic() + self._read_ttl_seconds, + ) + 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) + self._publish(EventType.L1_WRITE_FINISHED_AND_READ_RESERVED, 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], force: bool = False + ) -> 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. + + DIVERGENCE(L1Manager.delete): ``force`` is accepted for interface + parity but ignored — locked keys are ALWAYS refused. On the shared + pool a lock may be a pin held by another process, and the handler has + no force-delete RPC; freeing a pinned page would corrupt in-flight + reads/writes across the pool. Callers see the refusal as + KEY_IS_LOCKED (reported upstream as "skipped"), never a silent + success. + + Args: + keys: The list of object keys to delete. + force: Ignored (see DIVERGENCE above). Present so this backend + satisfies ``L1ManagerInterface.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). Returned regardless of + ``force``. + """ + 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. + 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 + successful_keys.append(k) + 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 + # 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) + self._publish(EventType.L1_KEYS_EVICTED, successful_keys) + return ret + + def touch_keys(self, keys: list[ObjectKey]) -> None: + """Mark keys as accessed, feeding the eviction LRU recency. + + 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: + """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 + 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) + self._publish(EventType.L1_KEYS_EVICTED, dropped) + + def _drain_staging(self, force: bool) -> list[ObjectKey]: + """Unpin staged reads (``pinned`` times) and abort staged writes. + + Args: + force: If True, log the dropped in-flight staging as a warning. + + Returns: + The keys dropped from staging (staged reads then staged writes). + """ + to_unpin: list[str] = [] + for k, entry in self._pending_read.items(): + # 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) + 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) + 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). + + 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 for the eviction watermark. + + 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 + 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). + """ + if not self._allocator.is_initialized: + return 0, self._config.pool_size_bytes + try: + handler = self._allocator.handler + 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() + 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") + 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 + + def get_l1_memory_desc(self) -> L1MemoryDesc | None: + """Return None: the shared pool has no single registerable region.""" + return None + + def close(self) -> None: + """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. + 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], + 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. + 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/memory_manager/maru_memory_allocator.py b/lmcache/v1/distributed/memory_manager/maru_memory_allocator.py new file mode 100644 index 00000000000..375112aba71 --- /dev/null +++ b/lmcache/v1/distributed/memory_manager/maru_memory_allocator.py @@ -0,0 +1,312 @@ +# 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 -- 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. +""" + +# 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, + auto_expand=self._config.auto_expand, + ) + 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]: + """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 + ) + + 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]]: + """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 + ) + + def free(self, memory_obj: MemoryObj, allocator_type: Optional[str] = None) -> None: + """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, + memory_objs: List[MemoryObj], + allocator_type: Optional[str] = None, + update_stats: bool = True, + ) -> None: + """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) + + 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: + """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 -- 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() + self._cxl_adapter = None + if self._handler is not None: + self._handler.close() + 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 diff --git a/lmcache/v1/distributed/storage_controllers/eviction_controller.py b/lmcache/v1/distributed/storage_controllers/eviction_controller.py index 8624207de38..936f60d074e 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 2138a33aa01..d5b8b78ea6b 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 ( @@ -210,7 +210,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, diff --git a/lmcache/v1/distributed/storage_manager.py b/lmcache/v1/distributed/storage_manager.py index 3054e2441c5..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 @@ -25,6 +29,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 +38,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 ( @@ -66,7 +72,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 +94,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 +183,48 @@ def __init__(self, config: StorageManagerConfig): ) # External APIs for serving engine integration code to call + def register_kv_layout( + self, + layout_desc: MemoryLayoutDesc, + 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, 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. + 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. + + Raises: + ValueError: The maru L1 tier supports a single object group only + (multi-model support is a maru TODO). + """ + # 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: + 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, + engine_kv_format, + chunk_size_in_tokens, + ) + @enable_tracing() def reserve_write( self, @@ -777,7 +835,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/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/lmcache/v1/multiprocess/modules/lmcache_driven_transfer.py b/lmcache/v1/multiprocess/modules/lmcache_driven_transfer.py index 89a7d9f8e01..eaf3481bb91 100644 --- a/lmcache/v1/multiprocess/modules/lmcache_driven_transfer.py +++ b/lmcache/v1/multiprocess/modules/lmcache_driven_transfer.py @@ -890,6 +890,24 @@ 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. 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, + ) + 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, diff --git a/tests/v1/distributed/conftest.py b/tests/v1/distributed/conftest.py index d8bbf8ca386..bebb204f6bf 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: @@ -91,6 +94,29 @@ class 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) + + # --------------------------------------------------------------------------- # CLI options for live-server integration tests (e.g. the Valkey adapter). # diff --git a/tests/v1/distributed/maru_fakes.py b/tests/v1/distributed/maru_fakes.py new file mode 100644 index 00000000000..ed46a9c8fc1 --- /dev/null +++ b/tests/v1/distributed/maru_fakes.py @@ -0,0 +1,263 @@ +# 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.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).""" + + 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 + 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: + 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), + }, + "cxl_pool": {"free_size": self.cxl_free}, + } + + 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, + auto_expand: bool = True, +) -> 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", + auto_expand=auto_expand, + ), + ), + 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_distributed_storage_manager.py b/tests/v1/distributed/test_distributed_storage_manager.py index cffd0f85c90..e801b665b3a 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,6 +22,7 @@ EvictionConfig, L1ManagerConfig, L1MemoryManagerConfig, + MaruL1Config, StorageManagerConfig, ) from lmcache.v1.distributed.l2_adapters.config import ( @@ -32,6 +34,8 @@ 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 +123,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.""" @@ -907,6 +932,71 @@ def test_sparse_from_l2_loads_all_found( 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] + 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, engine_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, + MagicMock(name="engine_kv_format"), + 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, format never dereferenced) even for a + # would-be-rejected group count. + sm.register_kv_layout( + basic_layout, + MagicMock(name="engine_kv_format"), + 16, + num_object_groups=2, + ) + finally: + sm.close() + + class TestStorageManagerDelete: """Tests for StorageManager.delete_l1_keys().""" 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..f14d0f027c6 --- /dev/null +++ b/tests/v1/distributed/test_l1_manager_conformance.py @@ -0,0 +1,356 @@ +# 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. +""" + +# Standard +import inspect + +# 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 + 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 +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) + + +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_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. + + 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") + + +# ========================================================================= +# 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_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 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..9c7eca4e21f --- /dev/null +++ b/tests/v1/distributed/test_maru_l1_manager.py @@ -0,0 +1,868 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""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 +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, + parse_args, + ) + from lmcache.v1.distributed.error import L1Error + 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 + import lmcache.c_ops as lmc_ops + + # Local + from .maru_fakes import ( + FakeCxlAdapter, + FakeMaruHandler, + RecordingListener, + make_maru_manager, + ) +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) + + +# ========================================================================= +# allocator (MaruMemoryAllocator): URL normalization, init, delegation +# (absorbed from the former test_maru_memory_allocator.py, mocked runtime) +# ========================================================================= + + +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 + ) + + 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() + + +# ========================================================================= +# config / startup guards: flag parsing + rejected backend combinations +# (absorbed from the former test_l1_config_maru.py) +# ========================================================================= + + +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 _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, 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_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")) + + +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_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=""), + ) + + +# ========================================================================= +# 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() + 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_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_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: OOM, mode guard, and page-reclaim safety +# ========================================================================= + + +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_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 + + +# ========================================================================= +# delete: cross-node pinned-key refusal +# ========================================================================= + + +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_force_is_ignored_for_locked_keys(): + """DIVERGENCE(L1Manager.delete): force never bypasses locks on maru. + + A lock may be a pin held by another process on the shared pool, so + force=True must still refuse both remotely pinned and locally staged + keys with KEY_IS_LOCKED (surfaced upstream as "skipped"). + """ + manager, handler, _ = make_maru_manager() + + pinned = _key(1) + _seed(handler, pinned) + handler.pins[object_key_to_string(pinned)] = 1 # pinned by another instance + + staged = _key(2) + manager.reserve_write([staged], [True], _LAYOUT, mode="new") # in-flight write + + results = manager.delete([pinned, staged], force=True) + assert results[pinned] == L1Error.KEY_IS_LOCKED + assert results[staged] == L1Error.KEY_IS_LOCKED + assert staged in manager._pending_write # staging untouched, no page freed + + +# ========================================================================= +# get_memory_usage (device-fill watermark) / is_key_evictable +# ========================================================================= + + +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_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_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) + _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) + + +# ========================================================================= +# finish_write_and_reserve_read (L2->L1 promote): retained vs temporary +# ========================================================================= + + +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_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_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 + + +# ========================================================================= +# TTL sweeper: reclaim orphan write pages / read pins on expiry +# ========================================================================= + + +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 + + +# ========================================================================= +# notifications: maru-specific event_bus / listener nuances +# (the shared listener lifecycle is asserted in the conformance suite) +# ========================================================================= + + +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_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_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]] + + +# ========================================================================= +# 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 _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() + + +# ========================================================================= +# 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 caa415856e8..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).""" 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..1bf5b3d959e --- /dev/null +++ b/tests/v1/multiprocess/test_maru_register_kv_layout.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +"""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 +from unittest.mock import MagicMock +import threading + +# Third Party +import pytest + +# First Party +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_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 + # (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): + """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()) + 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_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() + + _register(module) # must not raise + + assert 1 in module._cache_contexts