Refactoring wait_for_slices

PiperOrigin-RevId: 798375418

Author: lukebaumann

pathwaysutils/elastic/manager.py

@@ -644,18 +644,19 @@ def maybe_reshard_up(
   def wait_for_slices(
       self,
       slice_count: int | None = None,
-      wait_period: float | int = 10,
+      poll_interval: float | int = 10,
       timeout: float | int | None = None,
   ) -> set[int]:
     """Waits until after at least `slice_count` slices become available.
 
     Args:
       slice_count: The number of slices to wait for. If None, waits for all
         slices to become available.
-      wait_period: The number of seconds to wait between availability checks.
-      Defaults to 10 seconds.
-      timeout: The maximum number of seconds to wait. If None, there is
-        no timeout.
+      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.
+      timeout: The maximum number of seconds to wait. If None, there is no
+        timeout.
 
     Returns:
       The good slice indices
@@ -669,42 +670,52 @@ def wait_for_slices(
 
     start_time = time.time()
 
-    while (
-        len(good_slice_indices := self.get_slice_availability()) < slice_count
-    ):
-      elapsed_time = time.time() - start_time
-      if timeout is not None and elapsed_time + wait_period >= timeout:
-        raise TimeoutError(
-            f"Timed out waiting for {slice_count} slices. Only"
-            f" {len(good_slice_indices)} available after"
-            f" {elapsed_time:.2f} seconds."
-            f" Next check would occur after the timeout of {timeout}"
-            " seconds."
+    while True:
+      check_start_time = time.time()
+
+      if (
+          len(good_slice_indices := self.get_slice_availability())
+          >= slice_count
+      ):
+        _logger.info(
+            "%s/%s slices are available",
+            len(good_slice_indices),
+            self.total_slice_count,
         )
+        return good_slice_indices
 
       _logger.info(
-          "%s/%s slices available. Wanting at least %s/%s. Sleeping for %s"
-          " seconds.",
+          "%s/%s slices available. Wanting at least %s/%s.",
           len(good_slice_indices),
           self.total_slice_count,
           slice_count,
           self.total_slice_count,
-          wait_period,
       )
-      time.sleep(wait_period)
 
-    _logger.info(
-        "%s/%s slices are available",
-        len(good_slice_indices),
-        self.total_slice_count,
-    )
+      time_to_sleep = max(0, poll_interval - (time.time() - check_start_time))
 
-    return good_slice_indices
+      if (
+          timeout is not None
+          and (elapsed_time := time.time() - start_time) + time_to_sleep
+          >= timeout
+      ):
+        raise TimeoutError(
+            f"Timed out waiting for {slice_count} slices. Only"
+            f" {len(good_slice_indices)} available after"
+            f" {elapsed_time:.2f} seconds."
+            f" Next check would occur after the timeout of {timeout}"
+            " seconds."
+        )
+
+      if time_to_sleep > 0:
+        _logger.info("Sleeping for %s seconds.", time_to_sleep)
+
+        time.sleep(time_to_sleep)
 
   def pause_resume(
       self,
       max_retries: int,
-      wait_period: float | int = 10,
+      poll_interval: float | int = 10,
       timeout: float | None = None,
   ) -> Any:
     """Retries a function with pause/resume fault tolerance.
@@ -723,7 +734,7 @@ def pause_resume(
 
     Args:
       max_retries: The maximum number of times to retry the function.
-      wait_period: The number of seconds to wait between availability checks.
+      poll_interval: The number of seconds to wait between availability checks.
         Defaults to 10 seconds.
       timeout: The maximum number of seconds to wait for slices to become
         available before each retry attempt. If None, there is no timeout.
@@ -745,7 +756,7 @@ def wrapper(*args, **kwargs):
                 "Elastic attempt %d out of %d", retry_index + 1, max_retries
             )
 
-            self.wait_for_slices(wait_period=wait_period, timeout=timeout)
+            self.wait_for_slices(poll_interval=poll_interval, timeout=timeout)
 
             return func(*args, **kwargs)
           except jax.errors.JaxRuntimeError as error: