Merge pull request #1397 from CliMA/js/checkpoints

delete previous checkpoint when saving a new one

Author: juliasloan25

NEWS.md

@@ -6,14 +6,21 @@ ClimaCoupler.jl Release Notes
 
 ### ClimaCoupler features
 
-#### Remove `dt_save_state_to_disk` and `dt_save_to_sol` options
+#### Remove intermediate checkpoints PR[#1397](https://github.com/CliMA/ClimaCoupler.jl/pull/1397)
+
+Throughout the simulation, the previous checkpoint is now deleted whenever a new
+one is saved. The most recent checkpoint will always be available, so restarting
+is still supported. The field `prev_checkpoint_t` in the CoupledSimulation object
+is used to remove intermediate checkpoints.
+
+#### Remove `dt_save_state_to_disk` and `dt_save_to_sol` options PR[#1394](https://github.com/CliMA/ClimaCoupler.jl/pull/1394)
 
 `dt_save_state_to_disk` was unused and is removed from all configs in this
 commit. Note that ClimaAtmos does have an option with this name, but we
 pass `checkpoint_dt` to it. `dt_save_to_sol` is also removed as an option,
 in favor of using our more robust checkpointing infrastructure via `checkpoint_dt`.
 
-#### Misc. interface cleanup
+#### Misc. interface cleanup PR[#1341](https://github.com/CliMA/ClimaCoupler.jl/pull/1341)
 
 Including:
 - Remove `ρ_sfc` from surface model caches

experiments/ClimaEarth/setup_run.jl

@@ -455,13 +455,15 @@ function CoupledSimulation(config_dict::AbstractDict)
     specifics, the callbacks, the directory paths, and diagnostics for AMIP simulations.
     =#
 
