No public description

checkpoint/orbax/checkpoint/_src/asyncio_utils.py

@@ -40,6 +40,30 @@ def _run_event_loop(loop: asyncio.AbstractEventLoop) -> None:
   loop.close()
 
 
+def _run_sync_via_new_loop(coro: Coroutine[Any, Any, _T]) -> _T:
+  """Runs a coroutine in a fresh event loop without asyncio.run().
+
+  Args:
+    coro: The coroutine object to run.
+
+  Returns:
+    The result of the coroutine.
+
+  asyncio.run() tears down the event loop via _cancel_all_tasks(), which
+  iterates asyncio's global WeakSet of tasks. This WeakSet is not thread-safe:
+  concurrent asyncio.run() calls from multiple threads can race on iteration
+  vs. mutation, raising "RuntimeError: Set changed size during iteration".
+
+  Using run_until_complete() + close() achieves the same functional result
+  without the dangerous teardown.
+  """
+  loop = asyncio.new_event_loop()
+  try:
+    return loop.run_until_complete(coro)
+  finally:
+    loop.close()
+
+
 def run_sync(coro: Coroutine[Any, Any, _T]) -> _T:
   """Runs a coroutine and returns the result."""
   try:
@@ -49,8 +73,9 @@ def run_sync(coro: Coroutine[Any, Any, _T]) -> _T:
     loop = None
 
   if loop is None:
-    # No event loop is running, so we can safely use asyncio.run.
-    return asyncio.run(coro)
+    # No event loop is running. Use a fresh loop without asyncio.run() to
+    # avoid thread-unsafe teardown in _cancel_all_tasks().
+    return _run_sync_via_new_loop(coro)
   else:
     # An event loop is already running.
     if uvloop is None:
@@ -60,7 +85,7 @@ def run_sync(coro: Coroutine[Any, Any, _T]) -> _T:
             ' with an existing event loop.'
         )
       nest_asyncio.apply()
-      return asyncio.run(coro)
+      return _run_sync_via_new_loop(coro)
     else:
       event_loop = uvloop.new_event_loop()
       thread = threading.Thread(