pathwaysutils 0.1.6__tar.gz → 0.1.8__tar.gz

{pathwaysutils-0.1.6 → pathwaysutils-0.1.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pathwaysutils
-Version: 0.1.6
+Version: 0.1.8
 Summary: Package of Pathways-on-Cloud utilities.
 Keywords: 
 Author-email: Pathways-on-Cloud Utilities Developers <pathwaysutils-dev@google.com>
@@ -10,7 +10,7 @@ Classifier: License :: OSI Approved :: Apache Software License
 License-File: LICENSE
 Requires-Dist: absl-py
 Requires-Dist: fastapi
-Requires-Dist: jax>=0.7.2
+Requires-Dist: jax>=0.8.0
 Requires-Dist: orbax-checkpoint
 Requires-Dist: uvicorn
 Requires-Dist: requests

{pathwaysutils-0.1.6 → pathwaysutils-0.1.8}/pathwaysutils/__init__.py

@@ -12,12 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Package of Pathways-on-Cloud utilities."""
+from collections.abc import Callable
 from pathwaysutils import _initialize
 
-initialize = _initialize.initialize
-is_pathways_backend_used = _initialize.is_pathways_backend_used
+initialize: Callable[[], None] = _initialize.initialize
+is_pathways_backend_used: Callable[[], bool] = _initialize.is_pathways_backend_used
 
 del _initialize
 
 # When changing this, also update the CHANGELOG.md.
-__version__ = "v0.1.6"
+__version__ = "v0.1.8"

{pathwaysutils-0.1.6 → pathwaysutils-0.1.8}/pathwaysutils/collect_profile.py

@@ -26,7 +26,7 @@ _logger = logging.getLogger(__name__)
 _logger.setLevel(logging.INFO)
 
 
-_DESCRIPTION = """
+_DESCRIPTION: str = """
 To profile running JAX programs, you first need to start the profiler server
 in the program of interest. You can do this via
 `jax.profiler.start_server(<port>)`. Once the program is running and the
@@ -36,7 +36,7 @@ for a provided duration. The trace file will be dumped into a GCS bucket
 """
 
 
-def _get_parser():
+def _get_parser() -> argparse.ArgumentParser:
   """Returns an argument parser for the collect_profile script."""
   parser = argparse.ArgumentParser(description=_DESCRIPTION)
   parser.add_argument(
@@ -62,7 +62,7 @@ def _get_parser():
   return parser
 
 
-def main():
+def main() -> None:
   parser = _get_parser()
   args = parser.parse_args()
 

{pathwaysutils-0.1.6 → pathwaysutils-0.1.8}/pathwaysutils/debug/watchdog.py

@@ -43,12 +43,17 @@ def _log_thread_stack(thread: threading.Thread):
 
 
 @contextlib.contextmanager
-def watchdog(timeout: float, repeat: bool = True):
+def watchdog(
+    name: str,
+    timeout: float,
+    repeat: bool = True,
+    ):
   """Watchdog context manager.
 
   Prints the stack trace of all threads after `timeout` seconds.
 
   Args:
+    name: The name of the watchdog, used in log messages.
     timeout: The timeout in seconds. If the timeout is reached, the stack trace
       of all threads will be printed.
     repeat: Whether to repeat the watchdog after the timeout. If False, the
@@ -63,7 +68,10 @@ def watchdog(timeout: float, repeat: bool = True):
     count = 0
     while not event.wait(timeout):
       _logger.debug(
-          "Watchdog thread dump every %s seconds. Count: %s", timeout, count
+          "'%s' watchdog thread stack dump every %s seconds. Count: %s",
+          name,
+          timeout,
+          count,
       )
       try:
         for thread in threading.enumerate():
@@ -73,17 +81,22 @@ def watchdog(timeout: float, repeat: bool = True):
             _logger.debug("Error print traceback for thread: %s", thread.ident)
       finally:
         if not repeat:
-          _logger.critical("Timeout from watchdog!")
+          _logger.critical("Timeout from watchdog '%s'", name)
           os.abort()
 
       count += 1
 
-  _logger.debug("Registering watchdog")
-  watchdog_thread = threading.Thread(target=handler, name="watchdog")
+  _logger.debug(
+      "Registering '%s' watchdog with timeout %s seconds and repeat=%s",
+      name,
+      timeout,
+      repeat,
+  )
+  watchdog_thread = threading.Thread(target=handler, name=name)
   watchdog_thread.start()
   try:
     yield
   finally:
     event.set()
     watchdog_thread.join()
-    _logger.debug("Deregistering watchdog")
+    _logger.debug("Deregistering '%s' watchdog", name)

{pathwaysutils-0.1.6 → pathwaysutils-0.1.8}/pathwaysutils/elastic/elastic.py

@@ -19,7 +19,7 @@ events. It also provides a utility for waiting for slices to become active.
 """
 
 import collections
-from collections.abc import Mapping, Sequence
+from collections.abc import Mapping, Sequence, Set
 import logging
 import time
 
@@ -83,7 +83,7 @@ def _simple_execution(devices: Sequence[jax.Device]) -> jax.Array:
 
 def get_slice_to_devices(
     devices: Sequence[jax.Device],
-) -> dict[int, Sequence[jax.Device]]:
+) -> Mapping[int, Sequence[jax.Device]]:
   """Returns the mapping from slice index to devices."""
   slice_to_devices = collections.defaultdict(list)
   for d in devices:
@@ -94,7 +94,7 @@ def get_slice_to_devices(
 @timing.timeit
 def get_active_slice_indices(
     slice_to_devices: Mapping[int, Sequence[jax.Device]] | None = None,
-) -> set[int]:
+) -> Set[int]:
   """Returns the set of active slices indices.
 
   Args:
@@ -106,8 +106,14 @@ def get_active_slice_indices(
     A set of integers representing the indices of the active slices.
   """
   if slice_to_devices is None:
+    _logger.debug("slice_to_devices is None. Getting from jax.devices().")
     slice_to_devices = get_slice_to_devices(tuple(jax.devices()))
 
+  _logger.debug(
+      "Getting active slice indices for slices: %s",
+      sorted(list(slice_to_devices.keys())),
+  )
+
   active_slice_indices = set()
 
   results = {
@@ -116,17 +122,19 @@ def get_active_slice_indices(
   }
 
   for slice_index, x in results.items():
-    _logger.info("Checking slice_index=%s", slice_index)
+    _logger.debug("Checking slice_index=%s", slice_index)
     expected = (
         np.zeros(len(slice_to_devices[slice_index]), dtype=float)
         + _SIMPLE_EXECUTION_TEST_VALUE
     )
     try:
       with timing.Timer(f"Checking {slice_index=}"):
+        _logger.debug("Blocking until ready for slice_index=%s", slice_index)
         jax.block_until_ready(x)
+        _logger.debug("Execution finished for slice_index=%s", slice_index)
         if np.allclose(x, expected):
           active_slice_indices.add(slice_index)
-          _logger.info("slice_index=%s active", slice_index)
+          _logger.debug("slice_index=%s active", slice_index)
         else:
           _logger.error(
               "Error with _simple_execution for slice_index=%s. "
@@ -139,11 +147,15 @@ def get_active_slice_indices(
               f"Error with _simple_execution for slice_index={slice_index}."
           )
     except jax.errors.JaxRuntimeError as error:
+      _logger.debug(
+          "Caught JaxRuntimeError for slice_index=%s: %s", slice_index, error
+      )
       if not is_error_due_to_slice_down(error):
+        _logger.info("Re-raising error for slice_index=%s", slice_index)
         raise
-      _logger.info("slice_index=%s bad", slice_index)
+      _logger.debug("slice_index=%s bad", slice_index)
 
-  _logger.info("active_slice_indices=%s", active_slice_indices)
+  _logger.debug("active_slice_indices=%s", active_slice_indices)
 
   return active_slice_indices
 
@@ -153,7 +165,7 @@ def wait_for_slices(
     poll_interval: float | int = 10,
     timeout: float | int | None = None,
     slice_to_devices: Mapping[int, Sequence[jax.Device]] | None = None,
-) -> set[int]:
+) -> Set[int]:
   """Waits until after at least `slice_count` slices become active.
 
   Args:
@@ -174,22 +186,36 @@ def wait_for_slices(
       active.
   """
   if slice_to_devices is None:
+    _logger.debug("slice_to_devices is None. Getting from jax.devices().")
     slice_to_devices = get_slice_to_devices(jax.devices())
 
+  _logger.info(
+      "Waiting for %s slices. Poll interval: %s, Timeout: %s",
+      slice_count,
+      poll_interval,
+      timeout,
+  )
   start_time = time.time()
 
   while True:
     check_start_time = time.time()
 
+    _logger.debug("Checking active slices...")
     active_slice_indices = get_active_slice_indices(slice_to_devices)
     if len(active_slice_indices) >= slice_count:
-      _logger.info("%s slices active.", len(active_slice_indices))
+      _logger.info(
+          "Sufficient slices active: %s >= %s. Active indices: %s",
+          len(active_slice_indices),
+          slice_count,
+          active_slice_indices,
+      )
       return active_slice_indices
 
     _logger.info(
-        "%s slices active. Wanting at least %s.",
+        "%s slices active. Wanting at least %s. Active indices: %s",
         len(active_slice_indices),
         slice_count,
+        active_slice_indices,
     )
 
     time_to_sleep = max(0, poll_interval - (time.time() - check_start_time))
@@ -206,7 +232,7 @@ def wait_for_slices(
         )
 
     if time_to_sleep > 0:
-      _logger.info("Sleeping for %.2f seconds.", time_to_sleep)
+      _logger.debug("Sleeping for %.2f seconds.", time_to_sleep)
 
       time.sleep(time_to_sleep)
 
@@ -228,10 +254,14 @@ def is_error_due_to_slice_down(error: Exception) -> bool:
   traceback_logging_level = logging.DEBUG
 
   if isinstance(error, jax.errors.JaxRuntimeError):
+    _logger.debug("Checking if JaxRuntimeError is due to slice down: %s", error)
     if any(
         error_type in str(error) for error_type in _ELASTIC_DOWN_ERROR_TYPES
     ):
-      _logger.info("Caught an error due to slice down")
+      _logger.debug(
+          "Caught an error due to slice down (matched"
+          " _ELASTIC_DOWN_ERROR_TYPES)"
+      )
 
       error_due_to_slice_down = True
 
@@ -240,15 +270,16 @@ def is_error_due_to_slice_down(error: Exception) -> bool:
         for error_type in _ELASTIC_DOWN_ADDITIONAL_ERROR_TYPES
     ):
       _logger.warning(
-          "Caught an error due that may or may not be due to slice down. This"
-          " error will be treated as due to slice down."
+          "Caught an error that may or may not be due to slice down (matched"
+          " _ELASTIC_DOWN_ADDITIONAL_ERROR_TYPES). This error will be treated"
+          " as due to slice down."
       )
       traceback_logging_level = logging.WARNING
 
       error_due_to_slice_down = True
 
   if not error_due_to_slice_down:
-    _logger.info("Caught an error not due to slice down")
+    _logger.debug("Caught an error not due to slice down")
 
   _logger.log(traceback_logging_level, "Error details:", exc_info=True)
 

pathwaysutils-0.1.8/pathwaysutils/elastic/manager.py

@@ -0,0 +1,398 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Elasticity manager.
+
+This class provides a utility for elastic training. It provides a decorator that
+retries a function in case of `jax.errors.JaxRuntimeError` caused by slice down
+events. It also provides a utility for waiting for slices to become active.
+"""
+
+from collections.abc import Callable, Mapping, Sequence, Set
+import functools
+import logging
+import threading
+from typing import Any, TypeVar
+import warnings
+
+import jax
+from pathwaysutils.elastic import elastic
+
+
+_logger = logging.getLogger(__name__)
+
+
+class ElasticRuntimeError(RuntimeError):
+  """Error raised when elasticity cannot continue."""
+
+
+class ScaleUpSignalError(Exception):
+  """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() -> None:
+  """Cleans up JAX profiles, caches, and live arrays."""
+  try:
+    _logger.info("Cleaning up any ongoing traces")
+    jax.profiler.stop_trace()
+  except (RuntimeError, ValueError):
+    _logger.info("No ongoing traces to clean up")
+  except Exception:
+    _logger.exception("Error cleaning up ongoing traces")
+    raise
+
+  jax.clear_caches()
+  for array in jax.live_arrays():
+    array.delete()
+
+
+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.
+  """
+
+  slice_to_devices: Mapping[int, Sequence[jax.Device]]
+  all_slice_indices: Set[int]
+  active_slice_indices: Set[int]
+  new_slice_event: threading.Event
+
+  def __init__(self, devices: Sequence[jax.Device] | None = None) -> None:
+    """Initializes the manager.
+
+    Args:
+      devices: The devices to use. If None, jax.devices() is used.
+    """
+    if devices is None:
+      devices = jax.devices()
+    self.slice_to_devices = elastic.get_slice_to_devices(devices)
+
+    self.all_slice_indices = set(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()
+
+  @functools.cached_property
+  def total_slice_count(self) -> int:
+    """The total number of slices."""
+    return len(self.slice_to_devices)
+
+  @property
+  def default_device(self) -> jax.Device:
+    """The device that should be set to the default device.
+
+    This will be from one of the slices in `active_slice_indices`.
+    """
+    try:
+      return self.slice_to_devices[next(iter(self.active_slice_indices))][0]
+    except StopIteration as error:
+      raise ValueError("No active slices") from error
+
+  @property
+  def active_slice_count(self) -> int:
+    """The number of active slices."""
+    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 scale_by_active_slices(self, x: int | float) -> int | float:
+    """Scale x by the number of active slices."""
+    if isinstance(x, int):
+      quotient, remainder = divmod(
+          x * self.active_slice_count, self.total_slice_count
+      )
+      if remainder:
+        raise ValueError(
+            f"Cannot scale {x=} by active slices because it will result in a "
+            f"remainder of {remainder=}."
+        )
+      return quotient
+    elif isinstance(x, float):
+      return x * self.active_slice_count / self.total_slice_count
+    else:
+      raise ValueError(f"Unsupported type: {type(x)=}")
+
+  def _cleanup_on_retry(self):
+    """Cleans up JAX caches and traces on retry."""
+    try:
+      _logger.debug("Cleaning up any ongoing traces")
+      jax.profiler.stop_trace()
+    except (RuntimeError, ValueError):
+      _logger.debug("No ongoing traces to clean up")
+    except Exception:  # pylint: disable=broad-exception-caught
+      _logger.exception("Error cleaning up ongoing traces")
+
+    jax.clear_caches()
+    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
+
+        _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
+        )
+
+        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("No new slices found.")
+      except Exception:  # pylint: disable=broad-exception-caught
+        _logger.exception("Error in monitor thread")
+
+  def elastic_retry(
+      self,
+      max_retries: int,
+      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,
+  ) -> Callable[[_F], _F]:
+    """Retries a function with elasticity fault tolerance.
+
+    This decorator wraps a function to automatically retry execution in case of
+    `jax.errors.JaxRuntimeError` caused by slice down events. It waits for
+    `minimum_slice_count` active slices before each attempt and cleans up JAX
+    caches on failure.
+
+    If `minimum_slice_count` is not met, the function will wait until at least
+    `minimum_slice_count` slices are active before execution. If
+    `minimum_slice_count` is None, it defaults to the total number of slices
+    (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.
+
+    Often, the function will dispatch JAX operations and wait for them to
+    complete while creating a log message. If using Python logging, it is
+    recommended to set `logging.raiseExceptions=True` to ensure that the
+    `jax.errors.JaxRuntimeError` is not silently ignored within the logging
+    call.
+
+    Args:
+      max_retries: The maximum number of times to retry the function.
+      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.
+        Defaults to 10 seconds.
+      timeout: The maximum number of seconds to wait for slices to become active
+        before each retry attempt. If None, there is no timeout.
+      pre_callback: A callback to call before the function is attempted.
+      on_elastic_event_callback: A callback to call after an elastic failure
+        occurs.
+
+    Returns:
+      A decorator that retries the wrapped function.
+
+    Raises:
+      ElasticRuntimeError: If all retry attempts fail.
+      Exception: Any other exception raised by the wrapped function that is not
+        due to a slice down event.
+    """
+    target_slice_count = (
+        self.total_slice_count
+        if minimum_slice_count is None
+        else minimum_slice_count
+    )
+
+    if max_retries <= 0:
+      raise ValueError("max_retries must be positive.")
+
+    def decorator(func: _F) -> _F:
+      @functools.wraps(func)
+      def wrapper(*args: Any, **kwargs: Any) -> Any:
+
+        def attempt_execution(retry_index: int) -> Any:
+          _logger.info(
+              "Elastic attempt %d out of %d", retry_index + 1, max_retries
+          )
+          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,
+          )
+          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
+
+            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 decorator
+
+  def pause_resume(
+      self,
+      max_retries: int,
+      poll_interval: float | int = 10,
+      timeout: float | None = None,
+      pre_callback: Callable[..., Any] | None = None,
+      on_elastic_event_callback: Callable[..., Any] | None = None,
+  ) -> Callable[[_F], _F]:
+    """Retries a function with pause/resume fault tolerance.
+
+    DEPRECATED: Use `elastic_retry` instead.
+
+    Args:
+      max_retries: The maximum number of times to retry the function.
+      poll_interval: The number of seconds to wait between activity checks.
+        Defaults to 10 seconds.
+      timeout: The maximum number of seconds to wait for slices to become active
+        before each retry attempt. If None, there is no timeout.
+      pre_callback: A callback to call before the function is attempted.
+      on_elastic_event_callback: A callback to call after an elastic failure
+        occurs.
+
+    Returns:
+      A decorator that retries the wrapped function.
+    """
+    warnings.warn(
+        "`pause_resume` is deprecated. Please use `elastic_retry` instead.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
+    return self.elastic_retry(
+        max_retries=max_retries,
+        minimum_slice_count=None,
+        poll_interval=poll_interval,
+        timeout=timeout,
+        pre_callback=pre_callback,
+        on_elastic_event_callback=on_elastic_event_callback,
+    )
+
+  def replica_resize(
+      self,
+      max_resizes: int,
+      poll_interval: float = 10,
+      pre_callback: Callable[..., Any] | None = None,
+      on_elastic_event_callback: Callable[..., Any] | None = None,
+  ) -> Callable[[_F], _F]:
+    """Retries a function with replica/resize fault tolerance.
+
+    DEPRECATED: Use `elastic_retry` instead.
+
+    Args:
+      max_resizes: The maximum number of times to retry the function after
+        resizing the replica count.
+      poll_interval: The number of seconds to wait between active slice checks.
+        Defaults to 10 seconds.
+      pre_callback: A callback to call before the function is attempted.
+      on_elastic_event_callback: A callback to call after an elastic failure
+        occurs.
+
+    Returns:
+      A decorator that retries the wrapped function.
+    """
+    warnings.warn(
+        "`replica_resize` is deprecated. Please use `elastic_retry` instead.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
+    return self.elastic_retry(
+        max_retries=max_resizes,
+        minimum_slice_count=1,
+        poll_interval=poll_interval,
+        pre_callback=pre_callback,
+        on_elastic_event_callback=on_elastic_event_callback,
+    )