From edc02cb895520a4041549b5dcb19d86b64cec582 Mon Sep 17 00:00:00 2001
From: Orbax Authors <no-reply@google.com>
Date: Mon, 13 Jul 2026 12:35:15 -0700
Subject: [PATCH] No public description

PiperOrigin-RevId: 947185453
---
 .../orbax/checkpoint/_src/asyncio_utils.py    | 31 +++++++++++++++++--
 1 file changed, 28 insertions(+), 3 deletions(-)

diff --git a/checkpoint/orbax/checkpoint/_src/asyncio_utils.py b/checkpoint/orbax/checkpoint/_src/asyncio_utils.py
index 24caad861..64f8849b5 100644
--- a/checkpoint/orbax/checkpoint/_src/asyncio_utils.py
+++ b/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(