UCC/CORE: Added automation to loading local rank from topo if not provided by user

Author: MaayanGadishNvidia

src/components/topo/ucc_topo.h

@@ -290,6 +290,21 @@ static inline ucc_rank_t ucc_topo_nnodes(ucc_topo_t *topo)
     return sbgp->group_size;
 }
 
+static inline ucc_rank_t ucc_topo_node_local_rank(ucc_topo_t *topo)
+{
+    ucc_sbgp_t *sbgp;
+
+    if (!topo) {
+        return UCC_RANK_INVALID;
+    }
+    sbgp = ucc_topo_get_sbgp(topo, UCC_SBGP_NODE);
+
+    if (sbgp->status != UCC_SBGP_ENABLED) {
+        return UCC_RANK_INVALID;
+    }
+    return sbgp->group_rank;
+}
+
 /* Returns node leaders array - array that maps each rank to the TEAM RANK that
    is the leader of that rank's node. Also returns per-node leaders array - array
    mapping node_id to the TEAM RANK of that node's leader */

src/core/ucc_context.c

@@ -6,6 +6,7 @@
 
 #include "config.h"
 #include "ucc_context.h"
+#include "components/topo/ucc_topo.h"
 #include "utils/ucc_proc_info.h"
 #include "components/cl/ucc_cl.h"
 #include "components/tl/ucc_tl.h"
@@ -60,7 +61,12 @@ static ucc_config_field_t ucc_context_config_table[] = {
      ucc_offsetof(ucc_context_config_t, net_devices), UCC_CONFIG_TYPE_STRING_ARRAY},
 
     {"NODE_LOCAL_ID", "auto",
-     "An optimization hint for the local identificator on a single node.",
+     "An optimization hint for the local rank index on a single node.\n"
+     "With OOB and more than one endpoint, \"auto\" runs one prefix allgather "
+     "(ctx id / host), then the main address exchange (two allgathers: "
+     "lengths, then full addresses). Custom OOB backends must support three "
+     "sequential allgather operations on the same coll_info (each completes "
+     "before the next begins).",
      ucc_offsetof(ucc_context_config_t, node_local_id), UCC_CONFIG_TYPE_ULUNITS},
 
     {NULL}};
@@ -599,6 +605,98 @@ ucc_status_t ucc_core_addr_exchange(
     return UCC_OK;
 }
 
