Optional garbage collection and CheckpointManager._global_deps (#187)

* Optional gc_collect

* Add global_deps

* Add _adj_deps_cleaned into TimeStep
---------

Co-authored-by: Connor Ward <c.ward20@imperial.ac.uk>

Author: Ig-dolci

pyadjoint/checkpointing.py

@@ -1,5 +1,7 @@
 from enum import Enum
 import sys
+import gc
+import logging
 from functools import singledispatchmethod
 from checkpoint_schedules import Copy, Move, EndForward, EndReverse, \
     Forward, Reverse, StorageType, SingleMemoryStorageSchedule
@@ -37,6 +39,14 @@ class CheckpointManager:
     Args:
         schedule (checkpoint_schedules.schedule): A schedule provided by the `checkpoint_schedules` package.
         tape (Tape): A list of blocks :class:`Block` instances.
+        gc_timestep_frequency (None or int): The number of timesteps between garbage collections. The default
+        is `None`, which means no invoking the garbage collector during the executions. If an integer is
+        provided, the garbage collector is applied every `gc_timestep_frequency` timestep. That is useful when
+        being affected by the Python fails to track and clean all checkpoint objects in memory properly.
+        gc_generation (int): The generation for garbage collection. Default is 2 that runs a full collection.
+        To have more information about the garbage collector generation,
+        please refer to the `documentation
+        <https://docs.python.org/3/library/gc.html#gc.collect>`_.
 
     Attributes:
         tape (Tape): A list of blocks :class:`Block` instances.
@@ -52,7 +62,7 @@ class CheckpointManager:
         _current_action (checkpoint_schedules.CheckpointAction): The current `checkpoint_schedules` action.
 
     """
-    def __init__(self, schedule, tape):
+    def __init__(self, schedule, tape, gc_timestep_frequency=None, gc_generation=2):
         if (
             schedule.uses_storage_type(StorageType.DISK)
             and not tape._package_data
@@ -78,9 +88,15 @@ def __init__(self, schedule, tape):
         self.forward_schedule.append(self._current_action)
         # Tell the tape to only checkpoint input data until told otherwise.
         self.tape.latest_checkpoint = 0
-        self.end_timestep(-1)
         self._keep_init_state_in_work = False
-        self._adj_deps_cleaned = False
+        self._gc_timestep_frequency = gc_timestep_frequency
+        self._gc_generation = gc_generation
+        # ``self._global_deps`` stores checkpoint dependencies that remain unchanged across
+        # timesteps (``self.tape.timesteps``). During the forward taping process, the code
+        # checks whether a dependency is in ``self._global_deps`` to avoid unnecessary clearing
+        # and recreation of its checkpoint data.
+        self._global_deps = set()
+        self.end_timestep(-1)
 
     def end_timestep(self, timestep):
         """Mark the end of one timestep when taping the forward model.
@@ -164,11 +180,30 @@ def _(self, cp_action, timestep):
             ):
                 for package in self.tape._package_data.values():
                     package.continue_checkpointing()
+            if timestep == 1:
+                # Store the possible global dependencies.
+                for deps in self.tape.timesteps[timestep - 1].checkpointable_state:
+                    self._global_deps.add(deps)
+            else:
+                # Check if the block variables stored in `self._global_deps` are still
+                # dependencies in the previous timestep. If not, remove them from the
+                # global dependencies.
+                deps_to_clear = self._global_deps.difference(
+                    self.tape.timesteps[timestep - 1].checkpointable_state)
+
+                # Remove the block variables that are not global dependencies.
+                self._global_deps.difference_update(deps_to_clear)
+
+                # For no global dependencies, checkpoint storage occurs at a self.tape
+                # timestep only when required by an action from the schedule. Thus, we
+                # have to clear the checkpoint of block variables excluded from the self._global_deps.
+                for deps in deps_to_clear:
+                    deps._checkpoint = None
 
             self.tape.timesteps[timestep - 1].checkpoint(
-                _store_checkpointable_state, _store_adj_dependencies)
+                _store_checkpointable_state, _store_adj_dependencies, self._global_deps)
             # Remove unnecessary variables in working memory from previous steps.
-            for var in self.tape.timesteps[timestep - 1].checkpointable_state:
+            for var in self.tape.timesteps[timestep - 1].checkpointable_state - self._global_deps:
                 var._checkpoint = None
             for block in self.tape.timesteps[timestep - 1]:
                 for out in block.get_outputs():
@@ -182,6 +217,10 @@ def _(self, cp_action, timestep):
                 # Activate disk checkpointing only in the checkpointing process.
                 for package in self.tape._package_data.values():
                     package.pause_checkpointing()
+
+            if isinstance(self._gc_timestep_frequency, int) and timestep % self._gc_timestep_frequency == 0:
+                logging.info("Running a garbage collection cycle")
+                gc.collect(self._gc_generation)
             return True
         else:
             return False
@@ -300,7 +339,7 @@ def _(self, cp_action, progress_bar, functional=None, **kwargs):
                     for package in self.tape._package_data.values():
                         package.continue_checkpointing()
                 current_step.checkpoint(
-                    _store_checkpointable_state, _store_adj_dependencies)
+                    _store_checkpointable_state, _store_adj_dependencies, self._global_deps)
 
             to_keep = set()
             if step < (self.total_timesteps - 1):
