Add save_cache option

Author: ph-kev

experiments/ClimaEarth/cli_options.jl

@@ -104,6 +104,10 @@ function argparse_settings()
         help = "Boolean flag indicating whether to read the cache from the restart file if available [`true` (default), `false`]"
         arg_type = Bool
         default = true
+        "--save_cache"
+        help = "Boolean flag indicating whether to save the state and cache or only the state when checkpointing [`true` (default), `false`]"
+        arg_type = Bool
+        default = true
         # Diagnostics information
         "--use_coupler_diagnostics"
         help = "Boolean flag indicating whether to compute and output coupler diagnostics [`true` (default), `false`]"

experiments/ClimaEarth/setup_run.jl

@@ -131,6 +131,7 @@ function CoupledSimulation(config_dict::AbstractDict)
         restart_dir,
         restart_t,
         restart_cache,
+        save_cache,
         use_land_diagnostics,
         diagnostics_dt,
         evolving_ocean,
@@ -493,7 +494,8 @@ function CoupledSimulation(config_dict::AbstractDict)
     =#
     schedule_checkpoint =
         EveryCalendarDtSchedule(TimeManager.time_to_period(checkpoint_dt); start_date)
-    checkpoint_cb = TimeManager.Callback(schedule_checkpoint, Checkpointer.checkpoint_sims)
+    checkpoint_cb =
+        TimeManager.Callback(schedule_checkpoint, sim -> Checkpointer.checkpoint_sims(sim))
 
     callbacks = (checkpoint_cb,)
 
@@ -537,6 +539,7 @@ function CoupledSimulation(config_dict::AbstractDict)
         dir_paths,
         thermo_params,
         diags_handler,
+        save_cache,
     )
 
     #=

experiments/ClimaEarth/test/restart_state_only.jl

@@ -49,10 +49,18 @@ two_steps["dt_rad"] = "180secs"
 two_steps["checkpoint_dt"] = "360secs"
 two_steps["coupler_output_dir"] = tmpdir
 two_steps["job_id"] = "two_steps"
+two_steps["save_cache"] = false
 
 println("Simulating two steps")
 cs_two_steps = setup_and_run(two_steps)
 
+@testset "Cache files should not exist" begin
+    checkpoints_dir = joinpath(tmpdir, two_steps["job_id"], "checkpoints")
+    checkpoint_files = filter!(isfile, readdir(checkpoints_dir, join = true))
+    are_not_cache_files = .!occursin.("cache", basename.(checkpoint_files))
+    @test all(are_not_cache_files)
+end
+
 # Check that we can pick up a simulation by providing t_restart and restart_dir
 println("Simulating two steps, options from command line")
 two_steps_reading = deepcopy(two_steps)
@@ -63,6 +71,7 @@ two_steps_reading["restart_dir"] = cs_two_steps.dir_paths.checkpoints_dir
 two_steps_reading["restart_t"] = 360
 two_steps_reading["restart_cache"] = false
 two_steps_reading["job_id"] = "two_steps_reading"
+two_steps_reading["save_cache"] = false
 
 cs_two_steps_reading = setup_and_run(two_steps_reading)
 @testset "Restarts from command line arguments" begin

experiments/ClimaEarth/user_io/arg_parsing.jl

@@ -144,6 +144,7 @@ function get_coupler_args(config_dict::Dict)
     restart_dir = config_dict["restart_dir"]
     restart_t = config_dict["restart_t"]
     restart_cache = config_dict["restart_cache"]
+    save_cache = config_dict["save_cache"]
 
     # Diagnostics information
     use_coupler_diagnostics = config_dict["use_coupler_diagnostics"]
@@ -192,6 +193,7 @@ function get_coupler_args(config_dict::Dict)
         restart_dir,
         restart_t,
         restart_cache,
+        save_cache,
         use_coupler_diagnostics,
         diagnostics_dt,
         evolving_ocean,

src/Checkpointer.jl

@@ -128,7 +128,8 @@ end
 """
     checkpoint_sims(cs::CoupledSimulation)
 
-This is a callback function that checkpoints all simulations defined in the current coupled simulation.
+This is a callback function that checkpoints all simulations defined in the
+current coupled simulation.
 """
 function checkpoint_sims(cs::Interfacer.CoupledSimulation)
     time = Int(round(float(cs.t[])))
@@ -148,7 +149,7 @@ function checkpoint_sims(cs::Interfacer.CoupledSimulation)
                 output_dir,
             )
         end
-        if !isnothing(Checkpointer.get_model_cache(sim))
+        if !isnothing(Checkpointer.get_model_cache(sim)) && cs.save_cache
             Checkpointer.checkpoint_model_cache(
                 sim,
                 comms_ctx,

src/Interfacer.jl

@@ -60,6 +60,7 @@ struct CoupledSimulation{
     NTP <: NamedTuple,
     TP,
     DH,
+    SC <: Bool,
 }
     start_date::D
     fields::FV
@@ -73,6 +74,7 @@ struct CoupledSimulation{
     dir_paths::NTP
     thermo_params::TP
     diags_handler::DH
+    save_cache::SC
 end
 
 CoupledSimulation{FT}(args...) where {FT} = CoupledSimulation{FT, typeof.(args)...}(args...)

test/conservation_checker_tests.jl

@@ -85,6 +85,7 @@ for FT in (Float32, Float64)
             (;), # dir_paths
             nothing, # thermo_params
             nothing, # diags_handler
+            true, # save_cache
         )
 
         # set non-zero radiation and precipitation
@@ -187,6 +188,7 @@ for FT in (Float32, Float64)
             (;), # dir_paths
             nothing, # thermo_params
             nothing, # diags_handler
+            true, # save_cache
         )
 
         tot_energy, tot_water = ConservationChecker.check_conservation!(cs)

test/interfacer_tests.jl

@@ -58,6 +58,7 @@ for FT in (Float32, Float64)
             (;), # dir_paths
             nothing, # thermo_params
             nothing, # diags_handler
+            true, # save_cache
         )
         @test CC.Spaces.undertype(Interfacer.boundary_space(cs)) == FT
     end