fix(backend/python): dispatch sync servicer behaviors off the event loop

Addresses review feedback: awaiting only awaitable results removed the
TypeError, but still ran a sync LoadModel/Embedding -- and stepped a sync stream
via next() -- on the asyncio event-loop thread, so a slow load/inference/stream
could freeze all aio RPC handling.

Route sync behavior through run_in_executor (a worker thread) while awaiting
native async behavior directly. A callable wrapper that returns an awaitable is
run in the thread and its awaitable awaited back on the loop. Sync streaming
pulls each item via the executor with a done sentinel, so StopIteration cannot
escape through a Future.

Adds regression tests that record the handler thread id and assert it differs
from the event-loop thread, for LoadModel, a guarded unary RPC and a sync stream.

Signed-off-by: stefanwalcz <stefan.walcz@walcz.de>

Author: walcz-de

backend/python/common/model_identity.py

@@ -20,6 +20,7 @@
 the model itself.
 """
 
+import asyncio
 import inspect
 import threading
 
@@ -182,6 +183,40 @@ def guard(request, context):
         return _rebuild(handler, guard)
 
 
+_STREAM_DONE = object()
+
+
+def _next_or_done(iterator):
+    """next(iterator), returning the _STREAM_DONE sentinel at exhaustion.
+
+    StopIteration must not propagate out of a function run via run_in_executor:
+    it cannot travel through a Future and would surface as an opaque error.
+    """
+    try:
+        return next(iterator)
+    except StopIteration:
+        return _STREAM_DONE
+
+
+async def _call_behavior(behavior, request, context):
+    """Invoke a unary servicer behavior without blocking the event loop.
+
+    Native async behavior is awaited directly. A sync behavior -- many backends
+    define `def LoadModel` / `def Embedding`, not `async def` -- is dispatched to
+    a worker thread so a slow load/inference cannot freeze all aio RPC handling,
+    mirroring grpc.aio's own sync-handler adaptation. A callable wrapper that
+    returns an awaitable is supported too: the (cheap) call runs in the thread,
+    then the awaitable is awaited back on the loop.
+    """
+    if inspect.iscoroutinefunction(behavior):
+        return await behavior(request, context)
+    loop = asyncio.get_running_loop()
+    result = await loop.run_in_executor(None, behavior, request, context)
+    if inspect.isawaitable(result):
+        result = await result
+    return result
+
+
 class AsyncModelIdentityInterceptor(grpc.aio.ServerInterceptor):
     """Async counterpart for backends running grpc.aio servers."""
 
@@ -202,13 +237,10 @@ async def intercept_service(self, continuation, handler_call_details):
 
             async def record(request, context):
                 # A backend's LoadModel may be a plain sync method (many define
-                # `def LoadModel`, not `async def`). grpc.aio's own dispatch
-                # adapts both, but this interceptor calls the behavior directly,
-                # so it must not await a non-awaitable return -- otherwise a sync
-                # backend fails with "object <T> can't be used in 'await'".
-                result = original(request, context)
-                if inspect.isawaitable(result):
-                    result = await result
+                # `def LoadModel`, not `async def`). Dispatch it so it neither
+                # crashes with "object <T> can't be used in 'await'" nor runs its
+                # (potentially slow) body on the event loop thread.
+                result = await _call_behavior(original, request, context)
                 if getattr(result, "success", True):
                     self.state.record(getattr(request, "Model", ""))
                 return result
@@ -223,14 +255,20 @@ async def guard_stream(request, context):
                 if message is not None:
                     await context.abort(grpc.StatusCode.NOT_FOUND, message)
                 # A sync backend yields a plain generator, an async one an async
-                # generator; iterate whichever this is.
+                # generator. Async: iterate directly. Sync: pull each item via a
+                # worker thread so a slow producer doesn't block the event loop
+                # (and so StopIteration can't escape through a Future).
                 stream = original_stream(request, context)
                 if hasattr(stream, "__aiter__"):
                     async for response in stream:
                         yield response
                 else:
-                    for response in stream:
-                        yield response
+                    loop = asyncio.get_running_loop()
+                    while True:
+                        item = await loop.run_in_executor(None, _next_or_done, stream)
+                        if item is _STREAM_DONE:
+                            break
+                        yield item
 
             return _rebuild(handler, guard_stream)
 
@@ -240,9 +278,6 @@ async def guard(request, context):
             message = self.state.mismatch(getattr(request, "ModelIdentity", ""))
             if message is not None:
                 await context.abort(grpc.StatusCode.NOT_FOUND, message)
-            result = original_unary(request, context)
-            if inspect.isawaitable(result):
-                result = await result
-            return result
+            return await _call_behavior(original_unary, request, context)
 
         return _rebuild(handler, guard)

backend/python/common/model_identity_test.py

@@ -11,6 +11,7 @@
 
 import asyncio
 import os
+import threading
 import unittest
 
 import grpc
@@ -65,6 +66,13 @@ def _handler(behavior, response_streaming=False):
     return grpc.unary_unary_rpc_method_handler(behavior)
 
 
+def _const_continuation(handler):
+    async def continuation(_):
+        return handler
+
+    return continuation
+
+
 class TestInterceptorInstalled(unittest.TestCase):
     """The wiring, which is where this can silently do nothing.
 
