UCP: Exclude device atomics from failover endpoints.

src/ucp/core/ucp_context.c

@@ -482,11 +482,13 @@ static ucs_config_field_t ucp_context_config_table[] = {
    "Failover method for endpoints created with\n"
    "UCP_ERR_HANDLING_MODE_FAILOVER.\n"
    " auto  - auto selected recovery protocol. Transport selection is\n"
-   "         unchanged.\n"
+   "         unchanged. Selects software PSN when UCX_IB_INDIRECT_ATOMIC\n"
+   "         is enabled.\n"
    " sw    - force software PSN based recovery protocol. Transport\n"
    "         selection is unchanged.\n"
    " token - force token based retransmits. Transports that do not support\n"
-   "         the required capabilities are disabled.",
+   "         the required capabilities are disabled. Rejected when\n"
+   "         UCX_IB_INDIRECT_ATOMIC is enabled.",
    ucs_offsetof(ucp_context_config_t, failover_mode),
    UCS_CONFIG_TYPE_ENUM(ucp_failover_modes)},
 
@@ -2213,6 +2215,67 @@ static void ucp_fill_resources_reg_md_map_update(ucp_context_h context)
     }
 }
 
+/* Effective UCX_IB_INDIRECT_ATOMIC on a selected IB memory domain. IB MD
+ * default and environment, then cached ucp_config_modify() overrides.
+ * Returns 0 when no IB domain was selected. */
+static int ucp_context_query_ib_indirect_atomic(ucp_context_h context)
+{
+    int enabled = 1;
+    ucp_md_index_t md_index;
+    ucp_rsc_index_t cmpt_index;
+    uct_md_config_t *md_config;
+    ucs_status_t status;
+    char value[8];
+
+    for (md_index = 0; md_index < context->num_mds; ++md_index) {
+        cmpt_index = context->tl_mds[md_index].cmpt_index;
+        if (strcmp(context->tl_cmpts[cmpt_index].attr.name, "ib") != 0) {
+            continue;
+        }
+
+        status = uct_md_config_read(context->tl_cmpts[cmpt_index].cmpt,
+                                    NULL, NULL, &md_config);
+        if (status != UCS_OK) {
+            return 1;
+        }
+
+        ucp_apply_uct_config_list(context, md_config);
+        status = uct_config_get(md_config, "INDIRECT_ATOMIC", value,
+                                sizeof(value));
+        if ((status != UCS_OK) ||
+            (ucs_config_sscanf_bool(value, &enabled, NULL) == 0)) {
+            enabled = 1;
+        }
+
+        uct_config_release(md_config);
+        return enabled;
+    }
+
+    return 0;
+}
+
+/* Indirect atomics have no token replay. */
+static ucs_status_t ucp_context_adjust_failover_mode(ucp_context_h context)
+{
+    int enabled = ucp_context_query_ib_indirect_atomic(context);
+
+    if (!enabled) {
+        return UCS_OK;
+    }
+
+    if (context->config.ext.failover_mode == UCP_FAILOVER_MODE_TOKEN) {
+        ucs_error("UCX_FAILOVER_MODE=token requires "
+                  "UCX_IB_INDIRECT_ATOMIC=n");
+        return UCS_ERR_INVALID_PARAM;
+    } else if (context->config.ext.failover_mode == UCP_FAILOVER_MODE_AUTO) {
+        ucs_debug("ib indirect atomic is enabled, failover mode auto "
+                  "selects software psn recovery");
+        context->config.ext.failover_mode = UCP_FAILOVER_MODE_SW;
+    }
+
+    return UCS_OK;
+}
+
 static ucs_status_t
 ucp_fill_resources(ucp_context_h context, const ucp_config_t *config)
 {
@@ -2368,6 +2431,11 @@ ucp_fill_resources(ucp_context_h context, const ucp_config_t *config)
 
     ucp_context_log_tl_info(context, &all_rscs);
 
+    status = ucp_context_adjust_failover_mode(context);
+    if (status != UCS_OK) {
+        goto err_free_resources;
+    }
+
 out_release_components:
     ucs_array_cleanup_dynamic(&all_rscs);
     uct_release_component_list(uct_components);
@@ -2733,13 +2801,6 @@ static ucs_status_t ucp_fill_config(ucp_context_h context,
         context->config.worker_fence_mode = context->config.ext.fence_mode;
     }
 
-    /* Token based recovery is not implemented yet, so automatic selection
-     * uses the software PSN recovery protocol. */
-    if (context->config.ext.failover_mode == UCP_FAILOVER_MODE_AUTO) {
-        context->config.ext.failover_mode = UCP_FAILOVER_MODE_SW;
-        ucs_debug("failover mode auto selects software psn recovery");
-    }
-
     context->config.progress_wrapper_enabled =
             ucs_log_is_enabled(UCS_LOG_LEVEL_TRACE_REQ) ||
             ucp_context_usage_tracker_enabled(context);

src/ucp/rma/amo_sw.c

@@ -239,10 +239,14 @@ UCS_PROFILE_FUNC(ucs_status_t, ucp_atomic_req_handler, (arg, data, length, am_fl
                       UCT_IFACE_FLAG_ATOMIC_DEVICE))) {
         ucs_error("Unsupported: got software atomic request while device atomics are selected on worker %p",
                   worker);
-        /* TODO: this situation will be possible then CM wireup is implemented
-         *       and CM lane is bound to suboptimal device, then need to execute
-         *       AMO on fastest resource from worker->atomic_tls using loopback
-         *       EP and continue SW AMO protocol */
+        /* TODO: This situation may occur when:
+         *       - The CM lane is bound to a suboptimal device
+         *       - An EP with error handling mode failover is used alongside
+         *         other EPs
+         *       In this case, the AMO should be executed on the fastest
+         *       resource from worker->atomic_tls using the loopback EP, and
+         *       the SW AMO protocol should continue.
+         */
     }
 
     if (atomicreqh->req.req_id == UCS_PTR_MAP_KEY_INVALID) {
@@ -430,7 +434,8 @@ static void ucp_proto_amo_sw_probe(const ucp_proto_init_params_t *init_params,
         .super.send_op       = UCT_EP_OP_AM_BCOPY,
         .super.memtype_op    = UCT_EP_OP_GET_SHORT,
         .super.flags         = flags | UCP_PROTO_COMMON_INIT_FLAG_SINGLE_FRAG |
-                               UCP_PROTO_COMMON_INIT_FLAG_CAP_SEG_SIZE,
+                               UCP_PROTO_COMMON_INIT_FLAG_CAP_SEG_SIZE |
+                               UCP_PROTO_COMMON_INIT_FLAG_FAILOVER,
         .super.exclude_map   = ~UCP_MAX_FAST_PATH_LANES_MASK,
         .super.reg_mem_info  = ucp_mem_info_unknown,
         .lane_type           = UCP_LANE_TYPE_AM,
@@ -478,7 +483,7 @@ ucp_proto_t ucp_get_amo_post_proto = {
     .probe    = ucp_proto_amo_sw_post_probe,
     .query    = ucp_proto_single_query,
     .progress = {ucp_proto_amo_sw_progress_post},
-    .abort    = ucp_proto_abort_fatal_not_implemented,
+    .abort    = ucp_proto_request_bcopy_abort,
     .reset    = ucp_proto_request_bcopy_reset
 };
 
@@ -507,6 +512,6 @@ ucp_proto_t ucp_get_amo_fetch_proto = {
     .probe    = ucp_proto_amo_sw_fetch_probe,
     .query    = ucp_proto_single_query,
     .progress = {ucp_proto_amo_sw_progress_fetch},
-    .abort    = ucp_proto_abort_fatal_not_implemented,
+    .abort    = ucp_proto_request_bcopy_id_abort,
     .reset    = ucp_proto_request_bcopy_id_reset
 };

src/ucp/wireup/select.c

@@ -1332,6 +1332,27 @@ double ucp_wireup_amo_score_func(const ucp_worker_iface_t *wiface,
             wiface->attr.overhead);
 }
 
+/* Failover has no replay path for NIC hardware atomics. */
+static void
+ucp_wireup_exclude_device_amo_tls(ucp_worker_h worker,
+                                  ucp_tl_bitmap_t *tl_bitmap)
+{
+    ucp_rsc_index_t rsc_index;
+
+    UCP_TL_BITMAP_FOR_EACH_BIT(rsc_index, tl_bitmap) {
+        if (!(ucp_worker_iface_get_attr(worker, rsc_index)->cap.flags &
+              UCT_IFACE_FLAG_ATOMIC_DEVICE)) {
+            continue;
+        }
+
+        ucs_trace(UCT_TL_RESOURCE_DESC_FMT
+                  " : disabled, failover does not use device atomics",
+                  UCT_TL_RESOURCE_DESC_ARG(
+                          &worker->context->tl_rscs[rsc_index].tl_rsc));
+        UCS_STATIC_BITMAP_RESET(tl_bitmap, rsc_index);
+    }
+}
+
 static ucs_status_t
 ucp_wireup_add_amo_lanes(const ucp_wireup_select_params_t *select_params,
                          ucp_wireup_select_context_t *select_ctx)
@@ -1371,6 +1392,10 @@ ucp_wireup_add_amo_lanes(const ucp_wireup_select_params_t *select_params,
         }
     }
 
