Parallelize many directory creations using asyncio. Also reduce the n…

…umber of `asyncio.run` calls by moving `async` functions higher in the stack. PiperOrigin-RevId: 663767760
google · Aug 16, 2024 · 5046217 · 5046217
1 parent c36b3c6
commit 5046217
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/paxml/BUILD b/paxml/BUILD
@@ -167,7 +167,6 @@ pytype_strict_test(
     deps = [
         ":checkpoint_metadata",
         ":checkpoint_paths",
-        ":checkpoint_types",
         ":checkpoints",
         ":train_states",
         # Implicit absl.testing.absltest.absltest dependency.

diff --git a/paxml/checkpoints_test.py b/paxml/checkpoints_test.py
@@ -39,7 +39,11 @@
 
 
 def _create_tmp_directory(final_dir):
-  if hasattr(ocp.path, 'atomicity'):
+  if hasattr(ocp.test_utils, 'create_tmp_directory'):
+    tmp_dir = ocp.path.atomicity._get_tmp_directory(final_dir)
+    ocp.test_utils.create_tmp_directory(tmp_dir, final_dir)
+    return tmp_dir
+  elif hasattr(ocp.path.atomicity, '_create_tmp_directory'):
     tmp_dir = ocp.path.atomicity._get_tmp_directory(final_dir)
     ocp.path.atomicity._create_tmp_directory(tmp_dir, final_dir)
     return tmp_dir