mirror of
https://github.com/mudler/LocalAI.git
synced 2026-07-30 09:57:57 -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:
committed by
GitHub
parent
1618c2e445
commit
465d488c90
@@ -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 @@ public:
|
||||
}
|
||||
|
||||
result->set_success(true);
|
||||
g_loaded_model_identity = request->model();
|
||||
result->set_message("loaded " + model_path);
|
||||
return GStatus::OK;
|
||||
}
|
||||
@@ -724,6 +748,7 @@ public:
|
||||
backend::TokenizationResponse *response) override {
|
||||
std::lock_guard<std::mutex> 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 @@ public:
|
||||
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 @@ public:
|
||||
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);
|
||||
}
|
||||
|
||||
@@ -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 @@ public:
|
||||
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<backend::Reply>* 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 @@ public:
|
||||
|
||||
|
||||
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 @@ public:
|
||||
|
||||
/// 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 @@ public:
|
||||
}
|
||||
|
||||
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<llama_token> tokens = llama.tokenize(data["prompt"],false);
|
||||
|
||||
@@ -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 @@ public:
|
||||
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 @@ public:
|
||||
grpc::Status PredictStream(grpc::ServerContext* context, const backend::PredictOptions* request, grpc::ServerWriter<backend::Reply>* 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 @@ public:
|
||||
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 @@ public:
|
||||
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 @@ public:
|
||||
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");
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user