Checkpointer revival #326

ali-ramadhan · 2019-08-02T00:11:04Z

This PR reintroduces checkpointing. This time the checkpointer saves the bare essentials to a JLD2 file.

Checkpointer will not save forcing functions as before (see #141).

Need to test GPU checkpointing, checkpointing with LES closures (I assumed model.diffusivities does not need to be checkpointed), and checkpointing with fancier boundary conditions.

A major issue with this checkpointer is that it can't restore the largest models because it creates a Model and then copies the data from the checkpoint file into the model fields. But if the model is using up all memory, then there's no room for reading data from disk to memory.

It should feed the data directly through the Model constructor, but this would require refactoring some of the model and field code.

This can be addressed in this PR or in a future PR, although it's kind of useless because checkpointing becomes more important the larger the model/simulation...

Resolves #324

ali-ramadhan · 2019-08-02T00:13:19Z

I think output_writers.jl is getting pretty cluttered (almost 400 lines with functionality shared between output writers) so I might try to move them to individual files in a future PR.

glwagner · 2019-08-02T17:09:45Z

I’m happy with output_writers.jl; I think all the code in that file is appropriately related.

The checkpointer could provide its own constructor to avoid excess memory allocation.

For GPU problems I don’t think there is an issue: checkpointed arrays can b loaded into CPU memory rather than GPU memory, and then the data can be copied into the fields allocated by the model constructor. So at first glance the excess memory allocation does not seem like a major issue on modern CPUs.

I am particularly concerned about the maintainability of the checkpointer, since it will need to be updated every time a new feature is added. Let’s make sure the design is easy to maintain before merging.

glwagner

Some work to do I think but this is a good start.

glwagner · 2019-08-02T17:30:12Z

src/output_writers.jl

+
+function savesubfields!(file, model, name, flds=propertynames(getproperty(model, name)))
+    for f in flds
+        file["$name/$f"] = Array(getproperty(getproperty(model, name), f).data.parent)


Put an if statement that does not save a field if its type is Function?

Sounds like a good check. Won't be needed as long we maintain checkpointed_fieldsets but good check either way.

I think that we should allow users to supply the modelfields that are to be checkpointed. In that case, such a check will be important.

Continuing from my comment below, the if-statement can also emit a warning that field x will not be saved.

@ali-ramadhan was this resolved?

glwagner · 2019-08-02T17:31:14Z

src/output_writers.jl

+    output_frequency :: Int
+end
+
+function Checkpointer(; output_frequency, dir=".", prefix="checkpoint", force=false)


The kwarg 'force' is not used here. It should either be used, or removed from the function signature.

In the JLD2OutputWriter, the kwarg force indicates whether file creation should be 'forced' (it corresponds to the same keyword passed to mkpath.

glwagner · 2019-08-02T17:31:27Z

src/output_writers.jl

+checkpointed_fieldsets = [:velocities, :tracers, :G, :Gp]
+
+function write_output(model, c::Checkpointer)
+    @warn "Checkpointer will not save forcing functions, output writers, or diagnostics. They will need to be " *


Does it make sense to have a warning that is always printed? Perhaps is it simply better to document this aspect of the checkpointer?

Hmmm, might then be good to just print the warning if there is a non-zero forcing function, or an output writer or diagnostic included.

Then yeah warning may be removed if the checkpointer is well-documented.

Here is what I propose:

Allow users to set the modelfields that are to be checkpoints as an argument to the Checkpointer constructor

If one of those fields that the user has asked to be saved contains a function, emit a warning.

@ali-ramadhan was this resolved?

glwagner · 2019-08-02T17:32:10Z

src/output_writers.jl

+_arr(::GPU, a) = CuArray(a)
+
+function restore_from_checkpoint(filepath)
+    @warn "Checkpointer cannot restore forcing functions, output writers, or diagnostics. They will need to be " *


See above. What we need is documentation about proper checkpointing, which will include information about how to restore a model from a checkpoint that includes functions. In fact, we could even provide features in the checkpointer that streamline this process (by indicating parts of the model structure that are associated with functions, and asking the user to provide those functions during checkpoint restoration).

Yeah warning doesn't have be in restore_from_checkpoint. Passing forcing functions is a good idea.

glwagner · 2019-08-02T17:34:09Z

test/test_output_writers.jl

+    k1, k2 = round(Int, Nz/4), round(Int, 3Nz/4)
+    true_model.tracers.T.data[i1:i2, j1:j2, k1:k2] .+= 0.01
+
+    checkpointed_model = deepcopy(true_model)


Why is it necessary to copy the model before checkpointing?

It's creating a second model that will be checkpointed as opposed to true_model which isn't.

Probably over paranoid but I wanted the two models to be time-stepped separately.

Are you worried that checkpointing will modify model?

Since this behavior is undesirable / unexpected, perhaps it should simply be tested for, so that we can assume that the model is not modified.

glwagner · 2019-08-02T17:35:16Z

src/output_writers.jl

+end
+
+checkpointed_structs   = [:arch, :boundary_conditions, :grid, :clock, :eos, :constants, :closure]
+checkpointed_fieldsets = [:velocities, :tracers, :G, :Gp]


G and Gp are not fields of model.

True. This was before PR #325 was merged and included the tendencies with AdamsBashforthTimestepper. Will update.

@ali-ramadhan was this resolved?

glwagner · 2019-08-02T17:41:56Z

src/output_writers.jl

+_arr(::CPU, a) = a
+_arr(::GPU, a) = CuArray(a)
+
+function restore_from_checkpoint(filepath)


This function must have a ; kwargs... argument which is then merged with the retrieved kwargs from the checkpoint file before being passed to the model constructor. This is needed to restore models with forcing functions or non-default boundary conditions.

glwagner

Second review: add a field (perhaps checkpointed_fields) to the Checkpointer that allows the user to control which subfields of model are checkpointed. Use a kwarg to set the default subfields to current list. The subfields that contain Fields (like velocities) should be included the same list as subfields like constants (for simplicity, one list is best), and a function should be used to properly save a struct depending on whether it contains Fields, Functions, or neither.

glwagner · 2019-08-02T17:42:54Z

src/output_writers.jl

+        # Checkpointing model properties that we can just serialize.
+        [file["$p"] = getproperty(model, p) for p in checkpointed_structs]
+
+        # Checkpointing structs containing fields.


Perhaps this step can be combined with the first, along with an if statement that inspects the content of the struct to be saved and determines whether it contains Fields or not.

glwagner · 2019-08-02T17:44:05Z

src/output_writers.jl

+    return nothing
+end
+
+checkpointed_structs   = [:arch, :boundary_conditions, :grid, :clock, :eos, :constants, :closure]


These fields could be made properties of the Checkpointer. This will allow the user to add/remove fields from checkpointing, which may be generally useful in the future and is currently useful for omitting boundary conditions from checkpointing.

@ali-ramadhan was this resolved?

ali-ramadhan · 2019-08-05T13:44:49Z

I'll work on this until we're happy with the interace but will leave restoring massive models (that fill up memory) properly for another PR.

I’m happy with output_writers.jl; I think all the code in that file is appropriately related.

I still find it a little messy and hard to navigate. Maybe a good solution for now would be to reorganize the file into sections, and shared functions can be moved to the top.

I am particularly concerned about the maintainability of the checkpointer, since it will need to be updated every time a new feature is added. Let’s make sure the design is easy to maintain before merging.

Eventually the Model struct will become stable but until then we'll have to modify the checkpointer to account for that. But yeah, having something flexible and easy to modify is important.

add a field (perhaps checkpointed_fields) to the Checkpointer that allows the user to control which subfields of model are checkpointed.

Might make more sense. In general, the checkpointer should be flexible but you'd probably only need to change what's checkpointed for esoteric cases. E.g. if you remove something important, you'll

glwagner · 2019-08-05T20:50:50Z

Eventually the Model struct will become stable but until then we'll have to modify the checkpointer to account for that.

This is not the only issue. Another issue is if new subfields of model are added that contain non-checkpointable elements, like forcing functions. In that sense the Model struct will never be "stable" with respect to the types of data it references. For example, there could someday be a grid that contains forcing functions. Also, we should checkpoint boundary conditions by default. However, because boundary conditions can contain forcing functions, we need to be able to robustly/programmatically save only those boundary conditions that do not contain references to functions.

To address this, we need a generalized checkpointer that works in a wide range of circumstances, is relatively automated, modular, and can be user-modified.

I'm happy to help develop these features. One key is a function that can query a struct recursively to determine if it contains any references to Functions at any level. A second element is control to save an entire object in the case that the object contains no references to functions, versus a different but useful secondary behavior if the object does contain a reference to functions.

codecov · 2019-08-22T16:17:42Z

Codecov Report

Merging #326 into master will decrease coverage by 13.3%.
The diff coverage is 96.15%.

@@             Coverage Diff             @@
##           master     #326       +/-   ##
===========================================
- Coverage   74.74%   61.43%   -13.31%     
===========================================
  Files          22       24        +2     
  Lines        1176     1281      +105     
===========================================
- Hits          879      787       -92     
- Misses        297      494      +197

Impacted Files	Coverage Δ
src/Oceananigans.jl	`62.5% <ø> (-37.5%)`	⬇️
src/output_writers.jl	`46.97% <96.15%> (-17.77%)`	⬇️
src/poisson_solvers.jl	`40.65% <0%> (-56.97%)`	⬇️
src/utils.jl	`16.21% <0%> (-40.93%)`	⬇️
src/turbulence_closures/constant_smagorinsky.jl	`51.42% <0%> (-37.15%)`	⬇️
src/turbulence_closures/closure_operators.jl	`42.62% <0%> (-22.55%)`	⬇️
src/fields.jl	`42% <0%> (-20.03%)`	⬇️
src/boundary_conditions.jl	`59.15% <0%> (-10.85%)`	⬇️
... and 8 more

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 66670ad...40fd4b0. Read the comment docs.

codecov · 2019-08-22T16:17:42Z

Codecov Report

Merging #326 into master will increase coverage by 1.31%.
The diff coverage is 81.1%.

@@            Coverage Diff             @@
##           master     #326      +/-   ##
==========================================
+ Coverage   74.74%   76.06%   +1.31%     
==========================================
  Files          22       22              
  Lines        1176     1224      +48     
==========================================
+ Hits          879      931      +52     
+ Misses        297      293       -4

Impacted Files	Coverage Δ
src/Oceananigans.jl	`100% <ø> (ø)`	⬆️
src/output_writers.jl	`75.93% <81.1%> (+11.18%)`	⬆️

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 66670ad...4e72d35. Read the comment docs.

ali-ramadhan · 2019-08-22T21:00:03Z

I've changed things a bit in response to your suggestions.

Mainly, the structs and fieldsets to checkpoint can be specified through the Checkpointer constructor. If an incompatible struct is provided, e.g. :forcing, an exception is raised. Validation only occurs in the constructor so no need for if-statements and warnings anywhere else.

@glwagner Lemme know if this looks okay to merge.

glwagner · 2019-08-23T17:05:39Z

src/output_writers.jl

    mkpath(dir)
    return Checkpointer(dir, prefix, output_frequency)
 end

 function savesubfields!(file, model, name, flds=propertynames(getproperty(model, name)))
-    for f in flds
-        file["$name/$f"] = Array(getproperty(getproperty(model, name), f).data.parent)
+    if name ∉ (:forcing)


I'm a bit confused about how this function works:

why restrict this functionality to Field (as it looks to me), but provide a hook when name=:forcing?

can't we check the type of the object being saved? Other fields (like boundary conditions, or future fields that we add to model) may be functions.

I think it'd be better to use a function that dispatches on the type of the object being saved, eg

savefield(file, location, field) = file[location] = field savefield(file, location, field::AbstractArray) = file[location] = Array(field) savefield(file, location, field::Field) = file[location] = Array(field.data.parent) savefield(file, location, field::Function) = warn("Cannot checkpoint Function object into $location!")

Using such a function may be more robust and general.

This reverts commit cbd4a8d.

ali-ramadhan · 2019-08-26T17:02:53Z

Good suggestion. Actually, this should help fix one problem I've been having with JLD2 output: Julia ranges like grid.xC are serialized to the JLD2 file, and cannot be read outside of Julia.

You might want to serialize the grid when checkpointing to easily restore from a checkpoint file. But when saving the grid to a JLD2 output file, which may be read by a language other than Julia, the grid properties should be saved in a language-agnostic manner. Same for boundary conditions.

So I changed the way structs are saved to disk for both the JLD2 output writer and the checkpointer. It's all done recursively via multiple dispatch so it should be flexible enough to work for all current Model properties and it should accomodate future changes to Model with minor changes.

When saving stuff to disk like a JLD2 file, saveproperty! is used, which converts Julia objects to language-agnostic objects.

When checkpointing, serializeproperty! is used, which serializes objects, with fields and boundary conditions require special treatment.

We checkpoint structs that are important for timestepping. Diagnostics and output writers are not checkpointed, as they are not essential and can be added in any time after model constructions. But if one or more boundary conditions contain a function, model.boundary_conditions are not serialized and must be manually restored.

There is one mess bits associated with restoring from a checkpoint:

Fields cannot be passed to the Model constructor. When restoring fields we want to avoid loading a field from disk and allocating Model memory for it at the same time, as we won't be able to restore models whose memory footprint exceed ~half the CPU/GPU memory. Thus restoring fields is treated as a special case (see restore_fields!). It is done after model creation where fields are read from disk and used to fill up existing model fields. Unfortunately, model.timestepper doesn't fit the pattern and is treated as an extra special case.

@glwagner Let me know what you think. This PR has been open for a while so I'd like to merge it ASAP and work on more pressing issues.

Checkpointer revival Former-commit-id: 35a6a05

ali-ramadhan added 2 commits August 1, 2019 20:04

Checkpointer that saves the bare essentials using JLD2.

1212c82

Revive checkpointer test

d434e55

ali-ramadhan added the feature 🌟 Something new and shiny label Aug 2, 2019

ali-ramadhan added this to the v1.0 milestone Aug 2, 2019

ali-ramadhan self-assigned this Aug 2, 2019

ali-ramadhan requested a review from glwagner August 2, 2019 00:11

glwagner reviewed Aug 2, 2019

View reviewed changes

Merge branch 'master' into checkpointer

40fd4b0

ali-ramadhan added 7 commits August 22, 2019 12:52

Do not save forcing functions.

b47fda8

Move output utils to the bottom.

f7476b4

Validate frequency and interval.

171f582

Remove redundant warnings.

14c221b

User can select kwargs to pass to restored model (e.g. for forcing).

823d310

Clean up checkpointer utils.

7793b1e

Update checkpointer test.

02cb710

ali-ramadhan requested a review from glwagner August 22, 2019 21:00

Forgot to test checkpointer on GPU.

ddbc1d6

glwagner reviewed Aug 23, 2019

View reviewed changes

ali-ramadhan added 3 commits August 25, 2019 15:38

Switch to saveproperty! for JLD2OutputWriter.

8aeacbd

Analogous serializeproperty! methods for checkpointing.

0f933c8

Saving/serializing boundary conditions.

ebb0f34

ali-ramadhan added 7 commits August 25, 2019 18:30

Refactor Checkpointer.

6935cf6

Properly restore fields and boundary conditions.

c96175b

Index named tuples by name, rather than number/range.

cbd4a8d

hasfunction utility.

64db4c7

Revert "Index named tuples by name, rather than number/range."

c2c01ee

This reverts commit cbd4a8d.

Checkpointer will serialize bcs as a whole.

3552f77

Cleanup, comments, and update test.

89518cf

ali-ramadhan and others added 7 commits August 26, 2019 14:46

Fix typo.

cffc11d

Nuke HorizontalAverages and VerticalPlanes scratch structs.

bcfaff8

Simpler validate interval and frequency function

ea91e5f

Clean up checkpointer constructor

517b32b

adds array_refs property to checkpointer and generalizes hasfunction

898216b

gets has_reference function working with a little elbow grease

00b6eba

moves using Distributed to Oceananigans.jl from output_writers

4e72d35

ali-ramadhan merged commit 35a6a05 into master Aug 28, 2019

ali-ramadhan deleted the checkpointer branch August 28, 2019 00:02

arcavaliere pushed a commit to arcavaliere/Oceananigans.jl that referenced this pull request Nov 6, 2019

Merge pull request CliMA#326 from climate-machine/checkpointer

4540394

Checkpointer revival Former-commit-id: 35a6a05

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Checkpointer revival #326

Checkpointer revival #326

ali-ramadhan commented Aug 2, 2019

ali-ramadhan commented Aug 2, 2019

glwagner commented Aug 2, 2019 •

edited

Loading

glwagner left a comment

glwagner Aug 2, 2019

ali-ramadhan Aug 5, 2019

glwagner Aug 5, 2019

glwagner Aug 5, 2019

glwagner Aug 23, 2019

glwagner Aug 2, 2019

glwagner Aug 2, 2019

ali-ramadhan Aug 5, 2019

glwagner Aug 5, 2019

glwagner Aug 23, 2019

glwagner Aug 2, 2019

ali-ramadhan Aug 5, 2019

glwagner Aug 2, 2019

ali-ramadhan Aug 5, 2019

glwagner Aug 5, 2019

glwagner Aug 2, 2019

ali-ramadhan Aug 5, 2019 •

edited

Loading

glwagner Aug 23, 2019

glwagner Aug 2, 2019

glwagner left a comment •

edited

Loading

glwagner Aug 2, 2019

glwagner Aug 2, 2019

glwagner Aug 23, 2019

ali-ramadhan commented Aug 5, 2019

glwagner commented Aug 5, 2019

codecov bot commented Aug 22, 2019

codecov bot commented Aug 22, 2019 •

edited

Loading

ali-ramadhan commented Aug 22, 2019 •

edited

Loading

glwagner Aug 23, 2019 •

edited

Loading

ali-ramadhan commented Aug 26, 2019

Checkpointer revival #326

Checkpointer revival #326

Conversation

ali-ramadhan commented Aug 2, 2019

ali-ramadhan commented Aug 2, 2019

glwagner commented Aug 2, 2019 • edited Loading

glwagner left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ali-ramadhan Aug 5, 2019 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

glwagner left a comment • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ali-ramadhan commented Aug 5, 2019

glwagner commented Aug 5, 2019

codecov bot commented Aug 22, 2019

Codecov Report

codecov bot commented Aug 22, 2019 • edited Loading

Codecov Report

ali-ramadhan commented Aug 22, 2019 • edited Loading

glwagner Aug 23, 2019 • edited Loading

Choose a reason for hiding this comment

ali-ramadhan commented Aug 26, 2019

glwagner commented Aug 2, 2019 •

edited

Loading

ali-ramadhan Aug 5, 2019 •

edited

Loading

glwagner left a comment •

edited

Loading

codecov bot commented Aug 22, 2019 •

edited

Loading

ali-ramadhan commented Aug 22, 2019 •

edited

Loading

glwagner Aug 23, 2019 •

edited

Loading