Internal Change.

checkpoint/orbax/checkpoint/_src/asyncio_utils.py

@@ -33,6 +33,9 @@
 
 _T = TypeVar('_T')
 
+# Marks "this attribute did not exist", which `None` cannot express.
+_UNSET = object()
+
 
 async def cancellable(
     coro: Any,
@@ -118,19 +121,28 @@ def run_sync(coro: Coroutine[Any, Any, _T]) -> _T:
 
   async def _coro_with_registration():
     current_thread = threading.current_thread()
+    # `loop` and `main_task` are plain attributes on a shared thread object,
+    # not names Orbax owns. anyio's worker threads keep their own `loop` there
+    # and read it back after the callable returns, so deleting ours would
+    # destroy theirs: the worker then dies before resolving its future and the
+    # awaiting coroutine hangs forever. Put back whatever was there before.
+    previous = {
+        name: getattr(current_thread, name, _UNSET)
+        for name in ('loop', 'main_task')
+    }
     current_thread.loop = asyncio.get_running_loop()  # pyrefly: ignore[missing-attribute]
     current_thread.main_task = asyncio.current_task()  # pyrefly: ignore[missing-attribute]
     try:
       return await coro
     finally:
-      try:
-        delattr(current_thread, 'loop')
-      except AttributeError:
-        pass
-      try:
-        delattr(current_thread, 'main_task')
-      except AttributeError:
-        pass
+      for name, value in previous.items():
+        if value is _UNSET:
+          try:
+            delattr(current_thread, name)
+          except AttributeError:
+            pass
+        else:
+          setattr(current_thread, name, value)
 
   if loop is None:
     # No event loop is running. Use a fresh loop without asyncio.run() to

checkpoint/orbax/checkpoint/_src/asyncio_utils_test.py

@@ -23,6 +23,8 @@
 from absl import logging
 from absl.testing import absltest
 from absl.testing import parameterized
+import anyio
+import anyio.to_thread
 from orbax.checkpoint._src import asyncio_utils
 
 
@@ -339,6 +341,55 @@ async def _test():
     logging.info("time: run_sync_time=%s", run_sync_time)
 
 
+class ThreadAnnotationTest(absltest.TestCase):
+  """`run_sync` must leave the calling thread as it found it.
+
+  `run_sync` annotates `threading.current_thread()` with `loop` and
+  `main_task`. Those attribute names are not private to Orbax: anyio's worker
+  threads keep their own `loop` there and read it back after the callable
+  returns. Clearing ours used to clear theirs, killing the worker before it
+  resolved its future and hanging the awaiting coroutine forever.
+  """
+
+  def test_absent_attributes_stay_absent(self):
+    thread = threading.current_thread()
+    for name in ("loop", "main_task"):
+      if hasattr(thread, name):
+        delattr(thread, name)
+
+    asyncio_utils.run_sync(one())
+
+    self.assertFalse(hasattr(thread, "loop"))
+    self.assertFalse(hasattr(thread, "main_task"))
+
+  def test_pre_existing_attributes_are_restored(self):
+    thread = threading.current_thread()
+    sentinel_loop = object()
+    sentinel_task = object()
+    setattr(thread, "loop", sentinel_loop)
+    setattr(thread, "main_task", sentinel_task)
+    self.addCleanup(lambda: delattr(thread, "loop"))
+    self.addCleanup(lambda: delattr(thread, "main_task"))
+
+    asyncio_utils.run_sync(one())
+
+    self.assertIs(getattr(thread, "loop"), sentinel_loop)
+    self.assertIs(getattr(thread, "main_task"), sentinel_task)
+
+  def test_runs_inside_an_anyio_worker_thread(self):
+    """End-to-end guard: this used to hang forever rather than fail."""
+
+    def _blocking() -> int:
+      return asyncio_utils.run_sync(one())
+
+    async def _main() -> int:
+      with anyio.move_on_after(30):
+        return await anyio.to_thread.run_sync(_blocking)
+      raise AssertionError("anyio.to_thread.run_sync did not return")
+
+    self.assertEqual(asyncio.run(_main()), 1)
+
+
 class AsyncRunnerTest(absltest.TestCase):
 
   _TIMEOUT = 2