pathwaysutils 0.1.12__tar.gz → 0.1.13__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pathwaysutils
-Version: 0.1.12
+Version: 0.1.13
 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.8.3
+Requires-Dist: jax>=0.10.0
 Requires-Dist: orbax-checkpoint>=0.12.3
 Requires-Dist: uvicorn
 Requires-Dist: requests

{pathwaysutils-0.1.12 → 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.12"
+__version__ = "v0.1.13"

{pathwaysutils-0.1.12 → 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."""
@@ -299,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
@@ -313,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