Merge pull request #971 from CliMA/js/prints

print more cleanly

Project.toml

@@ -9,6 +9,7 @@ ClimaCore = "d414da3d-4745-48bb-8d80-42e94e092884"
 ClimaCoreTempestRemap = "d934ef94-cdd4-4710-83d6-720549644b70"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
+Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
@@ -21,6 +22,7 @@ ClimaCore = "0.14.6"
 ClimaCoreTempestRemap = "0.3"
 Dates = "1"
 JLD2 = "0.4, 0.5"
+Logging = "1"
 NCDatasets = "0.11, 0.12, 0.13, 0.14.2"
 SciMLBase = "1, 2"
 StaticArrays = "1"

experiments/ClimaCore/heat-diffusion/run.jl

@@ -87,7 +87,6 @@
 ## import Pkg; Pkg.add(url="https://github.com/CliMA/ClimaCore.jl",rev="main")
 
 #  - load external packages:
-import Logging
 import SciMLBase
 
 # load CliMA packages:

experiments/ClimaEarth/components/atmosphere/climaatmos.jl

@@ -302,27 +302,21 @@ function get_atmos_config_dict(coupler_dict::Dict, job_id::String)
     atmos_config_file = coupler_dict["atmos_config_file"]
     atmos_config_repo = coupler_dict["atmos_config_repo"]
     # override default or specified configs with coupler arguments, and set the correct atmos config_file
-    comms_ctx = ClimaComms.context()
     if atmos_config_repo == "ClimaCoupler"
         @assert !isnothing(atmos_config_file) "Must specify `atmos_config_file` within ClimaCoupler."
-        if ClimaComms.iamroot(comms_ctx)
-            @info "Using Atmos configuration from ClimaCoupler in $atmos_config_file"
-        end
+
+        @info "Using Atmos configuration from ClimaCoupler in $atmos_config_file"
         atmos_config = merge(
             CA.override_default_config(joinpath(pkgdir(ClimaCoupler), atmos_config_file)),
             coupler_dict,
             Dict("config_file" => atmos_config_file),
         )
     elseif atmos_config_repo == "ClimaAtmos"
         if isnothing(atmos_config_file)
-            if ClimaComms.iamroot(comms_ctx)
-                @info "Using Atmos default configuration"
-            end
+            @info "Using Atmos default configuration"
             atmos_config = merge(CA.default_config_dict(), coupler_dict, Dict("config_file" => atmos_config_file))
         else