@@ -303,6 +311,8 @@ def test_every_modality_rpc_serves_without_an_identity(self):
     def test_codec_rpcs_stay_unguarded(self):
         for method in ("/backend.Backend/AudioEncode", "/backend.Backend/AudioDecode"):
             self.assertNotIn(method, model_identity._GUARDED_METHODS)
+
+
 class TestAsyncInterceptorBehavior(unittest.TestCase):
     """The grpc.aio counterpart, which had no behavioral coverage.
 
@@ -407,6 +417,72 @@ async def behavior(request, context):
 
         self.assertEqual(self._drain_stream(behavior, "a.gguf"), ["a", "b"])
 
+    # --- sync behavior must not run on the event-loop thread ---
+    #
+    # Awaiting a sync method's return fixed the TypeError, but calling the
+    # (possibly slow) sync behavior on the event loop still froze all aio RPC
+    # handling. These record the thread each behavior runs on and assert it is a
+    # worker thread, not the loop thread.
+
+    def _run_capturing_loop_thread(self, method, handler, request):
+        captured = {}
+
+        async def run():
+            captured["loop"] = threading.get_ident()
+            wrapped = await self.interceptor.intercept_service(
+                _const_continuation(handler), _FakeCallDetails(method)
+            )
+            behavior = wrapped.unary_stream if handler.response_streaming else wrapped.unary_unary
+            if handler.response_streaming:
+                async for _ in behavior(request, _FakeContext()):
+                    pass
+            else:
+                await behavior(request, _FakeContext())
+
+        asyncio.run(run())
+        return captured["loop"]
+
+    def test_sync_load_runs_off_the_event_loop(self):
+        ran = {}
+
+        def behavior(request, context):
+            ran["thread"] = threading.get_ident()
+            return _Result()
+
+        loop_thread = self._run_capturing_loop_thread(
+            "/backend.Backend/LoadModel", _handler(behavior), _Request(Model="a.gguf")
+        )
+        self.assertIn("thread", ran)
+        self.assertNotEqual(ran["thread"], loop_thread)
+
+    def test_sync_guarded_unary_runs_off_the_event_loop(self):
+        self._load(lambda request, context: _Result(), "a.gguf")
+        ran = {}
+
+        def behavior(request, context):
+            ran["thread"] = threading.get_ident()
+            return "served"
+
+        loop_thread = self._run_capturing_loop_thread(
+            "/backend.Backend/Predict", _handler(behavior), _Request(ModelIdentity="a.gguf")
+        )
+        self.assertNotEqual(ran["thread"], loop_thread)
+
+    def test_sync_stream_next_runs_off_the_event_loop(self):
+        self._load(lambda request, context: _Result(), "a.gguf")
+        ran = {}
+
+        def behavior(request, context):
+            ran["thread"] = threading.get_ident()
+            yield "a"
+
+        loop_thread = self._run_capturing_loop_thread(
+            "/backend.Backend/PredictStream",
+            _handler(behavior, response_streaming=True),
+            _Request(ModelIdentity="a.gguf"),
+        )
+        self.assertNotEqual(ran["thread"], loop_thread)
+
 
 if __name__ == "__main__":
     unittest.main()