OSS async_utils_test.py and use operationId generator instead of passing operation id around.

PiperOrigin-RevId: 842618298

Author: mridul-sahu

checkpoint/orbax/checkpoint/_src/path/atomicity.py

@@ -675,7 +675,6 @@ def create_all_async(
     *,
     multiprocessing_options: options_lib.MultiprocessingOptions | None = None,
     subdirectories: Sequence[str] | None = None,
-    operation_id: str | None = None,
 ) -> future.Future:
   """Creates all temporary paths in parallel asynchronously.
 
@@ -688,8 +687,6 @@ def create_all_async(
     subdirectories: Sequence of subdirectories to create under `paths`. If not
       provided, no subdirectories will be created. The same set of
       subdirectories will be created under each path in `paths`.
-    operation_id: The operation id to use for the barrier keys. If None, the
-      current operation id is used.
 
   Returns:
     A future that which sends the completion signals when all paths are created.
@@ -720,7 +717,6 @@ def create_all_async(
         ),
         send_signals=completion_signals,
         timeout_secs=multihost.coordination_timeout(),
-        operation_id=operation_id,
     )
     future.AwaitableSignalsContract.add_to_awaitable_signals_contract(
         completion_signals

checkpoint/orbax/checkpoint/experimental/v1/_src/path/async_utils.py

@@ -56,7 +56,6 @@ async def await_creation(self) -> Path:
 def start_async_mkdir(
     path: atomicity_types.TemporaryPath,
     subdirectories: Iterable[str] = (),
-    operation_id: str | None = None,
 ) -> PathAwaitingCreation:
   """Starts async directory creation on a TemporaryPath.
 
@@ -76,8 +75,6 @@ def start_async_mkdir(
   Args:
     path: The path to create. May be an instance of `TemporaryPath`.
     subdirectories: A sequence of subdirectories to create under `path`.
-    operation_id: The operation id to use for the barrier keys. If None, the
-      current operation id is used.
 
   Returns:
     A PathAwaitingCreation object.
@@ -97,6 +94,5 @@ def start_async_mkdir(
       completion_signals=completion_signals,
       multiprocessing_options=context.multiprocessing_options.v0(),
       subdirectories=[name for name in subdirectories],
-      operation_id=operation_id,
   )
   return _PathAwaitingCreation(path.get(), f)

checkpoint/orbax/checkpoint/experimental/v1/_src/path/async_utils_test.py

@@ -0,0 +1,132 @@
+# Copyright 2025 The Orbax Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import asyncio
+import time
+import unittest
+from unittest import mock
+from absl.testing import absltest
+from etils import epath
+from orbax.checkpoint._src.path import async_path
+from orbax.checkpoint._src.path import atomicity
+from orbax.checkpoint.experimental.v1._src.context import context as context_lib
+from orbax.checkpoint.experimental.v1._src.path import async_utils
+
+
+class AsyncUtilsTest(absltest.TestCase, unittest.IsolatedAsyncioTestCase):
+
+  def setUp(self):
+    super().setUp()
+
+    orig_create_paths = atomicity._create_paths
+
+    async def _sleep_and_create_paths(
+        *args,
+        **kwargs,
+    ):
+      await asyncio.sleep(1)
+      return await orig_create_paths(*args, **kwargs)
+
+    self.enter_context(
+        mock.patch.object(
+            atomicity, '_create_paths', new=_sleep_and_create_paths
+        )
+    )
+    self.directory = epath.Path(self.create_tempdir().full_path)
+
+  async def assertExists(self, path: epath.Path):
+    self.assertTrue(await async_path.exists(path))
+
+  async def assertNotExists(self, path: epath.Path):
+    self.assertFalse(await async_path.exists(path))
+
+  def assertBetween(self, a, b, c):
+    self.assertGreater(b, a)
+    self.assertGreater(c, b)
+
+  async def test_async_mkdir(self):
+    await context_lib.synchronize_next_operation_id()
+    tmpdir = atomicity.AtomicRenameTemporaryPath(
+        self.directory / 'tmp', self.directory / 'final'
+    )
+    start = time.time()
+    p = async_utils.start_async_mkdir(tmpdir)
+    await p.await_creation()
+    self.assertBetween(1, time.time() - start, 2)
+    await self.assertExists(self.directory / 'tmp')
+
+  async def test_async_mkdir_with_subdirectories(self):
+    await context_lib.synchronize_next_operation_id()
+    tmpdir = atomicity.AtomicRenameTemporaryPath(
+        self.directory / 'tmp', self.directory / 'final'
+    )
+    start = time.time()
+    p = async_utils.start_async_mkdir(tmpdir, ['a', 'b'])
+    await p.await_creation()
+    self.assertBetween(1, time.time() - start, 2)
+    await self.assertExists(self.directory / 'tmp')
+    await self.assertExists(self.directory / 'tmp' / 'a')
+    await self.assertExists(self.directory / 'tmp' / 'b')
+
+  async def test_async_mkdir_sequential(self):
+    tmpdir1 = atomicity.AtomicRenameTemporaryPath(
+        self.directory / 'tmp1', self.directory / 'final1'
+    )
+    tmpdir2 = atomicity.AtomicRenameTemporaryPath(
+        self.directory / 'tmp2', self.directory / 'final2'
+    )
+    start = time.time()
+    await context_lib.synchronize_next_operation_id()
+    p1 = async_utils.start_async_mkdir(tmpdir1)
+    await context_lib.synchronize_next_operation_id()
+    p2 = async_utils.start_async_mkdir(tmpdir2)
+    await p1.await_creation()
+    await p2.await_creation()
+    # Awaiting sequentially does not take any longer than awaiting in parallel,
+    # because the operations are already in progress.
+    self.assertBetween(1, time.time() - start, 2)
+    await self.assertExists(self.directory / 'tmp1')
+    await self.assertExists(self.directory / 'tmp2')
+
+  async def test_async_mkdir_parallel(self):
+    tmpdir1 = atomicity.AtomicRenameTemporaryPath(
+        self.directory / 'tmp1', self.directory / 'final1'
+    )
+    tmpdir2 = atomicity.AtomicRenameTemporaryPath(
+        self.directory / 'tmp2', self.directory / 'final2'
+    )
+    start = time.time()
+    await context_lib.synchronize_next_operation_id()
+    p1 = async_utils.start_async_mkdir(tmpdir1)
+    await context_lib.synchronize_next_operation_id()
+    p2 = async_utils.start_async_mkdir(tmpdir2)
+    await asyncio.gather(p1.await_creation(), p2.await_creation())
+    self.assertBetween(1, time.time() - start, 2)
+    await self.assertExists(self.directory / 'tmp1')
+    await self.assertExists(self.directory / 'tmp2')
+
+  async def test_async_mkdir_with_delayed_wait(self):
+    await context_lib.synchronize_next_operation_id()
+    tmpdir = atomicity.AtomicRenameTemporaryPath(
+        self.directory / 'tmp', self.directory / 'final'
+    )
+    p = async_utils.start_async_mkdir(tmpdir)
+    await self.assertNotExists(self.directory / 'tmp')
+    await asyncio.sleep(1)
+    await p.await_creation()
+    await self.assertExists(self.directory / 'tmp')
+
+
+if __name__ == '__main__':
+  absltest.main()