-            if ClimaComms.iamroot(comms_ctx)
-                @info "Using Atmos configuration from $atmos_config_file"
-            end
+            @info "Using Atmos configuration from $atmos_config_file"
             atmos_config = merge(
                 CA.override_default_config(joinpath(pkgdir(CA), atmos_config_file)),
                 coupler_dict,
@@ -343,9 +337,7 @@ function get_atmos_config_dict(coupler_dict::Dict, job_id::String)
     toml_file = isnothing(toml_file) ? joinpath(pkgdir(ClimaCoupler), default_toml_file) : toml_file
 
     if !isnothing(toml_file)
-        if ClimaComms.iamroot(comms_ctx)
-            @info "Overwriting Atmos parameters from $toml_file"
-        end
+        @info "Overwriting Atmos parameters from $toml_file"
         atmos_config = merge(atmos_config, Dict("toml" => [toml_file]))
     end
 

experiments/ClimaEarth/run_amip.jl

@@ -108,6 +108,8 @@ if isinteractive()
     parsed_args["job_id"] = "interactive_debug"
 end
 
+comms_ctx = Utilities.get_comms_context(parsed_args)
+
 ## the unique job id should be passed in via the command line
 job_id = parsed_args["job_id"]
 @assert !isnothing(job_id) "job_id must be passed in via the command line"
@@ -116,9 +118,6 @@ job_id = parsed_args["job_id"]
 config_dict = YAML.load_file(parsed_args["config_file"])
 config_dict = merge(parsed_args, config_dict)
 
-comms_ctx = Utilities.get_comms_context(parsed_args)
-ClimaComms.init(comms_ctx)
-
 ## get component model dictionaries (if applicable)
 atmos_config_dict, config_dict = get_atmos_config_dict(config_dict, job_id)
 atmos_config_object = CA.AtmosConfig(atmos_config_dict)
@@ -165,11 +164,8 @@ temporary files will be saved.
 COUPLER_OUTPUT_DIR = joinpath(config_dict["coupler_output_dir"], joinpath(mode_name, job_id))
 dir_paths = setup_output_dirs(output_dir = COUPLER_OUTPUT_DIR, comms_ctx = comms_ctx)
 
-
-if ClimaComms.iamroot(comms_ctx)
-    @info(dir_paths.output)
-    config_dict["print_config_dict"] && @info(config_dict)
-end
+@info(dir_paths.output)
+config_dict["print_config_dict"] && @info(config_dict)
 
 #=
 ## Data File Paths
@@ -193,11 +189,11 @@ returns a `ComponentModelSimulation` object (see `Interfacer` docs for more deta
 This uses the `ClimaAtmos.jl` model, with parameterization options specified in the `atmos_config_object` dictionary.
 =#
 
-Utilities.show_memory_usage(comms_ctx)
+Utilities.show_memory_usage()
 
 ## init atmos model component
 atmos_sim = atmos_init(atmos_config_object);
-Utilities.show_memory_usage(comms_ctx)
+Utilities.show_memory_usage()
 
 thermo_params = get_thermo_params(atmos_sim) # TODO: this should be shared by all models #342
 
@@ -226,7 +222,7 @@ land_area_fraction = SpaceVaryingInput(land_mask_data, "LSMASK", boundary_space)
 if !mono_surface
     land_area_fraction = Regridder.binary_mask.(land_area_fraction)
 end
-Utilities.show_memory_usage(comms_ctx)
+Utilities.show_memory_usage()
 
 #=
 ### Surface Models: AMIP and SlabPlanet Modes
@@ -245,9 +241,9 @@ In this section of the code, we initialize all component models and read in the
 The specific models and data that are set up depend on which mode we're running.
 =#
 
-ClimaComms.iamroot(comms_ctx) && @info(mode_name)
+@info(mode_name)
 if mode_name == "amip"
-    ClimaComms.iamroot(comms_ctx) && @info("AMIP boundary conditions - do not expect energy conservation")
+    @info("AMIP boundary conditions - do not expect energy conservation")
 
     ## land model
     land_sim = bucket_init(
@@ -327,7 +323,7 @@ if mode_name == "amip"
         SIC_timevaryinginput = SIC_timevaryinginput,
         CO2_timevaryinginput = CO2_timevaryinginput,
     )
-    Utilities.show_memory_usage(comms_ctx)
+    Utilities.show_memory_usage()
 
 elseif mode_name in ("slabplanet", "slabplanet_aqua", "slabplanet_terra")
 
@@ -379,7 +375,7 @@ elseif mode_name in ("slabplanet", "slabplanet_aqua", "slabplanet_terra")
     ))
 
     mode_specifics = (; name = mode_name, SST_timevaryinginput = nothing, SIC_timevaryinginput = nothing)
-    Utilities.show_memory_usage(comms_ctx)
+    Utilities.show_memory_usage()
 
 elseif mode_name == "slabplanet_eisenman"
 
@@ -423,7 +419,7 @@ elseif mode_name == "slabplanet_eisenman"
     )
 
     mode_specifics = (; name = mode_name, SST_timevaryinginput = nothing, SIC_timevaryinginput = nothing)
-    Utilities.show_memory_usage(comms_ctx)
+    Utilities.show_memory_usage()
 end
 
 #=
@@ -458,7 +454,7 @@ coupler_field_names = (
 )
 coupler_fields =
     NamedTuple{coupler_field_names}(ntuple(i -> CC.Fields.zeros(boundary_space), length(coupler_field_names)))
-Utilities.show_memory_usage(comms_ctx)
+Utilities.show_memory_usage()
 
 ## model simulations
 model_sims = (atmos_sim = atmos_sim, ice_sim = ice_sim, land_sim = land_sim, ocean_sim = ocean_sim);
@@ -492,7 +488,7 @@ if use_coupler_diagnostics
     )
 
     diagnostics = (monthly_3d_diags, monthly_2d_diags)
