#p2p Implement CheckpointManager.reload() properly

PiperOrigin-RevId: 876307566

Author: Orbax Authors

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

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Composite Checkpoint Manager handling P2P syncing with optional Persistent Fallback."""
+"""Composite Checkpoint Manager with P2P syncing and Persistent Fallback."""
 
 import shutil
 import threading
@@ -178,7 +178,7 @@ def get_all_steps_from_peers(self) -> list[int]:
     return self._peer_selector.get_all_steps()
 
   def has_shard_for_step(self, step: int) -> bool:
-    """Checks if this process's shard for a given step exists in the P2P network."""
+    """Checks if this process's shard for a step exists in P2P."""
     assert self._peer_selector is not None
     return (
         self._peer_selector.get_source_peer(step, self._process_index)
@@ -212,7 +212,7 @@ def close(self):
 class CheckpointManager(
     abstract_checkpoint_manager.AbstractCheckpointManager, epy.ContextManager
 ):
-  """Orchestrates P2P local checkpointing with optional persistent storage failover.
+  """P2P local checkpointing with persistent storage failover.
 
   Restoration Strategy:
     1. Check Local Disk (Fastest)
@@ -333,7 +333,15 @@ def best_step(self) -> int | None:
 
   @override
   def reload(self):
+    """Reloads the checkpoint manager and its components.
+
+    This method refreshes the local and persistent managers and marks the P2P
+    registry as stale, forcing a re-sync on the next access.
+    """
+    self._p2p.mark_registry_stale()
     self._local_manager.reload()
+    if self._persistent_manager:
+      self._persistent_manager.reload()
 
   @override
   def reached_preemption(self, step: int) -> bool:

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

@@ -215,5 +215,9 @@ def wait_until_finished(self):
   def check_for_errors(self):
     self._manager.check_for_errors()
 
+  def reload(self):
+    """Reloads the internal checkpoint manager."""
+    self._manager.reload()
+
   def close(self):
     self._manager.close()

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

@@ -18,6 +18,7 @@
 import json
 import socket
 import struct
+import typing
 from typing import Any, Final
 from absl import logging
 from etils import epath
@@ -60,6 +61,16 @@ def from_dict(cls, data: dict[str, Any]) -> Self:
     )
 
 
+class ManifestEntry(typing.TypedDict):
+  """Type definition for a single file entry in a manifest."""
+
+  rel_path: str
+  size: int
+
+
+Manifest = list[ManifestEntry]
+
+
 def optimize_socket(sock: socket.socket) -> None:
   try:
     sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)

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

@@ -143,9 +143,7 @@ def stop(self):
     self._thread.join(timeout=2.0)
     self._thread = None
 
-  def handle_get_manifest(
-      self, payload: dict[str, Any]
-  ) -> list[dict[str, Any]]:
+  def handle_get_manifest(self, payload: dict[str, Any]) -> protocol.Manifest:
     """Handles GET_MANIFEST request.
 
     Args:
@@ -194,6 +192,7 @@ def handle_get_manifest(
           step,
           req_process_index,
       )
+
     return files
 
   def handle_download(self, sock, payload: dict[str, Any]):
@@ -237,7 +236,7 @@ def fetch_shard_from_peer(
     """
     logging.info('Requesting manifest from %s:%d for step %d', ip, port, step)
 
-    manifest = protocol.TCPClient.request(
+    manifest: protocol.Manifest = protocol.TCPClient.request(
         ip,
         port,
         protocol.OP_GET_MANIFEST,