UCP/CORE: Ignore DPU in topology assignments

src/ucp/core/ucp_context.c

@@ -629,7 +629,10 @@ static ucs_config_field_t ucp_context_config_table[] = {
   {"GPU_NIC_ASSIGNMENT_MODE", "auto",
    "Assign NICs to GPUs within each topology group, and restrict the lanes\n"
    "for a GPU's memory to the NICs assigned to that GPU.\n"
-   "All ports of a NIC are assigned together.\n"
+   "All ports of a NIC are assigned together, and all GPUs with the same PCI\n"
+   "address are assigned together (e.g. MLOPart partitions of a GPU).\n"
+   "NICs without memory registration and DPUs are skipped during assignment.\n"
+   "NICs excluded by UCX_NET_DEVICES or UCX_TLS are ignored.\n"
    "The 'flip', 'round_robin' and 'shared' modes apply on any hardware.\n"
    "With 'flip' and 'round_robin', a group with fewer NICs than GPUs leaves\n"
    "some GPUs without any NICs.\n"
@@ -1397,13 +1400,52 @@ ucp_add_tl_resource(ucp_context_h context, ucp_md_index_t md_index,
     return UCS_OK;
 }
 
+static int
+ucp_groups_is_net_device_allowed(const ucp_tl_md_t *md,
+                                 const uct_tl_resource_desc_t *resource,
+                                 uint8_t rsc_flags)
+{
+    if ((resource->dev_type != UCT_DEVICE_TYPE_NET) ||
+        (resource->sys_device == UCS_SYS_DEVICE_ID_UNKNOWN)) {
+        return 0;
+    }
+
+    if (rsc_flags & UCP_TL_RSC_FLAG_AUX) {
+        ucs_debug(UCT_TL_RESOURCE_DESC_FMT
+                  " is not allowed for gpu-nic assignment: auxiliary transport",
+                  UCT_TL_RESOURCE_DESC_ARG(resource));
+        return 0;
+    }
+
+    if (!(md->attr.flags & UCT_MD_FLAG_REG)) {
+        ucs_debug(UCT_TL_RESOURCE_DESC_FMT
+                  " is not allowed for gpu-nic assignment: no memory "
+                  "registration",
+                  UCT_TL_RESOURCE_DESC_ARG(resource));
+        return 0;
+    }
+
+    if (ucs_topo_sys_device_get_flags(resource->sys_device) &
+        UCS_TOPO_DEVICE_FLAG_DPU) {
+        ucs_debug(UCT_TL_RESOURCE_DESC_FMT
+                  " is not allowed for gpu-nic assignment: dpu device",
+                  UCT_TL_RESOURCE_DESC_ARG(resource));
+        return 0;
+    }
+
+    ucs_trace(UCT_TL_RESOURCE_DESC_FMT " is allowed for gpu-nic assignment",
+              UCT_TL_RESOURCE_DESC_ARG(resource));
+    return 1;
+}
+
 static ucs_status_t
 ucp_add_tl_resources(ucp_context_h context, ucp_md_index_t md_index,
                      const ucp_config_t *config,
                      const ucs_string_set_t *aux_tls, unsigned *num_resources_p,
                      ucs_string_set_t avail_devices[],
                      ucs_string_set_t *avail_tls, uint64_t dev_cfg_masks[],
-                     uint64_t *tl_cfg_mask, ucp_tl_info_array_t *all_rscs)
+                     uint64_t *tl_cfg_mask, ucp_tl_info_array_t *all_rscs,
+                     ucs_sys_device_bitmap_t *allowed_net_devices)
 {
     ucp_tl_md_t *md                 = &context->tl_mds[md_index];
     ucp_tl_info_entry_t *added_rscs = NULL;
@@ -1481,6 +1523,11 @@ ucp_add_tl_resources(ucp_context_h context, ucp_md_index_t md_index,
             goto free_resources;
         }
 
+        if (ucp_groups_is_net_device_allowed(md, &tl_resources[i], rsc_flags)) {
+            UCS_STATIC_BITMAP_SET(allowed_net_devices,
+                                  tl_resources[i].sys_device);
+        }
+
         ++(*num_resources_p);
 
         if (added_rscs != NULL) {
@@ -1808,7 +1855,8 @@ ucp_add_component_resources(ucp_context_h context, ucp_rsc_index_t cmpt_index,
                             uint64_t dev_cfg_masks[], uint64_t *tl_cfg_mask,
                             const ucp_config_t *config,
                             const ucs_string_set_t *aux_tls,
-                            ucp_tl_info_array_t *all_rscs)
+                            ucp_tl_info_array_t *all_rscs,
+                            ucs_sys_device_bitmap_t *allowed_net_devices)
 {
     const ucp_tl_cmpt_t *tl_cmpt = &context->tl_cmpts[cmpt_index];
     size_t avail_mds             = config->max_component_mds;
@@ -1858,7 +1906,7 @@ ucp_add_component_resources(ucp_context_h context, ucp_rsc_index_t cmpt_index,
         status = ucp_add_tl_resources(context, md_index, config, aux_tls,
                                       &num_tl_resources, avail_devices,
                                       avail_tls, dev_cfg_masks, tl_cfg_mask,
-                                      all_rscs);
+                                      all_rscs, allowed_net_devices);
         if (status != UCS_OK) {
             uct_md_close(context->tl_mds[md_index].md);
             goto out;
@@ -2185,7 +2233,8 @@ static void ucp_fill_resources_reg_md_map_update(ucp_context_h context)
 }
 
 static ucs_status_t
-ucp_fill_resources(ucp_context_h context, const ucp_config_t *config)
+ucp_fill_resources(ucp_context_h context, const ucp_config_t *config,
+                   ucs_sys_device_bitmap_t *allowed_net_devices)
 {
     ucp_tl_info_array_t all_rscs = UCS_ARRAY_DYNAMIC_INITIALIZER;
     uint64_t dev_cfg_masks[UCT_DEVICE_TYPE_LAST] = {};
@@ -2294,7 +2343,7 @@ ucp_fill_resources(ucp_context_h context, const ucp_config_t *config)
         status = ucp_add_component_resources(context, i, avail_devices,
                                              &avail_tls, dev_cfg_masks,
                                              &tl_cfg_mask, config, &aux_tls,
-                                             &all_rscs);
+                                             &all_rscs, allowed_net_devices);
         if (status != UCS_OK) {
             goto err_free_resources;
         }
@@ -2783,9 +2832,10 @@ ucp_version_check(unsigned api_major_version, unsigned api_minor_version)
     ucs_debug("Configured with: %s", UCX_CONFIGURE_FLAGS);
 }
 
-static ucs_status_t
-ucp_context_gpu_nic_assignment_init(ucp_gpu_nic_assignment_mode_t mode,
-                                    ucp_gpu_nic_assignment_t **assignment_p)
+static ucs_status_t ucp_context_gpu_nic_assignment_init(
+        ucp_gpu_nic_assignment_mode_t mode,
+        const ucs_sys_device_bitmap_t *allowed_net_devices,
+        ucp_gpu_nic_assignment_t **assignment_p)
 {
     ucp_gpu_nic_assignment_t *assignment;
     ucs_topo_groups_t groups;
@@ -2810,7 +2860,7 @@ ucp_context_gpu_nic_assignment_init(ucp_gpu_nic_assignment_mode_t mode,
 
     ucs_debug("gpu-nic assignment mode %s", ucp_gpu_nic_assignment_modes[mode]);
 
-    status = ucs_topo_build_groups(&groups);
+    status = ucs_topo_build_groups(allowed_net_devices, &groups);
     if (status != UCS_OK) {
         return status;
     }
@@ -2855,6 +2905,8 @@ ucs_status_t ucp_init_version(unsigned api_major_version, unsigned api_minor_ver
                               const ucp_params_t *params, const ucp_config_t *config,
                               ucp_context_h *context_p)
 {
+    ucs_sys_device_bitmap_t allowed_net_devices =
+            UCS_STATIC_BITMAP_ZERO_INITIALIZER;
     ucp_config_t *dfl_config = NULL;
     ucp_context_t *context;
     ucs_status_t status;
@@ -2889,13 +2941,13 @@ ucs_status_t ucp_init_version(unsigned api_major_version, unsigned api_minor_ver
     UCP_THREAD_LOCK_INIT(&context->mt_lock);
 
     /* fill resources we should use */
-    status = ucp_fill_resources(context, config);
+    status = ucp_fill_resources(context, config, &allowed_net_devices);
     if (status != UCS_OK) {
         goto err_thread_lock_finalize;
     }
 
     status = ucp_context_gpu_nic_assignment_init(
-            context->config.ext.gpu_nic_assignment_mode,
+            context->config.ext.gpu_nic_assignment_mode, &allowed_net_devices,
             &context->gpu_nic_assignment);
     if (status != UCS_OK) {
         goto err_free_res;

src/ucp/core/ucp_gpu_nic_assignment.c

@@ -57,7 +57,7 @@ ucp_gpu_nic_assignment_get_gpu_idx(ucp_gpu_nic_assignment_mode_t mode,
 }
 
 static void
-ucp_gpu_nic_bitmap_add_nic(ucp_gpu_nic_sys_dev_bitmap_t *nic_sys_dev_bitmap,
+ucp_gpu_nic_bitmap_add_nic(ucs_sys_device_bitmap_t *nic_sys_dev_bitmap,
                            const ucs_topo_group_element_t *nic)
 {
     const ucs_sys_device_t *nic_sys_dev;
@@ -72,7 +72,7 @@ ucp_gpu_nic_bitmap_add_nic(ucp_gpu_nic_sys_dev_bitmap_t *nic_sys_dev_bitmap,
     }
 }
 
-const ucp_gpu_nic_sys_dev_bitmap_t *
+const ucs_sys_device_bitmap_t *
 ucp_gpu_nic_assignment_lookup(const ucp_gpu_nic_assignment_t *assignment,
                               ucs_sys_device_t gpu_sys_dev)
 {
@@ -135,12 +135,12 @@ static void
 ucp_gpu_nic_assignment_log_gpu(const ucp_gpu_nic_assignment_t *assignment,
                                const ucs_topo_group_t *group,
                                const ucs_topo_group_element_t *gpu,
-                               ucp_gpu_nic_sys_dev_bitmap_t *assigned_nics)
+                               ucs_sys_device_bitmap_t *assigned_nics)
 {
     UCS_STRING_BUFFER_ONSTACK(gpu_strb, 128);
     UCP_GPU_NIC_STRING_BUFFERS_ONSTACK(nic_strbs, 256);
     char bdf_name[UCS_SYS_BDF_NAME_MAX];
-    const ucp_gpu_nic_sys_dev_bitmap_t *nic_sys_dev_bitmap;
+    const ucs_sys_device_bitmap_t *nic_sys_dev_bitmap;
     const ucs_topo_group_element_t *nic;
 
     ucs_assert((gpu->num_sys_devs > 0) &&
@@ -182,7 +182,7 @@ ucp_gpu_nic_assignment_log_gpu(const ucp_gpu_nic_assignment_t *assignment,
 
 static void ucp_gpu_nic_assignment_log_unassigned_nics(
         const ucs_topo_groups_t *groups,
-        const ucp_gpu_nic_sys_dev_bitmap_t *assigned_nics)
+        const ucs_sys_device_bitmap_t *assigned_nics)
 {
     UCP_GPU_NIC_STRING_BUFFERS_ONSTACK(nic_strbs, 256);
     const ucs_topo_group_element_t *nic;
@@ -211,8 +211,7 @@ static void
 ucp_gpu_nic_assignment_log(const ucp_gpu_nic_assignment_t *assignment,
                            const ucs_topo_groups_t *groups)
 {
-    ucp_gpu_nic_sys_dev_bitmap_t assigned_nics =
-            UCS_STATIC_BITMAP_ZERO_INITIALIZER;
+    ucs_sys_device_bitmap_t assigned_nics = UCS_STATIC_BITMAP_ZERO_INITIALIZER;
     const ucs_topo_group_element_t *gpu;
     const ucs_topo_group_t *group;
 
@@ -331,13 +330,13 @@ ucp_gpu_nic_assignment_add_nic_to_gpu(ucp_gpu_nic_assignment_t *assignment,
                                       const ucs_topo_group_element_t *gpu,
                                       const ucs_topo_group_element_t *nic)
 {
-    ucp_gpu_nic_sys_dev_bitmap_t *nic_sys_dev_bitmap;
+    ucs_sys_device_bitmap_t *nic_sys_dev_bitmap;
 
     ucs_assert((gpu->num_sys_devs > 0) &&
                (gpu->sys_devs[0] != UCS_SYS_DEVICE_ID_UNKNOWN));
 
     nic_sys_dev_bitmap = ucs_const_cast(
-            ucp_gpu_nic_sys_dev_bitmap_t*,
+            ucs_sys_device_bitmap_t*,
             ucp_gpu_nic_assignment_lookup(assignment, gpu->sys_devs[0]));
     ucs_assert(nic_sys_dev_bitmap != NULL);
 
@@ -404,7 +403,7 @@ ucp_gpu_nic_assignment_build(const ucs_topo_groups_t *groups,
     return UCS_OK;
 }
 
-int ucp_gpu_nic_bitmap_get(const ucp_gpu_nic_sys_dev_bitmap_t *bitmap,
+int ucp_gpu_nic_bitmap_get(const ucs_sys_device_bitmap_t *bitmap,
                            ucs_sys_device_t net_sys_dev)
 {
     ucs_assert(bitmap != NULL);

src/ucp/core/ucp_gpu_nic_assignment.h

@@ -21,13 +21,9 @@ BEGIN_C_DECLS
 #define UCP_GPU_NIC_BITMAP_INDEX_INVALID UCS_SYS_DEVICE_ID_COUNT
 
 
-typedef ucs_static_bitmap_s(UCS_SYS_DEVICE_ID_COUNT)
-        ucp_gpu_nic_sys_dev_bitmap_t;
-
-
 struct ucp_gpu_nic_assignment {
     /* NIC sys_dev bitmaps referenced by bitmap_idx_by_gpu_sys_dev. */
-    ucp_gpu_nic_sys_dev_bitmap_t *nic_sys_dev_bitmaps;
+    ucs_sys_device_bitmap_t      *nic_sys_dev_bitmaps;
 
     /* Number of entries in nic_sys_dev_bitmaps. */
     size_t                       num_bitmaps;
@@ -66,7 +62,7 @@ ucp_gpu_nic_assignment_build(const ucs_topo_groups_t *groups,
  *         has no assigned NICs. Returns NULL if @a gpu_sys_dev is unknown or
  *         is not represented in the assignment.
  */
-const ucp_gpu_nic_sys_dev_bitmap_t *
+const ucs_sys_device_bitmap_t *
 ucp_gpu_nic_assignment_lookup(const ucp_gpu_nic_assignment_t *assignment,
                               ucs_sys_device_t gpu_sys_dev);
 
@@ -80,7 +76,7 @@ ucp_gpu_nic_assignment_lookup(const ucp_gpu_nic_assignment_t *assignment,
  * @return Nonzero if @a net_sys_dev is present in @a bitmap, or zero if the
  *         system device is unknown or the bit is not set.
  */
-int ucp_gpu_nic_bitmap_get(const ucp_gpu_nic_sys_dev_bitmap_t *bitmap,
+int ucp_gpu_nic_bitmap_get(const ucs_sys_device_bitmap_t *bitmap,
                            ucs_sys_device_t net_sys_dev);
 
 

src/ucp/core/ucp_types.h

@@ -249,6 +249,9 @@ typedef enum {
      * across the N GPUs of the group.
      * The group's NICs are assigned to the following GPU indices in order:
      * 0, 1, ..., N-1, N-1, ..., 1, 0, 0, 1, ...
+     * This targets architectures where only some of the NICs share a PCIe
+     * switch with a local SSD, so that each GPU gets NICs both with and without
+     * direct SSD access.
      */
     UCP_GPU_NIC_ASSIGNMENT_MODE_FLIP,
 

src/ucs/sys/topo/base/topo.c

@@ -517,6 +517,7 @@ ucs_topo_find_device_by_bus_id_value(const ucs_sys_bus_id_t *bus_id,
         device->pci_id          = pci_id;
         device->user_value      = user_value;
         device->device_class    = UCS_TOPO_DEVICE_CLASS_UNKNOWN;
+        device->flags           = 0;
         device->class_ordinal   = UCS_SYS_DEVICE_ORDINAL_INVALID;
         device->sibling_role    = UCS_TOPO_SIBLING_ROLE_NONE;
         device->sibling_sys_dev = UCS_SYS_DEVICE_ID_UNKNOWN;
@@ -1119,6 +1120,51 @@ ucs_status_t ucs_topo_sys_device_set_class(ucs_sys_device_t sys_dev,
     return status;
 }
 
+ucs_status_t
+ucs_topo_sys_device_add_flags(ucs_sys_device_t sys_dev, unsigned flags)
+{
+    ucs_status_t status = UCS_OK;
+
+    if (sys_dev == UCS_SYS_DEVICE_ID_UNKNOWN) {
+        ucs_error("system device %d is unknown", sys_dev);
+        return UCS_ERR_INVALID_PARAM;
+    }
+
+    ucs_spin_lock(&ucs_topo_global_ctx.lock);
+
+    if (sys_dev >= ucs_topo_global_ctx.num_devices) {
+        ucs_error("system device %d is invalid (max: %u)", sys_dev,
+                  ucs_topo_global_ctx.num_devices);
+        status = UCS_ERR_INVALID_PARAM;
+        goto out_unlock;
+    }
+
+    ucs_topo_global_ctx.devices[sys_dev].flags |= flags;
+
+out_unlock:
+    ucs_spin_unlock(&ucs_topo_global_ctx.lock);
+    return status;
+}
+
+unsigned ucs_topo_sys_device_get_flags(ucs_sys_device_t sys_dev)
+{
+    unsigned flags;
+
+    if (sys_dev == UCS_SYS_DEVICE_ID_UNKNOWN) {
+        return 0;
+    }
+
+    ucs_spin_lock(&ucs_topo_global_ctx.lock);
+    if (sys_dev < ucs_topo_global_ctx.num_devices) {
+        flags = ucs_topo_global_ctx.devices[sys_dev].flags;
+    } else {
+        flags = 0;
+    }
+    ucs_spin_unlock(&ucs_topo_global_ctx.lock);
+
+    return flags;
+}
+
 ucs_status_t
 ucs_topo_device_class_mark_incomplete(ucs_topo_device_class_t device_class)
 {
@@ -1446,14 +1492,16 @@ static void ucs_topo_release_devices()
     }
 }
 
-ucs_status_t ucs_topo_build_groups(ucs_topo_groups_t *groups_p)
+ucs_status_t
+ucs_topo_build_groups(const ucs_sys_device_bitmap_t *allowed_net_devices,
+                      ucs_topo_groups_t *groups_p)
 {
     ucs_status_t status;
 
     ucs_spin_lock(&ucs_topo_global_ctx.lock);
     status = ucs_topo_build_groups_inner(ucs_topo_global_ctx.devices,
                                          ucs_topo_global_ctx.num_devices,
-                                         groups_p);
+                                         allowed_net_devices, groups_p);
     ucs_spin_unlock(&ucs_topo_global_ctx.lock);
 
     return status;