From 7de3aea714efd49712cac96bcb7ee2482f592701 Mon Sep 17 00:00:00 2001 From: Luke Baumann Date: Sat, 14 Feb 2026 00:41:49 -0800 Subject: [PATCH] Update active_slice_indices after waiting for slices within pause-resume. The result of `wait_for_slices` is now assigned to `self.active_slice_indices` within the retry loop. Also, updated the docstring for `wait_for_slices` to clarify the return value. PiperOrigin-RevId: 870064706 --- pathwaysutils/elastic/manager.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/pathwaysutils/elastic/manager.py b/pathwaysutils/elastic/manager.py index ce57424..8bd712e 100644 --- a/pathwaysutils/elastic/manager.py +++ b/pathwaysutils/elastic/manager.py @@ -293,7 +293,7 @@ def wait_for_slices( timeout. Returns: - The good slice indices + The active slice indices Raises: TimeoutError: If the timeout is reached before the slices become @@ -388,7 +388,9 @@ def wrapper(*args, **kwargs): "Elastic attempt %d out of %d", retry_index + 1, max_retries ) - self.wait_for_slices(poll_interval=poll_interval, timeout=timeout) + self.active_slice_indices = self.wait_for_slices( + poll_interval=poll_interval, timeout=timeout + ) return func(*args, **kwargs) except jax.errors.JaxRuntimeError as error: