Move temporary path cleanup to a separate thread and not block checkpoint manager init. We now wait for it when we try to do the first save post init.

PiperOrigin-RevId: 842605753

Author: mridul-sahu

checkpoint/orbax/checkpoint/_src/path/temporary_paths.py

@@ -156,11 +156,3 @@ async def cleanup_temporary_paths(
     await asyncio.gather(
         *[async_path.rmtree(tmp_path.get()) for tmp_path in tmp_paths]
     )
-  multihost.sync_global_processes(
-      multihost.unique_barrier_key(
-          'cleanup_tmp_dirs',
-          prefix=multiprocessing_options.barrier_sync_key_prefix,
-      ),
-      timeout=multihost.coordination_timeout(),
-      processes=multiprocessing_options.active_processes,
-  )

checkpoint/orbax/checkpoint/checkpoint_manager.py

@@ -34,7 +34,6 @@
 from orbax.checkpoint import checkpoint_args
 from orbax.checkpoint import options as options_lib
 from orbax.checkpoint import utils
-from orbax.checkpoint._src import asyncio_utils
 from orbax.checkpoint._src import threading as threading_lib
 from orbax.checkpoint._src.checkpoint_managers import policy_checkpoint_info
 from orbax.checkpoint._src.checkpoint_managers import preservation_policy as preservation_policy_lib
@@ -858,8 +857,9 @@ def __init__(
 
 
     # Cleanup directories from previous runs that may not have been finalized.
+    self._cleanup_tmp_directory_future = None
     if self._options.cleanup_tmp_directories:
-      asyncio_utils.run_sync(
+      self._cleanup_tmp_directory_future = future.CommitFuture(
           temporary_paths.cleanup_temporary_paths(
               self._directory,
               multiprocessing_options=self._options.multiprocessing_options,
@@ -928,6 +928,14 @@ def __init__(
         self,
     )
 
+  def _maybe_await_cleanup_tmp_directory(self):
+    if self._cleanup_tmp_directory_future is None:
+      return
+
+    self._cleanup_tmp_directory_future.result()
+    # Reset the future to None to avoid waiting for cleanup again.
+    self._cleanup_tmp_directory_future = None
+
   def _configure_checkpointer_common(
       self,
       handler: CompositeCheckpointHandler,
@@ -1388,7 +1396,8 @@ def save(
     self._validate_args(items, args)
     if not force and not self.should_save(step):
       return False
-
+    # Wait for any ongoing temporary path cleanup before starting the save.
+    self._maybe_await_cleanup_tmp_directory()
     multihost.sync_global_processes(
         multihost.unique_barrier_key(
             'CheckpointManager:save_start',

checkpoint/orbax/checkpoint/checkpoint_manager_test.py

@@ -928,20 +928,6 @@ def test_save_preempted(self, enable_async, prefix):
             cleanup_tmp_directories=True,
         ),
     ) as manager:
-      # Temp checkpoints cleaned up at creation.
-      self.assertFalse(tmp_dir.exists())
-      self.assertSameElements([], manager.all_steps())
-
-      # Sync to check directories before a new tmp dir is created.
-      test_utils.sync_global_processes('test_check_dirs_after_cleanup')
-
-      tmp_dir = test_utils.save_fake_tmp_dir(
-          self.directory, 0, 'params', subdirs=['subdir'], step_prefix=prefix
-      )
-      self.assertTrue(tmp_dir.exists())
-      tmp_dir_items = list(tmp_dir.iterdir())
-      self.assertLen(tmp_dir_items, 1)
-      self.assertIn('subdir', tmp_dir_items[0].name)
       self.assertSameElements(
           [], manager.all_steps()
       )  # Only picks up finalized.
@@ -1130,11 +1116,8 @@ def test_save_preempted_gcs(self, enable_async, is_gcs_path):
             enable_async_checkpointing=enable_async,
         ),
     ) as manager:
-      # Temp checkpoints cleaned up at creation.
-      self.assertFalse(tmp_dir.exists())
-      self.assertFalse((tmp_dir / 'subdir').exists())
       self.assertSameElements([], manager.all_steps())
-
+      manager._maybe_await_cleanup_tmp_directory()
       # Sync to check directories before a new tmp dir is created.
       test_utils.sync_global_processes('test_check_dirs_after_cleanup')