Merge pull request #1548 from CliMA/js/zs/soft-restarts

add an option `restart_cache`

Author: juliasloan25

.buildkite/pipeline.yml

@@ -127,6 +127,15 @@ steps:
           slurm_gres: "gpu:1"
           slurm_mem: 32GB
 
+      - label: "GPU restarts"
+        command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/test/restart_state_only.jl"
+        env:
+          CLIMACOMMS_DEVICE: "CUDA"
+        agents:
+          slurm_ntasks: 1
+          slurm_gres: "gpu:1"
+          slurm_mem: 32GB
+
   - group: "Integration Tests"
     steps:
       # SLABPLANET EXPERIMENTS

NEWS.md

@@ -6,6 +6,11 @@ ClimaCoupler.jl Release Notes
 
 ### ClimaCoupler features
 
+#### Add an option `restart_cache` PR[#1548](https://github.com/CliMA/ClimaCoupler.jl/pull/1548)
+Adds an option so that if restart files are available, we can choose to
+restart the state only (`restart_cache` false), or to restart both the state
+and the cache (`restart_cache` true). `restart_cache` is true by default.
+
 #### Remove `FluxCalculator.surface_inputs` helper function PR[#1543](https://github.com/CliMA/ClimaCoupler.jl/pull/1543)
 We can simplify the flux calculation by calling `SF.ValuesOnly` directly.
 Since we now remap all quantities onto the boundary space when we compute

experiments/ClimaEarth/cli_options.jl

@@ -100,6 +100,10 @@ function argparse_settings()
         help = "Time in seconds rounded to the nearest index to use at `t_start` for restarted simulation [nothing (default)]"
         arg_type = Int
         default = nothing
+        "--restart_cache"
+        help = "Boolean flag indicating whether to read the cache from the restart file if available [`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/components/atmosphere/climaatmos.jl

@@ -152,6 +152,7 @@ function Checkpointer.restore_cache!(sim::ClimaAtmosSimulation, new_cache)
             :hyperdiffusion_ghost_buffer,
             :data_handler,
             :graph_context,
+            :dt,
         ]),
     )
     return nothing

experiments/ClimaEarth/setup_run.jl

@@ -127,6 +127,7 @@ function CoupledSimulation(config_dict::AbstractDict)
         detect_restart_files,
         restart_dir,
         restart_t,
+        restart_cache,
         use_land_diagnostics,
         diagnostics_dt,
         evolving_ocean,
@@ -524,12 +525,19 @@ function CoupledSimulation(config_dict::AbstractDict)
 
     #=
     ## Restart component model states if specified
-    If a restart directory is specified and contains output files from the `checkpoint_cb` callback, the component model states are restarted from those files. The restart directory
-    is specified in the `config_dict` dictionary. The `restart_t` field specifies the time step at which the restart is performed.
+    If a restart directory is specified and contains output files from the `checkpoint_cb` callback,
+    the component model states are restarted from those files. The restart directory is specified in
+    the `config_dict` dictionary. The `restart_t` field specifies the time step at which the restart
+    is performed.
+
+    If `restart_cache` is true, the caches will be read from the restart file using `restore_cache!`.
+    Otherwise, the caches will be initialized in each component model's constructor.
+    When the caches are not read from the restart file, we have to perform the initial component
+    model exchange so that `set_caches!` can be called to initialize the caches.
     =#
-    should_restart && Checkpointer.restart!(cs, restart_dir, restart_t)
+    should_restart && Checkpointer.restart!(cs, restart_dir, restart_t, restart_cache)
 
-    if !should_restart
+    if !should_restart || !restart_cache
         #=
         ## Initialize Component Model Exchange
 

experiments/ClimaEarth/test/restart_state_only.jl

