mirror of
https://github.com/mudler/LocalAI.git
synced 2026-08-04 12:22:22 -04:00
Compare commits
2 Commits
bot/galler
...
harden/dis
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
240b59f257 | ||
|
|
fd809c23f9 |
@@ -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
|
||||
|
||||
@@ -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"))
|
||||
})
|
||||
})
|
||||
|
||||
@@ -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 "<artifact name>:<value>" 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)
|
||||
if configured != nil {
|
||||
if _, exists := configured[artifact.Name]; exists {
|
||||
xlog.Debug("keeping the configured companion option over the managed snapshot", "artifact", artifact.Name)
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
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)
|
||||
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 combined
|
||||
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 ManagedCompanionOptions(c.Artifacts, configured)
|
||||
}
|
||||
|
||||
func grpcModelOpts(c config.ModelConfig, modelPath string) *pb.ModelOptions {
|
||||
|
||||
58
core/services/nodes/file_staging_loadmodel_test.go
Normal file
58
core/services/nodes/file_staging_loadmodel_test.go
Normal file
@@ -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"))
|
||||
})
|
||||
})
|
||||
@@ -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)
|
||||
|
||||
61
core/services/nodes/router_companion_replay_test.go
Normal file
61
core/services/nodes/router_companion_replay_test.go
Normal file
@@ -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"}))
|
||||
})
|
||||
})
|
||||
Reference in New Issue
Block a user