From d9041d73de25f9a5c5ce69d4c04e394619f140b1 Mon Sep 17 00:00:00 2001 From: Abhinav Singh Date: Fri, 9 Oct 2026 13:09:31 -0700 Subject: [PATCH] Stop clearing JAX caches and deleting live arrays on elastic events. PiperOrigin-RevId: 996711731 --- pathwaysutils/elastic/manager.py | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/pathwaysutils/elastic/manager.py b/pathwaysutils/elastic/manager.py index 1f48d5c..5bf3544 100644 --- a/pathwaysutils/elastic/manager.py +++ b/pathwaysutils/elastic/manager.py @@ -52,7 +52,14 @@ class ScaleUpSignalError(Exception): def _elastic_event_cleanup() -> None: - """Cleans up JAX profiles, caches, and live arrays.""" + """Stops any ongoing JAX profiler trace. + + Live arrays are not deleted and the JAX caches are not cleared: JAX frees a + buffer once its last reference is dropped, which happens when the failed + attempt unwinds, and compiled programs for slices that went away are never + called again. Anything that must not keep the previous attempt's arrays + alive, such as a snapshot still being copied, is the caller's responsibility. + """ try: _logger.info("Cleaning up any ongoing traces") jax.profiler.stop_trace() @@ -62,10 +69,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 +302,7 @@ def elastic_retry( 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 +315,12 @@ def elastic_retry( 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