add snapshot_controller in elastic manager for resume dataloading with grain

PiperOrigin-RevId: 747478080

Author: aireenmei

pathwaysutils/elastic/manager.py

@@ -25,8 +25,10 @@
 - Resharding the snapshot.
 """
 
+import sys
 import collections
 from collections.abc import Callable, Mapping, Sequence
+import copy
 import itertools
 import logging
 import traceback
@@ -357,7 +359,7 @@ def _slice_down(self, reshard_retry: bool = False) -> None:
       )
 
   # TODO b/407772100 - Support multiple snapshots.
-  def pop_snapshot(self) -> tuple[int, PyTree]:
+  def pop_snapshot(self) -> tuple[int, PyTree, Any]:
     """Pops next snapshot.
 
     This function is used to get the next snapshot and remove it from
@@ -375,9 +377,13 @@ def pop_snapshot(self) -> tuple[int, PyTree]:
 
     step = self._snapshot.pop("step")
     snapshot = self._snapshot.pop("snapshot")
+    if "snapshot_controller" in self._snapshot:
+      snapshot_controller = self._snapshot.pop("snapshot_controller")
+    else:
+      snapshot_controller = None
     self._snapshot = None
 
-    return step, snapshot
+    return step, snapshot, snapshot_controller
 
   @staticmethod
   def _get_snapshot_size(snapshot: PyTree) -> int:
@@ -412,11 +418,18 @@ def _put_snapshot_on_host(
         may_alias=False,
     )
 
+  @staticmethod
+  def _put_snapshot_on_controller(
+      snapshot: PyTree,
+  ) -> PyTree:
+    return copy.deepcopy(snapshot)
+
   @timing.timeit
   def maybe_snapshot(
       self,
       step: int,
       snapshot: PyTree,
+      snapshot_controller: Mapping[str, Any] | None = None,
       force: bool = False,
       block: bool = False,
   ) -> None:
@@ -438,7 +451,7 @@ def maybe_snapshot(
 
     total_nbytes = self._get_snapshot_size(snapshot)
 
-    _logger.info("Saving a snapshot of %s bytes", total_nbytes)
+    _logger.info("Saving a snapshot of %s bytes on host", total_nbytes)
 
     snapshot_host = self._put_snapshot_on_host(snapshot)
     _logger.info("Snapshot dispatched")
@@ -448,12 +461,22 @@ def maybe_snapshot(
       _logger.info("Snapshot completed")
 
     # TODO b/407772100 - Support multiple snapshots.
-    self._snapshot = {"step": step, "snapshot": snapshot_host}
+    self._snapshot = {
+        "step": step,
+        "snapshot": snapshot_host,
+    }
+    if snapshot_controller is not None:
+      total_nbytes = sys.getsizeof(snapshot_controller)
+      _logger.info("Saving a snapshot of %s bytes on controller", total_nbytes)
+      snapshot_on_controller = self._put_snapshot_on_controller(
+          snapshot_controller
+      )
+      self._snapshot["snapshot_controller"] = snapshot_on_controller
 
   @timing.timeit
   def get_resharded_snapshot(
       self, mesh: jax.sharding.Mesh
-  ) -> tuple[int, Mapping[str, int | PyTree]]:
+  ) -> tuple[int, Mapping[str, int | PyTree], Any]:
     """Get the resharded snapshot.
 
     The snapshot on pinned memory is resharded to the new mesh. This snapshot is
@@ -466,7 +489,7 @@ def get_resharded_snapshot(
     Returns:
       The next step and snapshot resharded to the new mesh.
     """
-    step, snapshot = self.pop_snapshot()
+    step, snapshot, snapshot_controller = self.pop_snapshot()
 
     sharding_pinned_host = jax.tree.map(
         lambda x: jax.sharding.NamedSharding(
@@ -480,7 +503,11 @@ def get_resharded_snapshot(
         donate=True,
         may_alias=False,
     )
-    self._snapshot = {"step": step, "snapshot": resharded_snapshot_pinned_host}
+    self._snapshot = {
+        "step": step,
+        "snapshot": resharded_snapshot_pinned_host,
+        "snapshot_controller": snapshot_controller,
+    }
 
     sharding_device = jax.tree.map(
         lambda x: x.sharding.with_memory_kind("device"),
@@ -493,7 +520,7 @@ def get_resharded_snapshot(
         may_alias=False,
     )
 
-    return step, resharded_snapshot_device
+    return step, resharded_snapshot_device, snapshot_controller
 
   @timing.timeit
   def maybe_reshard_down(
@@ -561,6 +588,7 @@ def maybe_reshard_up(
       step: int,
       snapshot: Mapping[str, int | PyTree],
       elastic_handler: Callable[..., Any],
+      snapshot_controller: Mapping[str, Any] | None = None,
       handler_args: tuple[Any, ...] | None = None,
       handler_kwargs: Mapping[str, Any] | None = None,
   ) -> Any:
@@ -596,6 +624,7 @@ def maybe_reshard_up(
     self.maybe_snapshot(
         step=step,
         snapshot=snapshot,
+        snapshot_controller=snapshot_controller,
         force=True,
         block=True,
     )