diff --git a/core/application/distributed.go b/core/application/distributed.go index 750b73b02437..29daf36a96e4 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/backend/companion_options_test.go b/core/backend/companion_options_test.go index c4b6f3953a66..123fa597c73d 100644 --- a/core/backend/companion_options_test.go +++ b/core/backend/companion_options_test.go @@ -113,11 +113,34 @@ var _ = Describe("companion artifact backend options", func() { Expect(opts.Options).To(Equal([]string{"attention_backend:sdpa"})) }) - It("skips a companion that has not been resolved yet", func() { + It("names the source repository when the companion is not resolved yet", func() { + // A companion that reaches load time WITHOUT a resolved snapshot must not + // vanish silently: emitting no option lets the backend fall back to its own + // hardcoded default, which is how a distributed longcat-video worker ended + // up trying to load the wrong base model and failing "base_model must point + // to a LongCat-Video checkpoint". Naming the DECLARED repository instead + // points the backend at the artifact the config actually asked for. The + // snapshot path (the staged, no-download fast path) is still preferred + // whenever the companion IS resolved. cfg := configWithCompanion() cfg.Artifacts[1].Resolved = nil opts := grpcModelOpts(cfg, "/models") - _, found := optionValue(opts.Options, "base_model") - Expect(found).To(BeFalse()) + + value, found := optionValue(opts.Options, "base_model") + Expect(found).To(BeTrue()) + Expect(value).To(Equal("meituan-longcat/LongCat-Video")) + // The fallback is a repo reference, never a models-relative snapshot path. + Expect(value).ToNot(ContainSubstring(".artifacts")) + }) + + It("prefers the resolved snapshot path over the source repository", func() { + opts := grpcModelOpts(configWithCompanion(), "/models") + value, found := optionValue(opts.Options, "base_model") + Expect(found).To(BeTrue()) + expected, err := modelartifacts.RelativeSnapshotPath(companionKey) + Expect(err).NotTo(HaveOccurred()) + Expect(value).To(Equal(expected)) + // The resolved fast path must never degrade to a bare repo id. + Expect(value).ToNot(Equal("meituan-longcat/LongCat-Video")) }) }) diff --git a/core/backend/options.go b/core/backend/options.go index 72f49f1eda47..e0a2ced30475 100644 --- a/core/backend/options.go +++ b/core/backend/options.go @@ -294,6 +294,13 @@ func EffectiveBatchSize(c config.ModelConfig) int { // // An option the author set explicitly always wins: pinning a companion to a // local checkout has to beat the managed snapshot. +// +// A companion that is declared but NOT resolved falls back to its source +// repository id rather than being dropped: a dropped companion is invisible to +// the backend, which then loads its own hardcoded default and fails far away +// from the cause. The repo-id fallback trades the staging fast path (the weights +// are fetched on the worker) for correctness, and logs a warning so the missing +// controller-side resolution is diagnosable. func withCompanionArtifactOptions(options []string, artifacts []modelartifacts.Spec) []string { configured := make(map[string]struct{}, len(options)) for _, option := range options { @@ -301,26 +308,71 @@ func withCompanionArtifactOptions(options []string, artifacts []modelartifacts.S configured[name] = struct{}{} } } - // Copy before appending: opts.Options would otherwise share (and could // reallocate away from) the config's own slice. - combined := slices.Clone(options) + return append(slices.Clone(options), ManagedCompanionOptions(artifacts, configured)...) +} + +// ManagedCompanionOptions synthesizes the ":" option for +// every companion artifact whose name is not already in `configured`, so callers +// can append them to a ModelOptions.Options slice. It is exported because the +// distributed reconciler replays a stored proto blob that never runs +// grpcModelOpts, and it must re-derive the same companion options from the +// current config so a scaled-up replica gets them too. +// +// A resolved companion is surfaced as its staged, models-relative snapshot +// directory (the no-download fast path a remote worker resolves under its own +// ModelPath). A companion that reached load time WITHOUT a resolved snapshot +// falls back to its source repository id rather than being dropped: dropping it +// is invisible to the backend, which then loads its OWN hardcoded default and +// fails far from the cause (a distributed longcat-video worker fetched the wrong +// base model and failed "base_model must point to a LongCat-Video checkpoint"). +func ManagedCompanionOptions(artifacts []modelartifacts.Spec, configured map[string]struct{}) []string { + var out []string for _, artifact := range artifacts { - if artifact.Target != modelartifacts.TargetCompanion || artifact.Resolved == nil { + if artifact.Target != modelartifacts.TargetCompanion { continue } - if _, exists := configured[artifact.Name]; exists { - xlog.Debug("keeping the configured companion option over the managed snapshot", "artifact", artifact.Name) - continue + if configured != nil { + if _, exists := configured[artifact.Name]; exists { + xlog.Debug("keeping the configured companion option over the managed snapshot", "artifact", artifact.Name) + continue + } } - snapshot, err := modelartifacts.RelativeSnapshotPath(artifact.Resolved.CacheKey) - if err != nil { - xlog.Warn("skipping companion artifact with an unusable cache key", "artifact", artifact.Name, "error", err) + + if artifact.Resolved != nil { + if snapshot, err := modelartifacts.RelativeSnapshotPath(artifact.Resolved.CacheKey); err == nil { + out = append(out, artifact.Name+":"+snapshot) + continue + } else { + xlog.Warn("companion artifact has an unusable cache key; falling back to its source repository", "artifact", artifact.Name, "error", err) + } + } + + if repo := strings.TrimSpace(artifact.Source.Repo); repo != "" { + xlog.Warn("companion artifact is not resolved on the controller; the backend will fetch it by repository id (no staging fast path)", + "artifact", artifact.Name, "repo", repo) + out = append(out, artifact.Name+":"+repo) continue } - combined = append(combined, artifact.Name+":"+snapshot) + xlog.Warn("companion artifact is neither resolved nor has a source repository; the backend will get no option for it", "artifact", artifact.Name) + } + return out +} + +// CompanionArtifactOptions returns the managed companion options a model's +// current config would contribute, skipping any companion name the config +// already pins in Options. It is the reconciler's entry point for re-deriving +// companion options when it replays a stored ModelOptions blob (which predates, +// and so cannot carry, a companion resolved after that blob was captured). +func CompanionArtifactOptions(c config.ModelConfig) []string { + configured := make(map[string]struct{}, len(c.Options)) + for _, option := range c.Options { + if name, _, found := strings.Cut(option, ":"); found { + configured[name] = struct{}{} + } } - return combined + return ManagedCompanionOptions(c.Artifacts, configured) } func grpcModelOpts(c config.ModelConfig, modelPath string) *pb.ModelOptions { 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 000000000000..1f8ff9ec968c --- /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 5b6aed953df0..0d5b03b5360a 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 000000000000..746841a55953 --- /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"})) + }) +})