Add a replica_resize decorator for fault tolerance in elastic Pathways.

PiperOrigin-RevId: 857246882

Author: lukebaumann

pathwaysutils/elastic/manager.py

@@ -36,14 +36,21 @@ class ElasticRuntimeError(RuntimeError):
   """Error raised when elasticity cannot continue."""
 
 
-class NewSliceAvailableError(RuntimeError):
-  """Error raised when a new slice is available."""
+class ScaleUpError(RuntimeError):
+  """Signals that the workload is ready to scale up.
+
+  This exception should be raised by user code when it detects that new hardware
+  is available and it wants to restart computation to make use of it.
+  Raising this exception will interrupt the current computation and cause the
+  elasticity manager to retry it with an updated slice configuration that
+  includes the new hardware.
+  """
 
 
 _F = TypeVar("_F", bound=Callable[..., Any])
 
 
-def _elastic_event_cleanup():
+def _elastic_event_cleanup() -> None:
   """Cleans up JAX profiles, caches, and live arrays."""
   try:
     _logger.info("Cleaning up any ongoing traces")
@@ -60,10 +67,19 @@ def _elastic_event_cleanup():
 
 
 class Manager:
-  """Utility class for elastic training."""
+  """Utility class for elastic training.
+
+  Attributes:
+    slice_to_devices: A mapping from slice index to a sequence of `jax.Device`
+      objects for that slice.
+    all_slice_indices: A set of all possible slice indices.
+    active_slice_indices: A set of indices of the currently active slices.
+    new_slice_event: A `threading.Event` that is set when new slices become
+      available during replica/resize mode.
+  """
 
-  _total_slice_count: int | None = None
   slice_to_devices: Mapping[int, Sequence[jax.Device]]
+  all_slice_indices: set[int]
   active_slice_indices: set[int]
   new_slice_event: threading.Event
 
@@ -84,16 +100,14 @@ def __init__(self, devices: Sequence[jax.Device] | None = None) -> None:
     )
     self.new_slice_event = threading.Event()
 
-  @property
+  @functools.cached_property
   def total_slice_count(self) -> int:
-    """Returns the total number of slices."""
-    if self._total_slice_count is None:
-      self._total_slice_count = len(self.slice_to_devices)
-    return self._total_slice_count
+    """The total number of slices."""
+    return len(self.slice_to_devices)
 
   @property
   def default_device(self) -> jax.Device:
-    """Returns the device that should be set to the default device.
+    """The device that should be set to the default device.
 
     This will be from one of the slices in `active_slice_indices`.
     """
@@ -104,12 +118,12 @@ def default_device(self) -> jax.Device:
 
   @property
   def active_slice_count(self) -> int:
-    """Returns the number of slices."""
+    """The number of active slices."""
     return len(self.active_slice_indices)
 
   @property
   def inactive_slice_indices(self) -> set[int]:
-    """Returns the set of inactive slice indices."""
+    """The set of inactive slice indices."""
     return self.all_slice_indices - self.active_slice_indices
 
   def scale_by_active_slices(self, x: int | float) -> int | float:
@@ -149,11 +163,24 @@ def _elasticity_retry_decorator(
       pre_callback: Callable[..., Any] | None = None,
       on_elastic_event_callback: Callable[..., Any] | None = None,
   ) -> Callable[[_F], _F]:
-    """Retries a function with elasticity fault tolerance."""
+    """Retries a function with elasticity fault tolerance.
 
-    def decorator(func):
+    Args:
+      max_retries: The maximum number of times to retry the function.
+      pre_callback: A callback to call before each attempt of the wrapped
+        function.
+      on_elastic_event_callback: A callback to call after an elastic failure
+        occurs.
+
+    Returns:
+      A function decorator.
+    """
+
+    if max_retries <= 0:
+      raise ValueError("max_retries must be positive.")
+    def decorator(func: _F) -> _F:
       @functools.wraps(func)
-      def wrapper(*args, **kwargs):
+      def wrapper(*args: Any, **kwargs: Any) -> Any:
         for retry_index in range(max_retries):
           try:
             _logger.info(
@@ -164,8 +191,8 @@ def wrapper(*args, **kwargs):
 
             with jax.default_device(self.default_device):
               return func(*args, **kwargs)
-          except NewSliceAvailableError:
-            _logger.info("New slice available. Retrying.")
+          except ScaleUpError:
+            _logger.info("Scale up requested. Retrying.")
             _elastic_event_cleanup()
 
             if on_elastic_event_callback is not None:
@@ -185,9 +212,10 @@ def wrapper(*args, **kwargs):
 
             if on_elastic_event_callback is not None:
               on_elastic_event_callback()
-        raise ElasticRuntimeError(
-            f"Elastic attempt {max_retries} out of {max_retries} failed."
-        )
+        else:
+          raise ElasticRuntimeError(
+              f"Elastic attempt {max_retries} out of {max_retries} failed."
+          )
 
       return wrapper
 
@@ -226,7 +254,7 @@ def pause_resume(
         occurs.
 
     Returns:
-      The result of the wrapped function.
+      A decorator that retries the wrapped function.
 
     Raises:
       ElasticRuntimeError: If all retry attempts fail.
@@ -300,7 +328,7 @@ def replica_resize(
         occurs.
 
     Returns:
-      The result of the wrapped function.
+      A decorator that retries the wrapped function.
 
     Raises:
       ElasticRuntimeError: If all retry attempts fail.