diff --git a/core/application/distributed.go b/core/application/distributed.go index 40cc9fff3dfa..0519d96e787e 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -356,12 +356,16 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade PrefixConfig: prefixCfg, Pressure: pressure, SharedModels: cfg.Distributed.SharedModels, - // Cap how long a cold load may hold the per-model advisory lock: the - // configured backend.install deadline plus a margin for file staging and - // the remote LoadModel. Derived from the install timeout so raising it - // (for slow links pulling multi-GB images) widens the ceiling too, - // instead of letting the static default cut a legitimately slow load. - ModelLoadCeiling: cfg.Distributed.BackendInstallTimeoutOrDefault() + 10*time.Minute, + ModelLoadTimeout: cfg.Distributed.ModelLoadTimeoutOrDefault(), + // Cap how long a cold load may hold the per-model advisory lock. Derived + // from BOTH configured budgets it has to cover, so raising either the + // install timeout (slow links pulling multi-GB images) or the model load + // timeout (very large checkpoints) widens the ceiling too, instead of + // letting a stale bound cut a legitimately slow load short. + ModelLoadCeiling: nodes.ModelLoadCeilingFor( + cfg.Distributed.BackendInstallTimeoutOrDefault(), + cfg.Distributed.ModelLoadTimeoutOrDefault(), + ), }) // Wire staging-progress broadcasting so file-staging shows up on every diff --git a/core/cli/run.go b/core/cli/run.go index 6c4b1c5a16f1..659977cc1dd0 100644 --- a/core/cli/run.go +++ b/core/cli/run.go @@ -173,6 +173,7 @@ type RunCMD struct { DistributedPrefixCacheTTL string `env:"LOCALAI_DISTRIBUTED_PREFIX_CACHE_TTL" help:"Idle-timeout for prefix-cache index entries; also drives the background eviction cadence (every TTL/2). Default 5m." group:"distributed"` BackendInstallTimeout string `env:"LOCALAI_NATS_BACKEND_INSTALL_TIMEOUT" help:"NATS round-trip timeout for backend.install requests sent to worker nodes (default 15m). Increase for slow links pulling multi-GB images." group:"distributed"` 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:"gRPC deadline for the remote LoadModel call sent to a worker node once its backend is installed and model files are staged (default 5m). Increase for very large checkpoints (multi-tens-of-GB diffusion/video models) whose load and pipeline init exceed 5 minutes. Raising it also widens the cold-load lock ceiling." 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"` @@ -219,6 +220,18 @@ func DefaultUploadPath() string { return filepath.Join(userScopedTempDir(), "upload") } +// parseDistributedDuration parses an operator-supplied duration for a +// distributed timeout knob. The env var and the rejected value are both named +// in the error: these knobs are usually set from a compose file or a K8s +// manifest, where a bare "invalid duration" gives the operator nothing to grep. +func parseDistributedDuration(envVar, raw string) (time.Duration, error) { + d, err := time.ParseDuration(raw) + if err != nil { + return 0, fmt.Errorf("invalid %s %q: %w", envVar, raw, err) + } + return d, nil +} + func (r *RunCMD) Run(ctx *cliContext.Context) error { warnDeprecatedFlags() @@ -326,19 +339,26 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error { opts = append(opts, config.WithStorageSecretKey(r.StorageSecretKey)) } if r.BackendInstallTimeout != "" { - d, err := time.ParseDuration(r.BackendInstallTimeout) + d, err := parseDistributedDuration("LOCALAI_NATS_BACKEND_INSTALL_TIMEOUT", r.BackendInstallTimeout) if err != nil { - return fmt.Errorf("invalid LOCALAI_NATS_BACKEND_INSTALL_TIMEOUT %q: %w", r.BackendInstallTimeout, err) + return err } opts = append(opts, config.WithBackendInstallTimeout(d)) } if r.BackendUpgradeTimeout != "" { - d, err := time.ParseDuration(r.BackendUpgradeTimeout) + d, err := parseDistributedDuration("LOCALAI_NATS_BACKEND_UPGRADE_TIMEOUT", r.BackendUpgradeTimeout) if err != nil { - return fmt.Errorf("invalid LOCALAI_NATS_BACKEND_UPGRADE_TIMEOUT %q: %w", r.BackendUpgradeTimeout, err) + return err } opts = append(opts, config.WithBackendUpgradeTimeout(d)) } + if r.ModelLoadTimeout != "" { + d, err := parseDistributedDuration("LOCALAI_NATS_MODEL_LOAD_TIMEOUT", r.ModelLoadTimeout) + if err != nil { + return err + } + opts = append(opts, config.WithModelLoadTimeout(d)) + } if r.RegistrationToken != "" { opts = append(opts, config.WithRegistrationToken(r.RegistrationToken)) } diff --git a/core/cli/run_distributed_timeouts_test.go b/core/cli/run_distributed_timeouts_test.go new file mode 100644 index 000000000000..5f54c9efc77d --- /dev/null +++ b/core/cli/run_distributed_timeouts_test.go @@ -0,0 +1,23 @@ +package cli + +import ( + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("parseDistributedDuration", func() { + It("parses a valid operator-supplied duration", func() { + d, err := parseDistributedDuration("LOCALAI_NATS_MODEL_LOAD_TIMEOUT", "45m") + Expect(err).ToNot(HaveOccurred()) + Expect(d).To(Equal(45 * time.Minute)) + }) + + It("names the env var and the bad value so a typo is actionable", func() { + _, err := parseDistributedDuration("LOCALAI_NATS_MODEL_LOAD_TIMEOUT", "45mins") + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("LOCALAI_NATS_MODEL_LOAD_TIMEOUT")) + Expect(err.Error()).To(ContainSubstring("45mins")) + }) +}) diff --git a/core/config/distributed_config.go b/core/config/distributed_config.go index 7cb0a7c57ea0..f486ba196d17 100644 --- a/core/config/distributed_config.go +++ b/core/config/distributed_config.go @@ -76,6 +76,12 @@ type DistributedConfig struct { BackendInstallTimeout time.Duration // NATS round-trip timeout for backend.install (default 15m) BackendUpgradeTimeout time.Duration // NATS round-trip timeout for backend.upgrade (default 15m) + // ModelLoadTimeout is the gRPC deadline for the remote LoadModel call the + // router issues once a worker has the backend installed and the model files + // staged. It therefore covers only the backend's own checkpoint load and + // pipeline init, which for a multi-tens-of-GB diffusion/video checkpoint on + // unified memory can far exceed the 5m default. + ModelLoadTimeout time.Duration // gRPC deadline for remote LoadModel (default 5m) MaxUploadSize int64 // Maximum upload body size in bytes (default 50 GB) @@ -142,6 +148,7 @@ func (c DistributedConfig) Validate() error { FlagMCPCIJobTimeout: c.MCPCIJobTimeout, FlagBackendInstallTimeout: c.BackendInstallTimeout, FlagBackendUpgradeTimeout: c.BackendUpgradeTimeout, + FlagModelLoadTimeout: c.ModelLoadTimeout, } { if d < 0 { return fmt.Errorf("%s must not be negative", name) @@ -286,6 +293,12 @@ func WithBackendUpgradeTimeout(d time.Duration) AppOption { } } +func WithModelLoadTimeout(d time.Duration) AppOption { + return func(o *ApplicationConfig) { + o.Distributed.ModelLoadTimeout = d + } +} + var EnableAutoApproveNodes = func(o *ApplicationConfig) { o.Distributed.AutoApproveNodes = true } @@ -342,6 +355,7 @@ const ( FlagMCPCIJobTimeout = "mcp-ci-job-timeout" FlagBackendInstallTimeout = "backend-install-timeout" FlagBackendUpgradeTimeout = "backend-upgrade-timeout" + FlagModelLoadTimeout = "model-load-timeout" ) // Defaults for distributed timeouts. @@ -355,6 +369,7 @@ const ( DefaultMCPCIJobTimeout = 10 * time.Minute DefaultBackendInstallTimeout = 15 * time.Minute DefaultBackendUpgradeTimeout = 15 * time.Minute + DefaultModelLoadTimeout = 5 * time.Minute ) // DefaultMaxUploadSize is the default maximum upload body size (50 GB). @@ -408,6 +423,11 @@ func (c DistributedConfig) BackendUpgradeTimeoutOrDefault() time.Duration { return cmp.Or(c.BackendUpgradeTimeout, DefaultBackendUpgradeTimeout) } +// ModelLoadTimeoutOrDefault returns the configured timeout or the default. +func (c DistributedConfig) ModelLoadTimeoutOrDefault() time.Duration { + return cmp.Or(c.ModelLoadTimeout, DefaultModelLoadTimeout) +} + // MCPToolTimeoutOrDefault returns the configured timeout or the default. func (c DistributedConfig) MCPToolTimeoutOrDefault() time.Duration { return cmp.Or(c.MCPToolTimeout, DefaultMCPToolTimeout) diff --git a/core/config/distributed_config_test.go b/core/config/distributed_config_test.go index 377c397dee95..ec7fbe8dc7e9 100644 --- a/core/config/distributed_config_test.go +++ b/core/config/distributed_config_test.go @@ -33,6 +33,18 @@ var _ = Describe("DistributedConfig backend NATS timeouts", func() { Expect(c.BackendUpgradeTimeoutOrDefault()).To(Equal(30 * time.Minute)) }) }) + + Context("ModelLoadTimeoutOrDefault", func() { + It("returns 5 minutes when unset so existing clusters keep today's behaviour", func() { + c := config.DistributedConfig{} + Expect(c.ModelLoadTimeoutOrDefault()).To(Equal(5 * time.Minute)) + }) + + It("returns the configured value when set", func() { + c := config.DistributedConfig{ModelLoadTimeout: 45 * time.Minute} + Expect(c.ModelLoadTimeoutOrDefault()).To(Equal(45 * time.Minute)) + }) + }) }) var _ = Describe("DistributedConfig flag-name constants", func() { @@ -53,6 +65,7 @@ var _ = Describe("DistributedConfig flag-name constants", func() { 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"), + Entry("model load timeout", config.FlagModelLoadTimeout, "model-load-timeout"), ) }) @@ -80,6 +93,18 @@ var _ = Describe("DistributedConfig.Validate negative-duration errors", func() { Expect(err.Error()).To(ContainSubstring(config.FlagBackendUpgradeTimeout)) }) + It("rejects a negative ModelLoadTimeout with the flag name in the error", func() { + c := config.DistributedConfig{ + Enabled: true, + NatsURL: "nats://localhost:4222", + ModelLoadTimeout: -1 * time.Second, + } + err := c.Validate() + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring(config.FlagModelLoadTimeout)) + Expect(err.Error()).To(ContainSubstring("must not be negative")) + }) + It("accepts all-zero durations as valid (defaults apply)", func() { c := config.DistributedConfig{ Enabled: true, diff --git a/core/services/nodes/router.go b/core/services/nodes/router.go index 01100123401a..145880c7f174 100644 --- a/core/services/nodes/router.go +++ b/core/services/nodes/router.go @@ -18,9 +18,12 @@ import ( "github.com/mudler/LocalAI/pkg/distributedhdr" grpc "github.com/mudler/LocalAI/pkg/grpc" pb "github.com/mudler/LocalAI/pkg/grpc/proto" + "github.com/mudler/LocalAI/pkg/model" "github.com/mudler/LocalAI/pkg/vram" "github.com/mudler/xlog" "golang.org/x/sync/singleflight" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" "google.golang.org/protobuf/proto" "gorm.io/gorm" "gorm.io/gorm/clause" @@ -72,9 +75,46 @@ type SmartRouterOptions struct { // attempt (node selection -> backend install -> file staging -> LoadModel) // may run while holding the per-model advisory lock. It backstops every // sub-step's own timeout so a wedged worker can never pin the lock - and - // every other replica's request for that model - indefinitely. Zero selects - // defaultModelLoadCeiling. + // every other replica's request for that model - indefinitely. Zero derives + // it from the default install budget and ModelLoadTimeout via + // ModelLoadCeilingFor. ModelLoadCeiling time.Duration + // ModelLoadTimeout is the gRPC deadline for the remote LoadModel call, which + // runs after the backend install and file staging have already completed and + // so covers only the worker backend's checkpoint load and pipeline init. + // Zero selects config.DefaultModelLoadTimeout. Operators raise it via + // LOCALAI_NATS_MODEL_LOAD_TIMEOUT for very large checkpoints; ModelLoadCeiling + // must be widened in step (see ModelLoadCeilingFor) or the ceiling clips the + // longer deadline before it can be used. + ModelLoadTimeout time.Duration +} + +// modelLoadStagingMargin is the slack ModelLoadCeilingFor adds on top of the +// install + remote-load budgets to cover the steps that have no deadline of +// their own: node selection, replica allocation and model file staging. +const modelLoadStagingMargin = 5 * time.Minute + +// minModelLoadCeiling floors the derived ceiling at the historical 25m constant +// so shrinking the install or load budgets can never tighten the hold ceiling +// below what clusters relied on before it became derived. +const minModelLoadCeiling = 25 * time.Minute + +// ModelLoadCeilingFor derives the cold-load hold ceiling from the budgets it has +// to cover. The ceiling bounds how long one cold load may hold the per-model +// advisory lock; it must comfortably exceed the slowest legitimate load (backend +// install + staging + remote LoadModel) so it only ever fires when a step is +// genuinely wedged - e.g. a worker that died mid-install. Deriving it means +// raising LOCALAI_NATS_MODEL_LOAD_TIMEOUT for an 80 GB video checkpoint actually +// takes effect, instead of being cut short by a constant that silently went +// stale. Non-positive inputs fall back to their package defaults. +func ModelLoadCeilingFor(installTimeout, loadTimeout time.Duration) time.Duration { + if installTimeout <= 0 { + installTimeout = config.DefaultBackendInstallTimeout + } + if loadTimeout <= 0 { + loadTimeout = config.DefaultModelLoadTimeout + } + return max(installTimeout+loadTimeout+modelLoadStagingMargin, minModelLoadCeiling) } // SmartRouter routes inference requests to the best available backend node. @@ -111,15 +151,11 @@ type SmartRouter struct { // modelLoadCeiling bounds how long a cold load may hold the per-model // advisory lock (see SmartRouterOptions.ModelLoadCeiling). modelLoadCeiling time.Duration + // modelLoadTimeout is the deadline for the remote LoadModel gRPC call + // (see SmartRouterOptions.ModelLoadTimeout). + modelLoadTimeout time.Duration } -// defaultModelLoadCeiling is the fallback hold ceiling for a cold model load. -// It must comfortably exceed the slowest legitimate load - a multi-GB backend -// install (DefaultBackendInstallTimeout, 15m) plus staging and the remote -// LoadModel (5m) - so it never cuts a real load short; it only ever fires when -// a step is genuinely wedged (e.g. a worker that died mid-install). -const defaultModelLoadCeiling = 25 * time.Minute - // probeCacheTTL is how long a successful gRPC HealthCheck on a backend is // trusted before the next request re-probes. Matches healthCheckTTL in // pkg/model/model.go so the single-process and distributed paths share a @@ -134,9 +170,13 @@ func NewSmartRouter(registry ModelRouter, opts SmartRouterOptions) *SmartRouter if factory == nil { factory = &tokenClientFactory{token: opts.AuthToken} } + loadTimeout := opts.ModelLoadTimeout + if loadTimeout <= 0 { + loadTimeout = config.DefaultModelLoadTimeout + } ceiling := opts.ModelLoadCeiling if ceiling <= 0 { - ceiling = defaultModelLoadCeiling + ceiling = ModelLoadCeilingFor(config.DefaultBackendInstallTimeout, loadTimeout) } return &SmartRouter{ registry: registry, @@ -153,6 +193,7 @@ func NewSmartRouter(registry ModelRouter, opts SmartRouterOptions) *SmartRouter pressure: opts.Pressure, sharedModels: opts.SharedModels, modelLoadCeiling: ceiling, + modelLoadTimeout: loadTimeout, } } @@ -250,11 +291,21 @@ func (r *SmartRouter) scheduleAndLoad(ctx context.Context, backendType, tracking if loadOpts != nil { xlog.Info("Loading model on remote node", "node", node.Name, "model", modelName, "addr", backendAddr) - loadCtx, cancel := context.WithTimeout(ctx, 5*time.Minute) + loadCtx, cancel := context.WithTimeout(ctx, r.modelLoadTimeout) defer cancel() res, err := client.LoadModel(loadCtx, loadOpts) if err != nil { + // A gRPC deadline only cancels the CLIENT side of the call. A + // backend blocked in a synchronous weight load never observes its + // cancelled handler context, so it keeps loading with nobody + // waiting: an 83GB checkpoint was seen still downloading 30 + // minutes past the client timeout, and each retry stacked another + // multi-GB loader process on the worker. Reap the replica we just + // abandoned before handing the failure back. + if loadAbandonedOnWorker(err) { + r.reapAbandonedLoad(node, trackingKey, replicaIndex) + } return nil, fmt.Errorf("loading model %s on node %s: %w", modelName, node.Name, err) } if !res.Success { @@ -285,6 +336,51 @@ func (r *SmartRouter) scheduleAndLoad(ctx context.Context, backendType, tracking return &scheduleLoadResult{Node: node, Client: client, BackendAddr: backendAddr, ReplicaIndex: replicaIndex}, nil } +// loadAbandonedOnWorker reports whether a failed remote LoadModel left the +// worker process still running the load. +// +// Only a deadline or a cancellation qualifies: those are OUR side giving up +// while the backend keeps going. Every other failure (an unsupported model, a +// bad option, an OOM) is the backend answering, which means its handler +// returned and the process is idle — stopping it there would throw away a +// warm process and its downloaded weights for the next attempt. +func loadAbandonedOnWorker(err error) bool { + if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) { + return true + } + // The deadline usually reaches us as a gRPC status rather than a wrapped + // context error. errors.As (not status.Code) so a wrapped status is still + // recognised. + var st interface{ GRPCStatus() *status.Status } + if errors.As(err, &st) { + switch st.GRPCStatus().Code() { + case codes.DeadlineExceeded, codes.Canceled: + return true + } + } + return false +} + +// reapAbandonedLoad tells the worker to stop the one replica whose load we +// just abandoned. The exact `modelID#replicaIndex` process key matters: the +// bare model ID would stop every replica of the model on that node, including +// healthy ones serving traffic. +// +// Best-effort by design — the caller is waiting on the load failure, and a +// failed reap must be visible in the logs, never substituted for it. +func (r *SmartRouter) reapAbandonedLoad(node *BackendNode, trackingKey string, replicaIndex int) { + if r.unloader == nil { + return + } + processKey := model.BackendProcessKey(trackingKey, replicaIndex) + xlog.Warn("Reaping abandoned model load on worker", + "node", node.Name, "model", trackingKey, "replica", replicaIndex) + if err := r.unloader.StopBackend(node.ID, processKey); err != nil { + xlog.Warn("Failed to reap abandoned model load; the worker may still be loading", + "node", node.Name, "backend", processKey, "error", err) + } +} + // ScheduleAndLoadModel implements ModelScheduler for the reconciler. // It retrieves stored model options from an existing replica and performs the // full load sequence (stage files, LoadModel, SetNodeModel) on a new node. diff --git a/core/services/nodes/router_load_timeout_test.go b/core/services/nodes/router_load_timeout_test.go new file mode 100644 index 000000000000..d6ba09fa0b51 --- /dev/null +++ b/core/services/nodes/router_load_timeout_test.go @@ -0,0 +1,127 @@ +package nodes + +import ( + "context" + "errors" + "sync" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/services/messaging" + grpc "github.com/mudler/LocalAI/pkg/grpc" + pb "github.com/mudler/LocalAI/pkg/grpc/proto" + ggrpc "google.golang.org/grpc" +) + +// deadlineBackend records the deadline of the context handed to LoadModel so a +// spec can assert what budget the remote gRPC load actually got. A hardcoded +// deadline in scheduleAndLoad is invisible to every other assertion in this +// package — only the context the client sees proves it. +type deadlineBackend struct { + grpc.Backend // embedded: unused methods panic if called + + mu sync.Mutex + deadline time.Time + hadOne bool +} + +func (b *deadlineBackend) HealthCheck(_ context.Context) (bool, error) { return true, nil } + +func (b *deadlineBackend) IsBusy() bool { return false } + +func (b *deadlineBackend) LoadModel(ctx context.Context, _ *pb.ModelOptions, _ ...ggrpc.CallOption) (*pb.Result, error) { + b.mu.Lock() + defer b.mu.Unlock() + b.deadline, b.hadOne = ctx.Deadline() + return &pb.Result{Success: true}, nil +} + +// budget returns how much time LoadModel was given, rounded to the nearest +// second so scheduling jitter between the router's WithTimeout and the stub's +// clock read doesn't make the assertion flaky. +func (b *deadlineBackend) budget() time.Duration { + b.mu.Lock() + defer b.mu.Unlock() + Expect(b.hadOne).To(BeTrue(), "LoadModel must be called with a deadline") + return time.Until(b.deadline).Round(time.Second) +} + +type deadlineClientFactory struct{ client *deadlineBackend } + +func (f *deadlineClientFactory) NewClient(_ string, _ bool) grpc.Backend { return f.client } + +var _ = Describe("remote LoadModel deadline", func() { + var ( + reg *fakeModelRouter + backend *deadlineBackend + factory *deadlineClientFactory + unloader *fakeUnloader + ) + + BeforeEach(func() { + reg = &fakeModelRouter{findAndLockErr: errors.New("not found")} + reg.findIdleNode = &BackendNode{ID: "n1", Name: "worker", Address: "10.0.0.1:50051"} + backend = &deadlineBackend{} + factory = &deadlineClientFactory{client: backend} + unloader = &fakeUnloader{ + installReply: &messaging.BackendInstallReply{Success: true, Address: "10.0.0.1:9001"}, + } + }) + + route := func(router *SmartRouter) { + _, err := router.Route(context.Background(), "big-model", "models/big.gguf", "llama-cpp", + &pb.ModelOptions{Model: "models/big.gguf"}, false) + Expect(err).ToNot(HaveOccurred()) + } + + It("gives LoadModel the configured budget instead of the hardcoded 5 minutes", func() { + router := NewSmartRouter(reg, SmartRouterOptions{ + Unloader: unloader, + ClientFactory: factory, + ModelLoadTimeout: 45 * time.Minute, + // Ceiling must not be the thing that clips the deadline here. + ModelLoadCeiling: 2 * time.Hour, + }) + route(router) + Expect(backend.budget()).To(Equal(45 * time.Minute)) + }) + + It("falls back to 5 minutes when no load timeout is configured", func() { + router := NewSmartRouter(reg, SmartRouterOptions{ + Unloader: unloader, + ClientFactory: factory, + }) + route(router) + Expect(backend.budget()).To(Equal(5 * time.Minute)) + }) +}) + +var _ = Describe("ModelLoadCeilingFor", func() { + // The ceiling bounds how long a cold load may hold the per-model advisory + // lock. It is derived from the sub-step budgets it must cover, so raising + // the load timeout for a huge checkpoint cannot be silently clipped by a + // stale constant. + It("keeps today's 25m behaviour at the default install + load budgets", func() { + Expect(ModelLoadCeilingFor(config.DefaultBackendInstallTimeout, config.DefaultModelLoadTimeout)). + To(Equal(25 * time.Minute)) + }) + + It("widens when the load timeout is raised", func() { + Expect(ModelLoadCeilingFor(15*time.Minute, 45*time.Minute)).To(Equal(65 * time.Minute)) + }) + + It("widens when the install timeout is raised", func() { + Expect(ModelLoadCeilingFor(40*time.Minute, 5*time.Minute)).To(Equal(50 * time.Minute)) + }) + + It("never drops below the 25m floor when the budgets are shrunk", func() { + Expect(ModelLoadCeilingFor(1*time.Minute, 1*time.Minute)).To(Equal(25 * time.Minute)) + }) + + It("treats non-positive budgets as their defaults", func() { + Expect(ModelLoadCeilingFor(0, 0)).To(Equal(25 * time.Minute)) + }) +}) diff --git a/core/services/nodes/router_reap_load_test.go b/core/services/nodes/router_reap_load_test.go new file mode 100644 index 000000000000..4c8e7cf8bb45 --- /dev/null +++ b/core/services/nodes/router_reap_load_test.go @@ -0,0 +1,126 @@ +package nodes + +import ( + "context" + "errors" + "sync" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/messaging" + grpc "github.com/mudler/LocalAI/pkg/grpc" + pb "github.com/mudler/LocalAI/pkg/grpc/proto" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + ggrpc "google.golang.org/grpc" +) + +// failingLoadBackend returns a canned error from LoadModel so a spec can drive +// scheduleAndLoad down its failure path. The gRPC deadline that fires in +// production only cancels the client side of the call, so the error the router +// sees is the only signal that a worker was left loading. +type failingLoadBackend struct { + grpc.Backend // embedded: unused methods panic if called + + mu sync.Mutex + err error + seen int +} + +func (b *failingLoadBackend) HealthCheck(_ context.Context) (bool, error) { return true, nil } + +func (b *failingLoadBackend) IsBusy() bool { return false } + +func (b *failingLoadBackend) LoadModel(_ context.Context, _ *pb.ModelOptions, _ ...ggrpc.CallOption) (*pb.Result, error) { + b.mu.Lock() + defer b.mu.Unlock() + b.seen++ + return nil, b.err +} + +type failingClientFactory struct{ client *failingLoadBackend } + +func (f *failingClientFactory) NewClient(_ string, _ bool) grpc.Backend { return f.client } + +// replicaSlotRouter pins the replica slot scheduleAndLoad allocates so a spec +// can assert the reaped process key carries the real index, not a hardcoded 0. +type replicaSlotRouter struct { + *fakeModelRouter + replica int +} + +func (r *replicaSlotRouter) NextFreeReplicaIndex(_ context.Context, _, _ string, _ int) (int, error) { + return r.replica, nil +} + +var _ = Describe("reaping an abandoned remote load", func() { + var ( + reg *replicaSlotRouter + backend *failingLoadBackend + unloader *fakeUnloader + ) + + BeforeEach(func() { + base := &fakeModelRouter{findAndLockErr: errors.New("not found")} + base.findIdleNode = &BackendNode{ID: "n1", Name: "worker", Address: "10.0.0.1:50051", MaxReplicasPerModel: 4} + reg = &replicaSlotRouter{fakeModelRouter: base, replica: 2} + backend = &failingLoadBackend{} + unloader = &fakeUnloader{ + installReply: &messaging.BackendInstallReply{Success: true, Address: "10.0.0.1:9001"}, + } + }) + + route := func() error { + router := NewSmartRouter(reg, SmartRouterOptions{ + Unloader: unloader, + ClientFactory: &failingClientFactory{client: backend}, + ModelLoadTimeout: time.Minute, + ModelLoadCeiling: time.Hour, + }) + _, err := router.Route(context.Background(), "big-model", "models/big.gguf", "llama-cpp", + &pb.ModelOptions{Model: "models/big.gguf"}, false) + return err + } + + It("stops exactly the abandoned replica when LoadModel hits its deadline", func() { + backend.err = status.Error(codes.DeadlineExceeded, "context deadline exceeded") + + err := route() + + Expect(err).To(HaveOccurred()) + Expect(unloader.stopCalls).To(ConsistOf("n1:big-model#2"), + "the abandoned replica must be reaped by its exact worker process key") + }) + + It("still returns the load failure to the caller, not the stop result", func() { + backend.err = status.Error(codes.DeadlineExceeded, "context deadline exceeded") + unloader.stopErr = errors.New("nats publish failed") + + err := route() + + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("context deadline exceeded")) + Expect(err.Error()).ToNot(ContainSubstring("nats publish failed"), + "a failed reap must be logged, never substituted for the load error") + }) + + It("reaps when the deadline surfaces as a wrapped context error", func() { + backend.err = errors.Join(errors.New("loading model"), context.DeadlineExceeded) + + Expect(route()).To(HaveOccurred()) + Expect(unloader.stopCalls).To(ConsistOf("n1:big-model#2")) + }) + + It("does not reap when the backend rejected the model outright", func() { + // A clean application error means the handler returned: the worker is + // idle, not stuck mid-load, and stopping it would throw away a usable + // process (and its downloaded weights) for the next attempt. + backend.err = status.Error(codes.InvalidArgument, "unsupported model architecture") + + Expect(route()).To(HaveOccurred()) + Expect(unloader.stopCalls).To(BeEmpty()) + }) +}) diff --git a/docs/content/features/distributed-mode.md b/docs/content/features/distributed-mode.md index 6dfddff2db33..cdeea890e115 100644 --- a/docs/content/features/distributed-mode.md +++ b/docs/content/features/distributed-mode.md @@ -72,8 +72,11 @@ The frontend is a standard LocalAI instance with distributed mode enabled. These | `--auth-database-url` | `LOCALAI_AUTH_DATABASE_URL` | *(required)* | PostgreSQL connection URL | | `--backend-install-timeout` | `LOCALAI_NATS_BACKEND_INSTALL_TIMEOUT` | `15m` | How long the frontend waits for a worker to acknowledge a backend install before considering the request stalled. Raise it when workers pull large backend images over slow links. If a worker takes longer than this, the operation shows as "still installing in background" in the admin UI and clears once the worker finishes. | | `--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` | `5m` | Deadline for the `LoadModel` gRPC call the frontend issues to a worker. It starts *after* the backend is installed and the model files are staged, so it covers only the worker backend's own checkpoint load and pipeline init. Raise it for very large checkpoints: a multi-tens-of-GB diffusion or video model on unified memory routinely needs more than 5 minutes, and the load then fails with `rpc error: code = DeadlineExceeded` at exactly the timeout. Raising this also widens the cold-load lock ceiling below, so the two knobs never fight. | | `--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 router also bounds how long a single cold load may hold the per-model advisory lock, so a worker that dies mid-install cannot pin every other replica's request for that model. That ceiling is **derived**, not configured: `max(backend-install-timeout + model-load-timeout + 5m, 25m)`. With the defaults that is `15m + 5m + 5m = 25m`, matching previous releases; raising either timeout widens the ceiling in step, so a longer load deadline is never clipped by it. + ### NATS JWT authentication (recommended for production) By default, NATS connections are anonymous: any client that can reach port `4222` may publish control-plane subjects such as `nodes..backend.install`. Enable JWT auth to scope workers to their own node subjects and give the frontend a dedicated service credential. diff --git a/pkg/model/backend_log_store.go b/pkg/model/backend_log_store.go index 808b0d75c831..a0988790ff9e 100644 --- a/pkg/model/backend_log_store.go +++ b/pkg/model/backend_log_store.go @@ -2,6 +2,7 @@ package model import ( "sort" + "strconv" "strings" "sync" "time" @@ -15,6 +16,17 @@ import ( // package free of CLI imports. const replicaSeparator = "#" +// BackendProcessKey builds the worker supervisor's process key for one replica +// of a model (e.g. "qwen3-0.6b#0"). The worker keys its process map by this +// string and resolves an exact key to exactly that replica, so any caller that +// wants to address a single replica over NATS must format it identically — +// a mismatch degrades silently into "no process found" and leaks the process. +// This package is the lowest common dependency of the router and the worker, +// so the format lives here instead of being hand-rolled at each call site. +func BackendProcessKey(modelID string, replicaIndex int) string { + return modelID + replicaSeparator + strconv.Itoa(replicaIndex) +} + // BackendLogLine represents a single line of output from a backend process. type BackendLogLine struct { Timestamp time.Time `json:"timestamp"`