Add WallTimeReport callback

Add CLI arg for walltime logging

Make suggested changtes; move cb into src; add tests

Author: imreddyTeja

experiments/ClimaEarth/cli_options.jl

@@ -145,6 +145,10 @@ function argparse_settings()
         help = "An optional YAML file used to overwrite the default model parameters."
         arg_type = String
         default = nothing
+        "--atmos_log_progress"
+        help = "Use the ClimaAtmos walltime logging callback instead of the default ClimaCoupler one [`false` (default), `true`]"
+        arg_type = Bool
+        default = false
         "--albedo_model"
         help = "Type of albedo model. [`ConstantAlbedo`, `RegressionFunctionAlbedo`, `CouplerAlbedo` (default)]"
         arg_type = String

experiments/ClimaEarth/components/atmosphere/climaatmos.jl

@@ -626,6 +626,7 @@ function get_atmos_config_dict(
     # can pick up from where we have left. NOTE: This should not be needed, but
     # there is no easy way to initialize ClimaAtmos with a different t_start
     atmos_config["dt_save_state_to_disk"] = coupler_config["checkpoint_dt"]
+    atmos_config["log_progress"] = coupler_config["atmos_log_progress"]
 
     # The Atmos `get_simulation` function expects the atmos config to contains its timestep size
     # in the `dt` field. If there is a `dt_atmos` field in coupler_config, we add it to the atmos config as `dt`

experiments/ClimaEarth/setup_run.jl

@@ -497,7 +497,13 @@ function CoupledSimulation(config_dict::AbstractDict)
     checkpoint_cb =
         TimeManager.Callback(schedule_checkpoint, sim -> Checkpointer.checkpoint_sims(sim))
 
-    callbacks = (checkpoint_cb,)
+    # Don't use coupler walltime logging if atmos is using its own walltime logging is true
+    if config_dict["atmos_log_progress"]
+        callbacks = (checkpoint_cb,)
+    else
+        walltime_cb = TimeManager.capped_geometric_walltime_cb(t_start, t_end, Δt_cpl)
+        callbacks = (checkpoint_cb, walltime_cb)
+    end
 
     #= Set up default AMIP diagnostics
     Use ClimaDiagnostics for default AMIP diagnostics, which currently include turbulent energy fluxes.

src/TimeManager.jl

@@ -9,6 +9,7 @@ module TimeManager
 import Dates
 import ..Interfacer
 import ..Utilities: time_to_seconds
+import ClimaUtilities.OnlineLogging: WallTimeInfo, report_walltime
 
 """
     time_to_period(s::String)
@@ -100,4 +101,28 @@ function (::NeverSchedule)(args...)
     return false
 end
 
+"""
+    capped_geometric_walltime_cb(t_start, t_end, Δt_cpl)
+
+Create a callback that reports walltime at when the number of steps taken is a power of 2, or
+when the percent of the simulation that is completed is a multiple of 5. This skips the
+first two steps to avoid compilation time noise.
+"""
+function capped_geometric_walltime_cb(t_start, t_end, Δt_cpl)
+    tot_steps = Int(ceil(float(t_end - t_start) / float(Δt_cpl)))
+    five_percent_steps = ceil(Int, 0.05 * tot_steps)
+    steps_taken = (integrator) -> float(integrator.t - t_start) / float(Δt_cpl)
+    walltime_report_cond =
+        (integrator) -> begin
+            nsteps = steps_taken(integrator)
+            # skip first two steps for compilation
+            (nsteps <= 2) && return false
+            return nsteps % five_percent_steps == 0 || ispow2(nsteps)
+        end
+    walltime_affect! = let wt = WallTimeInfo()
+        (coupled_sim) -> report_walltime(wt, coupled_sim.model_sims.atmos_sim.integrator)
+    end
+    return TimeManager.Callback(walltime_report_cond, walltime_affect!)
+end
+
 end

test/utilities_tests.jl

@@ -5,6 +5,7 @@ import Test: @testset, @test
 import ClimaComms
 ClimaComms.@import_required_backends
 import ClimaCoupler: Utilities
+import ClimaCoupler: TimeManager
 import ClimaCore as CC
 
 # Initialize MPI context, in case
@@ -71,4 +72,31 @@ for FT in (Float32, Float64)
         )
         @test Utilities.integral(ones(space3d)) == sum(ones(space3d))
     end
+
+    @testset "WallTime Callback" begin
+        t_start = 0.0
+        t_end = 10.0
+        Δt_cpl = 0.1
+
+        cb = TimeManager.capped_geometric_walltime_cb(t_start, t_end, Δt_cpl)
+
+        # First two steps should not trigger
+        fake_integrator = (; t = t_start + Δt_cpl)
+        @test !cb.schedule(fake_integrator)
+        fake_integrator = (; t = t_start + Δt_cpl * 2)
+        @test !cb.schedule(fake_integrator)
+        # step 4, 8, 16 should trigger
+        fake_integrator = (; t = t_start + Δt_cpl * 4)
+        @test cb.schedule(fake_integrator)
+        fake_integrator = (; t = t_start + Δt_cpl * 8)
+        @test cb.schedule(fake_integrator)
+        fake_integrator = (; t = t_start + Δt_cpl * 14)
+        @test !cb.schedule(fake_integrator)
+        fake_integrator = (; t = t_start + Δt_cpl * 16)
+        @test cb.schedule(fake_integrator)
+
+        # 20% should trigger
+        fake_integrator = (; t = t_start + Δt_cpl * 20)
+        @test cb.schedule(fake_integrator)
+    end
 end