+    prev_checkpoint_t = Ref(-1) # no checkpoint taken yet
     cs = CoupledSimulation{FT}(
         Ref(start_date),
         coupler_fields,
         conservation_checks,
         [tspan[1], tspan[2]],
         Δt_cpl,
         Ref(tspan[1]),
+        prev_checkpoint_t,
         model_sims,
         callbacks,
         dir_paths,

experiments/ClimaEarth/test/component_model_tests/climaland_tests.jl

@@ -83,6 +83,7 @@ end
         tspan,
         dt,
         tspan[1],
+        Ref(-1), # prev_checkpoint_t
         model_sims,
         (;), # callbacks
         (;), # dirs

experiments/ClimaEarth/test/debug_plots_tests.jl

@@ -76,6 +76,7 @@ plot_field_names(sim::Interfacer.SurfaceStub) = (:stub_field,)
         (Int(0), Int(1)), # tspan
         Int(200), # Δt_cpl
         Ref(Int(0)), # t
+        Ref(-1), # prev_checkpoint_t
         model_sims, # model_sims
         (;), # callbacks
         (;), # dirs

experiments/ClimaEarth/test/restart.jl

@@ -46,7 +46,7 @@ four_steps["dt"] = "180secs"
 four_steps["dt_cpl"] = "180secs"
 four_steps["t_end"] = "720secs"
 four_steps["dt_rad"] = "180secs"
-four_steps["checkpoint_dt"] = "720secs"
+four_steps["checkpoint_dt"] = "360secs"
 four_steps["coupler_output_dir"] = tmpdir
 four_steps["job_id"] = "four_steps"
 

src/Checkpointer.jl

@@ -33,14 +33,24 @@ This is a template function that should be implemented for each component model.
 get_model_cache(sim::Interfacer.ComponentModelSimulation) = nothing
 
 """
-    checkpoint_model_state(sim::Interfacer.ComponentModelSimulation, comms_ctx::ClimaComms.AbstractCommsContext, t::Int; output_dir = "output")
+    checkpoint_model_state(
+        sim::Interfacer.ComponentModelSimulation,
+        comms_ctx::ClimaComms.AbstractCommsContext,
+        t::Int,
+        prev_checkpoint_t::Int;
+        output_dir = "output")
 
 Checkpoint the model state of a simulation to a HDF5 file at a given time, t (in seconds).
+
+If a previous checkpoint exists, it is removed. This is to avoid accumulating
+many checkpoint files in the output directory. A value of -1 for `prev_checkpoint_t`
+is used to indicate that there is no previous checkpoint to remove.
 """
 function checkpoint_model_state(
     sim::Interfacer.ComponentModelSimulation,
     comms_ctx::ClimaComms.AbstractCommsContext,
-    t::Int;
+    t::Int,
+    prev_checkpoint_t::Int;
     output_dir = "output",
 )
     Y = get_model_prog_state(sim)
@@ -52,23 +62,36 @@ function checkpoint_model_state(
     CC.InputOutput.HDF5.write_attribute(checkpoint_writer.file, "time", t)
     CC.InputOutput.write!(checkpoint_writer, Y, "model_state")
     Base.close(checkpoint_writer)
-    return nothing
 
+    # Remove previous checkpoint if it exists
+    prev_checkpoint_file = joinpath(output_dir, "checkpoint_$(nameof(sim))_$(prev_checkpoint_t).hdf5")
+    remove_checkpoint(prev_checkpoint_file, prev_checkpoint_t, comms_ctx)
+    return nothing
 end
 
 """
-    checkpoint_model_cache(sim::Interfacer.ComponentModelSimulation, comms_ctx::ClimaComms.AbstractCommsContext, t::Int; output_dir = "output")
+    checkpoint_model_cache(
+        sim::Interfacer.ComponentModelSimulation,
+        comms_ctx::ClimaComms.AbstractCommsContext,
+        t::Int,
+        prev_checkpoint_t::Int;
+        output_dir = "output")
 
 Checkpoint the model cache to N JLD2 files at a given time, t (in seconds),
 where N is the number of MPI ranks.
 
 Objects are saved to JLD2 files because caches are generally not ClimaCore
 objects (and ClimaCore.InputOutput can only save `Field`s or `FieldVector`s).
+
+If a previous checkpoint exists, it is removed. This is to avoid accumulating
+many checkpoint files in the output directory. A value of -1 for `prev_checkpoint_t`
+is used to indicate that there is no previous checkpoint to remove.
 """
 function checkpoint_model_cache(
     sim::Interfacer.ComponentModelSimulation,
     comms_ctx::ClimaComms.AbstractCommsContext,
-    t::Int;
+    t::Int,
+    prev_checkpoint_t::Int;
     output_dir = "output",
 )
     # Move p to CPU (because we cannot save CUArrays)
@@ -79,6 +102,10 @@ function checkpoint_model_cache(
     pid = ClimaComms.mypid(comms_ctx)
     output_file = joinpath(output_dir, "checkpoint_cache_$(pid)_$(nameof(sim))_$t.jld2")
     JLD2.jldsave(output_file, cache = p)
+
+    # Remove previous checkpoint if it exists
+    prev_checkpoint_file = joinpath(output_dir, "checkpoint_cache_$(pid)_$(nameof(sim))_$(prev_checkpoint_t).jld2")
+    remove_checkpoint(prev_checkpoint_file, prev_checkpoint_t, comms_ctx)
     return nothing
 end
 
@@ -104,21 +131,31 @@ function checkpoint_sims(cs::Interfacer.CoupledSimulation)
     day = floor(Int, time / (60 * 60 * 24))
     sec = floor(Int, time % (60 * 60 * 24))
     output_dir = cs.dirs.checkpoints
+    prev_checkpoint_t = cs.prev_checkpoint_t[]
     comms_ctx = ClimaComms.context(cs)
+
     for sim in cs.model_sims
         if !isnothing(Checkpointer.get_model_prog_state(sim))
-            Checkpointer.checkpoint_model_state(sim, comms_ctx, time; output_dir)
+            Checkpointer.checkpoint_model_state(sim, comms_ctx, time, prev_checkpoint_t; output_dir)
         end
         if !isnothing(Checkpointer.get_model_cache(sim))
-            Checkpointer.checkpoint_model_cache(sim, comms_ctx, time; output_dir)
+            Checkpointer.checkpoint_model_cache(sim, comms_ctx, time, prev_checkpoint_t; output_dir)
         end
     end
+
     # Checkpoint the Coupler fields
     pid = ClimaComms.mypid(comms_ctx)
     @info "Saving coupler fields to JLD2 on day $day second $sec"
     output_file = joinpath(output_dir, "checkpoint_coupler_fields_$(pid)_$time.jld2")
     # Adapt to Array move fields to the CPU
     JLD2.jldsave(output_file, coupler_fields = CC.Adapt.adapt(Array, cs.fields))
+
+    # Remove previous Coupler fields checkpoint if it exists
+    prev_checkpoint_file = joinpath(output_dir, "checkpoint_coupler_fields_$(pid)_$(prev_checkpoint_t).jld2")
+    remove_checkpoint(prev_checkpoint_file, prev_checkpoint_t, comms_ctx)
+
+    # Update previous checkpoint time stored in the coupled simulation
+    cs.prev_checkpoint_t[] = time
 end
 
 """
@@ -211,4 +248,19 @@ function t_start_from_checkpoint(checkpoint_dir)
     return parse(Int, match(restart_file_rx, latest_restart)[2])
 end
 
+"""
+    remove_checkpoint(prev_checkpoint_file, prev_checkpoint_t, comms_ctx)
+
+Delete the provided checkpoint file on the root process and print a helpful
+info message. This can be used to remove intermediate checkpoints, to prevent
+saving excessively large amounts of output.
+"""
+function remove_checkpoint(prev_checkpoint_file, prev_checkpoint_t, comms_ctx)
+    if ClimaComms.iamroot(comms_ctx) && prev_checkpoint_t != -1 && isfile(prev_checkpoint_file)
+        @info "Removing previous checkpoint file: $prev_checkpoint_file"
+        rm(prev_checkpoint_file)
+    end
+    return nothing
+end
+
 end # module

src/Interfacer.jl

@@ -46,13 +46,28 @@ abstract type AbstractSimulation{FT} end
     CoupledSimulation
 Stores information needed to run a simulation with the coupler.
 """
-struct CoupledSimulation{FT <: Real, D, FV, E, TS, DTI, TT, NTMS <: NamedTuple, CALLBACKS, NTP <: NamedTuple, TP, DH}
+struct CoupledSimulation{
+    FT <: Real,
+    D,
+    FV,
+    E,
+    TS,
+    DTI,
+    TT,
+    CTT,
+    NTMS <: NamedTuple,
+    CALLBACKS,
+    NTP <: NamedTuple,
+    TP,
+    DH,
+}
     start_date::D
     fields::FV
     conservation_checks::E
     tspan::TS
     Δt_cpl::DTI
     t::TT
+    prev_checkpoint_t::CTT
     model_sims::NTMS
     callbacks::CALLBACKS
     dirs::NTP
@@ -63,7 +78,7 @@ end
 CoupledSimulation{FT}(args...) where {FT} = CoupledSimulation{FT, typeof.(args)...}(args...)
 
 function Base.show(io::IO, sim::CoupledSimulation)
-    device_type = nameof(typeof(ClimaComms.device(sim.comms_ctx)))
+    device_type = nameof(typeof(ClimaComms.device(sim)))
     return print(
         io,
         "Coupled Simulation\n",

test/checkpointer_tests.jl

@@ -24,9 +24,10 @@ end
     comms_ctx = ClimaComms.context(ClimaComms.CPUSingleThreaded())
     boundary_space = CC.CommonSpaces.CubedSphereSpace(FT; comms_ctx, radius = FT(6371e3), n_quad_points = 4, h_elem = 4)
     t = 1
+    prev_checkpoint_t = -1
     # old sim run
     sim = DummySimulation(CC.Fields.FieldVector(T = ones(boundary_space)))
-    Checkpointer.checkpoint_model_state(sim, comms_ctx, t, output_dir = "test_checkpoint")
+    Checkpointer.checkpoint_model_state(sim, comms_ctx, t, prev_checkpoint_t, output_dir = "test_checkpoint")
 
     # new sim run
     sim_new = DummySimulation(CC.Fields.FieldVector(T = zeros(boundary_space)))

test/conservation_checker_tests.jl

@@ -64,6 +64,7 @@ for FT in (Float32, Float64)
             (Int(0), Int(1000)), # tspan
             Int(200), # Δt_cpl
             Ref(Int(0)), # t
+            Ref(-1), # prev_checkpoint_t
             model_sims, # model_sims
             (;), # callbacks
             (;), # dirs
@@ -152,6 +153,7 @@ for FT in (Float32, Float64)
             (Int(0), Int(1000)), # tspan
             Int(200), # Δt_cpl
             Ref(Int(0)), # t
+            Ref(-1), # prev_checkpoint_t
             model_sims, # model_sims
             (;), # callbacks
             (;), # dirs

test/field_exchanger_tests.jl

@@ -141,6 +141,7 @@ for FT in (Float32, Float64)
             (Int(0), Int(1000)), # tspan
             Int(200), # Δt_cpl
             Ref(Int(0)), # t
+            Ref(-1), # prev_checkpoint_t
             (;
                 ice_sim = DummyStub((; area_fraction = ice_d)),
                 ocean_sim = Interfacer.SurfaceStub((; area_fraction = ocean_d)),
@@ -382,7 +383,8 @@ for FT in (Float32, Float64)
             nothing, # conservation_checks
             nothing, # tspan
             nothing, # dt
-            nothing, # t_start
+            nothing, # t
+            Ref(-1), # prev_checkpoint_t
             model_sims,
             (;), # callbacks
             (;), # dirs