[Feature] Add RDMA auto-detection for kvcache (#1194)

- add RDMA auto-detection
- enable RDMA auto-detection in InfiniStore

Signed-off-by: Haiyang Shi <[email protected]>
Co-authored-by: Haiyang Shi <[email protected]>

Author: DwyaneShi

python/aibrix_kvcache/aibrix_kvcache/envs.py

@@ -101,7 +101,7 @@
     # Since 0.2.42, InfiniStore supports RDMA GID index in client config,
     # users can specify the GID index of each device in this format:
     # "mlx5_0:gid0,mlx5_1:gid1,mlx5_2:gid2"
-    AIBRIX_KV_CACHE_OL_INFINISTORE_VISIBLE_DEV_LIST: List[str] = ["mlx5_0"]
+    AIBRIX_KV_CACHE_OL_INFINISTORE_VISIBLE_DEV_LIST: List[str] = []
     AIBRIX_KV_CACHE_OL_INFINISTORE_USE_GDR: bool = True
 
     # HPKV Env Vars
@@ -111,6 +111,11 @@
     AIBRIX_KV_CACHE_OL_HPKV_LOCAL_PORT: int = 12345
     AIBRIX_KV_CACHE_OL_HPKV_USE_GDR: bool = True
 
+    # RDMA Auto-Detection Env Vars
+    # Defines the range of valid GIDs. Similar to NVSHMEM_IB_ADDR_RANGE
+    # for NVSHMEM. It must be a valid CIDR.
+    AIBRIX_KV_CACHE_OL_TRANSPORT_RDMA_ADDR_RANGE: str = "::/0"
+
 # The begin-* and end* here are used by the documentation generator
 # to extract the used env vars.
 
@@ -289,10 +294,13 @@
             "AIBRIX_KV_CACHE_OL_INFINISTORE_LINK_TYPE", "Ethernet"
         ).strip()
     ),
-    "AIBRIX_KV_CACHE_OL_INFINISTORE_VISIBLE_DEV_LIST": lambda: (
-        os.getenv("AIBRIX_KV_CACHE_OL_INFINISTORE_VISIBLE_DEV_LIST", "mlx5_0")
-        .strip()
-        .split(",")
+    "AIBRIX_KV_CACHE_OL_INFINISTORE_VISIBLE_DEV_LIST": lambda: list(
+        filter(
+            None,
+            os.getenv("AIBRIX_KV_CACHE_OL_INFINISTORE_VISIBLE_DEV_LIST", "")
+            .strip()
+            .split(","),
+        )
     ),
     "AIBRIX_KV_CACHE_OL_INFINISTORE_USE_GDR": lambda: (
         os.getenv("AIBRIX_KV_CACHE_OL_INFINISTORE_USE_GDR", "1").strip().lower()
@@ -315,6 +323,12 @@
         os.getenv("AIBRIX_KV_CACHE_OL_HPKV_USE_GDR", "1").strip().lower()
         in ("1", "true")
     ),
+    # ================== RDMA Auto-Detection Env Vars ==================
+    "AIBRIX_KV_CACHE_OL_TRANSPORT_RDMA_ADDR_RANGE": lambda: (
+        os.getenv(
+            "AIBRIX_KV_CACHE_OL_TRANSPORT_RDMA_ADDR_RANGE", "::/0"
+        ).strip()
+    ),
 }
 
 # end-env-vars-definition

python/aibrix_kvcache/aibrix_kvcache/l2/connectors/infinistore.py

@@ -18,13 +18,18 @@
 
 import infinistore
 import torch
+from validators import ipv4, ipv6
 
 from ... import envs
 from ...common import AsyncBase
+from ...common.absl_logging import getLogger
 from ...memory import MemoryRegion
 from ...status import Status, StatusCodes
+from ...transport import AddrFamily, DeviceRequest, GIDType, RDMATransport
 from . import Connector, ConnectorFeature
 
+logger = getLogger(__name__)
+
 
 @AsyncBase.async_wrap(delete="_delete")
 class InfiniStoreConnector(Connector[bytes, torch.Tensor], AsyncBase):
