Make cleanup_tmp_directories method a class method for CheckpointMana…

…ger to allow for customization.

PiperOrigin-RevId: 474902053

orbax/checkpoint/checkpoint_manager.py

@@ -152,7 +152,7 @@ def __init__(
     if self._options.best_mode not in ['min', 'max']:
       raise ValueError('`best_mode` must be one of: "min", "max"')
     # Cleanup directories from previous runs that may not have been finalized.
-    utils.cleanup_tmp_directories(self._directory)
+    self._cleanup_tmp_directories()
     self._checkpoints = self._create_checkpoints()
     # The distinction between _last_checkpoint and _last_preserved_checkpoint is
     # necessary because some checkpoints may not be kept in the long run, in
@@ -539,6 +539,9 @@ def _sort_checkpoints_by_metrics(
         key=lambda info: self._options.best_fn(info.metrics),
         reverse=(self._options.best_mode == 'min'))
 
+  def _cleanup_tmp_directories(self):
+    utils.cleanup_tmp_directories(self.directory)
+
   def _delete_directory(self, step: int):
     if jax.process_index() == 0:
       utils.rmtree(self._get_save_directory(step, self.directory))

setup.py

@@ -22,7 +22,7 @@
 
 setuptools.setup(
     name='orbax',
-    version='0.0.9',
+    version='0.0.10',
     description='Orbax',
     long_description=_LONG_DESCRIPTION,
     long_description_content_type='text/markdown',