ucx: check supported transports for setting priority

Author: yosefe

contrib/platform/mellanox/optimized.conf

@@ -61,6 +61,7 @@
 coll = ^ml
 hwloc_base_binding_policy = core
 btl = self
+pml_ucx_priority = 60
 # Basic behavior to smooth startup
 mca_base_component_show_load_errors = 0
 orte_abort_timeout = 10

ompi/mca/pml/ucx/pml_ucx.h

@@ -57,6 +57,7 @@ struct mca_pml_ucx_module {
     mca_pml_ucx_freelist_t    convs;
 
     int                       priority;
+    int                       max_priority;
     bool                      cuda_initialized;
 };
 

ompi/mca/pml/ucx/pml_ucx_component.c

@@ -49,14 +49,22 @@ mca_pml_base_component_2_0_0_t mca_pml_ucx_component = {
 
 static int mca_pml_ucx_component_register(void)
 {
-    ompi_pml_ucx.priority = 51;
+    ompi_pml_ucx.priority = 25;
     (void) mca_base_component_var_register(&mca_pml_ucx_component.pmlm_version, "priority",
                                            "Priority of the UCX component",
                                            MCA_BASE_VAR_TYPE_INT, NULL, 0, 0,
                                            OPAL_INFO_LVL_3,
                                            MCA_BASE_VAR_SCOPE_LOCAL,
                                            &ompi_pml_ucx.priority);
 
+    ompi_pml_ucx.max_priority = 51;
+    (void) mca_base_component_var_register(&mca_pml_ucx_component.pmlm_version, "max_priority",
+                                           "Priority of the UCX component if suitable hardware is detected",
+                                           MCA_BASE_VAR_TYPE_INT, NULL, 0, 0,
+                                           OPAL_INFO_LVL_3,
+                                           MCA_BASE_VAR_SCOPE_LOCAL,
+                                           &ompi_pml_ucx.max_priority);
+
     ompi_pml_ucx.num_disconnect = 1;
     (void) mca_base_component_var_register(&mca_pml_ucx_component.pmlm_version, "num_disconnect",
                                            "How may disconnects go in parallel",
@@ -98,7 +106,9 @@ mca_pml_ucx_component_init(int* priority, bool enable_progress_threads,
         return NULL;
     }
 
-    *priority = ompi_pml_ucx.priority;
+    *priority = opal_common_ucx_priority(ompi_pml_ucx.ucp_context,
+                                         ompi_pml_ucx.priority,
+                                         ompi_pml_ucx.max_priority);
     return &ompi_pml_ucx.super;
 }
 

opal/mca/common/ucx/common_ucx.c

@@ -19,6 +19,7 @@
 #include "opal/memoryhooks/memory.h"
 
 #include <ucm/api/ucm.h>
+#include <stdio.h>
 
 /***********************************************************************/
 
@@ -126,6 +127,59 @@ OPAL_DECLSPEC void opal_common_ucx_mca_deregister(void)
     opal_output_close(opal_common_ucx.output);
 }
 
+OPAL_DECLSPEC int opal_common_ucx_priority(ucp_context_h context, int min, int max)
+{
+#if HAVE_DECL_OPEN_MEMSTREAM
+    static const char *transports[] = {
+        "cuda_ipc", "ud_verbs", "rc_verbs", "dc_mlx5"
+    };
+    char *match, *ptr, *buffer;
+    char needle[32];
+    FILE *stream;
+    size_t size;
+    size_t i;
+
+    if (min >= max) {
+        /* no need to check anything - minimal priority is already high enough */
+        return min;
+    }
+
+    stream = open_memstream(&buffer, &size);
+    if (stream == NULL) {
+        MCA_COMMON_UCX_VERBOSE(1, "failed to open memory stream for ucx info (%s), "
+                               "setting default ucx priority to %d",
+                               strerror(errno), min);
+        return min;
+    }
+
+    ucp_context_print_info(context, stream);
+
+    fclose(stream);
+
+    for (i = 0; i < sizeof(transports) / sizeof(transports[0]); ++i) {
+        /* Match "resource 6  :  md 5  dev 4  flags -- rc_verbs/mlx5_0:1" */
+        snprintf(needle, sizeof(needle), " %s/", transports[i]);
+        match = strstr(buffer, needle);
+        if (match != NULL) {
+            MCA_COMMON_UCX_VERBOSE(2, "found '%s' transport on %s, setting priority to %d",
+                                   transports[i], strtok_r(match, " \n", &ptr), max);
+            free(buffer);
+            return max;
+        }
+    }
+
+    free(buffer);
+
+    MCA_COMMON_UCX_VERBOSE(2, "didn't find matching transports, setting priority to %d",
+                           min);
+#else
+    MCA_COMMON_UCX_VERBOSE(2, "open_memstream() was not found, setting priority to %d",
+                           min);
+#endif
+
+    return min;
+}
+
 void opal_common_ucx_empty_complete_cb(void *request, ucs_status_t status)
 {
 }

opal/mca/common/ucx/common_ucx.h

@@ -107,6 +107,7 @@ extern opal_common_ucx_module_t opal_common_ucx;
 
 OPAL_DECLSPEC void opal_common_ucx_mca_register(void);
 OPAL_DECLSPEC void opal_common_ucx_mca_deregister(void);
+OPAL_DECLSPEC int opal_common_ucx_priority(ucp_context_h context, int min, int max);
 OPAL_DECLSPEC void opal_common_ucx_mca_proc_added(void);
 OPAL_DECLSPEC void opal_common_ucx_empty_complete_cb(void *request, ucs_status_t status);
 OPAL_DECLSPEC int opal_common_ucx_mca_pmix_fence(ucp_worker_h worker);

opal/mca/common/ucx/configure.m4

@@ -18,6 +18,8 @@ AC_DEFUN([MCA_opal_common_ucx_CONFIG],[
                [common_ucx_happy="yes"],
                [common_ucx_happy="no"])
 
+    AC_CHECK_DECLS([open_memstream], [], [], [[#include <stdio.h>]])
+
     AS_IF([test "$common_ucx_happy" = "yes"],
           [$1],
           [$2])