@@ -52,18 +57,73 @@ def from_envs(
         service_port = kwargs.get(
             "port", envs.AIBRIX_KV_CACHE_OL_INFINISTORE_SERVICE_PORT
         )
+        assert ipv4(host_addr) or ipv6(host_addr), "Invalid host_addr"
         dev_list = envs.AIBRIX_KV_CACHE_OL_INFINISTORE_VISIBLE_DEV_LIST
-        assert (
-            len(dev_list) > 0
-        ), "AIBRIX_KV_CACHE_OL_INFINISTORE_VISIBLE_DEV_LIST is empty"
+        connection_type = envs.AIBRIX_KV_CACHE_OL_INFINISTORE_CONNECTION_TYPE
+        ib_port = envs.AIBRIX_KV_CACHE_OL_INFINISTORE_IB_PORT
+        link_type = envs.AIBRIX_KV_CACHE_OL_INFINISTORE_LINK_TYPE
+
+        if connection_type == "TCP":
+            config = infinistore.ClientConfig(
+                host_addr=host_addr,
+                service_port=service_port,
+                connection_type=connection_type,
+                link_type=link_type,
+            )
+            return cls(config, conn_id, executor)
+
+        # RDMA
+        addr_family = (
+            AddrFamily.AF_INET if ipv4(host_addr) else AddrFamily.AF_INET6
+        )
+        gid_type = GIDType.ROCE_V2 if link_type != "IB" else GIDType.IB_ROCE_V1
+        if len(dev_list) == 0:
+            logger.info(
+                "AIBRIX_KV_CACHE_OL_INFINISTORE_VISIBLE_DEV_LIST is not set, "
+                "trying to auto-detect visible devices"
+            )
+            addr_range = envs.AIBRIX_KV_CACHE_OL_TRANSPORT_RDMA_ADDR_RANGE
+            rdma = RDMATransport(
+                addr_range=addr_range,
+                addr_family=addr_family,
+                gid_type=gid_type,
+            )
+            status = rdma.get_device_list()
+            assert status.is_ok(), f"Failed to get device list: {status}"
+
+            devices = status.get()
+            for d in devices:
+                dev_list.append(f"{d.device_name}:{d.port_attrs.gid_index}")
+        else:
+            requests: List[DeviceRequest] = []
+            for dev_name in dev_list:
+                if ":" in dev_name:
+                    splits = dev_name.split(":")
+                    request = DeviceRequest(
+                        device_name=splits[0],
+                        gid_index=int(splits[1]),
+                    )
+                else:
+                    request = DeviceRequest(device_name=dev_name)
+                requests.append(request)
+            rdma = RDMATransport(
+                request=requests,
+                addr_family=addr_family,
+                gid_type=gid_type,
+            )
+            status = rdma.get_device_list()
+            # Only update dev_list if we got a new list
+            if status.is_ok():
+                devices = status.get()
+                dev_list.clear()
+                for d in devices:
+                    dev_list.append(f"{d.device_name}:{d.port_attrs.gid_index}")
 
         num_visible_gpus = torch.cuda.device_count()
 
-        dev_list = dev_list[:num_visible_gpus]
-        assert num_visible_gpus % len(dev_list) == 0, (
-            "AIBRIX_KV_CACHE_OL_INFINISTORE_VISIBLE_DEV_LIST is not a "
-            "multiple of num. of visible GPUs"
-        )
+        dev_list = [
+            dev_list[i % len(dev_list)] for i in range(num_visible_gpus)
+        ]
 
         # For InfiniStore RDMA, we need to map the GPU index to the RNIC
         # index to support multi-GPU per RNIC. For example, if we have 8
@@ -75,6 +135,8 @@ def from_envs(
         dev_name = dev_list[rnic_idx]
         hint_gid_index_str = ""
 
+        logger.info(f"InfiniStore selects {dev_name}")
+
         # If dev_name is in the format of "mlx5_i:xxx", then we need to
         # extract the dev_name and hint_gid_index from the dev_name.
         if ":" in dev_name:
@@ -85,9 +147,9 @@ def from_envs(
         config = infinistore.ClientConfig(
             host_addr=host_addr,
             service_port=service_port,
-            connection_type=envs.AIBRIX_KV_CACHE_OL_INFINISTORE_CONNECTION_TYPE,
-            ib_port=envs.AIBRIX_KV_CACHE_OL_INFINISTORE_IB_PORT,
-            link_type=envs.AIBRIX_KV_CACHE_OL_INFINISTORE_LINK_TYPE,
+            connection_type=connection_type,
+            ib_port=ib_port,
+            link_type=link_type,
             dev_name=dev_name,
         )
 

python/aibrix_kvcache/aibrix_kvcache/transport/__init__.py

@@ -0,0 +1,17 @@
+# Copyright 2024 The Aibrix Team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# 	http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .rdma import AddrFamily, DeviceRequest, GIDType, RDMATransport
+
+__all__ = ["AddrFamily", "DeviceRequest", "GIDType", "RDMATransport"]