diff --git a/core/application/distributed.go b/core/application/distributed.go index 750b73b02..29daf36a9 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -10,6 +10,7 @@ import ( "time" "github.com/google/uuid" + "github.com/mudler/LocalAI/core/backend" "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/services/agents" "github.com/mudler/LocalAI/core/services/distributed" @@ -377,6 +378,22 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade cfg.Distributed.BackendInstallTimeoutOrDefault(), cfg.Distributed.ModelLoadTimeoutOrDefault(), ), + // Re-derive managed companion options from the model's current config when + // the reconciler replays a stored ModelOptions blob (which never runs + // grpcModelOpts). Without this, a replica scaled up from a blob captured + // before the companion resolved loads without the companion option and the + // backend fetches its own wrong default. nil-safe: an unknown model or a + // config with no companions yields no options. + CompanionOptionsFor: func(modelName string) []string { + if configLoader == nil { + return nil + } + modelCfg, ok := configLoader.GetModelConfig(modelName) + if !ok { + return nil + } + return backend.CompanionArtifactOptions(modelCfg) + }, }) // Wire staging-progress broadcasting so file-staging shows up on every diff --git a/core/services/nodes/file_staging_loadmodel_test.go b/core/services/nodes/file_staging_loadmodel_test.go new file mode 100644 index 000000000..1f8ff9ec9 --- /dev/null +++ b/core/services/nodes/file_staging_loadmodel_test.go @@ -0,0 +1,58 @@ +package nodes + +import ( + "context" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + grpc "github.com/mudler/LocalAI/pkg/grpc" + pb "github.com/mudler/LocalAI/pkg/grpc/proto" + ggrpc "google.golang.org/grpc" +) + +// capturingLoadBackend records the ModelOptions handed to LoadModel, standing in +// for the remote gRPC backend the FileStagingClient forwards to. +type capturingLoadBackend struct { + grpc.Backend + got *pb.ModelOptions +} + +func (b *capturingLoadBackend) LoadModel(_ context.Context, in *pb.ModelOptions, _ ...ggrpc.CallOption) (*pb.Result, error) { + b.got = in + return &pb.Result{Success: true}, nil +} + +// FileStagingClient is the concrete type of `client` at router.go's +// client.LoadModel call site (buildClientForAddr wraps the gRPC backend with it +// whenever a file stager is configured, i.e. every distributed load). This pins +// the boundary the companion-option investigation kept returning to: the wrapper +// must forward ModelOptions.Options (the channel a managed companion option like +// base_model rides on) to the backend byte-for-byte. If a companion option is +// present in what the controller sends here but absent at the Python backend, +// this test proves the loss is NOT in this Go handoff. +var _ = Describe("FileStagingClient LoadModel option pass-through", func() { + It("forwards ModelOptions.Options to the backend unchanged", func(ctx SpecContext) { + backend := &capturingLoadBackend{} + client := NewFileStagingClient(backend, &fakeFileStager{}, "worker-1") + + sent := &pb.ModelOptions{ + Model: "longcat-video-avatar-1.5", + ModelPath: "/models/longcat-video-avatar-1.5", + Options: []string{ + "attention_backend:sdpa", + "use_distill:true", + "base_model:.artifacts/huggingface/deadbeef/snapshot", + }, + } + + _, err := client.LoadModel(ctx, sent) + Expect(err).NotTo(HaveOccurred()) + Expect(backend.got).NotTo(BeNil()) + Expect(backend.got.Options).To(Equal(sent.Options)) + Expect(backend.got.Options).To(ContainElement("base_model:.artifacts/huggingface/deadbeef/snapshot")) + // The nested staged ModelPath the backend joins the relative option against + // must survive too. + Expect(backend.got.ModelPath).To(Equal("/models/longcat-video-avatar-1.5")) + }) +}) diff --git a/core/services/nodes/router.go b/core/services/nodes/router.go index 5b6aed953..0d5b03b53 100644 --- a/core/services/nodes/router.go +++ b/core/services/nodes/router.go @@ -105,6 +105,12 @@ type SmartRouterOptions struct { // arriving, so a peer trickling bytes forever cannot pin the advisory lock // indefinitely. Zero selects modelLoadAbsoluteMax (24h). ModelLoadAbsoluteMax time.Duration + // CompanionOptionsFor re-derives a model's managed companion options from its + // current config (backend.CompanionArtifactOptions). The reconciler replays a + // stored ModelOptions blob that never runs grpcModelOpts; without this a + // companion resolved after the blob was captured never reaches a scaled-up + // replica. nil disables re-derivation. + CompanionOptionsFor func(modelName string) []string } // modelLoadStagingMargin is the slack ModelLoadCeilingFor adds on top of the @@ -187,6 +193,13 @@ type SmartRouter struct { // hard countdown into a progress-extended hold (see load_deadline.go). stagingStallWindow time.Duration modelLoadAbsoluteMax time.Duration + // companionOptionsFor re-derives a model's managed companion options from its + // CURRENT config. The reconciler's scale-up path replays a stored ModelOptions + // proto blob that never runs grpcModelOpts, so a companion resolved AFTER that + // blob was captured would otherwise never reach the replica. nil disables the + // re-derivation (single-node and tests that don't wire it). See + // SmartRouterOptions.CompanionOptionsFor. + companionOptionsFor func(modelName string) []string } // probeCacheTTL is how long a successful gRPC HealthCheck on a backend is @@ -238,6 +251,7 @@ func NewSmartRouter(registry ModelRouter, opts SmartRouterOptions) *SmartRouter // ceiling, so nothing to normalize here. stagingStallWindow: opts.StagingStallWindow, modelLoadAbsoluteMax: opts.ModelLoadAbsoluteMax, + companionOptionsFor: opts.CompanionOptionsFor, } } @@ -448,6 +462,39 @@ func (r *SmartRouter) reapAbandonedLoad(node *BackendNode, trackingKey string, r } } +// applyCompanionOptions merges a model's freshly re-derived managed companion +// options into opts.Options, adding only those the blob does not already carry. +// Merge, not replace: an option the blob already has (including an author-pinned +// one) wins, so this only ever ADDS a missing companion. No-op when no resolver +// is wired (single-node) or opts is nil. +func (r *SmartRouter) applyCompanionOptions(modelName string, opts *pb.ModelOptions) { + if r.companionOptionsFor == nil || opts == nil { + return + } + present := make(map[string]struct{}, len(opts.Options)) + for _, opt := range opts.Options { + if name, _, found := strings.Cut(opt, ":"); found { + present[name] = struct{}{} + } + } + for _, opt := range r.companionOptionsFor(modelName) { + name, _, found := strings.Cut(opt, ":") + if !found { + continue + } + if _, exists := present[name]; exists { + continue + } + // Debug (not info): a correctly-captured blob already carries the option, + // so this only fires on the rarer stale-blob recovery — worth surfacing to + // an operator tracing a companion, but not on every replay. + xlog.Debug("reconciler injecting companion option missing from the stored blob", + "model", modelName, "option", opt) + opts.Options = append(opts.Options, opt) + present[name] = struct{}{} + } +} + // 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. @@ -470,6 +517,12 @@ func (r *SmartRouter) ScheduleAndLoadModel(ctx context.Context, modelName string return nil, fmt.Errorf("unmarshalling stored model options for %s: %w", modelName, err) } + // Re-derive managed companion options from the CURRENT config and merge in any + // the stored blob is missing. This path replays a proto blob captured at first + // load and never runs grpcModelOpts, so a companion resolved after the blob was + // captured would otherwise never reach the scaled-up replica. + r.applyCompanionOptions(modelName, &modelOpts) + // initialInFlight=0: reconciler is pre-loading, not serving a request. // scheduleAndLoad picks both the node and the replica slot internally. result, err := r.scheduleAndLoad(ctx, backendType, modelName, modelName, &modelOpts, false, 0) diff --git a/core/services/nodes/router_companion_replay_test.go b/core/services/nodes/router_companion_replay_test.go new file mode 100644 index 000000000..746841a55 --- /dev/null +++ b/core/services/nodes/router_companion_replay_test.go @@ -0,0 +1,61 @@ +package nodes + +import ( + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + pb "github.com/mudler/LocalAI/pkg/grpc/proto" +) + +// The reconciler's scale-up path replays a stored ModelOptions proto blob and +// never runs grpcModelOpts, so a managed companion option (e.g. longcat-video's +// base_model) is present only if it was captured in that blob. A blob stored +// before the companion resolved permanently lacks it, and the replica loads +// without the companion — the backend then fetches its own wrong default and +// fails "base_model must point to a LongCat-Video checkpoint". applyCompanionOptions +// re-derives the companion from the CURRENT config and injects any that are +// missing, so a remote replica gets the same option a fresh grpcModelOpts would. +var _ = Describe("reconciler companion option re-derivation", func() { + It("injects a companion option the stored blob is missing", func() { + router := &SmartRouter{ + companionOptionsFor: func(_ string) []string { + return []string{"base_model:.artifacts/huggingface/deadbeef/snapshot"} + }, + } + // A blob captured before the companion resolved: no base_model. + opts := &pb.ModelOptions{Options: []string{"attention_backend:sdpa"}} + + router.applyCompanionOptions("longcat-video-avatar-1.5", opts) + + Expect(opts.Options).To(ContainElement("base_model:.artifacts/huggingface/deadbeef/snapshot")) + Expect(opts.Options).To(ContainElement("attention_backend:sdpa")) + }) + + It("never overrides a companion option the blob already carries", func() { + router := &SmartRouter{ + companionOptionsFor: func(_ string) []string { + return []string{"base_model:.artifacts/huggingface/managed/snapshot"} + }, + } + // An author-pinned (or already-captured) base_model must win. + opts := &pb.ModelOptions{Options: []string{"base_model:/opt/checkouts/LongCat-Video"}} + + router.applyCompanionOptions("longcat-video-avatar-1.5", opts) + + Expect(opts.Options).To(ContainElement("base_model:/opt/checkouts/LongCat-Video")) + base := 0 + for _, o := range opts.Options { + if len(o) >= len("base_model:") && o[:len("base_model:")] == "base_model:" { + base++ + } + } + Expect(base).To(Equal(1)) + }) + + It("is a no-op when no resolver is wired (single-node)", func() { + router := &SmartRouter{} + opts := &pb.ModelOptions{Options: []string{"attention_backend:sdpa"}} + router.applyCompanionOptions("m", opts) + Expect(opts.Options).To(Equal([]string{"attention_backend:sdpa"})) + }) +})