Increase t5x AsyncCheckpointer timeout to 600s.

PiperOrigin-RevId: 629170639
kkpan11 · Apr 29, 2024 · 130910a · 130910a
1 parent e4e262b
commit 130910a
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/t5x/checkpoints.py b/t5x/checkpoints.py
@@ -2310,7 +2310,9 @@ def __init__(
     )
     # TODO(b/273803615) Enable OCDBT.
     self._state_handler = ocp.PyTreeCheckpointHandler(use_ocdbt=False)
-    checkpointers = {_STATE_KEY: ocp.AsyncCheckpointer(self._state_handler)}
+    checkpointers = {
+        _STATE_KEY: ocp.AsyncCheckpointer(self._state_handler, timeout_secs=600)
+    }
     if self._should_write_dataset_ckpt:
       checkpointers[_DATASET_KEY] = ocp.Checkpointer(
           DatasetCheckpointHandler(checkpoint_filename=dataset_ckpt_name)