UCP/PROTO/RNDV: Isolate remote system device namespace

AUTHORS

@@ -59,6 +59,7 @@ Ivan Kochin <ikochin@nvidia.com>
 Jakir Kham <jakirkham@gmail.com>
 Jason Gunthorpe <jgg@mellanox.com>
 Jeff Daily <jeff.daily@amd.com>
+Jeffrey Mahou <jmahou@nvidia.com>
 Jianxin Xiong <jianxin.xiong@intel.com>
 John Snyder <jms285@duke.edu>
 Jonas Zhou <JonasZhou@zhaoxin.com>

src/ucp/proto/proto_common.c

@@ -496,6 +496,7 @@ ucp_proto_common_filter_min_frag(const ucp_proto_init_params_t *params,
     ucp_md_index_t md_index         = context->tl_rscs[rsc_index].md_index;
     const uct_md_attr_v2_t *md_attr = &context->tl_mds[md_index].attr;
     const uct_iface_attr_t *iface_attr;
+    ucs_sys_device_t lane_sys_dev;
     size_t max_iov, tl_min_frag, tl_max_frag;
 
     /* Check memory registration capabilities for zero-copy case */
@@ -532,6 +533,19 @@ ucp_proto_common_filter_min_frag(const ucp_proto_init_params_t *params,
                       ucs_memory_type_names[reg_mem_type]);
             return 0;
         }
+
+        /* The lane's device and the memory that will actually be registered
+         * on it (which may differ from select_param's buffer, e.g. a staging
+         * fragment for mtype protocols) must be topologically reachable. */
+        lane_sys_dev = ucp_proto_common_get_sys_dev(params, lane);
+        if (!ucs_topo_is_reachable(lane_sys_dev,
+                                   common_params->reg_mem_info.sys_dev)) {
+            ucs_trace("%s: no reachability between lane_sys_dev=%u and "
+                      "reg_mem_sys_dev=%u",
+                      lane_desc, lane_sys_dev,
+                      common_params->reg_mem_info.sys_dev);
+            return 0;
+        }
     }
 
     iface_attr = ucp_proto_common_get_iface_attr(params, lane);
@@ -593,7 +607,6 @@ ucp_proto_common_find_lanes(const ucp_proto_init_params_t *params,
     ucp_md_index_t md_index;
     ucp_lane_map_t lane_map;
     char lane_desc[64];
-    ucs_sys_device_t lane_sys_dev;
     ucs_status_t status;
 
     if (max_lanes == 0) {
@@ -716,15 +729,6 @@ ucp_proto_common_find_lanes(const ucp_proto_init_params_t *params,
             }
         }
 
-        /* The two devices must also have internal reachability */
-        lane_sys_dev = context->tl_rscs[rsc_index].tl_rsc.sys_device;
-        if (!ucs_topo_is_reachable(lane_sys_dev, select_param->sys_dev)) {
-            ucs_trace("%s: no reachability between lane_sys_dev=%u and "
-                      "sys_dev=%u",
-                      lane_desc, lane_sys_dev, select_param->sys_dev);
-            continue;
-        }
-
         ucs_trace("%s: added as lane %d", lane_desc, lane);
         lanes[num_lanes++] = lane;
     }

src/ucp/rndv/proto_rndv.c

@@ -368,11 +368,12 @@ static ucp_proto_select_param_t ucp_proto_rndv_remote_select_param_init(
                                     select_param->dt_class,
                                     &mem_info, select_param->sg_count);
     } else {
-        /* If we know the remote buffer parameters, these are actually the local
-         * parameters for the remote protocol
-         */
-        mem_info.sys_dev = init_params->rkey_config_key->sys_dev;
         mem_info.type    = init_params->rkey_config_key->mem_type;
+        /* If the local and remote memory types match, use the local system
+         * device as a namespace-safe proxy for remote protocol estimation. */
+        mem_info.sys_dev = (mem_info.type == select_param->mem_type) ?
+                                   select_param->sys_dev :
+                                   UCS_SYS_DEVICE_ID_UNKNOWN;
         mem_info.flags   = ucp_proto_rndv_rkey_mem_flags_estimate(init_params);
         ucp_proto_select_param_init(&remote_select_param, params->remote_op_id,
                                     op_attr_mask, params->remote_op_flags,

test/gtest/ucp/test_ucp_peer_failure.cc

@@ -1013,8 +1013,11 @@ UCS_TEST_P(test_ucp_peer_failure_rndv_put_ppln_abort, rtr_mtype)
     rndv_progress_failure_test(rndv_mode::put_ppln, true);
 }
 
+/* Use host staging so this test still selects rndv/put/mtype when CUDA
+ * fragments are unreachable from the network lane. The user buffer remains
+ * CUDA. */
 UCS_TEST_P(test_ucp_peer_failure_rndv_put_ppln_abort, pipeline,
-           "RNDV_FRAG_SIZE=host:8K,cuda:8K")
+           "RNDV_FRAG_MEM_TYPES=host", "RNDV_FRAG_SIZE=host:8K")
 {
     rndv_progress_failure_test(rndv_mode::put_ppln, true);
 }