#p2p Cleanup P2P restore directories after use

PiperOrigin-RevId: 866984150

Author: Orbax Authors

checkpoint/orbax/checkpoint/experimental/emergency/p2p/checkpoint_manager.py

@@ -14,6 +14,7 @@
 
 """Composite Checkpoint Manager handling P2P syncing with optional Persistent Fallback."""
 
+import shutil
 import threading
 import time
 from typing import Any, Iterable, Mapping, Optional, Sequence, Union, final
@@ -366,15 +367,27 @@ def _restore_from_local_or_p2p(
       return self._local_manager.restore(step, args=args)
     else:
       logging.info('Step %d not found locally, fetching from P2P.', step)
-      if self._p2p.fetch(step):
-        p2p_restore_dir = self._local_directory / constants.P2P_RESTORE_DIR_NAME
-        return self._local_manager.restore(
-            step, args=args, directory=p2p_restore_dir
-        )
-      else:
-        raise FileNotFoundError(
-            f'Failed to fetch step {step} from P2P network.'
-        )
+      p2p_restore_dir = self._local_directory / constants.P2P_RESTORE_DIR_NAME
+      try:
+        if self._p2p.fetch(step):
+          return self._local_manager.restore(
+              step, args=args, directory=p2p_restore_dir
+          )
+        else:
+          raise FileNotFoundError(
+              f'Failed to fetch step {step} from P2P network.'
+          )
+      finally:
+        if p2p_restore_dir.exists():
+          logging.info(
+              'Removing P2P restore directory: %s after restoration is'
+              ' complete',
+              p2p_restore_dir,
+          )
+          try:
+            shutil.rmtree(str(p2p_restore_dir))
+          except OSError as e:
+            logging.exception('Failed to remove P2P restore directory: %s', e)
 
   @override
   def restore(

checkpoint/orbax/checkpoint/experimental/emergency/p2p/checkpoint_manager_test.py

@@ -402,6 +402,38 @@ def test_restore_no_step_in_p2p_but_in_persistent(self, _):
 
     manager.close()
 
+  @mock.patch.object(p2p_cm.shutil, 'rmtree', autospec=True)
+  @mock.patch.object(multihost, 'process_index', return_value=0)
+  def test_restore_p2p_cleanup(self, unused_process_index, mock_rmtree):
+    """Tests that P2P restore directory is cleaned up after restore."""
+    self.local_manager_instance.scan_stored_steps.return_value = (0, [])
+    self.local_manager_instance.all_steps.return_value = []
+    self.mock_sync_global_data.return_value = []
+
+    # P2P fetch succeeds
+    self.peer_selector_instance.get_source_peer.return_value = (
+        protocol.PeerDiscoveryInfo(
+            ip='1.2.3.4', port=5678, process_index=1, steps=[1]
+        )
+    )
+    self.p2p_node_instance.fetch_shard_from_peer.return_value = True
+    self.local_manager_instance.restore.return_value = {'a': 1}
+
+    manager = p2p_cm.CheckpointManager(
+        self.mesh,
+        self.abstract_state,
+        self.local_dir,
+    )
+
+    # Make p2p_restore_dir exist so cleanup is triggered
+    p2p_restore_dir = self.local_dir / service.constants.P2P_RESTORE_DIR_NAME
+    p2p_restore_dir.mkdir()
+
+    manager.restore(1, args=self.restore_args)
+
+    mock_rmtree.assert_called_once_with(str(p2p_restore_dir))
+    manager.close()
+
 
 if __name__ == '__main__':
   absltest.main()

checkpoint/orbax/checkpoint/experimental/emergency/p2p/service.py

@@ -256,8 +256,6 @@ def fetch_shard_from_peer(
       )
       return False
 
-    # TODO(exlin): Remove this directory once the transfer is globally completed
-    # to save memory space.
     stage_dir = self.directory / f'stage_{step}_{stored_process_index}'
     if stage_dir.exists():
       shutil.rmtree(str(stage_dir))

checkpoint/orbax/checkpoint/experimental/emergency/p2p/service_test.py

@@ -333,6 +333,24 @@ def download_side_effect(unused_ip, unused_port, rel_path, dest_path):
     mock_move.assert_called_once_with(str(stage_dir / '1'), str(final_dir))
     mock_rmtree.assert_called_with(str(stage_dir), ignore_errors=True)
 
+  @mock.patch.object(service.shutil, 'rmtree', autospec=True)
+  @mock.patch.object(service.protocol.TCPClient, 'request', autospec=True)
+  @mock.patch.object(service.protocol.TCPClient, 'download', autospec=True)
+  def test_fetch_shard_from_peer_exception_cleanup(
+      self,
+      mock_download,
+      mock_request,
+      mock_rmtree,
+  ):
+    """Tests that stage_dir is cleaned up if an exception occurs."""
+    mock_request.return_value = [{'rel_path': '1/file1', 'size': 10}]
+    mock_download.side_effect = OSError('Download failed')
+
+    self.assertFalse(self.node.fetch_shard_from_peer('peer', 123, 1, 10))
+
+    stage_dir = self.temp_dir / 'stage_1_10'
+    mock_rmtree.assert_called_with(str(stage_dir), ignore_errors=True)
+
 
 if __name__ == '__main__':
   absltest.main()