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