diff --git a/.agents/api-endpoints-and-auth.md b/.agents/api-endpoints-and-auth.md index fffee21c93f4..e3a63b1e5ced 100644 --- a/.agents/api-endpoints-and-auth.md +++ b/.agents/api-endpoints-and-auth.md @@ -236,6 +236,58 @@ Use these HTTP status codes: If your endpoint should be tracked for usage (token counts, request counts), add the `usageMiddleware` to its middleware chain. See `core/http/middleware/usage.go` and how it's applied in `routes/openai.go`. +## Control-plane database health metrics + +In distributed mode the frontend registers three OpenTelemetry gauges over the +PostgreSQL control-plane database (`core/services/monitoring/control_plane_db.go`, +wired in `core/application/distributed.go`). They reach `/metrics` through the +same Prometheus exporter as the rest of the API metrics. + +| Metric | Meaning | Page when | +|--------|---------|-----------| +| `localai_control_plane_oldest_xmin_age` | Transactions elapsed since the oldest snapshot any backend still holds | above a few million, and rising | +| `localai_control_plane_longest_transaction_seconds` | Age of the longest open transaction | above 3600 | +| `localai_control_plane_dead_tuple_ratio` | Dead tuples per live tuple, labelled by `table`, on `backend_nodes`, `node_models` and `gallery_operations` | sustained above ~10 on a small table | + +A sustained high `localai_control_plane_oldest_xmin_age` is the one to page on. +While it grows, autovacuum can reclaim nothing anywhere in the database no +matter how often it runs, so the dead tuple ratio keeps climbing and a six-row +registry table can reach hundreds of megabytes. Tuning autovacuum does not help. +The fix is to find the transaction holding the horizon open and clear it: + +```sql +SELECT pid, state, age(backend_xmin) AS xmin_age, now() - xact_start AS xact_age, query +FROM pg_stat_activity +WHERE backend_xmin IS NOT NULL +ORDER BY age(backend_xmin) DESC; +``` + +Then `pg_terminate_backend(pid)` on the offenders, and `VACUUM (VERBOSE)` the +bloated tables once the horizon has moved. + +**A healthy-looking xmin age does not on its own prove the horizon is free.** +The gauge reads `pg_stat_activity`, which only sees live backends. Two other +things pin the very same horizon and are invisible there, so either one can hold +vacuum back while the gauge reads 0: + +```sql +SELECT gid, prepared, database, transaction FROM pg_prepared_xacts; +SELECT slot_name, active, xmin, catalog_xmin FROM pg_replication_slots; +``` + +An orphaned prepared transaction is cleared with `ROLLBACK PREPARED ''`, +and a stale slot with `pg_drop_replication_slot('')`. Check both +before concluding that a bloated table has some other cause. + +Sampling is scrape-driven behind a 30 second cache, so scrape frequency does not +translate into database load. Failed and timed-out samples cost the same interval +as successful ones, so a database that is already struggling is not retried on +every scrape. A failed sample reports the last good values rather than failing the +scrape, because these gauges matter most when the database is struggling. Before +the first successful sample the gauges are absent rather than zero, since a zero +xmin age would read as a healthy horizon: alert on `absent()` too if you need to +distinguish "healthy" from "never sampled". + ## Advertising surfaces — where to register a new capability Beyond routing and auth, LocalAI publishes its capability surface in **four independent places**. When you add an endpoint — especially one introducing a net-new capability like a new media type or a new auth-gated feature — you must update every relevant surface. These aren't optional: missing them means the endpoint works but is invisible to clients, admins, and the UI. diff --git a/backend/cpp/audio-cpp/Makefile b/backend/cpp/audio-cpp/Makefile index 231ab85a10fe..4c836a0cb24d 100644 --- a/backend/cpp/audio-cpp/Makefile +++ b/backend/cpp/audio-cpp/Makefile @@ -9,7 +9,7 @@ # recipe is a make target (not a prepare.sh) so 'make purge && make' is a clean # rebuild and so the bump bot can see the pin. -AUDIO_CPP_VERSION?=3497b7cc44753e2c141d8fe60ac42cec433e3281 +AUDIO_CPP_VERSION?=f334cff70a68ea3d2e40d6638733e8c1ec434164 AUDIO_CPP_REPO?=https://github.com/0xShug0/audio.cpp CURRENT_MAKEFILE_DIR := $(dir $(abspath $(lastword $(MAKEFILE_LIST)))) diff --git a/backend/cpp/ds4/generation_limits.h b/backend/cpp/ds4/generation_limits.h new file mode 100644 index 000000000000..0985b059072f --- /dev/null +++ b/backend/cpp/ds4/generation_limits.h @@ -0,0 +1,27 @@ +// SPDX-License-Identifier: MIT +#pragma once + +#include + +namespace ds4cpp { + +inline int EffectiveGenerationLimit(int requested, int context_size, + int session_position) { + const int limit = requested > 0 ? requested : 256; + const int room = context_size - session_position; + if (room <= 1) return 0; + return std::min(limit, room - 1); +} + +inline int RemainingGenerationBudget(int effective_limit, int produced) { + if (effective_limit <= produced) return 0; + return effective_limit - produced; +} + +inline int SpeculativeAcceptedCapacity(int remaining, int draft_allowance, + int buffer_capacity) { + if (remaining <= 0 || draft_allowance < 0 || buffer_capacity <= 0) return 0; + return std::min({remaining, draft_allowance + 1, buffer_capacity}); +} + +} // namespace ds4cpp diff --git a/backend/cpp/ds4/generation_limits_test.cpp b/backend/cpp/ds4/generation_limits_test.cpp new file mode 100644 index 000000000000..ee38251f655e --- /dev/null +++ b/backend/cpp/ds4/generation_limits_test.cpp @@ -0,0 +1,92 @@ +// SPDX-License-Identifier: MIT + +#include "generation_limits.h" + +#include + +namespace { + +int failures = 0; + +void check_equal(int got, int want, const char *name) { + if (got == want) return; + std::fprintf(stderr, "FAIL %s: got %d, want %d\n", name, got, want); + failures++; +} + +// Mutation caught: treating omitted or negative max_tokens as unlimited instead +// of preserving DS4's legacy 256-token default. +void test_nonpositive_uses_legacy_default_when_space_permits() { + check_equal(ds4cpp::EffectiveGenerationLimit(0, 4096, 100), 256, + "zero max_tokens uses legacy default"); + check_equal(ds4cpp::EffectiveGenerationLimit(-1, 4096, 100), 256, + "negative max_tokens uses legacy default"); +} + +// Mutation caught: applying the legacy default without clamping it to the +// post-prefill context room and reserved slot. +void test_legacy_default_is_clamped_by_context() { + check_equal(ds4cpp::EffectiveGenerationLimit(0, 300, 100), 199, + "legacy default is context-clamped"); +} + +// Mutation caught: allowing an explicitly large request to overrun the +// post-prefill context boundary. +void test_large_positive_limit_is_clamped_to_context() { + check_equal(ds4cpp::EffectiveGenerationLimit(32768, 32768, 100), 32667, + "large positive is context-clamped"); +} + +// Mutation caught: replacing every positive request with the legacy default +// rather than preserving a smaller configured limit. +void test_smaller_positive_limit_is_preserved() { + check_equal(ds4cpp::EffectiveGenerationLimit(64, 4096, 100), 64, + "smaller positive is preserved"); +} + +// Mutation caught: consuming the final context slot instead of reserving it as +// required by DS4's generation loop. +void test_no_usable_room_returns_zero() { + check_equal(ds4cpp::EffectiveGenerationLimit(32, 100, 99), 0, + "one remaining context slot is not usable"); +} + +// Mutation caught: sending the original generation limit to a later +// speculative cycle instead of subtracting tokens already produced. +void test_remaining_budget_accounts_for_produced_tokens() { + check_equal(ds4cpp::RemainingGenerationBudget(10, 4), 6, + "remaining budget subtracts produced tokens"); + check_equal(ds4cpp::RemainingGenerationBudget(10, 12), 0, + "remaining budget never becomes negative"); +} + +// Mutation caught: giving speculative evaluation capacity beyond either the +// output budget, the draft allowance plus its first target token, or the fixed +// accepted-token buffer. +void test_speculative_capacity_obeys_all_bounds() { + check_equal(ds4cpp::SpeculativeAcceptedCapacity(3, 8, 8), 3, + "capacity respects remaining output budget"); + check_equal(ds4cpp::SpeculativeAcceptedCapacity(20, 4, 8), 5, + "capacity includes one target token beyond draft allowance"); + check_equal(ds4cpp::SpeculativeAcceptedCapacity(20, 8, 6), 6, + "capacity respects fixed buffer"); +} + +} // namespace + +int main() { + test_nonpositive_uses_legacy_default_when_space_permits(); + test_legacy_default_is_clamped_by_context(); + test_large_positive_limit_is_clamped_to_context(); + test_smaller_positive_limit_is_preserved(); + test_no_usable_room_returns_zero(); + test_remaining_budget_accounts_for_produced_tokens(); + test_speculative_capacity_obeys_all_bounds(); + + if (failures == 0) { + std::fprintf(stderr, "all generation limit checks passed\n"); + return 0; + } + std::fprintf(stderr, "%d check(s) failed\n", failures); + return 1; +} diff --git a/backend/cpp/ds4/grpc-server.cpp b/backend/cpp/ds4/grpc-server.cpp index 2118fd1cb270..68ebdd3e3551 100644 --- a/backend/cpp/ds4/grpc-server.cpp +++ b/backend/cpp/ds4/grpc-server.cpp @@ -10,6 +10,7 @@ #include "dsml_parser.h" // populated in Task 12 #include "dsml_renderer.h" // populated in Task 16 +#include "generation_limits.h" #include "kv_cache.h" // populated in Task 17 extern "C" { @@ -769,7 +770,6 @@ class DS4Backend final : public backend::Backend::Service { } ds4_tokens prompt = {}; build_prompt(g_engine, request, &prompt); - int n_predict = request->tokens() > 0 ? request->tokens() : 256; const bool think_enabled = ds4_think_mode_enabled(parse_think_mode(request)); const bool starts_in_thinking = think_enabled && @@ -792,6 +792,9 @@ class DS4Backend final : public backend::Backend::Service { int prompt_len = prompt.len; ds4_tokens_free(&prompt); if (rc == 0) { + const int n_predict = ds4cpp::EffectiveGenerationLimit( + request->tokens(), ds4_session_ctx(g_session), + ds4_session_pos(g_session)); const int eos = ds4_token_eos(g_engine); const int draft_max = ds4_engine_mtp_draft_tokens(g_engine); int produced = 0; @@ -810,9 +813,12 @@ class DS4Backend final : public backend::Backend::Service { if (draft_max > 0 && sp.temperature <= 0.0f) { constexpr int kAcceptedMax = 8; int accepted[kAcceptedMax]; - int cap = std::min(kAcceptedMax, draft_max + 1); + const int remaining = ds4cpp::RemainingGenerationBudget( + n_predict, produced); + const int cap = ds4cpp::SpeculativeAcceptedCapacity( + remaining, draft_max, kAcceptedMax); int n = ds4_session_eval_speculative_argmax( - g_session, first, draft_max, eos, + g_session, first, remaining, eos, accepted, cap, err, sizeof(err)); if (n < 0) { rc = -1; break; } bool stop = false; @@ -873,7 +879,6 @@ class DS4Backend final : public backend::Backend::Service { } ds4_tokens prompt = {}; build_prompt(g_engine, request, &prompt); - int n_predict = request->tokens() > 0 ? request->tokens() : 256; const bool think_enabled = ds4_think_mode_enabled(parse_think_mode(request)); const bool starts_in_thinking = think_enabled && @@ -891,6 +896,9 @@ class DS4Backend final : public backend::Backend::Service { int rc = ds4_session_sync(g_session, &prompt, err, sizeof(err)); ds4_tokens_free(&prompt); if (rc == 0) { + const int n_predict = ds4cpp::EffectiveGenerationLimit( + request->tokens(), ds4_session_ctx(g_session), + ds4_session_pos(g_session)); const int eos = ds4_token_eos(g_engine); const int draft_max = ds4_engine_mtp_draft_tokens(g_engine); int produced = 0; @@ -908,9 +916,12 @@ class DS4Backend final : public backend::Backend::Service { if (draft_max > 0 && sp.temperature <= 0.0f) { constexpr int kAcceptedMax = 8; int accepted[kAcceptedMax]; - int cap = std::min(kAcceptedMax, draft_max + 1); + const int remaining = ds4cpp::RemainingGenerationBudget( + n_predict, produced); + const int cap = ds4cpp::SpeculativeAcceptedCapacity( + remaining, draft_max, kAcceptedMax); int n = ds4_session_eval_speculative_argmax( - g_session, first, draft_max, eos, + g_session, first, remaining, eos, accepted, cap, err, sizeof(err)); if (n < 0) { rc = -1; break; } bool stop = false; diff --git a/backend/go/vllm-cpp/Makefile b/backend/go/vllm-cpp/Makefile index e73976023648..598d814c3624 100644 --- a/backend/go/vllm-cpp/Makefile +++ b/backend/go/vllm-cpp/Makefile @@ -11,7 +11,7 @@ JOBS?=$(shell nproc --ignore=1 2>/dev/null || sysctl -n hw.ncpu 2>/dev/null || e # vllm.cpp version VLLM_CPP_REPO?=https://github.com/mudler/vllm.cpp -VLLM_CPP_VERSION?=839ea1ceddb787778b6bd86a38a917a1aab74d8f +VLLM_CPP_VERSION?=6bf3abb580982f4fd2e4525ef37802ee0ce28981 # MLX GEMM provider (darwin/metal only; see the metal branch below for why). # Consumed as the prebuilt pip wheel: building MLX from source needs `xcrun diff --git a/core/application/distributed.go b/core/application/distributed.go index b7dc0bf91351..1fceee1317d2 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -15,6 +15,7 @@ import ( "github.com/mudler/LocalAI/core/services/distributed" "github.com/mudler/LocalAI/core/services/jobs" "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/monitoring" "github.com/mudler/LocalAI/core/services/nodes" "github.com/mudler/LocalAI/core/services/nodes/prefixcache" "github.com/mudler/LocalAI/core/services/storage" @@ -162,6 +163,17 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade } xlog.Info("Node registry initialized") + // Bound durable heartbeat writes: a beat that only carries a fresher + // timestamp is what turned backend_nodes into a 460 MB six-row table. + registry.SetHeartbeatCheckpoint(cfg.Distributed.NodeHeartbeatCheckpointOrDefault()) + + // Measure the vacuum horizon. The 42 days it stayed open went unnoticed + // because no gauge reported it until models started failing to load. + if err := monitoring.RegisterControlPlaneDBMetrics(authDB, 30*time.Second); err != nil { + // Metrics are diagnostic; a failure here must not stop the frontend. + xlog.Warn("Control-plane database metrics unavailable", "error", err) + } + // Let scheduling rules be keyed by a model alias. The registry resolves a // rule's name through the config loader to find the model it governs, so an // operator can pin placement to a stable name like "production" and have it diff --git a/core/cli/run.go b/core/cli/run.go index 6b9b3e4dc0b9..6bf377682e0e 100644 --- a/core/cli/run.go +++ b/core/cli/run.go @@ -182,6 +182,8 @@ type RunCMD struct { BackendUpgradeTimeout string `env:"LOCALAI_NATS_BACKEND_UPGRADE_TIMEOUT" help:"NATS round-trip timeout for backend.upgrade requests (default 15m)." group:"distributed"` ModelLoadTimeout string `env:"LOCALAI_NATS_MODEL_LOAD_TIMEOUT" help:"Fixed gRPC deadline for the remote LoadModel call sent to a worker node once its backend is installed and model files are staged. Unset (the default), the deadline is derived from the checkpoint size instead: 5m plus 20s per GiB, capped at 6h, so multi-tens-of-GB diffusion/video checkpoints get the minutes they need without a fixed cliff. Set this only to pin a specific budget; the value is used verbatim, including when it is shorter than the derived one." group:"distributed"` ModelLoadWait string `env:"LOCALAI_MODEL_LOAD_WAIT" help:"How long an inference request waits for a model that is still cold-loading onto a worker before it is answered with 503, a Retry-After header and live staging progress (default 60s). The request is served the moment the model becomes ready, so a model already most of the way staged needs no client retry. Set to 0 to wait as long as the load takes — only safe when no ingress or load balancer with an idle timeout sits in front." group:"distributed"` + StaleNodeThreshold string `env:"LOCALAI_STALE_NODE_THRESHOLD" help:"How long a worker node may go without a durable heartbeat before the health monitor marks it offline (default 5m). Because a beat that only carries a fresher timestamp is held back by --node-heartbeat-checkpoint, this must stay comfortably wider than that interval; raise both together. Dead-node detection through the per-model gRPC health check and through request-time failure is unaffected by this knob." group:"distributed"` + NodeHeartbeatCheckpoint string `env:"LOCALAI_NODE_HEARTBEAT_CHECKPOINT" help:"Minimum gap between durable heartbeat writes for a worker node (default 60s). A beat that only carries a fresher timestamp is dropped until this interval elapses; every field is compared against the value last written, so a node's first beat, a changed total VRAM/total disk/GPU vendor, and a free VRAM/RAM/disk reading that has moved more than 256 MiB from the written value all still write immediately, and a node that is not active is never suppressed. Set below the worker heartbeat interval to write on every beat." group:"distributed"` NatsAccountSeed string `env:"LOCALAI_NATS_ACCOUNT_SEED" help:"NATS account signing seed (SU...) used to mint per-node worker JWTs at registration" group:"distributed"` NatsServiceJWT string `env:"LOCALAI_NATS_SERVICE_JWT" help:"NATS user JWT for the frontend (and agent workers) to publish control-plane messages" group:"distributed"` NatsServiceSeed string `env:"LOCALAI_NATS_SERVICE_SEED" help:"NATS user signing seed (SU...) paired with LOCALAI_NATS_SERVICE_JWT" group:"distributed"` @@ -397,6 +399,20 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error { } opts = append(opts, config.WithModelLoadWait(d)) } + if r.StaleNodeThreshold != "" { + d, err := parseDistributedDuration("LOCALAI_STALE_NODE_THRESHOLD", r.StaleNodeThreshold) + if err != nil { + return err + } + opts = append(opts, config.WithStaleNodeThreshold(d)) + } + if r.NodeHeartbeatCheckpoint != "" { + d, err := parseDistributedDuration("LOCALAI_NODE_HEARTBEAT_CHECKPOINT", r.NodeHeartbeatCheckpoint) + if err != nil { + return err + } + opts = append(opts, config.WithNodeHeartbeatCheckpoint(d)) + } if r.RegistrationToken != "" { opts = append(opts, config.WithRegistrationToken(r.RegistrationToken)) } diff --git a/core/config/distributed_config.go b/core/config/distributed_config.go index 5a48a84e9b44..ef7e01bfd18f 100644 --- a/core/config/distributed_config.go +++ b/core/config/distributed_config.go @@ -57,12 +57,13 @@ type DistributedConfig struct { StorageSecretKey string // --storage-secret-key / LOCALAI_STORAGE_SECRET_KEY // Timeout configuration (all have sensible defaults — zero means use default) - MCPToolTimeout time.Duration // MCP tool execution timeout (default 360s) - MCPDiscoveryTimeout time.Duration // MCP discovery timeout (default 60s) - WorkerWaitTimeout time.Duration // Max wait for healthy worker at startup (default 5m) - DrainTimeout time.Duration // Time to wait for in-flight requests during drain (default 30s) - HealthCheckInterval time.Duration // Health monitor check interval (default 15s) - StaleNodeThreshold time.Duration // Time before a node is considered stale (default 60s) + MCPToolTimeout time.Duration // MCP tool execution timeout (default 360s) + MCPDiscoveryTimeout time.Duration // MCP discovery timeout (default 60s) + WorkerWaitTimeout time.Duration // Max wait for healthy worker at startup (default 5m) + DrainTimeout time.Duration // Time to wait for in-flight requests during drain (default 30s) + HealthCheckInterval time.Duration // Health monitor check interval (default 15s) + StaleNodeThreshold time.Duration // Time before a node is considered stale (default 5m) + NodeHeartbeatCheckpoint time.Duration // Minimum gap between durable heartbeat writes (default 60s, 0 = every beat) // DisablePerModelHealthCheck turns off the health monitor's per-model // gRPC probe. When enabled (the default), the monitor pings each model's // gRPC address and removes stale node_models rows whose backend has @@ -165,16 +166,17 @@ func (c DistributedConfig) Validate() error { c.NatsAuthConfig().WarnIfInsecure(true) // Check for negative durations for name, d := range map[string]time.Duration{ - FlagMCPToolTimeout: c.MCPToolTimeout, - FlagMCPDiscoveryTimeout: c.MCPDiscoveryTimeout, - FlagWorkerWaitTimeout: c.WorkerWaitTimeout, - FlagDrainTimeout: c.DrainTimeout, - FlagHealthCheckInterval: c.HealthCheckInterval, - FlagStaleNodeThreshold: c.StaleNodeThreshold, - FlagMCPCIJobTimeout: c.MCPCIJobTimeout, - FlagBackendInstallTimeout: c.BackendInstallTimeout, - FlagBackendUpgradeTimeout: c.BackendUpgradeTimeout, - FlagModelLoadTimeout: c.ModelLoadTimeout, + FlagMCPToolTimeout: c.MCPToolTimeout, + FlagMCPDiscoveryTimeout: c.MCPDiscoveryTimeout, + FlagWorkerWaitTimeout: c.WorkerWaitTimeout, + FlagDrainTimeout: c.DrainTimeout, + FlagHealthCheckInterval: c.HealthCheckInterval, + FlagStaleNodeThreshold: c.StaleNodeThreshold, + FlagNodeHeartbeatCheckpoint: c.NodeHeartbeatCheckpoint, + FlagMCPCIJobTimeout: c.MCPCIJobTimeout, + FlagBackendInstallTimeout: c.BackendInstallTimeout, + FlagBackendUpgradeTimeout: c.BackendUpgradeTimeout, + FlagModelLoadTimeout: c.ModelLoadTimeout, } { if d < 0 { return fmt.Errorf("%s must not be negative", name) @@ -337,6 +339,27 @@ func WithModelLoadWait(d time.Duration) AppOption { } } +// WithStaleNodeThreshold sets how long a node may go without a durable +// heartbeat before the health monitor marks it offline. It has to be raised +// alongside WithNodeHeartbeatCheckpoint: a checkpoint interval wider than this +// threshold makes every healthy node look dead the moment its beats start +// being suppressed. +func WithStaleNodeThreshold(d time.Duration) AppOption { + return func(o *ApplicationConfig) { + o.Distributed.StaleNodeThreshold = d + } +} + +// WithNodeHeartbeatCheckpoint bounds durable heartbeat writes. A zero d is +// deliberately not special-cased into "unbounded": NodeHeartbeatCheckpointOrDefault +// reads zero as unset, and an operator who wants a write per beat sets a value +// below the worker's heartbeat interval instead. +func WithNodeHeartbeatCheckpoint(d time.Duration) AppOption { + return func(o *ApplicationConfig) { + o.Distributed.NodeHeartbeatCheckpoint = d + } +} + var EnableAutoApproveNodes = func(o *ApplicationConfig) { o.Distributed.AutoApproveNodes = true } @@ -391,17 +414,18 @@ func WithModelSchedulingConfigPath(path string) AppOption { // them as constants prevents the string from drifting from the actual // flag a future rename would produce. const ( - FlagMCPToolTimeout = "mcp-tool-timeout" - FlagMCPDiscoveryTimeout = "mcp-discovery-timeout" - FlagWorkerWaitTimeout = "worker-wait-timeout" - FlagDrainTimeout = "drain-timeout" - FlagHealthCheckInterval = "health-check-interval" - FlagStaleNodeThreshold = "stale-node-threshold" - FlagMCPCIJobTimeout = "mcp-ci-job-timeout" - FlagBackendInstallTimeout = "backend-install-timeout" - FlagBackendUpgradeTimeout = "backend-upgrade-timeout" - FlagModelLoadTimeout = "model-load-timeout" - FlagModelLoadWait = "model-load-wait" + FlagMCPToolTimeout = "mcp-tool-timeout" + FlagMCPDiscoveryTimeout = "mcp-discovery-timeout" + FlagWorkerWaitTimeout = "worker-wait-timeout" + FlagDrainTimeout = "drain-timeout" + FlagHealthCheckInterval = "health-check-interval" + FlagStaleNodeThreshold = "stale-node-threshold" + FlagNodeHeartbeatCheckpoint = "node-heartbeat-checkpoint" + FlagMCPCIJobTimeout = "mcp-ci-job-timeout" + FlagBackendInstallTimeout = "backend-install-timeout" + FlagBackendUpgradeTimeout = "backend-upgrade-timeout" + FlagModelLoadTimeout = "model-load-timeout" + FlagModelLoadWait = "model-load-wait" // FlagDiskHeadroomCheck names the disk-headroom toggle. It is quoted in // the warning the check emits while disabled, so the operator reading a // log line knows exactly which knob produced it. @@ -410,16 +434,22 @@ const ( // Defaults for distributed timeouts. const ( - DefaultMCPToolTimeout = 360 * time.Second - DefaultMCPDiscoveryTimeout = 60 * time.Second - DefaultWorkerWaitTimeout = 5 * time.Minute - DefaultDrainTimeout = 30 * time.Second - DefaultHealthCheckInterval = 15 * time.Second - DefaultStaleNodeThreshold = 60 * time.Second - DefaultMCPCIJobTimeout = 10 * time.Minute - DefaultBackendInstallTimeout = 15 * time.Minute - DefaultBackendUpgradeTimeout = 15 * time.Minute - DefaultModelLoadTimeout = 5 * time.Minute + DefaultMCPToolTimeout = 360 * time.Second + DefaultMCPDiscoveryTimeout = 60 * time.Second + DefaultWorkerWaitTimeout = 5 * time.Minute + DefaultDrainTimeout = 30 * time.Second + DefaultHealthCheckInterval = 15 * time.Second + // A beat that only refreshes the timestamp is now dropped until the + // checkpoint interval elapses, so the persisted column is up to one + // interval stale by design. The threshold covers that plus jitter. + // A genuinely dead node is still caught sooner by the per-model gRPC + // health check and by request-time failure, neither of which reads this. + DefaultStaleNodeThreshold = 5 * time.Minute + DefaultNodeHeartbeatCheckpoint = 60 * time.Second + DefaultMCPCIJobTimeout = 10 * time.Minute + DefaultBackendInstallTimeout = 15 * time.Minute + DefaultBackendUpgradeTimeout = 15 * time.Minute + DefaultModelLoadTimeout = 5 * time.Minute // DefaultModelLoadWait is how long a request waits for a cold-loading model // before it is answered with 503 and live progress. Chosen to sit under the // idle timeout of typical ingress/LB defaults, so the answer comes from @@ -519,6 +549,14 @@ func (c DistributedConfig) StaleNodeThresholdOrDefault() time.Duration { return cmp.Or(c.StaleNodeThreshold, DefaultStaleNodeThreshold) } +// NodeHeartbeatCheckpointOrDefault returns the configured interval or the +// default. A configured zero is indistinguishable from unset here, which is +// intentional: cmp.Or falls back to the default, and an operator who wants a +// write per beat sets a value below the heartbeat interval instead. +func (c DistributedConfig) NodeHeartbeatCheckpointOrDefault() time.Duration { + return cmp.Or(c.NodeHeartbeatCheckpoint, DefaultNodeHeartbeatCheckpoint) +} + // MCPCIJobTimeoutOrDefault returns the configured MCP CI job timeout or the default. func (c DistributedConfig) MCPCIJobTimeoutOrDefault() time.Duration { return cmp.Or(c.MCPCIJobTimeout, DefaultMCPCIJobTimeout) diff --git a/core/config/distributed_config_test.go b/core/config/distributed_config_test.go index ec7fbe8dc7e9..dc98c09bb4e9 100644 --- a/core/config/distributed_config_test.go +++ b/core/config/distributed_config_test.go @@ -47,6 +47,27 @@ var _ = Describe("DistributedConfig backend NATS timeouts", func() { }) }) +// Heartbeat checkpointing makes last_heartbeat up to one checkpoint interval +// stale by design, which is why the threshold defaults to 5 minutes. An +// operator who widens the checkpoint has to widen this to match, so it has to +// be reachable from the CLI rather than being a compile-time constant. +var _ = Describe("DistributedConfig stale node threshold", func() { + It("defaults to 5 minutes, wide enough to cover a suppressed beat", func() { + Expect(config.DistributedConfig{}.StaleNodeThresholdOrDefault()). + To(Equal(5 * time.Minute)) + Expect(config.DefaultStaleNodeThreshold). + To(BeNumerically(">", config.DefaultNodeHeartbeatCheckpoint), + "a threshold at or below the checkpoint interval marks healthy, "+ + "beating nodes offline every cycle") + }) + + It("is configurable, so a widened checkpoint can be matched", func() { + o := config.NewApplicationConfig(config.WithStaleNodeThreshold(20 * time.Minute)) + Expect(o.Distributed.StaleNodeThreshold).To(Equal(20 * time.Minute)) + Expect(o.Distributed.StaleNodeThresholdOrDefault()).To(Equal(20 * time.Minute)) + }) +}) + var _ = Describe("DistributedConfig flag-name constants", func() { // Pin the kebab-case strings so a rename of the Go field name (or a // CLI flag naming convention change) forces the constant to update, @@ -62,6 +83,7 @@ var _ = Describe("DistributedConfig flag-name constants", func() { Entry("drain timeout", config.FlagDrainTimeout, "drain-timeout"), Entry("health check interval", config.FlagHealthCheckInterval, "health-check-interval"), Entry("stale node threshold", config.FlagStaleNodeThreshold, "stale-node-threshold"), + Entry("node heartbeat checkpoint", config.FlagNodeHeartbeatCheckpoint, "node-heartbeat-checkpoint"), Entry("MCP CI job timeout", config.FlagMCPCIJobTimeout, "mcp-ci-job-timeout"), Entry("backend install timeout", config.FlagBackendInstallTimeout, "backend-install-timeout"), Entry("backend upgrade timeout", config.FlagBackendUpgradeTimeout, "backend-upgrade-timeout"), diff --git a/core/http/endpoints/openai/chat.go b/core/http/endpoints/openai/chat.go index fcec18932a6a..72fe59b0e536 100644 --- a/core/http/endpoints/openai/chat.go +++ b/core/http/endpoints/openai/chat.go @@ -4,6 +4,7 @@ import ( "encoding/json" "fmt" "net/http" + "strings" "time" "github.com/google/uuid" @@ -24,18 +25,47 @@ import ( "github.com/mudler/xlog" ) +// messageText returns the textual content of a message, preferring the +// middleware-populated StringContent and falling back to a string Content. +func messageText(m schema.Message) string { + if m.StringContent != "" { + return m.StringContent + } + if s, ok := m.Content.(string); ok { + return s + } + return "" +} + // hasSystemMessage reports whether the message slice already contains a -// system-role message — used to avoid clobbering a caller-supplied system -// prompt when the LocalAI Assistant modality is on. +// non-empty system-role message — used to avoid clobbering a caller-supplied +// system prompt when the LocalAI Assistant modality is on. Empty / whitespace +// system turns (historically sent by the web Chat UI) are ignored so they do +// not suppress the model config system_prompt. func hasSystemMessage(messages []schema.Message) bool { for _, m := range messages { - if m.Role == "system" { + if m.Role == "system" && strings.TrimSpace(messageText(m)) != "" { return true } } return false } +// stripEmptySystemMessages drops system-role messages whose content is empty +// or whitespace-only. An explicit blank system turn would otherwise satisfy +// tokenizer chat templates' `messages[0].role == "system"` check and suppress +// both the model's configured system_prompt and any template default. +func stripEmptySystemMessages(messages []schema.Message) []schema.Message { + out := messages[:0:0] + for _, m := range messages { + if m.Role == "system" && strings.TrimSpace(messageText(m)) == "" { + continue + } + out = append(out, m) + } + return out +} + // mergeToolCallDeltas merges streaming tool call deltas into complete tool calls. // In SSE streaming, a single tool call arrives as multiple chunks sharing the same Index: // the first chunk carries the ID, Type, and Name; subsequent chunks append to Arguments. @@ -149,6 +179,18 @@ func ChatEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator xlog.Debug("Chat endpoint configuration read", "config", config) + // Drop blank system turns from the web UI (and similar clients) so they + // cannot suppress the model YAML system_prompt / tokenizer defaults. + input.Messages = stripEmptySystemMessages(input.Messages) + + // Tokenizer-template models pass messages through to the backend as-is, + // so apply the configured system_prompt when the request did not supply + // one. Go-template models already receive SystemPrompt via PromptTemplateData. + if config.TemplateConfig.UseTokenizerTemplate && config.SystemPrompt != "" && !hasSystemMessage(input.Messages) { + prompt := config.SystemPrompt + input.Messages = append([]schema.Message{{Role: "system", Content: prompt, StringContent: prompt}}, input.Messages...) + } + // Cloud-proxy bail. Bypasses the local pipeline (templating, // MCP injection, gRPC backend) and forwards via the cloud- // proxy backend, which does the outbound HTTP. Request-side PII diff --git a/core/http/endpoints/openai/chat_test.go b/core/http/endpoints/openai/chat_test.go index ffb3086fbdb2..0ef0991dd1c4 100644 --- a/core/http/endpoints/openai/chat_test.go +++ b/core/http/endpoints/openai/chat_test.go @@ -357,3 +357,47 @@ var _ = Describe("mergeToolCallDeltas", func() { }) }) }) + +var _ = Describe("system message helpers", func() { + Describe("hasSystemMessage", func() { + It("ignores empty and whitespace-only system turns", func() { + Expect(hasSystemMessage([]schema.Message{ + {Role: "system", Content: "", StringContent: ""}, + {Role: "user", Content: "hi", StringContent: "hi"}, + })).To(BeFalse()) + Expect(hasSystemMessage([]schema.Message{ + {Role: "system", Content: " ", StringContent: " "}, + })).To(BeFalse()) + }) + + It("detects a real system prompt", func() { + Expect(hasSystemMessage([]schema.Message{ + {Role: "system", Content: "You are helpful.", StringContent: "You are helpful."}, + {Role: "user", Content: "hi", StringContent: "hi"}, + })).To(BeTrue()) + }) + }) + + Describe("stripEmptySystemMessages", func() { + It("removes blank system turns and keeps the rest", func() { + in := []schema.Message{ + {Role: "system", Content: "", StringContent: ""}, + {Role: "system", Content: " ", StringContent: " "}, + {Role: "user", Content: "Explain how this works", StringContent: "Explain how this works"}, + } + out := stripEmptySystemMessages(in) + Expect(out).To(HaveLen(1)) + Expect(out[0].Role).To(Equal("user")) + }) + + It("keeps a non-empty system turn", func() { + in := []schema.Message{ + {Role: "system", Content: "You are LocalAI.", StringContent: "You are LocalAI."}, + {Role: "user", Content: "hi", StringContent: "hi"}, + } + out := stripEmptySystemMessages(in) + Expect(out).To(HaveLen(2)) + Expect(out[0].StringContent).To(Equal("You are LocalAI.")) + }) + }) +}) diff --git a/core/http/react-ui/src/hooks/useChat.js b/core/http/react-ui/src/hooks/useChat.js index 567cf7f55cfa..16f99f64932b 100644 --- a/core/http/react-ui/src/hooks/useChat.js +++ b/core/http/react-ui/src/hooks/useChat.js @@ -1,6 +1,7 @@ import { useState, useCallback, useRef } from 'react' import { API_CONFIG } from '../utils/config' import { apiUrl } from '../utils/basePath' +import { effectiveSystemPrompt } from '../utils/systemPrompt' import { useDebouncedEffect } from './useDebounce' const thinkingTagRegex = /([\s\S]*?)<\/thinking>|([\s\S]*?)<\/think>|<\|channel>thought([\s\S]*?)/g @@ -348,8 +349,12 @@ export function useChat(initialModel = '') { // Build messages array for API const chat = chats.find(c => c.id === chatId) const messages = [] - if (chat?.systemPrompt) { - messages.push({ role: 'system', content: chat.systemPrompt }) + // Omit empty/whitespace system prompts so the model YAML system_prompt + // (and tokenizer chat-template defaults) are not suppressed by a blank + // system turn from Chat Settings. + const systemPrompt = effectiveSystemPrompt(chat?.systemPrompt) + if (systemPrompt) { + messages.push({ role: 'system', content: systemPrompt }) } // Filter out thinking/reasoning/tool_call/tool_result messages. // options.baseHistory lets callers (e.g. mid-conversation retry) pass the @@ -358,6 +363,7 @@ export function useChat(initialModel = '') { const baseHistory = options.baseHistory || chat?.history || [] const historyForApi = baseHistory.filter(m => m.role !== 'thinking' && m.role !== 'reasoning' && m.role !== 'tool_call' && m.role !== 'tool_result' + && !(m.role === 'system' && !effectiveSystemPrompt(typeof m.content === 'string' ? m.content : '')) ) messages.push(...historyForApi, { role: 'user', content: messageContent }) diff --git a/core/http/react-ui/src/utils/systemPrompt.js b/core/http/react-ui/src/utils/systemPrompt.js new file mode 100644 index 000000000000..541e72a5b7a6 --- /dev/null +++ b/core/http/react-ui/src/utils/systemPrompt.js @@ -0,0 +1,17 @@ +/** + * Normalize a chat UI system prompt for the /v1/chat/completions payload. + * Empty / whitespace-only values must be omitted so the backend can apply the + * model's configured system_prompt (and tokenizer chat templates are not fed + * an explicit blank system turn). + */ +export function effectiveSystemPrompt(prompt) { + if (typeof prompt !== 'string') return '' + return prompt.trim() +} + +/** + * True when the UI should include a system message in the request. + */ +export function shouldSendSystemPrompt(prompt) { + return effectiveSystemPrompt(prompt) !== '' +} diff --git a/core/http/react-ui/src/utils/systemPrompt.test.js b/core/http/react-ui/src/utils/systemPrompt.test.js new file mode 100644 index 000000000000..8f9b80f712d6 --- /dev/null +++ b/core/http/react-ui/src/utils/systemPrompt.test.js @@ -0,0 +1,25 @@ +import assert from 'node:assert/strict' +import test from 'node:test' + +import { effectiveSystemPrompt, shouldSendSystemPrompt } from './systemPrompt.js' + +test('empty string is omitted', () => { + assert.equal(effectiveSystemPrompt(''), '') + assert.equal(shouldSendSystemPrompt(''), false) +}) + +test('whitespace-only is omitted', () => { + assert.equal(effectiveSystemPrompt(' \n\t '), '') + assert.equal(shouldSendSystemPrompt(' \t'), false) +}) + +test('non-empty prompt is kept trimmed', () => { + assert.equal(effectiveSystemPrompt(' You are helpful. '), 'You are helpful.') + assert.equal(shouldSendSystemPrompt('You are helpful.'), true) +}) + +test('non-strings are treated as empty', () => { + assert.equal(effectiveSystemPrompt(null), '') + assert.equal(effectiveSystemPrompt(undefined), '') + assert.equal(shouldSendSystemPrompt(0), false) +}) diff --git a/core/http/static/chat.js b/core/http/static/chat.js index 59add0f09136..34d6c0395efa 100644 --- a/core/http/static/chat.js +++ b/core/http/static/chat.js @@ -1149,13 +1149,21 @@ async function promptGPT(systemPrompt, input) { messages = chatStore.messages(); // Exclude thinking/reasoning from API payload (backend chat templates expect only system/user/assistant) - messages = messages.filter((m) => m.role !== "thinking" && m.role !== "reasoning"); + // Also drop blank system turns so they cannot suppress the model default. + messages = messages.filter((m) => + m.role !== "thinking" && + m.role !== "reasoning" && + !(m.role === "system" && !(typeof m.content === "string" && m.content.trim())) + ); - // if systemPrompt isn't empty, push it at the start of messages - if (systemPrompt) { + // Omit empty/whitespace system prompts so the model YAML system_prompt + // (and tokenizer chat-template defaults) are not suppressed by a blank + // system turn from Chat Settings. + const trimmedSystemPrompt = typeof systemPrompt === "string" ? systemPrompt.trim() : ""; + if (trimmedSystemPrompt) { messages.unshift({ role: "system", - content: systemPrompt + content: trimmedSystemPrompt }); } @@ -2275,7 +2283,8 @@ storesystemPrompt = localStorage.getItem("system_prompt"); if (storesystemPrompt) { document.getElementById("systemPrompt").value = storesystemPrompt; } else { - document.getElementById("systemPrompt").value = null; + // Use "" — assigning null stringifies to "null" on textarea.value. + document.getElementById("systemPrompt").value = ""; } marked.setOptions({ diff --git a/core/services/monitoring/control_plane_db.go b/core/services/monitoring/control_plane_db.go new file mode 100644 index 000000000000..1f420ca98b6b --- /dev/null +++ b/core/services/monitoring/control_plane_db.go @@ -0,0 +1,204 @@ +package monitoring + +import ( + "context" + "fmt" + "sync" + "time" + + "github.com/mudler/LocalAI/core/services/distributed" + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/xlog" + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/metric" + "gorm.io/gorm" +) + +// controlPlaneModels are the registry records whose table bloat degrades +// routing. The table NAMES are resolved from gorm rather than written out +// here, because the three do not agree on how they get one: BackendNode and +// NodeModel take gorm's default pluralisation, while GalleryOperationRecord +// overrides TableName. A literal list would keep compiling and silently stop +// matching the day any of them gains or changes a TableName, and a gauge that +// reads zero because it matched no rows looks exactly like a healthy cluster. +var controlPlaneModels = []any{ + &nodes.BackendNode{}, + &nodes.NodeModel{}, + &distributed.GalleryOperationRecord{}, +} + +// controlPlaneTableNames asks gorm what each model is actually stored as. +func controlPlaneTableNames(db *gorm.DB) ([]string, error) { + names := make([]string, 0, len(controlPlaneModels)) + for _, model := range controlPlaneModels { + stmt := &gorm.Statement{DB: db} + if err := stmt.Parse(model); err != nil { + return nil, fmt.Errorf("resolving control-plane table name for %T: %w", model, err) + } + names = append(names, stmt.Table) + } + return names, nil +} + +// controlPlaneSampleTimeout bounds one catalog sample. A scrape drives the +// collection, so an unbounded query on a wedged database would hold the +// Prometheus scrape open for as long as the database stays wedged, which is +// precisely the situation these gauges exist to report. +const controlPlaneSampleTimeout = 5 * time.Second + +// ControlPlaneDBStats are the PostgreSQL numbers that predict a control-plane +// outage before it is visible as failed model loads. +type ControlPlaneDBStats struct { + // OldestXminAge is transactions elapsed since the oldest snapshot any + // backend still holds. While it grows, autovacuum can reclaim nothing in + // the database, however often it runs. It was 21,002,291 during the + // incident that motivated this. + OldestXminAge int64 + // LongestTransactionSeconds is the age of the longest open transaction. + LongestTransactionSeconds float64 + // DeadTupleRatio is dead tuples per live tuple, per control-plane table. + DeadTupleRatio map[string]float64 +} + +// SampleControlPlaneDB reads the stats in two cheap catalog queries. +func SampleControlPlaneDB(ctx context.Context, db *gorm.DB) (ControlPlaneDBStats, error) { + stats := ControlPlaneDBStats{DeadTupleRatio: map[string]float64{}} + + var horizon struct { + XminAge int64 + LongestS float64 + } + if err := db.WithContext(ctx).Raw(` + SELECT + COALESCE(MAX(age(backend_xmin)), 0) AS xmin_age, + COALESCE(MAX(EXTRACT(EPOCH FROM now() - xact_start)), 0) AS longest_s + FROM pg_stat_activity + WHERE datname = current_database()`).Scan(&horizon).Error; err != nil { + return stats, err + } + stats.OldestXminAge = horizon.XminAge + stats.LongestTransactionSeconds = horizon.LongestS + + var rows []struct { + Relname string + DeadTup int64 + LiveTup int64 + } + tables, err := controlPlaneTableNames(db) + if err != nil { + return stats, err + } + if err := db.WithContext(ctx).Raw(` + SELECT relname, n_dead_tup AS dead_tup, n_live_tup AS live_tup + FROM pg_stat_user_tables + WHERE relname IN ?`, tables).Scan(&rows).Error; err != nil { + return stats, err + } + for _, r := range rows { + live := r.LiveTup + if live < 1 { + live = 1 + } + stats.DeadTupleRatio[r.Relname] = float64(r.DeadTup) / float64(live) + } + return stats, nil +} + +// cachedDBSampler serves the catalog sample to scrape-driven collection. The +// cache bounds how often a scrape can reach the database, whether or not the +// sample succeeds, and it keeps the last good values available when a later +// sample fails. +type cachedDBSampler struct { + db *gorm.DB + minInterval time.Duration + + mu sync.Mutex + cached ControlPlaneDBStats + // hasSample records whether cached ever held a real reading. It is kept + // apart from lastAttempt so that a failing database does not blank the + // gauges, and so that a failed attempt still costs the interval. + hasSample bool + // lastAttempt times every attempt, successful or not. Gating on the last + // success instead would leave a failing database open to one query per + // scrape, which is a retry storm aimed at a database already in trouble. + lastAttempt time.Time +} + +// stats returns the values to report and whether any good sample exists yet. +// A failed sample is not fatal: these gauges matter most when the database is +// struggling, which is exactly when this query can fail, so the last good +// values keep being reported instead of failing the whole scrape. Before the +// first good sample there is nothing honest to report, and reporting zero +// would read as a healthy horizon, so the caller reports nothing at all. +func (c *cachedDBSampler) stats(ctx context.Context) (ControlPlaneDBStats, bool) { + c.mu.Lock() + defer c.mu.Unlock() + + if c.lastAttempt.IsZero() || time.Since(c.lastAttempt) >= c.minInterval { + sampleCtx, cancel := context.WithTimeout(ctx, controlPlaneSampleTimeout) + s, err := SampleControlPlaneDB(sampleCtx, c.db) + cancel() + // Timed from completion, and recorded whatever the outcome, so that a + // slow or failing sample is rate limited exactly like a good one. + c.lastAttempt = time.Now() + if err != nil { + xlog.Debug("Control-plane database sample failed, reporting the last good values", "error", err) + } else { + c.cached, c.hasSample = s, true + } + } + + return c.cached, c.hasSample +} + +// RegisterControlPlaneDBMetrics installs observable gauges backed by a cached +// sample. Collection is scrape-driven, so the cache bounds how often a scrape +// can reach the database. +func RegisterControlPlaneDBMetrics(db *gorm.DB, minInterval time.Duration) error { + // ORDERING DEPENDENCY, deliberately recorded because nothing enforces it: + // this reads the GLOBAL meter provider, so it must run after + // monitoring.NewLocalAIMetricsService has called otel.SetMeterProvider in + // core/application/startup.go. Called before that, the gauges bind to the + // no-op global and never reach /metrics, silently. Today the distributed + // wiring in core/application/distributed.go runs after that point, which is + // why the global is safe here rather than injected the way billing, pii and + // agentpool take an explicit meter. This repo has been bitten by that race + // three times already: see the comments at core/application/startup.go, + // core/http/app.go and core/application/application.go. Anyone moving this + // call earlier must switch it to an injected meter instead. + meter := otel.Meter("github.com/mudler/LocalAI") + + xminAge, err := meter.Int64ObservableGauge("localai_control_plane_oldest_xmin_age", + metric.WithDescription("Transactions elapsed since the oldest snapshot still held. While this grows, autovacuum can reclaim nothing.")) + if err != nil { + return err + } + longest, err := meter.Float64ObservableGauge("localai_control_plane_longest_transaction_seconds", + metric.WithDescription("Age of the longest open transaction, in seconds.")) + if err != nil { + return err + } + deadRatio, err := meter.Float64ObservableGauge("localai_control_plane_dead_tuple_ratio", + metric.WithDescription("Dead tuples per live tuple on the control-plane registry tables.")) + if err != nil { + return err + } + + sampler := &cachedDBSampler{db: db, minInterval: minInterval} + + _, err = meter.RegisterCallback(func(ctx context.Context, o metric.Observer) error { + cached, ok := sampler.stats(ctx) + if !ok { + return nil + } + + o.ObserveInt64(xminAge, cached.OldestXminAge) + o.ObserveFloat64(longest, cached.LongestTransactionSeconds) + for table, ratio := range cached.DeadTupleRatio { + o.ObserveFloat64(deadRatio, ratio, metric.WithAttributes(attribute.String("table", table))) + } + return nil + }, xminAge, longest, deadRatio) + return err +} diff --git a/core/services/monitoring/control_plane_db_test.go b/core/services/monitoring/control_plane_db_test.go new file mode 100644 index 000000000000..0aca60c174ce --- /dev/null +++ b/core/services/monitoring/control_plane_db_test.go @@ -0,0 +1,286 @@ +package monitoring + +import ( + "context" + "runtime" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "go.opentelemetry.io/otel" + sdkmetric "go.opentelemetry.io/otel/sdk/metric" + "go.opentelemetry.io/otel/sdk/metric/metricdata" + "gorm.io/gorm" + + "github.com/mudler/LocalAI/core/services/testutil" +) + +// collectedMetricNames returns the names of every metric the reader collected, +// so a spec can assert that a gauge is present or absent from a scrape. +func collectedMetricNames(reader sdkmetric.Reader) []string { + var rm metricdata.ResourceMetrics + ExpectWithOffset(1, reader.Collect(context.Background(), &rm)).To(Succeed()) + names := []string{} + for _, sm := range rm.ScopeMetrics { + for _, m := range sm.Metrics { + names = append(names, m.Name) + } + } + return names +} + +// collectedInt64Gauge returns the single data point of an int64 gauge. +func collectedInt64Gauge(reader sdkmetric.Reader, name string) int64 { + var rm metricdata.ResourceMetrics + ExpectWithOffset(1, reader.Collect(context.Background(), &rm)).To(Succeed()) + for _, sm := range rm.ScopeMetrics { + for _, m := range sm.Metrics { + if m.Name != name { + continue + } + gauge, ok := m.Data.(metricdata.Gauge[int64]) + ExpectWithOffset(1, ok).To(BeTrue(), "%s is not an int64 gauge", name) + ExpectWithOffset(1, gauge.DataPoints).To(HaveLen(1)) + return gauge.DataPoints[0].Value + } + } + Fail("gauge "+name+" was not collected", 1) + return 0 +} + +// closeDB drops the connection pool, which is the closest a test can get to +// the database being unreachable mid-scrape. +func closeDB(db *gorm.DB) { + sqlDB, err := db.DB() + ExpectWithOffset(1, err).ToNot(HaveOccurred()) + ExpectWithOffset(1, sqlDB.Close()).To(Succeed()) +} + +// Four wedged transactions pinned the vacuum horizon for 42 days and nothing +// measured it. These are the numbers that would have caught it on day one. +var _ = Describe("control plane database stats", func() { + var db *gorm.DB + + BeforeEach(func() { + if runtime.GOOS == "darwin" { + Skip("testcontainers requires Docker, not available on macOS CI") + } + db = testutil.SetupTestDB() + }) + + It("samples an idle database without error and reports a small horizon", func() { + stats, err := SampleControlPlaneDB(context.Background(), db) + Expect(err).ToNot(HaveOccurred()) + Expect(stats.OldestXminAge).To(BeNumerically(">=", 0)) + Expect(stats.LongestTransactionSeconds).To(BeNumerically("<", 60), + "an idle test database must not hold a long transaction") + }) + + It("reports a long-running transaction that is holding the horizon open", func() { + held := make(chan struct{}) + done := make(chan struct{}) + go func() { + defer close(done) + _ = db.Transaction(func(tx *gorm.DB) error { + var one int + _ = tx.Raw("SELECT 1").Scan(&one).Error + close(held) + time.Sleep(3 * time.Second) + return nil + }) + }() + <-held + time.Sleep(1500 * time.Millisecond) + + stats, err := SampleControlPlaneDB(context.Background(), db) + Expect(err).ToNot(HaveOccurred()) + Expect(stats.LongestTransactionSeconds).To(BeNumerically(">=", 1)) + <-done + }) + + It("reports the dead tuple ratio of a control-plane table that is accumulating dead rows", func() { + Expect(db.Exec(`CREATE TABLE backend_nodes (id serial PRIMARY KEY, name text)`).Error).To(Succeed()) + Expect(db.Exec(`INSERT INTO backend_nodes (name) SELECT 'node-' || g FROM generate_series(1, 200) g`).Error).To(Succeed()) + Expect(db.Exec(`DELETE FROM backend_nodes WHERE id % 2 = 0`).Error).To(Succeed()) + + // PostgreSQL flushes table statistics asynchronously, so poll rather + // than assume the delete is visible in the catalog straight away. + Eventually(func() float64 { + stats, err := SampleControlPlaneDB(context.Background(), db) + Expect(err).ToNot(HaveOccurred()) + return stats.DeadTupleRatio["backend_nodes"] + }, 30*time.Second, 500*time.Millisecond).Should(BeNumerically(">", 0), + "a table with 100 deleted rows must report dead tuples") + }) + + It("fails the sample rather than hanging when the caller's context is already done", func() { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + _, err := SampleControlPlaneDB(ctx, db) + Expect(err).To(HaveOccurred(), "the sample must run under the caller's context") + }) + + It("registers gauges without error", func() { + Expect(RegisterControlPlaneDBMetrics(db, time.Minute)).To(Succeed()) + }) + + Describe("cached sampling", func() { + It("does not re-read the database inside the minimum interval", func() { + sampler := &cachedDBSampler{db: db, minInterval: time.Hour} + first, ok := sampler.stats(context.Background()) + Expect(ok).To(BeTrue()) + Expect(first.LongestTransactionSeconds).To(BeNumerically("<", 1)) + + held := make(chan struct{}) + done := make(chan struct{}) + go func() { + defer close(done) + _ = db.Transaction(func(tx *gorm.DB) error { + var one int + _ = tx.Raw("SELECT 1").Scan(&one).Error + close(held) + time.Sleep(3 * time.Second) + return nil + }) + }() + <-held + time.Sleep(1500 * time.Millisecond) + + // A fresh read sees the open transaction, so a cached read that + // still reports the old value proves the database was not touched. + fresh, err := SampleControlPlaneDB(context.Background(), db) + Expect(err).ToNot(HaveOccurred()) + Expect(fresh.LongestTransactionSeconds).To(BeNumerically(">=", 1)) + + cached, ok := sampler.stats(context.Background()) + Expect(ok).To(BeTrue()) + Expect(cached.LongestTransactionSeconds).To(Equal(first.LongestTransactionSeconds)) + <-done + }) + + It("keeps reporting the last good values once the database is unreachable", func() { + Expect(db.Exec(`CREATE TABLE node_models (id serial PRIMARY KEY, name text)`).Error).To(Succeed()) + Expect(db.Exec(`INSERT INTO node_models (name) VALUES ('m1')`).Error).To(Succeed()) + + sampler := &cachedDBSampler{db: db, minInterval: 0} + Eventually(func() bool { + stats, ok := sampler.stats(context.Background()) + _, seen := stats.DeadTupleRatio["node_models"] + return ok && seen + }, 30*time.Second, 500*time.Millisecond).Should(BeTrue()) + good, _ := sampler.stats(context.Background()) + + closeDB(db) + + // minInterval is zero, so this call does attempt a fresh sample and + // that sample fails. The last good values must survive it. + after, ok := sampler.stats(context.Background()) + Expect(ok).To(BeTrue(), "a failed sample must not blank out the gauges") + Expect(after.OldestXminAge).To(Equal(good.OldestXminAge)) + Expect(after.DeadTupleRatio).To(HaveKey("node_models")) + }) + + It("rate-limits retries against a failing database to one attempt per interval", func() { + attempts := 0 + // Raw().Scan() runs through gorm's row callback, so this counts one + // tick per statement the sampler actually sends to the database. + Expect(db.Callback().Row().After("gorm:row").Register("count_attempts", func(tx *gorm.DB) { + attempts++ + })).To(Succeed()) + closeDB(db) + + sampler := &cachedDBSampler{db: db, minInterval: time.Hour} + for i := 0; i < 5; i++ { + _, ok := sampler.stats(context.Background()) + Expect(ok).To(BeFalse()) + } + + Expect(attempts).To(Equal(1), + "a failing sample must cost the cache interval too, or a struggling database is retried on every scrape") + }) + + It("reports nothing at all before the first successful sample", func() { + closeDB(db) + + sampler := &cachedDBSampler{db: db, minInterval: 0} + stats, ok := sampler.stats(context.Background()) + Expect(ok).To(BeFalse(), + "a zero xmin age would read as a healthy horizon, so nothing must be reported") + Expect(stats.OldestXminAge).To(BeZero()) + }) + }) + + Describe("scraping the gauges", func() { + var reader *sdkmetric.ManualReader + + BeforeEach(func() { + reader = sdkmetric.NewManualReader() + otel.SetMeterProvider(sdkmetric.NewMeterProvider(sdkmetric.WithReader(reader))) + }) + + It("exports the three gauges on a scrape", func() { + Expect(db.Exec(`CREATE TABLE backend_nodes (id serial PRIMARY KEY, name text)`).Error).To(Succeed()) + Expect(db.Exec(`INSERT INTO backend_nodes (name) VALUES ('n1')`).Error).To(Succeed()) + Expect(RegisterControlPlaneDBMetrics(db, 0)).To(Succeed()) + + Eventually(func() []string { + return collectedMetricNames(reader) + }, 30*time.Second, 500*time.Millisecond).Should(ContainElements( + "localai_control_plane_oldest_xmin_age", + "localai_control_plane_longest_transaction_seconds", + "localai_control_plane_dead_tuple_ratio", + )) + }) + + It("serves a scrape from the last good sample while the database is down", func() { + Expect(RegisterControlPlaneDBMetrics(db, 0)).To(Succeed()) + before := collectedInt64Gauge(reader, "localai_control_plane_oldest_xmin_age") + + closeDB(db) + + Expect(collectedMetricNames(reader)).To(ContainElement("localai_control_plane_oldest_xmin_age"), + "a scrape during a database outage must not drop the gauges") + Expect(collectedInt64Gauge(reader, "localai_control_plane_oldest_xmin_age")).To(Equal(before)) + }) + + It("omits the gauges from a scrape taken before any sample succeeded", func() { + closeDB(db) + Expect(RegisterControlPlaneDBMetrics(db, 0)).To(Succeed()) + + Expect(collectedMetricNames(reader)).ToNot(ContainElement("localai_control_plane_oldest_xmin_age")) + }) + }) +}) + +// The table names the gauges query used to be a hardcoded literal list. That +// compiles forever and matches nothing the moment a model's table name moves, +// and a dead-tuple gauge that matched no rows is indistinguishable from a +// healthy cluster. These specs pin that the names come from gorm instead. +var _ = Describe("control plane table names", func() { + var db *gorm.DB + + BeforeEach(func() { + if runtime.GOOS == "darwin" { + Skip("testcontainers requires Docker, not available on macOS CI") + } + db = testutil.SetupTestDB() + }) + + It("resolves the registry tables the gauges report on", func() { + names, err := controlPlaneTableNames(db) + Expect(err).ToNot(HaveOccurred()) + Expect(names).To(ConsistOf("backend_nodes", "node_models", "gallery_operations")) + }) + + It("honours a TableName override rather than guessing from the type", func() { + names, err := controlPlaneTableNames(db) + Expect(err).ToNot(HaveOccurred()) + + // GalleryOperationRecord overrides TableName. Naive pluralisation of the + // type would yield gallery_operation_records, so this assertion fails if + // the resolution ever stops consulting the model. + Expect(names).To(ContainElement("gallery_operations")) + Expect(names).ToNot(ContainElement("gallery_operation_records")) + }) +}) diff --git a/core/services/monitoring/monitoring_suite_test.go b/core/services/monitoring/monitoring_suite_test.go new file mode 100644 index 000000000000..05c5218333c9 --- /dev/null +++ b/core/services/monitoring/monitoring_suite_test.go @@ -0,0 +1,13 @@ +package monitoring + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestMonitoring(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Monitoring test suite") +} diff --git a/core/services/nodes/health.go b/core/services/nodes/health.go index ffe1cfa0e2e5..b82e57f91201 100644 --- a/core/services/nodes/health.go +++ b/core/services/nodes/health.go @@ -7,6 +7,7 @@ import ( "sync" "time" + "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/services/advisorylock" "github.com/mudler/xlog" "gorm.io/gorm" @@ -51,7 +52,11 @@ type HealthMonitor struct { // If clientFactory is nil, a default factory using the given authToken is used. func NewHealthMonitor(registry NodeHealthStore, db *gorm.DB, checkInterval, staleThreshold time.Duration, authToken string, perModelHealthCheck bool, clientFactory ...BackendClientFactory) *HealthMonitor { checkInterval = cmp.Or(checkInterval, 15*time.Second) - staleThreshold = cmp.Or(staleThreshold, 60*time.Second) + // Heartbeat checkpointing lets last_heartbeat sit up to one checkpoint + // interval behind by design, so a hardcoded 60s fallback here would mark + // every healthy, beating node offline. Track the shared default instead, + // which is derived from that interval. + staleThreshold = cmp.Or(staleThreshold, config.DefaultStaleNodeThreshold) var factory BackendClientFactory if len(clientFactory) > 0 && clientFactory[0] != nil { factory = clientFactory[0] diff --git a/core/services/nodes/heartbeat_checkpoint_test.go b/core/services/nodes/heartbeat_checkpoint_test.go new file mode 100644 index 000000000000..5d4dd0bc0578 --- /dev/null +++ b/core/services/nodes/heartbeat_checkpoint_test.go @@ -0,0 +1,338 @@ +package nodes + +import ( + "context" + "runtime" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/gorm" + + "github.com/mudler/LocalAI/core/services/testutil" +) + +// The per-beat UPDATE on backend_nodes was ~52k writes a day against a +// six-row table. With autovacuum healthy that is merely wasteful; with +// autovacuum blocked it grew the table to 460 MB and a six-row scan to 867 ms. +// A beat that carries only a fresher timestamp does not need to reach disk. +var _ = Describe("heartbeat checkpointing", func() { + var ( + db *gorm.DB + registry *NodeRegistry + ctx context.Context + nodeID string + ) + + // writtenAt reads the persisted timestamp, which is the only evidence of + // whether a beat actually reached the database. + writtenAt := func() time.Time { + var n BackendNode + Expect(db.First(&n, "id = ?", nodeID).Error).ToNot(HaveOccurred()) + return n.LastHeartbeat + } + + u64 := func(v uint64) *uint64 { return &v } + + // workerBeat mirrors what core/services/worker/registration.go heartbeatBody + // actually sends on EVERY beat: free VRAM, free RAM, and both disk figures. + // The disk block is unconditional by design, so total_disk is present on + // every real backend beat even though it never changes. + workerBeat := func(freeVRAM uint64) *HeartbeatUpdate { + return &HeartbeatUpdate{ + AvailableVRAM: u64(freeVRAM), + AvailableRAM: u64(16 << 30), + TotalDisk: u64(500 << 30), + AvailableDisk: u64(200 << 30), + } + } + + BeforeEach(func() { + if runtime.GOOS == "darwin" { + Skip("testcontainers requires Docker, not available on macOS CI") + } + db = testutil.SetupTestDB() + var err error + registry, err = NewNodeRegistry(db) + Expect(err).ToNot(HaveOccurred()) + ctx = context.Background() + + nodeID = "hb-node" + Expect(db.Create(&BackendNode{ + ID: nodeID, Name: "hb-box", NodeType: "backend", + Status: StatusHealthy, LastHeartbeat: time.Now().Add(-time.Hour), + }).Error).ToNot(HaveOccurred()) + }) + + It("writes the first beat and suppresses the ones inside the interval", func() { + registry.SetHeartbeatCheckpoint(time.Hour) + + Expect(registry.Heartbeat(ctx, nodeID, nil)).To(Succeed()) + first := writtenAt() + Expect(first).To(BeTemporally(">", time.Now().Add(-time.Minute))) + + for range 5 { + Expect(registry.Heartbeat(ctx, nodeID, nil)).To(Succeed()) + } + Expect(writtenAt()).To(BeTemporally("==", first), + "beats carrying only a timestamp must not reach the database") + }) + + It("writes every beat when checkpointing is disabled", func() { + registry.SetHeartbeatCheckpoint(0) + + Expect(registry.Heartbeat(ctx, nodeID, nil)).To(Succeed()) + first := writtenAt() + + time.Sleep(10 * time.Millisecond) + Expect(registry.Heartbeat(ctx, nodeID, nil)).To(Succeed()) + Expect(writtenAt()).To(BeTemporally(">", first)) + }) + + It("writes immediately when a reading moves materially", func() { + registry.SetHeartbeatCheckpoint(time.Hour) + + Expect(registry.Heartbeat(ctx, nodeID, nil)).To(Succeed()) + first := writtenAt() + + time.Sleep(10 * time.Millisecond) + vram := uint64(8) << 30 + Expect(registry.Heartbeat(ctx, nodeID, &HeartbeatUpdate{AvailableVRAM: &vram})).To(Succeed()) + Expect(writtenAt()).To(BeTemporally(">", first), + "the scheduler places against free VRAM, so a real change must not wait") + }) + + It("never suppresses beats from a node that is not active", func() { + registry.SetHeartbeatCheckpoint(time.Hour) + Expect(db.Model(&BackendNode{}).Where("id = ?", nodeID). + Update("status", StatusOffline).Error).ToNot(HaveOccurred()) + + Expect(registry.Heartbeat(ctx, nodeID, nil)).To(Succeed()) + first := writtenAt() + + time.Sleep(10 * time.Millisecond) + Expect(registry.Heartbeat(ctx, nodeID, nil)).To(Succeed()) + Expect(writtenAt()).To(BeTemporally(">", first), + "an offline node recovers only when the health monitor sees a fresh "+ + "timestamp, so suppressing its beats would strand it offline forever") + }) + + It("lets a node parked offline mid-interval write its next beat at once", func() { + registry.SetHeartbeatCheckpoint(time.Hour) + + Expect(registry.Heartbeat(ctx, nodeID, nil)).To(Succeed()) + first := writtenAt() + + // Taken offline while the checkpoint is still fresh, which is what a + // graceful shutdown or an admin action looks like. The Heartbeat path + // only forgets the checkpoint after a beat has already been written and + // found no active row, so the suppression state has to be dropped here. + Expect(registry.MarkOffline(ctx, nodeID)).To(Succeed()) + + time.Sleep(10 * time.Millisecond) + Expect(registry.Heartbeat(ctx, nodeID, nil)).To(Succeed()) + Expect(writtenAt()).To(BeTemporally(">", first), + "a node parked offline must not wait out the checkpoint before the "+ + "health monitor can see it is alive again") + }) + + // A backend worker never sends an empty body, so suppression that only + // holds for nil updates buys the widened staleness threshold and reduces + // nothing. These specs beat with the payload a real worker sends. + It("suppresses a realistic worker beat whose readings have not moved", func() { + registry.SetHeartbeatCheckpoint(time.Hour) + + Expect(registry.Heartbeat(ctx, nodeID, workerBeat(8<<30))).To(Succeed()) + first := writtenAt() + + for range 5 { + time.Sleep(5 * time.Millisecond) + Expect(registry.Heartbeat(ctx, nodeID, workerBeat(8<<30))).To(Succeed()) + } + Expect(writtenAt()).To(BeTemporally("==", first), + "re-reporting an unchanged reading must not be treated as a change: "+ + "total_disk rides on every backend beat, so testing presence "+ + "rather than movement suppresses nothing at all") + }) + + It("suppresses a re-reported total VRAM and GPU vendor that have not changed", func() { + registry.SetHeartbeatCheckpoint(time.Hour) + + beat := func() *HeartbeatUpdate { + u := workerBeat(8 << 30) + u.TotalVRAM = u64(24 << 30) + u.GPUVendor = "nvidia" + return u + } + + Expect(registry.Heartbeat(ctx, nodeID, beat())).To(Succeed()) + first := writtenAt() + + time.Sleep(10 * time.Millisecond) + Expect(registry.Heartbeat(ctx, nodeID, beat())).To(Succeed()) + Expect(writtenAt()).To(BeTemporally("==", first), + "hardware facts repeated verbatim carry nothing the database needs") + }) + + It("suppresses a reading that moves less than the material delta", func() { + registry.SetHeartbeatCheckpoint(time.Hour) + + Expect(registry.Heartbeat(ctx, nodeID, workerBeat(8<<30))).To(Succeed()) + first := writtenAt() + + time.Sleep(10 * time.Millisecond) + Expect(registry.Heartbeat(ctx, nodeID, workerBeat(8<<30+(1<<20)))).To(Succeed()) + Expect(writtenAt()).To(BeTemporally("==", first), + "1 MiB of drift cannot change a placement decision") + }) + + It("writes once accumulated drift passes the delta, measuring from the last persisted value", func() { + registry.SetHeartbeatCheckpoint(time.Hour) + + base := uint64(8) << 30 + Expect(registry.Heartbeat(ctx, nodeID, workerBeat(base))).To(Succeed()) + first := writtenAt() + + // Steps of 100 MiB. Each is under the 256 MiB delta, so measuring + // against the PREVIOUS BEAT would let free VRAM drift arbitrarily far + // from the persisted column without ever writing. Measuring against the + // last PERSISTED value means the third step crosses and writes. + step := uint64(100) << 20 + time.Sleep(5 * time.Millisecond) + Expect(registry.Heartbeat(ctx, nodeID, workerBeat(base-step))).To(Succeed()) + Expect(writtenAt()).To(BeTemporally("==", first), "100 MiB is under the delta") + + time.Sleep(5 * time.Millisecond) + Expect(registry.Heartbeat(ctx, nodeID, workerBeat(base-2*step))).To(Succeed()) + Expect(writtenAt()).To(BeTemporally("==", first), "200 MiB is still under the delta") + + time.Sleep(5 * time.Millisecond) + Expect(registry.Heartbeat(ctx, nodeID, workerBeat(base-3*step))).To(Succeed()) + Expect(writtenAt()).To(BeTemporally(">", first), + "300 MiB of drift from the persisted value must write, or the "+ + "scheduler places against a figure that silently walked away") + }) + + It("writes at once when a realistic beat moves free VRAM past the delta", func() { + registry.SetHeartbeatCheckpoint(time.Hour) + + Expect(registry.Heartbeat(ctx, nodeID, workerBeat(8<<30))).To(Succeed()) + first := writtenAt() + + time.Sleep(10 * time.Millisecond) + Expect(registry.Heartbeat(ctx, nodeID, workerBeat(4<<30))).To(Succeed()) + Expect(writtenAt()).To(BeTemporally(">", first), + "4 GiB less free VRAM changes where the scheduler can place") + }) + + // A node with a VRAM budget persists capAvailable(reported, ceiling), not + // the raw reading. Comparing the raw reading against the snapshot therefore + // measures a different quantity from the one the column holds: on a budgeted + // node whose actual free VRAM swings well above its ceiling, every beat looks + // material while the written value never moves, and suppression is defeated + // on exactly the nodes an operator has bothered to configure. + Describe("on a node with a VRAM budget", func() { + const ceiling = uint64(8) << 30 + + BeforeEach(func() { + Expect(db.Model(&BackendNode{}).Where("id = ?", nodeID). + Updates(map[string]any{ + ColTotalVRAM: uint64(24) << 30, + ColVRAMBudget: "8GB", + ColVRAMBudgetBytes: ceiling, + }).Error).ToNot(HaveOccurred()) + }) + + persistedVRAM := func() uint64 { + var n BackendNode + Expect(db.First(&n, "id = ?", nodeID).Error).ToNot(HaveOccurred()) + return n.AvailableVRAM + } + + It("suppresses beats whose raw reading oscillates above the ceiling", func() { + registry.SetHeartbeatCheckpoint(time.Hour) + + Expect(registry.Heartbeat(ctx, nodeID, workerBeat(20<<30))).To(Succeed()) + first := writtenAt() + Expect(persistedVRAM()).To(Equal(ceiling), "the column stores the capped figure") + + // Multi-gigabyte swings, all of them above the 8 GiB ceiling, so the + // capped value the column holds is unchanged every time. + for _, raw := range []uint64{12 << 30, 18 << 30, 9 << 30, 24 << 30} { + time.Sleep(5 * time.Millisecond) + Expect(registry.Heartbeat(ctx, nodeID, workerBeat(raw))).To(Succeed()) + } + + Expect(writtenAt()).To(BeTemporally("==", first), + "free VRAM above the budget ceiling cannot change a placement "+ + "decision, so a reading that only moves above it must not write") + Expect(persistedVRAM()).To(Equal(ceiling)) + }) + + It("still writes when the reading drops below the ceiling", func() { + registry.SetHeartbeatCheckpoint(time.Hour) + + Expect(registry.Heartbeat(ctx, nodeID, workerBeat(20<<30))).To(Succeed()) + first := writtenAt() + + time.Sleep(10 * time.Millisecond) + Expect(registry.Heartbeat(ctx, nodeID, workerBeat(2<<30))).To(Succeed()) + Expect(writtenAt()).To(BeTemporally(">", first), + "below the ceiling the reading is the scheduler's figure, and "+ + "6 GiB less of it changes where a model can be placed") + Expect(persistedVRAM()).To(Equal(uint64(2) << 30)) + }) + }) + + It("keeps a checkpointing node healthy while it beats normally", func() { + registry.SetHeartbeatCheckpoint(200 * time.Millisecond) + // perModelHealthCheck off: this spec is about liveness, not backends. + hm := NewHealthMonitor(registry, db, time.Minute, 5*time.Second, "", false) + + for range 6 { + Expect(registry.Heartbeat(ctx, nodeID, nil)).To(Succeed()) + time.Sleep(100 * time.Millisecond) + hm.doCheckAll(ctx) + } + + var n BackendNode + Expect(db.First(&n, "id = ?", nodeID).Error).ToNot(HaveOccurred()) + Expect(n.Status).To(Equal(StatusHealthy), + "suppressed beats must not read as a dead node") + }) + + It("does not mark a beating node offline when built with no explicit threshold", func() { + registry.SetHeartbeatCheckpoint(time.Minute) + // Two minutes since the last durable write is normal under + // checkpointing: the node may have beaten seconds ago and been + // suppressed. It is well inside the 5 minute default and well outside + // the 60 seconds the zero-threshold fallback used to resolve to, which + // would have flapped every node in the cluster offline each cycle. + Expect(db.Model(&BackendNode{}).Where("id = ?", nodeID). + Update("last_heartbeat", time.Now().Add(-2*time.Minute)).Error).ToNot(HaveOccurred()) + + // Zero staleThreshold: the constructor's fallback is what is under test. + hm := NewHealthMonitor(registry, db, time.Minute, 0, "", false) + hm.doCheckAll(ctx) + + var n BackendNode + Expect(db.First(&n, "id = ?", nodeID).Error).ToNot(HaveOccurred()) + Expect(n.Status).To(Equal(StatusHealthy), + "the default threshold has to track the checkpoint interval, or a "+ + "monitor built without one reaps every healthy node it sees") + }) + + It("still marks a genuinely silent node offline once the threshold elapses", func() { + registry.SetHeartbeatCheckpoint(time.Hour) + Expect(db.Model(&BackendNode{}).Where("id = ?", nodeID). + Update("last_heartbeat", time.Now().Add(-10*time.Minute)).Error).ToNot(HaveOccurred()) + + hm := NewHealthMonitor(registry, db, time.Minute, 5*time.Minute, "", false) + hm.doCheckAll(ctx) + + var n BackendNode + Expect(db.First(&n, "id = ?", nodeID).Error).ToNot(HaveOccurred()) + Expect(n.Status).To(Equal(StatusOffline), + "widening the threshold must not disable dead-node detection") + }) +}) diff --git a/core/services/nodes/registry.go b/core/services/nodes/registry.go index 4b4f7f1c8b32..f799f291f902 100644 --- a/core/services/nodes/registry.go +++ b/core/services/nodes/registry.go @@ -4,6 +4,7 @@ import ( "context" "errors" "fmt" + "sync" "sync/atomic" "time" @@ -366,6 +367,12 @@ type NodeRegistry struct { // Held in an atomic.Pointer for the same reason as the hooks above: the // startup wiring writes it while request handling reads it. aliasResolver atomic.Pointer[AliasResolver] + + // heartbeatCheckpoint bounds how often a beat that carries only a fresher + // timestamp reaches the database. Zero writes every beat. + heartbeatCheckpoint time.Duration + hbMu sync.Mutex + hbLastWrite map[string]heartbeatSnapshot } // AddReplicaRemovedHook registers a callback invoked after a replica row for @@ -437,6 +444,28 @@ func (r *NodeRegistry) nodeModelNames(ctx context.Context, db *gorm.DB, nodeID s return names } +// heartbeatSnapshot is the last state actually persisted for a node, so the +// next beat can tell an idle refresh from a real change. +type heartbeatSnapshot struct { + writtenAt time.Time + // availableVRAM is stored CAPPED by vramCeiling, because that is what the + // available_vram column holds. vramCeiling is the node's resolved VRAM + // budget (0 = none) as read on the last durable write, cached so a + // suppressed beat costs no query at all. + availableVRAM uint64 + vramCeiling uint64 + availableRAM uint64 + availableDisk uint64 + totalVRAM uint64 + totalDisk uint64 + gpuVendor string +} + +// heartbeatMaterialDelta is how far a reading must move before it is worth a +// write on its own. The scheduler places against free VRAM, so drift smaller +// than this cannot change a placement decision. +const heartbeatMaterialDelta = 256 << 20 // 256 MiB + // NewNodeRegistry creates a NodeRegistry and auto-migrates the schema. // Uses a PostgreSQL advisory lock to prevent concurrent migration races // when multiple instances (frontend + workers) start at the same time. @@ -478,7 +507,13 @@ func NewNodeRegistry(db *gorm.DB) (*NodeRegistry, error) { return nil }) - return &NodeRegistry{db: db}, nil + // heartbeatCheckpoint stays zero here: the registry writes every beat + // until the application wires the configured interval, so embedders and + // tests keep the historical behaviour. + return &NodeRegistry{ + db: db, + hbLastWrite: make(map[string]heartbeatSnapshot), + }, nil } // resolveVRAMBudgetBytes turns a budget string into an absolute byte ceiling @@ -695,6 +730,11 @@ func (r *NodeRegistry) MarkOffline(ctx context.Context, nodeID string) error { if err := r.setStatus(ctx, nodeID, StatusOffline); err != nil { return err } + // An offline node comes back only when the health monitor sees a fresh + // last_heartbeat, so its next beat must reach the database even if the + // checkpoint interval has not elapsed. The Heartbeat path forgets the + // checkpoint too, but only after a beat has already been suppressed. + r.forgetHeartbeatCheckpoint(nodeID) // Clear model records — node is shutting down. Capture the distinct models // and run the bulk delete inside a single transaction so the set of fired // hooks equals exactly the set of rows deleted: a SetNodeModel landing @@ -845,9 +885,144 @@ func (r *NodeRegistry) Deregister(ctx context.Context, nodeID string) error { for _, m := range removedModels { r.fireReplicaRemoved(m, nodeID, -1) } + // The node row is gone, so its checkpoint state is dead weight. Dropping + // it keeps the map bounded by the number of live nodes on a cluster that + // churns workers, and makes a re-registration under the same ID write its + // first beat immediately. + r.forgetHeartbeatCheckpoint(nodeID) return nil } +// SetHeartbeatCheckpoint bounds durable heartbeat writes. Zero restores a +// write per beat. +func (r *NodeRegistry) SetHeartbeatCheckpoint(d time.Duration) { + r.hbMu.Lock() + defer r.hbMu.Unlock() + r.heartbeatCheckpoint = d +} + +// skipHeartbeatWrite reports whether this beat carries nothing the database +// needs yet. It only decides: the caller commits the beat with +// recordHeartbeatWrite once it knows the budget ceiling the columns are +// actually written with. +func (r *NodeRegistry) skipHeartbeatWrite(nodeID string, update *HeartbeatUpdate) bool { + r.hbMu.Lock() + defer r.hbMu.Unlock() + + if r.heartbeatCheckpoint <= 0 { + return false + } + + prev, seen := r.hbLastWrite[nodeID] + if !seen { + return false + } + if heartbeatMaterial(prev, update) { + return false + } + return time.Since(prev.writtenAt) < r.heartbeatCheckpoint +} + +// heartbeatMaterial reports whether a beat carries something the database needs +// now, given what was last persisted for the node. +// +// Every field is compared against what was last PERSISTED, never merely tested +// for presence. A backend worker's heartbeat body carries total_disk (and +// available_disk, and free RAM) on every single beat by design, so presence +// would make every real beat material and suppress nothing at all: the +// empty-bodied agent workers would be the only nodes that ever benefited. +// +// Comparing against the last persisted value, rather than against the previous +// beat, is also what stops small moves accumulating: drift is always measured +// from the figure the scheduler is actually reading, so a reading that walks +// away in sub-delta steps still writes once the total distance crosses the +// threshold. +func heartbeatMaterial(prev heartbeatSnapshot, update *HeartbeatUpdate) bool { + if update == nil { + return false + } + if update.GPUVendor != "" && update.GPUVendor != prev.gpuVendor { + return true + } + // A total is a hardware fact, not a fluctuating reading, so any change at + // all is worth a write and no delta applies. + if update.TotalVRAM != nil && *update.TotalVRAM != prev.totalVRAM { + return true + } + if update.TotalDisk != nil && *update.TotalDisk != prev.totalDisk { + return true + } + // The CAPPED reading is compared, because capAvailable is what the column + // stores. On a node with a VRAM budget whose raw free reading oscillates + // above the ceiling, comparing the raw value makes every beat look material + // while the persisted value never moves, so suppression is defeated on + // precisely the nodes that have a budget set. + if update.AvailableVRAM != nil && + absDiff(capAvailable(*update.AvailableVRAM, prev.vramCeiling), prev.availableVRAM) > heartbeatMaterialDelta { + return true + } + if update.AvailableRAM != nil && absDiff(*update.AvailableRAM, prev.availableRAM) > heartbeatMaterialDelta { + return true + } + if update.AvailableDisk != nil && absDiff(*update.AvailableDisk, prev.availableDisk) > heartbeatMaterialDelta { + return true + } + return false +} + +// recordHeartbeatWrite snapshots what a beat is about to persist, so the next +// beat compares like with like. ceiling is the VRAM budget the write resolved; +// it is kept so the next beat can cap its own reading without re-reading the +// column, and free VRAM is stored capped for the same reason the column is. +func (r *NodeRegistry) recordHeartbeatWrite(nodeID string, update *HeartbeatUpdate, ceiling uint64) { + r.hbMu.Lock() + defer r.hbMu.Unlock() + + if r.heartbeatCheckpoint <= 0 { + return + } + + next := r.hbLastWrite[nodeID] + next.writtenAt = time.Now() + next.vramCeiling = ceiling + if update != nil { + if update.GPUVendor != "" { + next.gpuVendor = update.GPUVendor + } + if update.TotalVRAM != nil { + next.totalVRAM = *update.TotalVRAM + } + if update.TotalDisk != nil { + next.totalDisk = *update.TotalDisk + } + if update.AvailableVRAM != nil { + next.availableVRAM = capAvailable(*update.AvailableVRAM, ceiling) + } + if update.AvailableRAM != nil { + next.availableRAM = *update.AvailableRAM + } + if update.AvailableDisk != nil { + next.availableDisk = *update.AvailableDisk + } + } + r.hbLastWrite[nodeID] = next +} + +// forgetHeartbeatCheckpoint drops a node's suppression state so its next beat +// writes unconditionally. +func (r *NodeRegistry) forgetHeartbeatCheckpoint(nodeID string) { + r.hbMu.Lock() + defer r.hbMu.Unlock() + delete(r.hbLastWrite, nodeID) +} + +func absDiff(a, b uint64) uint64 { + if a > b { + return a - b + } + return b - a +} + // HeartbeatUpdate contains optional fields to update on heartbeat. type HeartbeatUpdate struct { AvailableVRAM *uint64 `json:"available_vram,omitempty"` @@ -867,16 +1042,28 @@ type HeartbeatUpdate struct { func (r *NodeRegistry) Heartbeat(ctx context.Context, nodeID string, update *HeartbeatUpdate) error { db := r.db.WithContext(ctx) + // Decided BEFORE the updates map is built, because building that map costs + // a SELECT for the node's VRAM budget ceiling, and a beat the database does + // not need must not pay for a query: per-beat control-plane queries are the + // load this checkpointing exists to remove. The decision reuses the ceiling + // cached on the last durable write, so it can be at most one checkpoint + // interval out of date. That costs at most one extra or one late write; it + // cannot persist a wrong figure, because the write path below re-reads the + // ceiling before it caps anything. + if r.skipHeartbeatWrite(nodeID, update) { + return nil + } + updates := map[string]any{ ColLastHeartbeat: time.Now(), } + var ceiling uint64 if update != nil { if update.AvailableVRAM != nil { // Cap the reported available against the node's resolved budget // ceiling (0 = none) so the SQL scheduler only ever sees budgeted // capacity. TotalVRAM stays raw (written below). - var ceiling uint64 db.Model(&BackendNode{}). Select(ColVRAMBudgetBytes).Where("id = ?", nodeID).Scan(&ceiling) updates[ColAvailableVRAM] = capAvailable(*update.AvailableVRAM, ceiling) @@ -904,6 +1091,10 @@ func (r *NodeRegistry) Heartbeat(ctx context.Context, nodeID string, update *Hea } } + // Recorded with the ceiling this write actually resolved, so the snapshot + // and the column always measure the same quantity. + r.recordHeartbeatWrite(nodeID, update, ceiling) + // Only update all fields (including status promotion) for active nodes. // Pending and offline nodes must go through approval or re-registration. result := db.Model(&BackendNode{}). @@ -913,6 +1104,9 @@ func (r *NodeRegistry) Heartbeat(ctx context.Context, nodeID string, update *Hea return fmt.Errorf("heartbeat for %s: %w", nodeID, result.Error) } if result.RowsAffected == 0 { + // Pending or offline. Its recovery depends on the health monitor + // seeing a fresh timestamp, so this node must never be suppressed. + r.forgetHeartbeatCheckpoint(nodeID) // May be pending or offline — still update heartbeat timestamp result = db.Model(&BackendNode{}).Where("id = ?", nodeID).Update(ColLastHeartbeat, time.Now()) if result.Error != nil { diff --git a/core/services/nodes/router.go b/core/services/nodes/router.go index 11ec2d6c639a..44acede1159a 100644 --- a/core/services/nodes/router.go +++ b/core/services/nodes/router.go @@ -1098,9 +1098,16 @@ func (r *SmartRouter) scheduleNewModel(ctx context.Context, backendType, modelID // carries the install. A worker that has died stops answering on the bus at // once but stays healthy in the database until its heartbeat ages out, so // without this the scheduler could commit to a node it cannot reach. + // + // The last selection error is kept because eviction below fires on a nil + // node, and a lookup that failed is not the same answer as a cluster with + // no room: a control-plane database slow enough to time out these queries + // read as "everybody is full" and cost a healthy model its place. + var selectErr error selectNode := func() *BackendNode { var candidate *BackendNode var selErr error + selectErr = nil if estimatedVRAM > 0 { if candidateNodeIDs != nil { candidate, selErr = r.registry.FindNodeWithVRAMFromSet(ctx, estimatedVRAM, candidateNodeIDs) @@ -1117,20 +1124,31 @@ func (r *SmartRouter) scheduleNewModel(ctx context.Context, backendType, modelID if candidateNodeIDs != nil { candidate, selErr = r.registry.FindIdleNodeFromSet(ctx, candidateNodeIDs) if selErr != nil { - candidate, _ = r.registry.FindLeastLoadedNodeFromSet(ctx, candidateNodeIDs) + candidate, selErr = r.registry.FindLeastLoadedNodeFromSet(ctx, candidateNodeIDs) } } else { candidate, selErr = r.registry.FindIdleNode(ctx) if selErr != nil { - candidate, _ = r.registry.FindLeastLoadedNode(ctx) + candidate, selErr = r.registry.FindLeastLoadedNode(ctx) } } + if candidate == nil { + selectErr = selErr + } } return candidate } node := r.pickReachableNode(ctx, selectNode) + // Same reasoning as the replica-slot guard further down: only + // gorm.ErrRecordNotFound is a verdict that the cluster has no node to give. + // Any other error left the question unanswered, and evicting on it costs a + // healthy model its place for no evidence. + if node == nil && selectErr != nil && !errors.Is(selectErr, gorm.ErrRecordNotFound) { + return nil, "", 0, fmt.Errorf("selecting a node for %s: %w", modelID, selectErr) + } + // 4. Preemptive eviction: if no suitable node found, evict the LRU model with zero in-flight if node == nil { evictedNode, evictErr := r.evictLRUAndFreeNodeFrom(ctx, candidateNodeIDs) @@ -1152,6 +1170,15 @@ func (r *SmartRouter) scheduleNewModel(ctx context.Context, backendType, modelID } replicaIdx, slotErr := r.registry.NextFreeReplicaIndex(ctx, node.ID, modelID, maxSlots) if slotErr != nil { + // Only ErrNoFreeSlot means "this node is full". Any other error means + // we could not find out, and evicting on a guess costs a healthy model + // its place: a control-plane database slow enough to time out this + // lookup made the scheduler evict loaded models it had no evidence to + // evict, and they thrashed. + if !errors.Is(slotErr, ErrNoFreeSlot) { + return nil, "", 0, fmt.Errorf("determining free replica slot on %s: %w", node.Name, slotErr) + } + // All slots on this node are taken — fall back to eviction. This is // rare in practice because FindNodesWithFreeSlot already filtered; // it can race with another concurrent scheduler. diff --git a/core/services/nodes/router_slot_uncertainty_test.go b/core/services/nodes/router_slot_uncertainty_test.go new file mode 100644 index 000000000000..b32649197761 --- /dev/null +++ b/core/services/nodes/router_slot_uncertainty_test.go @@ -0,0 +1,129 @@ +package nodes + +import ( + "context" + "errors" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/messaging" + pb "github.com/mudler/LocalAI/pkg/grpc/proto" + "gorm.io/gorm" +) + +// A slow control-plane database must never cost a loaded model its place. +// +// scheduleNewModel asks the registry for a free replica slot and, when that +// call fails, evicts the least-recently-used model to make room. The fallback +// is correct for "this node is full" and catastrophic for "I could not reach +// the database": a timeout evicted a healthy model belonging to another +// request, whose frontend then dialled the dead port and retried, so the model +// thrashed. Only ErrNoFreeSlot is evidence that the node is actually full. +var _ = Describe("replica slot lookup under database latency", func() { + var ( + reg *fakeModelRouter + backend *stubBackend + factory *stubClientFactory + unloader *fakeUnloader + ) + + BeforeEach(func() { + reg = &fakeModelRouter{ + findAndLockErr: errors.New("not found"), + findIdleNode: &BackendNode{ID: "n1", Name: "gpu-box", Address: "10.0.0.10:50051"}, + nextFreeReplicaErr: context.DeadlineExceeded, + } + backend = &stubBackend{loadResult: &pb.Result{Success: true}} + factory = &stubClientFactory{client: backend} + unloader = &fakeUnloader{ + installReply: &messaging.BackendInstallReply{Success: true, Address: "10.0.0.10:9001"}, + } + }) + + It("fails the load instead of evicting when the slot lookup times out", func() { + router := NewSmartRouter(reg, SmartRouterOptions{ + Unloader: unloader, + ClientFactory: factory, + }) + + _, err := router.Route(context.Background(), "new-model", "models/new.gguf", "llama-cpp", "", nil, false) + + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, context.DeadlineExceeded)).To(BeTrue(), + "the deadline must survive to the caller so the failure is diagnosable") + Expect(err.Error()).To(ContainSubstring("determining free replica slot")) + Expect(err.Error()).ToNot(ContainSubstring("eviction"), + "a timed-out lookup is not evidence that the node is full") + }) + + It("still evicts when the registry reports the node is genuinely full", func() { + reg.nextFreeReplicaErr = ErrNoFreeSlot + + router := NewSmartRouter(reg, SmartRouterOptions{ + Unloader: unloader, + ClientFactory: factory, + }) + + _, err := router.Route(context.Background(), "new-model", "models/new.gguf", "llama-cpp", "", nil, false) + + // The fake router has no DB, so evictLRUAndFreeNodeFrom reports + // ErrEvictionBusy. Reaching that error proves the eviction path ran, + // which is the behaviour this spec is pinning. + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("eviction failed")) + }) +}) + +// The same uncertainty applies one step earlier, when the scheduler asks which +// node should take the model. Both finder queries failing leaves it with no +// candidate, which the code read as "the cluster is full" and evicted for. +// gorm.ErrRecordNotFound is the only answer that means no node was available. +var _ = Describe("node selection under database latency", func() { + var ( + reg *fakeModelRouter + factory *stubClientFactory + unloader *fakeUnloader + ) + + BeforeEach(func() { + reg = &fakeModelRouter{findAndLockErr: errors.New("not found")} + factory = &stubClientFactory{client: &stubBackend{loadResult: &pb.Result{Success: true}}} + unloader = &fakeUnloader{ + installReply: &messaging.BackendInstallReply{Success: true, Address: "10.0.0.10:9001"}, + } + }) + + newRouter := func() *SmartRouter { + return NewSmartRouter(reg, SmartRouterOptions{ + Unloader: unloader, + ClientFactory: factory, + }) + } + + It("fails the load instead of evicting when the node lookup times out", func() { + reg.findIdleErr = context.DeadlineExceeded + reg.findLeastLoadedErr = context.DeadlineExceeded + + _, err := newRouter().Route(context.Background(), "new-model", "models/new.gguf", "llama-cpp", "", nil, false) + + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, context.DeadlineExceeded)).To(BeTrue(), + "the deadline must survive to the caller so the failure is diagnosable") + Expect(err.Error()).To(ContainSubstring("selecting a node")) + Expect(err.Error()).ToNot(ContainSubstring("eviction"), + "a timed-out lookup is not evidence that the cluster is full") + }) + + It("still evicts when the cluster genuinely has no node to give", func() { + reg.findIdleErr = gorm.ErrRecordNotFound + reg.findLeastLoadedErr = gorm.ErrRecordNotFound + + _, err := newRouter().Route(context.Background(), "new-model", "models/new.gguf", "llama-cpp", "", nil, false) + + // No DB behind the fake, so eviction reports ErrEvictionBusy. Reaching + // it proves the eviction path still runs for a genuinely full cluster. + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, ErrEvictionBusy)).To(BeTrue()) + }) +}) diff --git a/core/services/nodes/router_test.go b/core/services/nodes/router_test.go index 015cacb3040b..10c64632911d 100644 --- a/core/services/nodes/router_test.go +++ b/core/services/nodes/router_test.go @@ -93,6 +93,10 @@ type fakeModelRouter struct { findLRUModel *NodeModel findLRUErr error + // NextFreeReplicaIndex returns + nextFreeReplicaIdx int + nextFreeReplicaErr error + // Get returns getNode *BackendNode getErr error @@ -317,7 +321,7 @@ func (f *fakeModelRouter) ListModelCleanupRetries(_ context.Context, _ time.Time } func (f *fakeModelRouter) NextFreeReplicaIndex(_ context.Context, _, _ string, _ int) (int, error) { - return 0, nil + return f.nextFreeReplicaIdx, f.nextFreeReplicaErr } func (f *fakeModelRouter) CountReplicasOnNode(_ context.Context, _, _ string) (int, error) { diff --git a/core/services/nodes/worker_readiness.go b/core/services/nodes/worker_readiness.go index 89530abc827c..a0d5b3a19490 100644 --- a/core/services/nodes/worker_readiness.go +++ b/core/services/nodes/worker_readiness.go @@ -2,6 +2,7 @@ package nodes import ( "errors" + "fmt" "sync/atomic" ) @@ -73,3 +74,53 @@ func NATSReadiness(conn natsConn) func() error { return nil } } + +// ErrBackendUnreachable is reported when a worker holds a backend process it +// can no longer reach. +var ErrBackendUnreachable = errors.New("backend process is unreachable") + +// BackendAddressLister reports the gRPC addresses of the backend processes a +// worker currently believes it is running. +type BackendAddressLister interface { + LoadedBackendAddresses() []string +} + +// CompositeReadiness reports the first failure among its probes, so /readyz +// names the specific reason rather than a generic not-ready. +func CompositeReadiness(probes ...func() error) func() error { + return func() error { + for _, p := range probes { + if p == nil { + continue + } + if err := p(); err != nil { + return err + } + } + return nil + } +} + +// BackendDataPathReadiness closes the gap NATSReadiness leaves open: a worker +// whose NATS link is fine but whose backend processes have died is up and +// useless, and the scheduler cannot see the difference. It kept routing loads +// to a node that answered 200 while its backend port refused connections. +// +// A worker holding no backends is ready. That is the normal idle state, not a +// fault, and failing it would take every idle worker out of rotation. +func BackendDataPathReadiness(lister BackendAddressLister, dial func(string) error) func() error { + return func() error { + if lister == nil || dial == nil { + return nil + } + for _, addr := range lister.LoadedBackendAddresses() { + if addr == "" { + continue + } + if err := dial(addr); err != nil { + return fmt.Errorf("%w: %s: %v", ErrBackendUnreachable, addr, err) + } + } + return nil + } +} diff --git a/core/services/nodes/worker_readiness_datapath_test.go b/core/services/nodes/worker_readiness_datapath_test.go new file mode 100644 index 000000000000..7c5798c678de --- /dev/null +++ b/core/services/nodes/worker_readiness_datapath_test.go @@ -0,0 +1,57 @@ +package nodes + +import ( + "errors" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +type fakeAddressLister struct{ addrs []string } + +func (f *fakeAddressLister) LoadedBackendAddresses() []string { return f.addrs } + +// A worker that holds backend processes it can no longer reach is up and +// useless, but /readyz only tracked the NATS link, so it answered 200 and the +// scheduler kept routing loads to it. +var _ = Describe("worker readiness data path", func() { + It("is ready when a worker holds no backends", func() { + probe := BackendDataPathReadiness(&fakeAddressLister{}, func(string) error { + Fail("must not dial when there are no backends") + return nil + }) + Expect(probe()).To(Succeed()) + }) + + It("is ready when every held backend is dialable", func() { + probe := BackendDataPathReadiness( + &fakeAddressLister{addrs: []string{"10.0.0.1:9001", "10.0.0.1:9002"}}, + func(string) error { return nil }, + ) + Expect(probe()).To(Succeed()) + }) + + It("is not ready when a held backend cannot be reached", func() { + probe := BackendDataPathReadiness( + &fakeAddressLister{addrs: []string{"10.0.0.1:9001"}}, + func(string) error { return errors.New("connection refused") }, + ) + err := probe() + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, ErrBackendUnreachable)).To(BeTrue()) + Expect(err.Error()).To(ContainSubstring("10.0.0.1:9001")) + }) + + It("reports the first failing probe in a composite", func() { + boom := errors.New("nats is down") + probe := CompositeReadiness( + func() error { return nil }, + func() error { return boom }, + ) + Expect(probe()).To(MatchError(boom)) + }) + + It("is ready when a composite has no failing probe", func() { + Expect(CompositeReadiness(func() error { return nil }, func() error { return nil })()).To(Succeed()) + }) +}) diff --git a/core/services/worker/loaded_backend_addresses_test.go b/core/services/worker/loaded_backend_addresses_test.go new file mode 100644 index 000000000000..ed7a13287d8d --- /dev/null +++ b/core/services/worker/loaded_backend_addresses_test.go @@ -0,0 +1,66 @@ +package worker + +import ( + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// LoadedBackendAddresses feeds the worker's /readyz data-path probe, so every +// address it returns gets dialled. It must therefore report only processes +// that are actually listening: addr is recorded when the process is spawned, +// well before its gRPC server binds, and a cold start that took 10 to 15 +// seconds would otherwise pull the worker out of rotation every time. +var _ = Describe("loaded backend addresses", func() { + It("includes a backend that has passed the health-check gate", func() { + s := &backendSupervisor{processes: map[string]*backendProcess{ + "model#0": {addr: "127.0.0.1:50051", port: 50051, serving: true}, + }} + + Expect(s.LoadedBackendAddresses()).To(ConsistOf("127.0.0.1:50051")) + }) + + It("excludes a backend that is still starting", func() { + // The supervisor inserts the entry with addr already set, then polls + // for up to 30s waiting for the gRPC server to bind. Dialling during + // that window gets connection refused from a perfectly healthy start. + s := &backendSupervisor{processes: map[string]*backendProcess{ + "model#0": {addr: "127.0.0.1:50051", port: 50051}, + }} + + Expect(s.LoadedBackendAddresses()).To(BeEmpty()) + }) + + It("excludes a backend that is stopping", func() { + s := &backendSupervisor{processes: map[string]*backendProcess{ + "model#0": {addr: "127.0.0.1:50051", port: 50051, serving: true, stopping: true}, + }} + + Expect(s.LoadedBackendAddresses()).To(BeEmpty()) + }) + + It("reports only the serving members of a mixed set", func() { + s := &backendSupervisor{processes: map[string]*backendProcess{ + "serving#0": {addr: "127.0.0.1:50051", port: 50051, serving: true}, + "serving#1": {addr: "127.0.0.1:50052", port: 50052, serving: true}, + "starting#0": {addr: "127.0.0.1:50053", port: 50053}, + "stopping#0": {addr: "127.0.0.1:50054", port: 50054, serving: true, stopping: true}, + }} + + Expect(s.LoadedBackendAddresses()).To(ConsistOf("127.0.0.1:50051", "127.0.0.1:50052")) + }) + + It("holds no addresses when the worker holds no processes", func() { + s := &backendSupervisor{processes: map[string]*backendProcess{}} + + Expect(s.LoadedBackendAddresses()).To(BeEmpty()) + }) + + It("marks a backend serving only once it has answered a health check", func() { + bp := &backendProcess{addr: "127.0.0.1:50051", port: 50051} + s := &backendSupervisor{processes: map[string]*backendProcess{"model#0": bp}} + + Expect(s.LoadedBackendAddresses()).To(BeEmpty()) + Expect(s.markBackendServing("model#0", bp)).To(BeTrue()) + Expect(s.LoadedBackendAddresses()).To(ConsistOf("127.0.0.1:50051")) + }) +}) diff --git a/core/services/worker/replica_test.go b/core/services/worker/replica_test.go index a8cbfdc6c5f0..d7340ff32710 100644 --- a/core/services/worker/replica_test.go +++ b/core/services/worker/replica_test.go @@ -151,9 +151,9 @@ var _ = Describe("Worker per-replica process keying", func() { processes: map[string]*backendProcess{"model#0": bp}, } - Expect(s.backendStartStillValid("model#0", bp)).To(BeTrue()) + Expect(s.markBackendServing("model#0", bp)).To(BeTrue()) bp.stopping = true - Expect(s.backendStartStillValid("model#0", bp)).To(BeFalse()) + Expect(s.markBackendServing("model#0", bp)).To(BeFalse()) }) It("recycles a failed startup port at most once", func() { diff --git a/core/services/worker/supervisor.go b/core/services/worker/supervisor.go index cf95e8b63aaa..18441945785c 100644 --- a/core/services/worker/supervisor.go +++ b/core/services/worker/supervisor.go @@ -24,9 +24,19 @@ import ( // backendProcess represents a single gRPC backend process. type backendProcess struct { - proc *process.Process - addr string // gRPC address (host:port) - port int + proc *process.Process + addr string // gRPC address (host:port) + port int + // serving marks a process that has answered a gRPC health check and is + // therefore actually listening on addr. It is the opening bracket of the + // lifecycle that stopping closes. + // + // addr is recorded when the process is spawned, but the gRPC server can + // take 10 to 15 seconds to bind on a slow node (see the readiness poll in + // startBackend), so between those two points addr refuses connections. + // Anything that dials held backends must wait for this flag, or a cold + // start reads as a dead backend. + serving bool stopping bool // backendName is the gallery backend this process was started for (e.g. // "cuda13-nvidia-l4t-arm64-longcat-video"). It is NOT derivable from the @@ -511,7 +521,7 @@ func (s *backendSupervisor) startBackend(backend, backendName, backendPath strin // Verify the process wasn't stopped/replaced while health-checking. // A stopping entry remains in the map until process termination so its // port stays reserved, but it must not be advertised as ready. - if !s.backendStartStillValid(backend, bp) { + if !s.markBackendServing(backend, bp) { return "", fmt.Errorf("backend %s was stopped during startup", backend) } xlog.Debug("Backend gRPC server is ready", "backend", backend, "addr", clientAddr) @@ -545,14 +555,23 @@ func (s *backendSupervisor) startBackend(backend, backendName, backendPath strin return "", fmt.Errorf("backend %s did not become ready within %s. Last stderr:\n%s", backend, readinessTimeout, stderrTail) } -// backendStartStillValid verifies that a successful readiness probe still -// belongs to the active startup attempt. Stop keeps an entry tracked while it -// terminates, so pointer identity alone is not enough. -func (s *backendSupervisor) backendStartStillValid(key string, bp *backendProcess) bool { +// markBackendServing verifies that a successful readiness probe still belongs +// to the active startup attempt and, when it does, records the process as +// serving. Stop keeps an entry tracked while it terminates, so pointer +// identity alone is not enough. +// +// The check and the mark share one lock hold on purpose: serving is what +// LoadedBackendAddresses dials, so it must only ever be set on the entry this +// key currently owns and only once that entry has answered a health check. +func (s *backendSupervisor) markBackendServing(key string, bp *backendProcess) bool { s.mu.Lock() defer s.mu.Unlock() current, exists := s.processes[key] - return exists && current == bp && !current.stopping + if !exists || current != bp || current.stopping { + return false + } + current.serving = true + return true } // reapDeadProcess drops the bookkeeping for a process that exited without @@ -980,3 +999,26 @@ func (s *backendSupervisor) getAddr(backend string) string { } return "" } + +// LoadedBackendAddresses returns the gRPC addresses of the backend processes +// this worker is currently serving, for the /readyz data-path probe to dial. +// +// Only the middle of the lifecycle counts. A process is skipped until it has +// answered a health check, because addr is recorded at spawn time and refuses +// connections until the gRPC server binds, which takes 10 to 15 seconds on a +// slow node. It is skipped again once stopping is set. Reporting either end +// would make a routine cold start or shutdown read as a data-path fault, and a +// Kubernetes readinessProbe would pull the worker out of rotation for it. +func (s *backendSupervisor) LoadedBackendAddresses() []string { + s.mu.Lock() + defer s.mu.Unlock() + + addrs := make([]string, 0, len(s.processes)) + for _, bp := range s.processes { + if bp == nil || !bp.serving || bp.stopping || bp.addr == "" { + continue + } + addrs = append(addrs, bp.addr) + } + return addrs +} diff --git a/core/services/worker/worker.go b/core/services/worker/worker.go index 6434c3cd6b69..f57a5a4ae446 100644 --- a/core/services/worker/worker.go +++ b/core/services/worker/worker.go @@ -5,6 +5,7 @@ import ( "context" "encoding/json" "fmt" + "net" "os" "os/signal" "path/filepath" @@ -149,7 +150,8 @@ func Run(ctx *cliContext.Context, cfg *Config) error { httpAddr := cfg.resolveHTTPAddr() stagingDir := filepath.Join(cfg.ModelsPath, "..", "staging") dataDir := filepath.Join(cfg.ModelsPath, "..", "data") - // The readiness gate is created here but only armed once NATS is up, below. + // The readiness gate is created here but only armed once NATS is up and the + // backend supervisor exists, below, because the gate probes both. // Until then /readyz reports ready, which is correct: reaching this line // means the worker has already registered with the frontend, so it is // mid-startup rather than broken. @@ -172,11 +174,6 @@ func Run(ctx *cliContext.Context, cfg *Config) error { } defer natsClient.Close() - // Arm the readiness gate now that the worker can actually receive work. - // From here /readyz tracks the live NATS link, so a worker that is up but - // cut off from the bus reports 503 instead of a meaningless 200 (#10987). - readiness.Set(nodes.NATSReadiness(natsClient)) - // Start heartbeat goroutine (after NATS is connected so IsConnected check works) go func() { ticker := time.NewTicker(heartbeatInterval) @@ -226,6 +223,25 @@ func Run(ctx *cliContext.Context, cfg *Config) error { minPort: basePort, maxPort: cfg.effectiveMaxPort(basePort), } + + // Arm the readiness gate now that the worker can actually receive work. + // NATS is already connected at this point, so a worker that is up but cut + // off from the bus reports 503 instead of a meaningless 200 (#10987). + // + // Readiness also covers the data path: a worker whose NATS link is fine + // but whose backend processes have died is up and useless, and the + // scheduler cannot tell the difference from the bus alone. + readiness.Set(nodes.CompositeReadiness( + nodes.NATSReadiness(natsClient), + nodes.BackendDataPathReadiness(supervisor, func(addr string) error { + conn, err := net.DialTimeout("tcp", addr, 2*time.Second) + if err != nil { + return err + } + return conn.Close() + }), + )) + if err := supervisor.subscribeLifecycleEvents(); err != nil { nodes.ShutdownFileTransferServer(httpServer) return fmt.Errorf("subscribing to worker lifecycle events: %w", err) diff --git a/core/templates/evaluator.go b/core/templates/evaluator.go index 36e491b7de24..f4d40ac71e88 100644 --- a/core/templates/evaluator.go +++ b/core/templates/evaluator.go @@ -192,8 +192,10 @@ func (e *Evaluator) TemplateMessages(input schema.OpenAIRequest, messages []sche marshalAny(i.ToolCalls) } } - // Special Handling: System. We care if it was printed at all, not the r branch, so check separately - if contentExists && role == "system" { + // Special Handling: System. We care if it was printed at all, not the r branch, so check separately. + // Whitespace-only system content must not suppress the model config system_prompt + // (web Chat UI historically sent a blank system turn). + if contentExists && role == "system" && strings.TrimSpace(i.StringContent) != "" { suppressConfigSystemPrompt = true } } diff --git a/docker-compose.distributed.yaml b/docker-compose.distributed.yaml index 3387e313415b..014a74bf91a4 100644 --- a/docker-compose.distributed.yaml +++ b/docker-compose.distributed.yaml @@ -106,9 +106,11 @@ services: - worker # No HEALTHCHECK_ENDPOINT override is needed: the image's healthcheck # detects worker mode and derives the port from LOCALAI_SERVE_ADDR below - # (gRPC base port - 1 = 50050). The worker's /readyz reports 503 while its - # NATS connection is down, so `unhealthy` here means the worker genuinely - # cannot receive work. + # (gRPC base port - 1 = 50050). The worker's /readyz reports 503 when its + # NATS connection is down, and also when a backend process it is already + # serving no longer answers a TCP dial on its gRPC address. So `unhealthy` + # here means the worker genuinely cannot receive work, or cannot serve the + # models it is holding. environment: LOCALAI_SERVE_ADDR: "0.0.0.0:50051" LOCALAI_ADVERTISE_ADDR: "worker-1:50051" diff --git a/docs/content/advanced/model-configuration.md b/docs/content/advanced/model-configuration.md index a162a6225319..8e0ca855b0a1 100644 --- a/docs/content/advanced/model-configuration.md +++ b/docs/content/advanced/model-configuration.md @@ -173,6 +173,14 @@ These settings will be used as defaults for all the API calls to the model. | `tfz` | float | `1.0` | Tail free z parameter | | `keep` | int | `0` | Number of tokens to keep from the prompt | +{{% notice note %}} +The DS4 backend preserves its legacy behavior for omitted or non-positive +`max_tokens` values by generating at most 256 tokens. Set `max_tokens` to a +positive value when you need a specific DS4 output limit. After processing the +prompt, DS4 clamps that limit to the available context space and reserves one +context slot for safe generation. +{{% /notice %}} + ### Language and Translation | Field | Type | Description | diff --git a/docs/content/features/distributed-mode.md b/docs/content/features/distributed-mode.md index 0231c2dc4a52..faf721c054f4 100644 --- a/docs/content/features/distributed-mode.md +++ b/docs/content/features/distributed-mode.md @@ -76,6 +76,8 @@ The frontend is a standard LocalAI instance with distributed mode enabled. These | `--backend-upgrade-timeout` | `LOCALAI_NATS_BACKEND_UPGRADE_TIMEOUT` | `15m` | Same as the install timeout, applied to backend upgrades (force-reinstall). | | `--model-load-timeout` | `LOCALAI_NATS_MODEL_LOAD_TIMEOUT` | *(derived from checkpoint size)* | Pins the deadline for the `LoadModel` gRPC call the frontend issues to a worker. Leave it unset: by default the deadline is **derived from the checkpoint's on-disk size** (see below), which is what the worker actually spends its load time reading. Set it only to pin a specific budget — the value is then used verbatim, including when it is *shorter* than the derived one, so an operator who wants fast failure gets it. | | *(env only)* | `LOCALAI_MODEL_LOAD_WAIT` | `60s` | How long an inference request waits for a model that is still cold-loading onto a worker before it is answered with `503`, a `Retry-After` header and live staging progress. The request is served the moment the model becomes ready, so a model already most of the way staged needs no client retry. Set to `0` to wait as long as the load takes — only safe when no ingress or load balancer with an idle timeout sits in front. See [Requests for a model that is still loading](#requests-for-a-model-that-is-still-loading). | +| `--node-heartbeat-checkpoint` | `LOCALAI_NODE_HEARTBEAT_CHECKPOINT` | `60s` | Minimum gap between **durable** heartbeat writes for a worker node. A beat that only carries a fresher timestamp is kept in memory until this interval elapses instead of being written to PostgreSQL; every reported field is compared against the value last written rather than merely tested for presence, so a node's first beat, a changed total VRAM / total disk / GPU vendor, and a free VRAM / RAM / disk reading that has moved more than 256 MiB from the written value all still write immediately, and a node that is not active is never suppressed. Set it below the worker's `--heartbeat-interval` to restore a write per beat. See [Heartbeat writes and stale-node detection](#heartbeat-writes-and-stale-node-detection). | +| `--stale-node-threshold` | `LOCALAI_STALE_NODE_THRESHOLD` | `5m` | How long a node may go without a **durable** heartbeat before the health monitor marks it `offline`. Because `--node-heartbeat-checkpoint` holds back a beat that only carries a fresher timestamp, this has to stay comfortably wider than that interval: raising the checkpoint without raising this marks healthy, beating nodes offline. Neither the per-model gRPC health check nor request-time failure reads `last_heartbeat`, so neither is affected by this knob. See [Heartbeat writes and stale-node detection](#heartbeat-writes-and-stale-node-detection). | | `--expose-node-header` | `LOCALAI_EXPOSE_NODE_HEADER` | `false` | When enabled, inference responses carry an `X-LocalAI-Node` header with the ID of the worker node that served the request. Coverage spans the OpenAI-compatible endpoints (chat completions, completions, embeddings, audio transcriptions, audio speech / TTS, image generations, image inpainting), the Jina rerank endpoint (`/v1/rerank`), the VAD endpoints (`/v1/vad`, `/vad`), and the Anthropic Messages (`/v1/messages`) and Ollama (`/api/chat`, `/api/generate`, `/api/embed`) shims. Useful for debugging, observability and load-balancer attribution. Off by default: the node ID reveals internal cluster topology and should not be exposed on a public endpoint. Best-effort: under heavy concurrency for the same model across multiple replicas, the header may reflect a recent routing decision rather than this exact request's. Acceptable for observability and debugging. | ### The model load deadline scales with the checkpoint @@ -325,10 +327,14 @@ The worker's HTTP server (base port - 1, default 50050) exposes two unauthentica | Endpoint | Meaning | |----------|---------| | `/healthz` | **Liveness.** 200 whenever the process is up and serving. Deliberately independent of readiness, so a brief NATS outage does not trigger a restart storm across every worker. | -| `/readyz` | **Readiness.** 200 only when the worker is registered *and* its NATS connection is live; 503 otherwise. | +| `/readyz` | **Readiness.** 200 only when the worker is registered, its NATS connection is live, *and* every backend process it is currently serving answers a short TCP dial on its gRPC address; 503 otherwise. A worker holding no backends is ready, because idle is a healthy state, and so is one whose backends are still starting up. | `/readyz` reports something the frontend cannot see on its own. The node registry's `status` and `last_heartbeat` are driven by an HTTP heartbeat to the frontend, which is a different network path from NATS — a worker can keep heartbeating while its NATS link is dead, and so appear `healthy` in the registry while being unable to receive any work. The local probe closes that gap. +The same applies to the data path. A worker can hold a live NATS link while the backend processes it believes it is running have died, so it reports healthy while every load routed to it fails. `/readyz` therefore also dials the recorded gRPC address of each backend the worker is serving, and a worker whose backend port refuses connections drops out of rotation instead of absorbing work it cannot serve. + +Only backends in the middle of their lifecycle are dialled. A backend that is still starting is skipped until its gRPC server has answered a health check, which can take 10 to 15 seconds on a slow node, and a backend that is stopping is skipped from the moment shutdown begins. Neither a cold start nor an ordinary shutdown makes a worker report 503, so a Kubernetes `readinessProbe` at the usual 10s period does not pull a worker out of rotation every time it loads a model. + The container image's `HEALTHCHECK` detects worker mode and probes this endpoint automatically; no `HEALTHCHECK_ENDPOINT` override is needed. Set `HEALTHCHECK_ENDPOINT` only to pin an explicit URL. ### Worker Address Configuration @@ -642,6 +648,105 @@ To skip manual approval and let nodes join immediately, set `--auto-approve-node | `offline` | Node is temporarily offline (graceful shutdown or stale heartbeat). The node row is preserved so re-registration restores the previous approval status without requiring re-approval | | `draining` | Node is shutting down gracefully - no new requests are routed to it, existing in-flight requests are allowed to complete | +### Heartbeat writes and stale-node detection + +Workers beat every `--heartbeat-interval` (default `10s`). Writing each beat straight +to PostgreSQL means roughly 52,000 `UPDATE`s a day against a table that holds one row +per node. With autovacuum healthy that is merely wasteful. With autovacuum blocked -- +by a long-lived idle transaction, for example -- the dead tuples accumulate, and a +six-row table has been observed growing to 460 MB, at which point scanning it cost +867 ms and the queries that place models began timing out. + +So the frontend **checkpoints** the write. A beat that carries nothing but a fresher +timestamp is held in memory until `--node-heartbeat-checkpoint` (default `60s`) has +elapsed since that node's last durable write. These beats still reach the database +without waiting: + +- the node's first beat after the frontend starts, or after it was seen offline +- any beat from a node that is not active (`pending`, `offline`), because such a node + recovers only when the health monitor sees a fresh timestamp +- a GPU vendor, total VRAM or total disk that **differs** from the stored value, since + those are hardware facts and a change to one is a real event +- a free VRAM, free RAM or free disk reading that has moved more than 256 MiB, because + the scheduler places against those figures + +Every figure is compared against the value **last written**, not against the previous +beat. A worker reports its disk capacity on every single beat, so testing whether a +field is merely *present* would make every real beat look like a change and suppress +nothing. Measuring from the written value also means a reading that walks away in +sub-256 MiB steps still writes once the total distance crosses the threshold, rather +than drifting arbitrarily far from the figure the scheduler is reading. + +The consequence is that `last_heartbeat` is up to one checkpoint interval behind +reality **by design**. The stale-node threshold therefore defaults to **5 minutes** +(it was 60 seconds before checkpointing existed): the health monitor waits that long +without a fresh timestamp before it marks a node `offline`. It is configurable with +`--stale-node-threshold` / `LOCALAI_STALE_NODE_THRESHOLD`, and an operator who widens +`--node-heartbeat-checkpoint` must widen this to match, or the beats that checkpointing +suppresses will read as a dead node. + +{{% notice note %}} +Marking a node `offline` from a stale heartbeat now takes up to five minutes. This is +the slowest of the three ways a dead worker is noticed, not the only one. The per-model +gRPC health check still probes each loaded model on the health-monitor interval +(default `15s`) and removes replicas whose backend has died, and a request routed to a +gone worker still fails and is retried elsewhere at request time. Neither of those +paths reads `last_heartbeat`, so neither is slowed by this change. +{{% /notice %}} + +To go back to a durable write per beat -- on a database with plenty of write headroom, +or while debugging heartbeat delivery -- set `LOCALAI_NODE_HEARTBEAT_CHECKPOINT` to a +value below the worker's heartbeat interval, for example `1s`. + +### Operations: do not share a database with the vector store + +Give the control plane a PostgreSQL **database of its own**. Sharing one with the +agent vector store, or with anything else that holds long transactions, is the fastest +way to reproduce the 460 MB node registry described above. + +PostgreSQL computes the removable-tuple cutoff **per database**, not per table. One +transaction left open anywhere in the database -- a stalled embedding batch, an idle +`BEGIN` from a connection pool, an abandoned `psql` session -- pins that cutoff for +**every** table in it. Autovacuum still runs, finds nothing it is allowed to reclaim, +and moves on. The node registry is six rows rewritten tens of thousands of times a +day, so it is the table that pays: it bloats into hundreds of megabytes, a sequential +scan starts costing the best part of a second, and model placement begins timing out +while the vector store that caused it looks perfectly healthy. + +Concretely, these two must point at different databases: + +| Variable | What it holds | +|----------|---------------| +| `LOCALAI_AUTH_DATABASE_URL` | Auth **and the distributed control plane** - nodes, replicas, load jobs | +| `LOCALAI_AGENT_POOL_DATABASE_URL` | Agent collections and their embeddings | + +Different databases on the same PostgreSQL server is enough; they do not need separate +servers. Different *schemas* in one database is **not** enough, because the cutoff is +per database. + +To detect it before placement starts failing, watch the +`localai_control_plane_oldest_xmin_age` gauge, exported on the frontend's OpenTelemetry +meter. It reports how many transactions have elapsed +since the oldest snapshot still held open against the control plane's database. Under +normal load it stays small and flat. A line that climbs without coming back down means +something is holding a transaction open and autovacuum has stopped reclaiming the node +registry; find it with: + +```sql +SELECT pid, state, age(backend_xmin) AS xmin_age, query + FROM pg_stat_activity + WHERE backend_xmin IS NOT NULL + ORDER BY age(backend_xmin) DESC + LIMIT 5; +``` + +Grant `pg_read_all_stats` to the role LocalAI connects as (`GRANT pg_read_all_stats TO +localai;`), or make it a superuser. PostgreSQL blanks `backend_xmin` and `xact_start` in +`pg_stat_activity` for sessions owned by **other** roles, so without that grant both the +gauge and the query above see only LocalAI's own sessions -- and the transaction that +wedges the horizon is typically the co-located vector store connecting as a different +role, which is exactly the case they exist to catch. + ## Agent Workers Agent workers are dedicated processes for executing agent chats and MCP CI jobs. Unlike backend workers (which run gRPC model inference), agent workers use cogito to orchestrate multi-step conversations with tool calls. diff --git a/tests/e2e/distributed/db_latency_resilience_test.go b/tests/e2e/distributed/db_latency_resilience_test.go new file mode 100644 index 000000000000..471beda7b6af --- /dev/null +++ b/tests/e2e/distributed/db_latency_resilience_test.go @@ -0,0 +1,174 @@ +package distributed_test + +import ( + "context" + "fmt" + "net/url" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + pgdriver "gorm.io/driver/postgres" + "gorm.io/gorm" + "gorm.io/gorm/logger" + + "github.com/mudler/LocalAI/core/services/nodes" +) + +// routerRole is the unprivileged login the router's own handle uses. The +// control plane must be exercised as a role whose reads can actually be +// refused: the test container's owner is a PostgreSQL superuser, and a +// superuser bypasses every privilege check, so a REVOKE against it is recorded +// and then ignored. Injecting the failure through a separate role is the only +// way this spec observes a refusal at all. +const ( + routerRole = "slot_blind_router" + routerPassword = "slot_blind" +) + +// A control-plane database that fails its queries must cost the cluster +// throughput, never loaded models. Before the eviction guard, a slot lookup +// that failed evicted a healthy model that had done nothing wrong. +var _ = Describe("Control plane under a failing database", Label("Distributed"), func() { + var ( + infra *TestInfra + db *gorm.DB + routerDB *gorm.DB + ctx context.Context + ) + + openDB := func(dsn string) *gorm.DB { + GinkgoHelper() + h, err := gorm.Open(pgdriver.Open(dsn), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)}) + Expect(err).ToNot(HaveOccurred()) + return h + } + + // asRole rewrites the container's admin DSN to log in as the unprivileged + // router role, so the router's connections carry its restrictions. + asRole := func(adminDSN, role, password string) string { + GinkgoHelper() + u, err := url.Parse(adminDSN) + Expect(err).ToNot(HaveOccurred()) + u.User = url.UserPassword(role, password) + return u.String() + } + + BeforeEach(func() { + infra = SetupInfra("localai_db_latency_test") + ctx = context.Background() + db = openDB(infra.PGURL) + + _, err := nodes.NewNodeRegistry(db) + Expect(err).ToNot(HaveOccurred()) + + Expect(db.Create(&nodes.BackendNode{ + ID: "n-keep", Name: "keeper", NodeType: "backend", + Address: "10.0.0.1:50051", Status: nodes.StatusHealthy, + LastHeartbeat: time.Now(), MaxReplicasPerModel: 1, + }).Error).ToNot(HaveOccurred()) + + // A loaded model with nothing in flight: the exact row the old code + // would evict when it could not read the slot table. + Expect(db.Create(&nodes.NodeModel{ + ID: "nm-victim", NodeID: "n-keep", ModelName: "victim", ReplicaIndex: 0, + State: "loaded", InFlight: 0, Address: "10.0.0.1:9001", + LastUsed: time.Now().Add(-time.Hour), + }).Error).ToNot(HaveOccurred()) + + Expect(db.Exec(fmt.Sprintf( + `CREATE ROLE %s LOGIN PASSWORD '%s'`, routerRole, routerPassword)).Error).ToNot(HaveOccurred()) + for _, grant := range []string{ + fmt.Sprintf(`GRANT USAGE ON SCHEMA public TO %s`, routerRole), + fmt.Sprintf(`GRANT ALL PRIVILEGES ON ALL TABLES IN SCHEMA public TO %s`, routerRole), + fmt.Sprintf(`GRANT ALL PRIVILEGES ON ALL SEQUENCES IN SCHEMA public TO %s`, routerRole), + } { + Expect(db.Exec(grant).Error).ToNot(HaveOccurred()) + } + + routerDB = openDB(asRole(infra.PGURL, routerRole, routerPassword)) + + // Registered here rather than next to the REVOKE so the restore is + // unconditional: it runs whether the body revokes anything or not, and + // whether the body returns, fails an assertion or panics. A spec that + // aborts midway must not hand the next one a role that cannot read. + // A table-level grant covers every column, so this also supersedes the + // per-column grants the revoke leaves behind. + DeferCleanup(func() { + _ = db.Exec(fmt.Sprintf(`GRANT SELECT ON node_models TO %s`, routerRole)).Error + }) + }) + + // revokeSlotReads takes away the router's ability to read the replica-slot + // column, and only that column. + // + // The scope matters. Revoking the whole node_models table would break node + // selection too, which runs first and has a guard of its own, so the + // scheduler would never reach the slot lookup this spec is about. Column + // scope lets selection succeed (it reads node_id, state and in_flight, and + // counts rows) and lands the refusal exactly on NextFreeReplicaIndex, which + // plucks replica_index. A table-level grant covers every column, so the + // grant has to be re-issued column by column rather than carved out with a + // column-level REVOKE, which PostgreSQL ignores while the table grant + // stands. + revokeSlotReads := func() { + GinkgoHelper() + Expect(db.Exec(fmt.Sprintf(`REVOKE SELECT ON node_models FROM %s`, routerRole)).Error).ToNot(HaveOccurred()) + Expect(db.Exec(fmt.Sprintf(`DO $$ +DECLARE cols text; +BEGIN + SELECT string_agg(quote_ident(column_name), ', ') INTO cols + FROM information_schema.columns + WHERE table_schema = 'public' AND table_name = 'node_models' + AND column_name <> 'replica_index'; + EXECUTE format('GRANT SELECT (%%s) ON public.node_models TO %s', cols); +END $$`, routerRole)).Error).ToNot(HaveOccurred()) + } + + It("does not evict a loaded model when the database refuses the slot lookup", func() { + // Built before the revoke: AutoMigrate inspects the schema, and a + // migration that cannot run would degrade this spec into asserting + // almost nothing. + routerRegistry, err := nodes.NewNodeRegistry(routerDB) + Expect(err).ToNot(HaveOccurred()) + + revokeSlotReads() + + // Precondition: the injection actually bites. Without this the spec + // could pass against a database that answers every query. + _, slotErr := routerRegistry.NextFreeReplicaIndex(ctx, "n-keep", "newcomer", 1) + Expect(slotErr).To(HaveOccurred(), "the slot lookup must be refused for this spec to mean anything") + Expect(slotErr).ToNot(MatchError(nodes.ErrNoFreeSlot), "a refusal is not evidence that the node is full") + + router := nodes.NewSmartRouter(routerRegistry, nodes.SmartRouterOptions{DB: routerDB}) + _, routeErr := router.Route(ctx, "newcomer", "models/newcomer.gguf", "llama-cpp", "", nil, false) + + Expect(routeErr).To(HaveOccurred(), "a failing database cannot produce a successful placement") + + // These two carry the regression. Removing the guard makes the router + // fall through to eviction, and the message becomes "no replica slot on + // keeper and eviction failed", which trips both of them. + Expect(routeErr.Error()).To(ContainSubstring("determining free replica slot"), + "the failure must name the lookup that could not be answered") + Expect(routeErr.Error()).ToNot(ContainSubstring("eviction failed"), + "a failed lookup is not evidence that the node is full") + + // The row assertions below state the property the guard exists to + // protect, but under this particular injection they cannot fail on + // their own: the eviction path selects whole node_models rows, so the + // same revoke blinds it too and it deletes nothing. They are kept as + // the statement of intent, and would catch a regression that reaches a + // working eviction. The message assertions above are what actually + // pins the guard. + var victim nodes.NodeModel + Expect(db.First(&victim, "node_id = ? AND model_name = ?", "n-keep", "victim").Error). + ToNot(HaveOccurred(), "the loaded model must have survived the database outage") + Expect(victim.State).To(Equal("loaded")) + + var keeper nodes.BackendNode + Expect(db.First(&keeper, "id = ?", "n-keep").Error).ToNot(HaveOccurred()) + Expect(keeper.Status).To(Equal(nodes.StatusHealthy), + "a database failure must not change a node's liveness") + }) +})