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 cache or not 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,10 @@ 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, save_cache),
+    )
 
     callbacks = (checkpoint_cb,)
 

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

@@ -126,11 +126,13 @@ function restore_cache!(sim::Interfacer.ComponentModelSimulation, new_cache)
 end
 
 """
-    checkpoint_sims(cs::CoupledSimulation)
+    checkpoint_sims(cs::CoupledSimulation, save_cache)
 
-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. The argument `save_cache` determines whether to save
+the cache or not.
 """
-function checkpoint_sims(cs::Interfacer.CoupledSimulation)
+function checkpoint_sims(cs::Interfacer.CoupledSimulation, save_cache)
     time = Int(round(float(cs.t[])))
     day = floor(Int, time / (60 * 60 * 24))
     sec = floor(Int, time % (60 * 60 * 24))
@@ -148,7 +150,7 @@ function checkpoint_sims(cs::Interfacer.CoupledSimulation)
                 output_dir,
             )
         end
-        if !isnothing(Checkpointer.get_model_cache(sim))
+        if !isnothing(Checkpointer.get_model_cache(sim)) && save_cache
             Checkpointer.checkpoint_model_cache(
                 sim,
                 comms_ctx,