mirror of
https://github.com/mudler/LocalAI.git
synced 2026-07-30 09:57:57 -04:00
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 <stefan.walcz@walcz.de> * 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> --------- Signed-off-by: stefanwalcz <stefan.walcz@walcz.de>
This commit is contained in:
walcz-de
@@ -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 <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
|
||||
@@ -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)
|
||||
|
||||
@@ -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 <T> 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()
|
||||
|
||||
Reference in New Issue
Block a user