-    Utilities.show_memory_usage(comms_ctx)
+    Utilities.show_memory_usage()
 else
     diagnostics = ()
 end
@@ -598,7 +594,7 @@ cs = Interfacer.CoupledSimulation{FT}(
     turbulent_fluxes,
     thermo_params,
 );
-Utilities.show_memory_usage(comms_ctx)
+Utilities.show_memory_usage()
 
 #=
 ## Restart component model states if specified
@@ -683,16 +679,14 @@ function solve_coupler!(cs)
     (; model_sims, Δt_cpl, tspan, comms_ctx) = cs
     (; atmos_sim, land_sim, ocean_sim, ice_sim) = model_sims
 
-    ClimaComms.iamroot(comms_ctx) && @info("Starting coupling loop")
+    @info("Starting coupling loop")
     ## step in time
     for t in ((tspan[begin] + Δt_cpl):Δt_cpl:tspan[end])
 
         cs.dates.date[1] = TimeManager.current_date(cs, t)
 
         ## print date on the first of month
-        if cs.dates.date[1] >= cs.dates.date1[1]
-            ClimaComms.iamroot(comms_ctx) && @show(cs.dates.date[1])
-        end
+        cs.dates.date[1] >= cs.dates.date1[1] && @info(cs.dates.date[1])
 
         if cs.mode.name == "amip"
 
@@ -800,7 +794,7 @@ walltime = ClimaComms.@elapsed comms_ctx.device begin
         solve_coupler!(cs)
     end
 end
-ClimaComms.iamroot(comms_ctx) && @show(walltime)
+@info(walltime)
 
 ## Use ClimaAtmos calculation to show the simulated years per day of the simulation (SYPD)
 es = CA.EfficiencyStats(tspan, walltime)
@@ -821,7 +815,7 @@ if ClimaComms.iamroot(comms_ctx)
     end
 
     open(joinpath(dir_paths.artifacts, "max_rss_cpu.txt"), "w") do cpu_max_rss_filename
-        cpu_max_rss_GB = Utilities.show_memory_usage(comms_ctx)
+        cpu_max_rss_GB = Utilities.show_memory_usage()
         println(cpu_max_rss_filename, cpu_max_rss_GB)
     end
 end

experiments/ClimaEarth/run_cloudless_aquaplanet.jl

@@ -132,14 +132,13 @@ atmos_config_object = CA.AtmosConfig(atmos_config_dict)
 =#
 
 comms_ctx = Utilities.get_comms_context(Dict("device" => "auto"))
-ClimaComms.init(comms_ctx)
 
 #=
 ### I/O Directory Setup
 =#
 
 dir_paths = setup_output_dirs(output_dir = coupler_output_dir, comms_ctx = comms_ctx)
-ClimaComms.iamroot(comms_ctx) ? @info(config_dict) : nothing
+@info(config_dict)
 
 #=
 ## Component Model Initialization
@@ -204,7 +203,7 @@ coupler_field_names = (
 )
 coupler_fields =
     NamedTuple{coupler_field_names}(ntuple(i -> CC.Fields.zeros(boundary_space), length(coupler_field_names)))
-Utilities.show_memory_usage(comms_ctx)
+Utilities.show_memory_usage()
 
 ## model simulations
 model_sims = (atmos_sim = atmos_sim, ocean_sim = ocean_sim);
@@ -321,16 +320,14 @@ FieldExchanger.update_model_sims!(cs.model_sims, cs.fields, cs.turbulent_fluxes)
 function solve_coupler!(cs)
     (; Δt_cpl, tspan, comms_ctx) = cs
 
-    ClimaComms.iamroot(comms_ctx) && @info("Starting coupling loop")
+    @info("Starting coupling loop")
     ## step in time
     for t in ((tspan[begin] + Δt_cpl):Δt_cpl:tspan[end])
 
         cs.dates.date[1] = TimeManager.current_date(cs, t)
 
         ## print date on the first of month
-        if cs.dates.date[1] >= cs.dates.date1[1]
-            ClimaComms.iamroot(comms_ctx) && @show(cs.dates.date[1])
-        end
+        cs.dates.date[1] >= cs.dates.date1[1] && @info(cs.dates.date[1])
 
         ClimaComms.barrier(comms_ctx)
 

