model_identity.py

"""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 ""

    @property
    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)