fix(fs_native): give lookups and loads their own workers

.lil/changes/lmcache-107.json

@@ -0,0 +1,18 @@
+{
+  "schema": "local-inference-release-change/v1",
+  "id": "lmcache-107",
+  "category": "fix",
+  "summary": "Checkpoint restores from the disk tier no longer wait behind a burst of on-evict disk writes",
+  "models": [
+    "all"
+  ],
+  "compatibility": "No user action required. The fs_native disk tier now runs lookups (2 threads), loads (num_workers threads) and writes (max(2, num_workers / 4) threads) on separate workers. Set per_op_workers in the --l2-adapter JSON to choose other counts; per_op_workers {} restores the previous single pool.",
+  "details": [
+    "The disk tier served lookups, loads and writes from one queue. With on-evict writes, the first time RAM filled up it queued a write of every checkpoint page in RAM at once, and a restore that needed disk pages waited until all of them were written: 85-90 s for a 14 GiB burst on the test server, longer with a larger RAM tier. Such a restore now starts at once and no longer waits for the burst; in four runs of the same test it took 2, 18, 25 and 45 s on the test server's RAID5, whose reads still compete with the writes.",
+    "Fewer concurrent writes leave a busy disk room for restores; the same burst was written as fast as before."
+  ],
+  "pull_requests": [
+    107
+  ],
+  "requires": []
+}

csrc/storage_backends/fs/connector.cpp

@@ -297,8 +297,9 @@ static bool try_enable_odirect(int& flags, const void* buf, size_t len,
 
 FSConnector::FSConnector(std::string base_path, int num_workers,
                          std::string relative_tmp_dir, bool use_odirect,
-                         size_t read_ahead_size)
-    : ConnectorBase(num_workers),
+                         size_t read_ahead_size,
+                         WorkerPoolConfig worker_pool_config)
+    : ConnectorBase(num_workers, std::move(worker_pool_config)),
       base_path_(std::move(base_path)),
       relative_tmp_dir_(std::move(relative_tmp_dir)),
       use_odirect_(use_odirect),

csrc/storage_backends/fs/connector.h

