pathwaysutils 0.1.10__tar.gz → 0.1.12__tar.gz

{pathwaysutils-0.1.10 → pathwaysutils-0.1.12}/PKG-INFO

@@ -1,26 +1,41 @@
 Metadata-Version: 2.4
 Name: pathwaysutils
-Version: 0.1.10
+Version: 0.1.12
 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
 Requires-Dist: absl-py
 Requires-Dist: fastapi
 Requires-Dist: jax>=0.8.3
-Requires-Dist: orbax-checkpoint
+Requires-Dist: orbax-checkpoint>=0.12.3
 Requires-Dist: uvicorn
 Requires-Dist: requests
 Requires-Dist: packaging
 Requires-Dist: pylint>=2.6.0 ; extra == "dev"
 Requires-Dist: pyink ; extra == "dev"
+Requires-Dist: kubernetes ; extra == "gke"
+Requires-Dist: pyyaml ; extra == "gke"
+Requires-Dist: portpicker ; extra == "sps"
+Requires-Dist: google-cloud-monitoring ; extra == "sps"
+Requires-Dist: portpicker ; extra == "sps-without-monitoring"
+Requires-Dist: google-cloud-monitoring ; extra == "test"
+Requires-Dist: httpx ; extra == "test"
+Requires-Dist: kubernetes ; extra == "test"
+Requires-Dist: portpicker ; extra == "test"
+Requires-Dist: pyyaml ; extra == "test"
+Requires-Dist: pytest ; extra == "test"
 Project-URL: changelog, https://github.com/AI-Hypercomputer/pathways-utils/blob/main/CHANGELOG.md
 Project-URL: homepage, https://github.com/AI-Hypercomputer/pathways-utils
 Project-URL: repository, https://github.com/AI-Hypercomputer/pathways-utils
 Provides-Extra: dev
+Provides-Extra: gke
+Provides-Extra: sps
+Provides-Extra: sps-without-monitoring
+Provides-Extra: test
 
 # Package of Pathways-on-Cloud utilities
 

{pathwaysutils-0.1.10 → pathwaysutils-0.1.12}/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.10"
+__version__ = "v0.1.12"

{pathwaysutils-0.1.10 → pathwaysutils-0.1.12}/pathwaysutils/_initialize.py

@@ -18,10 +18,11 @@ import logging
 import os
 
 import jax
+from orbax.checkpoint import type_handlers
 from orbax.checkpoint._src.metadata import array_metadata_store as array_metadata_store_lib
+from orbax.checkpoint._src.serialization import cloud_pathways_array_handler
 from pathwaysutils import profiling
 from pathwaysutils import proxy_backend
-from pathwaysutils.persistence import orbax_handler
 
 
 _logger = logging.getLogger(__name__)
@@ -91,11 +92,14 @@ def initialize() -> None:
     _logger.debug("Detected Pathways-on-Cloud backend. Applying changes.")
     proxy_backend.register_backend_factory()
     profiling.monkey_patch_jax()
-    # TODO: b/365549911 - Remove when OCDBT-compatible
     if _is_persistence_enabled():
-      orbax_handler.register_pathways_handlers(
-          timeout=datetime.timedelta(hours=1),
-          array_metadata_store=array_metadata_store_lib.Store(),
+      type_handlers.register_type_handler(
+          jax.Array,
+          cloud_pathways_array_handler.CloudPathwaysArrayHandler(
+              timeout=datetime.timedelta(hours=1),
+              array_metadata_store=array_metadata_store_lib.Store(),
+          ),
+          override=True,
       )
 
     # Turn off JAX compilation cache because Pathways handles its own

{pathwaysutils-0.1.10 → pathwaysutils-0.1.12}/pathwaysutils/collect_profile.py

@@ -67,7 +67,10 @@ def main() -> None:
   args = parser.parse_args()
 
   if profiling.collect_profile(
-      args.port, args.duration_ms, args.host, args.log_dir
+      args.port,
+      args.duration_ms,
+      args.host,
+      args.log_dir,
   ):
     _logger.info("Dumped profiling information in: %s", args.log_dir)
   else:

{pathwaysutils-0.1.10 → pathwaysutils-0.1.12}/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.10 → pathwaysutils-0.1.12}/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)
 
 
 
