From 550abb6763f6d16f9ce4a77a443d7bc00da87b3d Mon Sep 17 00:00:00 2001 From: stefanwalcz Date: Mon, 20 Jul 2026 18:26:23 +0200 Subject: [PATCH 1/2] fix(backend/python): don't await sync servicer behaviors in AsyncModelIdentityInterceptor The model-identity interceptor (added for #10952) is installed on every Python backend's gRPC server. Its grpc.aio variant invokes the wrapped servicer behavior itself and awaits the result unconditionally: result = await original(request, context) # LoadModel return await original_unary(request, context) # guarded RPCs async for response in original_stream(request, context): # streaming But a backend's servicer methods may be plain sync functions. The transformers backend, for one, defines `def LoadModel` and `def Embedding` (not `async def`). grpc.aio's own dispatch adapts both shapes, but this interceptor calls the behavior directly and bypasses that. For a sync method `original(...)` returns a message object, not a coroutine, so the `await` raises: TypeError: object Result can't be used in 'await' expression The model loads, then the LoadModel RPC dies on return; the guarded sync Embedding fails the same way. It happens on every platform, not just one backend build. CI never caught it because AsyncModelIdentityInterceptor had no behavioral test -- only an "is it installed" assertion. Fix: await only when the behavior actually returned an awaitable (inspect.isawaitable), mirroring grpc.aio's own sync/async adaptation. The streaming guard iterates a sync generator with `for` and an async one with `async for`. Adds async-path coverage to model_identity_test.py exercising both sync and async LoadModel / guarded-unary / streaming behaviors. The sync cases fail on the current code with the TypeError above and pass with this fix. Signed-off-by: stefanwalcz --- backend/python/common/model_identity.py | 26 ++++- backend/python/common/model_identity_test.py | 104 +++++++++++++++++++ 2 files changed, 126 insertions(+), 4 deletions(-) diff --git a/backend/python/common/model_identity.py b/backend/python/common/model_identity.py index 256b73bb871e..99d69a31c66a 100644 --- a/backend/python/common/model_identity.py +++ b/backend/python/common/model_identity.py @@ -20,6 +20,7 @@ the model itself. """ +import inspect import threading import grpc @@ -200,7 +201,14 @@ async def intercept_service(self, continuation, handler_call_details): original = handler.unary_unary async def record(request, context): - result = await original(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 can't be used in 'await'". + result = original(request, context) + if inspect.isawaitable(result): + result = await result if getattr(result, "success", True): self.state.record(getattr(request, "Model", "")) return result @@ -214,8 +222,15 @@ async def guard_stream(request, context): message = self.state.mismatch(getattr(request, "ModelIdentity", "")) if message is not None: await context.abort(grpc.StatusCode.NOT_FOUND, message) - async for response in original_stream(request, context): - yield response + # A sync backend yields a plain generator, an async one an async + # generator; iterate whichever this is. + stream = original_stream(request, context) + if hasattr(stream, "__aiter__"): + async for response in stream: + yield response + else: + for response in stream: + yield response return _rebuild(handler, guard_stream) @@ -225,6 +240,9 @@ 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) - return await original_unary(request, context) + result = original_unary(request, context) + if inspect.isawaitable(result): + result = await result + return result return _rebuild(handler, guard) diff --git a/backend/python/common/model_identity_test.py b/backend/python/common/model_identity_test.py index 7d3019007eb5..23cda739d084 100644 --- a/backend/python/common/model_identity_test.py +++ b/backend/python/common/model_identity_test.py @@ -9,6 +9,7 @@ enforcement that rejects requests it should serve. """ +import asyncio import os import unittest @@ -302,6 +303,109 @@ 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. + + AsyncModelIdentityInterceptor wraps a backend's own servicer behavior. That + behavior may be sync or async: several backends define `def LoadModel` and + `def Embedding` (not `async def`), and grpc.aio's dispatch adapts both. The + interceptor invokes the behavior itself, so if it awaits unconditionally it + breaks every sync method it guards with "object can't be used in + 'await'". These tests exercise both shapes; the sync ones are the regression. + """ + + def setUp(self): + self.interceptor = model_identity.AsyncModelIdentityInterceptor() + + def _wrap(self, method, handler): + async def continuation(_): + return handler + + return asyncio.run( + self.interceptor.intercept_service(continuation, _FakeCallDetails(method)) + ) + + def _load(self, behavior, model): + wrapped = self._wrap("/backend.Backend/LoadModel", _handler(behavior)) + return asyncio.run(wrapped.unary_unary(_Request(Model=model), _FakeContext())) + + def _call_unary(self, behavior, identity): + wrapped = self._wrap("/backend.Backend/Predict", _handler(behavior)) + return asyncio.run( + wrapped.unary_unary(_Request(ModelIdentity=identity), _FakeContext()) + ) + + def _drain_stream(self, behavior, identity): + wrapped = self._wrap( + "/backend.Backend/PredictStream", _handler(behavior, response_streaming=True) + ) + + async def drain(): + out = [] + async for item in wrapped.unary_stream( + _Request(ModelIdentity=identity), _FakeContext() + ): + out.append(item) + return out + + return asyncio.run(drain()) + + # --- LoadModel: sync behavior is the regression, async must still work --- + + def test_load_records_with_sync_behavior(self): + self._load(lambda request, context: _Result(), "a.gguf") + self.assertEqual(self.interceptor.state.loaded, "a.gguf") + + def test_load_records_with_async_behavior(self): + async def behavior(request, context): + return _Result() + + self._load(behavior, "a.gguf") + self.assertEqual(self.interceptor.state.loaded, "a.gguf") + + def test_failed_sync_load_records_nothing(self): + self._load(lambda request, context: _Result(success=False), "a.gguf") + self.assertEqual(self.interceptor.state.loaded, "") + + # --- guarded unary: sync and async behaviors both served / rejected --- + + def test_guard_serves_sync_behavior(self): + self._load(lambda request, context: _Result(), "a.gguf") + result = self._call_unary(lambda request, context: "served", "a.gguf") + self.assertEqual(result, "served") + + def test_guard_serves_async_behavior(self): + self._load(lambda request, context: _Result(), "a.gguf") + + async def behavior(request, context): + return "served" + + self.assertEqual(self._call_unary(behavior, "a.gguf"), "served") + + def test_guard_rejects_mismatch(self): + self._load(lambda request, context: _Result(), "a.gguf") + with self.assertRaises(_Aborted): + self._call_unary(lambda request, context: "served", "b.gguf") + + # --- guarded stream: sync generator and async generator both work --- + + def test_guard_stream_serves_sync_generator(self): + self._load(lambda request, context: _Result(), "a.gguf") + + def behavior(request, context): + yield "a" + yield "b" + + self.assertEqual(self._drain_stream(behavior, "a.gguf"), ["a", "b"]) + + def test_guard_stream_serves_async_generator(self): + self._load(lambda request, context: _Result(), "a.gguf") + + async def behavior(request, context): + yield "a" + yield "b" + + self.assertEqual(self._drain_stream(behavior, "a.gguf"), ["a", "b"]) if __name__ == "__main__": From 7754116f406ca131f0497b87d3950fcb11a9a441 Mon Sep 17 00:00:00 2001 From: stefanwalcz Date: Mon, 20 Jul 2026 22:31:02 +0200 Subject: [PATCH 2/2] 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 --- backend/python/common/model_identity.py | 63 ++++++++++++---- backend/python/common/model_identity_test.py | 76 ++++++++++++++++++++ 2 files changed, 125 insertions(+), 14 deletions(-) diff --git a/backend/python/common/model_identity.py b/backend/python/common/model_identity.py index 99d69a31c66a..4e2ccb0218e5 100644 --- a/backend/python/common/model_identity.py +++ b/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 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 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) diff --git a/backend/python/common/model_identity_test.py b/backend/python/common/model_identity_test.py index 23cda739d084..24eb20028578 100644 --- a/backend/python/common/model_identity_test.py +++ b/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()