+ucc_status_t ucc_core_ctx_id_exchange(ucc_context_t *context, ucc_oob_coll_t *oob,
+                                      ucc_addr_storage_t *addr_storage)
+{
+    /* Caller must pass addr_storage cleared for a new exchange: at entry,
+       addr_len == 0 iff storage == NULL. */
+    ucc_status_t      status;
+    ucc_rank_t        i;
+    const size_t      ctx_addr_prefix_len =
+        ucc_offsetof(ucc_context_addr_header_t, n_components);
+    ucc_context_addr_header_t *h;
+    void              *scratch;
+
+poll:
+    if (addr_storage->oob_req) {
+        status = oob->req_test(addr_storage->oob_req);
+        if (status < 0) {
+            oob->req_free(addr_storage->oob_req);
+            addr_storage->oob_req = NULL;
+            ucc_free(addr_storage->storage);
+            addr_storage->storage = NULL;
+            addr_storage->addr_len = 0;
+            ucc_error("oob req test failed during ctx id prefix exchange");
+            return status;
+        } else if (UCC_INPROGRESS == status) {
+            return status;
+        }
+        oob->req_free(addr_storage->oob_req);
+        addr_storage->oob_req = NULL;
+    }
+    /* If addr_len is still 0, we have not started the prefix allgather yet.
+     * Once it is started, addr_len is set and we goto poll until completion;
+     * we never reach the rank loop below with addr_len == 0. */
+    if (0 == addr_storage->addr_len) {
+        /* addr_len == 0 iff storage == NULL at entry / before prefix allgather. */
+        ucc_assert(NULL == addr_storage->storage);
+        addr_storage->size = oob->n_oob_eps;
+
+        addr_storage->storage = ucc_malloc(
+            (addr_storage->size + 1) * ctx_addr_prefix_len, "ctx_ids_storage");
+        if (!addr_storage->storage) {
+            ucc_error(
+                "failed to allocate %zd bytes for ctx_ids storage",
+                (addr_storage->size + 1) * ctx_addr_prefix_len);
+            return UCC_ERR_NO_MEMORY;
+        }
+        addr_storage->addr_len = ctx_addr_prefix_len;
+
+        scratch = PTR_OFFSET(addr_storage->storage,
+                             ctx_addr_prefix_len * addr_storage->size);
+        memset(scratch, 0, ctx_addr_prefix_len);
+        h                         = (ucc_context_addr_header_t *)scratch;
+        h->ctx_id                 = context->id;
+        h->host_info              = ucc_local_host;
+
+        status = oob->allgather(scratch, addr_storage->storage,
+                                ctx_addr_prefix_len, oob->coll_info,
+                                &addr_storage->oob_req);
+        if (UCC_OK != status) {
+            ucc_free(addr_storage->storage);
+            addr_storage->storage = NULL;
+            addr_storage->addr_len = 0;
+            ucc_error("failed to start oob allgather for ctx_ids");
+            return status;
+        }
+        goto poll;
+    }
+    /* addr_len is the per-rank prefix stride; set with storage when the prefix
+       allgather is started (see allocation block above). */
+    {
+        ucc_rank_t r = UCC_RANK_MAX;
+
+        for (i = 0; i < addr_storage->size; i++) {
+            h = (ucc_context_addr_header_t *)PTR_OFFSET(
+                addr_storage->storage, addr_storage->addr_len * i);
+            if (UCC_CTX_ID_EQUAL(context->id, h->ctx_id)) {
+                if (r != UCC_RANK_MAX) {
+                    ucc_error("ctx_id collision: %d %d", r, i);
+                    ucc_free(addr_storage->storage);
+                    addr_storage->storage = NULL;
+                    addr_storage->addr_len = 0;
+                    return UCC_ERR_NO_MESSAGE;
+                }
+                r = i;
+            }
+        }
+
+        addr_storage->flags = 0;
+        addr_storage->rank = r;
+    }
+    return UCC_OK;
+}
+
 static void remove_tl_ctx_from_array(ucc_tl_context_t **array, unsigned *size,
                                      ucc_tl_context_t *tl_ctx)
 {
@@ -669,6 +767,102 @@ ucc_status_t ucc_context_create_proc_info(
         ucc_check_wait_for_debugger(ctx->rank);
 #endif
     }
+
+    ctx->id.pi      = *proc_info;
+    ctx->id.seq_num = ucc_atomic_fadd32(&ucc_context_seq_num, 1);
+
+    if (config->node_local_id == UCC_ULUNITS_AUTO) {
+        if ((params->mask & UCC_CONTEXT_PARAM_FIELD_OOB) &&
+            params->oob.n_oob_eps == 1) {
+            /* Single-rank OOB: no allgather; local rank on the node is 0. */
+            b_params.node_local_id = 0;
+        } else if ((params->mask & UCC_CONTEXT_PARAM_FIELD_OOB) &&
+                   params->oob.n_oob_eps > 1) {
+            status = ucc_sysinfo_get_host_info(&ucc_local_host);
+            if (UCC_OK != status) {
+                ucc_error("failed to refresh host info for ctx topo exchange");
+                goto error_ctx_create;
+            }
+            do {
+                /* UCC context create is blocking fn, so we can wait here for the
+                completion of addr exchange */
+                status = ucc_core_ctx_id_exchange(ctx, &ctx->params.oob,
+                                                &ctx->addr_storage);
+                if (status < 0) {
+                    ucc_error("failed to exchange addresses during context "
+                            "creation with status: %s",
+                            ucc_status_string(status));
+                    goto error_ctx_create;
+                }
+            } while (status == UCC_INPROGRESS);
+            /* Validate rank before ucc_context_topo_init allocates ctx->topo. */
+            if (ctx->addr_storage.rank == UCC_RANK_MAX) {
+                ucc_error("ctx_id not found in exchanged ids, rank detection "
+                          "failed");
+                status = UCC_ERR_NO_MESSAGE;
+                ucc_free(ctx->addr_storage.storage);
+                ctx->addr_storage.storage = NULL;
+                ctx->addr_storage.addr_len = 0;
+                goto error_ctx_create;
+            }
+            if (ctx->addr_storage.rank != params->oob.oob_ep) {
+                ucc_error("ctx_id exchange rank %d does not match OOB ep %d",
+                          ctx->addr_storage.rank, params->oob.oob_ep);
+                status = UCC_ERR_NO_MESSAGE;
+                ucc_free(ctx->addr_storage.storage);
+                ctx->addr_storage.storage = NULL;
+                ctx->addr_storage.addr_len = 0;
+                goto error_ctx_create;
+            }
+            status = ucc_context_topo_init(&ctx->addr_storage, &ctx->topo);
+            if (UCC_OK != status) {
+                ucc_free(ctx->addr_storage.storage);
+                ctx->addr_storage.storage = NULL;
+                ctx->addr_storage.addr_len = 0;
+                ucc_error("failed to init ctx topo");
+                goto error_ctx_create;
+            }
+            ucc_assert(ctx->topo != NULL);
+
+            {
+                ucc_subset_t set;
+                ucc_topo_t  *topo = NULL;
+
+                memset(&set.map, 0, sizeof(ucc_ep_map_t));
+                set.map.type   = UCC_EP_MAP_FULL;
+                set.myrank     = params->oob.oob_ep;
+                set.map.ep_num = params->oob.n_oob_eps;
+
+                status = ucc_topo_init(set, ctx->topo, &topo);
+                if (UCC_OK != status) {
+                    ucc_error("failed to init topo for computing local rank");
+                    ucc_free(ctx->addr_storage.storage);
+                    ctx->addr_storage.storage = NULL;
+                    ctx->addr_storage.addr_len = 0;
+                    goto error_ctx_create;
+                }
+                b_params.node_local_id = ucc_topo_node_local_rank(topo);
+                ucc_topo_cleanup(topo);
+            }
+
+            /* clean up addr_storage */
+            ucc_free(ctx->addr_storage.storage);
+            ctx->addr_storage.storage = NULL;
+            ctx->addr_storage.addr_len = 0;
+            ctx->addr_storage.size = 0;
+            ctx->addr_storage.rank = UCC_RANK_MAX;
+            ctx->addr_storage.flags = 0;
+            ctx->addr_storage.oob_req = NULL;
+
+            /* clean up topo */
+            ucc_context_topo_cleanup(ctx->topo);
+            ctx->topo = NULL;
+        } else if (!(params->mask & UCC_CONTEXT_PARAM_FIELD_OOB)) {
+            /* No OOB: cannot run topo exchange; sole implicit rank on the node. */
+            b_params.node_local_id = 0;
+        }
+    }
+
     status = ucc_create_tl_contexts(ctx, config, b_params);
     if (UCC_OK != status) {
         /* only critical error could have happened - bail */
@@ -748,8 +942,7 @@ ucc_status_t ucc_context_create_proc_info(
         ucc_error("failed to init progress queue for context %p", ctx);
         goto error_ctx_create;
     }
-    ctx->id.pi      = *proc_info;
-    ctx->id.seq_num = ucc_atomic_fadd32(&ucc_context_seq_num, 1);
+
     if (params->mask & UCC_CONTEXT_PARAM_FIELD_OOB &&
         params->oob.n_oob_eps > 1) {
         do {
@@ -765,15 +958,23 @@ ucc_status_t ucc_context_create_proc_info(
             }
         } while (status == UCC_INPROGRESS);
 
-        if (topo_required) {
+        if (topo_required && !ctx->topo) {
             /* At least one available CL context reported it needs topo info */
             status = ucc_context_topo_init(&ctx->addr_storage, &ctx->topo);
             if (UCC_OK != status) {
                 ucc_error("failed to init ctx topo");
                 ucc_free(ctx->addr_storage.storage);
+                ctx->addr_storage.storage = NULL;
+                ctx->addr_storage.addr_len = 0;
                 goto error_ctx_create;
             }
         }
+        if (ctx->addr_storage.rank == UCC_RANK_MAX) {
+            ucc_error("ctx_id not found after address exchange, rank detection "
+                      "failed");
+            status = UCC_ERR_NO_MESSAGE;
+            goto error_ctx_create;
+        }
         ucc_assert(ctx->addr_storage.rank == params->oob.oob_ep);
     }
     if (config->internal_oob) {
@@ -898,6 +1099,10 @@ ucc_status_t ucc_context_create_proc_info(
     }
     ucc_free(ctx->cl_ctx);
 error_ctx:
+    ucc_context_topo_cleanup(ctx->topo);
+    ctx->topo = NULL;
+    ucc_free(ctx->addr_storage.storage);
+    ctx->addr_storage.storage = NULL;
     ucc_free(ctx);
 error:
     return status;

src/core/ucc_context.h

@@ -152,6 +152,13 @@ ucc_status_t ucc_context_progress_deregister(ucc_context_t *ctx,
 ucc_status_t ucc_core_addr_exchange(ucc_context_t *context, ucc_oob_coll_t *oob,
                                     ucc_addr_storage_t *addr_storage);
 
+/* Performs context id address exchange between the processes group defined by OOB.
+   This function is used to exchange the context ids between the processes in order
+   to find the local rank.
+*/
+ucc_status_t ucc_core_ctx_id_exchange(ucc_context_t *context, ucc_oob_coll_t *oob,
+                                      ucc_addr_storage_t *addr_storage);
+
 /* UCC context packed address layout:
    --------------------------------------------------------------------------
    |ctx_id|host_info|n_components|id0|offset0|id1|offset1|..|