@@ -171,7 +171,7 @@ def get_active_slice_indices(
 
 
 def wait_for_slices(
-    slice_count: int,
+    slice_count: int | None = None,
     poll_interval: float | int = 10,
     timeout: float | int | None = None,
     slice_to_devices: Mapping[int, Sequence[jax.Device]] | None = None,
@@ -180,7 +180,8 @@ def wait_for_slices(
   """Waits until after at least `slice_count` slices become active.
 
   Args:
-    slice_count: The number of slices to wait for.
+    slice_count: The number of slices to wait for. If None, defaults to the
+      total number of slices.
     poll_interval: The minimum number of seconds to wait between availability
       checks. If the check takes longer than this, the next check will start
       immediately after the current check completes. Defaults to 10 seconds.
@@ -201,6 +202,10 @@ def wait_for_slices(
     _logger.debug("slice_to_devices is None. Getting from jax.devices().")
     slice_to_devices = get_slice_to_devices(jax.devices())
 
+  if slice_count is None:
+    _logger.debug("slice_count is None. Using len(slice_to_devices).")
+    slice_count = len(slice_to_devices)
+
   _logger.info(
       "Waiting for %s slices. Poll interval: %s, Timeout: %s",
       slice_count,
@@ -249,7 +254,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 +268,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 +298,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.10 → pathwaysutils-0.1.12}/pathwaysutils/elastic/manager.py

@@ -86,17 +86,24 @@ class Manager:
   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.
 
@@ -107,12 +114,56 @@ 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 = frozenset()
+    self.available_inactive_slices = frozenset()
 
     self.active_slice_indices = elastic.get_active_slice_indices(
         slice_to_devices=self.slice_to_devices
     )
-    self.new_slice_event = threading.Event()
+
+    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:
@@ -130,16 +181,26 @@ class Manager:
     except StopIteration as error:
       raise ValueError("No active slices") from error
 
+  @property
+  def active_slice_indices(self) -> frozenset[int]:
+    """The indices of active slices."""
+    return self._active_slice_indices
+
+  @active_slice_indices.setter
+  def active_slice_indices(self, value: Set[int]) -> None:
+    self._active_slice_indices = frozenset(value)
+    self.inactive_slice_indices = (
+        self.all_slice_indices - self._active_slice_indices
+    )
+    self.available_inactive_slices = frozenset(
+        self.available_inactive_slices - self._active_slice_indices
+    )
+
   @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):
@@ -171,37 +232,57 @@ 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,
@@ -227,10 +308,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
@@ -286,6 +367,7 @@ class Manager:
     def decorator(func: _F) -> _F:
       @functools.wraps(func)
       def wrapper(*args: Any, **kwargs: Any) -> Any:
+        self.start_monitoring(poll_interval)
 
         def attempt_execution(attempt: int) -> Any:
           _logger.info("Elastic attempt %d", attempt)
@@ -295,75 +377,66 @@ class Manager:
               poll_interval=poll_interval,
               timeout=timeout,
           )
+          # 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()
-
-        attempt = 1
-        while True:
-          try:
-            return attempt_execution(attempt)
-          except ScaleUpSignalError as error:
-            _logger.info("Scale up requested.")
-            _elastic_event_cleanup()
-
-            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
-
-            _logger.info("Retrying.")
-          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.")
-            else:
-              _logger.info("Slice down event detected.")
-
-            _elastic_event_cleanup()
-
-            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 JaxRuntimeError.")
-              raise ElasticRuntimeError(
-                  f"Elastic attempt {attempt} failed."
-              ) from error
-
-            _logger.info("Retrying.")
-
-          attempt += 1
-
-      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