Internal change

PiperOrigin-RevId: 878104670

Author: JustinPan-goog

checkpoint/orbax/checkpoint/_src/checkpointers/async_checkpointer.py

@@ -63,13 +63,25 @@ def _on_commit_callback(
   )
 
 
+def _check_for_timeout(
+    thread_start_time: float,
+    timeout_secs: int,
+):
+  """Checks if the timeout has been exceeded."""
+  time_remaining_secs = timeout_secs - (time.time() - thread_start_time)
+  if time_remaining_secs <= 0:
+    raise TimeoutError()
+  return time_remaining_secs
+
+
 def _background_wait_for_commit_futures(
     directory: epath.Path,
     commit_futures: Sequence[future.Future],
     on_commit_callback: Callable[[], None],
     *,
     barrier_sync_key_prefix: str,
-    sync_fn: Callable[[str], None],
+    sync_fn: Callable[[str, int], None],
+    timeout_secs: int,
     primary_host: int | None,
 ):
   """A function to be run in a background thread that waits for futures."""
@@ -85,7 +97,8 @@ def _background_wait_for_commit_futures(
 
   # Wait for commit operations to complete.
   for commit_future in commit_futures:
-    commit_future.result()
+    time_remaining_secs = _check_for_timeout(thread_start_time, timeout_secs)
+    commit_future.result(timeout=int(time_remaining_secs))
   commit_duration_secs = time.time() - thread_start_time
   logging.info(
       '[process=%s][thread=%s] %d Handler Commit operations completed. Time'
@@ -111,12 +124,14 @@ def _background_wait_for_commit_futures(
     # All processes will wait at the barrier. When all processes are at the
     # barrier, the barrier will be satisfied. If not, then it will timeout.
     try:
+      time_remaining_secs = _check_for_timeout(thread_start_time, timeout_secs)
       sync_fn(
           multihost.unique_barrier_key(
               'async_write_complete',
               prefix=barrier_sync_key_prefix,
               suffix=f'{directory.name}',
-          )
+          ),
+          int(time_remaining_secs * 1000),
       )
     except jax.errors.JaxRuntimeError as e:
       if sys.version_info >= (3, 11):
@@ -128,12 +143,14 @@ def _background_wait_for_commit_futures(
     on_commit_callback()
   if process_count > 1:
     # Block until process 0 completes on_commit_callback.
+    time_remaining_secs = _check_for_timeout(thread_start_time, timeout_secs)
     sync_fn(
         multihost.unique_barrier_key(
             'async_commit_complete',
             prefix=barrier_sync_key_prefix,
             suffix=f'{directory.name}',
-        )
+        ),
+        int(time_remaining_secs * 1000),
     )
 
   thread_duration_secs = time.time() - thread_start_time
@@ -190,9 +207,8 @@ def __init__(
     self._thread = None
     self._exception = None
 
-    timeout_in_ms = self._timeout_secs * 1000
-    self._sync_fn: Callable[[str], None] = lambda key: barrier_sync_fn(
-        key=key, timeout_ms=timeout_in_ms
+    self._sync_fn: Callable[[str, int], None] = (
+        lambda key, timeout_ms: barrier_sync_fn(key=key, timeout_ms=timeout_ms)
     )
 
   def __del__(self):
@@ -218,6 +234,7 @@ def _thread_func(
           on_commit_callback,
           barrier_sync_key_prefix=self._barrier_sync_key_prefix,
           sync_fn=self._sync_fn,
+          timeout_secs=self._timeout_secs,
           primary_host=self._primary_host,
       )
     except Exception as e:  # pylint: disable=broad-exception-caught