pathwaysutils 0.1.9__tar.gz → 0.1.11__tar.gz

{pathwaysutils-0.1.9 → pathwaysutils-0.1.11}/PKG-INFO

@@ -1,10 +1,10 @@
 Metadata-Version: 2.4
 Name: pathwaysutils
-Version: 0.1.9
+Version: 0.1.11
 Summary: Package of Pathways-on-Cloud utilities.
 Keywords: 
 Author-email: Pathways-on-Cloud Utilities Developers <pathwaysutils-dev@google.com>
-Requires-Python: >=3.10
+Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 Classifier: License :: OSI Approved :: Apache Software License
 License-File: LICENSE

{pathwaysutils-0.1.9 → pathwaysutils-0.1.11}/pathwaysutils/__init__.py

@@ -21,4 +21,4 @@ is_pathways_backend_used: Callable[[], bool] = _initialize.is_pathways_backend_u
 del _initialize
 
 # When changing this, also update the CHANGELOG.md.
-__version__ = "v0.1.9"
+__version__ = "v0.1.11"

{pathwaysutils-0.1.9 → pathwaysutils-0.1.11}/pathwaysutils/debug/watchdog.py

@@ -34,7 +34,7 @@ def _log_thread_stack(thread: threading.Thread):
   _logger.debug(
       "".join(
           traceback.format_stack(
-              sys._current_frames().get(  # pylint: disable=protected-access
+              sys._current_frames().get(  # pylint: disable=protected-access  # pyrefly: ignore[no-matching-overload]
                   thread.ident, []
               )
           )

{pathwaysutils-0.1.9 → pathwaysutils-0.1.11}/pathwaysutils/elastic/elastic.py

@@ -121,9 +121,9 @@ class DefaultSliceHealthChecker(SliceHealthChecker):
         _logger.debug(
             "Caught JaxRuntimeError for slice_index=%s: %s", slice_index, error
         )
-        if not is_error_due_to_slice_down(error):
+        if not is_error_due_to_slice_down(error, log_traceback=False):
           raise
-    return active_slice_indices
+    return frozenset(active_slice_indices)
 
 
 
@@ -249,7 +249,9 @@ def wait_for_slices(
       time.sleep(time_to_sleep)
 
 
-def is_error_due_to_slice_down(error: Exception) -> bool:
+def is_error_due_to_slice_down(
+    error: Exception, log_traceback: bool = True
+) -> bool:
   """Returns True if the error is due to slice down.
 
   The error types that are considered due to slice down are
@@ -261,6 +263,7 @@ def is_error_due_to_slice_down(error: Exception) -> bool:
 
   Args:
     error: The error to check.
+    log_traceback: If True, log the traceback of the error.
   """
   error_due_to_slice_down = False
   traceback_logging_level = logging.DEBUG
@@ -290,9 +293,7 @@ def is_error_due_to_slice_down(error: Exception) -> bool:
 
       error_due_to_slice_down = True
 
-  if not error_due_to_slice_down:
-    _logger.debug("Caught an error not due to slice down")
-
-  _logger.log(traceback_logging_level, "Error details:", exc_info=True)
+  if not error_due_to_slice_down and log_traceback:
+    _logger.log(traceback_logging_level, "Error details:", exc_info=True)
 
   return error_due_to_slice_down

{pathwaysutils-0.1.9 → pathwaysutils-0.1.11}/pathwaysutils/elastic/manager.py

@@ -22,7 +22,7 @@ from collections.abc import Callable, Mapping, Sequence, Set
 import functools
 import logging
 import threading
-from typing import Any, TypeVar
+from typing import Any, TypeAlias, TypeVar
 import warnings
 
 import jax
@@ -48,6 +48,7 @@ class ScaleUpSignalError(Exception):
 
 
 _F = TypeVar("_F", bound=Callable[..., Any])
+RetryPolicy: TypeAlias = Callable[[int, Exception], bool]
 
 
 def _elastic_event_cleanup() -> None:
@@ -66,23 +67,43 @@ def _elastic_event_cleanup() -> None:
     array.delete()
 
 
+class ElasticRetryLimit:
+  """A retry callback that limits the number of attempts."""
+
+  def __init__(self, max_attempts: int):
+    if max_attempts <= 0:
+      raise ValueError("max_attempts must be positive.")
+    self.max_attempts = max_attempts
+
+  def __call__(self, attempt: int, error: Exception) -> bool:
+    del error  # Unused
+    return attempt < self.max_attempts
+
+
 class Manager:
   """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.
+    all_slice_indices: A set of all possible slice indices in the allocation.
+    active_slice_indices: A set of indices of the slices currently participating
+      in the JAX computation mesh.
+    inactive_slice_indices: A set of indices of the slices in the allocation
+      that are not currently participating in the JAX computation mesh (e.g.
+      because they are down or have not joined yet).
+    available_inactive_slices: A set of indices of the inactive slices that
+      have been detected as up and healthy by the background monitor thread,
+      but have not yet been joined to the active JAX computation mesh.
   """
 
   slice_to_devices: Mapping[int, Sequence[jax.Device]]
   all_slice_indices: Set[int]
   active_slice_indices: Set[int]
-  new_slice_event: threading.Event
-
+  inactive_slice_indices: Set[int]
+  available_inactive_slices: Set[int]
+  _stop_event: threading.Event | None
+  _monitor_thread: threading.Thread | None
   def __init__(self, devices: Sequence[jax.Device] | None = None) -> None:
     """Initializes the manager.
 
@@ -93,12 +114,55 @@ class Manager:
       devices = jax.devices()
     self.slice_to_devices = elastic.get_slice_to_devices(devices)
 
-    self.all_slice_indices = set(self.slice_to_devices.keys())
+    self.all_slice_indices = frozenset(self.slice_to_devices.keys())
 
     self.active_slice_indices = elastic.get_active_slice_indices(
         slice_to_devices=self.slice_to_devices
     )
-    self.new_slice_event = threading.Event()
+    self.inactive_slice_indices = self.all_slice_indices - self.active_slice_indices
+    self.available_inactive_slices = frozenset()
+
+    self._stop_event = None
+    self._monitor_thread = None
+
+  def start_monitoring(self, poll_interval: float | int = 10) -> None:
+    """Starts the background monitor thread.
+
+    Args:
+      poll_interval: The number of seconds to wait between activity checks.
+    """
+    if self._monitor_thread is not None and self._monitor_thread.is_alive():
+      _logger.warning("Monitor thread is already running.")
+      return
+
+    self._stop_event = threading.Event()
+    self._monitor_thread = threading.Thread(
+        target=self._monitor_new_slices,
+        args=(self._stop_event, poll_interval),
+        daemon=True,
+    )
+    self._monitor_thread.start()
+    _logger.info("Elastic monitor thread started with interval %s.", poll_interval)
+
+  def close(self) -> None:
+    """Stops the background monitor thread."""
+    if self._stop_event is not None:
+      self._stop_event.set()
+    if self._monitor_thread is not None:
+      _logger.info("Closing manager, waiting for monitor thread to stop...")
+      try:
+        self._monitor_thread.join(timeout=5)
+        if self._monitor_thread.is_alive():
+          _logger.warning(
+              "Elastic monitor thread failed to stop within 5s timeout."
+          )
+      except RuntimeError as e:
+        if "cannot join thread" in str(e):
+          pass
+        else:
+          raise
+      self._monitor_thread = None
+      self._stop_event = None
 
   @functools.cached_property
   def total_slice_count(self) -> int:
@@ -122,9 +186,15 @@ class Manager:
     return len(self.active_slice_indices)
 
   @property
-  def inactive_slice_indices(self) -> set[int]:
-    """The set of inactive slice indices."""
-    return self.all_slice_indices - self.active_slice_indices
+  def new_slice_event(self) -> threading.Event:
+    """Deprecated compatibility property for un-updated MaxText code.
+
+    TODO: b/527183831 - Remove this property once MaxText CL 2 is submitted.
+    """
+    event = threading.Event()
+    if self.available_inactive_slices:
+      event.set()
+    return event
 
   def scale_by_active_slices(self, x: int | float) -> int | float:
     """Scale x by the number of active slices."""
@@ -157,46 +227,68 @@ class Manager:
     for array in jax.live_arrays():
       array.delete()
 
-  def _monitor_new_slices(
-      self, stop_event: threading.Event, poll_interval: float | int
-  ) -> None:
-    """Monitors for new slices and sets the `new_slice_event` if found."""
-    while not stop_event.wait(poll_interval):
-      try:
-        if not self.inactive_slice_indices:
-          _logger.debug("No inactive slices to check.")
-          continue
+  def _check_inactive_slices(self) -> None:
+    """Checks inactive slices and updates available_inactive_slices."""
+    if not self.inactive_slice_indices:
+      _logger.debug("No inactive slices to check.")
+      if self.available_inactive_slices:
+        self.available_inactive_slices = frozenset()
+      return
 
-        _logger.debug(
-            "Checking inactive slices: %s", self.inactive_slice_indices
-        )
-        inactive_slice_to_devices = {
-            i: self.slice_to_devices[i] for i in self.inactive_slice_indices
-        }
-        newly_active_indices = elastic.get_active_slice_indices(
-            inactive_slice_to_devices
-        )
+    _logger.debug(
+        "Now checking inactive slices %s", self.inactive_slice_indices
+    )
+    inactive_slice_to_devices = {
+        i: self.slice_to_devices[i] for i in self.inactive_slice_indices
+    }
+    found_slices = elastic.get_active_slice_indices(
+        inactive_slice_to_devices
+    )
 
-        if newly_active_indices:
-          _logger.info(
-              "New slices found: %s. Setting new slice event.",
-              newly_active_indices,
-          )
-          self.new_slice_event.set()
-          return
+    _logger.debug(
+        "Found available and inactive slices %s", found_slices
+    )
 
-        _logger.debug("No new slices found.")
-      except Exception:  # pylint: disable=broad-exception-caught
-        _logger.exception("Error in monitor thread")
+    # Filter against active_slice_indices in case the main thread initiated scale-up
+    # and claimed these slices while this background health check was running.
+    found_slices = found_slices - self.active_slice_indices
+
+    if found_slices != self.available_inactive_slices:
+      _logger.info(
+          "Newly available but inactive slices %s", found_slices
+      )
+      self.available_inactive_slices = frozenset(found_slices)
+
+  def _monitor_new_slices(
+      self, stop_event: threading.Event, poll_interval: float | int
+  ) -> None:
+    """Monitors for new slices and updates available_inactive_slices."""
+    _logger.info("Elastic monitor thread started.")
+    try:
+      while not stop_event.wait(poll_interval):
+        try:
+          self._check_inactive_slices()
+        except Exception:  # pylint: disable=broad-exception-caught
+          _logger.exception("Error in monitor thread loop")
+    except BaseException as e:
+      _logger.critical(
+          "Catastrophic error in monitor thread, thread is dying!",
+          exc_info=True,
+      )
+      raise
+    finally:
+      _logger.info("Elastic monitor thread stopped.")
 
   def elastic_retry(
       self,
-      max_retries: int,
+      # TODO: b/523384843 - Remove max_retries parameter.
+      max_retries: int | None = None,
       minimum_slice_count: int | None = None,
       poll_interval: float | int = 10,
       timeout: float | None = None,
       pre_callback: Callable[..., Any] | None = None,
       on_elastic_event_callback: Callable[..., Any] | None = None,
+      retry_policy: RetryPolicy | None = None,
   ) -> Callable[[_F], _F]:
     """Retries a function with elasticity fault tolerance.
 
@@ -211,10 +303,10 @@ class Manager:
     (i.e., it waits for all slices to be active).
 
     When `minimum_slice_count` is less than the total number of slices, a
-    background thread will monitor for new slices becoming available and set
-    `self.new_slice_event`. The user code can then poll this event and raise
-    a `ScaleUpSignalError` to gracefully interrupt the current execution and
-    trigger a retry.
+    background thread will monitor for newly joined inactive slices and populate
+    `self.available_inactive_slices`. User code can check this set (e.g. at step
+    boundaries) and raise a `ScaleUpSignalError` to gracefully interrupt the
+    current execution and trigger a retry with the expanded hardware.
 
     Often, the function will dispatch JAX operations and wait for them to
     complete while creating a log message. If using Python logging, it is
@@ -224,6 +316,8 @@ class Manager:
 
     Args:
       max_retries: The maximum number of times to retry the function.
+        Deprecated: Use `retry_policy` instead.
+        TODO: b/523384843 - Remove max_retries parameter.
       minimum_slice_count: The minimum number of slices required to run the
         function. If None, defaults to the total number of slices.
       poll_interval: The number of seconds to wait between activity checks.
@@ -233,6 +327,11 @@ class Manager:
       pre_callback: A callback to call before the function is attempted.
       on_elastic_event_callback: A callback to call after an elastic failure
         occurs.
+      retry_policy: A policy (callable) to determine if a retry should be
+        attempted. It accepts the attempt number (1-indexed) and the exception
+        that triggered the retry. If it returns False, no more retries are
+        attempted. If neither `retry_policy` nor `max_retries` is specified, it
+        defaults to unlimited retries.
 
     Returns:
       A decorator that retries the wrapped function.
@@ -248,77 +347,94 @@ class Manager:
         else minimum_slice_count
     )
 
-    if max_retries <= 0:
-      raise ValueError("max_retries must be positive.")
+    if max_retries is not None and retry_policy is not None:
+      raise ValueError("Cannot specify both max_retries and retry_policy.")
+
+    if retry_policy is None:
+      if max_retries is None:
+        # Default to unlimited retries if neither parameter is supplied.
+        retry_policy = lambda attempt, error: True
+      else:
+        if max_retries <= 0:
+          raise ValueError("max_retries must be positive.")
+        retry_policy = ElasticRetryLimit(max_retries)
 
     def decorator(func: _F) -> _F:
       @functools.wraps(func)
       def wrapper(*args: Any, **kwargs: Any) -> Any:
+        self.start_monitoring(poll_interval)
 
-        def attempt_execution(retry_index: int) -> Any:
-          _logger.info(
-              "Elastic attempt %d out of %d", retry_index + 1, max_retries
-          )
+        def attempt_execution(attempt: int) -> Any:
+          _logger.info("Elastic attempt %d", attempt)
           self.active_slice_indices = elastic.wait_for_slices(
               slice_count=target_slice_count,
               slice_to_devices=self.slice_to_devices,
               poll_interval=poll_interval,
               timeout=timeout,
           )
+          self.inactive_slice_indices = (
+              self.all_slice_indices - self.active_slice_indices
+          )
+          # Reset available_inactive_slices at attempt start since
+          # active_slice_indices has just been updated by wait_for_slices.
+          self.available_inactive_slices = frozenset()
           if pre_callback is not None:
             pre_callback()
 
           with jax.default_device(self.default_device):
-            self.new_slice_event.clear()
-            stop_event = threading.Event()
-
-            if target_slice_count < self.total_slice_count:
-              monitor_thread = threading.Thread(
-                  target=self._monitor_new_slices,
-                  args=(stop_event, poll_interval),
-                  daemon=True,
-              )
-              monitor_thread.start()
-            else:
-              monitor_thread = None
-
+            return func(*args, **kwargs)
+
+        def handle_scale_up_error(attempt: int, error: ScaleUpSignalError) -> None:
+          _logger.info("Scale up requested.")
+          _elastic_event_cleanup()
+          # Reset available_inactive_slices before retry callback and next attempt.
+          self.available_inactive_slices = frozenset()
+
+          if on_elastic_event_callback is not None:
+            on_elastic_event_callback()
+
+          if not retry_policy(attempt, error):
+            _logger.info("Retry policy rejected retry after ScaleUpSignalError.")
+            raise ElasticRuntimeError(
+                f"Elastic attempt {attempt} failed."
+            ) from error
+
+        def handle_slice_down_error(
+            attempt: int, error: jax.errors.JaxRuntimeError
+        ) -> None:
+          if not elastic.is_error_due_to_slice_down(error):
+            raise
+
+          _logger.exception("Elastic event detected")
+          _elastic_event_cleanup()
+          # Reset available_inactive_slices on slice-down failure before retry.
+          self.available_inactive_slices = frozenset()
+
+          if on_elastic_event_callback is not None:
+            on_elastic_event_callback()
+
+          if not retry_policy(attempt, error):
+            _logger.info(
+                "Retry policy rejected retry after jax.errors.JaxRuntimeError."
+            )
+            raise ElasticRuntimeError(
+                f"Elastic attempt {attempt} failed."
+            ) from error
+
+        try:
+          attempt = 1
+          while True:
             try:
-              return func(*args, **kwargs)
-            finally:
-              stop_event.set()
-              if monitor_thread is not None:
-                monitor_thread.join()
-
-        for retry_index in range(max_retries):
-          try:
-            return attempt_execution(retry_index)
-          except ScaleUpSignalError:
-            _logger.info("Scale up requested. Retrying.")
-            _elastic_event_cleanup()
-
-            if on_elastic_event_callback is not None:
-              on_elastic_event_callback()
-          except jax.errors.JaxRuntimeError as error:
-            if not elastic.is_error_due_to_slice_down(error):
-              raise
-
-            if self.new_slice_event.is_set():
-              _logger.info(
-                  "Slice down event and new slice available detected. Retrying."
-              )
-            else:
-              _logger.info("Slice down event detected. Retrying.")
-
-            _elastic_event_cleanup()
-
-            if on_elastic_event_callback is not None:
-              on_elastic_event_callback()
-        else:
-          raise ElasticRuntimeError(
-              f"Elastic attempt {max_retries} out of {max_retries} failed."
-          )
-
-      return wrapper
+              return attempt_execution(attempt)
+            except ScaleUpSignalError as error:
+              handle_scale_up_error(attempt, error)
+            except jax.errors.JaxRuntimeError as error:
+              handle_slice_down_error(attempt, error)
+            attempt += 1
+        finally:
+          self.close()
+
+      return wrapper  # pyrefly: ignore[bad-return]
 
     return decorator