@@ -42,9 +42,12 @@ struct WorkerFSConn {
 
 class FSConnector : public ConnectorBase<WorkerFSConn> {
  public:
+  // worker_pool_config gives operation lanes (lookup, retrieve, store,
+  // delete) dedicated workers; the other operations share num_workers.
   FSConnector(std::string base_path, int num_workers,
               std::string relative_tmp_dir = "", bool use_odirect = false,
-              size_t read_ahead_size = 0);
+              size_t read_ahead_size = 0,
+              WorkerPoolConfig worker_pool_config = WorkerPoolConfig{});
   ~FSConnector() override;
 
  protected:

csrc/storage_backends/fs/pybind.cpp

@@ -1,4 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
+#include <string>
+#include <utility>
+
 #include <pybind11/pybind11.h>
 #include "../connector_pybind_utils.h"
 #include "connector.h"
@@ -7,9 +10,18 @@ namespace py = pybind11;
 
 PYBIND11_MODULE(lmcache_fs, m) {
   py::class_<lmcache::connector::FSConnector>(m, "LMCacheFSClient")
-      .def(py::init<std::string, int, std::string, bool, size_t>(),
+      .def(py::init([](std::string base_path, int num_workers,
+                       std::string relative_tmp_dir, bool use_odirect,
+                       size_t read_ahead_size, py::object per_op_workers) {
+             return new lmcache::connector::FSConnector(
+                 std::move(base_path), num_workers, std::move(relative_tmp_dir),
+                 use_odirect, read_ahead_size,
+                 lmcache::connector::pybind_utils::parse_per_op_workers(
+                     per_op_workers));
+           }),
            py::arg("base_path"), py::arg("num_workers"),
            py::arg("relative_tmp_dir") = "", py::arg("use_odirect") = false,
-           py::arg("read_ahead_size") = 0)
+           py::arg("read_ahead_size") = 0,
+           py::arg("per_op_workers") = py::none())
           LMCACHE_BIND_CONNECTOR_METHODS(lmcache::connector::FSConnector);
 }

docs/source/mp/l2_storage/fs_native.rst

@@ -27,6 +27,15 @@ I/O queue depth on a single Python thread.
 - ``max_capacity_gb`` (float, default ``0``): Maximum L2 capacity in GB
   for client-side usage enforcement and global eviction. Default ``0``
   leaves capacity unlimited; resident-byte accounting remains available.
+- ``per_op_workers`` (dict, default ``{"lookup": 2, "retrieve": num_workers,
+  "store": max(2, num_workers // 4)}``): Dedicated worker threads per
+  operation lane (``lookup``, ``retrieve``, ``store``, ``delete``);
+  operations without a lane share the ``num_workers`` pool.  By default
+  lookups and loads have their own workers, so a restore does not wait in
+  the queue behind every write submitted before it: a burst of pages written
+  when RAM evicts them (for example with ``checkpoint_on_evict``) can occupy
+  a slow disk for minutes.  Fewer concurrent stores also keep a saturated
+  disk responsive to loads.  ``{}`` runs every operation in the shared pool.
 
 Restart recovery
 ----------------

lmcache/v1/distributed/l2_adapters/fs_native_l2_adapter.py

@@ -45,6 +45,11 @@
 
 _FILE_EXT = ".data"
 _IGNORED_FILE_SAMPLE_LIMIT = 5
+# Operation lanes of the native connector, the default lookup workers (a
+# lookup only stats files) and the least default store workers.
+_LANES = frozenset({"lookup", "retrieve", "store", "delete"})
+_DEFAULT_LOOKUP_WORKERS = 2
+_MIN_STORE_WORKERS = 2
 
 
 def _scan_existing_key_sizes(base_path: str) -> dict[ObjectKey, int]:
@@ -205,6 +210,29 @@ def _raise_walk_error(exc: OSError) -> None:
     return key_sizes
 
 
+def default_fs_native_lanes(num_workers: int) -> dict[str, int]:
+    """Return the default dedicated worker lanes of the native FS connector.
+
+    Lookups and loads get their own workers, so a restore never waits in the
+    queue behind every write submitted before it (a write-on-evict burst can
+    take minutes on a slow disk). Stores get a quarter of ``num_workers``,
+    at least 2: on a disk that stores saturate, fewer concurrent writes keep
+    loads fast without writing more slowly. Deletes share the ``num_workers``
+    pool.
+
+    Args:
+        num_workers: Shared worker count of the connector.
+
+    Returns:
+        Lane key to dedicated worker count.
+    """
+    return {
+        "lookup": _DEFAULT_LOOKUP_WORKERS,
+        "retrieve": num_workers,
+        "store": max(_MIN_STORE_WORKERS, num_workers // 4),
+    }
+
+
 class FSNativeL2AdapterConfig(L2AdapterConfigBase):
     """
     Config for an L2 adapter backed by the native C++
@@ -217,6 +245,11 @@ class FSNativeL2AdapterConfig(L2AdapterConfigBase):
     - use_odirect: bypass page cache via O_DIRECT.
     - read_ahead_size: trigger filesystem readahead by
       reading this many bytes first (optional).
+    - per_op_workers: dedicated worker counts per operation lane
+      (``lookup``, ``retrieve``, ``store``, ``delete``); the other
+      operations share ``num_workers``. None selects
+      :func:`default_fs_native_lanes`; ``{}`` runs every operation in the
+      shared pool.
     """
 
     def __init__(
@@ -227,13 +260,25 @@ def __init__(
         use_odirect: bool = False,
         read_ahead_size: Optional[int] = None,
         max_capacity_gb: float = 0,
+        per_op_workers: dict[str, int] | None = None,
     ):
         self.base_path = base_path
         self.num_workers = num_workers
         self.relative_tmp_dir = relative_tmp_dir
         self.use_odirect = use_odirect
         self.read_ahead_size = read_ahead_size
         self.max_capacity_gb = max_capacity_gb
+        per_op_workers = L2AdapterConfigBase._validate_per_op_workers(per_op_workers)
+        if per_op_workers is not None and not set(per_op_workers) <= _LANES:
+            raise ValueError(
+                f"per_op_workers keys must be among {sorted(_LANES)}, "
+                f"got {sorted(per_op_workers)}"
+            )
+        self.per_op_workers: dict[str, int] = (
+            default_fs_native_lanes(num_workers)
+            if per_op_workers is None
+            else dict(per_op_workers)
+        )
 
     @classmethod
     def from_dict(cls, d: dict) -> "FSNativeL2AdapterConfig":
@@ -269,6 +314,7 @@ def from_dict(cls, d: dict) -> "FSNativeL2AdapterConfig":
             use_odirect=use_odirect,
             read_ahead_size=read_ahead_size,
             max_capacity_gb=float(max_capacity_gb),
+            per_op_workers=L2AdapterConfigBase._parse_per_op_workers_from_dict(d),
         )
 
     @classmethod
@@ -288,7 +334,14 @@ def help(cls) -> str:
             "first (optional)\n"
             "- max_capacity_gb (float): max L2 capacity "
             "in GB for usage tracking / eviction "
-            "(default 0 = disabled)"
+            "(default 0 = disabled)\n"
+            "- per_op_workers (dict[str, int]): dedicated "
+            "worker counts per lane (lookup, retrieve, "
+            "store, delete); other ops share num_workers "
+            "(default: lookup 2, retrieve num_workers, "
+            "store max(2, num_workers // 4), so restores "
+            "never queue behind writes; {} = one shared "
+            "pool)"
         )
 
 
@@ -344,6 +397,7 @@ def _create_fs_native_l2_adapter(
         config.relative_tmp_dir,
         config.use_odirect,
         config.read_ahead_size or 0,
+        per_op_workers=config.per_op_workers,
     )
     try:
         adapter = NativeConnectorL2Adapter(
@@ -354,6 +408,7 @@ def _create_fs_native_l2_adapter(
                 "base_path": config.base_path,
                 "use_odirect": config.use_odirect,
                 "num_workers": config.num_workers,
+                "per_op_workers": dict(config.per_op_workers),
                 "read_ahead_size": config.read_ahead_size,
             },
             initial_key_sizes=initial_key_sizes,
@@ -363,10 +418,11 @@ def _create_fs_native_l2_adapter(
         native_client.close()
         raise
     logger.info(
-        "Created FS native L2 adapter: %s (workers=%d, odirect=%s, "
+        "Created FS native L2 adapter: %s (workers=%d, lanes=%s, odirect=%s, "
         "read_ahead=%s, restored_keys=%d, restored_bytes=%d)",
         config.base_path,
         config.num_workers,
+        config.per_op_workers,
         config.use_odirect,
         config.read_ahead_size,
         len(initial_key_sizes),

tests/v1/distributed/test_l2_adapter_factory.py

@@ -1116,6 +1116,87 @@ def _capture_factory2(config, l1_memory_desc=None):
             _L2_ADAPTER_FACTORY_REGISTRY["fs_native"] = old
 
 
+class TestFSNativeLanes:
+    """Lookups and loads of the native FS connector get their own workers."""
+
+    def test_default_gives_reads_their_own_lanes(self):
+        # First Party
+        from lmcache.v1.distributed.l2_adapters.fs_native_l2_adapter import (
+            FSNativeL2AdapterConfig,
+        )
+
+        cfg = FSNativeL2AdapterConfig.from_dict(
+            {"base_path": "/tmp/lanes", "num_workers": 8}
+        )
+        assert cfg.per_op_workers == {"lookup": 2, "retrieve": 8, "store": 2}
+        wide = FSNativeL2AdapterConfig.from_dict(
+            {"base_path": "/tmp/lanes", "num_workers": 32}
+        )
+        assert wide.per_op_workers == {"lookup": 2, "retrieve": 32, "store": 8}
+
+    def test_explicit_lanes_replace_the_default(self):
+        # First Party
+        from lmcache.v1.distributed.l2_adapters.fs_native_l2_adapter import (
+            FSNativeL2AdapterConfig,
+        )
+
+        cfg = FSNativeL2AdapterConfig.from_dict(
+            {"base_path": "/tmp/lanes", "per_op_workers": {"store": 3}}
+        )
+        assert cfg.per_op_workers == {"store": 3}
+        shared = FSNativeL2AdapterConfig.from_dict(
+            {"base_path": "/tmp/lanes", "per_op_workers": {}}
+        )
+        assert shared.per_op_workers == {}
+
+    @pytest.mark.parametrize(
+        "per_op_workers",
+        [{"bogus": 1}, {"lookup": 0}, {"retrieve": True}, ["lookup"]],
+    )
+    def test_invalid_lanes_are_rejected(self, per_op_workers):
+        # First Party
+        from lmcache.v1.distributed.l2_adapters.fs_native_l2_adapter import (
+            FSNativeL2AdapterConfig,
+        )
+
+        with pytest.raises(ValueError, match="per_op_workers"):
+            FSNativeL2AdapterConfig.from_dict(
+                {"base_path": "/tmp/lanes", "per_op_workers": per_op_workers}
+            )
+
+    def test_factory_forwards_lanes_to_the_native_client(self, monkeypatch, tmp_path):
+        # Standard
+        import sys
+        import types
+
+        # First Party
+        from lmcache.v1.distributed.l2_adapters.fs_native_l2_adapter import (
+            FSNativeL2AdapterConfig,
+            _create_fs_native_l2_adapter,
+        )
+
+        captured: dict = {}
+
+        class _LaneFSClient(_FakeLMCacheFSClient):
+            def __init__(self, *args, per_op_workers=None):
+                super().__init__(*args)
+                captured["per_op_workers"] = per_op_workers
+
+        module = types.ModuleType("lmcache.lmcache_fs")
+        module.LMCacheFSClient = _LaneFSClient
+        monkeypatch.setitem(sys.modules, "lmcache.lmcache_fs", module)
+        cfg = FSNativeL2AdapterConfig(base_path=str(tmp_path), num_workers=4)
+        adapter = _create_fs_native_l2_adapter(cfg)
+        try:
+            assert captured["per_op_workers"] == {
+                "lookup": 2,
+                "retrieve": 4,
+                "store": 2,
+            }
+        finally:
+            adapter.close()
+
+
 def _patched_fs_factory(config, l1_memory_desc=None):
     """Factory that uses _FakeLMCacheFSClient instead
     of the real C++ extension."""

tests/v1/storage_backend/test_fs_native_connector.py

@@ -567,3 +567,45 @@ def test_odirect_falls_back_for_misaligned_buffer_address(tmp_path) -> None:
         assert bytes(dest) == bytes(source)
     finally:
         client.close()
+
+
+@pytest.mark.parametrize(
+    "per_op_workers", [None, {"lookup": 1, "retrieve": 1}], ids=["shared", "lanes"]
+)
+def test_reads_in_their_own_lanes_do_not_wait_behind_queued_writes(
+    tmp_path, per_op_workers: dict[str, int] | None
+) -> None:
+    """A lookup and a load submitted after a large write finish before it
+    with read lanes; one shared worker serves them only after the write."""
+    LMCacheFSClient = _import_fs_client()
+    existing = "model@00000000@0@aa"
+    payload = bytearray(b"restored-page")
+    client = LMCacheFSClient(str(tmp_path), 1, per_op_workers=per_op_workers)
+    try:
+        stored = _submit_and_wait(
+            client, "submit_batch_set", existing, memoryview(payload)
+        )
+        assert stored[1], stored[2]
+        page = bytearray(16 << 20)
+        backlog = [f"model@00000000@0@{i + 256:x}" for i in range(32)]
+        write = client.submit_batch_set(backlog, [memoryview(page)] * len(backlog))
+        lookup = client.submit_batch_exists([existing])
+        destination = bytearray(len(payload))
+        load = client.submit_batch_get([existing], [memoryview(destination)])
+        order: list[int] = []
+        deadline = time.monotonic() + 120
+        while len(order) < 3 and time.monotonic() < deadline:
+            order += [
+                completion[0]
+                for completion in client.drain_completions()
+                if completion[0] in (write, lookup, load)
+            ]
+            select.select([client.event_fd()], [], [], 0.1)
+        assert sorted(order) == sorted([write, lookup, load])
+        assert destination == payload
+        if per_op_workers:
+            assert order[-1] == write
+        else:
+            assert order == [write, lookup, load]
+    finally:
+        client.close()