fix: keep worker-local sync calls off the event loop

Author: KAJdev

runpod/apps/handles.py

@@ -18,6 +18,7 @@ async def transcribe(audio_url: str): ...
     def load_model(): ...                    # worker startup hook
 """
 
+import asyncio
 import inspect
 from typing import (
     TYPE_CHECKING,
@@ -30,7 +31,7 @@ def load_model(): ...                    # worker startup hook
     Type,
 )
 
-from .context import Context, current_context
+from .context import Context, block, current_context
 from .discovery_state import DiscoveryInvocationError, in_discovery
 from .errors import InvalidResourceError
 from .invoker import Invoker, StreamInvoker
@@ -64,7 +65,7 @@ def __init__(
         if stamped and not spec.schedule:
             spec.schedule = stamped
 
-        self.remote = Invoker(self._remote_async)
+        self.remote = Invoker(self._remote_async, sync_factory=self._remote_sync)
         self.stream = StreamInvoker(self._stream_async)
         self.spawn = Invoker(self._spawn_async)
         self.job = Invoker(self._job_async)
@@ -117,22 +118,46 @@ def local(self, *args: Any, **kwargs: Any) -> Any:
         function is async."""
         return self._fn(*args, **kwargs)
 
+    def _call_local_sync(self, *args: Any, **kwargs: Any) -> Any:
+        result = self._fn(*args, **kwargs)
+        return list(result) if inspect.isgenerator(result) else result
+
+    @staticmethod
+    async def _await_local_result(result: Any) -> Any:
+        if inspect.isawaitable(result):
+            result = await result
+        if inspect.isasyncgen(result):
+            return [chunk async for chunk in result]
+        return list(result) if inspect.isgenerator(result) else result
+
+    def _remote_sync(self, *args: Any, **kwargs: Any) -> Any:
+        if (
+            current_context() is Context.WORKER
+            and self._is_current_worker()
+            and not inspect.iscoroutinefunction(self._fn)
+            and not inspect.isasyncgenfunction(self._fn)
+        ):
+            self._guard_discovery()
+            self.spec.validate()
+            result = self._call_local_sync(*args, **kwargs)
+            if inspect.isawaitable(result) or inspect.isasyncgen(result):
+                return block(self._await_local_result(result))
+            return result
+        return block(self._remote_async(*args, **kwargs))
+
     async def _remote_async(self, *args: Any, **kwargs: Any) -> Any:
         self._guard_discovery()
         self.spec.validate()
         ctx = current_context()
 
         if ctx is Context.WORKER and self._is_current_worker():
-            result = self._fn(*args, **kwargs)
-            if inspect.isawaitable(result):
-                result = await result
-            # generators aggregate, matching the deployed worker's
-            # return_aggregate_stream output for .remote()
-            if inspect.isasyncgen(result):
-                return [chunk async for chunk in result]
-            if inspect.isgenerator(result):
-                return list(result)
-            return result
+            if inspect.iscoroutinefunction(self._fn) or inspect.isasyncgenfunction(
+                self._fn
+            ):
+                result = self._fn(*args, **kwargs)
+            else:
+                result = await asyncio.to_thread(self._call_local_sync, *args, **kwargs)
+            return await self._await_local_result(result)
 
         target = await self._app._resolve(self.spec)
         payload = target.build_payload(self._fn, self.spec, args, kwargs)

runpod/apps/invoker.py

@@ -10,7 +10,7 @@
 returns a sync iterator, `.aio(...)` returns the async iterator.
 """
 
-from typing import Any, AsyncIterator, Callable, Coroutine, Iterator
+from typing import Any, AsyncIterator, Callable, Coroutine, Iterator, Optional
 
 from .context import block
 
@@ -19,12 +19,17 @@
 
 
 class Invoker:
-    __slots__ = ("_factory",)
+    __slots__ = ("_factory", "_sync_factory")
 
-    def __init__(self, factory: CoroFactory):
+    def __init__(
+        self, factory: CoroFactory, *, sync_factory: Optional[Callable] = None
+    ):
         self._factory = factory
+        self._sync_factory = sync_factory
 
     def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        if self._sync_factory is not None:
+            return self._sync_factory(*args, **kwargs)
         return block(self._factory(*args, **kwargs))
 
     def aio(self, *args: Any, **kwargs: Any) -> Coroutine[Any, Any, Any]:

tests/test_apps/test_dispatch.py

@@ -1,5 +1,8 @@
 """tests for context detection and remote dispatch."""
 
+import os
+import subprocess
+import sys
 from unittest.mock import AsyncMock, patch
 
 import pytest
@@ -210,6 +213,39 @@ def q(x):
 
         assert q.remote(4) == 40
 
+    def test_worker_local_recursion_from_sync_and_async_callers(self):
+        script = """
+import asyncio
+import runpod
+
+app = runpod.App("self-smoke")
+
+@app.queue(cpu="cpu3c-2-4")
+def countdown(n):
+    return 0 if n == 0 else 1 + countdown.remote(n - 1)
+
+assert countdown.remote(3) == 3
+
+async def main():
+    results = await asyncio.gather(*(countdown.remote.aio(4) for _ in range(40)))
+    assert results == [4] * 40
+
+asyncio.run(main())
+"""
+        subprocess.run(
+            [sys.executable, "-c", script],
+            env={
+                **os.environ,
+                "RUNPOD_POD_ID": "local-worker",
+                "RUNPOD_DEV_APP": "self-smoke",
+                "RUNPOD_DEV_RESOURCE": "countdown",
+            },
+            check=True,
+            capture_output=True,
+            text=True,
+            timeout=10,
+        )
+
     def test_worker_calls_sibling_via_sentinel(self, monkeypatch):
         monkeypatch.setenv("RUNPOD_ENDPOINT_ID", "ep-1")
         monkeypatch.setenv("FLASH_RESOURCE_NAME", "other")
@@ -409,7 +445,6 @@ async def caller():
 class TestModuleSourceShipping:
     def test_whole_module_ships(self, tmp_path):
         import importlib.util
-        import sys
 
         mod_file = tmp_path / "shipmod.py"
         mod_file.write_text(