From 3fa923c7b1bcf6bb1b06f82464529f3a54389dca Mon Sep 17 00:00:00 2001 From: Henrik Finsberg Date: Tue, 25 Aug 2026 12:02:05 +0200 Subject: [PATCH 1/9] Support checkpointing of time-dependent adjoints 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 --- _toc.yml | 1 + .../time_distributed_control_checkpointing.py | 183 ++++++++++ docs/bibliography.bib | 36 ++ src/dolfinx_adjoint/__init__.py | 2 + src/dolfinx_adjoint/blocks/solvers.py | 30 +- src/dolfinx_adjoint/checkpointing.py | 338 ++++++++++++++++++ src/dolfinx_adjoint/types/function.py | 13 + tests/test_checkpointing.py | 291 +++++++++++++++ 8 files changed, 884 insertions(+), 10 deletions(-) create mode 100644 demos/time_distributed_control_checkpointing.py create mode 100644 src/dolfinx_adjoint/checkpointing.py create mode 100644 tests/test_checkpointing.py diff --git a/_toc.yml b/_toc.yml index 6c632a1..dbe1b93 100644 --- a/_toc.yml +++ b/_toc.yml @@ -6,6 +6,7 @@ parts: chapters: - file: "demos/poisson_mother.py" - file: "demos/time_distributed_control.py" + - file: "demos/time_distributed_control_checkpointing.py" - caption: Python API chapters: - file: "docs/api" diff --git a/demos/time_distributed_control_checkpointing.py b/demos/time_distributed_control_checkpointing.py new file mode 100644 index 0000000..8181ffc --- /dev/null +++ b/demos/time_distributed_control_checkpointing.py @@ -0,0 +1,183 @@ +# # Time-distributed control with checkpointing +# +# This is the [time-distributed control](./time_distributed_control.py) demo again, with +# checkpointing switched on. +# +# Taping a time-dependent model keeps every intermediate state alive, because the adjoint +# sweep needs each of them on the way back. For a long simulation that is the thing that +# exhausts memory first. Checkpointing trades that memory for repeated work: only some states +# are kept, and the rest are recomputed from the nearest stored one when the adjoint asks for +# them. A schedule decides which to keep and when to recompute. The schedules come from +# `checkpoint_schedules` {cite}`Dolci2024`; for how step-based checkpointing combines with +# high-level algorithmic differentiation, see {cite}`Maddison2024`. +# +# Everything here comes from `pyadjoint` and `checkpoint_schedules` directly. The only thing +# `dolfinx_adjoint` adds is `enable_disk_checkpointing`, used at the end. + +from collections import OrderedDict + +from mpi4py import MPI + +import dolfinx +import numpy as np +import pyadjoint +import ufl +from checkpoint_schedules import Revolve + +import dolfinx_adjoint + +# ## Enabling a schedule +# +# A schedule has to be enabled on an empty tape, before anything is recorded, so that every +# timestep is treated the same way. `Revolve(num_steps, snapshots)` keeps at most `snapshots` +# states in memory and recomputes whatever else the adjoint needs. + +num_steps = 10 +mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 8, 8) +V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) # type: ignore[arg-type] + +nu = dolfinx.fem.Constant(mesh, np.float64(1e-5)) +dt = dolfinx.fem.Constant(mesh, dolfinx.default_scalar_type(0.1)) + +x = ufl.SpatialCoordinate(mesh) + +petsc_options = { + "ksp_type": "preonly", + "pc_type": "lu", + "ksp_error_if_not_converged": True, +} + + +def solve_heat(schedule=None, disk=False): + """Tape the heat equation over `num_steps` timesteps, optionally under a schedule. + + Returns the reduced functional and the controls, one control per timestep. + """ + tape = pyadjoint.Tape() + pyadjoint.set_working_tape(tape) + # Both of these configure how the tape stores state, so both have to happen before + # anything is recorded on it. + if disk: + dolfinx_adjoint.enable_disk_checkpointing() + if schedule is not None: + tape.enable_checkpointing(schedule) + + t = dolfinx_adjoint.Constant(mesh, dolfinx.default_scalar_type(0.0)) + t.name = "time" + d = 16 * x[0] * (x[0] - 1) * x[1] * (x[1] - 1) * ufl.sin(ufl.pi * t) + + ctrls = OrderedDict() + for i in range(num_steps): + ctrls[i] = dolfinx_adjoint.Function(V, name=f"control_{i}") + + u = ufl.TrialFunction(V) + v = ufl.TestFunction(V) + f = dolfinx_adjoint.Function(V, name="source") + u_0 = dolfinx_adjoint.Function(V, name="solution") + + F = ((u - u_0) / dt * v + nu * ufl.inner(ufl.grad(u), ufl.grad(v)) - f * v) * ufl.dx + a, L = ufl.system(F) + + mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim) + exterior_facets = dolfinx.mesh.exterior_facet_indices(mesh.topology) + exterior_dofs = dolfinx.fem.locate_dofs_topological(V, mesh.topology.dim - 1, exterior_facets) + bc = dolfinx.fem.dirichletbc(0.0, exterior_dofs, V) + + problem = dolfinx_adjoint.LinearProblem( + a, + L, + u=u_0, + bcs=[bc], + petsc_options=petsc_options, + adjoint_petsc_options=petsc_options, + ) + + j = 0.5 * float(dt) * dolfinx_adjoint.assemble_scalar((u_0 - d) ** 2 * ufl.dx) + + # `tape.timestepper` marks the tape timesteps the schedule reasons about. Note the + # `iter(...)`: `Tape.timestepper` calls `next()` on whatever it is given, so passing a bare + # `range` raises TypeError. Do not "simplify" it away. + for i in tape.timestepper(iter(range(num_steps))): + t_val = float(dt) * (i + 1) + dolfinx_adjoint.assign(t_val, t) + dolfinx_adjoint.assign(ctrls[i], f) + + problem.solve() + + weight = 0.5 if i == num_steps - 1 else 1.0 + j += weight * float(dt) * dolfinx_adjoint.assemble_scalar((u_0 - d) ** 2 * ufl.dx) + + controls = list(ctrls.values()) + rf = pyadjoint.ReducedFunctional(j, [pyadjoint.Control(c) for c in controls]) + return rf, controls + + +# ## Checkpointing does not change the answer +# +# A schedule only changes when state is stored and recomputed. The functional and its gradient +# are unchanged, which is worth checking explicitly the first time you enable one. + +rf_plain, controls_plain = solve_heat() +J_plain = rf_plain(controls_plain) +grad_plain = [np.copy(g.x.array) for g in rf_plain.derivative()] + +rf_ckpt, controls_ckpt = solve_heat(Revolve(num_steps, 3)) +J_ckpt = rf_ckpt(controls_ckpt) +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): + assert np.allclose(a, e) + +if mesh.comm.rank == 0: + print(f"J without checkpointing: {J_plain:.12g}") + print(f"J with Revolve({num_steps}, 3): {J_ckpt:.12g}") + +# ## A Taylor test through the schedule +# +# Reproducing the un-checkpointed gradient shows the two agree, but not that either is right. +# A Taylor test does: the first-order remainder should converge at second order. + +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) + +rf_ckpt, controls_ckpt = solve_heat(Revolve(num_steps, 3)) +rate = pyadjoint.taylor_test(rf_ckpt, controls_ckpt, directions) +assert rate > 1.9 + +# ## Storing checkpoints on disk +# +# `Revolve` keeps its checkpoints in memory. When even those do not fit, a schedule can put +# them on disk instead, and `dolfinx_adjoint.enable_disk_checkpointing()` provides the storage. +# +# These are *snapshot* checkpoints: they hold just this process's values for the function, and +# assume the mesh and its partition are unchanged, so they are valid only within the run that +# wrote them. They are deleted automatically. For a checkpoint that outlives the run, or that +# can be read back on a different number of processes, use +# [adios4dolfinx](https://github.com/jorgensd/adios4dolfinx) instead. +# +# Like the schedule, it must be enabled before anything is recorded on the tape. + +from checkpoint_schedules import SingleDiskStorageSchedule # noqa: E402 + +rf_disk, controls_disk = solve_heat(SingleDiskStorageSchedule(), disk=True) +J_disk = rf_disk(controls_disk) +grad_disk = [np.copy(g.x.array) for g in rf_disk.derivative()] + +assert np.isclose(J_plain, J_disk) +for a, e in zip(grad_disk, grad_plain): + assert np.allclose(a, e) + +if mesh.comm.rank == 0: + print(f"J with checkpoints on disk: {J_disk:.12g}") + print("Gradients agree to machine precision in all three cases.") + +# ## References +# ```{bibliography} +# :filter: cited and ({"demos/time_distributed_control_checkpointing"} >= docnames) +# ``` diff --git a/docs/bibliography.bib b/docs/bibliography.bib index e74a9ab..03eabcc 100644 --- a/docs/bibliography.bib +++ b/docs/bibliography.bib @@ -18,3 +18,39 @@ @book{troltzsch2010optimal year={2010}, publisher={American Mathematical Soc.} } + + +@article{Maddison2024, + author = {Maddison, James R.}, + title = {Step-based checkpointing with high-level algorithmic differentiation}, + journal = {Journal of Computational Science}, + volume = {82}, + pages = {102405}, + year = {2024}, + doi = {10.1016/j.jocs.2024.102405} +} + +@article{Maddison2019, + author = {Maddison, James R. and Goldberg, D. N. and Goddard, B. D.}, + title = {Automated calculation of higher order partial differential equation + constrained derivative information}, + journal = {SIAM Journal on Scientific Computing}, + volume = {41}, + number = {5}, + pages = {C417--C445}, + year = {2019}, + doi = {10.1137/18M1209465} +} + +@article{Dolci2024, + author = {Dolci, Daiane I. and Maddison, James R. and Ham, David A. and + Pallez, Guillaume and Herrmann, Julien}, + title = {checkpoint\_schedules: schedules for incremental checkpointing of + adjoint simulations}, + journal = {Journal of Open Source Software}, + volume = {9}, + number = {94}, + pages = {6148}, + year = {2024}, + doi = {10.21105/joss.06148} +} diff --git a/src/dolfinx_adjoint/__init__.py b/src/dolfinx_adjoint/__init__.py index 29c77e2..3e1ff4d 100644 --- a/src/dolfinx_adjoint/__init__.py +++ b/src/dolfinx_adjoint/__init__.py @@ -6,6 +6,7 @@ import pyadjoint as _pyad from .assembly import assemble_scalar, error_norm +from .checkpointing import enable_disk_checkpointing from .function import assign from .interpolation import interpolate from .solvers import LinearProblem, NonlinearProblem @@ -30,6 +31,7 @@ "NonlinearProblem", "assemble_scalar", "assign", + "enable_disk_checkpointing", "error_norm", "__version__", "__author__", diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py index 1ff15be..8460833 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/src/dolfinx_adjoint/blocks/solvers.py @@ -8,12 +8,30 @@ import pyadjoint import ufl from dolfinx.fem.function import Function as _Function +from pyadjoint.tape import stop_annotating from ..petsc_utils import LinearAdjointProblem, solve_linear_problem from ..types import Function from .assembly import _create_vector, _SpecialVector, assemble_compiled_form +def _initial_guess_for( + u: _Function | typing.Sequence[_Function], +) -> _Function | typing.Sequence[_Function]: + """Build the solution vector a block solves into when its forward is replayed. + + 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 + one survived for so long. + """ + with stop_annotating(): + if isinstance(u, dolfinx.fem.Function): + return Function(u.function_space, name=u.name + "_initial_guess") + return [Function(ui.function_space, name=ui.name + "_initial_guess") for ui in u] + + class LinearProblemBlock(pyadjoint.Block): """A linear problem that can be used with adjoint methods. @@ -210,12 +228,8 @@ def _replace_coefficients_in_form(self, form: ufl.Form) -> ufl.Form: def prepare_recompute_component(self, inputs, relevant_outputs): """Prepare for recomputing the block with different control inputs.""" - # Create initial guess for the KSP solver # Form independnet compilation would make it possible to use the same KSP for all re-evaluations. - if isinstance(self._u, Function): - initial_guess = dolfinx.fem.Function(self._u.function_space, name=self._u.name + "_initial_guess") - else: - initial_guess = [dolfinx.fem.Function(u.function_space, name=u.name + "_initial_guess") for u in self._u] + initial_guess = _initial_guess_for(self._u) # Replace form coefficients with checkpointed values. # Loop through the dependencies of the lhs and rhs, check if they are in the respective form @@ -868,12 +882,8 @@ def _replace_coefficients_in_form(self, form: ufl.Form) -> ufl.Form: def prepare_recompute_component(self, inputs, relevant_outputs): """Prepare for recomputing the block with different control inputs.""" - # Create initial guess for the KSP solver # Form independnet compilation would make it possible to use the same KSP for all re-evaluations. - if isinstance(self._u, Function): - initial_guess = dolfinx.fem.Function(self._u.function_space, name=self._u.name + "_initial_guess") - else: - initial_guess = [dolfinx.fem.Function(u.function_space, name=u.name + "_initial_guess") for u in self._u] + initial_guess = _initial_guess_for(self._u) # Replace values in the DirichletBC if it is dependent on a control # NOTE: Currently assume that BCS are control independent. diff --git a/src/dolfinx_adjoint/checkpointing.py b/src/dolfinx_adjoint/checkpointing.py new file mode 100644 index 0000000..7a70125 --- /dev/null +++ b/src/dolfinx_adjoint/checkpointing.py @@ -0,0 +1,338 @@ +"""Snapshot checkpointing of DOLFINx functions to disk. + +A checkpoint schedule that uses :class:`checkpoint_schedules.StorageType.DISK` needs somewhere +to put a function's values. This module provides that as a *snapshot* checkpoint: it is written +and read within a single run, by the same processes, against an unchanged mesh and partition. +Under those assumptions the whole payload is a process's local values, so no mesh, geometry or +permutation data is stored and the file is a flat array per stored value. Ghost values are +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 ``adios4dolfinx``. +""" + +from __future__ import annotations + +import os +import tempfile +import typing +import weakref + +from mpi4py import MPI + +import dolfinx +import numpy as np +import pyadjoint.checkpointing +from pyadjoint.tape import TapePackageData, get_working_tape + +__all__ = ["enable_disk_checkpointing", "disable_disk_checkpointing", "SnapshotCheckpoint"] + +#: Key under which the disk checkpointer registers itself in ``Tape._package_data``. +_PACKAGE_KEY = "dolfinx_adjoint" + +#: The active checkpointer, or None when disk checkpointing is not enabled. +_checkpointer: typing.Optional["_DiskCheckpointer"] = None + +# Message pyadjoint shows when a schedule wants disk storage but none is configured. +pyadjoint.checkpointing.disk_checkpointing_callback[_PACKAGE_KEY] = ( + "Call dolfinx_adjoint.enable_disk_checkpointing() before enabling a schedule that uses disk storage." +) + + +def _import_h5py(): + try: + import h5py + except ImportError as e: # pragma: no cover - exercised only without h5py + raise ImportError("Disk checkpointing requires h5py. Install it with 'pip install h5py'.") from e + return h5py + + +def _layout(function: dolfinx.fem.Function, shared_file: bool, comm: MPI.Comm) -> tuple[int, int, int]: + """Describe where this process's values sit in a stored dataset. + + The whole local array is stored, ghost values included, not just the locally owned values. + Owned values alone would be smaller, but restoring them requires a forward scatter to + refill the ghosts, and that is collective. Restores are driven by whichever blocks happen + to need a value, and are additionally filtered by a cache whose lifetime depends on when + the garbage collector runs -- which is not the same moment on every process. A collective + call on that path deadlocks as soon as one process takes a cached value while another + reads. Storing the ghosts makes restoring purely local, so it cannot deadlock. + + Returns: + A tuple of the number of values this process stores, the length of the whole dataset, + and this process's offset into it. + """ + n_local = function.x.array.size + 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) + return n_local, sum(sizes), sum(sizes[: comm.rank]) + + +class _CheckpointFile: + """One HDF5 file holding snapshot checkpoints. + + The file is opened once and closed explicitly. It must not be closed from a finaliser: + with MPI-IO, opening and closing are collective, and Python's garbage collector does not + run at the same moment on every process, so a close driven by collection deadlocks. Every + call here therefore happens at a point all processes reach together -- creating the file, + 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): + h5py = _import_h5py() + self.path = path + self.comm = comm + # A shared file is a single file that every process writes its own slice of. Without + # MPI-IO each process gets a file to itself instead. + # One shared file that every process writes a slice of, or one file per process. + self.shared_file = use_mpio or comm.size == 1 + # Only a shared file is written by more than one process, so only then does deleting it + # belong to a single one of them. + self._deleted_by_this_process = cleanup and (comm.rank == 0 or not self.shared_file) + kwargs = {"driver": "mpio", "comm": comm} if use_mpio else {} + self._handle = h5py.File(path, "w", **kwargs) + self._next_index = 0 + self._closed = False + + def next_key(self) -> str: + """Return a dataset name that every process agrees on. + + Safe because checkpoints are taken in the same order on every process: pyadjoint holds + the checkpointable state in an insertion-ordered set, and all processes run the same + schedule. + """ + key = f"checkpoint_{self._next_index}" + self._next_index += 1 + return key + + def write(self, key: str, values: np.ndarray, n_global: int, offset: int) -> None: + dataset = self._handle.create_dataset(key, (n_global,), dtype=values.dtype) + dataset[offset : offset + values.size] = values + + def read(self, key: str, n_local: int, offset: int) -> np.ndarray: + return self._handle[key][offset : offset + n_local] + + def close(self) -> None: + """Close the file, deleting it unless it is being kept for inspection. + + Collective when the file was opened with MPI-IO, so every process must call it. + """ + if self._closed: + return + self._closed = True + self._handle.close() + if self._deleted_by_this_process: + try: + os.remove(self.path) + except OSError: # pragma: no cover - another process may have removed it first + pass + + +class SnapshotCheckpoint: + """A stored checkpoint, holding a reference to its data rather than the data itself. + + Returned by :meth:`Function._ad_create_checkpoint` while disk checkpointing is active, and + turned back into a function by :meth:`Function._ad_restore_at_checkpoint`. + """ + + __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): + # Holding the file keeps it alive for exactly as long as some checkpoint needs it. + self._file = file + self._key = key + self._space = function.function_space + self._cls = type(function) + self._n_local = n_local + self._offset = offset + self._name = function.name + # Weak, so that repeated restores during one block evaluation hand back the *same* + # object -- the blocks build replacement maps across several `saved_output` accesses + # and a fresh object each time makes those maps miss. Weak rather than strong so the + # values are released again once the block is done with them, which is the point of + # storing them on disk in the first place. + self._cache: typing.Optional[weakref.ReferenceType] = None + + def restore(self): + """Read the stored values back into a function of the original type.""" + from .types.function import Function + + if self._cache is not None: + cached = self._cache() + if cached is not None: + return cached + + # Mirrors Function._ad_new_like: going through __new__ preserves the concrete subclass + # (Constant takes a different constructor signature). + restored = self._cls.__new__(self._cls, self._space) + Function.__init__(restored, self._space) + restored.name = self._name + # Purely local: the stored array already includes the ghost values, so no scatter. + restored.x.array[:] = self._file.read(self._key, self._n_local, self._offset) + self._cache = weakref.ref(restored) + return restored + + +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): + self._directory = directory + self._comm = comm + self._use_mpio = use_mpio + self._cleanup = cleanup + self._owns_directory = owns_directory + self._generation = 0 + self._storing = False + self._file = self._roll_to_new_file() + + def _roll_to_new_file(self) -> _CheckpointFile: + # Reached by every process together (pyadjoint resets package data on all of them), so + # it is safe to close the superseded file here. + previous = getattr(self, "_file", None) + if previous is not None: + previous.close() + rank_suffix = "" if (self._use_mpio or self._comm.size == 1) else f"_rank{self._comm.rank}" + path = os.path.join(self._directory, f"checkpoint_{self._generation}{rank_suffix}.h5") + self._generation += 1 + return _CheckpointFile(path, self._comm, self._use_mpio, self._cleanup) + + @property + def storing(self) -> bool: + """Whether values should currently be written to disk rather than kept in memory.""" + return self._storing + + def store(self, function: dolfinx.fem.Function) -> SnapshotCheckpoint: + n_local, n_global, offset = _layout(function, self._file.shared_file, self._comm) + key = self._file.next_key() + self._file.write(key, function.x.array, n_global, offset) + return SnapshotCheckpoint(self._file, key, function, n_local, offset) + + # -- TapePackageData ------------------------------------------------------------------ + + def clear(self): + # The tape is being discarded, so no checkpoint taken so far can still be wanted. + self._file = self._roll_to_new_file() + + def reset(self): + # Deliberately not rolling to a new file. pyadjoint resets package data before + # recomputing the forward, but then restores the initial condition from a checkpoint + # written while taping, so data from before the reset is still live. Rolling the file + # here deletes it and the restore fails. Checkpoints therefore accumulate in one file + # for as long as disk checkpointing is enabled, and are removed together at teardown. + self._storing = False + + def checkpoint(self): + return self._file + + def restore_from_checkpoint(self, state): + self._file = state + + def copy(self): + other = _DiskCheckpointer.__new__(_DiskCheckpointer) + other.__dict__.update(self.__dict__) + return other + + 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 + pass + + def continue_checkpointing(self): + self._storing = True + + def pause_checkpointing(self): + self._storing = False + + +def maybe_disk_checkpoint(function: dolfinx.fem.Function) -> typing.Optional[SnapshotCheckpoint]: + """Store ``function`` on disk if disk checkpointing is active, otherwise return None. + + Returning None tells the caller to fall back to an in-memory copy. Disk storage is only + active inside the windows pyadjoint opens around writing checkpoint data, so most calls + return None even when disk checkpointing is enabled. + """ + if _checkpointer is None or not _checkpointer.storing: + return None + return _checkpointer.store(function) + + +def enable_disk_checkpointing( + dirname: typing.Optional[str] = None, + comm: typing.Optional[MPI.Comm] = None, + cleanup: bool = True, + use_mpio: typing.Optional[bool] = None, +) -> None: + """Store checkpoints on disk rather than in memory. + + Must be called before any operation is recorded on the working tape, and before enabling a + checkpoint schedule on it. + + Args: + dirname: Directory to hold the checkpoint files. A temporary directory is created if + this is not given. + comm: MPI communicator. Defaults to ``MPI.COMM_WORLD``. + cleanup: Whether to delete checkpoint files once nothing refers to them. Pass False to + keep them for inspection. + use_mpio: Whether to write one shared file with MPI-IO. The default chooses it when + running on more than one process with an MPI-enabled h5py, and falls back to one + file per process otherwise. Pass False to force the per-process layout. + """ + global _checkpointer + + if _checkpointer is not None: + # Enabling twice would otherwise strand the previous files, open and undeleted. + disable_disk_checkpointing() + + tape = get_working_tape() + if tape.get_blocks(): + raise RuntimeError( + "Disk checkpointing must be enabled before any blocks are added to the tape, " + "so that every checkpoint is stored the same way." + ) + + comm = MPI.COMM_WORLD if comm is None else comm + h5py = _import_h5py() + if use_mpio is None: + use_mpio = comm.size > 1 and h5py.get_config().mpi + elif use_mpio and not h5py.get_config().mpi: + raise RuntimeError( + "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 + if dirname is None: + # Every process must agree on the directory, even in the per-process layout. + created = tempfile.mkdtemp(prefix="dolfinx_adjoint_checkpoints_") if comm.rank == 0 else None + directory = typing.cast(str, comm.bcast(created, root=0)) + else: + directory = dirname + if comm.rank == 0: + os.makedirs(directory, exist_ok=True) + comm.Barrier() + + _checkpointer = _DiskCheckpointer(directory, comm, use_mpio, cleanup, owns_directory) + tape._package_data[_PACKAGE_KEY] = _checkpointer + + +def disable_disk_checkpointing() -> None: + """Stop storing checkpoints on disk and delete the checkpoint files. + + Collective: every process must call it, because closing a shared checkpoint file is. + """ + global _checkpointer + + tape = get_working_tape() + tape._package_data.pop(_PACKAGE_KEY, None) + if _checkpointer is not None: + _checkpointer.close() + _checkpointer = None diff --git a/src/dolfinx_adjoint/types/function.py b/src/dolfinx_adjoint/types/function.py index 5c34eaa..db03471 100644 --- a/src/dolfinx_adjoint/types/function.py +++ b/src/dolfinx_adjoint/types/function.py @@ -69,6 +69,15 @@ def _ad_init_object(cls, obj): @no_annotations def _ad_create_checkpoint(self): + from ..checkpointing import maybe_disk_checkpoint + + # While a schedule is storing to disk, hand back a reference to the stored values + # rather than the values themselves. This is the only seam pyadjoint offers for + # choosing where checkpoint data lives. + stored = maybe_disk_checkpoint(self) + if stored is not None: + return stored + # Note: self.copy() (dolfinx.fem.Function.copy) always returns a plain # dolfinx.fem.Function regardless of self's concrete type, so wrapping it with # create_overloaded_object would silently downcast a Constant checkpoint to a @@ -79,6 +88,10 @@ def _ad_create_checkpoint(self): return checkpoint def _ad_restore_at_checkpoint(self, checkpoint): + from ..checkpointing import SnapshotCheckpoint + + if isinstance(checkpoint, SnapshotCheckpoint): + return checkpoint.restore() return checkpoint def _ad_dot(self, other: typing.Self, options: typing.Optional[dict] = None): diff --git a/tests/test_checkpointing.py b/tests/test_checkpointing.py new file mode 100644 index 0000000..7c226a7 --- /dev/null +++ b/tests/test_checkpointing.py @@ -0,0 +1,291 @@ +"""Checkpointing of time-dependent adjoint computations. + +The forward model is a heat equation advanced over a number of tape timesteps, with one +control per timestep. Every test compares a checkpointed run against the same run with +checkpointing disabled: a checkpoint schedule only changes *when* forward state is stored +and recomputed, never the value of the derivative. +""" + +from mpi4py import MPI +from petsc4py import PETSc + +import dolfinx +import h5py +import numpy as np +import pyadjoint +import pytest +import ufl +from checkpoint_schedules import Revolve, SingleDiskStorageSchedule +from pyadjoint.checkpointing import CheckpointError + +import dolfinx_adjoint + +_PETSC_OPTIONS = { + "ksp_type": "preonly", + "pc_type": "lu", + "ksp_error_if_not_converged": True, +} + + +@pytest.fixture(autouse=True) +def isolated_tape(): + """Keep these tests from leaking tape state into the rest of the suite. + + Two things would otherwise escape. The working tape: pyadjoint's Tape.clear_tape() resets + the checkpoint manager but leaves `_eagerly_checkpoint_outputs` and `latest_checkpoint` set, + so a tape that has once been checkpointed keeps checkpointing outputs eagerly even after + being cleared. And the PETSc options database: LinearProblem writes its options there under + a fixed default prefix and does not remove them, so a later solver constructed without + explicit options silently inherits whatever these tests set. + """ + previous_tape = pyadjoint.get_working_tape() + previous_options = dict(PETSc.Options().getAll()) + try: + yield + finally: + dolfinx_adjoint.checkpointing.disable_disk_checkpointing() + pyadjoint.set_working_tape(previous_tape) + options = PETSc.Options() + for key in set(options.getAll()) - set(previous_options): + options.delValue(key) + + +def _perturbation_directions(V, n): + """Perturbation directions for a Taylor test. + + Built by interpolating analytic expressions rather than from random numbers: the + directions must agree across processes, and per-rank random values do not. + """ + directions = [] + for k in range(n): + h = dolfinx_adjoint.Function(V, name=f"direction_{k}") + h.interpolate(lambda x, k=k: np.sin((k + 1) * np.pi * x[0]) * np.cos(np.pi * x[1])) + directions.append(h) + return directions + + +def _tape_heat_equation(n_steps, schedule=None, disk=False, use_mpio=None): + """Tape a heat equation with one control per tape timestep. + + Args: + n_steps: Number of tape timesteps to advance. + schedule: A ``checkpoint_schedules`` schedule, or None to disable checkpointing. + disk: Whether to store checkpoints on disk. + use_mpio: Passed through to ``enable_disk_checkpointing``, selecting the file layout. + + Returns: + A tuple of the reduced functional, the controls, and perturbation directions. + """ + tape = pyadjoint.Tape() + pyadjoint.set_working_tape(tape) + # Both of these must happen before anything is recorded on this tape. + if disk: + dolfinx_adjoint.enable_disk_checkpointing(use_mpio=use_mpio) + if schedule is not None: + tape.enable_checkpointing(schedule) + + mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 8, 8) + V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) # type: ignore[arg-type] + + dt = 0.1 + nu = dolfinx.fem.Constant(mesh, dolfinx.default_scalar_type(1.0e-2)) + + controls = [] + for i in range(n_steps): + c = dolfinx_adjoint.Function(V, name=f"control_{i}") + c.interpolate(lambda x, i=i: 0.5 + 0.1 * (i + 1) * x[0]) + controls.append(c) + + u = ufl.TrialFunction(V) + v = ufl.TestFunction(V) + f = dolfinx_adjoint.Function(V, name="source") + uh = dolfinx_adjoint.Function(V, name="solution") + + F = ((u - uh) / dt * v + nu * ufl.inner(ufl.grad(u), ufl.grad(v)) - f * v) * ufl.dx + a, L = ufl.system(F) + + mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim) + boundary_facets = dolfinx.mesh.exterior_facet_indices(mesh.topology) + boundary_dofs = dolfinx.fem.locate_dofs_topological(V, mesh.topology.dim - 1, boundary_facets) + bc = dolfinx.fem.dirichletbc(0.0, boundary_dofs, V) + + problem = dolfinx_adjoint.LinearProblem( + a, + L, + u=uh, + bcs=[bc], + petsc_options=_PETSC_OPTIONS, + adjoint_petsc_options=_PETSC_OPTIONS, + ) + + J = dolfinx_adjoint.assemble_scalar(dt * uh**2 * ufl.dx) + # NOTE: `iter(...)` is required. Tape.timestepper calls next() on what it is given, + # so a bare range() raises TypeError even though pyadjoint's own docstring shows one. + for i in tape.timestepper(iter(range(n_steps))): + dolfinx_adjoint.assign(controls[i], f) + problem.solve() + J = J + dolfinx_adjoint.assemble_scalar(dt * uh**2 * ufl.dx) + + rf = pyadjoint.ReducedFunctional(J, [pyadjoint.Control(c) for c in controls]) + return rf, controls, _perturbation_directions(V, n_steps) + + +def _gradient(rf, controls): + rf(controls) + return [np.copy(g.x.array) for g in rf.derivative()] + + +@pytest.mark.parametrize("n_steps, snapshots", [(6, 2), (10, 3)]) +def test_gradient_matches_uncheckpointed(n_steps, snapshots): + """A checkpoint schedule does not change the gradient.""" + rf_plain, controls_plain, _ = _tape_heat_equation(n_steps) + expected = _gradient(rf_plain, controls_plain) + + rf_ckpt, controls_ckpt, _ = _tape_heat_equation(n_steps, Revolve(n_steps, snapshots)) + actual = _gradient(rf_ckpt, controls_ckpt) + + assert len(actual) == len(expected) + for i, (a, e) in enumerate(zip(actual, expected)): + np.testing.assert_allclose(a, e, rtol=1e-12, atol=1e-14, err_msg=f"control {i}") + + +@pytest.mark.parametrize("n_steps, snapshots", [(6, 2), (10, 3)]) +def test_taylor_test_under_checkpointing(n_steps, snapshots): + """The checkpointed gradient is the actual derivative, not merely a reproducible one.""" + rf, controls, directions = _tape_heat_equation(n_steps, Revolve(n_steps, snapshots)) + rate = pyadjoint.taylor_test(rf, controls, directions) + assert rate > 1.95 + + +def _tape_snes_heat_equation(n_steps, schedule=None, solution_dependent_diffusivity=False): + """Tape a heat equation solved as a residual problem via SNES. + + Unlike the linear model this cannot step in place: the unknown and the previous state + must be distinct functions, so the state update is an explicit assignment. + + Args: + n_steps: Number of tape timesteps to advance. + schedule: A ``checkpoint_schedules`` schedule, or None to disable checkpointing. + solution_dependent_diffusivity: If True the residual is genuinely nonlinear in the + unknown, which puts the unknown into the Jacobian and hence into the block's own + dependencies. See ``test_solution_dependent_jacobian_is_unsupported``. + """ + tape = pyadjoint.Tape() + pyadjoint.set_working_tape(tape) + if schedule is not None: + tape.enable_checkpointing(schedule) + + mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 6, 6) + V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) # type: ignore[arg-type] + dt = 0.1 + + controls = [] + for i in range(n_steps): + c = dolfinx_adjoint.Function(V, name=f"control_{i}") + c.interpolate(lambda x, i=i: 0.5 + 0.1 * (i + 1) * x[0]) + controls.append(c) + + v = ufl.TestFunction(V) + f = dolfinx_adjoint.Function(V, name="source") + uh = dolfinx_adjoint.Function(V, name="solution") + u_prev = dolfinx_adjoint.Function(V, name="previous") + + nu = (1 + uh**2) if solution_dependent_diffusivity else 1.0 + F = ((uh - u_prev) / dt * v + nu * ufl.inner(ufl.grad(uh), ufl.grad(v)) - f * v) * ufl.dx + + mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim) + boundary_facets = dolfinx.mesh.exterior_facet_indices(mesh.topology) + boundary_dofs = dolfinx.fem.locate_dofs_topological(V, mesh.topology.dim - 1, boundary_facets) + bc = dolfinx.fem.dirichletbc(0.0, boundary_dofs, V) + + snes_options = { + "snes_type": "newtonls", + "snes_linesearch_type": "none", + "snes_error_if_not_converged": True, + "snes_atol": 1e-14, + "snes_rtol": 1e-14, + } + snes_options.update(_PETSC_OPTIONS) + problem = dolfinx_adjoint.NonlinearProblem( + F, + uh, + bcs=[bc], + petsc_options=snes_options, + adjoint_petsc_options=_PETSC_OPTIONS, + ) + + J = dolfinx_adjoint.assemble_scalar(dt * uh**2 * ufl.dx) + for i in tape.timestepper(iter(range(n_steps))): + dolfinx_adjoint.assign(controls[i], f) + problem.solve() + dolfinx_adjoint.assign(uh, u_prev) + J = J + dolfinx_adjoint.assemble_scalar(dt * uh**2 * ufl.dx) + + rf = pyadjoint.ReducedFunctional(J, [pyadjoint.Control(c) for c in controls]) + return rf, controls, _perturbation_directions(V, n_steps) + + +@pytest.mark.xfail( + strict=True, + reason=( + "Pre-existing NonlinearProblemBlock defect, unrelated to checkpointing: a residual " + "problem advanced over several timesteps gets a wrong adjoint. The unknown is a " + "coefficient of the residual and so is registered as one of the block's own " + "dependencies; once its incoming value is itself control-dependent (which is what a " + "time loop creates) the adjoint contribution for it is computed against a residual in " + "which that value no longer appears, and ufl.adjoint raises IndexError on the " + "resulting argument-less form. Suppressing the error is not a fix: the gradient is " + "then silently wrong (observed Taylor rate -0.41 rather than 2). Checkpointing is not " + "involved -- this test does not enable a schedule. Until it is fixed, NonlinearProblem " + "cannot be used in a time loop and so cannot be covered by the checkpointing tests " + "above." + ), +) +def test_snes_time_loop_gradient_is_correct(): + """Records that NonlinearProblem cannot yet be advanced over timesteps.""" + rf, controls, directions = _tape_snes_heat_equation(4) + assert pyadjoint.taylor_test(rf, controls, directions) > 1.95 + + +@pytest.mark.parametrize( + "use_mpio", + [ + None, + False, + pytest.param( + True, + marks=pytest.mark.skipif(not h5py.get_config().mpi, reason="h5py is not built against MPI"), + ), + ], +) +def test_disk_gradient_matches_uncheckpointed(use_mpio): + """Storing checkpoints on disk does not change the gradient, in either file layout. + + `use_mpio=None` picks the layout automatically, and resolves to the per-process one on a + single process, so `True` is passed explicitly to reach the shared MPI-IO file as well. + """ + n_steps = 6 + rf_plain, controls_plain, _ = _tape_heat_equation(n_steps) + expected = _gradient(rf_plain, controls_plain) + + rf_disk, controls_disk, _ = _tape_heat_equation(n_steps, SingleDiskStorageSchedule(), disk=True, use_mpio=use_mpio) + actual = _gradient(rf_disk, controls_disk) + dolfinx_adjoint.checkpointing.disable_disk_checkpointing() + + for i, (a, e) in enumerate(zip(actual, expected)): + np.testing.assert_allclose(a, e, rtol=1e-12, atol=1e-14, err_msg=f"control {i}") + + +def test_disk_taylor_test(): + """The gradient from disk-stored checkpoints is the actual derivative.""" + n_steps = 6 + rf, controls, directions = _tape_heat_equation(n_steps, SingleDiskStorageSchedule(), disk=True) + rate = pyadjoint.taylor_test(rf, controls, directions) + dolfinx_adjoint.checkpointing.disable_disk_checkpointing() + assert rate > 1.95 + + +def test_disk_schedule_without_enabling_is_refused(): + """A disk-using schedule with no disk backend configured fails loudly, and says why.""" + with pytest.raises(CheckpointError, match="enable_disk_checkpointing"): + _tape_heat_equation(4, SingleDiskStorageSchedule()) From ee9525ff367b95e74d384706fd34d921808c36d4 Mon Sep 17 00:00:00 2001 From: Henrik Finsberg Date: Tue, 25 Aug 2026 12:29:33 +0200 Subject: [PATCH 2/9] Declare h5py, and wire up docs cross-references 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. --- _config.yml | 10 ++++++ .../time_distributed_control_checkpointing.py | 32 ++++++++++++------- pyproject.toml | 3 ++ src/dolfinx_adjoint/checkpointing.py | 2 +- 4 files changed, 35 insertions(+), 12 deletions(-) diff --git a/_config.yml b/_config.yml index e198ed7..ffc6477 100644 --- a/_config.yml +++ b/_config.yml @@ -41,9 +41,19 @@ sphinx: .py: - jupytext.reads - fmt: py + intersphinx_mapping: + checkpoint_schedules: ["https://www.firedrakeproject.org/checkpoint_schedules/", null] + pyadjoint: ["https://www.dolfin-adjoint.org/en/latest/", null] + dolfinx: ["https://docs.fenicsproject.org/dolfinx/main/python", null] + ufl: ["https://docs.fenicsproject.org/ufl/main/", null] + h5py: ["https://docs.h5py.org/en/stable/", null] + numpy: ["https://numpy.org/doc/stable/", null] + mpi4py: ["https://mpi4py.readthedocs.io/en/stable", null] + petsc4py: ["https://petsc.org/release/petsc4py", null] extra_extensions: - 'sphinx.ext.autodoc' - 'sphinx.ext.napoleon' - 'sphinx.ext.viewcode' + - 'sphinx.ext.intersphinx' exclude_patterns: [".pytest_cache/*"] diff --git a/demos/time_distributed_control_checkpointing.py b/demos/time_distributed_control_checkpointing.py index 8181ffc..bf52e02 100644 --- a/demos/time_distributed_control_checkpointing.py +++ b/demos/time_distributed_control_checkpointing.py @@ -8,8 +8,8 @@ # exhausts memory first. Checkpointing trades that memory for repeated work: only some states # are kept, and the rest are recomputed from the nearest stored one when the adjoint asks for # them. A schedule decides which to keep and when to recompute. The schedules come from -# `checkpoint_schedules` {cite}`Dolci2024`; for how step-based checkpointing combines with -# high-level algorithmic differentiation, see {cite}`Maddison2024`. +# `checkpoint_schedules` {cite}`tdcc-Dolci2024`; for how step-based checkpointing combines +# with high-level algorithmic differentiation, see {cite}`tdcc-Maddison2024`. # # Everything here comes from `pyadjoint` and `checkpoint_schedules` directly. The only thing # `dolfinx_adjoint` adds is `enable_disk_checkpointing`, used at the end. @@ -29,8 +29,8 @@ # ## Enabling a schedule # # A schedule has to be enabled on an empty tape, before anything is recorded, so that every -# timestep is treated the same way. `Revolve(num_steps, snapshots)` keeps at most `snapshots` -# states in memory and recomputes whatever else the adjoint needs. +# timestep is treated the same way. {py:class}`Revolve ` +# keeps at most `snapshots` states in memory and recomputes whatever else the adjoint needs. num_steps = 10 mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 8, 8) @@ -94,9 +94,9 @@ def solve_heat(schedule=None, disk=False): j = 0.5 * float(dt) * dolfinx_adjoint.assemble_scalar((u_0 - d) ** 2 * ufl.dx) - # `tape.timestepper` marks the tape timesteps the schedule reasons about. Note the - # `iter(...)`: `Tape.timestepper` calls `next()` on whatever it is given, so passing a bare - # `range` raises TypeError. Do not "simplify" it away. + # Tape.timestepper marks the tape timesteps the schedule reasons about. Note the + # `iter(...)`: it calls next() on whatever it is given, so passing a bare range raises + # TypeError. Do not "simplify" it away. for i in tape.timestepper(iter(range(num_steps))): t_val = float(dt) * (i + 1) dolfinx_adjoint.assign(t_val, t) @@ -152,14 +152,17 @@ def solve_heat(schedule=None, disk=False): # ## Storing checkpoints on disk # -# `Revolve` keeps its checkpoints in memory. When even those do not fit, a schedule can put -# them on disk instead, and `dolfinx_adjoint.enable_disk_checkpointing()` provides the storage. +# {py:class}`Revolve ` keeps its checkpoints in memory. +# When even those do not fit, a schedule such as +# {py:class}`SingleDiskStorageSchedule ` +# can put them on disk instead, and {py:func}`dolfinx_adjoint.enable_disk_checkpointing` +# provides the storage. # # These are *snapshot* checkpoints: they hold just this process's values for the function, and # assume the mesh and its partition are unchanged, so they are valid only within the run that # wrote them. They are deleted automatically. For a checkpoint that outlives the run, or that # can be read back on a different number of processes, use -# [adios4dolfinx](https://github.com/jorgensd/adios4dolfinx) instead. +# [io4dolfinx](https://github.com/scientificcomputing/io4dolfinx) instead. # # Like the schedule, it must be enabled before anything is recorded on the tape. @@ -177,7 +180,14 @@ def solve_heat(schedule=None, disk=False): print(f"J with checkpoints on disk: {J_disk:.12g}") print("Gradients agree to machine precision in all three cases.") +# Turning it off again deletes the checkpoint files. Every process must call it, because +# closing a shared checkpoint file is collective. + +dolfinx_adjoint.checkpointing.disable_disk_checkpointing() + # ## References # ```{bibliography} -# :filter: cited and ({"demos/time_distributed_control_checkpointing"} >= docnames) +# :filter: cited +# :labelprefix: +# :keyprefix: tdcc- # ``` diff --git a/pyproject.toml b/pyproject.toml index cf59b1d..30e6961 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,6 +14,9 @@ dependencies = [ "pyadjoint-ad>=2025.10.0", "typing_extensions; python_version < '3.11'", "packaging>=24.2", + # Storage for checkpoint schedules that keep state on disk. Build it against MPI to get + # one shared checkpoint file; without that each process writes its own. + "h5py", ] diff --git a/src/dolfinx_adjoint/checkpointing.py b/src/dolfinx_adjoint/checkpointing.py index 7a70125..5262b91 100644 --- a/src/dolfinx_adjoint/checkpointing.py +++ b/src/dolfinx_adjoint/checkpointing.py @@ -8,7 +8,7 @@ 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 ``adios4dolfinx``. +different number of processes. For a checkpoint that outlives the run, use ``io4dolfinx``. """ from __future__ import annotations From 6afa100be70a3eb28b752bed2db1c81fb26b95f6 Mon Sep 17 00:00:00 2001 From: Henrik Finsberg Date: Tue, 25 Aug 2026 12:33:04 +0200 Subject: [PATCH 3/9] Address review comments 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. --- .../time_distributed_control_checkpointing.py | 18 +++++++++--------- tests/test_checkpointing.py | 4 ++-- 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/demos/time_distributed_control_checkpointing.py b/demos/time_distributed_control_checkpointing.py index bf52e02..5b8dd26 100644 --- a/demos/time_distributed_control_checkpointing.py +++ b/demos/time_distributed_control_checkpointing.py @@ -94,9 +94,8 @@ def solve_heat(schedule=None, disk=False): j = 0.5 * float(dt) * dolfinx_adjoint.assemble_scalar((u_0 - d) ** 2 * ufl.dx) - # Tape.timestepper marks the tape timesteps the schedule reasons about. Note the - # `iter(...)`: it calls next() on whatever it is given, so passing a bare range raises - # TypeError. Do not "simplify" it away. + # `iter(...)` because timestepper calls next() on what it is given, and the default + # progress bar passes it straight through. Setting tape.progress_bar works too. for i in tape.timestepper(iter(range(num_steps))): t_val = float(dt) * (i + 1) dolfinx_adjoint.assign(t_val, t) @@ -126,8 +125,8 @@ def solve_heat(schedule=None, disk=False): 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): - assert np.allclose(a, e) +for a, e in zip(grad_ckpt, grad_plain, strict=True): + np.testing.assert_allclose(a, e) if mesh.comm.rank == 0: print(f"J without checkpointing: {J_plain:.12g}") @@ -135,8 +134,9 @@ def solve_heat(schedule=None, disk=False): # ## A Taylor test through the schedule # -# Reproducing the un-checkpointed gradient shows the two agree, but not that either is right. -# A Taylor test does: the first-order remainder should converge at second order. +# The check above shows the two gradients agree with each other. It does not show that either +# is correct, since both could be wrong in the same way. A Taylor test checks that directly: the +# first-order remainder must converge at second order. directions = [] for k in range(num_steps): @@ -173,8 +173,8 @@ def solve_heat(schedule=None, disk=False): grad_disk = [np.copy(g.x.array) for g in rf_disk.derivative()] assert np.isclose(J_plain, J_disk) -for a, e in zip(grad_disk, grad_plain): - assert np.allclose(a, e) +for a, e in zip(grad_disk, grad_plain, strict=True): + np.testing.assert_allclose(a, e) if mesh.comm.rank == 0: print(f"J with checkpoints on disk: {J_disk:.12g}") diff --git a/tests/test_checkpointing.py b/tests/test_checkpointing.py index 7c226a7..143c893 100644 --- a/tests/test_checkpointing.py +++ b/tests/test_checkpointing.py @@ -145,7 +145,7 @@ def test_gradient_matches_uncheckpointed(n_steps, snapshots): actual = _gradient(rf_ckpt, controls_ckpt) assert len(actual) == len(expected) - for i, (a, e) in enumerate(zip(actual, expected)): + for i, (a, e) in enumerate(zip(actual, expected, strict=True)): np.testing.assert_allclose(a, e, rtol=1e-12, atol=1e-14, err_msg=f"control {i}") @@ -272,7 +272,7 @@ def test_disk_gradient_matches_uncheckpointed(use_mpio): actual = _gradient(rf_disk, controls_disk) dolfinx_adjoint.checkpointing.disable_disk_checkpointing() - for i, (a, e) in enumerate(zip(actual, expected)): + for i, (a, e) in enumerate(zip(actual, expected, strict=True)): np.testing.assert_allclose(a, e, rtol=1e-12, atol=1e-14, err_msg=f"control {i}") From 65cdf7289e80ecd4da01de0d52674d9a5a4ba652 Mon Sep 17 00:00:00 2001 From: Henrik Finsberg Date: Tue, 25 Aug 2026 12:43:50 +0200 Subject: [PATCH 4/9] Require Python 3.10, and type-narrow the linear combination assignment 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. --- demos/time_distributed_control2.py | 139 +++++++++++++++++++++++++++++ pyproject.toml | 1 + src/dolfinx_adjoint/utils.py | 4 + 3 files changed, 144 insertions(+) create mode 100644 demos/time_distributed_control2.py diff --git a/demos/time_distributed_control2.py b/demos/time_distributed_control2.py new file mode 100644 index 0000000..d1b2e80 --- /dev/null +++ b/demos/time_distributed_control2.py @@ -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 + +from collections import OrderedDict + +from mpi4py import MPI + +import dolfinx +import numpy as np +import pyadjoint +import ufl + +import dolfinx_adjoint + +mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 8, 8) +x = ufl.SpatialCoordinate(mesh) + +nu = dolfinx.fem.Constant(mesh, np.float64(1e-5)) +nu.name = "nu" # type: ignore + +t = dolfinx_adjoint.Constant(mesh, dolfinx.default_scalar_type(0.0)) # type: ignore +t.name = "time" +d = 16 * x[0] * (x[0] - 1) * x[1] * (x[1] - 1) * ufl.sin(ufl.pi * t) + +dt = dolfinx.fem.Constant(mesh, dolfinx.default_scalar_type(0.1)) # type: ignore +T = 1 + +V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) # type: ignore[arg-type] +ctrls = OrderedDict() +t_val = float(dt) +while t_val <= T: + ctrls[t_val] = dolfinx_adjoint.Function(V, name=f"control_{t_val}") + t_val += float(dt) + + +def solve_heat(ctrls): + u = ufl.TrialFunction(V) + v = ufl.TestFunction(V) + + f = dolfinx_adjoint.Function(V, name="source") + u_0 = dolfinx_adjoint.Function(V, name="solution") + + F = ((u - u_0) / dt * v + nu * ufl.inner(ufl.grad(u), ufl.grad(v)) - f * v) * ufl.dx + a, L = ufl.system(F) + mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim) + exterior_facets = dolfinx.mesh.exterior_facet_indices(mesh.topology) + exterior_dofs = dolfinx.fem.locate_dofs_topological(V, mesh.topology.dim - 1, exterior_facets) + + bc = dolfinx.fem.dirichletbc(0.0, exterior_dofs, V) + + j = 0.5 * float(dt) * dolfinx_adjoint.assemble_scalar((u_0 - d) ** 2 * ufl.dx) + + t_val = float(dt) + problem = dolfinx_adjoint.LinearProblem( + a, + L, + u=u_0, + bcs=[bc], + petsc_options={ + "ksp_type": "preonly", + "pc_type": "lu", + "pc_factor_mat_solver_type": "mumps", + "ksp_error_if_not_converged": True, + }, + adjoint_petsc_options={ + "ksp_type": "preonly", + "pc_type": "lu", + "pc_factor_mat_solver_type": "mumps", + "ksp_error_if_not_converged": True, + }, + tlm_petsc_options={ + "ksp_type": "preonly", + "pc_type": "lu", + "pc_factor_mat_solver_type": "mumps", + "ksp_error_if_not_converged": True, + }, + ) + dolfinx_adjoint.assign(t_val, t) + + while t_val <= T: + # Update source term from control array + dolfinx_adjoint.assign(ctrls[t_val], f) + + # Update data function + + # Solve PDE + problem.solve() + + # Implement a trapezoidal rule + if t_val > T - float(dt): + weight = 0.5 + else: + weight = 1 + j += weight * float(dt) * dolfinx_adjoint.assemble_scalar((u_0 - d) ** 2 * ufl.dx) + # Update time + t_val += float(dt) + dolfinx_adjoint.assign(t_val, t) + + return u_0, d, j + + +u, d, j = solve_heat(ctrls) + + +alpha = dolfinx.fem.Constant(mesh, np.float64(1.0e-1)) +regularisation = ( + alpha + / 2 + * sum([1 / dt * (fb - fa) ** 2 * ufl.dx for fb, fa in zip(list(ctrls.values())[1:], list(ctrls.values())[:-1])]) +) + + +J = j + dolfinx_adjoint.assemble_scalar(regularisation) +m = [pyadjoint.Control(c) for c in ctrls.values()] + + +rf = pyadjoint.ReducedFunctional(J, m) + +tape = pyadjoint.get_working_tape() +total_steps = 10 +tape.timestepper(range(total_steps)) +tape.visualise_dot("test.dot") + +opt_ctrls = pyadjoint.minimize( + rf, + method="BFGS", + # method="Newton-CG", + options={"maxiter": 100, "disp": True}, +) + +out_ctrl = dolfinx.fem.Function(V, name="optimal_control") +with dolfinx.io.VTXWriter(mesh.comm, "opt_ctrl.bp", [out_ctrl]) as vtx: + for t_val, c in zip(ctrls.keys(), opt_ctrls): + out_ctrl.x.array[:] = c.x.array[:] + vtx.write(t_val) + + +assert np.isclose(np.linalg.norm(opt_ctrls[0].x.array), 4.930056079391683) +assert np.isclose(np.linalg.norm(opt_ctrls[-1].x.array), 2.8756312728703963) diff --git a/pyproject.toml b/pyproject.toml index 30e6961..a99319e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,6 +9,7 @@ authors = [{ name = "Jørgen S. Dokken", email = "dokken@simula.no" }] license = "MIT" license-files = ["LICENSE"] readme = "README.md" +requires-python = ">=3.10" dependencies = [ "fenics-dolfinx>=0.10.0", "pyadjoint-ad>=2025.10.0", diff --git a/src/dolfinx_adjoint/utils.py b/src/dolfinx_adjoint/utils.py index 5b5730e..91a60d0 100644 --- a/src/dolfinx_adjoint/utils.py +++ b/src/dolfinx_adjoint/utils.py @@ -59,6 +59,10 @@ def assign_linear_combination(value: ufl.core.expr.Expr, function: dolfinx.fem.F function.x.array[:] = 0.0 floatifier = Floatify() for weight, func in pairs: + # extract_linear_combination is typed against UFL, which knows nothing of degrees of + # freedom; assigning a linear combination needs the DOLFINx Function that carries them. + if not isinstance(func, dolfinx.fem.Function): + raise TypeError(f"Expected the linear combination to be over dolfinx Functions, got {type(func)}.") if not func.function_space == function.function_space: raise ValueError("Function spaces of all functions in the linear combination must match for assignment.") function.x.array[:] += floatifier.process(weight) * func.x.array[:] From d5d7befc09dc62990fb63ee9e7eefd719da4345d Mon Sep 17 00:00:00 2001 From: Henrik Finsberg Date: Tue, 25 Aug 2026 12:10:30 +0000 Subject: [PATCH 5/9] Remove time_distributed_control2.py --- demos/time_distributed_control2.py | 139 ----------------------------- 1 file changed, 139 deletions(-) delete mode 100644 demos/time_distributed_control2.py diff --git a/demos/time_distributed_control2.py b/demos/time_distributed_control2.py deleted file mode 100644 index d1b2e80..0000000 --- a/demos/time_distributed_control2.py +++ /dev/null @@ -1,139 +0,0 @@ -# # Time-distributed control -# Based on example from https://dolfin-adjoint.github.io/dolfin-adjoint/documentation/time-distributed-control/time-distributed-control.html - -from collections import OrderedDict - -from mpi4py import MPI - -import dolfinx -import numpy as np -import pyadjoint -import ufl - -import dolfinx_adjoint - -mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 8, 8) -x = ufl.SpatialCoordinate(mesh) - -nu = dolfinx.fem.Constant(mesh, np.float64(1e-5)) -nu.name = "nu" # type: ignore - -t = dolfinx_adjoint.Constant(mesh, dolfinx.default_scalar_type(0.0)) # type: ignore -t.name = "time" -d = 16 * x[0] * (x[0] - 1) * x[1] * (x[1] - 1) * ufl.sin(ufl.pi * t) - -dt = dolfinx.fem.Constant(mesh, dolfinx.default_scalar_type(0.1)) # type: ignore -T = 1 - -V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1)) # type: ignore[arg-type] -ctrls = OrderedDict() -t_val = float(dt) -while t_val <= T: - ctrls[t_val] = dolfinx_adjoint.Function(V, name=f"control_{t_val}") - t_val += float(dt) - - -def solve_heat(ctrls): - u = ufl.TrialFunction(V) - v = ufl.TestFunction(V) - - f = dolfinx_adjoint.Function(V, name="source") - u_0 = dolfinx_adjoint.Function(V, name="solution") - - F = ((u - u_0) / dt * v + nu * ufl.inner(ufl.grad(u), ufl.grad(v)) - f * v) * ufl.dx - a, L = ufl.system(F) - mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim) - exterior_facets = dolfinx.mesh.exterior_facet_indices(mesh.topology) - exterior_dofs = dolfinx.fem.locate_dofs_topological(V, mesh.topology.dim - 1, exterior_facets) - - bc = dolfinx.fem.dirichletbc(0.0, exterior_dofs, V) - - j = 0.5 * float(dt) * dolfinx_adjoint.assemble_scalar((u_0 - d) ** 2 * ufl.dx) - - t_val = float(dt) - problem = dolfinx_adjoint.LinearProblem( - a, - L, - u=u_0, - bcs=[bc], - petsc_options={ - "ksp_type": "preonly", - "pc_type": "lu", - "pc_factor_mat_solver_type": "mumps", - "ksp_error_if_not_converged": True, - }, - adjoint_petsc_options={ - "ksp_type": "preonly", - "pc_type": "lu", - "pc_factor_mat_solver_type": "mumps", - "ksp_error_if_not_converged": True, - }, - tlm_petsc_options={ - "ksp_type": "preonly", - "pc_type": "lu", - "pc_factor_mat_solver_type": "mumps", - "ksp_error_if_not_converged": True, - }, - ) - dolfinx_adjoint.assign(t_val, t) - - while t_val <= T: - # Update source term from control array - dolfinx_adjoint.assign(ctrls[t_val], f) - - # Update data function - - # Solve PDE - problem.solve() - - # Implement a trapezoidal rule - if t_val > T - float(dt): - weight = 0.5 - else: - weight = 1 - j += weight * float(dt) * dolfinx_adjoint.assemble_scalar((u_0 - d) ** 2 * ufl.dx) - # Update time - t_val += float(dt) - dolfinx_adjoint.assign(t_val, t) - - return u_0, d, j - - -u, d, j = solve_heat(ctrls) - - -alpha = dolfinx.fem.Constant(mesh, np.float64(1.0e-1)) -regularisation = ( - alpha - / 2 - * sum([1 / dt * (fb - fa) ** 2 * ufl.dx for fb, fa in zip(list(ctrls.values())[1:], list(ctrls.values())[:-1])]) -) - - -J = j + dolfinx_adjoint.assemble_scalar(regularisation) -m = [pyadjoint.Control(c) for c in ctrls.values()] - - -rf = pyadjoint.ReducedFunctional(J, m) - -tape = pyadjoint.get_working_tape() -total_steps = 10 -tape.timestepper(range(total_steps)) -tape.visualise_dot("test.dot") - -opt_ctrls = pyadjoint.minimize( - rf, - method="BFGS", - # method="Newton-CG", - options={"maxiter": 100, "disp": True}, -) - -out_ctrl = dolfinx.fem.Function(V, name="optimal_control") -with dolfinx.io.VTXWriter(mesh.comm, "opt_ctrl.bp", [out_ctrl]) as vtx: - for t_val, c in zip(ctrls.keys(), opt_ctrls): - out_ctrl.x.array[:] = c.x.array[:] - vtx.write(t_val) - - -assert np.isclose(np.linalg.norm(opt_ctrls[0].x.array), 4.930056079391683) -assert np.isclose(np.linalg.norm(opt_ctrls[-1].x.array), 2.8756312728703963) From 56f3cdae8a2fab1d940a108396877337dd612702 Mon Sep 17 00:00:00 2001 From: Henrik Finsberg Date: Tue, 25 Aug 2026 14:52:55 +0200 Subject: [PATCH 6/9] Fix the parallel deadlock, and address review comments 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. --- .../time_distributed_control_checkpointing.py | 15 +- src/dolfinx_adjoint/blocks/solvers.py | 42 ++++- src/dolfinx_adjoint/checkpointing.py | 175 +++++++++++++----- src/dolfinx_adjoint/types/function.py | 5 +- tests/test_checkpointing.py | 29 ++- 5 files changed, 202 insertions(+), 64 deletions(-) diff --git a/demos/time_distributed_control_checkpointing.py b/demos/time_distributed_control_checkpointing.py index 5b8dd26..9798b0f 100644 --- a/demos/time_distributed_control_checkpointing.py +++ b/demos/time_distributed_control_checkpointing.py @@ -139,12 +139,15 @@ def solve_heat(schedule=None, disk=False): # first-order remainder must converge at second order. 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) +# The directions are inputs to the test, not part of the model, so building them should not be +# recorded on the tape. +with pyadjoint.stop_annotating(): + 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) rf_ckpt, controls_ckpt = solve_heat(Revolve(num_steps, 3)) rate = pyadjoint.taylor_test(rf_ckpt, controls_ckpt, directions) diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py index 8460833..7cdd1f5 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/src/dolfinx_adjoint/blocks/solvers.py @@ -23,11 +23,11 @@ def _initial_guess_for( 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 + {py:class}`dolfinx.fem.Function` cannot do. Outside checkpointing nothing asks, which is why a plain one survived for so long. """ with stop_annotating(): - if isinstance(u, dolfinx.fem.Function): + if isinstance(u, _Function): return Function(u.function_space, name=u.name + "_initial_guess") return [Function(ui.function_space, name=ui.name + "_initial_guess") for ui in u] @@ -475,6 +475,21 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar return dudm + def _reusable_vector(self, attribute: str, template: dolfinx.la.Vector) -> dolfinx.la.Vector: + """Return a cached vector shaped like ``template``, allocating it on first use. + + Allocating a distributed vector is collective, and so is destroying the PETSc vector it + can hand out. Doing either per call leaves both at the mercy of the garbage collector, + which does not run at the same moment on every process: one process then enters a + collective the others have already left, and the adjoint deadlocks. Allocating once, at + a point every process reaches together, and reusing it keeps the collectives in step. + """ + cached = getattr(self, attribute, None) + if cached is None or cached.array.size != template.array.size: + cached = dolfinx.la.vector(template.index_map, template.block_size) + setattr(self, attribute, cached) + return cached + def prepare_evaluate_adj( self, inputs: typing.Sequence[Function], @@ -491,8 +506,8 @@ def prepare_evaluate_adj( # Extract dJ/du[v] from the adjoint inputs. assert len(adj_inputs) == 1 adj_rhs = adj_inputs[0] - dJdu = dolfinx.la.vector(adj_rhs.index_map, adj_rhs.block_size) - dJdu.array[:] = adj_rhs.array[:].copy() + dJdu = self._reusable_vector("_adjoint_rhs", adj_rhs) + dJdu.array[:] = adj_rhs.array[:] # Solve adjoint problem compiled_dFdu = dolfinx.fem.form( @@ -1093,6 +1108,21 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar solve_linear_problem(A_tlm, dudm.x, b_tlm, petsc_options=self._tlm_petsc_options) return dudm + def _reusable_vector(self, attribute: str, template: dolfinx.la.Vector) -> dolfinx.la.Vector: + """Return a cached vector shaped like ``template``, allocating it on first use. + + Allocating a distributed vector is collective, and so is destroying the PETSc vector it + can hand out. Doing either per call leaves both at the mercy of the garbage collector, + which does not run at the same moment on every process: one process then enters a + collective the others have already left, and the adjoint deadlocks. Allocating once, at + a point every process reaches together, and reusing it keeps the collectives in step. + """ + cached = getattr(self, attribute, None) + if cached is None or cached.array.size != template.array.size: + cached = dolfinx.la.vector(template.index_map, template.block_size) + setattr(self, attribute, cached) + return cached + def prepare_evaluate_adj( self, inputs: typing.Sequence[Function], @@ -1108,8 +1138,8 @@ def prepare_evaluate_adj( # Extract dJ/du[v] from the adjoint inputs. assert len(adj_inputs) == 1 adj_rhs = adj_inputs[0] - dJdu = dolfinx.la.vector(adj_rhs.index_map, adj_rhs.block_size) - dJdu.array[:] = adj_rhs.array[:].copy() + dJdu = self._reusable_vector("_adjoint_rhs", adj_rhs) + dJdu.array[:] = adj_rhs.array[:] # Solve adjoint problem compiled_dFdu = dolfinx.fem.form( diff --git a/src/dolfinx_adjoint/checkpointing.py b/src/dolfinx_adjoint/checkpointing.py index 5262b91..8278653 100644 --- a/src/dolfinx_adjoint/checkpointing.py +++ b/src/dolfinx_adjoint/checkpointing.py @@ -1,6 +1,6 @@ """Snapshot checkpointing of DOLFINx functions to disk. -A checkpoint schedule that uses :class:`checkpoint_schedules.StorageType.DISK` needs somewhere +A checkpoint schedule that uses {py:class}`checkpoint_schedules.schedule.StorageType` ``DISK`` needs somewhere to put a function's values. This module provides that as a *snapshot* checkpoint: it is written and read within a single run, by the same processes, against an unchanged mesh and partition. Under those assumptions the whole payload is a process's local values, so no mesh, geometry or @@ -8,30 +8,36 @@ 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``. +different number of processes. For a checkpoint that outlives the run, use {py:mod}`io4dolfinx`. """ from __future__ import annotations import os +import pathlib import tempfile import typing import weakref from mpi4py import MPI -import dolfinx +import h5py import numpy as np import pyadjoint.checkpointing from pyadjoint.tape import TapePackageData, get_working_tape +if typing.TYPE_CHECKING: + # Annotations only. Importing at runtime would be circular: dolfinx_adjoint.types imports + # this module to decide where a checkpoint goes. + from .types.function import Function + __all__ = ["enable_disk_checkpointing", "disable_disk_checkpointing", "SnapshotCheckpoint"] #: Key under which the disk checkpointer registers itself in ``Tape._package_data``. _PACKAGE_KEY = "dolfinx_adjoint" #: The active checkpointer, or None when disk checkpointing is not enabled. -_checkpointer: typing.Optional["_DiskCheckpointer"] = None +_checkpointer: "_DiskCheckpointer | None" = None # Message pyadjoint shows when a schedule wants disk storage but none is configured. pyadjoint.checkpointing.disk_checkpointing_callback[_PACKAGE_KEY] = ( @@ -39,15 +45,7 @@ ) -def _import_h5py(): - try: - import h5py - except ImportError as e: # pragma: no cover - exercised only without h5py - raise ImportError("Disk checkpointing requires h5py. Install it with 'pip install h5py'.") from e - return h5py - - -def _layout(function: dolfinx.fem.Function, shared_file: bool, comm: MPI.Comm) -> tuple[int, int, int]: +def _layout(function: Function, shared_file: bool, comm: MPI.Intracomm) -> tuple[int, int, int]: """Describe where this process's values sit in a stored dataset. The whole local array is stored, ghost values included, not just the locally owned values. @@ -58,6 +56,12 @@ def _layout(function: dolfinx.fem.Function, shared_file: bool, comm: MPI.Comm) - call on that path deadlocks as soon as one process takes a cached value while another reads. Storing the ghosts makes restoring purely local, so it cannot deadlock. + Args: + function: The function whose values are about to be stored. + shared_file: Whether the dataset spans every process's values (one shared file) or + only this process's (one file per process). + comm: The communicator the checkpoint files are shared over. + Returns: A tuple of the number of values this process stores, the length of the whole dataset, and this process's offset into it. @@ -65,9 +69,13 @@ def _layout(function: dolfinx.fem.Function, shared_file: bool, comm: MPI.Comm) - n_local = function.x.array.size 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) - return n_local, sum(sizes), sum(sizes[: comm.rank]) + # Collective, but reached only from the write path, which every process reaches together. + # An exclusive scan rather than gathering every size and summing a prefix: it is the + # operation this actually is, and its cost does not grow with the number of processes. + offset = comm.exscan(n_local, op=MPI.SUM) + if comm.rank == 0: + offset = 0 + return n_local, comm.allreduce(n_local, op=MPI.SUM), offset class _CheckpointFile: @@ -80,22 +88,43 @@ class _CheckpointFile: 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): - h5py = _import_h5py() - self.path = path - self.comm = comm - # A shared file is a single file that every process writes its own slice of. Without - # MPI-IO each process gets a file to itself instead. + def __init__(self, path: pathlib.Path, comm: MPI.Intracomm, use_mpio: bool, cleanup: bool): + """ + Args: + path: Where to create the file. + comm: The communicator the file is shared over. + use_mpio: Whether to open one shared file with MPI-IO, so that every process writes + its own slice of each dataset. Without it each process gets its own file. + cleanup: Whether to delete the file when it is closed. False keeps it on disk for + inspection, which is only useful for debugging. + """ + self._path = path + self._comm = comm # One shared file that every process writes a slice of, or one file per process. - self.shared_file = use_mpio or comm.size == 1 + self._shared_file = use_mpio or comm.size == 1 # Only a shared file is written by more than one process, so only then does deleting it # belong to a single one of them. - self._deleted_by_this_process = cleanup and (comm.rank == 0 or not self.shared_file) + self._deleted_by_this_process = cleanup and (comm.rank == 0 or not self._shared_file) kwargs = {"driver": "mpio", "comm": comm} if use_mpio else {} self._handle = h5py.File(path, "w", **kwargs) self._next_index = 0 self._closed = False + @property + def path(self) -> pathlib.Path: + """Where this file lives.""" + return self._path + + @property + def comm(self) -> MPI.Intracomm: + """The communicator this file is shared over.""" + return self._comm + + @property + def shared_file(self) -> bool: + """Whether one file holds every process's values, rather than one file per process.""" + return self._shared_file + def next_key(self) -> str: """Return a dataset name that every process agrees on. @@ -108,10 +137,28 @@ def next_key(self) -> str: return key def write(self, key: str, values: np.ndarray, n_global: int, offset: int) -> None: + """Store one process's values in a new dataset. + + Args: + key: Dataset name, from {py:meth}`next_key`. + values: The values this process contributes, ghost values included. + n_global: Length of the whole dataset, across every process. + offset: Where this process's values start in it. + """ dataset = self._handle.create_dataset(key, (n_global,), dtype=values.dtype) dataset[offset : offset + values.size] = values def read(self, key: str, n_local: int, offset: int) -> np.ndarray: + """Read this process's values back out of a dataset. + + Args: + key: Dataset name, as passed to {py:meth}`write`. + n_local: How many values this process stored. + offset: Where this process's values start in the dataset. + + Returns: + The stored values, ghost values included. + """ return self._handle[key][offset : offset + n_local] def close(self) -> None: @@ -125,7 +172,7 @@ def close(self) -> None: self._handle.close() if self._deleted_by_this_process: try: - os.remove(self.path) + os.remove(self._path) except OSError: # pragma: no cover - another process may have removed it first pass @@ -139,7 +186,7 @@ class SnapshotCheckpoint: __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): + def __init__(self, file: _CheckpointFile, key: str, function: Function, n_local: int, offset: int): # Holding the file keeps it alive for exactly as long as some checkpoint needs it. self._file = file self._key = key @@ -153,9 +200,9 @@ def __init__(self, file: _CheckpointFile, key: str, function: dolfinx.fem.Functi # and a fresh object each time makes those maps miss. Weak rather than strong so the # values are released again once the block is done with them, which is the point of # storing them on disk in the first place. - self._cache: typing.Optional[weakref.ReferenceType] = None + self._cache: weakref.ReferenceType | None = None - def restore(self): + def restore(self) -> Function: """Read the stored values back into a function of the original type.""" from .types.function import Function @@ -166,7 +213,7 @@ def restore(self): # Mirrors Function._ad_new_like: going through __new__ preserves the concrete subclass # (Constant takes a different constructor signature). - restored = self._cls.__new__(self._cls, self._space) + restored = self._cls.__new__(self._cls, self._space) # type: ignore[call-arg] Function.__init__(restored, self._space) restored.name = self._name # Purely local: the stored array already includes the ghost values, so no scatter. @@ -178,7 +225,23 @@ def restore(self): 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): + def __init__( + self, + directory: pathlib.Path, + comm: MPI.Intracomm, + use_mpio: bool, + cleanup: bool, + owns_directory: bool, + ): + """ + Args: + directory: Where the checkpoint files are written. + comm: The communicator the files are shared over. + use_mpio: Whether to write one shared file with MPI-IO. + cleanup: Whether to delete the files, and the directory, on teardown. + owns_directory: Whether this object created the directory and so should remove it. + Must agree across processes, or teardown deadlocks. + """ self._directory = directory self._comm = comm self._use_mpio = use_mpio @@ -195,7 +258,7 @@ def _roll_to_new_file(self) -> _CheckpointFile: if previous is not None: previous.close() rank_suffix = "" if (self._use_mpio or self._comm.size == 1) else f"_rank{self._comm.rank}" - path = os.path.join(self._directory, f"checkpoint_{self._generation}{rank_suffix}.h5") + path = self._directory / f"checkpoint_{self._generation}{rank_suffix}.h5" self._generation += 1 return _CheckpointFile(path, self._comm, self._use_mpio, self._cleanup) @@ -204,7 +267,15 @@ def storing(self) -> bool: """Whether values should currently be written to disk rather than kept in memory.""" return self._storing - def store(self, function: dolfinx.fem.Function) -> SnapshotCheckpoint: + def store(self, function: Function) -> SnapshotCheckpoint: + """Write a function's values to the current checkpoint file. + + Args: + function: The function to store. + + Returns: + A handle that reads the values back. + """ n_local, n_global, offset = _layout(function, self._file.shared_file, self._comm) key = self._file.next_key() self._file.write(key, function.x.array, n_global, offset) @@ -254,12 +325,18 @@ def pause_checkpointing(self): self._storing = False -def maybe_disk_checkpoint(function: dolfinx.fem.Function) -> typing.Optional[SnapshotCheckpoint]: +def maybe_disk_checkpoint(function: Function) -> SnapshotCheckpoint | None: """Store ``function`` on disk if disk checkpointing is active, otherwise return None. Returning None tells the caller to fall back to an in-memory copy. Disk storage is only active inside the windows pyadjoint opens around writing checkpoint data, so most calls return None even when disk checkpointing is enabled. + + Args: + function: The function pyadjoint is asking to checkpoint. + + Returns: + A handle to the stored values, or None to keep them in memory. """ if _checkpointer is None or not _checkpointer.storing: return None @@ -267,10 +344,10 @@ def maybe_disk_checkpoint(function: dolfinx.fem.Function) -> typing.Optional[Sna def enable_disk_checkpointing( - dirname: typing.Optional[str] = None, - comm: typing.Optional[MPI.Comm] = None, + dirname: str | os.PathLike | None = None, + comm: MPI.Intracomm | None = None, cleanup: bool = True, - use_mpio: typing.Optional[bool] = None, + use_mpio: bool | None = None, ) -> None: """Store checkpoints on disk rather than in memory. @@ -281,8 +358,9 @@ def enable_disk_checkpointing( dirname: Directory to hold the checkpoint files. A temporary directory is created if this is not given. comm: MPI communicator. Defaults to ``MPI.COMM_WORLD``. - cleanup: Whether to delete checkpoint files once nothing refers to them. Pass False to - keep them for inspection. + cleanup: Whether to delete the checkpoint files, and the temporary directory, on + teardown. Pass False to keep them for inspection; they are unreadable by any later + run either way. use_mpio: Whether to write one shared file with MPI-IO. The default chooses it when running on more than one process with an MPI-enabled h5py, and falls back to one file per process otherwise. Pass False to force the per-process layout. @@ -301,7 +379,6 @@ def enable_disk_checkpointing( ) comm = MPI.COMM_WORLD if comm is None else comm - h5py = _import_h5py() if use_mpio is None: use_mpio = comm.size > 1 and h5py.get_config().mpi elif use_mpio and not h5py.get_config().mpi: @@ -309,15 +386,25 @@ def enable_disk_checkpointing( "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 - if dirname is None: + # Whether we created the directory decides whether teardown removes it, and teardown + # synchronises the processes before doing so. Every process must therefore agree: if some + # were given a `dirname` and others were not, teardown would deadlock. + 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, " + f"but it was omitted on {without_dirname} of {comm.size}." + ) + + owns_directory = without_dirname == comm.size + if owns_directory: # Every process must agree on the directory, even in the per-process layout. created = tempfile.mkdtemp(prefix="dolfinx_adjoint_checkpoints_") if comm.rank == 0 else None - directory = typing.cast(str, comm.bcast(created, root=0)) + directory = pathlib.Path(comm.bcast(created, root=0)) else: - directory = dirname + directory = pathlib.Path(typing.cast("str | os.PathLike", dirname)) if comm.rank == 0: - os.makedirs(directory, exist_ok=True) + directory.mkdir(parents=True, exist_ok=True) comm.Barrier() _checkpointer = _DiskCheckpointer(directory, comm, use_mpio, cleanup, owns_directory) diff --git a/src/dolfinx_adjoint/types/function.py b/src/dolfinx_adjoint/types/function.py index db03471..54740c1 100644 --- a/src/dolfinx_adjoint/types/function.py +++ b/src/dolfinx_adjoint/types/function.py @@ -15,6 +15,7 @@ from pyadjoint.tape import no_annotations from ..blocks.assembly import assemble_compiled_form +from ..checkpointing import SnapshotCheckpoint, maybe_disk_checkpoint from ..utils import function_from_vector, gather @@ -69,8 +70,6 @@ def _ad_init_object(cls, obj): @no_annotations def _ad_create_checkpoint(self): - from ..checkpointing import maybe_disk_checkpoint - # While a schedule is storing to disk, hand back a reference to the stored values # rather than the values themselves. This is the only seam pyadjoint offers for # choosing where checkpoint data lives. @@ -88,8 +87,6 @@ def _ad_create_checkpoint(self): return checkpoint def _ad_restore_at_checkpoint(self, checkpoint): - from ..checkpointing import SnapshotCheckpoint - if isinstance(checkpoint, SnapshotCheckpoint): return checkpoint.restore() return checkpoint diff --git a/tests/test_checkpointing.py b/tests/test_checkpointing.py index 143c893..f916bc5 100644 --- a/tests/test_checkpointing.py +++ b/tests/test_checkpointing.py @@ -6,6 +6,8 @@ and recomputed, never the value of the derivative. """ +import gc + from mpi4py import MPI from petsc4py import PETSc @@ -40,9 +42,11 @@ def isolated_tape(): """ previous_tape = pyadjoint.get_working_tape() previous_options = dict(PETSc.Options().getAll()) + _collect() try: yield finally: + _collect() dolfinx_adjoint.checkpointing.disable_disk_checkpointing() pyadjoint.set_working_tape(previous_tape) options = PETSc.Options() @@ -50,17 +54,32 @@ def isolated_tape(): options.delValue(key) +def _collect(): + """Collect garbage now, so that it happens at the same moment on every process. + + Discarded tapes hold blocks, and each block owns a dolfinx LinearProblem whose __del__ + destroys PETSc objects -- which is collective. Blocks sit in reference cycles, so they are + freed by the cyclic collector rather than by refcounting, and that runs when each process + happens to cross an allocation threshold, not in step. Whichever process collects first + then enters a collective the others are not in, and the run deadlocks. Collecting + deliberately, at points every process reaches together, keeps those destructors in step. + """ + gc.collect() + + def _perturbation_directions(V, n): """Perturbation directions for a Taylor test. Built by interpolating analytic expressions rather than from random numbers: the directions must agree across processes, and per-rank random values do not. """ + # Not part of the model, so keep them off the tape. directions = [] - for k in range(n): - h = dolfinx_adjoint.Function(V, name=f"direction_{k}") - h.interpolate(lambda x, k=k: np.sin((k + 1) * np.pi * x[0]) * np.cos(np.pi * x[1])) - directions.append(h) + with pyadjoint.stop_annotating(): + for k in range(n): + h = dolfinx_adjoint.Function(V, name=f"direction_{k}") + h.interpolate(lambda x, k=k: np.sin((k + 1) * np.pi * x[0]) * np.cos(np.pi * x[1])) + directions.append(h) return directions @@ -76,6 +95,7 @@ def _tape_heat_equation(n_steps, schedule=None, disk=False, use_mpio=None): Returns: A tuple of the reduced functional, the controls, and perturbation directions. """ + _collect() tape = pyadjoint.Tape() pyadjoint.set_working_tape(tape) # Both of these must happen before anything is recorded on this tape. @@ -170,6 +190,7 @@ def _tape_snes_heat_equation(n_steps, schedule=None, solution_dependent_diffusiv unknown, which puts the unknown into the Jacobian and hence into the block's own dependencies. See ``test_solution_dependent_jacobian_is_unsupported``. """ + _collect() tape = pyadjoint.Tape() pyadjoint.set_working_tape(tape) if schedule is not None: From 84f20d8b9ff1a764cc2c5928bc1aa18c150751a9 Mon Sep 17 00:00:00 2001 From: Henrik Finsberg Date: Tue, 25 Aug 2026 16:16:47 +0200 Subject: [PATCH 7/9] Say why the checked type and the constructed type differ 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. --- src/dolfinx_adjoint/blocks/solvers.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/src/dolfinx_adjoint/blocks/solvers.py b/src/dolfinx_adjoint/blocks/solvers.py index 7cdd1f5..a19a0ea 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/src/dolfinx_adjoint/blocks/solvers.py @@ -17,14 +17,19 @@ def _initial_guess_for( u: _Function | typing.Sequence[_Function], -) -> _Function | typing.Sequence[_Function]: +) -> Function | typing.Sequence[Function]: """Build the solution vector a block solves into when its forward is replayed. 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 - {py:class}`dolfinx.fem.Function` cannot do. Outside checkpointing nothing asks, which is why a plain - one survived for so long. + {py:class}`dolfinx.fem.Function` cannot do. Outside checkpointing nothing asks, which is + why a plain one survived for so long. + + The two types are deliberately not the same. What comes in is only known to be a + {py:class}`dolfinx.fem.Function`, since the overloaded type is a subclass and either may be + passed, so that is what the check is against. What goes out is always the overloaded type, + because it ends up on the tape. """ with stop_annotating(): if isinstance(u, _Function): From dc2317aedde0a1c456cde4259fd1cb5402c4e77e Mon Sep 17 00:00:00 2001 From: Henrik Finsberg Date: Fri, 28 Aug 2026 09:58:00 +0000 Subject: [PATCH 8/9] Fix test_checkpointing so that uh doesn't enter the forms --- tests/test_checkpointing.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/tests/test_checkpointing.py b/tests/test_checkpointing.py index f916bc5..311e1a2 100644 --- a/tests/test_checkpointing.py +++ b/tests/test_checkpointing.py @@ -120,8 +120,14 @@ def _tape_heat_equation(n_steps, schedule=None, disk=False, use_mpio=None): v = ufl.TestFunction(V) f = dolfinx_adjoint.Function(V, name="source") uh = dolfinx_adjoint.Function(V, name="solution") + u_prev = dolfinx_adjoint.Function(V, name="previous") - F = ((u - uh) / dt * v + nu * ufl.inner(ufl.grad(u), ufl.grad(v)) - f * v) * ufl.dx + # The unknown and the previous state must be distinct functions: the block that solves + # for uh must not also depend on uh through its own form, or recompute under a checkpoint + # schedule would read whatever value happens to be in uh at replay time instead of the + # checkpointed one. The state update is therefore an explicit, tape-recorded assignment, + # matching _tape_snes_heat_equation below. + F = ((u - u_prev) / dt * v + nu * ufl.inner(ufl.grad(u), ufl.grad(v)) - f * v) * ufl.dx a, L = ufl.system(F) mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim) @@ -144,6 +150,7 @@ def _tape_heat_equation(n_steps, schedule=None, disk=False, use_mpio=None): for i in tape.timestepper(iter(range(n_steps))): dolfinx_adjoint.assign(controls[i], f) problem.solve() + dolfinx_adjoint.assign(uh, u_prev) J = J + dolfinx_adjoint.assemble_scalar(dt * uh**2 * ufl.dx) rf = pyadjoint.ReducedFunctional(J, [pyadjoint.Control(c) for c in controls]) From 0ccac31c545dba8783880e49dc0d2e34eb223616 Mon Sep 17 00:00:00 2001 From: Henrik Finsberg Date: Fri, 28 Aug 2026 10:22:15 +0000 Subject: [PATCH 9/9] Ignore .worktrees/ used for isolated feature branches Co-Authored-By: Claude Sonnet 5 --- .gitignore | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.gitignore b/.gitignore index 0e95bcd..2029210 100644 --- a/.gitignore +++ b/.gitignore @@ -132,3 +132,6 @@ _build/ *.dot *.bp + +# git worktrees used for isolated feature branches +/.worktrees/