PML/UCX: Use net worker address for remote peers

For remote node peers pack smaller worker address, which contains
network device addresses only. This would reduce amount of OOB traffic
during startup.

Signed-off-by: Mikhail Brinskii <[email protected]>

Author: brminich

config/ompi_check_ucx.m4

@@ -120,6 +120,10 @@ AC_DEFUN([OMPI_CHECK_UCX],[
                                   UCP_ATOMIC_FETCH_OP_FXOR],
                                  [], [],
                                  [#include <ucp/api/ucp.h>])
+                  AC_CHECK_DECLS([UCP_WORKER_ATTR_FIELD_ADDRESS_FLAGS],
+                                 [AC_DEFINE([HAVE_UCP_WORKER_ADDRESS_FLAGS], [1],
+                                            [have worker address attribute])], [],
+                                 [#include <ucp/api/ucp.h>])
                   CPPFLAGS=$old_CPPFLAGS
 
                   OPAL_SUMMARY_ADD([[Transports]],[[Open UCX]],[$1],[$ompi_check_ucx_happy])])])

ompi/mca/pml/ucx/pml_ucx.c

@@ -82,11 +82,46 @@ mca_pml_ucx_module_t ompi_pml_ucx = {
 #define PML_UCX_REQ_ALLOCA() \
     ((char *)alloca(ompi_pml_ucx.request_size) + ompi_pml_ucx.request_size);
 
+#if HAVE_UCP_WORKER_ADDRESS_FLAGS
+static int mca_pml_ucx_send_worker_address_type(int addr_flags, int modex_scope)
+{
+    ucs_status_t status;
+    ucp_worker_attr_t attrs;
+    int rc;
+
+    attrs.field_mask    = UCP_WORKER_ATTR_FIELD_ADDRESS |
+                          UCP_WORKER_ATTR_FIELD_ADDRESS_FLAGS;
+    attrs.address_flags = addr_flags;
+
+    status = ucp_worker_query(ompi_pml_ucx.ucp_worker, &attrs);
+    if (UCS_OK != status) {
+        PML_UCX_ERROR("Failed to query UCP worker address");
+        return OMPI_ERROR;
+    }
+
+    OPAL_MODEX_SEND(rc, modex_scope, &mca_pml_ucx_component.pmlm_version,
+                    (void*)attrs.address, attrs.address_length);
+
+    ucp_worker_release_address(ompi_pml_ucx.ucp_worker, attrs.address);
+
+    if (OMPI_SUCCESS != rc) {
+        return OMPI_ERROR;
+    }
+
+    PML_UCX_VERBOSE(2, "Pack %s worker address, size %ld",
+                    (modex_scope == OPAL_PMIX_LOCAL) ? "local" : "remote",
+                    attrs.address_length);
+
+    return OMPI_SUCCESS;
+}
+#endif
 
 static int mca_pml_ucx_send_worker_address(void)
 {
-    ucp_address_t *address;
     ucs_status_t status;
+
+#if !HAVE_UCP_WORKER_ADDRESS_FLAGS
+    ucp_address_t *address;
     size_t addrlen;
     int rc;
 
@@ -96,16 +131,35 @@ static int mca_pml_ucx_send_worker_address(void)
         return OMPI_ERROR;
     }
 
+    PML_UCX_VERBOSE(2, "Pack worker address, size %ld", addrlen);
+
     OPAL_MODEX_SEND(rc, OPAL_PMIX_GLOBAL,
                     &mca_pml_ucx_component.pmlm_version, (void*)address, addrlen);
+
+    ucp_worker_release_address(ompi_pml_ucx.ucp_worker, address);
+
     if (OMPI_SUCCESS != rc) {
-        PML_UCX_ERROR("Open MPI couldn't distribute EP connection details");
-        return OMPI_ERROR;
+        goto err;
+    }
+#else
+    /* Pack just network device addresses for remote node peers */
+    status = mca_pml_ucx_send_worker_address_type(UCP_WORKER_ADDRESS_FLAG_NET_ONLY,
+                                                  OPAL_PMIX_REMOTE);
+    if (UCS_OK != status) {
+        goto err;
     }
 
-    ucp_worker_release_address(ompi_pml_ucx.ucp_worker, address);
+    status = mca_pml_ucx_send_worker_address_type(0, OPAL_PMIX_LOCAL);
+    if (UCS_OK != status) {
+        goto err;
+    }
+#endif
 
     return OMPI_SUCCESS;
+
+err:
+    PML_UCX_ERROR("Open MPI couldn't distribute EP connection details");
+    return OMPI_ERROR;
 }
 
 static int mca_pml_ucx_recv_worker_address(ompi_proc_t *proc,
@@ -121,6 +175,9 @@ static int mca_pml_ucx_recv_worker_address(ompi_proc_t *proc,
         PML_UCX_ERROR("Failed to receive UCX worker address: %s (%d)",
                       opal_strerror(ret), ret);
     }
+
+    PML_UCX_VERBOSE(2, "Got proc %d address, size %ld",
+                    proc->super.proc_name.vpid, *addrlen_p);
     return ret;
 }