Support checkpointing of time-dependent adjoints - #71
Conversation
pyadjoint already drives checkpoint schedules from checkpoint_schedules; what was missing was the DOLFINx side of the contract. Blocks must not put non-overloaded values on the tape. Both solver blocks built their replay solution vector as a plain dolfinx.fem.Function, which became the block's output. Outside checkpointing nothing asks such a value to checkpoint itself, so this went unnoticed; under a schedule a stored output is re-stored on a later pass and it fails. Extracted the construction so the rule lives in one place. Added a disk backend for schedules that store on disk, written with h5py rather than taking on adios4dolfinx: these are snapshot checkpoints, valid only within the run that wrote them and against an unchanged partition, so the payload is just a process's local values with no mesh or permutation data. Ghost values are stored alongside the owned ones so that restoring needs no communication -- restores are filtered by a cache whose lifetime depends on when the garbage collector runs, which is not the same moment on every process, and a collective call on that path deadlocks. Checkpoint data stays in one file until teardown, because pyadjoint resets package data before recomputing but then restores an initial condition written while taping. enable_disk_checkpointing is the only name this adds; schedules and the timestepping loop stay pure pyadjoint. Tests compare the checkpointed gradient against the un-checkpointed one and run a Taylor test, over both file layouts, serially and on two processes. NonlinearProblem cannot yet be advanced over timesteps: its adjoint is wrong with checkpointing disabled too (Taylor rate -0.41 against 2 for the linear model), which is pre-existing and unrelated. Recorded as a strict xfail so it reports itself when fixed. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The disk backend needs h5py, and the CI image does not ship it: both the test module and the demo's disk section failed to import. Declared as a dependency rather than an extra, since the demo exercises it on every docs build. The lazy import stays as a safety net, and an h5py without MPI support still works -- each process then writes its own checkpoint file. Point the demo's API references at the packages they belong to via intersphinx, and cite the checkpointing literature with a per-document key prefix so the labels stay unique across pages. Refer to io4dolfinx rather than its former name. The demo now also turns disk checkpointing off when it is done, which is what deletes the checkpoint files.
| grad_ckpt = [np.copy(g.x.array) for g in rf_ckpt.derivative()] | ||
|
|
||
| assert np.isclose(J_plain, J_ckpt) | ||
| for a, e in zip(grad_ckpt, grad_plain): |
There was a problem hiding this comment.
Done, in both places in the demo, and in the two equivalent zips in tests/test_checkpointing.py for consistency.
| directions = [] | ||
| for k in range(num_steps): | ||
| h = dolfinx_adjoint.Function(V, name=f"direction_{k}") | ||
| # Interpolated rather than random: the direction has to be the same on every process, and | ||
| # per-process random numbers are not. | ||
| h.interpolate(lambda x, k=k: np.sin((k + 1) * np.pi * x[0]) * np.cos(np.pi * x[1])) | ||
| directions.append(h) |
There was a problem hiding this comment.
Shouldn't annotate=False be set in initialization of these functions, or use with pyadjoint.pause_annotating():.... ?
There was a problem hiding this comment.
Right, they were being taped. Wrapped in pyadjoint.stop_annotating() in both the demo and the test helper: the directions are inputs to the test, not part of the model.
Simplify the note on tape.timestepper, and say why Firedrake can pass a bare range: it sets tape.progress_bar, whose iter() returns a real iterator, while the default passes the argument straight through for next() to choke on. Use strict zips and np.testing.assert_allclose so a mismatch reports what differed rather than just failing, and rewrite the sentence introducing the Taylor test to say what it means.
|
|
||
| assert np.isclose(J_plain, J_disk) | ||
| for a, e in zip(grad_disk, grad_plain): | ||
| assert np.allclose(a, e) |
There was a problem hiding this comment.
Strict in zip and assert allclose instead of allclose.
There was a problem hiding this comment.
Both done in that block already: zip(..., strict=True) and np.testing.assert_allclose.
| Overloaded rather than plain, because whatever the block returns from | ||
| `recompute_component` becomes its output on the tape: under a checkpoint schedule a stored | ||
| output is asked to checkpoint itself again on a later pass, which a plain | ||
| `dolfinx.fem.Function` cannot do. Outside checkpointing nothing asks, which is why a plain |
There was a problem hiding this comment.
Use {py:class}dolfinx.fem.Function here
| """ | ||
| with stop_annotating(): | ||
| if isinstance(u, dolfinx.fem.Function): | ||
| return Function(u.function_space, name=u.name + "_initial_guess") |
There was a problem hiding this comment.
Function is imported as _Function in line 10.
There was a problem hiding this comment.
Good spot — switched the isinstance check to _Function to match the annotation.
|
|
||
| @no_annotations | ||
| def _ad_create_checkpoint(self): | ||
| from ..checkpointing import maybe_disk_checkpoint |
There was a problem hiding this comment.
Can import be moved to the top of the file? Is there any reason not to do it?
There was a problem hiding this comment.
Moved to module scope. It was local to avoid a cycle, but checkpointing.py only needs Function for annotations now (under typing.TYPE_CHECKING), so there is no runtime cycle left.
| return checkpoint | ||
|
|
||
| def _ad_restore_at_checkpoint(self, checkpoint): | ||
| from ..checkpointing import SnapshotCheckpoint |
There was a problem hiding this comment.
Same comment regarding import. Why not at top of file?
There was a problem hiding this comment.
Same — moved to module scope.
zip(..., strict=True) needs 3.10, so declare the floor rather than leave it implicit. assign_linear_combination reached for .function_space and .x on whatever extract_linear_combination returned, which UFL types as BaseCoefficient. Newer UFL types that tightly enough for mypy to reject it, failing the formatting job on main since before this branch. Assigning a linear combination genuinely needs the DOLFINx Function that carries the degrees of freedom, so check for one and say so, rather than reaching for an attribute UFL does not promise.
| @@ -0,0 +1,139 @@ | |||
| # # Time-distributed control | |||
| # Based on example from https://dolfin-adjoint.github.io/dolfin-adjoint/documentation/time-distributed-control/time-distributed-control.html | |||
There was a problem hiding this comment.
Why do you need this demo, and the time_distributed_control checkpointing demo. Which should be kept and why?
There was a problem hiding this comment.
This remove pushed by accident. I have now removed it
There was a problem hiding this comment.
Understood — I picked up your removal when I merged main into the branch, so the file is gone here too. Nothing outstanding on this one.
| stored alongside the owned ones, which keeps restoring free of communication -- see `_layout`. | ||
|
|
||
| Snapshot checkpoints are therefore not portable. They cannot be reopened by a later run, or on a | ||
| different number of processes. For a checkpoint that outlives the run, use ``io4dolfinx``. |
There was a problem hiding this comment.
Done — {py:mod}io4dolfinx`` in the module docstring, and the demo links to the repository.
|
|
||
|
|
||
| def _import_h5py(): | ||
| try: |
There was a problem hiding this comment.
Add note regarding lazy import in python 3.15 https://docs.python.org/3.15/reference/simple_stmts.html#lazy-imports
There was a problem hiding this comment.
h5py is a declared dependency now, so the lazy import has gone entirely and it is imported at module scope.
| if not shared_file: | ||
| return n_local, n_local, 0 | ||
| # Collective, but called only from the write path, which every process reaches together. | ||
| sizes = comm.allgather(n_local) |
There was a problem hiding this comment.
Replace with exscan.
As the following example shows:
from mpi4py import MPI
import time
import numpy as np
comm = MPI.COMM_WORLD
rng = np.random.default_rng(seed=42)
z = np.random.randint(0, 100, size=1)[0]
start_time = time.perf_counter()
vals = comm.allgather(z)
offset = np.sum(vals[: comm.rank])
end_time = time.perf_counter()
start_ex = time.perf_counter()
offset_ex = comm.exscan(z, op=MPI.SUM)
if comm.rank == 0:
offset_ex = 0
end_ex = time.perf_counter()
np.testing.assert_array_equal(offset, offset_ex)
print(f"Rank {comm.rank}: offset = {offset}, time taken = {end_time - start_time:.6f} seconds")
print(f"Rank {comm.rank}: offset_ex = {offset_ex}, time taken = {end_ex - start_ex:.6f} seconds")
print(
f"Ratio allreduce/exscan: {end_time - start_time:.6f} / {end_ex - start_ex:.6f} = {(end_time - start_time) / (end_ex - start_ex):.2f}"
)root@docker-desktop:~/shared# mpirun --allow-run-as-root -n 8 python3 mwe_exscan.py
Rank 3: offset = 201, time taken = 0.001738 seconds
Rank 3: offset_ex = 201, time taken = 0.000080 seconds
Ratio allreduce/exscan: 0.001738 / 0.000080 = 21.78
Rank 1: offset = 95, time taken = 0.000637 seconds
Rank 1: offset_ex = 95, time taken = 0.000075 seconds
Ratio allreduce/exscan: 0.000637 / 0.000075 = 8.47
Rank 6: offset = 330, time taken = 0.001919 seconds
Rank 6: offset_ex = 330, time taken = 0.000077 seconds
Ratio allreduce/exscan: 0.001919 / 0.000077 = 24.76
Rank 2: offset = 145, time taken = 0.001058 seconds
Rank 2: offset_ex = 145, time taken = 0.000079 seconds
Ratio allreduce/exscan: 0.001058 / 0.000079 = 13.37
Rank 5: offset = 283, time taken = 0.001528 seconds
Rank 5: offset_ex = 283, time taken = 0.000081 seconds
Ratio allreduce/exscan: 0.001528 / 0.000081 = 18.83
Rank 4: offset = 254, time taken = 0.001921 seconds
Rank 4: offset_ex = 254, time taken = 0.000077 seconds
Ratio allreduce/exscan: 0.001921 / 0.000077 = 24.96
Rank 7: offset = 416, time taken = 0.000133 seconds
Rank 7: offset_ex = 416, time taken = 0.000075 seconds
Ratio allreduce/exscan: 0.000133 / 0.000075 = 1.77
Rank 0: offset = 0.0, time taken = 0.000638 seconds
Rank 0: offset_ex = 0, time taken = 0.000073 seconds
Ratio allreduce/exscan: 0.000638 / 0.000073 = 8.73exscan is always faster.
There was a problem hiding this comment.
Done — comm.exscan(n_local, op=MPI.SUM) with the rank-0 None handled, plus one allreduce for the total length, which the dataset has to be sized with. Thanks for the benchmark; the prefix-sum-over-allgather was me reaching for the first thing that worked rather than the operation this actually is.
| rolling to a new one when the tape resets, and tearing down. | ||
| """ | ||
|
|
||
| def __init__(self, path: str, comm: MPI.Comm, use_mpio: bool, cleanup: bool): |
There was a problem hiding this comment.
use pathlib.Path here as input instead.
There was a problem hiding this comment.
Done — pathlib.Path for the file path, the checkpointer directory, and the dirname argument.
| self.path = path | ||
| self.comm = comm |
There was a problem hiding this comment.
Store with _path and _comm and make @property decorators to fetch them.
same for shared_file below.
There was a problem hiding this comment.
Done — _path, _comm and _shared_file with @property accessors.
| rolling to a new one when the tape resets, and tearing down. | ||
| """ | ||
|
|
||
| def __init__(self, path: str, comm: MPI.Comm, use_mpio: bool, cleanup: bool): |
There was a problem hiding this comment.
Document what input args do. It is unclear what cleanup does.
There was a problem hiding this comment.
Done. cleanup now says what it controls: whether the files and the temporary directory are deleted on teardown, with the note that they are unreadable by a later run either way, so keeping them is only useful for debugging.
|
|
||
| __slots__ = ("_file", "_key", "_space", "_cls", "_n_local", "_offset", "_name", "_cache", "__weakref__") | ||
|
|
||
| def __init__(self, file: _CheckpointFile, key: str, function: dolfinx.fem.Function, n_local: int, offset: int): |
|
|
||
| __slots__ = ("_file", "_key", "_space", "_cls", "_n_local", "_offset", "_name", "_cache", "__weakref__") | ||
|
|
||
| def __init__(self, file: _CheckpointFile, key: str, function: dolfinx.fem.Function, n_local: int, offset: int): |
There was a problem hiding this comment.
Should the input be a dolfinx_adjoint.Function or the parent class dolfinx.Function?
There was a problem hiding this comment.
Changed to dolfinx_adjoint.Function — only overloaded types are ever checkpointed, so the parent class was too loose. It is imported under typing.TYPE_CHECKING since dolfinx_adjoint.types imports this module at runtime.
| class _DiskCheckpointer(TapePackageData): | ||
| """Tape-attached state owning the checkpoint files for one tape.""" | ||
|
|
||
| def __init__(self, directory: str, comm: MPI.Comm, use_mpio: bool, cleanup: bool, owns_directory: bool): |
| "use_mpio=True requires an MPI-enabled build of h5py. Use use_mpio=False to write " | ||
| "one checkpoint file per process instead." | ||
| ) | ||
| owns_directory = dirname is None |
There was a problem hiding this comment.
Here we need an all-gather to ensure that all processes owns directory or not, as it can cause a deadlock in _DiskCheckpointer at
def close(self) -> None:
"""Close the current file and remove the directory if this object created it."""
self._file.close()
self._storing = False
if self._owns_directory:
self._comm.Barrier()
if self._comm.rank == 0:
try:
os.rmdir(self._directory)
except OSError: # pragma: no cover - non-empty when cleanup was disabled
passif self._owns_directory is not synced across all processes.
There was a problem hiding this comment.
Good catch, and it is a real deadlock. Fixed: the processes now agree explicitly rather than each deciding for itself.
without_dirname = comm.allreduce(int(dirname is None), op=MPI.SUM)
if without_dirname not in (0, comm.size):
raise ValueError("dirname must be given on every process or on none of them, ...")
owns_directory = without_dirname == comm.sizeSo a mixed call fails immediately with a clear message instead of hanging in close().
| self._storing = False | ||
|
|
||
|
|
||
| def maybe_disk_checkpoint(function: dolfinx.fem.Function) -> typing.Optional[SnapshotCheckpoint]: |
There was a problem hiding this comment.
Is this the correct class, did you say that we always want to use dolfinx_adjoint.Function.
Furthermore, use SnapshotCheckpoint | None instead of typing.Optional.
There was a problem hiding this comment.
Changed to dolfinx_adjoint.Function, and to SnapshotCheckpoint | None. I have replaced typing.Optional throughout the module while I was there.
|
|
||
|
|
||
| def enable_disk_checkpointing( | ||
| dirname: typing.Optional[str] = None, |
The parallel test run hung intermittently -- about one run in three locally, and in CI it sat for 42 minutes before being cancelled. Two ranks were in different collectives at once: one inside dolfinx's mpi_jit, the other inside dolfinx.fem.petsc.LinearProblem.__del__, which destroys the KSP and matrices and is collective. Every problem.solve() creates a block owning its own LinearProblem, and pyadjoint blocks sit in reference cycles, so a discarded tape's solvers are freed by the cyclic garbage collector rather than by refcounting. That runs when a process crosses an allocation threshold, which is not the same moment on every process, so one process enters a collective destructor the others are not in. Checkpointing made it likely by recomputing the forward many times over. Collect deliberately at points every process reaches together -- around each test and before building a tape -- so those destructors stay in step. Twelve consecutive parallel runs of the full suite now pass where two in three hung before. The underlying hazard is architectural: a block should not own PETSc objects it cannot destroy deterministically. Sharing one solver across a problem's blocks is the real fix, and is the solver-reuse work already planned separately. Also stop allocating the adjoint right-hand side per evaluation and reuse it, which removes another per-call collective allocation from that path. Review comments: exclusive scan instead of gathering every size, pathlib for paths, private attributes behind properties, documented arguments, cross-reference roles, imports at module scope, dolfinx_adjoint.Function in annotations, and modern optional syntax. Whether the checkpoint directory was created here is now agreed across processes rather than assumed: disagreeing would deadlock teardown, which synchronises before removing it. Perturbation directions are built with annotation stopped; they are inputs to a test, not part of the model.
| """ | ||
| with stop_annotating(): | ||
| if isinstance(u, _Function): | ||
| return Function(u.function_space, name=u.name + "_initial_guess") |
There was a problem hiding this comment.
Shouldn't this and the next line then use _Function(...) as a constructor.
There was a problem hiding this comment.
No — and thank you for asking, because the asymmetry is deliberate and it read like an oversight, which means it needed saying in the code rather than in a thread.
_Function is dolfinx.fem.Function; Function is our overloaded subclass. The check has to be against the base class, because u may be either a plain function or an overloaded one and the overloaded type is a subclass of it. The construction has to be of the overloaded type, because what this returns is what recompute_component hands back, and that becomes the block's output on the tape. A plain dolfinx.fem.Function cannot checkpoint itself, so constructing one here reintroduces exactly the failure this helper exists to prevent — it is the bug this PR opened with.
I have narrowed the return annotation to the overloaded type and spelled the distinction out in the docstring, so the next reader gets the answer without having to ask.
Nonmatching interpolation landed on main while this branch was open, and the resulting conflict stopped GitHub computing a merge ref, which is why the pull-request workflows silently stopped running: only the push-triggered documentation build was left. Both sides added a demo to the table of contents; kept both, with the checkpointing demo next to the one it extends. The index_map property main added is annotated dolfinx.cpp.la.IndexMap, which does not exist -- the type is dolfinx.cpp.common.IndexMap, publicly dolfinx.common.IndexMap. mypy rejects the merged tree without this, so it is corrected here rather than left to fail the formatting job.
The isinstance check is against dolfinx's Function and the construction is of the overloaded one, which reads like an oversight when the two names are an underscore apart. It is not: what arrives is only known to be a dolfinx.fem.Function, because the overloaded type is a subclass and either may be passed, while what leaves ends up on the tape and so must be overloaded. Constructing the plain one would reintroduce the failure this helper exists to prevent. Return annotation narrowed to the overloaded type to match.
| if _checkpointer is None or not _checkpointer.storing: | ||
| return None | ||
| return _checkpointer.store(function) |
There was a problem hiding this comment.
Don't we need to add a
global _checkpointerhere?
There was a problem hiding this comment.
No — global is only needed to rebind a module-level name. maybe_disk_checkpoint only reads _checkpointer, and a read resolves through the module globals at call time, so it always sees the current value. The two functions that assign to it, enable_disk_checkpointing and disable_disk_checkpointing, do both declare global.
Adding it here would be harmless but misleading, since it would suggest this function reassigns the checkpointer when it does not.
Demonstrated against the built module:
before enable : None
after enable : _DiskCheckpointer
storing=False -> None
storing=True -> SnapshotCheckpoint
after disable : None -> maybe_disk_checkpoint: None
The disk tests are the standing proof of the same thing: they only pass because maybe_disk_checkpoint picks up the checkpointer that enable_disk_checkpointing installed.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
# Conflicts: # _config.yml # docs/bibliography.bib # src/dolfinx_adjoint/blocks/solvers.py
The demo solved into `u_0` while `u_0` was also a coefficient of `L`, which `LinearProblem` rejects, so `solve_heat` raised on the very first call and none of the demo ran -- including the disk-checkpointing section, the only documentation of `enable_disk_checkpointing`. It is executed by the docs build. Split the unknown from the previous state, as the sibling demo and the checkpointing tests already do, and update the state with an explicit tape-recorded assignment. Also return the `LinearProblem` so callers can keep it alive: without that it is collected when `solve_heat` returns, and replay rebuilds an equivalent one, warning as it goes. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
test_snes_time_loop_gradient_is_correct was marked xfail(strict=True) against a serial defect -- ufl.adjoint raising IndexError on an argument-less form -- that has since been fixed. It XPASSed, so pytest exited non-zero and CI was red. The marker also claimed NonlinearProblem could not be covered by the checkpointing tests at all. It can: the gradient now matches the uncheckpointed one and Taylor-tests at rate 2 under Revolve, for a constant diffusivity and for a solution-dependent one -- the latter being the case that puts the unknown into the block's own dependencies, which nothing exercised before. Both are now tested, and the docstring's dangling reference to a test_solution_dependent_jacobian_is_unsupported that was never written is gone. These run on one process only. On two the SNES time-loop replay is intermittently and structurally wrong (a functional several times too large, in 10-50% of runs). That defect is pre-existing -- it reproduces on main, more often than here -- and is not about checkpointing, since it needs no schedule. Skipping there keeps the mpirun -n 2 CI job honest rather than flaky; the defect is tracked in .scratch/snes-time-loop-parallel/, with the reproducer and what has been ruled out. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The active checkpointer was held in a module global as well as in the tape's package data, and the global drove the file's lifetime. Enabling disk checkpointing a second time closed and unlinked the HDF5 file that the first tape's live SnapshotCheckpoints still point into, so re-evaluating a reduced functional built on that tape died with a KeyError from h5py reading a closed handle. disable_disk_checkpointing() had the mirror-image problem: it popped the package key off whichever tape happened to be working rather than off the tape that registered the checkpointer, leaving the real owner holding a checkpointer with no file -- which still satisfies pyadjoint's "disk storage is configured" check, so the failure surfaced much later, at the first restore. The tape's package data is now the only place a checkpointer is held, and every lookup goes through the working tape. Enabling on a second tape leaves the first one's file open and readable. disable_disk_checkpointing takes an optional tape so a tape that is no longer current can still be torn down. Nothing closes a file implicitly, and that is deliberate rather than an oversight: closing a shared file is collective, so it cannot be driven by garbage collection. A tape whose checkpointing is never disabled keeps its files until the process exits, which the docstring now says. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The checkpoint file grew without bound: every evaluation of the reduced functional writes a fresh checkpoint for each state it recomputes and drops the previous one, but nothing unlinked the datasets those dropped checkpoints had been reading. Measured at 77 datasets and 80 KB after one evaluation of a six-step model, then +30 datasets and ~32 KB for each one after -- so an optimisation loop on a real mesh fills the disk, which is the opposite of what moving checkpoints off the heap is for. Each file now tracks its readers weakly, and Tape.reset -- the one hook pyadjoint calls on every process together, and so the only place a collective unlink from a shared file is safe -- drops the datasets none of them refer to any more. HDF5 does not shrink the file, but the freed space goes to its own allocator and later datasets reuse it. The same model now settles at 76 datasets and ~84 KB and stays there, measured over fifteen evaluations. Which checkpoints are dead is a question about each process's reference counts, which need not drop at the same moment, so a shared file takes the union across processes: a dataset goes only once no process can still read it, and every process drops exactly the same ones in the same order. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Review flagged that always going through _ad_create_checkpoint allocates a function on every recompute, where the previous in-place reuse did not, and that the reuse was safe with no schedule active. Both are true, but measurement says the regression is not worth acting on: over an evaluation and its adjoint of a 40-step heat equation, all 160 _ad_create_checkpoint calls together take 2.9 ms of 263 ms, and end-to-end this branch and main are indistinguishable. A second code path guarded on "no schedule is active" would buy about one percent in exchange for a correctness invariant about pyadjoint's internals, so it is written down here instead of being built. Also corrects the FunctionAssignBlock comment, which attributed the rule to ADR-0001. ADR-0001 says something narrower -- that a block must not put a non-overloaded value on the tape. The actual hazard is that TimeStep.checkpoint stores BlockVariable._checkpoint itself for a global dependency rather than a copy, and restore_from_checkpoint hands that same object back. Comment-only, so the mypy hook was skipped: it reports five errors in blocks/solvers.py that main reports identically (whole-project mypy is 10 errors in 4 files on both), and this commit changes no code. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
… pin Review read `requires-python = ">=3.12"` as an unexplained narrowing that contradicted the retained `typing_extensions; python_version < '3.11'` marker, and read the commit message that first added the floor as evidence it should be 3.10. The floor is right and the marker is what was stale. src/dolfinx_adjoint/typing_utils.py uses PEP 695 `type X[T] = ...` aliases, which are 3.12 syntax rather than a library feature, so on 3.10 or 3.11 the package raises SyntaxError on import. That file is on main and predates this branch, so the declaration was catching up with the code rather than restricting it. Written down in the file now, so it does not get corrected downwards again. typing_extensions is imported nowhere in the package, and under a 3.12 floor its `python_version < '3.11'` marker could never be satisfied anyway, so the dependency is removed. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2 participants
Adds DOLFINx-side support for checkpointing time-dependent adjoints.
pyadjointalready drives checkpoint schedules fromcheckpoint_schedules; this fills in the missing half of that contract so long time loops can trade memory for recomputation instead of keeping every state alive for the adjoint sweep.Usage is entirely through pyadjoint's own API —
tape.enable_checkpointing(...)andtape.timestepper(...). The one new name isenable_disk_checkpointing, for schedules that spill to disk.What changed
dolfinx.fem.Functiononto the tape. BothLinearProblemBlockandNonlinearProblemBlockbuilt their replay output as a non-overloadedFunction; under a schedule that value gets re-stored on a later pass and checkpointing fails. Fixed in one shared helper.h5py-based) for schedules that checkpoint to disk. These are run-local snapshot checkpoints, not portable artifacts (no mesh/partition data, unlikeadios4dolfinx), and are deleted at teardown. Ghost values are stored alongside owned ones so restoring needs no collective communication.use_mpiopicks one shared MPI-IO file vs. one file per process.Known limitation (pre-existing, not from this change)
NonlinearProblemcan't yet be advanced over timesteps — its adjoint is already wrong on unmodifiedmainwith checkpointing disabled (Taylor rate -0.41 vs. 2.0). Recorded as a strictxfail.Not in scope
Solver reuse across timesteps, and Hessians under checkpointing (blocked on an upstream pyadjoint change —
evaluate_tlm/evaluate_hessiandon't consult the checkpoint manager the wayevaluate_adjdoes).Verification
72 passed, 2 xfailed, serially and undermpirun -n 2.mypy/ruffclean. Tests check the checkpointed gradient against the un-checkpointed one and run a Taylor test, over both file layouts.🤖 Generated with Claude Code