@@ -0,0 +1,70 @@
+# This test runs a small AMIP simulation twice times.
+#
+# - The first time the simulation is run for two steps
+# - The second time the simulation is run for two steps, but restarting from the
+#   first simulation
+#
+# Since the caches are not read from the restart file, the results will not be
+# bit-wise identical to running the simulation without restarting. This test
+# only checks that a simulation can be restarted from the state only without erroring.
+#
+# The content of the simulation is not the most important, but it helps if it
+# has all of the complexity possible.
+
+import ClimaComms
+ClimaComms.@import_required_backends
+import ClimaUtilities.OutputPathGenerator: maybe_wait_filesystem
+import YAML
+import Logging
+using Test
+
+# Uncomment the following for cleaner output (but more difficult debugging)
+# Logging.disable_logging(Logging.Warn)
+
+include("compare.jl")
+include("../setup_run.jl")
+
+comms_ctx = ClimaComms.context()
+@info "Context: $(comms_ctx)"
+ClimaComms.init(comms_ctx)
+
+# Make sure that all MPI processes agree on the output_loc
+tmpdir = ClimaComms.iamroot(comms_ctx) ? mktempdir(pwd()) : ""
+tmpdir = ClimaComms.bcast(comms_ctx, tmpdir)
+# Sometimes the shared filesystem doesn't work properly and the folder is not
+# synced across MPI processes. Let's add an additional check here.
+maybe_wait_filesystem(ClimaComms.context(), tmpdir)
+
+# Parse the input config file as a dictionary
+config_file = joinpath(@__DIR__, "amip_test.yml")
+config_dict = get_coupler_config_dict(config_file)
+
+# Four steps
+two_steps = deepcopy(config_dict)
+
+two_steps["dt"] = "180secs"
+two_steps["dt_cpl"] = "180secs"
+two_steps["t_end"] = "360secs"
+two_steps["dt_rad"] = "180secs"
+two_steps["checkpoint_dt"] = "360secs"
+two_steps["coupler_output_dir"] = tmpdir
+two_steps["job_id"] = "two_steps"
+
+println("Simulating two steps")
+cs_two_steps = setup_and_run(two_steps)
+
+# 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)
+
+two_steps_reading["t_end"] = "540secs"
+two_steps_reading["detect_restart_files"] = true
+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"
+
+cs_two_steps_reading = setup_and_run(two_steps_reading)
+@testset "Restarts from command line arguments" begin
+    @test cs_two_steps_reading.tspan[1] == cs_two_steps.tspan[2]
+end

experiments/ClimaEarth/user_io/arg_parsing.jl

@@ -145,6 +145,7 @@ function get_coupler_args(config_dict::Dict)
     detect_restart_files = config_dict["detect_restart_files"]
     restart_dir = config_dict["restart_dir"]
     restart_t = config_dict["restart_t"]
+    restart_cache = config_dict["restart_cache"]
 
     # Diagnostics information
     use_coupler_diagnostics = config_dict["use_coupler_diagnostics"]
@@ -188,6 +189,7 @@ function get_coupler_args(config_dict::Dict)
         detect_restart_files,
         restart_dir,
         restart_t,
+        restart_cache,
         use_coupler_diagnostics,
         diagnostics_dt,
         evolving_ocean,

src/Checkpointer.jl

@@ -176,13 +176,16 @@ function checkpoint_sims(cs::Interfacer.CoupledSimulation)
 end
 
 """
-    restart!(cs::CoupledSimulation, checkpoint_dir, checkpoint_t)
+    restart!(cs::CoupledSimulation, checkpoint_dir, checkpoint_t, restart_cache)
 
 Overwrite the content of `cs` with checkpoints in `checkpoint_dir` at time `checkpoint_t`.
 
+If `restart_cache` is true, the cache will be read from the restart file using `restore_cache!`.
+Otherwise, the cache will be left unchanged.
+
 Return a true if the simulation was restarted.
 """
-function restart!(cs, checkpoint_dir, checkpoint_t)
+function restart!(cs, checkpoint_dir, checkpoint_t, restart_cache)
     @info "Restarting from time $(checkpoint_t) and directory $(checkpoint_dir)"
     pid = ClimaComms.mypid(ClimaComms.context(cs))
     for sim in cs.model_sims
@@ -194,7 +197,7 @@ function restart!(cs, checkpoint_dir, checkpoint_t)
                 )
             restart_model_state!(sim, input_file_state, ClimaComms.context(cs))
         end
-        if !isnothing(Checkpointer.get_model_cache(sim))
+        if !isnothing(Checkpointer.get_model_cache(sim)) && restart_cache
             input_file_cache = joinpath(
                 checkpoint_dir,
                 "checkpoint_cache_$(pid)_$(nameof(sim))_$(checkpoint_t).jld2",