Internal Change

PiperOrigin-RevId: 870544311

Author: JustinPan-goog

checkpoint/orbax/checkpoint/_src/handlers/composite_checkpoint_handler.py

@@ -1044,7 +1044,7 @@ def finalize(self, directory: epath.Path):
         # Not an error, as some items may not have been saved.
         continue
       handler.finalize(tmp_dir.get())
-      asyncio.run(
+      asyncio_utils.run_sync(
           tmp_dir.finalize(
           )
       )

checkpoint/orbax/checkpoint/checkpoint_utils.py

@@ -14,7 +14,6 @@
 
 """High-level checkpoint utils provided for user convenience."""
 
-import asyncio
 import contextlib
 import time
 from typing import Any, Callable, Iterator, Optional
@@ -25,6 +24,7 @@
 from jax.experimental import layout
 import numpy as np
 from orbax.checkpoint import utils
+from orbax.checkpoint._src import asyncio_utils
 from orbax.checkpoint._src.arrays import sharding as arrays_sharding_lib
 from orbax.checkpoint._src.metadata import tree as tree_metadata
 from orbax.checkpoint._src.metadata import value as value_metadata
@@ -121,7 +121,7 @@ def _snapshot_checkpoint(
           f'Ignoring error when snapshotting checkpoint for step: {step}'
       ),
   ):
-    asyncio.run(snapshot_impl.create_snapshot())
+    asyncio_utils.run_sync(snapshot_impl.create_snapshot())
   return True
 
 
@@ -146,7 +146,7 @@ def _release_snapshot(
             f'Ignoring error when releasing snapshot for step: {step}'
         ),
     ):
-      asyncio.run(snapshot_impl.release_snapshot())
+      asyncio_utils.run_sync(snapshot_impl.release_snapshot())
 
 
 def _reached_desired_step(step: int, until_step: Optional[int]) -> bool:
@@ -409,7 +409,7 @@ def checkpoints_iterator(
               f' {step_dir.name}'
           ),
       ):
-        asyncio.run(snapshot_impl.release_snapshot())
+        asyncio_utils.run_sync(snapshot_impl.release_snapshot())
   checkpoint_step = None
   while True:
     until_step = checkpoint_step + 1 if checkpoint_step is not None else None

checkpoint/orbax/checkpoint/experimental/emergency/checkpoint_manager.py

