From 6d1bbb74c4a2e9e7848ff4df2d7728458c7b5459 Mon Sep 17 00:00:00 2001 From: walcz-de Date: Wed, 22 Jul 2026 16:03:46 +0200 Subject: [PATCH] fix(backend/python): don't await sync servicer behaviors in AsyncModelIdentityInterceptor (#10980) * 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 * 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 --------- Signed-off-by: stefanwalcz --- backend/python/common/model_identity.py | 61 ++++++- backend/python/common/model_identity_test.py | 180 +++++++++++++++++++ 2 files changed, 237 insertions(+), 4 deletions(-) diff --git a/backend/python/common/model_identity.py b/backend/python/common/model_identity.py index 256b73bb8..4e2ccb021 100644 --- a/backend/python/common/model_identity.py +++ b/backend/python/common/model_identity.py @@ -20,6 +20,8 @@ Enforcement is deliberately narrow: it compares two strings and never inspects the model itself. """ +import asyncio +import inspect import threading import grpc @@ -181,6 +183,40 @@ class ModelIdentityInterceptor(grpc.ServerInterceptor): 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.""" @@ -200,7 +236,11 @@ class AsyncModelIdentityInterceptor(grpc.aio.ServerInterceptor): 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`). 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 @@ -214,8 +254,21 @@ class AsyncModelIdentityInterceptor(grpc.aio.ServerInterceptor): 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. 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: + 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) @@ -225,6 +278,6 @@ class AsyncModelIdentityInterceptor(grpc.aio.ServerInterceptor): 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) + 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 7d3019007..24eb20028 100644 --- a/backend/python/common/model_identity_test.py +++ b/backend/python/common/model_identity_test.py @@ -9,7 +9,9 @@ failure modes are silent: enforcement that is wired up but never installed, and enforcement that rejects requests it should serve. """ +import asyncio import os +import threading import unittest import grpc @@ -64,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. @@ -304,5 +313,176 @@ class TestModalityMethods(unittest.TestCase): 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"]) + + # --- 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()