NA UCX: add ucx log outlet and redirect ucx log

Use default HG log level if UCX_LOG_LEVEL is not set
mercury-hpc · Jun 20, 2024 · ea078c9 · ea078c9
1 parent 7b7f9ad
commit ea078c9
Show file tree

Hide file tree

Showing 3 changed files with 115 additions and 1 deletion.
diff --git a/src/na/na.c b/src/na/na.c
@@ -245,8 +245,9 @@ HG_LOG_SUBSYS_DECL_STATE_REGISTER(poll_loop, NA_SUBSYS_NAME, HG_LOG_OFF);
 HG_LOG_SUBSYS_DECL_STATE_REGISTER(ip, NA_SUBSYS_NAME, HG_LOG_OFF);
 HG_LOG_SUBSYS_DECL_STATE_REGISTER(perf, NA_SUBSYS_NAME, HG_LOG_OFF);
 
-/* Extra log outlet for libfabric */
+/* Extra log outlet for libfabric / ucx */
 HG_LOG_SUBSYS_DECL_STATE_REGISTER(libfabric, NA_SUBSYS_NAME, HG_LOG_OFF);
+HG_LOG_SUBSYS_DECL_STATE_REGISTER(ucx, NA_SUBSYS_NAME, HG_LOG_OFF);
 
 #ifdef NA_HAS_DYNAMIC_PLUGINS
 /* Initialize list of plugins etc */

diff --git a/src/na/na_error.h b/src/na/na_error.h
@@ -38,6 +38,8 @@ extern NA_PLUGIN_VISIBILITY HG_LOG_OUTLET_SUBSYS_DECL(
 /* Plugin specific log (must be declared here to prevent contructor issues) */
 extern NA_PLUGIN_VISIBILITY HG_LOG_OUTLET_SUBSYS_DECL(
     libfabric, na); /* Libfabric log */
+extern NA_PLUGIN_VISIBILITY HG_LOG_OUTLET_SUBSYS_DECL(
+    ucx, na); /* UCX log */
 
 /* Base log macros */
 #define NA_LOG_ERROR(...) HG_LOG_WRITE(na, HG_LOG_LEVEL_ERROR, __VA_ARGS__)

diff --git a/src/na/na_ucx.c b/src/na/na_ucx.c
@@ -17,6 +17,7 @@
 #include "mercury_thread_spin.h"
 
 #include <ucp/api/ucp.h>
+#include <ucs/debug/log_def.h>
 #include <uct/api/uct.h> /* To query component info */
 
 #include <stdalign.h>
@@ -268,6 +269,34 @@ enum na_ucp_type { NA_UCP_CONFIG, NA_UCP_CONTEXT, NA_UCP_WORKER };
 /* NA UCP helpers                                                            */
 /*---------------------------------------------------------------------------*/
 
+/**
+ * Import/close UCX log.
+ */
+static void
+na_ucs_log_import(void) NA_CONSTRUCTOR;
+static void
+na_ucs_log_close(void) NA_DESTRUCTOR;
+
+/**
+ * Print UCX log.
+ */
+static ucs_log_func_rc_t
+na_ucs_log_func(const char *file, unsigned line, const char *function,
+    ucs_log_level_t level, const ucs_log_component_config_t *comp_conf,
+    const char *message, va_list ap);
+
+/**
+ * Convert UCX log level to HG log level.
+ */
+static enum hg_log_level
+na_ucs_log_level_to_hg(ucs_log_level_t level);
+
+/**
+ * Convert HG log level to UCX log level string.
+ */
+static const char *
+na_ucs_log_level_to_string(enum hg_log_level level);
+
 /**
  * Convert UCX status to NA return values.
  */
@@ -1044,6 +1073,79 @@ na_ucs_status_to_na(ucs_status_t status)
     return ret;
 }
 
+/*---------------------------------------------------------------------------*/
+static void
+na_ucs_log_import(void)
+{
+    ucs_log_push_handler(na_ucs_log_func);
+}
+
+/*---------------------------------------------------------------------------*/
+static void
+na_ucs_log_close(void)
+{
+    ucs_log_pop_handler();
+}
+
+/*---------------------------------------------------------------------------*/
+static ucs_log_func_rc_t
+na_ucs_log_func(const char *file, unsigned line, const char *function,
+    ucs_log_level_t level, const ucs_log_component_config_t *comp_conf,
+    const char *message, va_list ap)
+{
+    HG_LOG_VWRITE_FUNC(na_ucx, na_ucs_log_level_to_hg(level), comp_conf->name,
+        file, line, function, false, message, ap);
+
+    return UCS_LOG_FUNC_RC_STOP;
+}
+
+/*---------------------------------------------------------------------------*/
+static enum hg_log_level
+na_ucs_log_level_to_hg(ucs_log_level_t level)
+{
+    switch (level) {
+        case UCS_LOG_LEVEL_FATAL:
+        case UCS_LOG_LEVEL_ERROR:
+            return HG_LOG_LEVEL_ERROR;
+        case UCS_LOG_LEVEL_WARN:
+            return HG_LOG_LEVEL_WARNING;
+        case UCS_LOG_LEVEL_DIAG:
+        case UCS_LOG_LEVEL_INFO:
+        case UCS_LOG_LEVEL_DEBUG:
+        case UCS_LOG_LEVEL_TRACE:
+        case UCS_LOG_LEVEL_TRACE_REQ:
+        case UCS_LOG_LEVEL_TRACE_DATA:
+        case UCS_LOG_LEVEL_TRACE_ASYNC:
+        case UCS_LOG_LEVEL_TRACE_FUNC:
+        case UCS_LOG_LEVEL_TRACE_POLL:
+            return HG_LOG_LEVEL_DEBUG;
+        case UCS_LOG_LEVEL_LAST:
+        case UCS_LOG_LEVEL_PRINT:
+        default:
+            return HG_LOG_LEVEL_MAX;
+    }
+}
+
+/*---------------------------------------------------------------------------*/
+static const char *
+na_ucs_log_level_to_string(enum hg_log_level level)
+{
+    switch (level) {
+        case HG_LOG_LEVEL_ERROR:
+            return "error";
+        case HG_LOG_LEVEL_WARNING:
+            return "warn";
+        case HG_LOG_LEVEL_MIN_DEBUG:
+            return "trace";
+        case HG_LOG_LEVEL_DEBUG:
+            return "debug";
+        case HG_LOG_LEVEL_NONE:
+        case HG_LOG_LEVEL_MAX:
+        default:
+            return "";
+    }
+}
+
 /*---------------------------------------------------------------------------*/
 static na_return_t
 na_uct_get_transport_alias(
@@ -1267,6 +1369,15 @@ na_ucp_config_init(
             ucs_status_string(status));
     }
 
+    /* Set matching log level by default */
+    if (getenv("UCX_LOG_LEVEL") == NULL) {
+        status = ucp_config_modify(config, "LOG_LEVEL",
+            na_ucs_log_level_to_string(hg_log_get_level()));
+        NA_CHECK_SUBSYS_ERROR(cls, status != UCS_OK, error, ret,
+            na_ucs_status_to_na(status), "ucp_config_modify() failed (%s)",
+            ucs_status_string(status));
+    }
+
     /* Reuse addr for tcp by default */
     if (getenv("UCX_CM_REUSEADDR") == NULL) {
         status = ucp_config_modify(config, "CM_REUSEADDR", "y");