mirror of
https://github.com/mudler/LocalAI.git
synced 2026-07-30 18:09:05 -04:00
fix(distributed): reject wrong-model requests at the backend (#10970)
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. Assisted-by: Claude Code:claude-opus-4-8[1m] [Read] [Edit] [Bash] Signed-off-by: Ettore Di Giacinto <mudler@localai.io> Co-authored-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
mudler's LocalAI [bot]
committed by
GitHub
GitHub
parent
1618c2e445
commit
465d488c90
@@ -11,6 +11,7 @@ import os
|
||||
|
||||
import grpc
|
||||
|
||||
from model_identity import AsyncModelIdentityInterceptor, ModelIdentityInterceptor
|
||||
from parent_watch import start_parent_death_watcher
|
||||
|
||||
|
||||
@@ -64,13 +65,15 @@ class AsyncTokenAuthInterceptor(grpc.aio.ServerInterceptor):
|
||||
|
||||
|
||||
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
|
||||
|
||||
192
backend/python/common/model_identity.py
Normal file
192
backend/python/common/model_identity.py
Normal file
@@ -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)
|
||||
222
backend/python/common/model_identity_test.py
Normal file
222
backend/python/common/model_identity_test.py
Normal file
@@ -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()
|
||||
Reference in New Issue
Block a user