Make the create_directories_asynchronously option default to True…

… if jax distributed client is initialized. PiperOrigin-RevId: 730743121
google · Feb 27, 2025 · cfcd833 · cfcd833
1 parent acec3f3
commit cfcd833
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/checkpoint/orbax/checkpoint/options.py b/checkpoint/orbax/checkpoint/options.py
@@ -31,7 +31,9 @@ class AsyncOptions:
   timeout_secs: int = 600  # 10 minutes. Same as default in `AsyncCheckpointer`.
   barrier_sync_fn: Optional[multihost.BarrierSyncFn] = None
   post_finalization_callback: Optional[Callable[[], None]] = None
-  create_directories_asynchronously: bool = False
+  create_directories_asynchronously: bool = (
+      multihost.is_jax_distributed_client_initialized()
+  )
 
 
 @dataclasses.dataclass