#p2p Refactor service and protocol for better type safety

PiperOrigin-RevId: 876445289

Author: Orbax Authors

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,