@@ -310,7 +349,7 @@ def _(self, cp_action, progress_bar, functional=None, **kwargs):
             if functional:
                 to_keep = to_keep.union([functional.block_variable])
 
-            for var in current_step.checkpointable_state - to_keep:
+            for var in current_step.checkpointable_state - to_keep.union(self._global_deps):
                 # Handle the case where step is 0
                 if step == 0 and var not in current_step._checkpoint:
                     # Ensure initialisation state is kept.
@@ -346,6 +385,10 @@ def _(self, cp_action, progress_bar, functional=None, **kwargs):
                         if bv not in current_step.adjoint_dependencies.union(to_keep):
                             bv._checkpoint = None
 
+            if self._gc_timestep_frequency and step % self._gc_timestep_frequency == 0:
+                logging.info("Running a garbage collection cycle")
+                gc.collect(self._gc_generation)
+
             step += 1
             if cp_action.storage == StorageType.DISK:
                 # Activate disk checkpointing only in the checkpointing process.
@@ -361,11 +404,11 @@ def _(self, cp_action, progress_bar, markings, functional=None, **kwargs):
             current_step = self.tape.timesteps[step]
             for block in reversed(current_step):
                 block.evaluate_adj(markings=markings)
-                if not self._adj_deps_cleaned:
+                if not current_step._adj_deps_cleaned:
                     for out in block._outputs:
                         if not out.marked_in_path:
                             current_step.adjoint_dependencies.discard(out)
-                    self._adj_deps_cleaned = True
+            current_step._adj_deps_cleaned = True
             # Output variables are used for the last time when running
             # backwards.
             to_keep = current_step.checkpointable_state
@@ -379,6 +422,9 @@ def _(self, cp_action, progress_bar, markings, functional=None, **kwargs):
                         out.reset_variables(("adjoint", "hessian"))
                     if cp_action.clear_adj_deps and out not in to_keep:
                         out._checkpoint = None
+            if self._gc_timestep_frequency and step % self._gc_timestep_frequency == 0:
+                logging.info("Running a garbage collection cycle")
+                gc.collect(self._gc_generation)
 
     @process_operation.register(Copy)
     def _(self, cp_action, progress_bar, **kwargs):

pyadjoint/tape.py

@@ -304,7 +304,7 @@ def add_to_adjoint_dependencies(self, block_var, last_used):
         for step in self.timesteps[last_used + 1:]:
             step.adjoint_dependencies.add(block_var)
 
-    def enable_checkpointing(self, schedule):
+    def enable_checkpointing(self, schedule, gc_timestep_frequency=None, gc_generation=2):
         """Enable checkpointing on the adjoint evaluation.
 
         A checkpoint manager able to execute the forward and adjoint computations
@@ -313,12 +313,23 @@ def enable_checkpointing(self, schedule):
         Args:
             schedule (checkpoint_schedules.schedule): A schedule provided by the
             checkpoint_schedules package.
+            gc_timestep_frequency (None or int): The timestep frequency for garbage collection.
+            For additional information, please refer to the :class:`CheckpointManager`
+            documentation.
+            gc_generation (int): The generation for garbage collection. For additional
+            information, please refer to the :class:`CheckpointManager` documentation.
         """
         if self._blocks:
             raise CheckpointError(
                 "Checkpointing must be enabled before any blocks are added to the tape."
             )
-        self._checkpoint_manager = CheckpointManager(schedule, self)
+
+        if gc_timestep_frequency is not None and not isinstance(gc_timestep_frequency, int):
+            raise CheckpointError("gc_timestep_frequency must be an integer.")
+
+        self._checkpoint_manager = CheckpointManager(
+            schedule, self, gc_timestep_frequency=gc_timestep_frequency,
+            gc_generation=gc_generation)
 
     def get_blocks(self, tag=None):
         """Returns a list of the blocks on the tape.
@@ -782,25 +793,36 @@ def __init__(self, blocks=()):
         # A dictionary mapping the block variables in the checkpointable state
         # to their checkpoint values.
         self._checkpoint = {}
+        # A flag to indicate whether the adjoint dependencies have been cleaned
+        # from the outputs not marked in the path.
+        self._adj_deps_cleaned = False
 
     def copy(self, blocks=None):
         out = TimeStep(blocks or self)
         out.checkpointable_state = self.checkpointable_state
         return out
 
-    def checkpoint(self, checkpointable_state, adj_dependencies):
+    def checkpoint(self, checkpointable_state, adj_dependencies, global_deps):
         """Store a copy of the checkpoints in the checkpointable state.
 
         Args:
             checkpointable_state (bool): If True, store the checkpointable state
             required to restart from the start of a timestep.
             adj_dependencies (bool): If True, store the adjoint dependencies required
             to compute the adjoint of a timestep.
+            global_deps (set): This set stores the common dependencies for all timesteps.
+            For additional information, please refer to the :class:`CheckpointManager`
+            documentation.
         """
         with stop_annotating():
             if checkpointable_state:
                 for var in self.checkpointable_state:
-                    self._checkpoint[var] = var.saved_output._ad_create_checkpoint()
+                    if var in global_deps:
+                        # Creating a new checkpoint object is not necessary here
+                        # because the global dependencies do not change.
+                        self._checkpoint[var] = var._checkpoint
+                    else:
+                        self._checkpoint[var] = var.saved_output._ad_create_checkpoint()
 
             if adj_dependencies:
                 for var in self.adjoint_dependencies: