Add CLI arg for walltime logging

Author: imreddyTeja

experiments/ClimaEarth/cli_options.jl

@@ -141,6 +141,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

@@ -510,7 +510,9 @@ function CoupledSimulation(config_dict::AbstractDict)
         (coupled_sim) -> report_walltime(wt, coupled_sim.model_sims.atmos_sim.integrator)
     end
     walltime_cb = TimeManager.Callback(walltime_report_cond, walltime_affect!)
-    callbacks = (checkpoint_cb, walltime_cb)
+    # Disable coupler walltime logging if atmos is using its own walltime logging is true
+    callbacks =
+        config_dict["atmos_log_progress"] ? (checkpoint_cb,) : (checkpoint_cb, walltime_cb)
 
     #= Set up default AMIP diagnostics
     Use ClimaDiagnostics for default AMIP diagnostics, which currently include turbulent energy fluxes.