Fix possible race condition in orbax root directory creation.

PiperOrigin-RevId: 649014114
google · Jul 11, 2024 · 557613f · 557613f
1 parent 1e06498
commit 557613f
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/checkpoint/orbax/checkpoint/checkpoint_manager.py b/checkpoint/orbax/checkpoint/checkpoint_manager.py
@@ -364,7 +364,9 @@ def _create_root_directory(
   if not directory.exists() and utils.is_primary_host(
       multiprocessing_options.primary_host
   ):
-    directory.mkdir(parents=True)
+    # We need exists_ok=True because the directory might have been created due
+    # to a race condition.
+    directory.mkdir(parents=True, exist_ok=True)
     logging.info('Created directory=%s', directory)
   multihost.sync_global_processes(
       multihost.unique_barrier_key(