pathwaysutils 0.1.11__tar.gz → 0.1.13__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pathwaysutils
-Version: 0.1.11
+Version: 0.1.13
 Summary: Package of Pathways-on-Cloud utilities.
 Keywords: 
 Author-email: Pathways-on-Cloud Utilities Developers <pathwaysutils-dev@google.com>
@@ -10,17 +10,32 @@ 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: jax>=0.10.0
+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.11 → pathwaysutils-0.1.13}/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.11"
+__version__ = "v0.1.13"

{pathwaysutils-0.1.11 → pathwaysutils-0.1.13}/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.11 → pathwaysutils-0.1.13}/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.11 → pathwaysutils-0.1.13}/pathwaysutils/elastic/elastic.py

@@ -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,

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

@@ -52,7 +52,7 @@ RetryPolicy: TypeAlias = Callable[[int, Exception], bool]
 
 
 def _elastic_event_cleanup() -> None:
-  """Cleans up JAX profiles, caches, and live arrays."""
+  """Stops any ongoing JAX profiler trace."""
   try:
     _logger.info("Cleaning up any ongoing traces")
     jax.profiler.stop_trace()
@@ -62,10 +62,6 @@ def _elastic_event_cleanup() -> None:
     _logger.exception("Error cleaning up ongoing traces")
     raise
 
-  jax.clear_caches()
-  for array in jax.live_arrays():
-    array.delete()
-
 
 class ElasticRetryLimit:
   """A retry callback that limits the number of attempts."""
@@ -116,11 +112,12 @@ class Manager:
 
     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.inactive_slice_indices = self.all_slice_indices - self.active_slice_indices
-    self.available_inactive_slices = frozenset()
 
     self._stop_event = None
     self._monitor_thread = None
@@ -180,22 +177,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 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."""
     if isinstance(x, int):
@@ -294,8 +295,7 @@ class Manager:
 
     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.
+    `minimum_slice_count` active slices before each attempt.
 
     If `minimum_slice_count` is not met, the function will wait until at least
     `minimum_slice_count` slices are active before execution. If
@@ -308,6 +308,12 @@ class Manager:
     boundaries) and raise a `ScaleUpSignalError` to gracefully interrupt the
     current execution and trigger a retry with the expanded hardware.
 
+    On an elastic event the arrays of the failed attempt are not deleted
+    explicitly; JAX frees them once the attempt has unwound and nothing
+    references them anymore. The wrapped function and the callbacks must not
+    keep references to them, or the memory stays allocated during the next
+    attempt.
+
     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
@@ -372,9 +378,6 @@ class Manager:
               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()