move restore into Checkpointer

Author: juliasloan25

docs/src/checkpointer.md

@@ -12,7 +12,7 @@ Checkpoints are a mix of HDF5 and JLD2 files and are typically saved in a
 !!! warning "Known limitations"
 
     - The number of MPI processes has to remain the same across checkpoints
-    - Restart files are generally not portable across machines, julia versions, 
+    - Restart files are generally not portable across machines, julia versions,
       and package versions
     - Adding/changing new component models will probably require adding/changing code
 
@@ -56,7 +56,7 @@ If the model does not support directly reading a checkpoint, the `Checkpointer`
 module provides a straightforward way to add this feature.
 [`Checkpointer.restart!`](@ref) takes a coupled simulation, a `restart_dir`, and
 a `restart_t` and overwrites the content of the coupled simulation with what is
-in the checkpoint. 
+in the checkpoint.
 
 ## Developer notes
 
@@ -135,4 +135,5 @@ Types to watch for:
     ClimaCoupler.Checkpointer.restart!
     ClimaCoupler.Checkpointer.checkpoint_sims
     ClimaCoupler.Checkpointer.t_start_from_checkpoint
+    ClimaCoupler.Checkpointer.restore!
 ```

experiments/ClimaEarth/components/atmosphere/climaatmos.jl

@@ -20,7 +20,6 @@ if pkgversion(CA) < v"0.28.6"
     CC.Adapt.@adapt_structure CA.RRTMGPInterface.RRTMGPModel
 end
 
-include("../shared/restore.jl")
 
 ###
 ### Functions required by ClimaCoupler.jl for an AtmosModelSimulation
@@ -140,7 +139,7 @@ end
 
 function Checkpointer.restore_cache!(sim::ClimaAtmosSimulation, new_cache)
     comms_ctx = ClimaComms.context(sim.integrator.u.c)
-    restore!(
+    Checkpointer.restore!(
         Checkpointer.get_model_cache(sim),
         new_cache,
         comms_ctx;

experiments/ClimaEarth/components/land/climaland_bucket.jl

@@ -13,7 +13,6 @@ import ClimaCoupler: Checkpointer, FluxCalculator, Interfacer, FieldExchanger
 using NCDatasets
 include("climaland_helpers.jl")
 
-include("../shared/restore.jl")
 
 ###
 ### Functions required by ClimaCoupler.jl for a SurfaceModelSimulation
@@ -411,7 +410,7 @@ end
 function Checkpointer.restore_cache!(sim::BucketSimulation, new_cache)
     old_cache = Checkpointer.get_model_cache(sim)
     comms_ctx = ClimaComms.context(sim.model)
-    restore!(
+    Checkpointer.restore!(
         old_cache,
         new_cache,
         comms_ctx,

experiments/ClimaEarth/components/land/climaland_integrated.jl

@@ -533,7 +533,7 @@ end
 function Checkpointer.restore_cache!(sim::ClimaLandSimulation, new_cache)
     old_cache = Checkpointer.get_model_cache(sim)
     comms_ctx = ClimaComms.context(sim.model.soil)
-    restore!(
+    Checkpointer.restore!(
         old_cache,
         new_cache,
         comms_ctx,

experiments/ClimaEarth/components/shared/restore.jl

@@ -9,12 +9,14 @@ import ClimaComms
 import ClimaCore as CC
 import ClimaUtilities.Utils: sort_by_creation_time
 import ClimaUtilities.TimeManager: ITime, seconds
+import ClimaUtilities.TimeVaryingInputs: AbstractTimeVaryingInput
 import ..Interfacer
 import Dates
+import StaticArrays
 
 import JLD2
 
-export get_model_prog_state, checkpoint_model_state, checkpoint_sims
+export get_model_prog_state, checkpoint_model_state, checkpoint_sims, restore!
 
 """
     get_model_prog_state(sim::Interfacer.ComponentModelSimulation)
@@ -295,4 +297,175 @@ function remove_checkpoint(prev_checkpoint_file, prev_checkpoint_t, comms_ctx)
     return nothing
 end
 
+"""
+    restore!(v1, v2, comms_ctx; ignore)
+
+Recursively traverse `v1` and `v2`, setting each field of `v1` with the
+corresponding field in `v2`. In this, ignore all the properties that have name
+within the `ignore` iterable.
+
+This is intended to be used when restarting a simulation's cache object
+from a checkpoint.
+
+`ignore` is useful when there are stateful properties, such as live pointers.
+"""
+function restore!(v1::T1, v2::T2, comms_ctx; name = "", ignore) where {T1, T2}
+    # We pick fieldnames(T2) because v2 tend to be simpler (Array as opposed
+    # to CuArray)
+    fields = filter(x -> !(x in ignore), fieldnames(T2))
+    # If there are no fields to restore, we check for consistency
+    if isempty(fields)
+        v1 == v2 || error("$v1 != $v2")
+    else
+        # Recursive case: restore each field
+        for p in fields
+            restore!(
+                getfield(v1, p),
+                getfield(v2, p),
+                comms_ctx;
+                name = "$(name).$(p)",
+                ignore,
+            )
+        end
+    end
+    return nothing
+end
+
+"""
+    restore!(
+        v1::Union{
+            AbstractTimeVaryingInput,
+            ClimaComms.AbstractCommsContext,
+            ClimaComms.AbstractDevice,
+            UnionAll,
+            DataType,
+        },
+        v2::Union{
+            AbstractTimeVaryingInput,
+            ClimaComms.AbstractCommsContext,
+            ClimaComms.AbstractDevice,
+            UnionAll,
+            DataType,
+        },
+        _comms_ctx;
+        name,
+        ignore,
+    )
+
+Ignore certain types that don't need to be restored.
+`UnionAll` and `DataType` are infinitely recursive, so we also ignore those.
+"""
+function restore!(
+    v1::Union{
+        AbstractTimeVaryingInput,
+        ClimaComms.AbstractCommsContext,
+        ClimaComms.AbstractDevice,
+        UnionAll,
+        DataType,
+    },
+    v2::Union{
+        AbstractTimeVaryingInput,
+        ClimaComms.AbstractCommsContext,
+        ClimaComms.AbstractDevice,
+        UnionAll,
+        DataType,
+    },
+    _comms_ctx;
+    name,
+    ignore,
+)
+    return nothing
+end
+
+"""
+    restore!(
+        v1::Union{CC.DataLayouts.AbstractData, AbstractArray},
+        v2::Union{CC.DataLayouts.AbstractData, AbstractArray},
+        comms_ctx;
+        name,
+        ignore,
+    )
+
+For array-like objects, we move the original data (v2) to the
+device of the new data (v1). Then we copy the original data to
+the new object.
+"""
+function restore!(
+    v1::Union{CC.DataLayouts.AbstractData, AbstractArray},
+    v2::Union{CC.DataLayouts.AbstractData, AbstractArray},
+    comms_ctx;
+    name,
+    ignore,
+)
+    ArrayType =
+        parent(v1) isa Array ? Array : ClimaComms.array_type(ClimaComms.device(comms_ctx))
+    moved_to_device = ArrayType(parent(v2))
+
+    parent(v1) .= moved_to_device
+    return nothing
+end
+
+"""
+    restore!(
+        v1::Union{StaticArrays.StaticArray, Number, UnitRange, LinRange, Symbol},
+        v2::Union{StaticArrays.StaticArray, Number, UnitRange, LinRange, Symbol},
+        comms_ctx;
+        name,
+        ignore,
+    )
+
+Ensure that immutable objects have been initialized correctly,
+as they cannot be restored from a checkpoint.
+"""
+function restore!(
+    v1::Union{StaticArrays.StaticArray, Number, UnitRange, LinRange, Symbol},
+    v2::Union{StaticArrays.StaticArray, Number, UnitRange, LinRange, Symbol},
+    comms_ctx;
+    name,
+    ignore,
+)
+    v1 == v2 || error("$name is a immutable but it inconsistent ($(v1) != $(v2))")
+    return nothing
+end
+
+"""
+    restore!(v1::Dict, v2::Dict, comms_ctx; name, ignore)
+
+RRTMGP has some internal dictionaries, which we check for consistency.
+"""
+function restore!(v1::Dict, v2::Dict, comms_ctx; name, ignore)
+    v1 == v2 || error("$name is inconsistent")
+    return nothing
+end
+
+"""
+    restore!(
+        v1::T1,
+        v2::T2,
+        comms_ctx;
+        name,
+        ignore,
+    ) where {
+        T1 <: Union{Dates.DateTime, Dates.UTInstant, Dates.Millisecond},
+        T2 <: Union{Dates.DateTime, Dates.UTInstant, Dates.Millisecond},
+    }
+
+Special case to compare time-related types to allow different timestamps during restore.
+"""
+function restore!(
+    v1::T1,
+    v2::T2,
+    comms_ctx;
+    name,
+    ignore,
+) where {
+    T1 <: Union{Dates.DateTime, Dates.UTInstant, Dates.Millisecond},
+    T2 <: Union{Dates.DateTime, Dates.UTInstant, Dates.Millisecond},
+}
+    if v1 != v2
+        @warn "Time value differs in restart" field = name original = v2 new = v1
+    end
+    return nothing
+end
+
 end # module