@@ -25,7 +25,6 @@
 always be called across all processes within the primary slice.
 """
 
-import asyncio
 import dataclasses
 import functools
 import time
@@ -40,6 +39,7 @@
 from orbax.checkpoint import args as args_lib
 from orbax.checkpoint import checkpoint_manager
 from orbax.checkpoint import checkpoint_utils
+from orbax.checkpoint._src import asyncio_utils
 from orbax.checkpoint._src.checkpoint_managers import preservation_policy as preservation_policy_lib
 from orbax.checkpoint._src.checkpoint_managers import save_decision_policy as save_decision_policy_lib
 from orbax.checkpoint._src.handlers import handler_registration
@@ -991,7 +991,7 @@ def _restore_from_local(
             1,
             'Debugging single-slice restore_args used for restoration.',
         )
-        asyncio.run(
+        asyncio_utils.run_sync(
             local_checkpoint_data_debugging.print_chunk_debug_info(
                 restore_directory / _STATE_ITEM_NAME,
                 single_slice_restore_args,

checkpoint/orbax/checkpoint/experimental/v1/_src/layout/orbax_layout.py

@@ -19,6 +19,7 @@
 from typing import Any, Awaitable
 
 from absl import logging
+from orbax.checkpoint._src import asyncio_utils
 from orbax.checkpoint._src.path import async_path
 from orbax.checkpoint._src.path import temporary_paths
 from orbax.checkpoint.experimental.v1._src.context import context as context_lib
@@ -75,7 +76,7 @@ def is_orbax_v1_checkpoint(path: path_types.PathLike) -> bool:
   ctx = context_lib.get_context()
   path = ctx.file_options.path_class(path)
   try:
-    asyncio.run(OrbaxLayout().validate(path))
+    asyncio_utils.run_sync(OrbaxLayout().validate(path))
     return True
   except InvalidLayoutError:
     return False

checkpoint/orbax/checkpoint/experimental/v1/_src/layout/orbax_v0_layout.py

@@ -18,6 +18,7 @@
 import logging
 from typing import Any, Awaitable
 
+from orbax.checkpoint._src import asyncio_utils
 from orbax.checkpoint._src.metadata import step_metadata_serialization
 from orbax.checkpoint._src.path import async_path
 from orbax.checkpoint._src.path import temporary_paths
@@ -49,7 +50,7 @@ def is_orbax_v0_checkpoint(path: path_types.PathLike) -> bool:
   ctx = context_lib.get_context()
   path = ctx.file_options.path_class(path)
   try:
-    asyncio.run(OrbaxV0Layout().validate(path))
+    asyncio_utils.run_sync(OrbaxV0Layout().validate(path))
     return True
   except InvalidLayoutError:
     return False

checkpoint/orbax/checkpoint/experimental/v1/_src/layout/registry.py

@@ -15,6 +15,8 @@
 """Registry for checkpoint layouts."""
 
 import asyncio
+
+from orbax.checkpoint._src import asyncio_utils
 from orbax.checkpoint.experimental.v1._src.context import context as context_lib
 from orbax.checkpoint.experimental.v1._src.context import options as options_lib
 from orbax.checkpoint.experimental.v1._src.layout import checkpoint_layout
@@ -25,6 +27,7 @@
 from orbax.checkpoint.experimental.v1._src.layout import safetensors_layout
 from orbax.checkpoint.experimental.v1._src.path import types as path_types
 
+
 InvalidLayoutError = checkpoint_layout.InvalidLayoutError
 CheckpointLayout = checkpoint_layout.CheckpointLayout
 CheckpointLayoutEnum = options_lib.CheckpointLayout
@@ -49,7 +52,7 @@ async def _is_orbax_checkpoint_async(path: path_types.PathLike) -> bool:
 
 def is_orbax_checkpoint(path: path_types.PathLike) -> bool:
   """Returns True if the path is an Orbax checkpoint."""
-  return asyncio.run(_is_orbax_checkpoint_async(path))
+  return asyncio_utils.run_sync(_is_orbax_checkpoint_async(path))
 
 
 async def get_layout_class(

checkpoint/orbax/checkpoint/experimental/v1/_src/loading/loading.py

@@ -14,12 +14,12 @@
 
 """Defines free-function interface for loading."""
 
-import asyncio
 import functools
 import time
 from typing import Any, Awaitable, Protocol
 
 from absl import logging
+from orbax.checkpoint._src import asyncio_utils
 from orbax.checkpoint._src.logging import event_tracking
 from orbax.checkpoint.experimental.v1._src.context import context as context_lib
 import orbax.checkpoint.experimental.v1._src.handlers.global_registration  # pylint: disable=unused-import
@@ -115,7 +115,7 @@ def load_pytree(
   logging.info('Loading checkpoint from %s.', path)
   ctx = context_lib.get_context()
   path = ctx.file_options.path_class(path)
-  layout = asyncio.run(
+  layout = asyncio_utils.run_sync(
       layout_registry.get_checkpoint_layout_pytree(
           path, ctx.checkpoint_layout, checkpointable_name
       )
@@ -204,7 +204,7 @@ def load_checkpointables(
   logging.info('Loading checkpoint from %s.', path)
   ctx = context_lib.get_context()
   path = ctx.file_options.path_class(path)
-  layout = asyncio.run(
+  layout = asyncio_utils.run_sync(
       layout_registry.get_checkpoint_layout(path, ctx.checkpoint_layout)
   )
 
@@ -264,7 +264,7 @@ async def _load() -> Any:
     )
     return result
 
-  result = asyncio.run(_load())
+  result = asyncio_utils.run_sync(_load())
 
   event_tracking.record_read_event(path)
 

checkpoint/orbax/checkpoint/experimental/v1/_src/metadata/loading.py

@@ -14,9 +14,9 @@
 
 """Functions for loading metadata from a checkpoint."""
 
-import asyncio
 from typing import Any
 
+from orbax.checkpoint._src import asyncio_utils
 from orbax.checkpoint.experimental.v1 import errors
 from orbax.checkpoint.experimental.v1._src.context import context as context_lib
 import orbax.checkpoint.experimental.v1._src.handlers.global_registration  # pylint: disable=unused-import
@@ -88,7 +88,7 @@ def _get_abstract_array(arr):
   ctx = context_lib.get_context()
   path = ctx.file_options.path_class(path)
 
-  layout = asyncio.run(
+  layout = asyncio_utils.run_sync(
       layout_registry.get_checkpoint_layout_pytree(
           path, ctx.checkpoint_layout, checkpointable_name
       )
@@ -140,7 +140,7 @@ def checkpointables_metadata(
   """
   ctx = context_lib.get_context()
   path = ctx.file_options.path_class(path)
