OSS async_utils_test.py and respect operation if passed to create_all_async when updating await-able signal contract.

PiperOrigin-RevId: 842618298

Author: mridul-sahu

checkpoint/orbax/checkpoint/_src/futures/future.py

@@ -102,27 +102,32 @@ def get_awaitable_signals_from_contract(
   def add_to_awaitable_signals_contract(
       cls,
       signals: Sequence[synchronization.HandlerAwaitableSignal],
+      operation_id: str | None = None,
   ):
     """Adds awaitable signals to `AWAITABLE_SIGNALS_CONTRACT` for lower checkpointing layers to wait on.
 
-    These signals are added to the list of awaitable signals for the current
-    operation id in `OperationIdGenerator`.
+    These signals are added to the list of awaitable signals for the given
+    operaion id or the current operation id in `OperationIdGenerator`.
 
     Args:
       signals: The signals to add to the list of awaitable signals.
+      operation_id: The operation id to use for the barrier keys. If None, the
+        current operation id is used.
     """
     if not signals:
       return
-
-    current_signals = list(cls.get_awaitable_signals_from_contract())
+    operation_id = (
+        operation_id
+        or synchronization.OperationIdGenerator.get_current_operation_id()
+    )
+    current_signals = list(
+        cls.get_awaitable_signals_from_contract(operation_id=operation_id)
+    )
     current_signals.extend(signals)
     keys = ','.join(
         [current_signal.value for current_signal in current_signals]
     )
     client = signaling_client.get_signaling_client()
-    operation_id = (
-        synchronization.OperationIdGenerator.get_current_operation_id()
-    )
     barrier_key = cls.get_unique_awaitable_singal_key(
         synchronization.HandlerAwaitableSignal.AWAITABLE_SIGNALS_CONTRACT,
         operation_id,

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

@@ -723,7 +723,7 @@ def create_all_async(
         operation_id=operation_id,
     )
     future.AwaitableSignalsContract.add_to_awaitable_signals_contract(
-        completion_signals
+        completion_signals, operation_id=operation_id
     )
 
   # Sync to enusre that all hosts have the same awaitable signals contract.

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

@@ -0,0 +1,124 @@
+# 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.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):
+    tmpdir = atomicity.AtomicRenameTemporaryPath(
+        self.directory / 'tmp', self.directory / 'final'
+    )
+    start = time.time()
+    p = async_utils.start_async_mkdir(tmpdir, operation_id='op1')
+    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):
+    tmpdir = atomicity.AtomicRenameTemporaryPath(
+        self.directory / 'tmp', self.directory / 'final'
+    )
+    start = time.time()
+    p = async_utils.start_async_mkdir(tmpdir, ['a', 'b'], operation_id='op1')
+    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()
+    p1 = async_utils.start_async_mkdir(tmpdir1, operation_id='op1')
+    p2 = async_utils.start_async_mkdir(tmpdir2, operation_id='op2')
+    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()
+    p1 = async_utils.start_async_mkdir(tmpdir1, operation_id='op1')
+    p2 = async_utils.start_async_mkdir(tmpdir2, operation_id='op2')
+    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):
+    tmpdir = atomicity.AtomicRenameTemporaryPath(
+        self.directory / 'tmp', self.directory / 'final'
+    )
+    p = async_utils.start_async_mkdir(tmpdir, operation_id='op1')
+    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()