Add overall timeout in AsyncCheckpointer.

PiperOrigin-RevId: 878104670

Author: JustinPan-goog

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

@@ -14,6 +14,7 @@
 
 """AsyncCheckpointer."""
 
+import datetime
 import sys
 import threading
 import time
@@ -69,23 +70,27 @@ def _background_wait_for_commit_futures(
     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."""
   current_process = multihost.process_index()
   current_thread_id = threading.current_thread().name
   process_count = jax.process_count()
   logging.info(
-      '[process=%s][thread=%s] Background save thread started.',
+      '[process=%s][thread=%s] Background save thread started. Deadline for'
+      ' this save operation is %s',
       current_process,
       current_thread_id,
+      datetime.datetime.now() + datetime.timedelta(seconds=timeout_secs),
   )
   thread_start_time = time.time()
 
   # Wait for commit operations to complete.
-  for commit_future in commit_futures:
-    commit_future.result()
+  future.ChainedFuture(commit_futures, cb=lambda: None).result(
+      timeout=timeout_secs
+  )
   commit_duration_secs = time.time() - thread_start_time
   logging.info(
       '[process=%s][thread=%s] %d Handler Commit operations completed. Time'
@@ -111,30 +116,48 @@ 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 = future.get_remaining_time(
+          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):
         if 'DEADLINE_EXCEEDED' in str(e):
           _add_deadline_exceeded_notes(e)
-      raise
+      raise TimeoutError(
+          'Timed out while waiting for async_write_complete barrier.'
+      ) from e
 
   if utils.is_primary_host(primary_host):
     on_commit_callback()
   if process_count > 1:
     # Block until process 0 completes on_commit_callback.
-    sync_fn(
-        multihost.unique_barrier_key(
-            'async_commit_complete',
-            prefix=barrier_sync_key_prefix,
-            suffix=f'{directory.name}',
-        )
-    )
+    try:
+      time_remaining_secs = future.get_remaining_time(
+          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),
+      )
+    except jax.errors.JaxRuntimeError as e:
+      if sys.version_info >= (3, 11):
+        if 'DEADLINE_EXCEEDED' in str(e):
+          _add_deadline_exceeded_notes(e)
+      raise TimeoutError(
+          'Timed out while waiting for async_commit_complete barrier.'
+      ) from e
 
   thread_duration_secs = time.time() - thread_start_time
   jax.monitoring.record_event_duration_secs(
@@ -190,9 +213,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 +240,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

checkpoint/orbax/checkpoint/_src/futures/future.py

@@ -170,14 +170,14 @@ class Future(Protocol):
   future, but merely wait for an ongoing operation to complete.
   """
 
-  def result(self, timeout: Optional[int] = None) -> Any:
+  def result(self, timeout: Optional[float] = None) -> Any:
     """Waits for the future to complete its operation."""
     ...
 
 
 class NoopFuture:
 
-  def result(self, timeout: Optional[int] = None) -> Any:
+  def result(self, timeout: Optional[float] = None) -> Any:
     del timeout
     return None
 
@@ -189,21 +189,18 @@ def __init__(self, futures: Sequence[Future], cb: Callable[[], None]):
     self._futures = futures
     self._cb = cb
 
-  def result(self, timeout: Optional[int] = None) -> Any:
+  def result(self, timeout: Optional[float] = None) -> Any:
     """Waits for all futures to complete."""
     n = len(self._futures)
     start = time.time()
-    time_remaining = timeout
     for k, f in enumerate(self._futures):
-      f.result(timeout=time_remaining)
-      if time_remaining is not None:
-        time_elapsed = time.time() - start
-        time_remaining -= time_elapsed
-        if time_remaining <= 0:
-          raise TimeoutError(
-              'ChainedFuture completed {:d}/{:d} futures but timed out after'
-              ' {:.2f} seconds.'.format(k, n, time_elapsed)
-          )
+      try:
+        f.result(timeout=get_remaining_time(start, timeout))
+      except TimeoutError as e:
+        raise TimeoutError(
+            f'ChainedFuture completed {k}/{n} futures but timed out after'
+            f' {time.time() - start:.2f} seconds.'
+        ) from e
     time_elapsed = time.time() - start
     logging.vlog(
         1,
@@ -215,6 +212,18 @@ def result(self, timeout: Optional[int] = None) -> Any:
     self._cb()
 
 
+def get_remaining_time(
+    start_time: float, timeout_secs: Optional[float]
+) -> Optional[float]:
+  """Returns remaining time in secs, or None if timeout_secs is None."""
+  if timeout_secs is None:
+    return None
+  elapsed = time.time() - start_time
+  if elapsed >= timeout_secs:
+    raise TimeoutError(f'Timed out after {elapsed} seconds.')
+  return timeout_secs - elapsed
+
+
 def wait_for_signals(
     receive_signals: Sequence[synchronization.HandlerAwaitableSignal],
     *,