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/ 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/_toc.yml b/_toc.yml index 582d81e..a0ad400 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" - file: "demos/demo_nonmatching_grids.py" - caption: Python API chapters: diff --git a/demos/time_distributed_control_checkpointing.py b/demos/time_distributed_control_checkpointing.py new file mode 100644 index 0000000..9798b0f --- /dev/null +++ b/demos/time_distributed_control_checkpointing.py @@ -0,0 +1,196 @@ +# # 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}`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. + +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. {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) +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) + + # `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) + 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, strict=True): + np.testing.assert_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 +# +# 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 = [] +# 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) +assert rate > 1.9 + +# ## Storing checkpoints on disk +# +# {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 +# [io4dolfinx](https://github.com/scientificcomputing/io4dolfinx) 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, strict=True): + np.testing.assert_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.") + +# 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 +# :labelprefix: +# :keyprefix: tdcc- +# ``` 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/pyproject.toml b/pyproject.toml index 3bc88d4..7ff7d0c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,11 +9,15 @@ authors = [{ name = "Jørgen S. Dokken", email = "dokken@simula.no" }] license = "MIT" license-files = ["LICENSE"] readme = "README.md" +requires-python = ">=3.12" dependencies = [ "fenics-dolfinx>=0.10.0", "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/__init__.py b/src/dolfinx_adjoint/__init__.py index 2e8ad47..9ff4b1a 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, interpolate_nonmatching 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 549c8be..89ae950 100644 --- a/src/dolfinx_adjoint/blocks/solvers.py +++ b/src/dolfinx_adjoint/blocks/solvers.py @@ -486,15 +486,20 @@ def recompute_component( block_variable: pyadjoint.block_variable.BlockVariable, idx: int, prepared: dolfinx.fem.Function | typing.Sequence[dolfinx.fem.Function], - ) -> dolfinx.fem.Function: + ) -> Function: """Recompute and return an isolated copy of the solution state.""" if isinstance(prepared, dolfinx.fem.Function): assert idx == 0 - # Return an explicit copy so each tape block gets an isolated state snapshot - return prepared.copy() + out = prepared else: assert isinstance(prepared, typing.Iterable) - return prepared[idx].copy() + out = prepared[idx] + # Function.copy() always returns a plain dolfinx.fem.Function, which cannot checkpoint + # itself under a schedule; _ad_new_like() keeps the overloaded type on the tape. + assert isinstance(out, Function) + isolated = out._ad_new_like() + isolated.x.array[:] = out.x.array[:] + return isolated def _should_compute_boundary_adjoint( self, relevant_dependencies: typing.List[tuple[int, pyadjoint.block_variable.BlockVariable]] @@ -1161,7 +1166,12 @@ def recompute_component( else: output = self._forward_solver._u assert isinstance(output, Function) - return output + # self._forward_solver._u is warm-started and solved into in place on every replay of + # this block, so returning it directly would alias the tape output of every recompute + # to the same mutable object; isolate a snapshot instead. + isolated = output._ad_new_like() + isolated.x.array[:] = output.x.array[:] + return isolated def _should_compute_boundary_adjoint( self, relevant_dependencies: typing.List[tuple[int, pyadjoint.block_variable.BlockVariable]] diff --git a/src/dolfinx_adjoint/checkpointing.py b/src/dolfinx_adjoint/checkpointing.py new file mode 100644 index 0000000..8278653 --- /dev/null +++ b/src/dolfinx_adjoint/checkpointing.py @@ -0,0 +1,425 @@ +"""Snapshot checkpointing of DOLFINx functions to disk. + +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 +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 {py:mod}`io4dolfinx`. +""" + +from __future__ import annotations + +import os +import pathlib +import tempfile +import typing +import weakref + +from mpi4py import MPI + +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: "_DiskCheckpointer | None" = 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 _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. + 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. + + 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. + """ + n_local = function.x.array.size + if not shared_file: + return n_local, n_local, 0 + # 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: + """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: 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 + # 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 + + @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. + + 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: + """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: + """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: 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: weakref.ReferenceType | None = None + + def restore(self) -> Function: + """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) # 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. + 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: 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 + 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 = 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: 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) + 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: 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 + return _checkpointer.store(function) + + +def enable_disk_checkpointing( + dirname: str | os.PathLike | None = None, + comm: MPI.Intracomm | None = None, + cleanup: bool = True, + use_mpio: bool | None = 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 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. + """ + 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 + 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." + ) + # 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 = pathlib.Path(comm.bcast(created, root=0)) + else: + directory = pathlib.Path(typing.cast("str | os.PathLike", dirname)) + if comm.rank == 0: + directory.mkdir(parents=True, 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 3f92f4a..1502a6a 100644 --- a/src/dolfinx_adjoint/types/function.py +++ b/src/dolfinx_adjoint/types/function.py @@ -16,6 +16,7 @@ from ..blocks._vector import _SpecialVector, _vector from ..blocks.assembly import assemble_compiled_form +from ..checkpointing import SnapshotCheckpoint, maybe_disk_checkpoint from ..utils import function_from_vector, gather @@ -87,12 +88,19 @@ def _ad_init_object(cls, obj): return cls(obj.function_space, obj.x, obj.name) @property - def index_map(self) -> dolfinx.cpp.la.IndexMap: # type: ignore [name-defined] + def index_map(self) -> dolfinx.common.IndexMap: """Return the index map of the function's vector.""" return self.x.index_map @no_annotations def _ad_create_checkpoint(self): + # 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 @@ -103,6 +111,8 @@ def _ad_create_checkpoint(self): return checkpoint def _ad_restore_at_checkpoint(self, checkpoint): + if isinstance(checkpoint, SnapshotCheckpoint): + return checkpoint.restore() return checkpoint def _ad_dot(self, other: typing.Self, options: typing.Optional[dict] = None): diff --git a/src/dolfinx_adjoint/utils.py b/src/dolfinx_adjoint/utils.py index 6b0b01f..91a60d0 100644 --- a/src/dolfinx_adjoint/utils.py +++ b/src/dolfinx_adjoint/utils.py @@ -59,7 +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: - assert isinstance(func, dolfinx.fem.Function), "All operands in the linear combination must be Functions." + # 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[:] diff --git a/tests/test_checkpointing.py b/tests/test_checkpointing.py new file mode 100644 index 0000000..311e1a2 --- /dev/null +++ b/tests/test_checkpointing.py @@ -0,0 +1,319 @@ +"""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. +""" + +import gc + +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()) + _collect() + try: + yield + finally: + _collect() + 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 _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 = [] + 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 + + +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. + """ + _collect() + 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") + u_prev = dolfinx_adjoint.Function(V, name="previous") + + # 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) + 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() + 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) + + +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, strict=True)): + 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``. + """ + _collect() + 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, strict=True)): + 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())