Checkpointer updates

.buildkite/pipeline.yml

@@ -1,5 +1,5 @@
 env:
-  JULIA_VERSION: "1.10.8"
+  JULIA_VERSION: "1.10.9"
   JULIA_MINOR_VERSION: "1.10"
   TARTARUS_HOME: "/storage5/buildkite-agent"
   JULIA_PKG_SERVER_REGISTRY_PREFERENCE: eager
@@ -209,7 +209,7 @@ steps:
       architecture: CPU
     retry:
       automatic:
-        - exit_status: 1      
+        - exit_status: 1
           limit: 1
     depends_on: "init"
 
@@ -224,7 +224,7 @@ steps:
       architecture: CPU
     retry:
       automatic:
-        - exit_status: 1      
+        - exit_status: 1
           limit: 1
     depends_on: "init"
 
@@ -337,7 +337,7 @@ steps:
 ##### Turbulence Closures
 #####
 
-            
+
   - label: "🎣 gpu turbulence closures"
     env:
       JULIA_DEPOT_PATH: "$TARTARUS_HOME/.julia-$BUILDKITE_BUILD_NUMBER"
@@ -683,7 +683,7 @@ steps:
       architecture: CPU
     retry:
       automatic:
-        - exit_status: 1      
+        - exit_status: 1
           limit: 1
     depends_on: "init"
 
@@ -699,7 +699,7 @@ steps:
       architecture: CPU
     retry:
       automatic:
-        - exit_status: 1      
+        - exit_status: 1
           limit: 1
     depends_on: "init"
 

Project.toml

@@ -1,7 +1,7 @@
 name = "Oceananigans"
 uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
 authors = ["Climate Modeling Alliance and contributors"]
-version = "0.96.2"
+version = "0.96.3"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/Models/Models.jl

@@ -20,7 +20,7 @@ using Oceananigans.Utils: Time
 
 import Oceananigans: initialize!
 import Oceananigans.Architectures: architecture
-import Oceananigans.TimeSteppers: reset!
+import Oceananigans.TimeSteppers: reset!, set_clock!
 import Oceananigans.Solvers: iteration
 
 # A prototype interface for AbstractModel.
@@ -116,6 +116,8 @@ const OceananigansModels = Union{HydrostaticFreeSurfaceModel,
                                  NonhydrostaticModel,
                                  ShallowWaterModel}
 
+set_clock!(model::OceananigansModels, new_clock) = set_clock!(model.clock, new_clock)
+
 """
     possible_field_time_series(model::HydrostaticFreeSurfaceModel)
 
@@ -128,10 +130,10 @@ function possible_field_time_series(model::OceananigansModels)
     # such as model.diffusivity_fields
     return tuple(model_fields, forcing)
 end
- 
-# Update _all_ `FieldTimeSeries`es in an `OceananigansModel`. 
+
+# Update _all_ `FieldTimeSeries`es in an `OceananigansModel`.
 # Extract `FieldTimeSeries` from all property names that might contain a `FieldTimeSeries`
-# Flatten the resulting tuple by extracting unique values and set! them to the 
+# Flatten the resulting tuple by extracting unique values and set! them to the
 # correct time range by looping over them
 function update_model_field_time_series!(model::OceananigansModels, clock::Clock)
     time = Time(clock.time)
@@ -146,7 +148,7 @@ function update_model_field_time_series!(model::OceananigansModels, clock::Clock
 
     return nothing
 end
-               
+
 import Oceananigans.TimeSteppers: reset!
 
 function reset!(model::OceananigansModels)
@@ -162,15 +164,15 @@ function reset!(model::OceananigansModels)
     for field in model.timestepper.Gⁿ
         fill!(field, 0)
     end
-    
+
     return nothing
 end
 
 # Check for NaNs in the first prognostic field (generalizes to prescribed velocities).
 function default_nan_checker(model::OceananigansModels)
     model_fields = prognostic_fields(model)
 
-    if isempty(model_fields) 
+    if isempty(model_fields)
         return nothing
     end
 
@@ -182,7 +184,7 @@ end
 
 using Oceananigans.Models.HydrostaticFreeSurfaceModels: OnlyParticleTrackingModel
 
-# Particle tracking models with prescribed velocities (and no tracers) 
+# Particle tracking models with prescribed velocities (and no tracers)
 # have no prognostic fields and no chance to producing a NaN.
 default_nan_checker(::OnlyParticleTrackingModel) = nothing
 

src/OutputWriters/checkpointer.jl

@@ -5,32 +5,60 @@ using Oceananigans: fields, prognostic_fields
 using Oceananigans.Fields: offset_data
 using Oceananigans.TimeSteppers: QuasiAdamsBashforth2TimeStepper
 
-import Oceananigans.Fields: set! 
+import Oceananigans.Fields: set!
 
-mutable struct Checkpointer{T, P} <: AbstractOutputWriter
+mutable struct Checkpointer{T} <: AbstractOutputWriter
     schedule :: T
     dir :: String
     prefix :: String
-    properties :: P
     overwrite_existing :: Bool
     verbose :: Bool
     cleanup :: Bool
 end
 
-function default_checkpointed_properties(model)
-    properties = [:grid, :particles, :clock, :timestepper]
-    #if has_ab2_timestepper(model)
-    #    push!(properties, :timestepper)
-    #end
-    return properties
-end
+default_checkpointed_properties(model) = [:grid, :particles, :clock, :timestepper]
 
 has_ab2_timestepper(model) = try
     model.timestepper isa QuasiAdamsBashforth2TimeStepper
 catch
     false
 end
 
+# Certain properties are required for `set!` to pickup from a checkpoint.
+function required_checkpointed_properties(model)
+    properties = [:grid, :particles, :clock]
+
+    if has_ab2_timestepper(model)
+        push!(properties, :timestepper)
+    end
+
+    return properties
+end
+
+function validate_checkpointed_properties(model, properties)
+    required_properties = required_checkpointed_properties(model)
+
+    for rp in required_properties
+        if rp ∉ properties
+            @warn "$rp is required for checkpointing. It will be added to checkpointed properties"
+            push!(properties, rp)
+        end
+    end
+
+    for p in properties
+        p isa Symbol || error("Property $p to be checkpointed must be a Symbol.")
+        p ∉ propertynames(model) && error("Cannot checkpoint $p, it is not a model property!")
+
+        if (p ∉ required_properties) && has_reference(Function, getproperty(model, p))
+            @warn "model.$p contains a function somewhere in its hierarchy and will not be checkpointed."
+            filter!(e -> e != p, properties)
+        end
+    end
+
+    return properties
+end
+
+
 """
     Checkpointer(model;
                  schedule,
@@ -79,36 +107,11 @@ function Checkpointer(model; schedule,
                       prefix = "checkpoint",
                       overwrite_existing = false,
                       verbose = false,
-                      cleanup = false,
-                      properties = default_checkpointed_properties(model))
-
-    # Certain properties are required for `set!` to pickup from a checkpoint.
-    required_properties = [:grid, :particles, :clock]
-
-    if has_ab2_timestepper(model)
-        push!(required_properties, :timestepper)
-    end
-
-    for rp in required_properties
-        if rp ∉ properties
-            @warn "$rp is required for checkpointing. It will be added to checkpointed properties"
-            push!(properties, rp)
-        end
-    end
-
-    for p in properties
-        p isa Symbol || error("Property $p to be checkpointed must be a Symbol.")
-        p ∉ propertynames(model) && error("Cannot checkpoint $p, it is not a model property!")
-
-        if (p ∉ required_properties) && has_reference(Function, getproperty(model, p))
-            @warn "model.$p contains a function somewhere in its hierarchy and will not be checkpointed."
-            filter!(e -> e != p, properties)
-        end
-    end
+                      cleanup = false)
 
     mkpath(dir)
 
