From 5046217cad19ba78373730b440bf965e3858d58f Mon Sep 17 00:00:00 2001 From: Colin Gaffney Date: Fri, 16 Aug 2024 10:08:41 -0700 Subject: [PATCH] Parallelize many directory creations using asyncio. Also reduce the number of `asyncio.run` calls by moving `async` functions higher in the stack. PiperOrigin-RevId: 663767760 --- paxml/BUILD | 1 - paxml/checkpoints_test.py | 6 +++++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/paxml/BUILD b/paxml/BUILD index 6d5d9df6..cb7def88 100644 --- 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 index d4470dea..d0045ced 100644 --- 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