UCT/CUDA-IPC: Fix MNNVL reachability check

config/m4/cuda.m4

@@ -66,6 +66,13 @@ AS_IF([test "x$cuda_checked" != "xyes"],
                                     [AC_MSG_ERROR([libnvidia-ml not found. Install appropriate nvidia-driver package])])
                               cuda_happy="no"])])
 
+         # Check for nvmlDeviceGetGpuFabricInfoV
+         AC_CHECK_DECLS([nvmlDeviceGetGpuFabricInfoV],
+                        [AC_DEFINE([HAVE_NVML_FABRIC_INFO], 1, [Enable NVML GPU fabric info support])],
+                        [AC_MSG_NOTICE([nvmlDeviceGetGpuFabricInfoV function not found in libnvidia-ml. MNNVL support will be disabled.])],
+                        [[#include <nvml.h>]])
+
+
          # Check for cuda static library
          have_cuda_static="no"
          AS_IF([test "x$cuda_happy" = "xyes"],

src/ucp/wireup/wireup.c

@@ -1249,10 +1249,12 @@ int ucp_wireup_is_reachable(ucp_ep_h ep, unsigned ep_init_flags,
     ucp_context_h context      = ep->worker->context;
     ucp_worker_iface_t *wiface = ucp_worker_iface(ep->worker, rsc_index);
     uct_iface_is_reachable_params_t params = {
-        .field_mask  = UCT_IFACE_IS_REACHABLE_FIELD_DEVICE_ADDR |
-                       UCT_IFACE_IS_REACHABLE_FIELD_IFACE_ADDR,
-        .device_addr = ae->dev_addr,
-        .iface_addr  = ae->iface_addr,
+        .field_mask         = UCT_IFACE_IS_REACHABLE_FIELD_DEVICE_ADDR |
+                              UCT_IFACE_IS_REACHABLE_FIELD_IFACE_ADDR |
+                              UCT_IFACE_IS_REACHABLE_FIELD_DEVICE_ADDR_LENGTH,
+        .device_addr        = ae->dev_addr,
+        .iface_addr         = ae->iface_addr,
+        .device_addr_length = ae->dev_addr_len
     };
 
     if (info_str != NULL) {

src/uct/api/v2/uct_v2.h

@@ -273,7 +273,8 @@ typedef enum {
     UCT_IFACE_IS_REACHABLE_FIELD_IFACE_ADDR         = UCS_BIT(1), /**< iface_addr field */
     UCT_IFACE_IS_REACHABLE_FIELD_INFO_STRING        = UCS_BIT(2), /**< info_string field */
     UCT_IFACE_IS_REACHABLE_FIELD_INFO_STRING_LENGTH = UCS_BIT(3), /**< info_string_length field */
-    UCT_IFACE_IS_REACHABLE_FIELD_SCOPE              = UCS_BIT(4) /**<  scope field */
+    UCT_IFACE_IS_REACHABLE_FIELD_SCOPE              = UCS_BIT(4), /**< scope field */
+    UCT_IFACE_IS_REACHABLE_FIELD_DEVICE_ADDR_LENGTH = UCS_BIT(5), /**< device_addr_length field */
 } uct_iface_is_reachable_field_mask_t;
 
 
@@ -611,6 +612,13 @@ typedef struct uct_iface_is_reachable_params {
      * Reachability scope.
      */
     uct_iface_reachability_scope_t scope;
+
+    /**
+     * Device address length. If not provided, the transport will assume a
+     * default minimal length according to the address buffer contents.
+     */
+    size_t                        device_addr_length;
+
 } uct_iface_is_reachable_params_t;
 
 

src/uct/cuda/base/cuda_md.c

@@ -101,6 +101,37 @@ uct_cuda_base_query_md_resources(uct_component_t *component,
                                            num_resources_p);
 }
 
+int uct_cuda_base_is_coherent()
+{
+#if HAVE_CUDA_FABRIC
+    static int coherent = -1;
+    CUdevice cu_device;
+    ucs_status_t status;
+
+    if (coherent != -1) {
+        goto out;
+    }
+
+    status = UCT_CUDADRV_FUNC_LOG_ERR(cuDeviceGet(&cu_device, 0));
+    if (status != UCS_OK) {
+        coherent = 0;
+        goto out;
+    }
+
+    status = UCT_CUDADRV_FUNC_LOG_ERR(
+            cuDeviceGetAttribute(&coherent,
+               CU_DEVICE_ATTRIBUTE_PAGEABLE_MEMORY_ACCESS_USES_HOST_PAGE_TABLES,
+               cu_device));
+    if (status != UCS_OK) {
+        coherent = 0;
+    }
+out:
+    return coherent;
+#else
+    return 0;
+#endif
+}
+
 UCS_STATIC_INIT
 {
     /* coverity[check_return] */

src/uct/cuda/base/cuda_md.h

@@ -42,4 +42,12 @@ uct_cuda_base_query_devices(uct_md_h md,
  */
 ucs_status_t uct_cuda_base_check_device_name(const uct_iface_params_t *params);
 
+
+/**
+ * Check whether the platform is coherent.
+ *
+ * @return 1 if coherent, or 0 otherwise.
+ */
+int uct_cuda_base_is_coherent();
+
 #endif

src/uct/cuda/cuda_ipc/cuda_ipc_iface.c

@@ -20,6 +20,18 @@
 #include <pthread.h>
 #include <nvml.h>
 
+
+typedef enum {
+    UCT_CUDA_IPC_DEVICE_ADDR_FLAG_MNNVL = UCS_BIT(0)
+} uct_cuda_ipc_device_addr_flags_t;
+
+
+typedef struct {
+    uint64_t     system_uuid;
+    uint8_t      flags; /* uct_cuda_ipc_device_addr_flags_t */
+} UCS_S_PACKED uct_cuda_ipc_device_addr_t;
+
+
 static ucs_config_field_t uct_cuda_ipc_iface_config_table[] = {
 
     {"", "", NULL,
@@ -63,7 +75,18 @@ static void UCS_CLASS_DELETE_FUNC_NAME(uct_cuda_ipc_iface_t)(uct_iface_t*);
 ucs_status_t uct_cuda_ipc_iface_get_device_address(uct_iface_t *tl_iface,
                                                    uct_device_addr_t *addr)
 {
-    *(uint64_t*)addr = ucs_get_system_id();
+    uct_cuda_ipc_device_addr_t *dev_addr = (uct_cuda_ipc_device_addr_t*)addr;
+    uct_cuda_ipc_iface_t *iface          = ucs_derived_of(tl_iface,
+                                                          uct_cuda_ipc_iface_t);
+    uct_cuda_ipc_md_t *md                = ucs_derived_of(iface->super.super.md,
+                                                           uct_cuda_ipc_md_t);
+
+    if (md->enable_mnnvl) {
+        dev_addr->flags = UCT_CUDA_IPC_DEVICE_ADDR_FLAG_MNNVL;
+    }
+
+    dev_addr->system_uuid = ucs_get_system_id();
+
     return UCS_OK;
 }
 
@@ -74,60 +97,53 @@ static ucs_status_t uct_cuda_ipc_iface_get_address(uct_iface_h tl_iface,
     return UCS_OK;
 }
 
-static int uct_cuda_ipc_iface_is_mnnvl_supported(uct_cuda_ipc_md_t *md)
+static int
+uct_cuda_ipc_iface_mnnvl_supported(uct_cuda_ipc_md_t *md,
+                                   const uct_cuda_ipc_device_addr_t *dev_addr,
+                                   size_t dev_addr_len)
 {
-#if HAVE_CUDA_FABRIC
-    CUdevice cu_device;
-    int coherent;
-    ucs_status_t status;
-
-    status = UCT_CUDADRV_FUNC_LOG_ERR(cuDeviceGet(&cu_device, 0));
-    if (status != UCS_OK) {
-        return 0;
-    }
-
-    status = UCT_CUDADRV_FUNC_LOG_ERR(
-            cuDeviceGetAttribute(&coherent,
-                                 CU_DEVICE_ATTRIBUTE_PAGEABLE_MEMORY_ACCESS_USES_HOST_PAGE_TABLES,
-                                 cu_device));
-    if (status != UCS_OK) {
-        return 0;
+    if (md->enable_mnnvl && (dev_addr_len != sizeof(uint64_t))) {
+        ucs_assertv(dev_addr_len >= sizeof(uct_cuda_ipc_device_addr_t),
+                    "dev_addr_len=%zu", dev_addr_len);
+        return (dev_addr->flags & UCT_CUDA_IPC_DEVICE_ADDR_FLAG_MNNVL);
     }
 
-    return coherent && (md->enable_mnnvl != UCS_NO);
-#else
     return 0;
-#endif
 }
 
 static int
 uct_cuda_ipc_iface_is_reachable_v2(const uct_iface_h tl_iface,
                                    const uct_iface_is_reachable_params_t *params)
 {
-    uct_base_iface_t *base_iface = ucs_derived_of(tl_iface, uct_base_iface_t);
-    uct_cuda_ipc_md_t *md        = ucs_derived_of(base_iface->md, uct_cuda_ipc_md_t);
+    uct_cuda_ipc_iface_t *iface = ucs_derived_of(tl_iface, uct_cuda_ipc_iface_t);
+    uct_cuda_ipc_md_t *md       = ucs_derived_of(iface->super.super.md,
+                                                 uct_cuda_ipc_md_t);
+    const uct_cuda_ipc_device_addr_t *dev_addr;
+    size_t dev_addr_len;
+    int same_uuid;
 
     if (!uct_iface_is_reachable_params_addrs_valid(params)) {
         return 0;
     }
 
-    if (getpid() == *(pid_t*)params->iface_addr) {
+    dev_addr_len = UCS_PARAM_VALUE(UCT_IFACE_IS_REACHABLE_FIELD, params,
+                                   device_addr_length, DEVICE_ADDR_LENGTH,
+                                   sizeof(dev_addr->system_uuid));
+    dev_addr     = (const uct_cuda_ipc_device_addr_t *)params->device_addr;
+    same_uuid    = (ucs_get_system_id() == dev_addr->system_uuid);
+
+    if ((getpid() == *(pid_t*)params->iface_addr) && same_uuid) {
         uct_iface_fill_info_str_buf(params, "same process");
         return 0;
     }
 
-    /* Either multi-node NVLINK should be supported or iface has to be on the
-     * same node for cuda-ipc to be reachable */
-    if ((ucs_get_system_id() != *((const uint64_t*)params->device_addr)) &&
-        !uct_cuda_ipc_iface_is_mnnvl_supported(md)) {
-        uct_iface_fill_info_str_buf(params,
-                                    "different system id %"PRIx64" vs %"PRIx64"",
-                                    ucs_get_system_id(),
-                                    *((const uint64_t*)params->device_addr));
-        return 0;
+    if (same_uuid ||
+        uct_cuda_ipc_iface_mnnvl_supported(md, dev_addr, dev_addr_len)) {
+        return uct_iface_scope_is_reachable(tl_iface, params);
     }
 
-    return uct_iface_scope_is_reachable(tl_iface, params);
+    uct_iface_fill_info_str_buf(params, "MNNVL is not supported");
+    return 0;
 }
 
 static double uct_cuda_ipc_iface_get_bw()
@@ -245,15 +261,17 @@ static ucs_status_t uct_cuda_ipc_iface_query(uct_iface_h tl_iface,
     uct_base_iface_query(&iface->super.super, iface_attr);
 
     iface_attr->iface_addr_len          = sizeof(pid_t);
-    iface_attr->device_addr_len         = sizeof(uint64_t);
+    iface_attr->device_addr_len         = md->enable_mnnvl ?
+                                          sizeof(uct_cuda_ipc_device_addr_t) :
+                                          sizeof(uint64_t);
     iface_attr->ep_addr_len             = 0;
     iface_attr->max_conn_priv           = 0;
     iface_attr->cap.flags               = UCT_IFACE_FLAG_ERRHANDLE_PEER_FAILURE |
                                           UCT_IFACE_FLAG_CONNECT_TO_IFACE |
                                           UCT_IFACE_FLAG_PENDING          |
                                           UCT_IFACE_FLAG_GET_ZCOPY        |
                                           UCT_IFACE_FLAG_PUT_ZCOPY;
-    if (uct_cuda_ipc_iface_is_mnnvl_supported(md)) {
+    if (md->enable_mnnvl) {
         iface_attr->cap.flags |= UCT_IFACE_FLAG_INTER_NODE;
     }
 
@@ -614,7 +632,7 @@ uct_cuda_ipc_query_devices(
     uct_device_type_t dev_type = UCT_DEVICE_TYPE_SHM;
     uct_cuda_ipc_md_t *md      = ucs_derived_of(uct_md, uct_cuda_ipc_md_t);
 
-    if (uct_cuda_ipc_iface_is_mnnvl_supported(md)) {
+    if (md->enable_mnnvl) {
         dev_type = UCT_DEVICE_TYPE_NET;
     }
 

src/uct/cuda/cuda_ipc/cuda_ipc_md.c

@@ -17,6 +17,7 @@
 #include <ucs/debug/memtrack_int.h>
 #include <ucs/type/class.h>
 #include <ucs/profile/profile.h>
+#include <ucs/sys/string.h>
 #include <uct/api/v2/uct_v2.h>
 #include <sys/types.h>
 #include <unistd.h>
@@ -407,6 +408,63 @@ uct_cuda_ipc_mem_dereg(uct_md_h md, const uct_md_mem_dereg_params_t *params)
     return UCS_OK;
 }
 
+static int
+uct_cuda_ipc_md_check_fabric_info(uct_cuda_ipc_md_t *md,
+                                  ucs_ternary_auto_value_t mnnvl_enable)
+{
+    int mnnvl_supported = 0;
+#if HAVE_NVML_FABRIC_INFO
+    nvmlGpuFabricInfoV_t fabric_info;
+    nvmlDevice_t device;
+    ucs_status_t status;
+    char buf[64];
+
+    if (!uct_cuda_base_is_coherent() || (mnnvl_enable == UCS_NO)) {
+        goto out;
+    }
+
+    status = UCT_NVML_FUNC(nvmlInit_v2(), UCS_LOG_LEVEL_DIAG);
+    if (status != UCS_OK) {
+        goto out;
+    }
+
+    status = UCT_NVML_FUNC_LOG_ERR(nvmlDeviceGetHandleByIndex(0, &device));
+    if (status != UCS_OK) {
+        goto out_sd;
+    }
+
+    fabric_info.version = nvmlGpuFabricInfo_v2;
+    status              = UCT_NVML_FUNC_LOG_ERR(
+                             nvmlDeviceGetGpuFabricInfoV(device, &fabric_info));
+    if (status != UCS_OK) {
+        goto out_sd;
+    }
+
+    ucs_debug("fabric_info: healthmask=%u state=%u status=%u clique=%u uuid=%s",
+              fabric_info.healthMask, fabric_info.state, fabric_info.status,
+              fabric_info.cliqueId,
+              ucs_str_dump_hex(
+                  fabric_info.clusterUuid, NVML_GPU_FABRIC_UUID_LEN, buf,
+                  sizeof(buf), SIZE_MAX));
+
+    if ((fabric_info.state != NVML_GPU_FABRIC_STATE_COMPLETED) ||
+        (fabric_info.status != NVML_SUCCESS)) {
+        goto out_sd;
+    }
+
+    mnnvl_supported = 1;
+
+out_sd:
+    UCT_NVML_FUNC_LOG_ERR(nvmlShutdown());
+out:
+#endif
+    if ((mnnvl_enable == UCS_YES) && !mnnvl_supported) {
+        ucs_error("multi-node NVLINK support is requested but not supported");
+    }
+
+    return mnnvl_supported;
+}
+
 static void uct_cuda_ipc_md_close(uct_md_h md)
 {
     ucs_free(md);
@@ -436,9 +494,10 @@ uct_cuda_ipc_md_open(uct_component_t *component, const char *md_name,
 
     md->super.ops       = &md_ops;
     md->super.component = &uct_cuda_ipc_component.super;
-    md->enable_mnnvl    = ipc_config->enable_mnnvl;
+    md->enable_mnnvl    = uct_cuda_ipc_md_check_fabric_info(
+                                                  md, ipc_config->enable_mnnvl);
     *md_p               = &md->super;
-  
+
     return UCS_OK;
 }
 

src/uct/cuda/cuda_ipc/cuda_ipc_md.h

@@ -42,7 +42,7 @@ typedef CUipcMemHandle uct_cuda_ipc_md_handle_t;
  */
 typedef struct uct_cuda_ipc_md {
     uct_md_t                 super;   /**< Domain info */
-    ucs_ternary_auto_value_t enable_mnnvl;
+    int                      enable_mnnvl; /**< Multi-node NVLINK support status */
 } uct_cuda_ipc_md_t;