-    return Checkpointer(schedule, dir, prefix, properties, overwrite_existing, verbose, cleanup)
+    return Checkpointer(schedule, dir, prefix, overwrite_existing, verbose, cleanup)
 end
 
 #####
@@ -177,29 +180,40 @@ end
 function write_output!(c::Checkpointer, model)
     filepath = checkpoint_path(model.clock.iteration, c)
     c.verbose && @info "Checkpointing to file $filepath..."
-    addr = checkpointer_address(model)
 
     t1 = time_ns()
 
-    jldopen(filepath, "w") do file
-        file["$addr/checkpointed_properties"] = c.properties
-        serializeproperties!(file, model, c.properties, addr)
-        model_fields = prognostic_fields(model)
-        field_names = keys(model_fields)
-        for name in field_names
-            full_address = "$addr/$name"
-            serializeproperty!(file, full_address, model_fields[name])
-        end
-    end
+    write_output!(c, model, filepath, "w")
 
     t2, sz = time_ns(), filesize(filepath)
+
     c.verbose && @info "Checkpointing done: time=$(prettytime((t2 - t1) * 1e-9)), size=$(pretty_filesize(sz))"
 
     c.cleanup && cleanup_checkpoints(c)
 
     return nothing
 end
 
+function write_output!(c, model, filepath::AbstractString, mode::AbstractString;
+                       properties = default_checkpointed_properties(model))
+    @show properties
+    @show model
+
+    properties = validate_checkpointed_properties(model, properties)
+    addr = checkpointer_address(model)
+
+    jldopen(filepath, mode) do file
+        file["$addr/checkpointed_properties"] =
+        serializeproperties!(file, model, properties, addr)
+        model_fields = prognostic_fields(model)
+        field_names = keys(model_fields)
+        for name in field_names
+            full_address = "$addr/$name"
+            serializeproperty!(file, full_address, model_fields[name])
+        end
+    end
+end
+
 function cleanup_checkpoints(checkpointer)
     filepaths = glob(checkpoint_superprefix(checkpointer.prefix) * "*.jld2", checkpointer.dir)
     latest_checkpoint_filepath = latest_checkpoint(checkpointer, filepaths)
@@ -212,13 +226,12 @@ end
 #####
 
 """
-    set!(model, filepath::AbstractString)
+    set!(model::AbstractModel, filepath::AbstractString)
 
 Set data in `model.velocities`, `model.tracers`, `model.timestepper.Gⁿ`, and
 `model.timestepper.G⁻` to checkpointed data stored at `filepath`.
 """
 function set!(model::AbstractModel, filepath::AbstractString)
-
     addr = checkpointer_address(model)
 
     jldopen(filepath, "r") do file
@@ -280,9 +293,8 @@ function set_time_stepper_tendencies!(timestepper, file, model_fields, addr)
     return nothing
 end
 
-# For self-starting timesteppers like RK3 we do nothing 
+# For self-starting timesteppers like RK3 we do nothing
 set_time_stepper!(timestepper, args...) = nothing
 
 set_time_stepper!(timestepper::QuasiAdamsBashforth2TimeStepper, args...) =
     set_time_stepper_tendencies!(timestepper, args...)
-