[TEMP][DO NOT REVIEW] Isolate PR 11685 mock topology state

AUTHORS

@@ -58,6 +58,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>

contrib/test_jenkins.sh

@@ -1024,7 +1024,7 @@ run_malloc_hook_gtest() {
 }
 
 set_gtest_common_test_flags() {
-	export GTEST_RANDOM_SEED=0
+	export GTEST_RANDOM_SEED=37290
 	export GTEST_SHUFFLE=1
 	# Run UCT tests for TCP over fastest device only
 	export GTEST_UCT_TCP_FASTEST_DEV=1
@@ -1040,7 +1040,7 @@ set_gtest_make_test_flags() {
 	# Report TOP-20 longest test at the end of testing
 	export GTEST_REPORT_LONGEST_TESTS=20
 
-	GTEST_EXTRA_ARGS=""
+	GTEST_EXTRA_ARGS="-s 11689"
 	if [ "$TEST_PERF" == 1 ] && [[ "$VALGRIND_CHECK" != "yes" ]]
 	then
 		# Check performance with 10 retries and 2 seconds interval

src/ucp/rndv/proto_rndv.c

@@ -331,9 +331,13 @@ static ucp_proto_select_param_t ucp_proto_rndv_remote_select_param_init(
                                     &mem_info, select_param->sg_count);
     } else {
         /* If we know the remote buffer parameters, these are actually the local
-         * parameters for the remote protocol
+         * parameters for the remote protocol. The rkey system device belongs
+         * to the remote topology namespace, so it cannot represent a local
+         * device while estimating the protocol that the peer will select.
+         * The peer will use its actual local system device when selecting the
+         * protocol.
          */
-        mem_info.sys_dev = init_params->rkey_config_key->sys_dev;
+        mem_info.sys_dev = UCS_SYS_DEVICE_ID_UNKNOWN;
         mem_info.type    = init_params->rkey_config_key->mem_type;
         mem_info.flags   = ucp_proto_rndv_rkey_mem_flags_estimate(init_params);
         ucp_proto_select_param_init(&remote_select_param, params->remote_op_id,

src/ucp/rndv/rndv_get.c

@@ -326,6 +326,7 @@ static void
 ucp_proto_rndv_get_mtype_probe(const ucp_proto_init_params_t *init_params)
 {
     ucp_context_t *context = init_params->worker->context;
+    ucp_proto_rndv_mtype_init_params_t mtype_init_params;
     ucs_memory_type_t frag_mem_type;
     ucp_md_map_t mdesc_md_map;
     ucs_status_t status;
@@ -347,7 +348,9 @@ ucp_proto_rndv_get_mtype_probe(const ucp_proto_init_params_t *init_params)
             continue;
         }
 
-        ucp_proto_rndv_get_common_probe(init_params,
+        ucp_proto_rndv_mtype_init_params_prepare(init_params, &frag_mem_info,
+                                                 &mtype_init_params);
+        ucp_proto_rndv_get_common_probe(&mtype_init_params.super,
                                         UCS_BIT(UCP_RNDV_MODE_GET_PIPELINE),
                                         frag_size, UCT_EP_OP_PUT_ZCOPY, 0,
                                         mdesc_md_map, 1, &frag_mem_info);

src/ucp/rndv/rndv_mtype.inl

@@ -44,6 +44,24 @@ ucp_proto_rndv_mtype_init(const ucp_proto_init_params_t *init_params,
     return UCS_OK;
 }
 
+typedef struct {
+    ucp_proto_init_params_t  super;
+    ucp_proto_select_param_t select_param;
+} ucp_proto_rndv_mtype_init_params_t;
+
+static UCS_F_ALWAYS_INLINE void ucp_proto_rndv_mtype_init_params_prepare(
+        const ucp_proto_init_params_t *init_params,
+        const ucp_memory_info_t *frag_mem_info,
+        ucp_proto_rndv_mtype_init_params_t *mtype_init_params)
+{
+    /* Use the staging buffer device for local lane modeling without changing
+     * the protocol selection lookup key. */
+    mtype_init_params->select_param         = *init_params->select_param;
+    mtype_init_params->select_param.sys_dev = frag_mem_info->sys_dev;
+    mtype_init_params->super                = *init_params;
+    mtype_init_params->super.select_param   = &mtype_init_params->select_param;
+}
+
 static UCS_F_ALWAYS_INLINE ucs_status_t
 ucp_proto_rndv_mtype_request_init(ucp_request_t *req,
                                   ucs_memory_type_t frag_mem_type,

src/ucp/rndv/rndv_put.c

@@ -597,6 +597,7 @@ static void
 ucp_proto_rndv_put_mtype_probe(const ucp_proto_init_params_t *init_params)
 {
     ucp_context_t *context = init_params->worker->context;
+    ucp_proto_rndv_mtype_init_params_t mtype_init_params;
     ucs_memory_type_t frag_mem_type;
     uct_completion_callback_t comp_cb;
     ucp_md_map_t mdesc_md_map;
@@ -641,10 +642,14 @@ ucp_proto_rndv_put_mtype_probe(const ucp_proto_init_params_t *init_params)
         comp_cb = ucp_proto_rndv_put_mtype_completion;
     }
 
-    ucp_proto_rndv_put_common_probe(
-            init_params, UCS_BIT(UCP_RNDV_MODE_PUT_PIPELINE), frag_size,
-            UCT_EP_OP_GET_ZCOPY, flags, mdesc_md_map, comp_cb, 1,
-            UCP_WORKER_STAT_RNDV_PUT_MTYPE_ZCOPY, &frag_mem_info);
+    ucp_proto_rndv_mtype_init_params_prepare(init_params, &frag_mem_info,
+                                             &mtype_init_params);
+    ucp_proto_rndv_put_common_probe(&mtype_init_params.super,
+                                    UCS_BIT(UCP_RNDV_MODE_PUT_PIPELINE),
+                                    frag_size, UCT_EP_OP_GET_ZCOPY, flags,
+                                    mdesc_md_map, comp_cb, 1,
+                                    UCP_WORKER_STAT_RNDV_PUT_MTYPE_ZCOPY,
+                                    &frag_mem_info);
 }
 
 static void

test/gtest/ucp/test_ucp_proto_mock.cc

@@ -13,6 +13,7 @@ extern "C" {
 #include <uct/base/uct_iface.h>
 #include <ucp/proto/proto_debug.h>
 #include <ucp/proto/proto_select.inl>
+#include <ucp/rndv/proto_rndv.h>
 #include <ucs/memory/numa.h>
 #include <ucs/sys/sys.h>
 #include <ucs/sys/topo/base/topo.h>
@@ -867,6 +868,70 @@ UCS_TEST_P(test_ucp_proto_mock_rcx, rma_put_2_lanes,
     }, key, 0);
 }
 
+UCS_TEST_P(test_ucp_proto_mock_rcx, rndv_remote_sys_dev_namespace,
+           "IB_NUM_PATHS?=1", "RNDV_SCHEME=put_zcopy")
+{
+    constexpr size_t select_size = UCS_MBYTE;
+    uint8_t remote               = 0;
+    auto memh           = mem_map(receiver(), &remote, sizeof(remote));
+    auto rkey_packed    = rkey_pack(receiver(), memh);
+    auto rkey           = rkey_unpack(sender().ep(), rkey_packed);
+    ucp_worker_h worker = sender().worker();
+    ucp_worker_cfg_index_t ep_cfg_index   = ep_config_index(sender());
+    const ucs_sys_device_t remote_sys_dev = get_mock_sys_dev_by_name(
+            "mock_0:1");
+    const ucp_rkey_config_t *base_rkey_config =
+            &ucs_array_elem(&worker->rkey_config, rkey->cfg_index);
+    const ucp_ep_config_t *ep_config      = ucp_worker_ep_config(worker,
+                                                                 ep_cfg_index);
+    ucp_rkey_config_key_t rkey_config_key = base_rkey_config->key;
+    ucs_sys_dev_distance_t lanes_distance[UCP_MAX_LANES];
+    ucp_worker_cfg_index_t rkey_cfg_index;
+    ucp_rkey_config_t *rkey_config;
+    ucp_proto_select_param_t select_param;
+    ucp_memory_info_t mem_info;
+    const ucp_proto_select_elem_t *select_elem;
+    const ucp_proto_threshold_elem_t *threshold;
+    const ucp_proto_rndv_ctrl_priv_t *rpriv;
+    ucp_lane_index_t lane;
+
+    ASSERT_EQ(UCS_SYS_DEVICE_ID_UNKNOWN, rkey_config_key.sys_dev);
+    ASSERT_NE(UCS_SYS_DEVICE_ID_UNKNOWN, remote_sys_dev);
+
+    for (lane = 0; lane < ep_config->key.num_lanes; ++lane) {
+        lanes_distance[lane] = base_rkey_config->lanes_distance[lane];
+    }
+
+    /* Simulate an rkey received from a peer whose process-local sys_dev value
+     * happens to name a device in this process as well. */
+    rkey_config_key.sys_dev = remote_sys_dev;
+    ASSERT_UCS_OK(ucp_worker_rkey_config_get(worker, &rkey_config_key,
+                                             lanes_distance, &rkey_cfg_index));
+
+    mem_info.type    = UCS_MEMORY_TYPE_HOST;
+    mem_info.sys_dev = get_mock_sys_dev_by_name("mock_1:1");
+    mem_info.flags   = UCS_MEM_FLAG_REGISTRABLE;
+    ucp_proto_select_param_init(&select_param, UCP_OP_ID_RNDV_RECV, 0, 0,
+                                UCP_DATATYPE_CONTIG, &mem_info, 1);
+
+    rkey_config = &ucs_array_elem(&worker->rkey_config, rkey_cfg_index);
+    select_elem = ucp_proto_select_lookup_slow(worker,
+                                               &rkey_config->proto_select, 1,
+                                               ep_cfg_index, rkey_cfg_index,
+                                               &select_param);
+    ASSERT_NE(nullptr, select_elem);
+
+    threshold = ucp_proto_thresholds_search_slow(select_elem->thresholds,
+                                                 select_size);
+    ASSERT_STREQ("rndv/rtr", threshold->proto_config.proto->name);
+    rpriv = static_cast<const ucp_proto_rndv_ctrl_priv_t*>(
+            threshold->proto_config.priv);
+
+    EXPECT_EQ(remote_sys_dev, rkey_config->key.sys_dev);
+    EXPECT_EQ(UCS_SYS_DEVICE_ID_UNKNOWN,
+              rpriv->remote_proto_config.select_param.sys_dev);
+}
+
 UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_proto_mock_rcx, rcx, "rc_x")
 
 class test_ucp_proto_mock_rcx2 : public test_ucp_proto_mock {
@@ -1280,6 +1345,148 @@ UCS_TEST_P(test_ucp_proto_mock_gpu, cuda_managed_ppln_host_frag,
 UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_proto_mock_gpu, rcx_gpu,
                               "rc_x,cuda,rocm")
 
+class test_ucp_proto_mock_mtype_sys_dev : public test_ucp_proto_mock {
+public:
+    test_ucp_proto_mock_mtype_sys_dev()
+    {
+        mock_transport("rc_mlx5");
+    }
+
+    virtual void init() override
+    {
+        auto iface_attr_func = [](uct_iface_attr_t &iface_attr) {
+            iface_attr.cap.am.max_short = 2000;
+            iface_attr.bandwidth.shared = 28e9;
+            iface_attr.latency.c        = 600e-9;
+            iface_attr.latency.m        = 1e-9;
+        };
+
+        add_mock_iface("mock_0:1", iface_attr_func);
+        add_mock_iface("mock_1:1", iface_attr_func);
+        test_ucp_proto_mock::init();
+
+        s_user_sys_dev = get_mock_sys_dev_by_name("mock_0:1");
+        s_frag_sys_dev = get_mock_sys_dev_by_name("mock_1:1");
+        EXPECT_UCS_OK(ucs_topo_sys_device_set_class(
+                s_user_sys_dev, UCS_TOPO_DEVICE_CLASS_ACC));
+        EXPECT_UCS_OK(ucs_topo_sys_device_set_class(
+                s_frag_sys_dev, UCS_TOPO_DEVICE_CLASS_ACC));
+
+        set_frag_sys_dev(sender());
+        set_frag_sys_dev(receiver());
+    }
+
+    virtual void cleanup() override
+    {
+        test_ucp_proto_mock::cleanup();
+
+        if (s_user_sys_dev != UCS_SYS_DEVICE_ID_UNKNOWN) {
+            EXPECT_UCS_OK(ucs_topo_sys_device_set_class(
+                    s_user_sys_dev, UCS_TOPO_DEVICE_CLASS_UNKNOWN));
+        }
+        if (s_frag_sys_dev != UCS_SYS_DEVICE_ID_UNKNOWN) {
+            EXPECT_UCS_OK(ucs_topo_sys_device_set_class(
+                    s_frag_sys_dev, UCS_TOPO_DEVICE_CLASS_UNKNOWN));
+        }
+
+        s_frag_sys_dev = UCS_SYS_DEVICE_ID_UNKNOWN;
+        s_user_sys_dev = UCS_SYS_DEVICE_ID_UNKNOWN;
+    }
+
+protected:
+    void check_mtype_uses_frag_sys_dev(ucp_operation_id_t op_id,
+                                       const char *proto_name)
+    {
+        static constexpr size_t msg_size = UCS_KBYTE;
+        uint8_t remote                   = 0;
+        ucp_worker_h worker              = sender().worker();
+
+        if (worker->mem_type_ep[UCS_MEMORY_TYPE_CUDA_MANAGED] == NULL) {
+            UCS_TEST_SKIP_R("CUDA managed memory type endpoint is unavailable");
+        }
+
+        auto memh           = mem_map(receiver(), &remote, sizeof(remote));
+        auto rkey_packed    = rkey_pack(receiver(), memh);
+        auto rkey           = rkey_unpack(sender().ep(), rkey_packed);
+        ucp_worker_cfg_index_t ep_cfg_index = ep_config_index(sender());
+        ucp_rkey_config_t *rkey_config = &ucs_array_elem(&worker->rkey_config,
+                                                         rkey->cfg_index);
+        ucp_memory_info_t mem_info;
+        ucp_proto_select_param_t select_param;
+        const ucp_proto_select_elem_t *select_elem;
+        const ucp_proto_threshold_elem_t *threshold;
+        ucp_proto_query_attr_t attr;
+        unsigned user_ordinal, frag_ordinal;
+
+        ASSERT_NE(s_user_sys_dev, s_frag_sys_dev);
+        user_ordinal = ucs_topo_sys_device_get_bdf_class_ordinal(
+                s_user_sys_dev);
+        frag_ordinal = ucs_topo_sys_device_get_bdf_class_ordinal(
+                s_frag_sys_dev);
+        ASSERT_NE(UCS_SYS_DEVICE_ORDINAL_INVALID, user_ordinal);
+        ASSERT_NE(UCS_SYS_DEVICE_ORDINAL_INVALID, frag_ordinal);
+        ASSERT_NE(user_ordinal % 2, frag_ordinal % 2);
+
+        mem_info.type    = UCS_MEMORY_TYPE_CUDA_MANAGED;
+        mem_info.sys_dev = s_user_sys_dev;
+        mem_info.flags   = UCS_MEM_FLAG_REGISTRABLE;
+        ucp_proto_select_param_init(&select_param, op_id, 0, 0,
+                                    UCP_DATATYPE_CONTIG, &mem_info, 1);
+        select_elem = ucp_proto_select_lookup_slow(worker,
+                                                   &rkey_config->proto_select,
+                                                   1, ep_cfg_index,
+                                                   rkey->cfg_index,
+                                                   &select_param);
+        ASSERT_NE(nullptr, select_elem);
+
+        threshold = ucp_proto_thresholds_search_slow(select_elem->thresholds,
+                                                     msg_size);
+        ASSERT_STREQ(proto_name, threshold->proto_config.proto->name);
+        ucp_proto_config_query(worker, &threshold->proto_config, msg_size,
+                               &attr);
+        EXPECT_STREQ((frag_ordinal % 2) ? "rc_mlx5/mock_1:1" :
+                                          "rc_mlx5/mock_0:1",
+                     attr.config);
+    }
+
+private:
+    static void set_frag_sys_dev(entity &e)
+    {
+        ucp_memory_info_t mem_info;
+        ucp_md_index_t md_index;
+
+        ASSERT_UCS_OK(ucp_mm_get_alloc_md_index(e.ucph(), UCS_MEMORY_TYPE_HOST,
+                                                UCS_SYS_DEVICE_ID_UNKNOWN,
+                                                &md_index, &mem_info));
+        e.ucph()->alloc_md[UCS_MEMORY_TYPE_HOST].sys_dev = s_frag_sys_dev;
+    }
+
+    static ucs_sys_device_t s_frag_sys_dev;
+    static ucs_sys_device_t s_user_sys_dev;
+};
+
+ucs_sys_device_t test_ucp_proto_mock_mtype_sys_dev::s_frag_sys_dev =
+        UCS_SYS_DEVICE_ID_UNKNOWN;
+ucs_sys_device_t test_ucp_proto_mock_mtype_sys_dev::s_user_sys_dev =
+        UCS_SYS_DEVICE_ID_UNKNOWN;
+
+UCS_TEST_P(test_ucp_proto_mock_mtype_sys_dev, get_uses_frag_sys_dev,
+           "RNDV_SCHEME=get_ppln", "RNDV_THRESH=1", "RNDV_FRAG_MEM_TYPES=host",
+           "RNDV_FRAG_SIZE=host:8K", "IB_NUM_PATHS?=1", "MAX_RNDV_LANES=1")
+{
+    check_mtype_uses_frag_sys_dev(UCP_OP_ID_RNDV_RECV, "rndv/get/mtype");
+}
+
+UCS_TEST_P(test_ucp_proto_mock_mtype_sys_dev, put_uses_frag_sys_dev,
+           "RNDV_SCHEME=put_ppln", "RNDV_THRESH=1", "RNDV_FRAG_MEM_TYPES=host",
+           "RNDV_FRAG_SIZE=host:8K", "IB_NUM_PATHS?=1", "MAX_RNDV_LANES=1")
+{
+    check_mtype_uses_frag_sys_dev(UCP_OP_ID_RNDV_SEND, "rndv/put/mtype");
+}
+
+UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_proto_mock_mtype_sys_dev, rcx_gpu,
+                              "rc_x,cuda,rocm")
+
 class test_ucp_proto_mock_cuda_ipc : public test_ucp_proto_mock {
 public:
     test_ucp_proto_mock_cuda_ipc()