experiments/ClimaEarth/run_cloudy_aquaplanet.jl

@@ -156,12 +156,13 @@ atmos_config_object.toml_dict["max_area_limiter_scale"]["value"] = 0
 comms_ctx = Utilities.get_comms_context(Dict("device" => "auto"))
 ClimaComms.init(comms_ctx)
 
+
 #=
 ### I/O Directory Setup
 =#
 
 dir_paths = setup_output_dirs(output_dir = coupler_output_dir, comms_ctx = comms_ctx)
-ClimaComms.iamroot(comms_ctx) ? @info(config_dict) : nothing
+@info(config_dict)
 
 #=
 ## Component Model Initialization
@@ -225,7 +226,7 @@ coupler_field_names = (
 )
 coupler_fields =
     NamedTuple{coupler_field_names}(ntuple(i -> CC.Fields.zeros(boundary_space), length(coupler_field_names)))
-Utilities.show_memory_usage(comms_ctx)
+Utilities.show_memory_usage()
 
 ## model simulations
 model_sims = (atmos_sim = atmos_sim, ocean_sim = ocean_sim);
@@ -333,16 +334,14 @@ FieldExchanger.update_model_sims!(cs.model_sims, cs.fields, cs.turbulent_fluxes)
 function solve_coupler!(cs)
     (; Δt_cpl, tspan, comms_ctx) = cs
 
-    ClimaComms.iamroot(comms_ctx) && @info("Starting coupling loop")
+    @info("Starting coupling loop")
     ## step in time
     for t in ((tspan[begin] + Δt_cpl):Δt_cpl:tspan[end])
 
         cs.dates.date[1] = TimeManager.current_date(cs, t)
 
         ## print date on the first of month
-        if cs.dates.date[1] >= cs.dates.date1[1]
-            ClimaComms.iamroot(comms_ctx) && @show(cs.dates.date[1])
-        end
+        cs.dates.date[1] >= cs.dates.date1[1] && @info(cs.dates.date[1])
 
         ClimaComms.barrier(comms_ctx)
 

experiments/ClimaEarth/run_cloudy_slabplanet.jl

@@ -154,14 +154,13 @@ atmos_config_object.toml_dict["max_area_limiter_scale"]["value"] = 0
 =#
 
 comms_ctx = Utilities.get_comms_context(Dict("device" => "auto"))
-ClimaComms.init(comms_ctx)
 
 #=
 ### I/O Directory Setup
 =#
 
 dir_paths = setup_output_dirs(output_dir = coupler_output_dir, comms_ctx = comms_ctx)
-ClimaComms.iamroot(comms_ctx) ? @info(config_dict) : nothing
+@info(config_dict)
 
 #=
 ## Data File Paths
@@ -275,7 +274,7 @@ coupler_field_names = (
 )
 coupler_fields =
     NamedTuple{coupler_field_names}(ntuple(i -> CC.Fields.zeros(boundary_space), length(coupler_field_names)))
-Utilities.show_memory_usage(comms_ctx)
+Utilities.show_memory_usage()
 
 ## model simulations
 model_sims = (atmos_sim = atmos_sim, ocean_sim = ocean_sim);
@@ -385,16 +384,14 @@ FieldExchanger.update_model_sims!(cs.model_sims, cs.fields, cs.turbulent_fluxes)
 function solve_coupler!(cs)
     (; Δt_cpl, tspan, comms_ctx) = cs
 
-    ClimaComms.iamroot(comms_ctx) && @info("Starting coupling loop")
+    @info("Starting coupling loop")
     ## step in time
     for t in ((tspan[begin] + Δt_cpl):Δt_cpl:tspan[end])
 
         cs.dates.date[1] = TimeManager.current_date(cs, t)
 
         ## print date on the first of month
-        if cs.dates.date[1] >= cs.dates.date1[1]
-            ClimaComms.iamroot(comms_ctx) && @show(cs.dates.date[1])
-        end
+        cs.dates.date[1] >= cs.dates.date1[1] && @info(cs.dates.date[1])
 
         ClimaComms.barrier(comms_ctx)