Add compatibility for missing uvloop in asyncio_utils

checkpoint/CHANGELOG.md

@@ -21,6 +21,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Enforce the array shape and type check during Array restoration when
   `ArrayRestoreArgs.strict` is set but shape/dtype is not provided.
+- On platforms where `uvloop` is not supported, fallback to `nest_asyncio`.
 
 ## [0.11.33] - 2026-02-17
 

checkpoint/orbax/checkpoint/_src/asyncio_utils.py

@@ -18,8 +18,15 @@
 import threading
 from typing import Any, Coroutine, TypeVar
 
-import uvloop
+try:
+  import uvloop  # pylint: disable=g-import-not-at-top
+except ImportError:
+  uvloop = None
 
+try:
+  import nest_asyncio  # pylint: disable=g-import-not-at-top # pytype: disable=import-error
+except ImportError:
+  nest_asyncio = None
 
 _T = TypeVar('_T')
 
@@ -43,13 +50,22 @@ def run_sync(coro: Coroutine[Any, Any, _T]) -> _T:
     return asyncio.run(coro)
   else:
     # An event loop is already running.
-    event_loop = uvloop.new_event_loop()
-    thread = threading.Thread(
-        target=_run_event_loop, args=(event_loop,), daemon=True
-    )
-    thread.start()
-    try:
-      return asyncio.run_coroutine_threadsafe(coro, event_loop).result()
-    finally:
-      event_loop.call_soon_threadsafe(event_loop.stop)
-      thread.join()
+    if uvloop is None:
+      if nest_asyncio is None:
+        raise RuntimeError(
+            'nest_asyncio is not installed. Please install it to use run_sync'
+            ' with an existing event loop.'
+        )
+      nest_asyncio.apply()
+      return asyncio.run(coro)
+    else:
+      event_loop = uvloop.new_event_loop()
+      thread = threading.Thread(
+          target=_run_event_loop, args=(event_loop,), daemon=True
+      )
+      thread.start()
+      try:
+        return asyncio.run_coroutine_threadsafe(coro, event_loop).result()
+      finally:
+        event_loop.call_soon_threadsafe(event_loop.stop)
+        thread.join()

checkpoint/pyproject.toml

@@ -35,7 +35,8 @@ dependencies = [
     'humanize',
     'simplejson >= 3.16.0',
     'psutil',
-    'uvloop',
+    'uvloop; platform_system != "Windows"',
+    'nest_asyncio; platform_system == "Windows"',
 ]
 
 dynamic = ['version']