+    if (ep_init_flags & UCP_EP_INIT_ERR_MODE_FAILOVER) {
+        ucp_wireup_exclude_device_amo_tls(worker, &tl_bitmap);
+    }
+
     return ucp_wireup_add_memaccess_lanes(select_params, ep_init_flags,
                                           &criteria, UCS_MEMORY_TYPE_HOST,
                                           tl_bitmap, UCP_LANE_TYPE_AMO,

test/gtest/ucp/test_ucp_fault_tolerance.cc

@@ -932,6 +932,8 @@ class DISABLED_test_ucp_fault_tolerance_tf : public test_ucp_fault_tolerance {
 public:
     DISABLED_test_ucp_fault_tolerance_tf() {
         modify_config("FAILOVER_MODE", "token");
+        /* IB MD key; set UCX_IB_INDIRECT_ATOMIC so token mode stays selected. */
+        modify_config("IB_INDIRECT_ATOMIC", "n", SETENV_IF_NOT_EXIST);
     }
 };
 

test/gtest/ucp/test_ucp_wireup.cc

@@ -88,6 +88,8 @@ class test_ucp_wireup : public ucp_test {
     bool ep_iface_has_caps(const entity& e, const std::string& tl,
                            uint64_t caps);
 
+    static void unmap_memh(ucp_mem_h memh, ucp_context_h context);
+
 protected:
     vec_type                               m_send_data;
     vec_type                               m_recv_data;
@@ -98,8 +100,6 @@ class test_ucp_wireup : public ucp_test {
 private:
     static void stream_recv_completion(void *request, ucs_status_t status,
                                        size_t length);
-
-    static void unmap_memh(ucp_mem_h memh, ucp_context_h context);
 };
 
 void test_ucp_wireup::get_test_variants(std::vector<ucp_test_variant>& variants,
@@ -1615,6 +1615,75 @@ UCS_TEST_P(test_ucp_wireup_amo, relese_key_after_flush) {
 
 UCP_INSTANTIATE_TEST_CASE(test_ucp_wireup_amo)
 
+static void test_ucp_wireup_failover_err_cb(void *arg, ucp_ep_h ep,
+                                            ucs_status_t status)
+{
+}
+
+UCS_TEST_P(test_ucp_wireup_fallback_amo, failover_no_device_amo)
+{
+    ucp_ep_params_t params = get_ep_params();
+    ucp_mem_map_params_t map_params;
+    ucp_request_param_t req_param;
+    ucp_mem_h memh;
+    ucp_rkey_h rkey;
+    void *req;
+    void *flush_req;
+    ucs_status_t status;
+
+    params.field_mask     |= UCP_EP_PARAM_FIELD_ERR_HANDLING_MODE |
+                             UCP_EP_PARAM_FIELD_ERR_HANDLER;
+    params.err_mode        = UCP_ERR_HANDLING_MODE_FAILOVER;
+    params.err_handler.cb  = test_ucp_wireup_failover_err_cb;
+    params.err_handler.arg = NULL;
+
+    create_entity(false);
+    create_entity(true);
+    sender().connect(&receiver(), params);
+    EXPECT_FALSE(use_device_amo(sender().ep()));
+
+    m_send_data[0] = 1;
+    clear_recv_data();
+
+    map_params.field_mask = UCP_MEM_MAP_PARAM_FIELD_ADDRESS |
+                            UCP_MEM_MAP_PARAM_FIELD_LENGTH |
+                            UCP_MEM_MAP_PARAM_FIELD_FLAGS;
+    map_params.address    = &m_recv_data[0];
+    map_params.length     = m_recv_data.size() * sizeof(m_recv_data[0]);
+    map_params.flags      = 0;
+    status                = ucp_mem_map(receiver().ucph(), &map_params, &memh);
+    ASSERT_UCS_OK(status);
+    m_memh_receiver.reset(memh, unmap_memh, receiver().ucph());
+
+    rkey = get_rkey(sender().ep(), m_memh_receiver);
+    m_rkeys.push_back(ucs::handle<ucp_rkey_h>(rkey, ucp_rkey_destroy));
+
+    {
+        scoped_log_handler slh(wrap_errors_logger);
+
+        req_param.op_attr_mask = UCP_OP_ATTR_FIELD_DATATYPE;
+        req_param.datatype     = ucp_dt_make_contig(sizeof(m_send_data[0]));
+        req                    = ucp_atomic_op_nbx(sender().ep(),
+                                                   UCP_ATOMIC_OP_ADD,
+                                                   &m_send_data[0], 1,
+                                                   (uint64_t)&m_recv_data[0],
+                                                   rkey, &req_param);
+        status                 = request_wait(req);
+        if (status == UCS_OK) {
+            /* Post completes locally; flush waits for the remote atomic. */
+            flush_req = sender().flush_ep_nb();
+            status    = request_wait(flush_req);
+        }
+    }
+
+    if (status == UCS_OK) {
+        EXPECT_EQ(m_send_data[0], m_recv_data[0]);
+    }
+
+    test_ucp_wireup::cleanup();
+    EXPECT_EQ(UCS_OK, status);
+}
+
 UCS_TEST_P(test_ucp_wireup_fallback_amo, different_amo_types) {
     std::vector<std::string> tls;