-  layout = asyncio.run(
+  layout = asyncio_utils.run_sync(
       layout_registry.get_checkpoint_layout(path, ctx.checkpoint_layout)
   )
   return _checkpointables_metadata_impl(layout, path)
@@ -157,4 +157,4 @@ async def _load_metadata() -> (
   ):
     return await layout.metadata(path)
 
-  return asyncio.run(_load_metadata())
+  return asyncio_utils.run_sync(_load_metadata())

checkpoint/orbax/checkpoint/experimental/v1/_src/saving/execution.py

@@ -16,7 +16,6 @@
 
 from __future__ import annotations
 
-import asyncio
 import hashlib
 import time
 from typing import Any, Awaitable, Iterable
@@ -25,6 +24,7 @@
 from absl import logging
 import jax
 import numpy as np
+from orbax.checkpoint._src import asyncio_utils
 from orbax.checkpoint._src.futures import future
 from orbax.checkpoint._src.logging import event_tracking
 from orbax.checkpoint._src.metadata import step_metadata_serialization
@@ -380,7 +380,7 @@ def save_checkpointables_impl(
   event_tracking.record_save_start(path, async_origin=async_origin)
   # Ensure the operation ID is incremented as soon as possible. This must be
   # done uniquely for each save operation.
-  asyncio.run(context_lib.synchronize_next_operation_id())
+  asyncio_utils.run_sync(context_lib.synchronize_next_operation_id())
   context = context_lib.get_context()
 
   path = context.file_options.path_class(path)
@@ -397,7 +397,7 @@ def save_checkpointables_impl(
       subdirectories=subdirectories,
       use_snapshot=path_exists,
   )
-  background_awaitable = asyncio.run(
+  background_awaitable = asyncio_utils.run_sync(
       _run_blocking_save(
           temporary_path,
           checkpointables,

checkpoint/orbax/checkpoint/experimental/v1/_src/testing/handler_utils.py

@@ -16,7 +16,6 @@
 
 from __future__ import annotations
 
-import asyncio
 import dataclasses
 import json
 from typing import Any, Awaitable, Generic, Type, TypeVar
@@ -25,6 +24,7 @@
 from etils import epath
 from orbax.checkpoint import checkpoint_args as v0_args
 from orbax.checkpoint import handlers as v0_handlers
+from orbax.checkpoint._src import asyncio_utils
 from orbax.checkpoint.experimental.v1._src.context import context as context_lib
 from orbax.checkpoint.experimental.v1._src.handlers import types as handler_types
 from orbax.checkpoint.experimental.v1._src.path import types as path_types
@@ -59,24 +59,26 @@ def __init__(
 
   def save(self, directory: Path, checkpointable: T):
     path = path_test_utils.PathAwaitingCreationWrapper(directory)
-    awaitable = asyncio.run(self._handler.save(path, checkpointable))
-    return asyncio.run(_run_awaitable(awaitable))
+    awaitable = asyncio_utils.run_sync(self._handler.save(path, checkpointable))
+    return asyncio_utils.run_sync(_run_awaitable(awaitable))
 
   def save_async(self, directory: Path, checkpointable: T):
     path = path_test_utils.PathAwaitingCreationWrapper(directory)
-    return asyncio.run(self._handler.save(path, checkpointable))
+    return asyncio_utils.run_sync(self._handler.save(path, checkpointable))
 
   def load(self, path: Path, abstract_checkpointable: AbstractT | None = None):
     awaitable = self.load_async(path, abstract_checkpointable)
-    return asyncio.run(_run_awaitable(awaitable))
+    return asyncio_utils.run_sync(_run_awaitable(awaitable))
 
   def load_async(
       self, path: Path, abstract_checkpointable: AbstractT | None = None
   ):
-    return asyncio.run(self._handler.load(path, abstract_checkpointable))
+    return asyncio_utils.run_sync(
+        self._handler.load(path, abstract_checkpointable)
+    )
 
   def metadata(self, path: Path) -> AbstractT:
-    return asyncio.run(self._handler.metadata(path))
+    return asyncio_utils.run_sync(self._handler.metadata(path))
 
   def is_handleable(self, checkpointable: Any) -> bool:
     return self._handler.is_handleable(checkpointable)