Buckets:
| """Model-identity enforcement for LocalAI Python backends. | |
| In distributed mode the controller caches a routing row naming a backend's | |
| host:port. A worker can recycle a stopped backend's gRPC port for a different | |
| model's backend, and the controller's health probe checks liveness rather than | |
| identity, so the request is dispatched to whatever now occupies the port and | |
| the caller gets a silent wrong-model answer (#10952). | |
| Every request message that reaches a backend through the distributed router | |
| carries a ModelIdentity field naming the model the request is for, so the | |
| backend can reject it at the point of use. This module enforces that for every | |
| Python backend at once: all 36 of them build their server through | |
| grpc_auth.get_auth_interceptors(), so wiring it there needs no per-backend | |
| change. There is no shared BackendServicer base class to hook instead, and the | |
| backends store their loaded model in wildly different attributes, so an | |
| interceptor is the only single point that sees both the LoadModel request and | |
| the inference requests. | |
| Enforcement is deliberately narrow: it compares two strings and never inspects | |
| the model itself. | |
| """ | |
| import asyncio | |
| import inspect | |
| import threading | |
| import grpc | |
| # Must match grpcerrors.ModelMismatchSentinel in pkg/grpc/grpcerrors/errors.go. | |
| # The router requires this substring AND the NOT_FOUND code before it treats a | |
| # reply as a mismatch, because NOT_FOUND alone is not exclusively ours on these | |
| # RPCs (insightface's Embedding returns it for "no face detected"). | |
| MODEL_MISMATCH_SENTINEL = "model identity mismatch" | |
| _LOAD_METHOD = "/backend.Backend/LoadModel" | |
| # Every RPC whose request message carries a ModelIdentity field. This set IS | |
| # the enforcement surface for all 36 Python backends: an RPC missing here is | |
| # silently unprotected, so model_identity_test.py pins the full list. | |
| # | |
| # The guard reads request.ModelIdentity generically, so nothing here is | |
| # modality-specific — a backend that does not implement an RPC simply never | |
| # sees it. | |
| # | |
| # TTS and SoundGeneration are guarded on ModelIdentity, NOT on their `model` | |
| # field: the controller's FileStagingClient rewrites `model` to a worker-local | |
| # absolute path, so comparing that would reject valid requests in distributed | |
| # mode. ModelIdentity is a separate, untranslated field for exactly that reason. | |
| # | |
| # AudioEncode/AudioDecode are absent deliberately: the opus codec backend they | |
| # target is loaded from a literal rather than a ModelConfig, so no value carries | |
| # the load-time/request-time equality guarantee this comparison depends on. | |
| _GUARDED_METHODS = frozenset( | |
| ( | |
| # PredictOptions RPCs (#10970) | |
| "/backend.Backend/Predict", | |
| "/backend.Backend/PredictStream", | |
| "/backend.Backend/Embedding", | |
| "/backend.Backend/TokenizeString", | |
| # Remaining modalities | |
| "/backend.Backend/GenerateImage", | |
| "/backend.Backend/GenerateVideo", | |
| "/backend.Backend/TTS", | |
| "/backend.Backend/TTSStream", | |
| "/backend.Backend/SoundGeneration", | |
| "/backend.Backend/AudioTranscription", | |
| "/backend.Backend/AudioTranscriptionStream", | |
| "/backend.Backend/Detect", | |
| "/backend.Backend/Depth", | |
| "/backend.Backend/FaceVerify", | |
| "/backend.Backend/FaceAnalyze", | |
| "/backend.Backend/VoiceVerify", | |
| "/backend.Backend/VoiceAnalyze", | |
| "/backend.Backend/VoiceEmbed", | |
| "/backend.Backend/Rerank", | |
| "/backend.Backend/TokenClassify", | |
| "/backend.Backend/Score", | |
| "/backend.Backend/VAD", | |
| "/backend.Backend/Diarize", | |
| "/backend.Backend/SoundDetection", | |
| "/backend.Backend/AudioTransform", | |
| ) | |
| ) | |
| class ModelIdentityState: | |
| """The identity this process loaded, and the rule for judging a request. | |
| A backend process serves exactly one model (worker process keys are | |
| model+backend+replica), so a single value is enough. | |
| """ | |
| def __init__(self): | |
| self._lock = threading.Lock() | |
| self._loaded = "" | |
| def record(self, model: str) -> None: | |
| with self._lock: | |
| self._loaded = model or "" | |
| def loaded(self) -> str: | |
| with self._lock: | |
| return self._loaded | |
| def mismatch(self, requested: str): | |
| """Return an error message when `requested` names another model. | |
| Either side being empty means "skip": the request side is empty for a | |
| controller that predates the field and for internally synthesized | |
| requests, and the loaded side is empty when such a controller performed | |
| the load. Neither can judge the other, and a false rejection is worse | |
| than the miss it prevents. | |
| """ | |
| if not requested: | |
| return None | |
| loaded = self.loaded | |
| if not loaded or loaded == requested: | |
| return None | |
| return "{}: loaded {!r}, requested {!r}".format( | |
| MODEL_MISMATCH_SENTINEL, loaded, requested | |
| ) | |
| def _rebuild(handler, behavior): | |
| """Return a copy of `handler` with its behavior replaced. | |
| Only unary-request handlers are ever passed here: LoadModel and every | |
| entry in _GUARDED_METHODS take a single request message. The bidirectional | |
| streams (AudioTranscriptionLive, AudioTransformStream, AudioToAudioStream, | |
| Forward) are not guarded and never reach this function. | |
| """ | |
| if handler.response_streaming: | |
| return grpc.unary_stream_rpc_method_handler( | |
| behavior, | |
| request_deserializer=handler.request_deserializer, | |
| response_serializer=handler.response_serializer, | |
| ) | |
| return grpc.unary_unary_rpc_method_handler( | |
| behavior, | |
| request_deserializer=handler.request_deserializer, | |
| response_serializer=handler.response_serializer, | |
| ) | |
| class ModelIdentityInterceptor(grpc.ServerInterceptor): | |
| """Sync interceptor that records the loaded model and guards inference.""" | |
| def __init__(self, state: ModelIdentityState = None): | |
| self.state = state or ModelIdentityState() | |
| def intercept_service(self, continuation, handler_call_details): | |
| method = handler_call_details.method | |
| if method != _LOAD_METHOD and method not in _GUARDED_METHODS: | |
| return continuation(handler_call_details) | |
| handler = continuation(handler_call_details) | |
| if handler is None: | |
| return handler | |
| if method == _LOAD_METHOD: | |
| original = handler.unary_unary | |
| def record(request, context): | |
| result = original(request, context) | |
| # Only a successful load owns the identity; a failed one leaves | |
| # no model, which the model-not-loaded signal already covers. | |
| if getattr(result, "success", True): | |
| self.state.record(getattr(request, "Model", "")) | |
| return result | |
| return _rebuild(handler, record) | |
| original = handler.unary_stream if handler.response_streaming else handler.unary_unary | |
| def guard(request, context): | |
| message = self.state.mismatch(getattr(request, "ModelIdentity", "")) | |
| if message is not None: | |
| # abort() raises, so the request never reaches the model. | |
| context.abort(grpc.StatusCode.NOT_FOUND, message) | |
| return original(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.""" | |
| def __init__(self, state: ModelIdentityState = None): | |
| self.state = state or ModelIdentityState() | |
| async def intercept_service(self, continuation, handler_call_details): | |
| method = handler_call_details.method | |
| if method != _LOAD_METHOD and method not in _GUARDED_METHODS: | |
| return await continuation(handler_call_details) | |
| handler = await continuation(handler_call_details) | |
| if handler is None: | |
| return handler | |
| if method == _LOAD_METHOD: | |
| original = handler.unary_unary | |
| async def record(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 | |
| return _rebuild(handler, record) | |
| if handler.response_streaming: | |
| original_stream = handler.unary_stream | |
| 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) | |
| # 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) | |
| original_unary = handler.unary_unary | |
| 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 _call_behavior(original_unary, request, context) | |
| return _rebuild(handler, guard) | |
Xet Storage Details
- Size:
- 11.6 kB
- Xet hash:
- 20c8de8a43f03a2d283de27d8a9e5367d0ec6f1647e56e860608eb41f147d01d
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.