From 8f28f4311a6cf234a7fc9a545d05a94ef8916070 Mon Sep 17 00:00:00 2001
From: Haakon Ludvig Langeland Ervik
 <45243236+haakon-e@users.noreply.github.com>
Date: Tue, 26 Nov 2024 13:47:34 -0800
Subject: [PATCH] fix: logging for preset comms device

---
 src/solver/type_getters.jl | 33 +++++++++++++++++++++------------
 src/solver/types.jl        | 16 ++++++++++++----
 2 files changed, 33 insertions(+), 16 deletions(-)

diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl
index ba3faf9c3b..498133a29e 100644
--- a/src/solver/type_getters.jl
+++ b/src/solver/type_getters.jl
@@ -589,18 +589,7 @@ function get_comms_context(parsed_args)
     end
     comms_ctx = ClimaComms.context(device)
     ClimaComms.init(comms_ctx)
-    if ClimaComms.iamroot(comms_ctx)
-        Logging.global_logger(Logging.ConsoleLogger(stderr, Logging.Info))
-    else
-        Logging.global_logger(Logging.NullLogger())
-    end
-    @info "Running on $(nameof(typeof(device)))."
-    if comms_ctx isa ClimaComms.SingletonCommsContext
-        @info "Setting up single-process ClimaAtmos run"
-    else
-        @info "Setting up distributed ClimaAtmos run" nprocs =
-            ClimaComms.nprocs(comms_ctx)
-    end
+
     if NVTX.isactive()
         # makes output on buildkite a bit nicer
         if ClimaComms.iamroot(comms_ctx)
@@ -613,6 +602,26 @@ function get_comms_context(parsed_args)
     return comms_ctx
 end
 
+"""
+    silence_non_root_processes(comms_ctx)
+
+Set the logging behavior based on the process rank within the given communication context `comms_ctx`.
+If the process is the root process, logging is set to display messages to the console with `Info` level.
+For all other processes, logging is silenced by setting it to a `NullLogger`.
+
+# Arguments
+- `comms_ctx`: The communication context used to determine the rank of the process.
+
+"""
+function silence_non_root_processes(comms_ctx)
+    # Set logging to only display for the root process
+    if ClimaComms.iamroot(comms_ctx)
+        Logging.global_logger(Logging.ConsoleLogger(stderr, Logging.Info))
+    else
+        Logging.global_logger(Logging.NullLogger())
+    end
+end
+
 function get_simulation(config::AtmosConfig)
     params = create_parameter_set(config)
     atmos = get_atmos(config, params)
diff --git a/src/solver/types.jl b/src/solver/types.jl
index 7afa3b7d11..b566d30419 100644
--- a/src/solver/types.jl
+++ b/src/solver/types.jl
@@ -647,17 +647,25 @@ function AtmosConfig(
         override_file = CP.merge_toml_files(config["toml"]),
     )
     comms_ctx = isnothing(comms_ctx) ? get_comms_context(config) : comms_ctx
+    device = ClimaComms.device(comms_ctx)
+    silence_non_root_processes(comms_ctx)
+    @info "Running on $(nameof(typeof(device)))"
+    if comms_ctx isa ClimaComms.SingletonCommsContext
+        @info "Setting up single-process ClimaAtmos run"
+    else
+        @info "Setting up distributed ClimaAtmos run" nprocs =
+            ClimaComms.nprocs(comms_ctx)
+    end
 
     config = config_with_resolved_and_acquired_artifacts(config, comms_ctx)
-    device = ClimaComms.device(comms_ctx)
     if device isa ClimaComms.CPUMultiThreaded
-        @info "Running ClimaCore in threaded mode, with $(Threads.nthreads()) threads."
+        @info "Running ClimaCore in threaded mode, with $(Threads.nthreads()) threads"
     else
-        @info "Running ClimaCore in unthreaded mode."
+        @info "Running ClimaCore in unthreaded mode"
     end
 
     isempty(job_id) &&
-        @warn "`job_id` is empty and likely not passed to AtmosConfig."
+        @warn "`job_id` is empty and likely not passed to AtmosConfig"
 
     @info "Making AtmosConfig with config files: $(sprint(config_summary, config_files))"