From a1e2921403eb6845720ccb2b0a078093b4ef1948 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Sun, 19 Jul 2026 23:55:33 +0000 Subject: [PATCH] fix(distributed): reject wrong-model requests at the backend (#10952) In distributed mode the controller caches a NodeModel row naming a backend's host:port. A worker can recycle a stopped backend's gRPC port for a different model's backend, and probeHealth verifies liveness rather than identity, so the probe succeeds against whatever now occupies the port and the request is dispatched to the wrong backend. The caller gets a silent wrong-model answer. Nothing in the request could catch this: PredictOptions had no model field, so model identity crossed the wire only in ModelOptions.Model at LoadModel time, and the cached-hit path issues no LoadModel. Every backend's "model not loaded" guard checks a nil handle, which a process holding a different model passes, so the stale row was never dropped either. Add PredictOptions.ModelIdentity and enforce it at the point of use: - The controller populates it in gRPCPredictOpts from ModelConfig.Model, the same expression ModelOptions feeds to model.WithModel and therefore the same value the backend received as ModelOptions.Model. Both are read from one config value in one function, so they are equal by construction and the comparison cannot false-reject. - Backends compare it against what they loaded and return NOT_FOUND with a fixed sentinel. Enforced in pkg/grpc/server.go (27 Go backends), an interceptor in backend/python/common (all 36 Python backends, no per-backend change), and the llama-cpp / ik-llama-cpp / ds4 C++ servers. That is every backend with real exposure: kokoros answers all four RPCs with unimplemented and privacy-filter implements none of them. - The router's reconcile drops the stale replica row on a mismatch, so the next request reloads somewhere correct. Empty means "skip the check" on both sides: a controller that predates the field sends nothing, a backend loaded by such a controller has nothing to compare, and the C++ server synthesizes PredictOptions internally for ASR. That keeps upgrades working in both directions. Scoped to the four PredictOptions RPCs. TTSRequest.model and SoundGenerationRequest.model are deliberately NOT validated: FileStagingClient already rewrites them to worker-local absolute paths, so in distributed mode they already differ from the load-time value and comparing them would reject valid requests. IsModelMismatch requires both the NOT_FOUND code and the sentinel, unlike the neighbouring helpers which accept either. insightface's Embedding returns NOT_FOUND "no face detected" on a PredictOptions RPC, and a code-only check would drop a healthy replica row on every faceless image. Signed-off-by: Ettore Di Giacinto Assisted-by: Claude Code:claude-opus-4-8[1m] [Read] [Edit] [Bash] --- backend/backend.proto | 26 ++ backend/cpp/ds4/grpc-server.cpp | 27 +++ backend/cpp/ik-llama-cpp/grpc-server.cpp | 35 +++ backend/cpp/llama-cpp/grpc-server.cpp | 35 +++ backend/python/common/grpc_auth.py | 23 +- backend/python/common/model_identity.py | 192 +++++++++++++++ backend/python/common/model_identity_test.py | 222 ++++++++++++++++++ core/backend/options.go | 7 + core/backend/options_internal_test.go | 41 ++++ core/services/nodes/inflight.go | 21 +- core/services/nodes/inflight_test.go | 40 ++++ .../services/nodes/model_identity_e2e_test.go | 91 +++++++ pkg/grpc/grpcerrors/errors.go | 37 +++ pkg/grpc/grpcerrors/errors_test.go | 33 +++ pkg/grpc/model_identity_test.go | 128 ++++++++++ pkg/grpc/server.go | 54 +++++ tests/e2e/mock-backend/main.go | 29 +++ 17 files changed, 1030 insertions(+), 11 deletions(-) create mode 100644 backend/python/common/model_identity.py create mode 100644 backend/python/common/model_identity_test.py create mode 100644 core/services/nodes/model_identity_e2e_test.go create mode 100644 pkg/grpc/model_identity_test.go diff --git a/backend/backend.proto b/backend/backend.proto index ad62c6df07ae..107e2a94e1bd 100644 --- a/backend/backend.proto +++ b/backend/backend.proto @@ -315,6 +315,32 @@ message PredictOptions { int32 TopLogprobs = 51; // Number of top logprobs to return per token (maps to OpenAI top_logprobs parameter) map Metadata = 52; // Generic per-request metadata (e.g., enable_thinking) float MinP = 53; // Minimum probability sampling threshold (0.0 = disabled) + + // ModelIdentity names the model this request is for, so a backend can reject + // a request that reached it by mistake instead of answering from whatever + // model it happens to hold. In distributed mode a worker can recycle a + // stopped backend's gRPC port for a different model's backend, and a + // liveness-only health probe cannot tell that apart from a valid cached + // route (#10952). + // + // The value is the controller's ModelConfig.Model, the SAME expression that + // produces ModelOptions.Model at LoadModel time, so the two are equal by + // construction rather than by convention. + // + // Empty means "no identity supplied": backends MUST skip the check. That + // keeps an old controller talking to a new backend working, and covers + // callers that legitimately synthesize a PredictOptions internally. + // + // Do NOT reuse TTSRequest.model or SoundGenerationRequest.model for this + // purpose. FileStagingClient already rewrites those to worker-local absolute + // paths (core/services/nodes/file_staging_client.go), so in distributed mode + // they already differ from the load-time value and comparing them would + // reject valid requests. Extending identity to those RPCs needs a separate + // field carrying the untranslated value. + string ModelIdentity = 54; + + // 24 was never assigned; reserve it so it is not silently reused. + reserved 24; } // ToolCallDelta represents an incremental tool call update from the C++ parser. diff --git a/backend/cpp/ds4/grpc-server.cpp b/backend/cpp/ds4/grpc-server.cpp index 3d324cd6a7cb..924da5476994 100644 --- a/backend/cpp/ds4/grpc-server.cpp +++ b/backend/cpp/ds4/grpc-server.cpp @@ -51,6 +51,11 @@ namespace { // Global state - ds4 is single-engine-per-process by design. std::mutex g_engine_mu; +// The ModelOptions.Model this process loaded, compared against +// PredictOptions.ModelIdentity so a request that arrived through a stale +// distributed route is rejected rather than answered from the wrong model +// (#10952). Guarded by g_engine_mu like the rest of the engine state. +std::string g_loaded_model_identity; ds4_engine *g_engine = nullptr; ds4_session *g_session = nullptr; int g_ctx_size = 32768; @@ -562,6 +567,24 @@ static void build_prompt(ds4_engine *engine, const backend::PredictOptions *requ ds4_chat_append_assistant_prefix(engine, out, think); } +// check_model_identity mirrors pkg/grpc/server.go and +// backend/python/common/model_identity.py. Either side empty means "skip": the +// request side is empty for a controller that predates the field, the loaded +// side when such a controller performed the load. A false rejection is worse +// than the miss it prevents. Callers must already hold g_engine_mu. +static GStatus check_model_identity(const backend::PredictOptions *request) { + if (request == nullptr || request->modelidentity().empty()) return GStatus::OK; + if (g_loaded_model_identity.empty() || + g_loaded_model_identity == request->modelidentity()) { + return GStatus::OK; + } + // NOT_FOUND plus this exact sentinel is the cross-language contract the + // router matches on (grpcerrors.ModelMismatchSentinel). + return GStatus(StatusCode::NOT_FOUND, + "ds4: model identity mismatch: loaded \"" + g_loaded_model_identity + + "\", requested \"" + request->modelidentity() + "\""); +} + class DS4Backend final : public backend::Backend::Service { public: GStatus Health(ServerContext *, const backend::HealthMessage *, @@ -716,6 +739,7 @@ class DS4Backend final : public backend::Backend::Service { } result->set_success(true); + g_loaded_model_identity = request->model(); result->set_message("loaded " + model_path); return GStatus::OK; } @@ -724,6 +748,7 @@ class DS4Backend final : public backend::Backend::Service { backend::TokenizationResponse *response) override { std::lock_guard lock(g_engine_mu); if (!g_engine) return GStatus(StatusCode::FAILED_PRECONDITION, "ds4: model not loaded"); + if (GStatus id = check_model_identity(request); !id.ok()) return id; ds4_tokens out = {}; ds4_tokenize_text(g_engine, request->prompt().c_str(), &out); for (int i = 0; i < out.len; ++i) response->add_tokens(out.v[i]); @@ -738,6 +763,7 @@ class DS4Backend final : public backend::Backend::Service { if (!g_engine || !g_session) { return GStatus(StatusCode::FAILED_PRECONDITION, "ds4: model not loaded"); } + if (GStatus id = check_model_identity(request); !id.ok()) return id; if (std::string route_err = wait_route_ready(lock); !route_err.empty()) { return GStatus(StatusCode::UNAVAILABLE, route_err); } @@ -837,6 +863,7 @@ class DS4Backend final : public backend::Backend::Service { if (!g_engine || !g_session) { return GStatus(StatusCode::FAILED_PRECONDITION, "ds4: model not loaded"); } + if (GStatus id = check_model_identity(request); !id.ok()) return id; if (std::string route_err = wait_route_ready(lock); !route_err.empty()) { return GStatus(StatusCode::UNAVAILABLE, route_err); } diff --git a/backend/cpp/ik-llama-cpp/grpc-server.cpp b/backend/cpp/ik-llama-cpp/grpc-server.cpp index 96eea3b58351..ce43d6d016d0 100644 --- a/backend/cpp/ik-llama-cpp/grpc-server.cpp +++ b/backend/cpp/ik-llama-cpp/grpc-server.cpp @@ -2412,7 +2412,33 @@ static void params_parse(const backend::ModelOptions* request, // GRPC Server start class BackendServiceImpl final : public backend::Backend::Service { +private: + // The ModelOptions.Model this process was loaded with. Compared against + // PredictOptions.ModelIdentity so a request that reached us through a stale + // distributed route is rejected instead of answered from the wrong model + // (#10952). + std::string loaded_model_identity; + public: + // checkModelIdentity mirrors pkg/grpc/server.go and + // backend/python/common/model_identity.py. Either side being empty means + // "skip": the request side is empty for a controller that predates the field, + // and the loaded side is empty when such a controller performed the load. A + // false rejection is worse than the miss it prevents. + grpc::Status checkModelIdentity(const backend::PredictOptions* request) { + if (request == nullptr || request->modelidentity().empty()) { + return grpc::Status::OK; + } + if (loaded_model_identity.empty() || loaded_model_identity == request->modelidentity()) { + return grpc::Status::OK; + } + // NOT_FOUND plus this exact sentinel is the cross-language contract the + // router matches on (grpcerrors.ModelMismatchSentinel). + return grpc::Status(grpc::StatusCode::NOT_FOUND, + "ik-llama-cpp: model identity mismatch: loaded \"" + loaded_model_identity + + "\", requested \"" + request->modelidentity() + "\""); + } + grpc::Status Health(ServerContext* context, const backend::HealthMessage* request, backend::Reply* reply) { // Implement Health RPC reply->set_message("OK"); @@ -2438,9 +2464,12 @@ class BackendServiceImpl final : public backend::Backend::Service { result->set_message("Loading succeeded"); result->set_success(true); loaded_model = true; + loaded_model_identity = request->model(); return Status::OK; } grpc::Status PredictStream(grpc::ServerContext* context, const backend::PredictOptions* request, grpc::ServerWriter* writer) override { + auto identity = checkModelIdentity(request); + if (!identity.ok()) return identity; json data = parse_options(true, request, llama); const int task_id = llama.queue_tasks.get_new_id(); llama.queue_results.add_waiting_task_id(task_id); @@ -2495,6 +2524,8 @@ class BackendServiceImpl final : public backend::Backend::Service { grpc::Status Predict(ServerContext* context, const backend::PredictOptions* request, backend::Reply* reply) { + auto identity = checkModelIdentity(request); + if (!identity.ok()) return identity; json data = parse_options(false, request, llama); const int task_id = llama.queue_tasks.get_new_id(); llama.queue_results.add_waiting_task_id(task_id); @@ -2532,6 +2563,8 @@ class BackendServiceImpl final : public backend::Backend::Service { /// https://github.com/ggerganov/llama.cpp/blob/aa2341298924ac89778252015efcb792f2df1e20/examples/server/server.cpp#L2969 grpc::Status Embedding(ServerContext* context, const backend::PredictOptions* request, backend::EmbeddingResult* embeddingResult) { + auto identity = checkModelIdentity(request); + if (!identity.ok()) return identity; json data = parse_options(false, request, llama); const int task_id = llama.queue_tasks.get_new_id(); llama.queue_results.add_waiting_task_id(task_id); @@ -2556,6 +2589,8 @@ class BackendServiceImpl final : public backend::Backend::Service { } grpc::Status TokenizeString(ServerContext* context, const backend::PredictOptions* request, backend::TokenizationResponse* response){ + auto identity = checkModelIdentity(request); + if (!identity.ok()) return identity; json data = parse_options(false, request, llama); std::vector tokens = llama.tokenize(data["prompt"],false); diff --git a/backend/cpp/llama-cpp/grpc-server.cpp b/backend/cpp/llama-cpp/grpc-server.cpp index cd4fb341d326..995b61affd8c 100644 --- a/backend/cpp/llama-cpp/grpc-server.cpp +++ b/backend/cpp/llama-cpp/grpc-server.cpp @@ -1401,10 +1401,36 @@ class BackendServiceImpl final : public backend::Backend::Service { private: server_context& ctx_server; common_params params_base; // Store copy of params_base, set after model load + // The ModelOptions.Model this process was loaded with. Compared against + // PredictOptions.ModelIdentity so a request that reached us through a stale + // distributed route is rejected instead of answered from the wrong model + // (#10952). Written under LoadModel, read by the inference RPCs. + std::string loaded_model_identity; public: BackendServiceImpl(server_context& ctx) : ctx_server(ctx) {} + // checkModelIdentity mirrors pkg/grpc/server.go and + // backend/python/common/model_identity.py. Either side being empty means + // "skip": the request side is empty for a controller that predates the + // field and for the synthetic PredictOptions this server builds internally + // for ASR, and the loaded side is empty when such a controller performed + // the load. A false rejection is worse than the miss it prevents. + grpc::Status checkModelIdentity(const backend::PredictOptions* request) { + if (request == nullptr || request->modelidentity().empty()) { + return grpc::Status::OK; + } + if (loaded_model_identity.empty() || loaded_model_identity == request->modelidentity()) { + return grpc::Status::OK; + } + // NOT_FOUND plus this exact sentinel is the cross-language contract the + // router matches on (grpcerrors.ModelMismatchSentinel). The code alone + // is not enough: NOT_FOUND is returned for unrelated reasons elsewhere. + return grpc::Status(grpc::StatusCode::NOT_FOUND, + "llama-cpp: model identity mismatch: loaded \"" + loaded_model_identity + + "\", requested \"" + request->modelidentity() + "\""); + } + grpc::Status Health(ServerContext* context, const backend::HealthMessage* /*request*/, backend::Reply* reply) override { auto auth = checkAuth(context); if (!auth.ok()) return auth; @@ -1535,6 +1561,7 @@ class BackendServiceImpl final : public backend::Backend::Service { result->set_message("Loading succeeded"); result->set_success(true); loaded_model = true; + loaded_model_identity = request->model(); // Store copy of params_base for use in parse_options and other methods params_base = params; @@ -1616,6 +1643,8 @@ class BackendServiceImpl final : public backend::Backend::Service { grpc::Status PredictStream(grpc::ServerContext* context, const backend::PredictOptions* request, grpc::ServerWriter* writer) override { auto auth = checkAuth(context); if (!auth.ok()) return auth; + auto identity = checkModelIdentity(request); + if (!identity.ok()) return identity; if (params_base.model.path.empty()) { return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION, "Model not loaded"); } @@ -2183,6 +2212,8 @@ class BackendServiceImpl final : public backend::Backend::Service { grpc::Status Predict(ServerContext* context, const backend::PredictOptions* request, backend::Reply* reply) override { auto auth = checkAuth(context); if (!auth.ok()) return auth; + auto identity = checkModelIdentity(request); + if (!identity.ok()) return identity; if (params_base.model.path.empty()) { return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION, "Model not loaded"); } @@ -2715,6 +2746,8 @@ class BackendServiceImpl final : public backend::Backend::Service { grpc::Status Embedding(ServerContext* context, const backend::PredictOptions* request, backend::EmbeddingResult* embeddingResult) override { auto auth = checkAuth(context); if (!auth.ok()) return auth; + auto identity = checkModelIdentity(request); + if (!identity.ok()) return identity; if (params_base.model.path.empty()) { return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION, "Model not loaded"); } @@ -3108,6 +3141,8 @@ class BackendServiceImpl final : public backend::Backend::Service { grpc::Status TokenizeString(ServerContext* context, const backend::PredictOptions* request, backend::TokenizationResponse* response) override { auto auth = checkAuth(context); if (!auth.ok()) return auth; + auto identity = checkModelIdentity(request); + if (!identity.ok()) return identity; if (params_base.model.path.empty()) { return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION, "Model not loaded"); } diff --git a/backend/python/common/grpc_auth.py b/backend/python/common/grpc_auth.py index 9ed866abb502..a765a78098e3 100644 --- a/backend/python/common/grpc_auth.py +++ b/backend/python/common/grpc_auth.py @@ -11,6 +11,7 @@ import grpc +from model_identity import AsyncModelIdentityInterceptor, ModelIdentityInterceptor from parent_watch import start_parent_death_watcher @@ -64,13 +65,15 @@ async def intercept_service(self, continuation, handler_call_details): def get_auth_interceptors(*, aio: bool = False): - """Return a list of gRPC interceptors for bearer token auth. + """Return the gRPC server interceptors every LocalAI Python backend installs. + + Always includes model-identity enforcement (model_identity.py), which is + unrelated to authentication. Bearer token auth is added on top only when + LOCALAI_GRPC_AUTH_TOKEN is set. Args: aio: If True, return async-compatible interceptors for grpc.aio.server(). If False (default), return sync interceptors for grpc.server(). - - Returns an empty list when LOCALAI_GRPC_AUTH_TOKEN is not set. """ # Arm the best-effort parent-death backstop here: this is the single helper # every LocalAI Python backend invokes exactly once while building its gRPC @@ -79,9 +82,17 @@ def get_auth_interceptors(*, aio: bool = False): # unsupported platforms — see parent_watch.py. start_parent_death_watcher() + # Model-identity enforcement is independent of authentication and must be + # installed BEFORE the token check returns. gRPC auth is off by default, so + # an identity interceptor added below the early return would never be + # installed on any Python backend, and nothing would report it. + interceptors = [AsyncModelIdentityInterceptor()] if aio else [ModelIdentityInterceptor()] + token = os.environ.get("LOCALAI_GRPC_AUTH_TOKEN", "") if not token: - return [] + return interceptors if aio: - return [AsyncTokenAuthInterceptor(token)] - return [TokenAuthInterceptor(token)] + interceptors.append(AsyncTokenAuthInterceptor(token)) + else: + interceptors.append(TokenAuthInterceptor(token)) + return interceptors diff --git a/backend/python/common/model_identity.py b/backend/python/common/model_identity.py new file mode 100644 index 000000000000..fb64f93add76 --- /dev/null +++ b/backend/python/common/model_identity.py @@ -0,0 +1,192 @@ +"""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). + +PredictOptions.ModelIdentity carries 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 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" + +# The four RPCs that carry PredictOptions. Nothing else has an identity field. +# TTS and SoundGeneration are excluded on purpose: their `model` field is +# already rewritten to a worker-local path by the controller's +# FileStagingClient, so comparing it would reject valid requests. +_GUARDED_METHODS = frozenset( + ( + "/backend.Backend/Predict", + "/backend.Backend/PredictStream", + "/backend.Backend/Embedding", + "/backend.Backend/TokenizeString", + ) +) + + +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 the four + guarded RPCs all take a single request message. + """ + 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) + + +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): + result = await 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) + async for response in original_stream(request, context): + yield response + + 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 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 new file mode 100644 index 000000000000..2cb7011170b8 --- /dev/null +++ b/backend/python/common/model_identity_test.py @@ -0,0 +1,222 @@ +"""Unit tests for model-identity enforcement (model_identity.py). + +Run inside any backend venv (needs grpcio, which every Python backend has): + python -m unittest model_identity_test + +Mirrors the Go coverage in pkg/grpc/model_identity_test.go and +pkg/grpc/grpcerrors/errors_test.go. The rules under test are the ones whose +failure modes are silent: enforcement that is wired up but never installed, and +enforcement that rejects requests it should serve. +""" + +import os +import unittest + +import grpc + +import grpc_auth +import model_identity + + +class _Aborted(Exception): + pass + + +class _FakeContext: + """Minimal ServicerContext: abort() records and raises, like the real one.""" + + def __init__(self): + self.code = None + self.details = None + + def abort(self, code, details): + self.code = code + self.details = details + raise _Aborted(details) + + +class _FakeCallDetails: + def __init__(self, method): + self.method = method + self.invocation_metadata = () + + +class _Request: + """Stands in for ModelOptions / PredictOptions. + + The generated protobuf classes are built at container-build time and are + not importable here, but the interceptor only ever reads two attributes. + """ + + def __init__(self, Model="", ModelIdentity=""): + self.Model = Model + self.ModelIdentity = ModelIdentity + + +class _Result: + def __init__(self, success=True): + self.success = success + + +def _handler(behavior, response_streaming=False): + if response_streaming: + return grpc.unary_stream_rpc_method_handler(behavior) + return grpc.unary_unary_rpc_method_handler(behavior) + + +class TestInterceptorInstalled(unittest.TestCase): + """The wiring, which is where this can silently do nothing. + + get_auth_interceptors() returns early when LOCALAI_GRPC_AUTH_TOKEN is + unset, which is the DEFAULT configuration. An identity interceptor added + after that return is never installed on any of the 36 Python backends, and + nothing else would notice. + """ + + def setUp(self): + self._saved = os.environ.get("LOCALAI_GRPC_AUTH_TOKEN") + os.environ.pop("LOCALAI_GRPC_AUTH_TOKEN", None) + + def tearDown(self): + if self._saved is None: + os.environ.pop("LOCALAI_GRPC_AUTH_TOKEN", None) + else: + os.environ["LOCALAI_GRPC_AUTH_TOKEN"] = self._saved + + def test_installed_when_auth_is_disabled(self): + interceptors = grpc_auth.get_auth_interceptors() + self.assertTrue( + any(isinstance(i, model_identity.ModelIdentityInterceptor) for i in interceptors), + "identity enforcement must be installed even with gRPC auth off " + "(the default); got {!r}".format(interceptors), + ) + + def test_installed_when_auth_is_disabled_aio(self): + interceptors = grpc_auth.get_auth_interceptors(aio=True) + self.assertTrue( + any( + isinstance(i, model_identity.AsyncModelIdentityInterceptor) + for i in interceptors + ), + "async identity enforcement must be installed with gRPC auth off", + ) + + def test_installed_alongside_auth_when_enabled(self): + os.environ["LOCALAI_GRPC_AUTH_TOKEN"] = "secret" + interceptors = grpc_auth.get_auth_interceptors() + self.assertTrue( + any(isinstance(i, model_identity.ModelIdentityInterceptor) for i in interceptors) + ) + self.assertTrue( + any(isinstance(i, grpc_auth.TokenAuthInterceptor) for i in interceptors) + ) + + +class TestMismatchRule(unittest.TestCase): + """The pure policy. Every 'serve' case here is a false-rejection guard.""" + + def setUp(self): + self.state = model_identity.ModelIdentityState() + + def test_rejects_a_different_model(self): + self.state.record("a.gguf") + message = self.state.mismatch("b.gguf") + self.assertIsNotNone(message) + self.assertIn(model_identity.MODEL_MISMATCH_SENTINEL, message) + self.assertIn("a.gguf", message) + self.assertIn("b.gguf", message) + + def test_serves_the_same_model(self): + self.state.record("a.gguf") + self.assertIsNone(self.state.mismatch("a.gguf")) + + def test_serves_when_the_request_has_no_identity(self): + self.state.record("a.gguf") + self.assertIsNone(self.state.mismatch("")) + + def test_serves_when_nothing_was_recorded(self): + self.assertIsNone(self.state.mismatch("b.gguf")) + + +class TestInterceptorBehavior(unittest.TestCase): + def setUp(self): + self.interceptor = model_identity.ModelIdentityInterceptor() + self.served = [] + + def _intercept(self, method, handler): + return self.interceptor.intercept_service( + lambda _: handler, _FakeCallDetails(method) + ) + + def _load(self, model, success=True): + handler = _handler(lambda request, context: _Result(success=success)) + wrapped = self._intercept("/backend.Backend/LoadModel", handler) + wrapped.unary_unary(_Request(Model=model), _FakeContext()) + + def _call(self, method, identity, response_streaming=False): + def behavior(request, context): + self.served.append(method) + return "served" + + handler = _handler(behavior, response_streaming=response_streaming) + wrapped = self._intercept(method, handler) + context = _FakeContext() + behavior_fn = wrapped.unary_stream if response_streaming else wrapped.unary_unary + return behavior_fn(_Request(ModelIdentity=identity), context), context + + def test_load_records_the_identity(self): + self._load("a.gguf") + self.assertEqual(self.interceptor.state.loaded, "a.gguf") + + def test_failed_load_records_nothing(self): + self._load("a.gguf", success=False) + self.assertEqual(self.interceptor.state.loaded, "") + + def test_rejects_every_guarded_rpc_on_mismatch(self): + self._load("a.gguf") + for method in sorted(model_identity._GUARDED_METHODS): + streaming = method.endswith("PredictStream") + with self.subTest(method=method): + with self.assertRaises(_Aborted): + self._call(method, "b.gguf", response_streaming=streaming) + self.assertEqual(self.served, [], "no request may reach the model") + + def test_reject_uses_not_found_and_the_sentinel(self): + self._load("a.gguf") + + def behavior(request, context): + return "served" + + wrapped = self._intercept( + "/backend.Backend/Predict", _handler(behavior) + ) + context = _FakeContext() + with self.assertRaises(_Aborted): + wrapped.unary_unary(_Request(ModelIdentity="b.gguf"), context) + self.assertEqual(context.code, grpc.StatusCode.NOT_FOUND) + self.assertIn(model_identity.MODEL_MISMATCH_SENTINEL, context.details) + + def test_serves_matching_identity(self): + self._load("a.gguf") + for method in sorted(model_identity._GUARDED_METHODS): + streaming = method.endswith("PredictStream") + self._call(method, "a.gguf", response_streaming=streaming) + self.assertEqual(len(self.served), len(model_identity._GUARDED_METHODS)) + + def test_serves_request_without_identity(self): + self._load("a.gguf") + self._call("/backend.Backend/Predict", "") + self.assertEqual(self.served, ["/backend.Backend/Predict"]) + + def test_serves_when_load_recorded_nothing(self): + self._call("/backend.Backend/Predict", "b.gguf") + self.assertEqual(self.served, ["/backend.Backend/Predict"]) + + def test_unguarded_rpcs_pass_through_untouched(self): + handler = _handler(lambda request, context: "served") + wrapped = self._intercept("/backend.Backend/TTS", handler) + self.assertIs(wrapped, handler) + + +if __name__ == "__main__": + unittest.main() diff --git a/core/backend/options.go b/core/backend/options.go index a7cf00dc3dd2..72f49f1eda47 100644 --- a/core/backend/options.go +++ b/core/backend/options.go @@ -514,6 +514,13 @@ func gRPCPredictOpts(c config.ModelConfig, modelPath string) *pb.PredictOptions } pbOpts := &pb.PredictOptions{ + // c.Model, not c.ModelID()/c.ModelFileName(): this must be the SAME + // expression ModelOptions feeds to model.WithModel, which is what the + // backend receives as ModelOptions.Model at LoadModel time. Both are + // read from this same config value, so the backend's equality check + // cannot false-reject. See PredictOptions.ModelIdentity in + // backend/backend.proto and #10952. + ModelIdentity: c.Model, Temperature: float32(*c.Temperature), TopP: float32(*c.TopP), NDraft: c.NDraft, diff --git a/core/backend/options_internal_test.go b/core/backend/options_internal_test.go index f8ac7244af64..bf4258bbfdc2 100644 --- a/core/backend/options_internal_test.go +++ b/core/backend/options_internal_test.go @@ -314,3 +314,44 @@ var _ = Describe("EffectiveContextSize", func() { }) }) }) + +// Guards the model identity carried in PredictOptions, which lets a backend +// reject a request that reached it through a stale distributed route (#10952). +// +// The safety of the whole mechanism rests on the predict-time value being the +// SAME expression as the load-time one: ModelOptions passes model.WithModel( +// c.Model), which becomes ModelOptions.Model at LoadModel, and gRPCPredictOpts +// receives the same config value in the same function a few lines later. If +// this ever starts sending ModelID()/Name or the resolved file path instead, +// every request to a correctly-routed backend gets rejected. The configs below +// deliberately give Model, Name and ModelID() three different values so that +// substituting any of them for the others fails here. +var _ = Describe("gRPCPredictOpts model identity", func() { + withModel := func(name, modelFile string) config.ModelConfig { + cfg := config.ModelConfig{} + cfg.SetDefaults() + cfg.Name = name + cfg.Model = modelFile + return cfg + } + + It("sends ModelConfig.Model, the value LoadModel receives", func() { + cfg := withModel("friendly-name", "qwen/actual-weights.gguf") + opts := gRPCPredictOpts(cfg, "/tmp/models") + Expect(opts.ModelIdentity).To(Equal("qwen/actual-weights.gguf")) + }) + + It("does not send ModelID(), which LoadModel never receives", func() { + cfg := withModel("friendly-name", "qwen/actual-weights.gguf") + Expect(cfg.ModelID()).To(Equal("friendly-name")) + opts := gRPCPredictOpts(cfg, "/tmp/models") + Expect(opts.ModelIdentity).ToNot(Equal(cfg.ModelID())) + }) + + // Configs with no model file cannot identify anything. Empty means "skip + // the check" on the backend, which is the safe direction. + It("leaves the identity empty when the config names no model", func() { + opts := gRPCPredictOpts(withModel("some-name", ""), "/tmp/models") + Expect(opts.ModelIdentity).To(BeEmpty()) + }) +}) diff --git a/core/services/nodes/inflight.go b/core/services/nodes/inflight.go index bfc71b999a89..f4bb8988f7d8 100644 --- a/core/services/nodes/inflight.go +++ b/core/services/nodes/inflight.go @@ -86,18 +86,29 @@ func (c *InFlightTrackingClient) track(ctx context.Context) func() { // request triggers a fresh load instead of routing back here. Without this the // model stays unreachable until the controller restarts. The original error is // returned unchanged. +// +// It also drops the row when the backend reports that it has a DIFFERENT model +// loaded. That happens when a worker recycles a stopped backend's gRPC port for +// another model's backend: the cached row still names that address, and a +// liveness-only probe cannot tell the stale row from a valid one, so the wrong +// backend answers (#10952). The two conditions differ in cause but not in cure. func (c *InFlightTrackingClient) reconcile(err error) error { - if !grpcerrors.IsModelNotLoaded(err) { + mismatch := grpcerrors.IsModelMismatch(err) + if !grpcerrors.IsModelNotLoaded(err) && !mismatch { return err } + reason := "model not loaded" + if mismatch { + reason = "wrong model loaded" + } rmCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() if rmErr := c.registry.RemoveNodeModel(rmCtx, c.nodeID, c.modelName, c.replicaIndex); rmErr != nil { - xlog.Warn("Failed to drop stale replica after model-not-loaded", - "node", c.nodeID, "model", c.modelName, "replica", c.replicaIndex, "error", rmErr) + xlog.Warn("Failed to drop stale replica", + "node", c.nodeID, "model", c.modelName, "replica", c.replicaIndex, "reason", reason, "error", rmErr) } else { - xlog.Warn("Backend reports model not loaded; dropped stale replica so the next request reloads", - "node", c.nodeID, "model", c.modelName, "replica", c.replicaIndex) + xlog.Warn("Dropped stale replica so the next request reloads", + "node", c.nodeID, "model", c.modelName, "replica", c.replicaIndex, "reason", reason) } return err } diff --git a/core/services/nodes/inflight_test.go b/core/services/nodes/inflight_test.go index 1b5755cd6769..120d2e3129ec 100644 --- a/core/services/nodes/inflight_test.go +++ b/core/services/nodes/inflight_test.go @@ -9,8 +9,11 @@ import ( . "github.com/onsi/gomega" grpc "github.com/mudler/LocalAI/pkg/grpc" + "github.com/mudler/LocalAI/pkg/grpc/grpcerrors" pb "github.com/mudler/LocalAI/pkg/grpc/proto" ggrpc "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" ) // --- Fakes --- @@ -449,4 +452,41 @@ var _ = Describe("InFlightTrackingClient", func() { Expect(tracker.removed).To(Equal(1)) }) }) + + // A wrong-model answer means the cached row points at a port that has been + // recycled for another model's backend (#10952). The row is stale for a + // different reason than "model not loaded", but the cure is the same: drop + // it so the next request reloads somewhere correct. + Describe("wrong-model self-heal", func() { + It("removes the replica when the backend reports a model mismatch", func() { + backend.predictErr = grpcerrors.ModelMismatch("llama-cpp", "a.gguf", "b.gguf") + _, err := client.Predict(context.Background(), &pb.PredictOptions{}) + Expect(err).To(HaveOccurred()) + Expect(tracker.removed).To(Equal(1)) + }) + + It("removes the replica on a streamed call too", func() { + backend.streamErr = grpcerrors.ModelMismatch("llama-cpp", "a.gguf", "b.gguf") + err := client.PredictStream(context.Background(), &pb.PredictOptions{}, func(*pb.Reply) {}) + Expect(err).To(HaveOccurred()) + Expect(tracker.removed).To(Equal(1)) + }) + + It("returns the mismatch error unchanged so the caller can tell why", func() { + backend.predictErr = grpcerrors.ModelMismatch("llama-cpp", "a.gguf", "b.gguf") + _, err := client.Predict(context.Background(), &pb.PredictOptions{}) + Expect(grpcerrors.IsModelMismatch(err)).To(BeTrue()) + }) + + // The sentinel is what separates our signal from an ordinary NotFound. + // insightface's Embedding answers NOT_FOUND "no face detected" on a + // PredictOptions RPC, and dropping a healthy row for that would evict + // a working replica every time a photo has no face in it. + It("keeps the replica on an unrelated NotFound", func() { + backend.predictErr = status.Error(codes.NotFound, "no face detected") + _, err := client.Predict(context.Background(), &pb.PredictOptions{}) + Expect(err).To(HaveOccurred()) + Expect(tracker.removed).To(Equal(0)) + }) + }) }) diff --git a/core/services/nodes/model_identity_e2e_test.go b/core/services/nodes/model_identity_e2e_test.go new file mode 100644 index 000000000000..9306d9ae6851 --- /dev/null +++ b/core/services/nodes/model_identity_e2e_test.go @@ -0,0 +1,91 @@ +package nodes + +import ( + "context" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + grpc "github.com/mudler/LocalAI/pkg/grpc" + "github.com/mudler/LocalAI/pkg/grpc/base" + "github.com/mudler/LocalAI/pkg/grpc/grpcerrors" + pb "github.com/mudler/LocalAI/pkg/grpc/proto" +) + +// staleRouteBackend stands in for the backend a recycled port now points at: +// a real, healthy, loaded process that simply holds a DIFFERENT model than the +// controller's cached row believes. +type staleRouteBackend struct { + base.SingleThread + + served int +} + +func (b *staleRouteBackend) Load(*pb.ModelOptions) error { return nil } + +func (b *staleRouteBackend) Predict(*pb.PredictOptions) (string, error) { + b.served++ + return "an answer from the WRONG model", nil +} + +// This joins the two halves of the #10952 fix that are unit-tested separately +// elsewhere: the backend rejecting a request whose identity does not match what +// it loaded (pkg/grpc), and the router dropping the stale row when it sees that +// rejection (inflight.go). Neither half is worth much without the other, and +// nothing else exercises them against a real gRPC server together. +var _ = Describe("stale distributed route is caught end to end", func() { + var ( + tracker *fakeInFlightTracker + llm *staleRouteBackend + tracked *InFlightTrackingClient + ) + + BeforeEach(func() { + tracker = &fakeInFlightTracker{} + llm = &staleRouteBackend{} + + addr := "test://stale-route" + grpc.Provide(addr, llm) + client := grpc.NewClient(addr, true, nil, false) + + // The worker loaded "model-a" on this port. + _, err := client.LoadModel(context.Background(), &pb.ModelOptions{Model: "model-a.gguf"}) + Expect(err).ToNot(HaveOccurred()) + + tracked = NewInFlightTrackingClient(client, tracker, "node-1", "model-b", 0) + }) + + It("rejects the wrong-model request and drops the stale replica row", func() { + // The controller's cached row for "model-b" points here because the + // port was recycled. Liveness alone cannot tell, so the identity does. + _, err := tracked.Predict(context.Background(), &pb.PredictOptions{ + Prompt: "hello", + ModelIdentity: "model-b.gguf", + }) + + Expect(err).To(HaveOccurred()) + Expect(grpcerrors.IsModelMismatch(err)).To(BeTrue(), "got %v", err) + Expect(llm.served).To(Equal(0), "the wrong model must never answer") + Expect(tracker.removed).To(Equal(1), "the stale row must be dropped so the next request reloads") + }) + + It("serves and keeps the row when the route is correct", func() { + _, err := tracked.Predict(context.Background(), &pb.PredictOptions{ + Prompt: "hello", + ModelIdentity: "model-a.gguf", + }) + + Expect(err).ToNot(HaveOccurred()) + Expect(llm.served).To(Equal(1)) + Expect(tracker.removed).To(Equal(0)) + }) + + // Every deployment that has not upgraded its controller sends no identity. + It("serves and keeps the row when the controller sends no identity", func() { + _, err := tracked.Predict(context.Background(), &pb.PredictOptions{Prompt: "hello"}) + + Expect(err).ToNot(HaveOccurred()) + Expect(llm.served).To(Equal(1)) + Expect(tracker.removed).To(Equal(0)) + }) +}) diff --git a/pkg/grpc/grpcerrors/errors.go b/pkg/grpc/grpcerrors/errors.go index 8cb57f41675f..4a0306338e25 100644 --- a/pkg/grpc/grpcerrors/errors.go +++ b/pkg/grpc/grpcerrors/errors.go @@ -34,6 +34,43 @@ func IsModelNotLoaded(err error) bool { return strings.Contains(strings.ToLower(err.Error()), "model not loaded") } +// ModelMismatchSentinel is the fixed marker every backend puts in a +// model-identity mismatch error, in every language. IsModelMismatch requires +// it, so it is part of the cross-language wire contract: changing it breaks +// detection for backends that have not been rebuilt. +const ModelMismatchSentinel = "model identity mismatch" + +// ModelMismatch returns the canonical error a backend returns when the request +// names a model other than the one it has loaded. That means the router's +// cached row is stale and points at a recycled port, so the caller should drop +// the row rather than retry against the same replica. +func ModelMismatch(backend, loaded, requested string) error { + return status.Errorf(codes.NotFound, "%s: %s: loaded %q, requested %q", + backend, ModelMismatchSentinel, loaded, requested) +} + +// IsModelMismatch reports whether err signals that the backend has a DIFFERENT +// model loaded than the request asked for. +// +// It requires BOTH the code and the sentinel, which inverts the "code OR +// message" style of the helpers above. That is deliberate, and must not be +// "simplified" to a code check: codes.NotFound is not exclusively ours on the +// PredictOptions RPCs. backend/python/insightface/backend.py:127 returns +// NOT_FOUND "no face detected" from Embedding, and a code-only check would +// make the router drop a healthy replica row on every faceless image. +// +// Unlike IsModelNotLoaded, a false positive here is NOT harmless in the same +// way, which is the other half of why the sentinel is mandatory. +func IsModelMismatch(err error) bool { + if err == nil { + return false + } + if status.Code(err) != codes.NotFound { + return false + } + return strings.Contains(strings.ToLower(err.Error()), ModelMismatchSentinel) +} + // LiveTranscriptionUnsupported returns the canonical error a backend returns // when it (or the loaded model) cannot serve the bidirectional // AudioTranscriptionLive RPC. It carries codes.Unimplemented deliberately: diff --git a/pkg/grpc/grpcerrors/errors_test.go b/pkg/grpc/grpcerrors/errors_test.go index 7ce668226920..cb7be2978e1a 100644 --- a/pkg/grpc/grpcerrors/errors_test.go +++ b/pkg/grpc/grpcerrors/errors_test.go @@ -35,6 +35,39 @@ var _ = Describe("grpcerrors", func() { Expect(status.Code(grpcerrors.ModelNotLoaded("whisper"))).To(Equal(codes.FailedPrecondition)) }) + DescribeTable("IsModelMismatch", + func(err error, want bool) { + Expect(grpcerrors.IsModelMismatch(err)).To(Equal(want)) + }, + Entry("nil", nil, false), + Entry("typed via constructor", grpcerrors.ModelMismatch("llama-cpp", "a.gguf", "b.gguf"), true), + // The sentinel is what makes this detectable, not the code alone. + // insightface's Embedding returns NOT_FOUND "no face detected" on a + // PredictOptions RPC (backend/python/insightface/backend.py:127). A + // code-only check would make the router drop a healthy replica row on + // every faceless image, so this entry must stay false. + Entry("insightface no-face NotFound", status.Error(codes.NotFound, "no face detected"), false), + Entry("insightface no-face in both images", + status.Error(codes.NotFound, "no face detected in one or both images"), false), + Entry("unrelated NotFound", status.Error(codes.NotFound, "Job 7 not found"), false), + // Must not overlap with the model-not-loaded signal, which the router + // handles identically today but for a different reason. + Entry("model not loaded is not a mismatch", grpcerrors.ModelNotLoaded("llama-cpp"), false), + Entry("unrelated error", errors.New("context deadline exceeded"), false), + ) + + It("ModelMismatch carries NotFound and names both models", func() { + err := grpcerrors.ModelMismatch("llama-cpp", "a.gguf", "b.gguf") + Expect(status.Code(err)).To(Equal(codes.NotFound)) + Expect(err.Error()).To(ContainSubstring(grpcerrors.ModelMismatchSentinel)) + Expect(err.Error()).To(ContainSubstring("a.gguf")) + Expect(err.Error()).To(ContainSubstring("b.gguf")) + }) + + It("IsModelNotLoaded does not claim a mismatch", func() { + Expect(grpcerrors.IsModelNotLoaded(grpcerrors.ModelMismatch("llama-cpp", "a", "b"))).To(BeFalse()) + }) + DescribeTable("IsLiveTranscriptionUnsupported", func(err error, want bool) { Expect(grpcerrors.IsLiveTranscriptionUnsupported(err)).To(Equal(want)) diff --git a/pkg/grpc/model_identity_test.go b/pkg/grpc/model_identity_test.go new file mode 100644 index 000000000000..ad63d1120c67 --- /dev/null +++ b/pkg/grpc/model_identity_test.go @@ -0,0 +1,128 @@ +package grpc + +import ( + "context" + + "github.com/mudler/LocalAI/pkg/grpc/base" + "github.com/mudler/LocalAI/pkg/grpc/grpcerrors" + pb "github.com/mudler/LocalAI/pkg/grpc/proto" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// identityBackend records what it was loaded with and answers every inference +// RPC successfully. Any request that reaches it has passed the identity guard, +// so `served` is the signal for "the guard let this through". +type identityBackend struct { + base.SingleThread + + loaded string + served int +} + +func (b *identityBackend) Load(opts *pb.ModelOptions) error { + b.loaded = opts.Model + return nil +} + +func (b *identityBackend) Predict(*pb.PredictOptions) (string, error) { + b.served++ + return "ok", nil +} + +func (b *identityBackend) PredictStream(_ *pb.PredictOptions, ch chan string) error { + b.served++ + ch <- "ok" + close(ch) + return nil +} + +func (b *identityBackend) Embeddings(*pb.PredictOptions) ([]float32, error) { + b.served++ + return []float32{1}, nil +} + +func (b *identityBackend) TokenizeString(*pb.PredictOptions) (pb.TokenizationResponse, error) { + b.served++ + return pb.TokenizationResponse{Length: 1}, nil +} + +var _ AIModel = (*identityBackend)(nil) + +// callAll exercises the four PredictOptions RPCs and returns the first error. +// All four share one guard, so all four must behave identically. +func callAll(c Backend, in *pb.PredictOptions) []error { + ctx := context.Background() + errs := []error{} + + _, err := c.Predict(ctx, in) + errs = append(errs, err) + + errs = append(errs, c.PredictStream(ctx, in, func(*pb.Reply) {})) + + _, err = c.Embeddings(ctx, in) + errs = append(errs, err) + + _, err = c.TokenizeString(ctx, in) + errs = append(errs, err) + + return errs +} + +var _ = Describe("PredictOptions model identity guard", func() { + newServed := func(addr, loadedModel string) (Backend, *identityBackend) { + b := &identityBackend{} + Provide(addr, b) + c := NewClient(addr, true, nil, false) + _, err := c.LoadModel(context.Background(), &pb.ModelOptions{Model: loadedModel}) + Expect(err).ToNot(HaveOccurred()) + Expect(b.loaded).To(Equal(loadedModel)) + return c, b + } + + It("rejects every PredictOptions RPC when the identity names another model", func() { + c, b := newServed("test://identity-mismatch", "a.gguf") + + for _, err := range callAll(c, &pb.PredictOptions{ModelIdentity: "b.gguf", Prompt: "hi"}) { + Expect(err).To(HaveOccurred()) + Expect(grpcerrors.IsModelMismatch(err)).To(BeTrue(), "want a mismatch error, got %v", err) + // The router reacts differently to the two signals, so a mismatch + // must never be mistaken for a not-loaded. + Expect(grpcerrors.IsModelNotLoaded(err)).To(BeFalse()) + } + Expect(b.served).To(Equal(0), "no request may reach the model on a mismatch") + }) + + It("serves when the identity matches the loaded model", func() { + c, b := newServed("test://identity-match", "a.gguf") + + for _, err := range callAll(c, &pb.PredictOptions{ModelIdentity: "a.gguf", Prompt: "hi"}) { + Expect(err).ToNot(HaveOccurred()) + } + Expect(b.served).To(Equal(4)) + }) + + // Compatibility, old controller -> new backend. Every existing deployment + // sends no identity, and tests/e2e-backends/backend_test.go drives real + // backends with bare PredictOptions at 8+ call sites. Tightening this + // breaks all of them, so it must fail here first. + It("serves when the request carries no identity", func() { + c, b := newServed("test://identity-empty-request", "a.gguf") + + for _, err := range callAll(c, &pb.PredictOptions{Prompt: "hi"}) { + Expect(err).ToNot(HaveOccurred()) + } + Expect(b.served).To(Equal(4)) + }) + + // The backend side of the same rule: a model loaded without an identity + // (an old controller did the load) cannot judge anything, so it must serve. + It("serves when the backend has no recorded identity", func() { + c, b := newServed("test://identity-empty-loaded", "") + + for _, err := range callAll(c, &pb.PredictOptions{ModelIdentity: "b.gguf", Prompt: "hi"}) { + Expect(err).ToNot(HaveOccurred()) + } + Expect(b.served).To(Equal(4)) + }) +}) diff --git a/pkg/grpc/server.go b/pkg/grpc/server.go index 0ed50360f14b..8bdd88f2fd95 100644 --- a/pkg/grpc/server.go +++ b/pkg/grpc/server.go @@ -10,7 +10,9 @@ import ( "net" "os" "strings" + "sync" + "github.com/mudler/LocalAI/pkg/grpc/grpcerrors" pb "github.com/mudler/LocalAI/pkg/grpc/proto" "google.golang.org/grpc" "google.golang.org/grpc/codes" @@ -31,6 +33,38 @@ import ( type server struct { pb.UnimplementedBackendServer llm AIModel + + // identityMu guards loadedIdentity: LoadModel writes it, the + // PredictOptions RPCs read it, and nothing serialises those against each + // other (llm.Locking() is the model's own lock, and it is optional). + identityMu sync.RWMutex + loadedIdentity string +} + +// checkModelIdentity reports an error when the request names a model other +// than the one this process loaded. It is the point-of-use half of the fix for +// #10952: in distributed mode a worker can recycle a stopped backend's gRPC +// port for another model's backend, and the controller's liveness-only health +// probe cannot tell a stale cached route from a valid one, so the backend has +// to be the one to catch it. +// +// Either side being empty means "skip the check". The request side is empty +// for a controller that predates the field and for internally synthesized +// requests; the loaded side is empty when such a controller performed the +// load. Neither can judge the other, and a false rejection is far worse than +// the miss. +func (s *server) checkModelIdentity(in *pb.PredictOptions) error { + if in == nil || in.ModelIdentity == "" { + return nil + } + s.identityMu.RLock() + loaded := s.loadedIdentity + s.identityMu.RUnlock() + + if loaded == "" || loaded == in.ModelIdentity { + return nil + } + return grpcerrors.ModelMismatch("grpc-server", loaded, in.ModelIdentity) } func (s *server) Health(ctx context.Context, in *pb.HealthMessage) (*pb.Reply, error) { @@ -38,6 +72,9 @@ func (s *server) Health(ctx context.Context, in *pb.HealthMessage) (*pb.Reply, e } func (s *server) Embedding(ctx context.Context, in *pb.PredictOptions) (*pb.EmbeddingResult, error) { + if err := s.checkModelIdentity(in); err != nil { + return nil, err + } if s.llm.Locking() { s.llm.Lock() defer s.llm.Unlock() @@ -60,10 +97,21 @@ func (s *server) LoadModel(ctx context.Context, in *pb.ModelOptions) (*pb.Result if err != nil { return &pb.Result{Message: fmt.Sprintf("Error loading model: %s", err.Error()), Success: false}, err } + + // Record what we loaded so the PredictOptions RPCs can reject requests + // meant for a different model. Only on success: a failed load leaves no + // model, which IsModelNotLoaded already covers. + s.identityMu.Lock() + s.loadedIdentity = in.Model + s.identityMu.Unlock() + return &pb.Result{Message: "Loading succeeded", Success: true}, nil } func (s *server) Predict(ctx context.Context, in *pb.PredictOptions) (*pb.Reply, error) { + if err := s.checkModelIdentity(in); err != nil { + return nil, err + } if s.llm.Locking() { s.llm.Lock() defer s.llm.Unlock() @@ -361,6 +409,9 @@ func (s *server) AudioTranscriptionLive(stream pb.Backend_AudioTranscriptionLive } func (s *server) PredictStream(in *pb.PredictOptions, stream pb.Backend_PredictStreamServer) error { + if err := s.checkModelIdentity(in); err != nil { + return err + } if s.llm.Locking() { s.llm.Lock() defer s.llm.Unlock() @@ -404,6 +455,9 @@ func (s *server) PredictStream(in *pb.PredictOptions, stream pb.Backend_PredictS } func (s *server) TokenizeString(ctx context.Context, in *pb.PredictOptions) (*pb.TokenizationResponse, error) { + if err := s.checkModelIdentity(in); err != nil { + return nil, err + } if s.llm.Locking() { s.llm.Lock() defer s.llm.Unlock() diff --git a/tests/e2e/mock-backend/main.go b/tests/e2e/mock-backend/main.go index 10fb25296e24..32d52cc126e7 100644 --- a/tests/e2e/mock-backend/main.go +++ b/tests/e2e/mock-backend/main.go @@ -15,6 +15,7 @@ import ( "strings" "sync" + "github.com/mudler/LocalAI/pkg/grpc/grpcerrors" pb "github.com/mudler/LocalAI/pkg/grpc/proto" "github.com/mudler/xlog" "google.golang.org/grpc" @@ -54,6 +55,22 @@ func snapshotLoadParams() *pb.ModelOptions { return lastLoadParams } +// checkModelIdentity mirrors the guard the real backends apply (pkg/grpc, +// backend/python/common/model_identity.py, backend/cpp/*/grpc-server.cpp) so +// the distributed e2e suite exercises the #10952 fix rather than only the +// unit tests. Empty on either side means skip, which is why every existing +// spec that sends a bare PredictOptions keeps working. +func checkModelIdentity(in *pb.PredictOptions) error { + if in == nil || in.ModelIdentity == "" { + return nil + } + opts := snapshotLoadParams() + if opts == nil || opts.Model == "" || opts.Model == in.ModelIdentity { + return nil + } + return grpcerrors.ModelMismatch("mock-backend", opts.Model, in.ModelIdentity) +} + // promptHasToolResults checks if the prompt contains evidence of prior tool // execution — specifically the output from the mock MCP server's get_weather tool. func promptHasToolResults(prompt string) bool { @@ -79,6 +96,9 @@ func (m *MockBackend) LoadModel(ctx context.Context, in *pb.ModelOptions) (*pb.R } func (m *MockBackend) Predict(ctx context.Context, in *pb.PredictOptions) (*pb.Reply, error) { + if err := checkModelIdentity(in); err != nil { + return nil, err + } xlog.Debug("Predict called", "prompt", in.Prompt) if strings.Contains(in.Prompt, "MOCK_ERROR") { return nil, fmt.Errorf("mock backend predict error: simulated failure") @@ -232,6 +252,9 @@ func (m *MockBackend) Predict(ctx context.Context, in *pb.PredictOptions) (*pb.R } func (m *MockBackend) PredictStream(in *pb.PredictOptions, stream pb.Backend_PredictStreamServer) error { + if err := checkModelIdentity(in); err != nil { + return err + } xlog.Debug("PredictStream called", "prompt", in.Prompt) if strings.Contains(in.Prompt, "MOCK_ERROR_IMMEDIATE") { return fmt.Errorf("mock backend stream error: simulated failure") @@ -391,6 +414,9 @@ func mockToolNameFromRequest(in *pb.PredictOptions) string { } func (m *MockBackend) Embedding(ctx context.Context, in *pb.PredictOptions) (*pb.EmbeddingResult, error) { + if err := checkModelIdentity(in); err != nil { + return nil, err + } xlog.Debug("Embedding called", "prompt", in.Prompt) // Return a mock embedding vector of 768 dimensions embeddings := make([]float32, 768) @@ -574,6 +600,9 @@ func (m *MockBackend) AudioTranscription(ctx context.Context, in *pb.TranscriptR } func (m *MockBackend) TokenizeString(ctx context.Context, in *pb.PredictOptions) (*pb.TokenizationResponse, error) { + if err := checkModelIdentity(in); err != nil { + return nil, err + } xlog.Debug("TokenizeString called", "prompt_len", len(in.Prompt)) // Approximate BPE: ~4 chars/token, minimum 1. Realistic enough for the // router's fitMessages to exercise the budget/rune-pretrim path with