Internal change

PiperOrigin-RevId: 877831005

Author: mxberlot

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

@@ -54,6 +54,7 @@
 
 import abc
 import asyncio
+import concurrent.futures
 import pickle
 import threading
 import time
@@ -228,6 +229,42 @@ def get_awaitable_path(self) -> path_types.PathAwaitingCreation:
     ...
 
 
+class DeferredPath(path_types.PathAwaitingCreation):
+  """A path that is created asynchronously and can be awaited.
+
+  Uses concurrent.futures.Future instead of asyncio.Task to avoid
+  event loop binding issues when create() runs in a different thread.
+  The Future is thread-safe and can be awaited from any event loop.
+  """
+
+  def __init__(self):
+    self._future_path: concurrent.futures.Future[epath.Path] = (
+        concurrent.futures.Future()
+    )
+
+  def set_path(self, path: epath.Path) -> None:
+    """Sets the path result. Called by create() when allocation completes."""
+    self._future_path.set_result(path)
+
+  def __truediv__(
+      self, other: path_types.PathLike
+  ) -> path_types.PathAwaitingCreation:
+    child = DeferredPath()
+    self._future_path.add_done_callback(
+        lambda f: child.set_path(f.result() / other)
+    )
+    return child
+
+  @property
+  def path(self) -> epath.Path:
+    if not self._future_path.done():
+      raise ValueError('Path has not been created yet. Call await_creation().')
+    return self._future_path.result()
+
+  async def await_creation(self) -> epath.Path:
+    return await asyncio.wrap_future(self._future_path)
+
+
 class ReadOnlyTemporaryPath(atomicity_types.TemporaryPath):
   """A read-only, serializable object providing path properties access.
 

checkpoint/orbax/checkpoint/_src/path/atomicity_test.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import asyncio
+import concurrent.futures
 import stat
 import unittest
 from absl.testing import absltest
@@ -205,6 +207,33 @@ async def test_finalize_raises(self):
       )
 
 
+class DeferredPathTest(absltest.TestCase):
+
+  def test_set_and_get_path(self):
+    dp = atomicity.DeferredPath()
+    test_path = epath.Path('/test/path')
+    dp.set_path(test_path)
+    self.assertEqual(dp.path, test_path)
+
+  def test_path_before_set_raises(self):
+    dp = atomicity.DeferredPath()
+    with self.assertRaises(ValueError):
+      _ = dp.path
+
+  def test_await_creation(self):
+    dp = atomicity.DeferredPath()
+    test_path = epath.Path('/test/path')
+    dp.set_path(test_path)
+    result = asyncio.run(dp.await_creation())
+    self.assertEqual(result, test_path)
+
+  def test_set_path_twice_raises(self):
+    dp = atomicity.DeferredPath()
+    dp.set_path(epath.Path('/first'))
+    with self.assertRaises(concurrent.futures.InvalidStateError):
+      dp.set_path(epath.Path('/second'))
+
+
 
 if __name__ == '__main__':
   absltest.main()