use FlexKV main + add test result

Author: Clebrate

corelib/recsys_kvcache_manager/recsys_kvcache_manager/default_kvcache_backend.py

@@ -313,6 +313,22 @@ def _build_host_kvstorage_manager_from_config(
                 }
             else:
                 flexkv_as_batch = bool(flexkv_as_batch_raw)
+            flexkv_enable_layerwise = extra.get("flexkv_enable_layerwise", None)
+            if isinstance(flexkv_enable_layerwise, str):
+                flexkv_enable_layerwise = flexkv_enable_layerwise.strip().lower() in {
+                    "1",
+                    "true",
+                    "yes",
+                    "on",
+                }
+            elif flexkv_enable_layerwise is not None:
+                flexkv_enable_layerwise = bool(flexkv_enable_layerwise)
+            flexkv_layerwise_eventfd_socket = extra.get(
+                "flexkv_layerwise_eventfd_socket", None
+            )
+            flexkv_layerwise_counter_id = int(
+                extra.get("flexkv_layerwise_counter_id", 0)
+            )
 
             return FlexKVStorage(
                 mode=flexkv_mode,
@@ -331,6 +347,9 @@ def _build_host_kvstorage_manager_from_config(
                 host_kvstorage_fail_policy=flexkv_host_kvstorage_fail_policy,
                 hostkv_wait_timeout_ms=int(kvcache_config.offload_timeout_ms),
                 config_path=flexkv_config_path,
+                enable_layerwise=flexkv_enable_layerwise,
+                layerwise_eventfd_socket=flexkv_layerwise_eventfd_socket,
+                layerwise_counter_id=flexkv_layerwise_counter_id,
             )
         else:
             raise NotImplementedError(

corelib/recsys_kvcache_manager/recsys_kvcache_manager/flex_kvcache_manager.py

@@ -21,6 +21,7 @@
 
 import numpy as np
 import torch
+from flexkv.common.config import LayerGroupSpec
 from flexkv.common.storage import KVCacheLayout, KVCacheLayoutType
 from flexkv.server.client import KVTPClient
 
@@ -31,6 +32,10 @@
 
     KVResponse = Any  # type: ignore
 
+from .flexkv_layerwise import (
+    FlexKVLayerwiseEventfdSender,
+    create_layerwise_eventfd_socket_path,
+)
 from .host_kvstorage_manager import (
     HostKVStorageBase,
     HostKVTaskHandle,
@@ -39,10 +44,6 @@
 )
 from .kvcache_metadata import KVCacheMetadata
 from .kvcache_utils import KVIndexMeta, KVLookupResult
-from .flexkv_layerwise import (
-    DEFAULT_LAYERWISE_EVENTFD_SOCKET,
-    FlexKVLayerwiseEventfdSender,
-)
 
 
 @dataclass
@@ -148,10 +149,8 @@ def __init__(
         self.enable_layerwise = enable_layerwise
         self.layerwise_eventfd_socket = (
             layerwise_eventfd_socket
-            or os.environ.get(
-                "FLEXKV_LAYERWISE_EVENTFD_SOCKET",
-                DEFAULT_LAYERWISE_EVENTFD_SOCKET,
-            )
+            or os.environ.get("FLEXKV_LAYERWISE_EVENTFD_SOCKET")
+            or create_layerwise_eventfd_socket_path()
         )
         self.layerwise_counter_id = int(layerwise_counter_id)
         self.backend_name = "flexkv"
@@ -190,16 +189,38 @@ def register_gpu_cache_tables(self, cache_table_list: List[torch.Tensor]) -> Non
             tokens_per_block=int(first_table.shape[2]),
             num_head=int(first_table.shape[3]),
             head_size=int(first_table.shape[4]),
-            is_mla=False,
         )
         tp_client = KVTPClient(
             gpu_register_port=self._gpu_register_port,
             dp_client_id=0,
             device_id=device_id,
         )
+        register_kwargs: Dict[str, Any] = {}
+        if self.enable_layerwise:
+            # FlexKV main's multi-group path is also its public representation
+            # for a uniform cache with one member per layer. Registering HSTU as
+            # one group keeps the layout unchanged while avoiding assumptions in
+            # the legacy single-group layerwise constructor.
+            register_kwargs.update(
+                layer_groups=[
+                    LayerGroupSpec(
+                        num_layers=gpu_layout.num_layer,
+                        num_kv_heads=gpu_layout.num_head,
+                        head_size=gpu_layout.head_size,
+                        layer_indices=list(range(gpu_layout.num_layer)),
+                        dtype=self.dtype,
+                    )
+                ],
+                gpu_layouts=[gpu_layout],
+                handles_per_group=[self._gpu_cache_table_list],
+            )
         tp_client.register_to_server(
-            kv_caches=self._gpu_cache_table_list, kv_layout=gpu_layout
+            kv_caches=self._gpu_cache_table_list,
+            kv_layout=gpu_layout,
+            **register_kwargs,
         )
+        if self.enable_layerwise and self._layerwise_eventfd_sender is not None:
+            self._layerwise_eventfd_sender.wait_until_ready()
         self._registered = True
 
         # Client becomes operational only after transfer manager is ready.
@@ -559,9 +580,6 @@ def offload_kvcache_launch(
             else self._build_slot_mappings(kvcache_metadata)
         )
 
-        # Keep offload on the direct put_async path while validating layerwise.
-        # This matches the known-good 2e259b0 behavior and isolates PR418's
-        # put_match/as_batch offload optimization from PR428 layerwise onboard.
         use_batch = False
         task_ids: List[int] = []
         batch_slot_mappings: List[torch.Tensor] = []

corelib/recsys_kvcache_manager/recsys_kvcache_manager/flexkv_layerwise.py

@@ -15,7 +15,9 @@
 
 import os
 import socket
+import stat
 import struct
+import tempfile
 import threading
 import time
 from array import array
@@ -26,8 +28,16 @@
 DEFAULT_LAYERWISE_EVENTFD_SOCKET = "/tmp/flexkv_layerwise_eventfd.sock"
 
 
+def create_layerwise_eventfd_socket_path() -> str:
+    socket_directory = tempfile.mkdtemp(
+        prefix=f"recsys-flexkv-layerwise-{os.geteuid()}-"
+    )
+    os.chmod(socket_directory, 0o700)
+    return os.path.join(socket_directory, "eventfd.sock")
+
+
 class FlexKVLayerwiseEventfdSender:
-    """Creates mock layerwise eventfds and sends them to FlexKV's worker."""
+    """Creates layerwise eventfds and sends them to FlexKV's worker."""
 
     def __init__(
         self,
@@ -42,6 +52,8 @@ def __init__(
         self.timeout_s = float(timeout_s)
         self._eventfds: Optional[List[List[int]]] = None
         self._thread: Optional[threading.Thread] = None
+        self._handoff_done = threading.Event()
+        self._handoff_error: Optional[BaseException] = None
 
     @staticmethod
     def _create_eventfd() -> int:
@@ -76,20 +88,87 @@ def start(self) -> None:
             return
         self.create_eventfds()
         self._thread = threading.Thread(
-            target=self._send_eventfds,
+            target=self._run_sender,
             name="flexkv-layerwise-eventfd-sender",
             daemon=True,
         )
         self._thread.start()
 
+    def wait_until_ready(self, timeout_s: Optional[float] = None) -> None:
+        timeout = self.timeout_s + 1.0 if timeout_s is None else float(timeout_s)
+        if not self._handoff_done.wait(timeout):
+            raise TimeoutError(
+                "Timed out waiting for FlexKV layerwise eventfd handoff "
+                f"on socket {self.socket_path}"
+            )
+        if self._handoff_error is not None:
+            raise RuntimeError(
+                "FlexKV layerwise eventfd handoff failed"
+            ) from self._handoff_error
+
+    def _run_sender(self) -> None:
+        try:
+            self._send_eventfds()
+        except BaseException as error:
+            self._handoff_error = error
+        finally:
+            self._handoff_done.set()
+
+    @staticmethod
+    def _is_process_descendant(process_id: int, ancestor_id: int) -> bool:
+        current_id = int(process_id)
+        ancestor_id = int(ancestor_id)
+        visited = set()
+        while current_id > 1 and current_id not in visited:
+            if current_id == ancestor_id:
+                return True
+            visited.add(current_id)
+            try:
+                with open(f"/proc/{current_id}/stat", encoding="utf-8") as stat_file:
+                    process_stat = stat_file.read()
+                fields_after_name = process_stat[process_stat.rfind(")") + 1 :].split()
+                current_id = int(fields_after_name[1])
+            except (OSError, IndexError, ValueError):
+                return False
+        return current_id == ancestor_id
+
+    def _authenticate_peer(self, sock: socket.socket) -> None:
+        socket_stat = os.stat(self.socket_path, follow_symlinks=False)
+        if not stat.S_ISSOCK(socket_stat.st_mode):
+            raise PermissionError(
+                f"Layerwise eventfd path is not a socket: {self.socket_path}"
+            )
+        if socket_stat.st_uid != os.geteuid():
+            raise PermissionError(
+                "FlexKV layerwise socket owner does not match the current user"
+            )
+
+        if not hasattr(socket, "SO_PEERCRED"):
+            raise RuntimeError("SO_PEERCRED is required for layerwise eventfd handoff")
+        credentials_size = struct.calcsize("3i")
+        peer_pid, peer_uid, peer_gid = struct.unpack(
+            "3i",
+            sock.getsockopt(socket.SOL_SOCKET, socket.SO_PEERCRED, credentials_size),
+        )
+        if peer_uid != os.geteuid() or peer_gid != os.getegid():
+            raise PermissionError(
+                "FlexKV layerwise peer credentials do not match the current process"
+            )
+        if not self._is_process_descendant(peer_pid, os.getpid()):
+            raise PermissionError(
+                f"FlexKV layerwise peer pid {peer_pid} is outside the current process tree"
+            )
+
     def _send_eventfds(self) -> None:
         eventfds = self.create_eventfds()
         deadline = time.time() + self.timeout_s
         last_error: Optional[Exception] = None
         while time.time() < deadline:
             sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
             try:
+                sock.settimeout(min(1.0, max(deadline - time.time(), 0.01)))
                 sock.connect(self.socket_path)
+                self._authenticate_peer(sock)
                 metadata = struct.pack(
                     "iiii",
                     0,
@@ -116,12 +195,12 @@ def _send_eventfds(self) -> None:
                         f"FlexKV layerwise eventfd receiver returned ack={ack!r}"
                     )
                 return
-            except (FileNotFoundError, ConnectionRefusedError, socket.timeout) as e:
+            except (OSError, RuntimeError) as e:
                 last_error = e
                 time.sleep(0.05)
             finally:
                 sock.close()
         raise RuntimeError(
-            "Timed out sending mock layerwise eventfds to FlexKV "
+            "Timed out sending layerwise eventfds to FlexKV "
             f"socket {self.socket_path}: {last_error}"
         )

corelib/recsys_kvcache_manager/recsys_kvcache_manager/host_kvstorage_manager.py

@@ -119,11 +119,14 @@ def onboard_kvcache_launch(
     def onboard_kvcache_wait(self, task_handle: HostKVTaskHandle) -> HostKVWaitResult:
         ...
 
-    @abstractmethod
     def onboard_kvcache_wait_by_layer(
         self, task_handle: HostKVTaskHandle, layer_idx: int
     ) -> HostKVWaitResult:
-        ...
+        return HostKVWaitResult(
+            status=HostKVTaskStatus.SKIPPED,
+            ready=False,
+            message="Layerwise onboard wait is not supported by this backend.",
+        )
 
     @abstractmethod
     def offload_kvcache_launch(

corelib/recsys_kvcache_manager/test/test_flexkv.py

@@ -22,17 +22,21 @@
 from recsys_kvcache_manager.kvcache_manager import KVCacheManager
 
 
-def _get_flexkv_enable_layerwise() -> str:
-    return os.environ.get("RECSYS_FLEXKV_ENABLE_LAYERWISE", "")
+def _get_flexkv_config_path() -> str:
+    return os.environ.get("RECSYS_FLEXKV_CONFIG_PATH", "")
 
 
 def create_testing_kvcache_manager() -> KVCacheManager:
-    flexkv_enable_layerwise = _get_flexkv_enable_layerwise()
+    flexkv_config_path = _get_flexkv_config_path()
+    flexkv_enable_layerwise = os.environ.get("RECSYS_FLEXKV_ENABLE_LAYERWISE", "")
     extra_configs = {
         "flexkv_mode": "direct",
         "flexkv_host_kvstorage_fail_policy": "fail_open",
         "flexkv_enable_mps": 0,
+        "flexkv_as_batch": 1,
     }
+    if flexkv_config_path:
+        extra_configs["flexkv_config_path"] = flexkv_config_path
     if flexkv_enable_layerwise:
         extra_configs["flexkv_enable_layerwise"] = flexkv_enable_layerwise
 
@@ -68,15 +72,27 @@ def create_testing_kvcache_manager() -> KVCacheManager:
     )
     print(f"[TEST] KVCache GPU Memory Usage: {gpu_gib} GiB.")
     print(f"[TEST] KVCache Host Memory Usage: {host_gib} GiB.")
+    if flexkv_config_path:
+        print(f"[TEST] FlexKV config path: {flexkv_config_path}")
     kvcache_mgr = KVCacheManager.from_config(kvcache_config)
     flexkv_mgr = kvcache_mgr.host_kvstorage_manager
+    cache_cfg = flexkv_mgr._client.cache_config
+    if cache_cfg.enable_ssd:
+        print(
+            "[TEST] Created KVCache Manager with FlexKV SSD tier: "
+            f"num_cpu_blocks={cache_cfg.num_cpu_blocks}, "
+            f"num_ssd_blocks={cache_cfg.num_ssd_blocks}, "
+            f"ssd_cache_dir={cache_cfg.ssd_cache_dir}, "
+            f"enable_gds={cache_cfg.enable_gds}"
+        )
+    else:
+        print("[TEST] Created KVCache Manager with FlexKV CPU tier only")
     if flexkv_mgr.enable_layerwise:
         print(
             "[TEST] FlexKV layerwise transfer enabled: "
             f"eventfd_socket={flexkv_mgr.layerwise_eventfd_socket}, "
             f"counter_id={flexkv_mgr.layerwise_counter_id}"
         )
-    print("[TEST] Created KVCache Manager with FlexKV CPU tier only")
     return kvcache_mgr