UCP/FT: probe-gated lane recovery via aux uct_ep_check

src/ucp/core/ucp_ep.h

@@ -169,6 +169,7 @@ enum {
                                                            transports for AM lane */
     UCP_EP_INIT_ERR_MODE_FAILOVER      = UCS_BIT(11), /**< Endpoint requires an
                                                            @ref UCP_ERR_HANDLING_MODE_FAILOVER */
+    UCP_EP_INIT_RECOVERY               = UCS_BIT(12),
 
     /**
      * For consistency with @ref UCP_SA_DATA_MASK_ERR_MODE_FAILOVER
@@ -492,10 +493,27 @@ typedef struct {
 } ucp_ep_flush_state_t;
 
 
+typedef struct ucp_ep_recovery_probe {
+    uct_completion_t  comp;
+    ucp_ep_h          ep;
+    ucp_lane_index_t  lane;
+} ucp_ep_recovery_probe_t;
+
+
+enum {
+    UCP_EP_RECOVERY_STATE_IDLE,
+    UCP_EP_RECOVERY_STATE_WAIT_REPLY,
+    UCP_EP_RECOVERY_STATE_PROBING,
+    UCP_EP_RECOVERY_STATE_PROBE_OK
+};
+
+
 /* Per-EP recovery retry state. */
 typedef struct ucp_ep_recovery_arg {
     /* number of retries left before giving up */
-    unsigned    retries_left;
+    unsigned                retries_left;
+    uint8_t                 state;
+    ucp_ep_recovery_probe_t probe[UCP_MAX_LANES];
 } ucp_ep_recovery_arg_t;
 
 
@@ -598,6 +616,8 @@ typedef struct ucp_ep {
         /* How many UCT EP discarding operations are in-progress scheduled for
          * the EP */
         unsigned                      discard;
+        /* How many recovery aux probes are in-progress on the EP */
+        unsigned                      probe;
     } refcounts;
 #endif
 
@@ -1005,6 +1025,12 @@ ucs_status_t ucp_ep_reconfig_clear_failed_lanes(ucp_ep_h ep,
 ucs_status_t ucp_ep_recovery_arm(ucp_ep_h ep);
 
 
+/**
+ * Notify recovery progress that a lanes-address reply was received.
+ */
+void ucp_ep_recovery_on_reply_received(ucp_ep_h ep);
+
+
 /**
  * Progress function for failed lanes recovery.
  *

src/ucp/core/ucp_ep.inl

@@ -142,6 +142,15 @@ static inline int ucp_ep_is_lane_p2p(ucp_ep_h ep, ucp_lane_index_t lane)
     return !!(ucp_ep_config(ep)->p2p_lanes & UCS_BIT(lane));
 }
 
+/* NULL-safe check whether the lane currently holds a failed-stub UCT EP. */
+static UCS_F_ALWAYS_INLINE int
+ucp_ep_is_lane_failed_stub(ucp_ep_h ep, ucp_lane_index_t lane)
+{
+    uct_ep_h uct_ep = ucp_ep_get_lane(ep, lane);
+
+    return (uct_ep != NULL) && ucp_is_uct_ep_failed(uct_ep);
+}
+
 static inline ucp_md_index_t ucp_ep_md_index(ucp_ep_h ep, ucp_lane_index_t lane)
 {
     return ucp_ep_config(ep)->md_index[lane];

src/ucp/wireup/address.c

@@ -364,8 +364,11 @@ static int ucp_address_pack_v1_extra_info(ucp_object_version_t addr_version,
            (pack_flags & UCP_ADDRESS_PACK_FLAG_RELEASE_VER_V1);
 }
 
+/* Failed stubs cannot produce ep addresses. Check UCT EP state because a
+ * recovering lane keeps its FAILED bit while holding a packable wireup proxy. */
 static ucs_status_t
-ucp_address_gather_devices(ucp_worker_h worker, const ucp_ep_config_key_t *key,
+ucp_address_gather_devices(ucp_worker_h worker, ucp_ep_h ep,
+                           const ucp_ep_config_key_t *key,
                            const ucp_tl_bitmap_t *tl_bitmap, uint64_t flags,
                            ucp_object_version_t addr_version,
                            unsigned max_num_paths,
@@ -403,7 +406,8 @@ ucp_address_gather_devices(ucp_worker_h worker, const ucp_ep_config_key_t *key,
              */
             for (lane = 0; lane < key->num_lanes; ++lane) {
                 if ((key->lanes[lane].rsc_index == rsc_index) &&
-                    ucp_ep_config_connect_p2p(worker, key, rsc_index)) {
+                    ucp_ep_config_connect_p2p(worker, key, rsc_index) &&
+                    !((ep != NULL) && ucp_ep_is_lane_failed_stub(ep, lane))) {
                     dev->tl_addrs_size += !ucp_worker_is_unified_mode(worker);
                     dev->tl_addrs_size += iface_attr->ep_addr_len;
                     dev->tl_addrs_size += sizeof(uint8_t); /* lane index */
@@ -1386,7 +1390,8 @@ ucp_address_do_pack(ucp_worker_h worker, ucp_ep_h ep, void *buffer, size_t size,
 
                 ucs_for_each_bit(lane, ucp_ep_config(ep)->p2p_lanes) {
                     ucs_assert(lane < UCP_MAX_LANES);
-                    if (ucp_ep_get_rsc_index(ep, lane) != rsc_index) {
+                    if ((ucp_ep_get_rsc_index(ep, lane) != rsc_index) ||
+                        ucp_ep_is_lane_failed_stub(ep, lane)) {
                         continue;
                     }
 
@@ -1514,7 +1519,7 @@ ucp_address_length(ucp_worker_h worker, const ucp_ep_config_key_t *key,
     ssize_t size;
 
     /* Collect all devices required to pack their address */
-    status = ucp_address_gather_devices(worker, key, tl_bitmap, pack_flags,
+    status = ucp_address_gather_devices(worker, NULL, key, tl_bitmap, pack_flags,
                                         addr_version, UINT_MAX, &devices,
                                         &num_devices);
     if (status != UCS_OK) {
@@ -1559,7 +1564,7 @@ ucs_status_t ucp_address_pack(ucp_worker_h worker, ucp_ep_h ep,
     }
 
     /* Collect all devices we want to pack */
-    status = ucp_address_gather_devices(worker, key, tl_bitmap, pack_flags,
+    status = ucp_address_gather_devices(worker, ep, key, tl_bitmap, pack_flags,
                                         addr_version, max_num_paths, &devices,
                                         &num_devices);
     if (status != UCS_OK) {

src/ucp/wireup/select.c

@@ -1162,6 +1162,22 @@ ucp_wireup_aux_seg_size_score_func(const ucp_worker_iface_t *wiface,
     return ucp_wireup_aux_seg_size(&wiface->attr, remote_addr);
 }
 
+static uint64_t
+ucp_wireup_aux_local_mandatory_flags(unsigned ep_init_flags)
+{
+    uint64_t flags = 0;
+
+    if (!ucp_ep_init_flags_has_cm(ep_init_flags)) {
+        flags |= UCT_IFACE_FLAG_CONNECT_TO_IFACE;
+    }
+
+    if (ep_init_flags & UCP_EP_INIT_RECOVERY) {
+        flags |= UCT_IFACE_FLAG_EP_CHECK;
+    }
+
+    return flags;
+}
+
 static void ucp_wireup_fill_aux_criteria(ucp_wireup_criteria_t *criteria,
                                          unsigned ep_init_flags,
                                          uint64_t mandatory_flags)
@@ -1174,13 +1190,15 @@ static void ucp_wireup_fill_aux_criteria(ucp_wireup_criteria_t *criteria,
     ucp_wireup_init_select_flags(&criteria->remote_iface_flags,
                                  UCP_ADDR_IFACE_FLAG_AM_SYNC, 0);
 
+    criteria->local_iface_flags.mandatory |=
+        ucp_wireup_aux_local_mandatory_flags(ep_init_flags);
+
     /* CM lane doesn't require to use CONNECT_TO_IFACE for auxiliary lane */
     if (!ucp_ep_init_flags_has_cm(ep_init_flags)) {
-        criteria->local_iface_flags.mandatory  |=
-                UCT_IFACE_FLAG_CONNECT_TO_IFACE;
         criteria->remote_iface_flags.mandatory |=
                 UCP_ADDR_IFACE_FLAG_CONNECT_TO_IFACE | mandatory_flags;
     }
+
     criteria->local_cmpt_flags   = 0;
     criteria->local_event_flags  = 0;
     criteria->remote_event_flags = 0;
@@ -3043,15 +3061,16 @@ static ucs_status_t
 ucp_wireup_select_aux_transport_by_seg_size(
         const ucp_wireup_select_context_t *select_ctx,
         const ucp_wireup_select_params_t *select_params,
-        const ucp_wireup_criteria_t *criteria, int show_error,
+        const ucp_wireup_criteria_t *criteria, uint64_t local_dev_bitmap,
+        uint64_t remote_dev_bitmap, int show_error,
         ucp_wireup_select_info_t *select_info)
 {
     ucs_status_t status;
 
     status = ucp_wireup_select_transport(select_ctx, select_params, criteria,
                                          ucp_tl_bitmap_max, UINT64_MAX,
-                                         UINT64_MAX, UINT64_MAX, show_error,
-                                         select_info);
+                                         local_dev_bitmap, remote_dev_bitmap,
+                                         show_error, select_info);
     if (status != UCS_OK) {
         return status;
     }
@@ -3070,6 +3089,8 @@ ucs_status_t
 ucp_wireup_select_aux_transport(ucp_ep_h ep, unsigned ep_init_flags,
                                 ucp_tl_bitmap_t tl_bitmap,
                                 const ucp_unpacked_address_t *remote_address,
+                                uint64_t local_dev_bitmap,
+                                uint64_t remote_dev_bitmap,
                                 ucp_wireup_select_info_t *select_info)
 {
     ucp_wireup_select_context_t select_ctx = {};
@@ -3085,7 +3106,9 @@ ucp_wireup_select_aux_transport(ucp_ep_h ep, unsigned ep_init_flags,
                                  UCP_ADDR_IFACE_FLAG_CB_ASYNC);
     status = ucp_wireup_select_aux_transport_by_seg_size(&select_ctx,
                                                          &select_params,
-                                                         &criteria, 0,
+                                                         &criteria,
+                                                         local_dev_bitmap,
+                                                         remote_dev_bitmap, 0,
                                                          select_info);
     if (status == UCS_OK) {
         return UCS_OK;
@@ -3096,6 +3119,8 @@ ucp_wireup_select_aux_transport(ucp_ep_h ep, unsigned ep_init_flags,
     ucp_wireup_fill_aux_criteria(&criteria, ep_init_flags, 0);
     return ucp_wireup_select_aux_transport_by_seg_size(&select_ctx,
                                                        &select_params,
-                                                       &criteria, 1,
+                                                       &criteria,
+                                                       local_dev_bitmap,
+                                                       remote_dev_bitmap, 1,
                                                        select_info);
 }

src/ucp/wireup/wireup.c

@@ -981,6 +981,42 @@ void ucp_wireup_process_ack(ucp_worker_h worker, ucp_ep_h ep,
     ucp_wireup_remote_connected(ep);
 }
 
+static void
+ucp_wireup_augment_aux_tls(ucp_ep_h ep, ucp_lane_map_t lane_map,
+                           ucp_tl_bitmap_t *tl_bitmap)
+{
+    ucp_context_h context = ep->worker->context;
+    ucp_rsc_index_t lane_rsc, aux_rsc;
+    ucp_lane_index_t lane;
+    uint64_t iface_flags;
+
+    ucs_for_each_bit(lane, lane_map) {
+        lane_rsc = ucp_ep_get_rsc_index(ep, lane);
+        if (lane_rsc == UCP_NULL_RESOURCE) {
+            continue;
+        }
+
+        for (aux_rsc = 0; aux_rsc < context->num_tls; ++aux_rsc) {
+            if (UCS_STATIC_BITMAP_GET(*tl_bitmap, aux_rsc)) {
+                continue;
+            }
+
+            if (context->tl_rscs[aux_rsc].dev_index !=
+                context->tl_rscs[lane_rsc].dev_index) {
+                continue;
+            }
+
+            iface_flags = ucp_worker_iface_get_attr(ep->worker,
+                                                    aux_rsc)->cap.flags;
+            if (ucs_test_all_flags(iface_flags,
+                                   UCT_IFACE_FLAG_CONNECT_TO_IFACE |
+                                   UCT_IFACE_FLAG_EP_CHECK)) {
+                UCS_STATIC_BITMAP_SET(tl_bitmap, aux_rsc);
+            }
+        }
+    }
+}
+
 void ucp_wireup_send_lanes_addr_msg(ucp_ep_h ep, uint8_t msg_type,
                                     ucp_lane_map_t requested_lane_map,
                                     ucp_lane_map_t provided_lane_map)
@@ -989,6 +1025,7 @@ void ucp_wireup_send_lanes_addr_msg(ucp_ep_h ep, uint8_t msg_type,
     ucs_status_t status;
 
     tl_bitmap = ucp_wireup_get_ep_tl_bitmap(ep, provided_lane_map);
+    ucp_wireup_augment_aux_tls(ep, provided_lane_map, &tl_bitmap);
 
     ucs_debug("ep %p: send %s requested=0x%" PRIx64 " provided=0x%" PRIx64, ep,
               ucp_wireup_msg_str(msg_type), (uint64_t)requested_lane_map,
@@ -1066,6 +1103,7 @@ ucp_wireup_process_lanes_addr_reply(
     rebuilt = ucp_ep_recovery_rebuild_lanes(
             ep, lanes_info->provided_lane_map & ucp_ep_get_failed_lanes(ep),
             remote_address);
+    ucp_ep_recovery_on_reply_received(ep);
 
     ucs_debug("ep %p: LANES_ADDR_REP requested=0x%" PRIx64
               " provided=0x%" PRIx64 " rebuilt=0x%" PRIx64,

src/ucp/wireup/wireup.h

@@ -178,6 +178,8 @@ ucs_status_t
 ucp_wireup_select_aux_transport(ucp_ep_h ep, unsigned ep_init_flags,
                                 ucp_tl_bitmap_t tl_bitmap,
                                 const ucp_unpacked_address_t *remote_address,
+                                uint64_t local_dev_bitmap,
+                                uint64_t remote_dev_bitmap,
                                 ucp_wireup_select_info_t *select_info);
 
 double ucp_wireup_amo_score_func(const ucp_worker_iface_t *wiface,

src/ucp/wireup/wireup_ep.c

@@ -277,6 +277,7 @@ ucp_wireup_ep_connect_aux(ucp_wireup_ep_t *wireup_ep, unsigned ep_init_flags,
      */
     status = ucp_wireup_select_aux_transport(ucp_ep, ep_init_flags,
                                              ucp_tl_bitmap_max, remote_address,
+                                             UINT64_MAX, UINT64_MAX,
                                              &select_info);
     if (status != UCS_OK) {
         return status;
@@ -335,6 +336,30 @@ void ucp_wireup_ep_discard_aux_ep(ucp_wireup_ep_t *wireup_ep,
     }
 }
 
+void ucp_wireup_ep_destroy_aux_ep(ucp_wireup_ep_t *wireup_ep)
+{
+    ucp_ep_h ucp_ep     = wireup_ep->super.ucp_ep;
+    ucp_worker_h worker = ucp_ep->worker;
+    uct_ep_h aux_ep     = wireup_ep->aux_ep;
+    ucp_rsc_index_t rsc_index;
+
+    if (aux_ep == NULL) {
+        return;
+    }
+
+    rsc_index                = wireup_ep->aux_rsc_index;
+    wireup_ep->aux_ep        = NULL;
+    wireup_ep->aux_rsc_index = UCP_NULL_RESOURCE;
+    wireup_ep->flags        &= ~UCP_WIREUP_EP_FLAG_AUX_P2P;
+
+    ucp_ep_unprogress_uct_ep(ucp_ep, aux_ep, rsc_index);
+    if (worker->context->config.ext.proto_enable) {
+        ucp_worker_iface_unprogress_ep(ucp_worker_iface(worker, rsc_index));
+    }
+
+    uct_ep_destroy(aux_ep);
+}
+
 static ucs_status_t ucp_wireup_ep_flush(uct_ep_h uct_ep, unsigned flags,
                                         uct_completion_t *comp)
 {

src/ucp/wireup/wireup_ep.h

@@ -104,6 +104,8 @@ void ucp_wireup_ep_discard_aux_ep(ucp_wireup_ep_t *wireup_ep,
                                   uct_pending_purge_callback_t purge_cb,
                                   void *purge_arg);
 
+void ucp_wireup_ep_destroy_aux_ep(ucp_wireup_ep_t *wireup_ep);
+
 int ucp_wireup_ep_has_next_ep(ucp_wireup_ep_t *wireup_ep);
 
 void ucp_wireup_ep_set_next_ep(uct_ep_h uct_ep, uct_ep_h next_ep,