[TEMP][DO NOT REVIEW] Diagnose PR 11685 Azure protocol selection

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/rndv/proto_rndv.c

@@ -341,9 +341,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

@@ -327,6 +327,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;
@@ -348,7 +349,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

@@ -598,6 +598,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;
@@ -642,10 +643,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_peer_failure.cc

@@ -840,12 +840,15 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure {
         ucs::mock mock;
 
         if (proto->name == self->m_proto_name) {
-            if (self->m_replace_ops) {
-                mock_rndv_ops(req->send.ep, mock);
-            }
+            self->m_target_proto_seen = true;
+            if (self->m_fail_target_proto) {
+                if (self->m_replace_ops) {
+                    mock_rndv_ops(req->send.ep, mock);
+                }
 
-            if (stage == self->m_proto_xfer_stage) {
-                self->close_peer();
+                if (stage == self->m_proto_xfer_stage) {
+                    self->close_peer();
+                }
             }
         }
 
@@ -904,6 +907,20 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure {
         setup_progress_mock(sender().worker(), m_progress_mock);
         setup_progress_mock(receiver().worker(), m_progress_mock);
 
+        if (mode == rndv_mode::put_ppln) {
+            /* The capability check in init() does not include staging-device
+             * reachability. Verify that this topology actually selected the
+             * protocol whose abort handling is under test. */
+            m_fail_target_proto = false;
+            m_target_proto_seen = false;
+            smoke_test(true);
+            if (!m_target_proto_seen) {
+                m_progress_mock.cleanup();
+                UCS_TEST_SKIP_R("rndv/put/mtype was not selected");
+            }
+        }
+
+        m_fail_target_proto = true;
         {
             scoped_log_handler err_wrapper(wrap_errors_logger);
             scoped_log_handler warn_wrapper(wrap_warns_logger);
@@ -920,6 +937,8 @@ class test_ucp_peer_failure_rndv_abort : public test_ucp_peer_failure {
 
     ucs::mock         m_progress_mock;
     bool              m_is_peer_closed{false};
+    bool              m_fail_target_proto{true};
+    bool              m_target_proto_seen{false};
     std::string       m_proto_name{};
     /* Protocol stage during which data transfer happens */
     uint8_t           m_proto_xfer_stage{};

test/gtest/ucp/test_ucp_proto.cc

@@ -400,6 +400,84 @@ class test_ucp_proto_am_rndv : public test_ucp_proto {
     }
 
 protected:
+    void dump_am_rndv_selection(const ucp_proto_select_elem_t *select_elem)
+    {
+        const ucp_proto_threshold_elem_t *thresh = select_elem->thresholds;
+        size_t min_msg_length                     = 0;
+
+        for (;;) {
+            const ucp_proto_config_t *proto_config = &thresh->proto_config;
+
+            UCS_TEST_MESSAGE << "PR11685_DIAG outer range=["
+                             << min_msg_length << ","
+                             << thresh->max_msg_length << "] proto="
+                             << proto_config->proto->name << " sys_dev="
+                             << static_cast<unsigned>(
+                                        proto_config->select_param.sys_dev)
+                             << " mem_type="
+                             << static_cast<unsigned>(
+                                        proto_config->select_param.mem_type)
+                             << " op_id_flags="
+                             << static_cast<unsigned>(
+                                        proto_config->select_param.op_id_flags);
+
+            if (std::strcmp(proto_config->proto->name, "am/rndv") == 0) {
+                const ucp_proto_rndv_ctrl_priv_t *rpriv =
+                        static_cast<const ucp_proto_rndv_ctrl_priv_t*>(
+                                proto_config->priv);
+                const ucp_proto_config_t *remote_proto_config =
+                        &rpriv->remote_proto_config;
+                const ucp_proto_select_param_t *remote_select_param =
+                        &remote_proto_config->select_param;
+
+                UCS_TEST_MESSAGE << "PR11685_DIAG remote outer_range=["
+                                 << min_msg_length << ","
+                                 << thresh->max_msg_length << "] sys_dev="
+                                 << static_cast<unsigned>(
+                                            remote_select_param->sys_dev)
+                                 << " mem_type="
+                                 << static_cast<unsigned>(
+                                            remote_select_param->mem_type)
+                                 << " mem_flags="
+                                 << static_cast<unsigned>(
+                                            remote_select_param->op.mem_flags)
+                                 << " op_id_flags="
+                                 << static_cast<unsigned>(
+                                            remote_select_param->op_id_flags)
+                                 << " rkey_cfg_index="
+                                 << remote_proto_config->rkey_cfg_index;
+
+                if (remote_proto_config->rkey_cfg_index !=
+                    UCP_WORKER_CFG_INDEX_NULL) {
+                    const ucp_rkey_config_key_t *rkey_key =
+                            &ucs_array_elem(
+                                     &worker()->rkey_config,
+                                     remote_proto_config->rkey_cfg_index).key;
+
+                    UCS_TEST_MESSAGE << "PR11685_DIAG rkey outer_range=["
+                                     << min_msg_length << ","
+                                     << thresh->max_msg_length << "] sys_dev="
+                                     << static_cast<unsigned>(rkey_key->sys_dev)
+                                     << " mem_type="
+                                     << static_cast<unsigned>(
+                                                rkey_key->mem_type)
+                                     << " flags="
+                                     << static_cast<unsigned>(rkey_key->flags)
+                                     << " md_map=" << rkey_key->md_map
+                                     << " unreachable_md_map="
+                                     << rkey_key->unreachable_md_map;
+                }
+            }
+
+            if (thresh->max_msg_length == SIZE_MAX) {
+                break;
+            }
+
+            min_msg_length = thresh->max_msg_length + 1;
+            ++thresh;
+        }
+    }
+
     const ucp_proto_config_t *select_am_rndv_remote_proto_config(
             uint8_t op_flags)
     {
@@ -428,6 +506,8 @@ class test_ucp_proto_am_rndv : public test_ucp_proto {
             return nullptr;
         }
 
+        dump_am_rndv_selection(select_elem);
+
         thresh = ucp_proto_thresholds_search_slow(select_elem->thresholds,
                                                   UCS_MBYTE);
         EXPECT_NE(nullptr, thresh);
@@ -778,7 +858,8 @@ UCS_TEST_P(test_ucp_proto_am_rndv,
 UCS_TEST_P(test_ucp_proto_am_rndv,
            am_explicit_scheme_disables_force_provenance,
            "RNDV_PIPELINE_SHM_CUDA_STAGING_FORCE=y",
-           "RNDV_SCHEME=get_zcopy")
+           "RNDV_SCHEME=get_zcopy",
+           "PROTO_INFO=y")
 {
     require_cuda_rndv_support();