diff --git a/core/application/distributed.go b/core/application/distributed.go index e22ac5534361..8389c5c9f3b8 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -41,6 +41,7 @@ type DistributedServices struct { FileStager nodes.FileStager ModelAdapter *nodes.ModelRouterAdapter Unloader *nodes.RemoteUnloaderAdapter + ModelCleanup *nodes.ModelCleanupService shutdownOnce sync.Once } @@ -346,8 +347,10 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade if configLoader != nil { conflictResolver = configLoader } + modelCleanup := nodes.NewModelCleanupService(registry, remoteUnloader) router := nodes.NewSmartRouter(registry, nodes.SmartRouterOptions{ Unloader: remoteUnloader, + ModelCleanup: modelCleanup, FileStager: fileStager, GalleriesJSON: routerGalleriesJSON, AuthToken: routerAuthToken, @@ -437,6 +440,7 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade FileStager: fileStager, ModelAdapter: modelAdapter, Unloader: remoteUnloader, + ModelCleanup: modelCleanup, }, nil } diff --git a/core/application/startup.go b/core/application/startup.go index b46af704618d..17144b7a512c 100644 --- a/core/application/startup.go +++ b/core/application/startup.go @@ -298,6 +298,7 @@ func New(opts ...config.AppOption) (*Application, error) { if distSvc.Reconciler != nil { go distSvc.Reconciler.Run(options.Context) } + go distSvc.ModelCleanup.Run(options.Context) // In distributed mode, MCP CI jobs are executed by agent workers (not the frontend) // because the frontend can't create MCP sessions (e.g., stdio servers using docker). // The dispatcher still subscribes to jobs.new for persistence (result/progress subs) @@ -370,13 +371,15 @@ func New(opts ...config.AppOption) (*Application, error) { gs := application.galleryService sys := options.SystemState cfgLoaderOpts := options.ToConfigLoaderOptions() + modelRevisionLifecycle := modeladmin.NewDistributedModelRevisionLifecycle(distSvc.Registry, distSvc.ModelCleanup) + gs.SetModelRevisionLifecycle(modelRevisionLifecycle) gs.OnModelsChanged = func(evt messaging.CacheInvalidateEvent) { // ApplyRemoteChange honors the op: a "delete" prunes the element // (a reload-from-path is additive and cannot drop it), anything // else reloads from disk; a named element's running instance is // shut down so the new config takes effect. The originating // replica reloads inline and never depends on this path. - if err := modeladmin.ApplyRemoteChange(application.ModelConfigLoader(), application.modelLoader, sys.Model.ModelsPath, evt, cfgLoaderOpts...); err != nil { + if err := modeladmin.ApplyRemoteChange(options.Context, application.ModelConfigLoader(), sys.Model.ModelsPath, evt, modelRevisionLifecycle, cfgLoaderOpts...); err != nil { xlog.Warn("Failed to apply peer model config change", "error", err) } } diff --git a/core/backend/ctx_propagation_test.go b/core/backend/ctx_propagation_test.go index 34f269aa3e14..5fc6b321ecbb 100644 --- a/core/backend/ctx_propagation_test.go +++ b/core/backend/ctx_propagation_test.go @@ -24,8 +24,8 @@ import ( "github.com/mudler/LocalAI/core/backend" "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/schema" - pbproto "github.com/mudler/LocalAI/pkg/grpc/proto" "github.com/mudler/LocalAI/pkg/distributedhdr" + pbproto "github.com/mudler/LocalAI/pkg/grpc/proto" "github.com/mudler/LocalAI/pkg/model" "github.com/mudler/LocalAI/pkg/system" @@ -41,7 +41,7 @@ import ( func newCapturingLoader() (*model.ModelLoader, *atomic.Value, func() context.Context) { loader := model.NewModelLoader(&system.SystemState{}) var captured atomic.Value - loader.SetModelRouter(func(ctx context.Context, _ string, _, _, _ string, _ *pbproto.ModelOptions, _ bool) (*model.Model, error) { + loader.SetModelRouter(func(ctx context.Context, _ string, _, _, _, _ string, _ *pbproto.ModelOptions, _ bool) (*model.Model, error) { captured.Store(ctx) // Return an error so the backend short-circuits before trying to // dial gRPC. We only care about the context-arrival contract. diff --git a/core/backend/model_identity_modalities_test.go b/core/backend/model_identity_modalities_test.go index 0087d84460ad..9f83f6da86c6 100644 --- a/core/backend/model_identity_modalities_test.go +++ b/core/backend/model_identity_modalities_test.go @@ -163,7 +163,7 @@ func (r *recordingBackend) VoiceEmbed(_ context.Context, in *pb.VoiceEmbedReques // backend, so every helper below reaches it through the real Load path. func newRecordingLoader(rec *recordingBackend) *model.ModelLoader { loader := model.NewModelLoader(&system.SystemState{}) - loader.SetModelRouter(func(_ context.Context, id string, _, _, _ string, _ *pb.ModelOptions, _ bool) (*model.Model, error) { + loader.SetModelRouter(func(_ context.Context, id string, _, _, _, _ string, _ *pb.ModelOptions, _ bool) (*model.Model, error) { return model.NewModelWithClient(id, "test://recording", rec), nil }) return loader diff --git a/core/backend/options.go b/core/backend/options.go index f8056d964dff..b56275c55a81 100644 --- a/core/backend/options.go +++ b/core/backend/options.go @@ -202,6 +202,11 @@ func ModelOptions(c config.ModelConfig, so *config.ApplicationConfig, opts ...mo model.WithContext(so.Context), model.WithModelID(c.ModelID()), } + if revision, err := config.ModelConfigRevision(&c); err == nil { + defOpts = append(defOpts, model.WithConfigRevision(revision)) + } else { + xlog.Warn("Failed to compute model configuration revision", "model", c.ModelID(), "error", err) + } managedPrimary := len(c.Artifacts) > 0 && c.Artifacts[0].Resolved != nil if managedPrimary { defOpts = append(defOpts, model.WithModelFile(c.ModelFileName())) diff --git a/core/config/application_config.go b/core/config/application_config.go index 1794cadc59f0..e6668f8822cc 100644 --- a/core/config/application_config.go +++ b/core/config/application_config.go @@ -1140,6 +1140,7 @@ func (o *ApplicationConfig) ToConfigLoaderOptions() []ConfigLoaderOption { LoadOptionF16(o.F16), LoadOptionThreads(o.Threads), ModelPath(o.SystemState.Model.ModelsPath), + LoadOptionGalleryFiles(o.Galleries...), } } diff --git a/core/config/model_config_loader.go b/core/config/model_config_loader.go index 788b4ce7fae7..d8bb02e406ff 100644 --- a/core/config/model_config_loader.go +++ b/core/config/model_config_loader.go @@ -6,6 +6,7 @@ import ( "errors" "fmt" "io/fs" + "net/url" "os" "path/filepath" "reflect" @@ -17,6 +18,7 @@ import ( "github.com/mudler/LocalAI/core/schema" "github.com/mudler/LocalAI/pkg/downloader" "github.com/mudler/LocalAI/pkg/modelartifacts" + "github.com/mudler/LocalAI/pkg/safefile" "github.com/mudler/LocalAI/pkg/utils" "github.com/mudler/xlog" "gopkg.in/yaml.v3" @@ -55,9 +57,20 @@ type ModelConfigLoader struct { artifactMaterializer ArtifactMaterializer preloadRenderMode string disablePreloadColor bool + mutationMu sync.Mutex sync.Mutex } +// WithModelConfigMutation serializes filesystem-backed configuration changes +// and their lifecycle publication for every service sharing this loader. It is +// deliberately separate from the loader's map lock: callbacks reload and +// replace loader state and would deadlock if that lock were held here. +func (bcl *ModelConfigLoader) WithModelConfigMutation(fn func() error) error { + bcl.mutationMu.Lock() + defer bcl.mutationMu.Unlock() + return fn() +} + func NewModelConfigLoader(modelPath string, options ...ModelConfigLoaderOption) *ModelConfigLoader { loader := &ModelConfigLoader{ configs: make(map[string]ModelConfig), @@ -76,6 +89,7 @@ type LoadOptions struct { debug bool threads, ctxSize int f16 bool + galleryFiles map[string]struct{} } func LoadOptionDebug(debug bool) ConfigLoaderOption { @@ -108,6 +122,30 @@ func LoadOptionF16(f16 bool) ConfigLoaderOption { } } +// LoadOptionGalleryFiles identifies local gallery sources that can legitimately +// live in the models directory. Exact paths provide provenance; document shape +// validation alone cannot distinguish an overrides-only gallery entry from a +// malformed runtime model configuration. +func LoadOptionGalleryFiles(galleries ...Gallery) ConfigLoaderOption { + return func(o *LoadOptions) { + if o.galleryFiles == nil { + o.galleryFiles = map[string]struct{}{} + } + for _, configured := range galleries { + for _, raw := range append([]string{configured.URL}, configured.Mirrors...) { + parsed, err := url.Parse(raw) + if err != nil || parsed.Scheme != "file" || parsed.Path == "" { + continue + } + absolute, err := filepath.Abs(filepath.FromSlash(parsed.Path)) + if err == nil { + o.galleryFiles[absolute] = struct{}{} + } + } + } + } +} + type ConfigLoaderOption func(*LoadOptions) func (lo *LoadOptions) Apply(options ...ConfigLoaderOption) { @@ -307,6 +345,18 @@ func (bcl *ModelConfigLoader) RemoveModelConfig(m string) { delete(bcl.configs, m) } +// ReplaceModelConfigs atomically replaces the in-memory configuration set with +// a previously parsed snapshot. +func (bcl *ModelConfigLoader) ReplaceModelConfigs(configs []ModelConfig) { + bcl.Lock() + defer bcl.Unlock() + replacement := make(map[string]ModelConfig, len(configs)) + for _, cfg := range configs { + replacement[cfg.Name] = cfg + } + bcl.configs = replacement +} + // GetModelsConflictingWith returns the names of every other configured (and // not-disabled) model that shares at least one concurrency group with the // named model. Returns nil if the named model has no groups, is unknown, or @@ -629,6 +679,16 @@ func (bcl *ModelConfigLoader) MITMHostOwners() MITMHostOwnership { // LoadModelConfigsFromPath reads all the configurations of the models from a path // (non-recursive) func (bcl *ModelConfigLoader) LoadModelConfigsFromPath(path string, opts ...ConfigLoaderOption) error { + return bcl.loadModelConfigsFromPath(path, false, opts...) +} + +// LoadModelConfigsFromPathStrict builds an authoritative snapshot and fails +// when any visible config cannot be parsed or validated. +func (bcl *ModelConfigLoader) LoadModelConfigsFromPathStrict(path string, opts ...ConfigLoaderOption) error { + return bcl.loadModelConfigsFromPath(path, true, opts...) +} + +func (bcl *ModelConfigLoader) loadModelConfigsFromPath(path string, strict bool, opts ...ConfigLoaderOption) error { bcl.Lock() defer bcl.Unlock() @@ -644,6 +704,8 @@ func (bcl *ModelConfigLoader) LoadModelConfigsFromPath(path string, opts ...Conf } files = append(files, info) } + loadOptions := &LoadOptions{} + loadOptions.Apply(opts...) for _, file := range files { // Only load real YAML config files and ignore dotfiles or backup variants ext := strings.ToLower(filepath.Ext(file.Name())) @@ -652,10 +714,35 @@ func (bcl *ModelConfigLoader) LoadModelConfigsFromPath(path string, opts ...Conf } filePath := filepath.Join(path, file.Name()) + absolutePath, absErr := filepath.Abs(filePath) + if absErr != nil { + return absErr + } + if _, gallerySource := loadOptions.galleryFiles[absolutePath]; gallerySource { + galleryDocument, err := classifyGalleryDocument(filePath) + if err != nil { + if strict { + return err + } + xlog.Error("LoadModelConfigsFromPath cannot validate gallery YAML file", "error", err, "File Name", file.Name()) + continue + } + if !galleryDocument { + if strict { + return fmt.Errorf("configured gallery source %q is not valid gallery metadata", filePath) + } + xlog.Error("Configured gallery source is not valid gallery metadata", "File Name", file.Name()) + continue + } + continue + } // Read config(s) - handles both single and array formats configs, err := readModelConfigsFromFile(filePath, opts...) if err != nil { + if strict { + return err + } xlog.Error("LoadModelConfigsFromPath cannot read config file", "error", err, "File Name", file.Name()) continue } @@ -665,6 +752,9 @@ func (bcl *ModelConfigLoader) LoadModelConfigsFromPath(path string, opts ...Conf if valid, validationErr := c.Validate(); valid { bcl.configs[c.Name] = *c } else { + if strict { + return fmt.Errorf("invalid model config %q: %w", c.Name, validationErr) + } xlog.Error("config is not valid", "error", validationErr, "Name", c.Name) } } @@ -688,3 +778,169 @@ func (bcl *ModelConfigLoader) LoadModelConfigsFromPath(path string, opts ...Conf return nil } + +var galleryMetadataKeys = map[string]struct{}{ + "name": {}, "description": {}, "license": {}, "icon": {}, "tags": {}, "size": {}, + "url": {}, "urls": {}, "config_file": {}, "overrides": {}, "files": {}, "variants": {}, "prompt_templates": {}, +} + +// classifyGalleryDocument recognizes the two gallery documents LocalAI writes +// beside model configurations: a GalleryModel catalogue sequence and the +// legacy downloadable ModelConfig mapping. A gallery discriminator makes the +// document subject to the complete shape check; malformed or mixed documents +// are errors rather than silently disappearing from an authoritative snapshot. +func classifyGalleryDocument(path string) (bool, error) { + data, _, err := safefile.ReadRegularAt(filepath.Dir(path), filepath.Base(path)) + if err != nil { + return false, fmt.Errorf("read YAML file %q for classification: %w", path, err) + } + var document yaml.Node + if err := yaml.Unmarshal(data, &document); err != nil || len(document.Content) != 1 { + return false, nil // The model-config parser supplies the syntax error. + } + root := document.Content[0] + switch root.Kind { + case yaml.SequenceNode: + looksGallery := false + for _, entry := range root.Content { + if entry.Kind == yaml.MappingNode && hasAnyMappingKey(entry, "url", "config_file", "variants", "files", "overrides") { + looksGallery = true + } + } + if !looksGallery { + return false, nil + } + if len(root.Content) == 0 { + return false, nil + } + for _, entry := range root.Content { + if err := validateGalleryCatalogueEntry(entry); err != nil { + return false, fmt.Errorf("invalid gallery catalogue %q: %w", path, err) + } + } + return true, nil + case yaml.MappingNode: + if !hasAnyMappingKey(root, "config_file", "prompt_templates") { + return false, nil + } + if err := validateLegacyGalleryModel(root); err != nil { + return false, fmt.Errorf("invalid gallery model metadata %q: %w", path, err) + } + return true, nil + default: + return false, nil + } +} + +func validateGalleryCatalogueEntry(entry *yaml.Node) error { + if entry.Kind != yaml.MappingNode { + return errors.New("entry must be a mapping") + } + if err := validateGalleryKeys(entry); err != nil { + return err + } + if !nonemptyScalar(galleryMappingValue(entry, "name")) { + return errors.New("entry name must be a non-empty string") + } + hasPayload := nonemptyScalar(galleryMappingValue(entry, "url")) + if node := galleryMappingValue(entry, "config_file"); node != nil { + if node.Kind != yaml.MappingNode { + return errors.New("config_file must be a mapping in a gallery catalogue") + } + hasPayload = true + } + for _, key := range []string{"overrides", "files", "variants"} { + if node := galleryMappingValue(entry, key); node != nil { + if err := validateGalleryPayload(key, node); err != nil { + return err + } + hasPayload = hasPayload || len(node.Content) > 0 + } + } + if !hasPayload { + return errors.New("entry has no installable gallery payload") + } + return nil +} + +func validateLegacyGalleryModel(entry *yaml.Node) error { + if err := validateGalleryKeys(entry); err != nil { + return err + } + if !nonemptyScalar(galleryMappingValue(entry, "name")) { + return errors.New("model name must be a non-empty string") + } + configFile := galleryMappingValue(entry, "config_file") + if !nonemptyScalar(configFile) { + return errors.New("config_file must be a non-empty YAML string") + } + for _, key := range []string{"files", "prompt_templates"} { + if node := galleryMappingValue(entry, key); node != nil && node.Kind != yaml.SequenceNode { + return fmt.Errorf("%s must be a sequence", key) + } + } + return nil +} + +func validateGalleryKeys(entry *yaml.Node) error { + for i := 0; i+1 < len(entry.Content); i += 2 { + key := entry.Content[i].Value + if _, ok := galleryMetadataKeys[key]; !ok { + return fmt.Errorf("field %q is not gallery metadata", key) + } + } + return nil +} + +func validateGalleryPayload(key string, node *yaml.Node) error { + switch key { + case "overrides": + if node.Kind != yaml.MappingNode { + return errors.New("overrides must be a mapping") + } + case "files": + if node.Kind != yaml.SequenceNode { + return errors.New("files must be a sequence") + } + for _, file := range node.Content { + if file.Kind != yaml.MappingNode || !nonemptyScalar(galleryMappingValue(file, "filename")) || !nonemptyScalar(galleryMappingValue(file, "uri")) { + return errors.New("each gallery file must have non-empty filename and uri strings") + } + } + case "variants": + if node.Kind != yaml.SequenceNode || len(node.Content) == 0 { + return errors.New("variants must be a non-empty sequence") + } + for _, variant := range node.Content { + if variant.Kind != yaml.MappingNode || len(variant.Content) != 2 || variant.Content[0].Value != "model" || !nonemptyScalar(variant.Content[1]) { + return errors.New("each variant must contain only a non-empty model string") + } + } + } + return nil +} + +func galleryMappingValue(mapping *yaml.Node, key string) *yaml.Node { + if mapping == nil || mapping.Kind != yaml.MappingNode { + return nil + } + for i := 0; i+1 < len(mapping.Content); i += 2 { + if mapping.Content[i].Value == key { + return mapping.Content[i+1] + } + } + return nil +} + +func hasAnyMappingKey(mapping *yaml.Node, keys ...string) bool { + for _, key := range keys { + if galleryMappingValue(mapping, key) != nil { + return true + } + } + return false +} + +func nonemptyScalar(node *yaml.Node) bool { + return node != nil && node.Kind == yaml.ScalarNode && node.Tag == "!!str" && strings.TrimSpace(node.Value) != "" +} diff --git a/core/config/model_config_revision.go b/core/config/model_config_revision.go new file mode 100644 index 000000000000..e8dd1bee5d1f --- /dev/null +++ b/core/config/model_config_revision.go @@ -0,0 +1,48 @@ +package config + +import ( + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + + pb "github.com/mudler/LocalAI/pkg/grpc/proto" + "google.golang.org/protobuf/proto" +) + +// ModelConfigRevision returns a stable revision of the persisted semantic +// configuration. ModelConfig's JSON tags exclude runtime-derived state and +// source bookkeeping, while encoding/json orders map keys deterministically. +func ModelConfigRevision(cfg *ModelConfig) (string, error) { + if cfg == nil { + return "", errors.New("model config is nil") + } + + canonical, err := json.Marshal(cfg) + if err != nil { + return "", err + } + + return sha256Hex(canonical), nil +} + +// EffectiveModelOptionsHash returns a deterministic hash of materialized +// backend options without allowing protobuf marshaling to mutate caller state. +func EffectiveModelOptionsHash(opts *pb.ModelOptions) (string, error) { + if opts == nil { + return "", errors.New("model options are nil") + } + + cloned := proto.Clone(opts).(*pb.ModelOptions) + canonical, err := (proto.MarshalOptions{Deterministic: true}).Marshal(cloned) + if err != nil { + return "", err + } + + return sha256Hex(canonical), nil +} + +func sha256Hex(value []byte) string { + digest := sha256.Sum256(value) + return hex.EncodeToString(digest[:]) +} diff --git a/core/config/model_config_revision_test.go b/core/config/model_config_revision_test.go new file mode 100644 index 000000000000..0b95d3549514 --- /dev/null +++ b/core/config/model_config_revision_test.go @@ -0,0 +1,143 @@ +package config_test + +import ( + "os" + "path/filepath" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/config" + pb "github.com/mudler/LocalAI/pkg/grpc/proto" + "gopkg.in/yaml.v3" +) + +var _ = Describe("Model configuration revisions", func() { + parse := func(document string) *config.ModelConfig { + cfg := &config.ModelConfig{} + Expect(yaml.Unmarshal([]byte(document), cfg)).To(Succeed()) + return cfg + } + + revision := func(cfg *config.ModelConfig) string { + value, err := config.ModelConfigRevision(cfg) + Expect(err).NotTo(HaveOccurred()) + return value + } + + It("is stable across equivalent YAML formatting and map order", func() { + first := parse("name: example\nbackend: llama-cpp\nroles: {user: USER, assistant: ASSISTANT}\n") + second := parse("# Comments and presentation details do not affect the typed configuration.\n" + + "backend: llama-cpp\nroles:\n assistant: ASSISTANT\n user: USER\nname: example\n") + + Expect(revision(first)).To(Equal(revision(second))) + }) + + It("changes for context and parallel options", func() { + base := parse("name: example\ncontext_size: 2048\noptions: [parallel:1]\n") + contextChanged := parse("name: example\ncontext_size: 4096\noptions: [parallel:1]\n") + parallelChanged := parse("name: example\ncontext_size: 2048\noptions: [parallel:2]\n") + + Expect(revision(contextChanged)).NotTo(Equal(revision(base))) + Expect(revision(parallelChanged)).NotTo(Equal(revision(base))) + }) + + It("preserves meaningful absence versus explicit zero", func() { + absent := parse("name: example\n") + explicitZero := parse("name: example\ncontext_size: 0\n") + + Expect(revision(explicitZero)).NotTo(Equal(revision(absent))) + }) + + It("excludes the configuration source path", func() { + dir := GinkgoT().TempDir() + paths := []string{filepath.Join(dir, "first.yaml"), filepath.Join(dir, "second.yaml")} + for _, path := range paths { + Expect(os.WriteFile(path, []byte("name: example\nbackend: llama-cpp\n"), 0o600)).To(Succeed()) + } + + loaded := make([]config.ModelConfig, 0, len(paths)) + for _, path := range paths { + loader := config.NewModelConfigLoader(dir) + Expect(loader.ReadModelConfig(path)).To(Succeed()) + cfg, found := loader.GetModelConfig("example") + Expect(found).To(BeTrue()) + loaded = append(loaded, cfg) + } + + Expect(loaded[0].GetModelConfigFile()).NotTo(Equal(loaded[1].GetModelConfigFile())) + Expect(revision(&loaded[0])).To(Equal(revision(&loaded[1]))) + }) + + It("hashes effective protobuf options deterministically without mutation", func() { + opts := &pb.ModelOptions{Model: "example", ContextSize: 2048, TensorParallelSize: 1} + original := opts.String() + + first, err := config.EffectiveModelOptionsHash(opts) + Expect(err).NotTo(HaveOccurred()) + second, err := config.EffectiveModelOptionsHash(opts) + Expect(err).NotTo(HaveOccurred()) + + Expect(second).To(Equal(first)) + Expect(opts.String()).To(Equal(original)) + + changed := &pb.ModelOptions{Model: "example", ContextSize: 4096, TensorParallelSize: 1} + different, err := config.EffectiveModelOptionsHash(changed) + Expect(err).NotTo(HaveOccurred()) + Expect(different).NotTo(Equal(first)) + }) +}) + +var _ = Describe("Strict model configuration snapshots", func() { + write := func(dir, name, body string) { + Expect(os.WriteFile(filepath.Join(dir, name), []byte(body), 0o600)).To(Succeed()) + } + + It("ignores valid catalogue and legacy gallery metadata", func() { + dir := GinkgoT().TempDir() + write(dir, "catalogue.yaml", "- name: downloadable\n url: github:example/model.yaml\n- name: inline\n config_file:\n backend: llama-cpp\n") + write(dir, "gallery_simple.yaml", "name: legacy\nconfig_file: |\n backend: llama-cpp\nfiles:\n- filename: model.gguf\n uri: https://example.invalid/model.gguf\n") + write(dir, "installed.yaml", "name: installed\nbackend: llama-cpp\n") + + loader := config.NewModelConfigLoader(dir) + galleryFiles := config.LoadOptionGalleryFiles( + config.Gallery{URL: "file://" + filepath.Join(dir, "catalogue.yaml")}, + config.Gallery{URL: "file://" + filepath.Join(dir, "gallery_simple.yaml")}, + ) + Expect(loader.LoadModelConfigsFromPathStrict(dir, galleryFiles)).To(Succeed()) + _, found := loader.GetModelConfig("installed") + Expect(found).To(BeTrue()) + _, found = loader.GetModelConfig("downloadable") + Expect(found).To(BeFalse()) + _, found = loader.GetModelConfig("legacy") + Expect(found).To(BeFalse()) + }) + + DescribeTable("rejects malformed gallery-looking documents", + func(body, message string) { + dir := GinkgoT().TempDir() + write(dir, "broken.yaml", body) + loader := config.NewModelConfigLoader(dir) + galleryFile := config.LoadOptionGalleryFiles(config.Gallery{URL: "file://" + filepath.Join(dir, "broken.yaml")}) + Expect(loader.LoadModelConfigsFromPathStrict(dir, galleryFile)).To(MatchError(ContainSubstring(message))) + }, + Entry("invalid variants", "- name: broken\n variants: []\n", "variants must be a non-empty sequence"), + Entry("malformed payload type", "- name: broken\n files: nope\n", "files must be a sequence"), + Entry("mixed runtime and gallery fields", "- name: broken\n backend: llama-cpp\n url: github:example/model.yaml\n", `field "backend" is not gallery metadata`), + Entry("malformed legacy config", "name: broken\nconfig_file: [not, yaml]\n", "config_file must be a non-empty YAML string"), + ) + + It("still rejects invalid runtime configuration sequences", func() { + dir := GinkgoT().TempDir() + write(dir, "broken.yaml", "- name: broken\n backend: [\n") + loader := config.NewModelConfigLoader(dir) + Expect(loader.LoadModelConfigsFromPathStrict(dir)).To(MatchError(ContainSubstring("cannot unmarshal config file"))) + }) + + It("does not skip a valid gallery-shaped document without configured provenance", func() { + dir := GinkgoT().TempDir() + write(dir, "runtime.yaml", "- name: ambiguous\n overrides:\n backend: llama-cpp\n") + loader := config.NewModelConfigLoader(dir) + Expect(loader.LoadModelConfigsFromPathStrict(dir)).ToNot(Succeed()) + }) +}) diff --git a/core/config/model_config_revision_transition.go b/core/config/model_config_revision_transition.go new file mode 100644 index 000000000000..81c6a507a8b7 --- /dev/null +++ b/core/config/model_config_revision_transition.go @@ -0,0 +1,9 @@ +package config + +// ModelConfigRevisionTransition describes one authoritative model identity. +// Related transitions are applied together by revision lifecycle services. +type ModelConfigRevisionTransition struct { + ModelName string + ConfigRevision string + Disabled bool +} diff --git a/core/http/endpoints/localai/config_meta.go b/core/http/endpoints/localai/config_meta.go index 3db694512a11..b891a764e75c 100644 --- a/core/http/endpoints/localai/config_meta.go +++ b/core/http/endpoints/localai/config_meta.go @@ -155,8 +155,8 @@ func AutocompleteEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, a // @Param name path string true "Model name" // @Success 200 {object} map[string]any "success message" // @Router /api/models/config-json/{name} [patch] -func PatchConfigEndpoint(cl *config.ModelConfigLoader, _ *model.ModelLoader, gs *galleryop.GalleryService, appConfig *config.ApplicationConfig) echo.HandlerFunc { - svc := modeladmin.NewConfigService(cl, appConfig) +func PatchConfigEndpoint(cl *config.ModelConfigLoader, gs *galleryop.GalleryService, appConfig *config.ApplicationConfig, lifecycle ...modeladmin.ModelRevisionLifecycle) echo.HandlerFunc { + svc := modeladmin.NewConfigService(cl, appConfig, lifecycle...) return func(c echo.Context) error { modelName := c.Param("name") if decoded, err := url.PathUnescape(modelName); err == nil { @@ -170,7 +170,8 @@ func PatchConfigEndpoint(cl *config.ModelConfigLoader, _ *model.ModelLoader, gs if err := json.Unmarshal(patchBody, &patchMap); err != nil { return c.JSON(http.StatusBadRequest, map[string]any{"error": "invalid JSON: " + err.Error()}) } - if _, err := svc.PatchConfig(c.Request().Context(), modelName, patchMap); err != nil { + result, err := svc.PatchConfig(c.Request().Context(), modelName, patchMap) + if err != nil { return c.JSON(httpStatusForModelAdminError(err), map[string]any{"error": err.Error()}) } @@ -178,12 +179,14 @@ func PatchConfigEndpoint(cl *config.ModelConfigLoader, _ *model.ModelLoader, gs // tell peers to refresh so the change is consistent across replicas. // No-op in standalone mode. if gs != nil { - gs.BroadcastModelsChanged(modelName, "install") + gs.BroadcastModelsChangedRevision(modelName, "install", result.ConfigRevision) } return c.JSON(http.StatusOK, map[string]any{ - "success": true, - "message": fmt.Sprintf("Model '%s' updated successfully", modelName), + "success": true, + "message": fmt.Sprintf("Model '%s' updated successfully", modelName), + "config_revision": result.ConfigRevision, + "pending_cleanup": result.PendingCleanup, }) } } diff --git a/core/http/endpoints/localai/config_meta_test.go b/core/http/endpoints/localai/config_meta_test.go index e60f7e08d1e4..db7e0fd5fbdb 100644 --- a/core/http/endpoints/localai/config_meta_test.go +++ b/core/http/endpoints/localai/config_meta_test.go @@ -2,6 +2,7 @@ package localai_test import ( "bytes" + "context" "encoding/json" "net/http" "net/http/httptest" @@ -11,12 +12,24 @@ import ( "github.com/labstack/echo/v4" "github.com/mudler/LocalAI/core/config" . "github.com/mudler/LocalAI/core/http/endpoints/localai" + "github.com/mudler/LocalAI/core/services/galleryop" + "github.com/mudler/LocalAI/core/services/modeladmin" "github.com/mudler/LocalAI/pkg/model" "github.com/mudler/LocalAI/pkg/system" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" ) +type endpointLifecycleRecorder struct { + batches [][]modeladmin.ModelRevisionTransition + pendingCleanup int +} + +func (r *endpointLifecycleRecorder) ApplyConfigRevisions(_ context.Context, transitions []modeladmin.ModelRevisionTransition) (int, error) { + r.batches = append(r.batches, append([]modeladmin.ModelRevisionTransition(nil), transitions...)) + return r.pendingCleanup, nil +} + var _ = Describe("Config Metadata Endpoints", func() { var ( app *echo.Echo @@ -45,7 +58,7 @@ var _ = Describe("Config Metadata Endpoints", func() { app = echo.New() app.GET("/api/models/config-metadata", ConfigMetadataEndpoint()) app.GET("/api/models/config-metadata/autocomplete/:provider", AutocompleteEndpoint(configLoader, modelLoader, appConfig)) - app.PATCH("/api/models/config-json/:name", PatchConfigEndpoint(configLoader, modelLoader, nil, appConfig)) + app.PATCH("/api/models/config-json/:name", PatchConfigEndpoint(configLoader, nil, appConfig)) }) AfterEach(func() { @@ -167,6 +180,39 @@ backend: llama-cpp }) Context("PATCH /api/models/config-json/:name", func() { + It("rejects a name change without disk, loader, lifecycle, or broadcast effects", func() { + seedConfig := "name: test-model\nbackend: llama-cpp\ncontext_size: 4096\n" + configPath := filepath.Join(tempDir, "test-model.yaml") + Expect(os.WriteFile(configPath, []byte(seedConfig), 0644)).To(Succeed()) + Expect(configLoader.LoadModelConfigsFromPath(tempDir)).To(Succeed()) + lifecycle := &endpointLifecycleRecorder{} + galleryService := galleryop.NewGalleryService(appConfig, nil) + client := &endpointRecordingClient{} + galleryService.SetNATSClient(client) + endpointApp := echo.New() + endpointApp.PATCH("/api/models/config-json/:name", PatchConfigEndpoint(configLoader, galleryService, appConfig, lifecycle)) + + body := bytes.NewBufferString(`{"name":"renamed","context_size":8192}`) + req := httptest.NewRequest(http.MethodPatch, "/api/models/config-json/test-model", body) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + endpointApp.ServeHTTP(rec, req) + + Expect(rec.Code).To(Equal(http.StatusBadRequest), rec.Body.String()) + Expect(rec.Body.String()).To(ContainSubstring("cannot rename")) + Expect(configPath).To(BeAnExistingFile()) + contents, err := os.ReadFile(configPath) + Expect(err).ToNot(HaveOccurred()) + Expect(string(contents)).To(Equal(seedConfig)) + loaded, ok := configLoader.GetModelConfig("test-model") + Expect(ok).To(BeTrue()) + Expect(loaded.ContextSize).To(HaveValue(Equal(4096))) + _, renamed := configLoader.GetModelConfig("renamed") + Expect(renamed).To(BeFalse()) + Expect(lifecycle.batches).To(BeEmpty()) + Expect(client.published).To(BeEmpty()) + }) + It("should return 404 for nonexistent model", func() { body := bytes.NewBufferString(`{"backend": "bar"}`) req := httptest.NewRequest(http.MethodPatch, "/api/models/config-json/nonexistent", body) @@ -228,6 +274,8 @@ backend: llama-cpp var resp map[string]any Expect(json.Unmarshal(rec.Body.Bytes(), &resp)).To(Succeed()) Expect(resp["success"]).To(BeTrue()) + Expect(resp["config_revision"]).ToNot(BeEmpty()) + Expect(resp["pending_cleanup"]).To(BeNumerically("==", 0)) // Verify the reloaded config has the updated value updatedConfig, exists := configLoader.GetModelConfig("test-model") @@ -240,6 +288,30 @@ backend: llama-cpp Expect(string(data)).To(ContainSubstring("vllm")) }) + It("reports lifecycle-backed pending cleanup with the saved revision", func() { + seedConfig := "name: test-model\nbackend: llama-cpp\ncontext_size: 4096\n" + Expect(os.WriteFile(filepath.Join(tempDir, "test-model.yaml"), []byte(seedConfig), 0o644)).To(Succeed()) + Expect(configLoader.LoadModelConfigsFromPath(tempDir)).To(Succeed()) + lifecycle := &endpointLifecycleRecorder{pendingCleanup: 4} + endpointApp := echo.New() + endpointApp.PATCH( + "/api/models/config-json/:name", + PatchConfigEndpoint(configLoader, nil, appConfig, lifecycle), + ) + + body := bytes.NewBufferString(`{"context_size":8192}`) + req := httptest.NewRequest(http.MethodPatch, "/api/models/config-json/test-model", body) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + endpointApp.ServeHTTP(rec, req) + + Expect(rec.Code).To(Equal(http.StatusOK), rec.Body.String()) + var response map[string]any + Expect(json.Unmarshal(rec.Body.Bytes(), &response)).To(Succeed()) + Expect(response).To(HaveKeyWithValue("config_revision", Not(BeEmpty()))) + Expect(response).To(HaveKeyWithValue("pending_cleanup", BeNumerically("==", 4))) + }) + It("should not persist runtime defaults (SetDefaults values) to disk", func() { // Create a minimal pipeline config - no sampling params seedConfig := `name: gpt-realtime diff --git a/core/http/endpoints/localai/edit_model.go b/core/http/endpoints/localai/edit_model.go index 5dd5737510ef..117d352617ad 100644 --- a/core/http/endpoints/localai/edit_model.go +++ b/core/http/endpoints/localai/edit_model.go @@ -13,7 +13,6 @@ import ( "github.com/mudler/LocalAI/core/services/galleryop" "github.com/mudler/LocalAI/core/services/modeladmin" "github.com/mudler/LocalAI/internal" - "github.com/mudler/LocalAI/pkg/model" ) // GetEditModelPage renders the edit model page with current configuration @@ -56,8 +55,8 @@ func GetEditModelPage(cl *config.ModelConfigLoader, appConfig *config.Applicatio } // EditModelEndpoint handles updating existing model configurations -func EditModelEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, gs *galleryop.GalleryService, appConfig *config.ApplicationConfig) echo.HandlerFunc { - svc := modeladmin.NewConfigService(cl, appConfig) +func EditModelEndpoint(cl *config.ModelConfigLoader, gs *galleryop.GalleryService, appConfig *config.ApplicationConfig, lifecycle ...modeladmin.ModelRevisionLifecycle) echo.HandlerFunc { + svc := modeladmin.NewConfigService(cl, appConfig, lifecycle...) return func(c echo.Context) error { modelName := c.Param("name") if decoded, err := url.PathUnescape(modelName); err == nil { @@ -67,7 +66,7 @@ func EditModelEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, gs * if err != nil { return c.JSON(http.StatusBadRequest, ModelResponse{Success: false, Error: "Failed to read request body: " + err.Error()}) } - result, err := svc.EditYAML(c.Request().Context(), modelName, body, ml) + result, err := svc.EditYAML(c.Request().Context(), modelName, body) if err != nil { return c.JSON(httpStatusForModelAdminError(err), ModelResponse{Success: false, Error: err.Error()}) } @@ -77,9 +76,9 @@ func EditModelEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, gs * // plus an install of the new one. No-op in standalone mode. if gs != nil { if result.Renamed { - gs.BroadcastModelsChanged(result.OldName, "delete") + gs.BroadcastModelsChangedRevision(result.OldName, "delete", modeladmin.DeletedModelConfigRevision(result.OldName)) } - gs.BroadcastModelsChanged(result.NewName, "install") + gs.BroadcastModelsChangedRevision(result.NewName, "install", result.ConfigRevision) } msg := fmt.Sprintf("Model '%s' updated successfully. Model has been reloaded with new configuration.", result.NewName) @@ -87,10 +86,12 @@ func EditModelEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, gs * msg = fmt.Sprintf("Model '%s' renamed to '%s' and updated successfully.", result.OldName, result.NewName) } return c.JSON(http.StatusOK, ModelResponse{ - Success: true, - Message: msg, - Filename: result.Filename, - Config: result.Config, + Success: true, + Message: msg, + Filename: result.Filename, + Config: result.Config, + ConfigRevision: result.ConfigRevision, + PendingCleanup: result.PendingCleanup, }) } } diff --git a/core/http/endpoints/localai/edit_model_test.go b/core/http/endpoints/localai/edit_model_test.go index 54ad2d5ec06b..17f7c4a7c8f3 100644 --- a/core/http/endpoints/localai/edit_model_test.go +++ b/core/http/endpoints/localai/edit_model_test.go @@ -2,23 +2,65 @@ package localai_test import ( "bytes" + "context" "encoding/json" + "errors" "io" "net/http" "net/http/httptest" "os" "path/filepath" + "time" "github.com/labstack/echo/v4" "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/gallery" . "github.com/mudler/LocalAI/core/http/endpoints/localai" - "github.com/mudler/LocalAI/pkg/model" + "github.com/mudler/LocalAI/core/services/galleryop" + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/modeladmin" + "github.com/mudler/LocalAI/pkg/modelartifacts" "github.com/mudler/LocalAI/pkg/system" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" ) +type endpointFailingMaterializer struct{} + +func (*endpointFailingMaterializer) Ensure(context.Context, string, modelartifacts.Spec) (modelartifacts.Result, error) { + return modelartifacts.Result{}, errors.New("artifact unavailable") +} + +type endpointRecordingClient struct { + published []messaging.CacheInvalidateEvent +} +type endpointSubscription struct{} + +func (*endpointSubscription) Unsubscribe() error { return nil } +func (c *endpointRecordingClient) Publish(subject string, data any) error { + if subject == messaging.SubjectCacheInvalidateModels { + c.published = append(c.published, data.(messaging.CacheInvalidateEvent)) + } + return nil +} +func (*endpointRecordingClient) Subscribe(string, func([]byte)) (messaging.Subscription, error) { + return &endpointSubscription{}, nil +} +func (*endpointRecordingClient) QueueSubscribe(string, string, func([]byte)) (messaging.Subscription, error) { + return &endpointSubscription{}, nil +} +func (*endpointRecordingClient) QueueSubscribeReply(string, string, func([]byte, func([]byte))) (messaging.Subscription, error) { + return &endpointSubscription{}, nil +} +func (*endpointRecordingClient) SubscribeReply(string, func([]byte, func([]byte))) (messaging.Subscription, error) { + return &endpointSubscription{}, nil +} +func (*endpointRecordingClient) Request(string, []byte, time.Duration) ([]byte, error) { + return nil, nil +} +func (*endpointRecordingClient) IsConnected() bool { return true } +func (*endpointRecordingClient) Close() {} + // testRenderer is a simple renderer for tests that returns JSON type testRenderer struct{} @@ -40,6 +82,91 @@ var _ = Describe("Edit Model test", func() { }) Context("Edit Model endpoint", func() { + DescribeTable("reports the saved revision and pending cleanup count", + func(pendingCleanup int) { + systemState, err := system.GetSystemState(system.WithModelPath(tempDir)) + Expect(err).ToNot(HaveOccurred()) + applicationConfig := config.NewApplicationConfig(config.WithSystemState(systemState)) + loader := config.NewModelConfigLoader(tempDir) + Expect(os.WriteFile( + filepath.Join(tempDir, "model.yaml"), + []byte("name: model\nbackend: llama-cpp\ncontext_size: 4096\n"), + 0o644, + )).To(Succeed()) + Expect(loader.LoadModelConfigsFromPath(tempDir)).To(Succeed()) + lifecycle := &endpointLifecycleRecorder{pendingCleanup: pendingCleanup} + app := echo.New() + app.POST("/models/edit/:name", EditModelEndpoint(loader, nil, applicationConfig, lifecycle)) + + req := httptest.NewRequest( + http.MethodPost, + "/models/edit/model", + bytes.NewBufferString("name: model\nbackend: llama-cpp\ncontext_size: 8192\n"), + ) + rec := httptest.NewRecorder() + app.ServeHTTP(rec, req) + + Expect(rec.Code).To(Equal(http.StatusOK), rec.Body.String()) + var response map[string]any + Expect(json.Unmarshal(rec.Body.Bytes(), &response)).To(Succeed()) + Expect(response).To(HaveKeyWithValue("config_revision", Not(BeEmpty()))) + Expect(response).To(HaveKeyWithValue("pending_cleanup", BeNumerically("==", pendingCleanup))) + }, + Entry("when no replicas need cleanup", 0), + Entry("when stale replicas remain queued for cleanup", 2), + ) + + It("does not broadcast an in-place edit that rolls back during preload", func() { + systemState, err := system.GetSystemState(system.WithModelPath(tempDir)) + Expect(err).ToNot(HaveOccurred()) + applicationConfig := config.NewApplicationConfig(config.WithSystemState(systemState)) + loader := config.NewModelConfigLoader(tempDir, config.WithArtifactMaterializer(&endpointFailingMaterializer{})) + path := filepath.Join(tempDir, "model.yaml") + Expect(os.WriteFile(path, []byte("name: model\nbackend: llama-cpp\ncontext_size: 4096\n"), 0644)).To(Succeed()) + Expect(loader.LoadModelConfigsFromPath(tempDir)).To(Succeed()) + galleryService := galleryop.NewGalleryService(applicationConfig, nil) + client := &endpointRecordingClient{} + galleryService.SetNATSClient(client) + + app := echo.New() + app.POST("/models/edit/:name", EditModelEndpoint(loader, galleryService, applicationConfig)) + body := "name: model\nbackend: llama-cpp\ncontext_size: 8192\nartifacts:\n - name: model\n target: model\n source: {type: huggingface, repo: owner/repo}\n" + req := httptest.NewRequest("POST", "/models/edit/model", bytes.NewBufferString(body)) + rec := httptest.NewRecorder() + app.ServeHTTP(rec, req) + + Expect(rec.Code).To(Equal(http.StatusInternalServerError)) + Expect(client.published).To(BeEmpty()) + bodyOnDisk, err := os.ReadFile(path) + Expect(err).ToNot(HaveOccurred()) + Expect(string(bodyOnDisk)).To(ContainSubstring("context_size: 4096")) + }) + + It("does not broadcast an edit that rolls back during preload", func() { + systemState, err := system.GetSystemState(system.WithModelPath(tempDir)) + Expect(err).ToNot(HaveOccurred()) + applicationConfig := config.NewApplicationConfig(config.WithSystemState(systemState)) + loader := config.NewModelConfigLoader(tempDir, config.WithArtifactMaterializer(&endpointFailingMaterializer{})) + path := filepath.Join(tempDir, "old.yaml") + Expect(os.WriteFile(path, []byte("name: old\nbackend: llama-cpp\ncontext_size: 4096\n"), 0644)).To(Succeed()) + Expect(loader.LoadModelConfigsFromPath(tempDir)).To(Succeed()) + galleryService := galleryop.NewGalleryService(applicationConfig, nil) + client := &endpointRecordingClient{} + galleryService.SetNATSClient(client) + + app := echo.New() + app.POST("/models/edit/:name", EditModelEndpoint(loader, galleryService, applicationConfig)) + body := "name: new\nbackend: llama-cpp\ncontext_size: 8192\nartifacts:\n - name: model\n target: model\n source: {type: huggingface, repo: owner/repo}\n" + req := httptest.NewRequest("POST", "/models/edit/old", bytes.NewBufferString(body)) + rec := httptest.NewRecorder() + app.ServeHTTP(rec, req) + + Expect(rec.Code).To(Equal(http.StatusInternalServerError)) + Expect(client.published).To(BeEmpty()) + Expect(path).To(BeAnExistingFile()) + Expect(filepath.Join(tempDir, "new.yaml")).NotTo(BeAnExistingFile()) + }) + It("should edit a model", func() { systemState, err := system.GetSystemState( system.WithModelPath(filepath.Join(tempDir)), @@ -92,7 +219,6 @@ var _ = Describe("Edit Model test", func() { config.WithSystemState(systemState), ) modelConfigLoader := config.NewModelConfigLoader(systemState.Model.ModelsPath) - modelLoader := model.NewModelLoader(systemState) oldYAML := "name: oldname\nbackend: llama\nmodel: foo\n" oldPath := filepath.Join(tempDir, "oldname.yaml") @@ -106,7 +232,7 @@ var _ = Describe("Edit Model test", func() { Expect(exists).To(BeTrue()) app := echo.New() - app.POST("/models/edit/:name", EditModelEndpoint(modelConfigLoader, modelLoader, nil, applicationConfig)) + app.POST("/models/edit/:name", EditModelEndpoint(modelConfigLoader, nil, applicationConfig)) newYAML := "name: newname\nbackend: llama\nmodel: foo\n" req := httptest.NewRequest("POST", "/models/edit/oldname", bytes.NewBufferString(newYAML)) @@ -139,6 +265,68 @@ var _ = Describe("Edit Model test", func() { Expect(modelConfigLoader.GetAllModelsConfigs()).To(HaveLen(1)) }) + It("broadcasts rename tombstone and install events with their own revisions", func() { + systemState, err := system.GetSystemState(system.WithModelPath(tempDir)) + Expect(err).ToNot(HaveOccurred()) + applicationConfig := config.NewApplicationConfig(config.WithSystemState(systemState)) + loader := config.NewModelConfigLoader(tempDir) + Expect(os.WriteFile(filepath.Join(tempDir, "old.yaml"), []byte("name: old\nbackend: llama-cpp\ncontext_size: 4096\n"), 0644)).To(Succeed()) + Expect(loader.LoadModelConfigsFromPath(tempDir)).To(Succeed()) + peerLoader := config.NewModelConfigLoader(tempDir) + Expect(peerLoader.LoadModelConfigsFromPath(tempDir)).To(Succeed()) + galleryService := galleryop.NewGalleryService(applicationConfig, nil) + client := &endpointRecordingClient{} + galleryService.SetNATSClient(client) + lifecycle := &endpointLifecycleRecorder{pendingCleanup: 2} + app := echo.New() + app.POST("/models/edit/:name", EditModelEndpoint(loader, galleryService, applicationConfig, lifecycle)) + + req := httptest.NewRequest(http.MethodPost, "/models/edit/old", bytes.NewBufferString("name: new\nbackend: llama-cpp\ncontext_size: 8192\n")) + rec := httptest.NewRecorder() + app.ServeHTTP(rec, req) + + Expect(rec.Code).To(Equal(http.StatusOK), rec.Body.String()) + var response map[string]any + Expect(json.Unmarshal(rec.Body.Bytes(), &response)).To(Succeed()) + Expect(response).To(HaveKeyWithValue("config_revision", Not(BeEmpty()))) + Expect(response).To(HaveKeyWithValue("pending_cleanup", BeNumerically("==", 2))) + Expect(client.published).To(HaveLen(2)) + Expect(client.published[0]).To(Equal(messaging.CacheInvalidateEvent{ + Element: "old", Op: "delete", ConfigRevision: modeladmin.DeletedModelConfigRevision("old"), + })) + newConfig, ok := loader.GetModelConfig("new") + Expect(ok).To(BeTrue()) + newRevision, err := config.ModelConfigRevision(&newConfig) + Expect(err).ToNot(HaveOccurred()) + Expect(client.published[1]).To(Equal(messaging.CacheInvalidateEvent{ + Element: "new", Op: "install", ConfigRevision: newRevision, + })) + Expect(lifecycle.batches).To(HaveLen(1)) + Expect(lifecycle.batches[0]).To(Equal([]modeladmin.ModelRevisionTransition{ + {ModelName: "old", ConfigRevision: modeladmin.DeletedModelConfigRevision("old"), Disabled: true}, + {ModelName: "new", ConfigRevision: newRevision}, + })) + + peerLifecycle := &endpointLifecycleRecorder{} + for _, event := range client.published { + Expect(modeladmin.ApplyRemoteChange(context.Background(), peerLoader, tempDir, event, peerLifecycle, applicationConfig.ToConfigLoaderOptions()...)).To(Succeed()) + } + _, oldOnPeer := peerLoader.GetModelConfig("old") + Expect(oldOnPeer).To(BeFalse()) + peerConfig, newOnPeer := peerLoader.GetModelConfig("new") + Expect(newOnPeer).To(BeTrue()) + peerRevision, err := config.ModelConfigRevision(&peerConfig) + Expect(err).ToNot(HaveOccurred()) + Expect(peerRevision).To(Equal(newRevision)) + Expect(peerLifecycle.batches).To(Equal([][]modeladmin.ModelRevisionTransition{ + { + {ModelName: "new", ConfigRevision: newRevision}, + {ModelName: "old", ConfigRevision: modeladmin.DeletedModelConfigRevision("old"), Disabled: true}, + }, + {{ModelName: "new", ConfigRevision: newRevision}}, + })) + }) + It("rejects a rename when the new name already exists", func() { systemState, err := system.GetSystemState( system.WithModelPath(tempDir), @@ -148,7 +336,6 @@ var _ = Describe("Edit Model test", func() { config.WithSystemState(systemState), ) modelConfigLoader := config.NewModelConfigLoader(systemState.Model.ModelsPath) - modelLoader := model.NewModelLoader(systemState) Expect(os.WriteFile( filepath.Join(tempDir, "oldname.yaml"), @@ -163,7 +350,7 @@ var _ = Describe("Edit Model test", func() { Expect(modelConfigLoader.LoadModelConfigsFromPath(tempDir)).To(Succeed()) app := echo.New() - app.POST("/models/edit/:name", EditModelEndpoint(modelConfigLoader, modelLoader, nil, applicationConfig)) + app.POST("/models/edit/:name", EditModelEndpoint(modelConfigLoader, nil, applicationConfig)) req := httptest.NewRequest( "POST", @@ -194,7 +381,6 @@ var _ = Describe("Edit Model test", func() { config.WithSystemState(systemState), ) modelConfigLoader := config.NewModelConfigLoader(systemState.Model.ModelsPath) - modelLoader := model.NewModelLoader(systemState) Expect(os.WriteFile( filepath.Join(tempDir, "oldname.yaml"), @@ -204,7 +390,7 @@ var _ = Describe("Edit Model test", func() { Expect(modelConfigLoader.LoadModelConfigsFromPath(tempDir)).To(Succeed()) app := echo.New() - app.POST("/models/edit/:name", EditModelEndpoint(modelConfigLoader, modelLoader, nil, applicationConfig)) + app.POST("/models/edit/:name", EditModelEndpoint(modelConfigLoader, nil, applicationConfig)) req := httptest.NewRequest( "POST", diff --git a/core/http/endpoints/localai/nodes_test.go b/core/http/endpoints/localai/nodes_test.go index 52cef6f03738..19e6a6b07eea 100644 --- a/core/http/endpoints/localai/nodes_test.go +++ b/core/http/endpoints/localai/nodes_test.go @@ -8,6 +8,7 @@ import ( "net/http" "net/http/httptest" "strings" + "time" "github.com/labstack/echo/v4" "github.com/mudler/LocalAI/core/services/nodes" @@ -447,4 +448,58 @@ var _ = Describe("Node HTTP handlers", func() { Expect(list[0].NodeID).To(Equal("n1")) }) }) + + Describe("GetNodeModelsEndpoint", func() { + It("returns revision and cleanup state without serialized model options", func() { + ctx := context.Background() + Expect(registry.Register(ctx, &nodes.BackendNode{ + ID: "n1", Name: "alpha", Address: "10.0.0.1:50051", Status: nodes.StatusHealthy, + }, true)).To(Succeed()) + + Expect(registry.EstablishModelConfigRevision(ctx, "current-model", "revision-current")).To(Succeed()) + Expect(registry.SetNodeModelRevision(ctx, "n1", "current-model", 0, "loaded", "10.0.0.1:50052", 0, "revision-current", "options-current")).To(Succeed()) + Expect(registry.SetNodeModelLoadInfoRevision(ctx, "n1", "current-model", 0, "llama-cpp", "revision-current", []byte("serialized-options"))).To(Succeed()) + + Expect(registry.EstablishModelConfigRevision(ctx, "changed-model", "revision-old")).To(Succeed()) + Expect(registry.SetNodeModelRevision(ctx, "n1", "changed-model", 0, "loaded", "10.0.0.1:50053", 0, "revision-old", "options-old")).To(Succeed()) + quarantined, err := registry.AdvanceModelConfigRevision(ctx, "changed-model", "revision-new") + Expect(err).ToNot(HaveOccurred()) + Expect(quarantined).To(HaveLen(1)) + retryAt := time.Now().UTC().Add(time.Minute).Truncate(time.Second) + Expect(registry.RecordModelCleanupFailure(ctx, "n1", "changed-model", 0, "worker unreachable", retryAt)).To(Succeed()) + + e := echo.New() + req := httptest.NewRequest(http.MethodGet, "/api/nodes/n1/models", nil) + rec := httptest.NewRecorder() + c := e.NewContext(req, rec) + c.SetPath("/api/nodes/:id/models") + c.SetParamNames("id") + c.SetParamValues("n1") + + Expect(GetNodeModelsEndpoint(registry)(c)).To(Succeed()) + Expect(rec.Code).To(Equal(http.StatusOK)) + + var modelsResponse []map[string]any + Expect(json.Unmarshal(rec.Body.Bytes(), &modelsResponse)).To(Succeed()) + Expect(modelsResponse).To(HaveLen(2)) + byName := map[string]map[string]any{} + for _, model := range modelsResponse { + byName[model["model_name"].(string)] = model + Expect(model).ToNot(HaveKey("model_opts_blob")) + } + + Expect(byName["current-model"]).To(SatisfyAll( + HaveKeyWithValue("state", "loaded"), + HaveKeyWithValue("config_revision", "revision-current"), + HaveKeyWithValue("effective_options_hash", "options-current"), + )) + Expect(byName["changed-model"]).To(SatisfyAll( + HaveKeyWithValue("state", "unloading"), + HaveKeyWithValue("config_revision", "revision-old"), + HaveKeyWithValue("effective_options_hash", "options-old"), + HaveKeyWithValue("cleanup_error", "worker unreachable"), + HaveKey("cleanup_next_retry_at"), + )) + }) + }) }) diff --git a/core/http/endpoints/localai/toggle_model.go b/core/http/endpoints/localai/toggle_model.go index 545fdc8af290..8b988c77a4cd 100644 --- a/core/http/endpoints/localai/toggle_model.go +++ b/core/http/endpoints/localai/toggle_model.go @@ -9,7 +9,6 @@ import ( "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/services/galleryop" "github.com/mudler/LocalAI/core/services/modeladmin" - "github.com/mudler/LocalAI/pkg/model" ) // ToggleModelEndpoint handles enabling or disabling a model from being loaded on demand. @@ -25,15 +24,15 @@ import ( // @Failure 404 {object} ModelResponse // @Failure 500 {object} ModelResponse // @Router /api/models/{name}/{action} [put] -func ToggleStateModelEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, gs *galleryop.GalleryService, appConfig *config.ApplicationConfig) echo.HandlerFunc { - svc := modeladmin.NewConfigService(cl, appConfig) +func ToggleStateModelEndpoint(cl *config.ModelConfigLoader, gs *galleryop.GalleryService, appConfig *config.ApplicationConfig, lifecycle ...modeladmin.ModelRevisionLifecycle) echo.HandlerFunc { + svc := modeladmin.NewConfigService(cl, appConfig, lifecycle...) return func(c echo.Context) error { modelName := c.Param("name") if decoded, err := url.PathUnescape(modelName); err == nil { modelName = decoded } action := modeladmin.Action(c.Param("action")) - result, err := svc.ToggleState(c.Request().Context(), modelName, action, ml) + result, err := svc.ToggleState(c.Request().Context(), modelName, action) if err != nil { return c.JSON(httpStatusForModelAdminError(err), ModelResponse{Success: false, Error: err.Error()}) } @@ -42,13 +41,13 @@ func ToggleStateModelEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoade // local loader; tell peers to refresh so the model's availability is // consistent across replicas. No-op in standalone mode. if gs != nil { - gs.BroadcastModelsChanged(modelName, "install") + gs.BroadcastModelsChangedRevision(modelName, "install", result.ConfigRevision) } msg := fmt.Sprintf("Model '%s' has been %sd successfully.", modelName, action) if action == modeladmin.ActionDisable { msg += " The model will not be loaded on demand until re-enabled." } - return c.JSON(http.StatusOK, ModelResponse{Success: true, Message: msg, Filename: result.Filename}) + return c.JSON(http.StatusOK, ModelResponse{Success: true, Message: msg, Filename: result.Filename, ConfigRevision: result.ConfigRevision, PendingCleanup: result.PendingCleanup}) } } diff --git a/core/http/endpoints/localai/toggle_model_test.go b/core/http/endpoints/localai/toggle_model_test.go new file mode 100644 index 000000000000..2a2ef31da66a --- /dev/null +++ b/core/http/endpoints/localai/toggle_model_test.go @@ -0,0 +1,54 @@ +package localai_test + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/config" + . "github.com/mudler/LocalAI/core/http/endpoints/localai" + "github.com/mudler/LocalAI/pkg/system" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Toggle model endpoint", func() { + It("always reports the saved revision and pending cleanup count", func() { + tempDir, err := os.MkdirTemp("", "toggle-model-test-*") + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(os.RemoveAll, tempDir) + + systemState, err := system.GetSystemState(system.WithModelPath(tempDir)) + Expect(err).NotTo(HaveOccurred()) + appConfig := config.NewApplicationConfig(config.WithSystemState(systemState)) + loader := config.NewModelConfigLoader(tempDir) + Expect(os.WriteFile(filepath.Join(tempDir, "model.yaml"), []byte("name: model\nbackend: llama-cpp\n"), 0o644)).To(Succeed()) + Expect(loader.LoadModelConfigsFromPath(tempDir)).To(Succeed()) + lifecycle := &endpointLifecycleRecorder{} + app := echo.New() + app.PUT("/api/models/:name/:action", ToggleStateModelEndpoint(loader, nil, appConfig, lifecycle)) + + request := func(action string) map[string]any { + req := httptest.NewRequest(http.MethodPut, "/api/models/model/"+action, nil).WithContext(context.Background()) + rec := httptest.NewRecorder() + app.ServeHTTP(rec, req) + Expect(rec.Code).To(Equal(http.StatusOK), rec.Body.String()) + var response map[string]any + Expect(json.Unmarshal(rec.Body.Bytes(), &response)).To(Succeed()) + return response + } + + disabled := request("disable") + Expect(disabled).To(HaveKeyWithValue("config_revision", Not(BeEmpty()))) + Expect(disabled).To(HaveKeyWithValue("pending_cleanup", BeNumerically("==", 0))) + + lifecycle.pendingCleanup = 3 + enabled := request("enable") + Expect(enabled).To(HaveKeyWithValue("config_revision", Not(BeEmpty()))) + Expect(enabled).To(HaveKeyWithValue("pending_cleanup", BeNumerically("==", 3))) + }) +}) diff --git a/core/http/endpoints/localai/types.go b/core/http/endpoints/localai/types.go index f1c507472496..93928ed9b2a7 100644 --- a/core/http/endpoints/localai/types.go +++ b/core/http/endpoints/localai/types.go @@ -2,10 +2,12 @@ package localai // ModelResponse represents the common response structure for model operations type ModelResponse struct { - Success bool `json:"success"` - Message string `json:"message"` - Filename string `json:"filename,omitempty"` - Config any `json:"config,omitempty"` - Error string `json:"error,omitempty"` - Details []string `json:"details,omitempty"` + Success bool `json:"success"` + Message string `json:"message"` + Filename string `json:"filename,omitempty"` + Config any `json:"config,omitempty"` + Error string `json:"error,omitempty"` + Details []string `json:"details,omitempty"` + ConfigRevision string `json:"config_revision,omitempty"` + PendingCleanup int `json:"pending_cleanup"` } diff --git a/core/http/routes/localai.go b/core/http/routes/localai.go index 1da4683db85c..2120a07f1f2e 100644 --- a/core/http/routes/localai.go +++ b/core/http/routes/localai.go @@ -86,13 +86,13 @@ func RegisterLocalAIRoutes(router *echo.Echo, router.POST("/models/import-uri", localai.ImportModelURIEndpoint(cl, appConfig, galleryService, opcache), adminMiddleware) // Custom model edit endpoint - router.POST("/models/edit/:name", localai.EditModelEndpoint(cl, ml, galleryService, appConfig), adminMiddleware) + router.POST("/models/edit/:name", localai.EditModelEndpoint(cl, galleryService, appConfig, modelRevisionLifecycleFor(app)), adminMiddleware) // List model aliases endpoint router.GET("/api/aliases", localai.ListAliasesEndpoint(cl), adminMiddleware) // Toggle model enable/disable endpoint - router.PUT("/models/toggle-state/:name/:action", localai.ToggleStateModelEndpoint(cl, ml, galleryService, appConfig), adminMiddleware) + router.PUT("/models/toggle-state/:name/:action", localai.ToggleStateModelEndpoint(cl, galleryService, appConfig, modelRevisionLifecycleFor(app)), adminMiddleware) // Toggle model pinned status endpoint router.PUT("/models/toggle-pinned/:name/:action", localai.TogglePinnedModelEndpoint(cl, appConfig, func() { diff --git a/core/http/routes/model_revision_lifecycle.go b/core/http/routes/model_revision_lifecycle.go new file mode 100644 index 000000000000..820844db4206 --- /dev/null +++ b/core/http/routes/model_revision_lifecycle.go @@ -0,0 +1,17 @@ +package routes + +import ( + "github.com/mudler/LocalAI/core/application" + "github.com/mudler/LocalAI/core/services/modeladmin" +) + +func modelRevisionLifecycleFor(app *application.Application) modeladmin.ModelRevisionLifecycle { + if app == nil || app.Distributed() == nil { + if app == nil { + return nil + } + return modeladmin.NewLocalModelRevisionLifecycle(app.ModelLoader()) + } + distributed := app.Distributed() + return modeladmin.NewDistributedModelRevisionLifecycle(distributed.Registry, distributed.ModelCleanup) +} diff --git a/core/http/routes/ui_api.go b/core/http/routes/ui_api.go index ad2ba25647be..d888dbde07b6 100644 --- a/core/http/routes/ui_api.go +++ b/core/http/routes/ui_api.go @@ -1181,7 +1181,7 @@ func RegisterUIAPIRoutes(app *echo.Echo, cl *config.ModelConfigLoader, ml *model app.GET("/api/models/config-metadata/autocomplete/:provider", localai.AutocompleteEndpoint(cl, ml, appConfig), adminMiddleware) // PATCH config endpoint - partial update using nested JSON merge - app.PATCH("/api/models/config-json/:name", localai.PatchConfigEndpoint(cl, ml, galleryService, appConfig), adminMiddleware) + app.PATCH("/api/models/config-json/:name", localai.PatchConfigEndpoint(cl, galleryService, appConfig, modelRevisionLifecycleFor(applicationInstance)), adminMiddleware) // VRAM estimation endpoint app.POST("/api/models/vram-estimate", localai.VRAMEstimateEndpoint(cl, appConfig), adminMiddleware) diff --git a/core/services/galleryop/model_revision_delete_test.go b/core/services/galleryop/model_revision_delete_test.go new file mode 100644 index 000000000000..7da910102ed6 --- /dev/null +++ b/core/services/galleryop/model_revision_delete_test.go @@ -0,0 +1,231 @@ +package galleryop + +import ( + "context" + "errors" + "os" + "path/filepath" + "time" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/gallery" + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/pkg/modelartifacts" + "github.com/mudler/LocalAI/pkg/system" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +type deleteRevisionLifecycle struct { + applied bool + revision string + err error +} + +func (l *deleteRevisionLifecycle) ApplyConfigRevisions(_ context.Context, transitions []config.ModelConfigRevisionTransition) (int, error) { + Expect(transitions).To(HaveLen(1)) + Expect(transitions[0].ModelName).To(Equal("doomed")) + Expect(transitions[0].Disabled).To(BeTrue()) + l.applied = true + l.revision = transitions[0].ConfigRevision + return 1, l.err +} + +type orderedDeleteManager struct { + called bool + path string +} + +type realDeletingManager struct { + state *system.SystemState + afterDelete func() error +} + +type countingMessagingClient struct{ subjects []string } + +func (c *countingMessagingClient) Publish(subject string, _ any) error { + c.subjects = append(c.subjects, subject) + return nil +} +func (c *countingMessagingClient) Subscribe(string, func([]byte)) (messaging.Subscription, error) { + return nil, nil +} +func (c *countingMessagingClient) QueueSubscribe(string, string, func([]byte)) (messaging.Subscription, error) { + return nil, nil +} +func (c *countingMessagingClient) QueueSubscribeReply(string, string, func([]byte, func([]byte))) (messaging.Subscription, error) { + return nil, nil +} +func (c *countingMessagingClient) SubscribeReply(string, func([]byte, func([]byte))) (messaging.Subscription, error) { + return nil, nil +} +func (c *countingMessagingClient) Request(string, []byte, time.Duration) ([]byte, error) { + return nil, nil +} +func (c *countingMessagingClient) IsConnected() bool { return true } +func (c *countingMessagingClient) Close() {} + +func (m *realDeletingManager) DeleteModel(name string) error { + if err := gallery.DeleteModelFromSystem(m.state, name); err != nil { + return err + } + if m.afterDelete != nil { + return m.afterDelete() + } + return nil +} + +func (m *realDeletingManager) InstallModel(context.Context, *ManagementOp[gallery.GalleryModel, gallery.ModelConfig], ProgressCallback) error { + return nil +} + +func (m *orderedDeleteManager) DeleteModel(name string) error { + m.called = true + Expect(name).To(Equal("doomed")) + if m.path != "" { + return os.Remove(m.path) + } + return nil +} + +type rejectingDeleteMaterializer struct{ calls int } + +func (m *rejectingDeleteMaterializer) Ensure(context.Context, string, modelartifacts.Spec) (modelartifacts.Result, error) { + m.calls++ + return modelartifacts.Result{}, errors.New("deleted config was preloaded") +} + +func (m *orderedDeleteManager) InstallModel(context.Context, *ManagementOp[gallery.GalleryModel, gallery.ModelConfig], ProgressCallback) error { + return nil +} + +var _ = Describe("model deletion revision lifecycle", func() { + It("deletes the real config, publishes its tombstone, and replaces the loader from disk", func() { + dir := GinkgoT().TempDir() + appConfig := &config.ApplicationConfig{SystemState: &system.SystemState{Model: system.Model{ModelsPath: dir}}} + loader := config.NewModelConfigLoader(dir) + path := filepath.Join(dir, "doomed.yaml") + Expect(os.WriteFile(path, []byte("name: doomed\nbackend: llama-cpp\n"), 0644)).To(Succeed()) + Expect(loader.LoadModelConfigsFromPath(dir)).To(Succeed()) + lifecycle := &deleteRevisionLifecycle{} + service := NewGalleryService(appConfig, nil) + service.SetModelRevisionLifecycle(lifecycle) + service.SetModelManager(&orderedDeleteManager{path: path}) + op := &ManagementOp[gallery.GalleryModel, gallery.ModelConfig]{ + ID: "delete-operation", GalleryElementName: "doomed", Delete: true, Context: context.Background(), + } + + Expect(service.modelHandler(op, loader, appConfig.SystemState)).To(Succeed()) + Expect(lifecycle.revision).To(HaveLen(64)) + Expect(path).NotTo(BeAnExistingFile()) + _, ok := loader.GetModelConfig("doomed") + Expect(ok).To(BeFalse()) + restarted := config.NewModelConfigLoader(dir) + Expect(restarted.LoadModelConfigsFromPath(dir)).To(Succeed()) + _, ok = restarted.GetModelConfig("doomed") + Expect(ok).To(BeFalse()) + }) + + It("keeps standalone deletion authoritative and never preloads the deleted config", func() { + dir := GinkgoT().TempDir() + appConfig := &config.ApplicationConfig{SystemState: &system.SystemState{Model: system.Model{ModelsPath: dir}}} + materializer := &rejectingDeleteMaterializer{} + loader := config.NewModelConfigLoader(dir, config.WithArtifactMaterializer(materializer)) + path := filepath.Join(dir, "doomed.yaml") + Expect(os.WriteFile(path, []byte("name: doomed\nbackend: llama-cpp\nartifacts:\n - name: model\n target: model\n source: {type: huggingface, repo: owner/repo}\n"), 0644)).To(Succeed()) + Expect(loader.LoadModelConfigsFromPath(dir)).To(Succeed()) + service := NewGalleryService(appConfig, nil) + service.SetModelManager(&orderedDeleteManager{path: path}) + op := &ManagementOp[gallery.GalleryModel, gallery.ModelConfig]{ + ID: "delete-operation", GalleryElementName: "doomed", Delete: true, Context: context.Background(), + } + + Expect(service.modelHandler(op, loader, appConfig.SystemState)).To(Succeed()) + Expect(materializer.calls).To(Equal(0)) + _, ok := loader.GetModelConfig("doomed") + Expect(ok).To(BeFalse()) + }) + + It("does not replace the authoritative loader when tombstone publication fails", func() { + dir := GinkgoT().TempDir() + appConfig := &config.ApplicationConfig{SystemState: &system.SystemState{Model: system.Model{ModelsPath: dir}}} + loader := config.NewModelConfigLoader(dir) + path := filepath.Join(dir, "doomed.yaml") + Expect(os.WriteFile(path, []byte("name: doomed\nbackend: llama-cpp\n"), 0644)).To(Succeed()) + Expect(loader.LoadModelConfigsFromPath(dir)).To(Succeed()) + lifecycle := &deleteRevisionLifecycle{err: errors.New("registry unavailable")} + manager := &realDeletingManager{state: appConfig.SystemState} + service := NewGalleryService(appConfig, nil) + service.SetModelRevisionLifecycle(lifecycle) + service.SetModelManager(manager) + op := &ManagementOp[gallery.GalleryModel, gallery.ModelConfig]{ + ID: "delete-operation", GalleryElementName: "doomed", Delete: true, Context: context.Background(), + } + + Expect(service.modelHandler(op, loader, appConfig.SystemState)).To(MatchError(ContainSubstring("registry unavailable"))) + Expect(path).To(BeAnExistingFile()) + _, ok := loader.GetModelConfig("doomed") + Expect(ok).To(BeTrue()) + restarted := config.NewModelConfigLoader(dir) + Expect(restarted.LoadModelConfigsFromPath(dir)).To(Succeed()) + _, ok = restarted.GetModelConfig("doomed") + Expect(ok).To(BeTrue()) + }) + + DescribeTable("rolls back real deletion before the commit boundary", + func(failurePoint string) { + dir := GinkgoT().TempDir() + materializer := &rejectingDeleteMaterializer{} + appConfig := &config.ApplicationConfig{ + SystemState: &system.SystemState{Model: system.Model{ModelsPath: dir}}, + ModelArtifactMaterializer: materializer, + } + loader := config.NewModelConfigLoader(dir, config.WithArtifactMaterializer(materializer)) + configPath := filepath.Join(dir, "doomed.yaml") + metadataPath := filepath.Join(dir, gallery.GalleryFileName("doomed")) + configData := []byte("name: doomed\nbackend: llama-cpp\n") + metadataData := []byte("files: []\n") + Expect(os.WriteFile(configPath, configData, 0640)).To(Succeed()) + Expect(os.WriteFile(metadataPath, metadataData, 0600)).To(Succeed()) + + manager := &realDeletingManager{state: appConfig.SystemState} + lifecycle := &deleteRevisionLifecycle{} + switch failurePoint { + case "parse": + manager.afterDelete = func() error { + return os.WriteFile(filepath.Join(dir, "broken.yaml"), []byte("name: ["), 0644) + } + case "preload": + Expect(os.WriteFile(filepath.Join(dir, "survivor.yaml"), []byte("name: survivor\nbackend: transformers\nartifacts:\n - name: model\n target: model\n source: {type: huggingface, repo: owner/repo}\n"), 0644)).To(Succeed()) + case "lifecycle": + lifecycle.err = errors.New("injected lifecycle failure") + } + Expect(loader.LoadModelConfigsFromPath(dir, appConfig.ToConfigLoaderOptions()...)).To(Succeed()) + + service := NewGalleryService(appConfig, nil) + bus := &countingMessagingClient{} + service.SetNATSClient(bus) + service.SetModelManager(manager) + service.SetModelRevisionLifecycle(lifecycle) + op := &ManagementOp[gallery.GalleryModel, gallery.ModelConfig]{ + ID: "delete-operation", GalleryElementName: "doomed", Delete: true, Context: context.Background(), + } + + Expect(service.modelHandler(op, loader, appConfig.SystemState)).ToNot(Succeed()) + Expect(bus.subjects).NotTo(ContainElement(messaging.SubjectCacheInvalidateModels)) + Expect(os.ReadFile(configPath)).To(Equal(configData)) + Expect(os.ReadFile(metadataPath)).To(Equal(metadataData)) + Expect(filepath.Join(dir, "broken.yaml")).NotTo(BeAnExistingFile()) + loaded, ok := loader.GetModelConfig("doomed") + Expect(ok).To(BeTrue()) + Expect(loaded.Name).To(Equal("doomed")) + fresh := config.NewModelConfigLoader(dir) + Expect(fresh.LoadModelConfigsFromPath(dir)).To(Succeed()) + _, ok = fresh.GetModelConfig("doomed") + Expect(ok).To(BeTrue()) + }, + Entry("when authoritative parsing fails", "parse"), + Entry("when preload fails", "preload"), + Entry("when lifecycle publication fails", "lifecycle"), + ) +}) diff --git a/core/services/galleryop/models.go b/core/services/galleryop/models.go index f9703797040d..c308215f189e 100644 --- a/core/services/galleryop/models.go +++ b/core/services/galleryop/models.go @@ -2,11 +2,14 @@ package galleryop import ( "context" + "crypto/sha256" "encoding/json" "errors" "fmt" "os" + "path/filepath" "slices" + "strings" "time" "github.com/mudler/LocalAI/core/config" @@ -14,8 +17,10 @@ import ( "github.com/mudler/LocalAI/core/services/messaging" "github.com/mudler/LocalAI/pkg/model" "github.com/mudler/LocalAI/pkg/modelartifacts" + "github.com/mudler/LocalAI/pkg/safefile" "github.com/mudler/LocalAI/pkg/system" "github.com/mudler/LocalAI/pkg/utils" + "github.com/mudler/xlog" "gopkg.in/yaml.v3" ) @@ -24,7 +29,37 @@ const ( ) func (g *GalleryService) modelHandler(op *ManagementOp[gallery.GalleryModel, gallery.ModelConfig], cl *config.ModelConfigLoader, systemState *system.SystemState) error { + if op.Delete && cl != nil { + return cl.WithModelConfigMutation(func() error { + return g.modelHandlerLocked(op, cl, systemState) + }) + } + return g.modelHandlerLocked(op, cl, systemState) +} + +func (g *GalleryService) modelHandlerLocked(op *ManagementOp[gallery.GalleryModel, gallery.ModelConfig], cl *config.ModelConfigLoader, systemState *system.SystemState) (returnErr error) { utils.ResetDownloadTimers() + var deleteSnapshot *modelConfigFilesSnapshot + deleteStarted := false + deleteCommitted := false + var priorConfigs []config.ModelConfig + if op.Delete && cl != nil && systemState != nil { + var err error + deleteSnapshot, err = snapshotModelConfigFiles(systemState.Model.ModelsPath) + if err != nil { + return err + } + priorConfigs = cl.GetAllModelsConfigs() + defer func() { + if !deleteStarted || deleteCommitted { + return + } + if err := deleteSnapshot.restore(); err != nil { + returnErr = errors.Join(returnErr, fmt.Errorf("restore model configuration after failed deletion: %w", err)) + } + cl.ReplaceModelConfigs(priorConfigs) + }() + } // Dedup check in distributed mode — skip if another instance is already processing this element if g.galleryStore != nil && op.GalleryElementName != "" && !op.Delete { @@ -101,7 +136,10 @@ func (g *GalleryService) modelHandler(op *ManagementOp[gallery.GalleryModel, gal } var err error + configRevision := "" if op.Delete { + configRevision = fmt.Sprintf("%x", sha256.Sum256([]byte("deleted\x00"+op.GalleryElementName))) + deleteStarted = true err = g.modelManager.DeleteModel(op.GalleryElementName) } else { err = g.modelManager.InstallModel(operationCtx, op, progressCallback) @@ -137,19 +175,39 @@ func (g *GalleryService) modelHandler(op *ManagementOp[gallery.GalleryModel, gal } } - // Reload models - err = cl.LoadModelConfigsFromPath(systemState.Model.ModelsPath, g.appConfig.ToConfigLoaderOptions()...) + // Parse a complete disk snapshot. LoadModelConfigsFromPath is additive on + // an existing loader, so using it directly would retain a just-deleted + // model and could preload artifacts for a config that no longer exists. + authoritative := config.NewModelConfigLoader(systemState.Model.ModelsPath) + err = authoritative.LoadModelConfigsFromPathStrict(systemState.Model.ModelsPath, g.appConfig.ToConfigLoaderOptions()...) if err != nil { return err } - + cl.ReplaceModelConfigs(authoritative.GetAllModelsConfigs()) err = cl.PreloadWithContext(operationCtx, systemState.Model.ModelsPath) if err != nil { return err } - // Tell peer replicas to refresh their own ModelConfigLoader. The local - // LoadModelConfigsFromPath above already covered THIS replica; without + // Lifecycle publication is the irreversible boundary. File mutation, + // authoritative parsing, loader replacement, and preload have all completed, + // so no later failure can roll local configuration back behind an accepted + // registry revision. + if op.Delete && g.modelRevisionLifecycle != nil { + pending, lifecycleErr := g.modelRevisionLifecycle.ApplyConfigRevisions(operationCtx, []config.ModelConfigRevisionTransition{{ + ModelName: op.GalleryElementName, ConfigRevision: configRevision, Disabled: true, + }}) + if lifecycleErr != nil { + return lifecycleErr + } + if pending > 0 { + xlog.Warn("Model deletion continuing with exact cleanup pending", "model", op.GalleryElementName, "configRevision", configRevision, "pendingCleanup", pending) + } + } + deleteCommitted = true + + // Tell peer replicas to refresh their own ModelConfigLoader. The + // authoritative replacement above already covered THIS replica; without // this broadcast a chat completion routed by the load balancer to a peer // would fail to find a model just installed. op2 := "install" @@ -157,8 +215,9 @@ func (g *GalleryService) modelHandler(op *ManagementOp[gallery.GalleryModel, gal op2 = "delete" } g.publishCacheInvalidate(messaging.SubjectCacheInvalidateModels, messaging.CacheInvalidateEvent{ - Element: op.GalleryElementName, - Op: op2, + Element: op.GalleryElementName, + Op: op2, + ConfigRevision: configRevision, }) legacyCoalescer.Close() @@ -174,6 +233,86 @@ func (g *GalleryService) modelHandler(op *ManagementOp[gallery.GalleryModel, gal return nil } +type savedModelConfigFile struct { + data []byte + mode os.FileMode +} + +type modelConfigFilesSnapshot struct { + dir string + files map[string]savedModelConfigFile +} + +func isModelConfigMetadata(name string) bool { + lower := strings.ToLower(name) + return strings.HasSuffix(lower, ".yaml") || strings.HasSuffix(lower, ".yml") +} + +func snapshotModelConfigFiles(dir string) (*modelConfigFilesSnapshot, error) { + entries, err := os.ReadDir(dir) + if err != nil { + return nil, fmt.Errorf("snapshot model configuration: %w", err) + } + snapshot := &modelConfigFilesSnapshot{dir: dir, files: map[string]savedModelConfigFile{}} + for _, entry := range entries { + if entry.IsDir() || !isModelConfigMetadata(entry.Name()) { + continue + } + data, mode, err := safefile.ReadRegularAt(dir, entry.Name()) + if err != nil { + return nil, fmt.Errorf("snapshot model configuration metadata %q: %w", entry.Name(), err) + } + snapshot.files[entry.Name()] = savedModelConfigFile{data: data, mode: mode} + } + return snapshot, nil +} + +func (s *modelConfigFilesSnapshot) restore() error { + entries, err := os.ReadDir(s.dir) + if err != nil { + return err + } + var restoreErr error + for _, entry := range entries { + if entry.IsDir() || !isModelConfigMetadata(entry.Name()) { + continue + } + if _, exists := s.files[entry.Name()]; exists { + continue + } + if err := os.Remove(filepath.Join(s.dir, entry.Name())); err != nil && !errors.Is(err, os.ErrNotExist) { + restoreErr = errors.Join(restoreErr, err) + } + } + for name, file := range s.files { + if err := writeRestoredConfigFile(filepath.Join(s.dir, name), file.data, file.mode); err != nil { + restoreErr = errors.Join(restoreErr, err) + } + } + return restoreErr +} + +func writeRestoredConfigFile(path string, data []byte, mode os.FileMode) error { + tmp, err := os.CreateTemp(filepath.Dir(path), ".model-config-restore-*") + if err != nil { + return err + } + tmpPath := tmp.Name() + defer func() { _ = os.Remove(tmpPath) }() + if err := tmp.Chmod(mode); err != nil { + _ = tmp.Close() + return err + } + if _, err := tmp.Write(data); err != nil { + _ = tmp.Close() + return err + } + if err := tmp.Close(); err != nil { + return err + } + return os.Rename(tmpPath, path) +} + func installModelFromRemoteConfig(ctx context.Context, systemState *system.SystemState, modelLoader *model.ModelLoader, req gallery.GalleryModel, downloadStatus func(string, string, string, float64), enforceScan, automaticallyInstallBackend bool, backendGalleries []config.Gallery, requireBackendIntegrity bool, options ...gallery.InstallOption) error { config, err := gallery.GetGalleryConfigFromURLWithContext[gallery.ModelConfig](ctx, req.URL, systemState.Model.ModelsPath) if err != nil { diff --git a/core/services/galleryop/service.go b/core/services/galleryop/service.go index 7a51fe440c1a..2a7410ab35a1 100644 --- a/core/services/galleryop/service.go +++ b/core/services/galleryop/service.go @@ -55,7 +55,20 @@ type GalleryService struct { // load-balances onto this replica can find the just-installed model. // The originating replica reloads inline (models.go) so it does not need // the hook. - OnModelsChanged func(messaging.CacheInvalidateEvent) + OnModelsChanged func(messaging.CacheInvalidateEvent) + modelRevisionLifecycle interface { + ApplyConfigRevisions(context.Context, []config.ModelConfigRevisionTransition) (int, error) + } +} + +// SetModelRevisionLifecycle wires the distributed config-generation boundary +// into gallery deletion without coupling gallery operations to node internals. +func (g *GalleryService) SetModelRevisionLifecycle(lifecycle interface { + ApplyConfigRevisions(context.Context, []config.ModelConfigRevisionTransition) (int, error) +}) { + g.Lock() + defer g.Unlock() + g.modelRevisionLifecycle = lifecycle } func NewGalleryService(appConfig *config.ApplicationConfig, ml *model.ModelLoader) *GalleryService { @@ -235,9 +248,16 @@ func (g *GalleryService) publishCacheInvalidate(subject string, evt messaging.Ca // disk) or "delete" for a removal (the element must be pruned from memory, // which a reload-from-path cannot do because the loader is additive). func (g *GalleryService) BroadcastModelsChanged(element, op string) { + g.BroadcastModelsChangedRevision(element, op, "") +} + +// BroadcastModelsChangedRevision includes the accepted semantic generation so +// peers can apply the same registry transition idempotently. +func (g *GalleryService) BroadcastModelsChangedRevision(element, op, configRevision string) { g.publishCacheInvalidate(messaging.SubjectCacheInvalidateModels, messaging.CacheInvalidateEvent{ - Element: element, - Op: op, + Element: element, + Op: op, + ConfigRevision: configRevision, }) } diff --git a/core/services/messaging/subjects.go b/core/services/messaging/subjects.go index 5f09adda5f50..c1f4cf8bfbab 100644 --- a/core/services/messaging/subjects.go +++ b/core/services/messaging/subjects.go @@ -296,6 +296,29 @@ func SubjectNodeBackendStop(nodeID string) string { return subjectNodePrefix + sanitizeSubjectToken(nodeID) + ".backend.stop" } +// SubjectNodeModelStop targets one supervisor process and acknowledges only +// after that process has exited and its worker-side resources are released. +func SubjectNodeModelStop(nodeID string) string { + return subjectNodePrefix + sanitizeSubjectToken(nodeID) + ".model.stop" +} + +type ModelStopRequest struct { + ModelName string `json:"model_name"` + ProcessKey string `json:"process_key"` + ExpectedAddress string `json:"expected_address"` + Force bool `json:"force,omitempty"` + ConfigRevision string `json:"config_revision,omitempty"` +} + +type ModelStopReply struct { + Matched bool `json:"matched"` + Freed bool `json:"freed"` + Terminated bool `json:"terminated"` + ProcessKey string `json:"process_key"` + Address string `json:"address,omitempty"` + Error string `json:"error,omitempty"` +} + // SubjectNodeBackendDelete tells a worker node to delete a backend (stop + remove files). // Uses NATS request-reply. func SubjectNodeBackendDelete(nodeID string) string { @@ -447,8 +470,9 @@ const ( // Element names a specific model/backend when known; empty means "the whole // set was touched, do a full reload." type CacheInvalidateEvent struct { - Element string `json:"element,omitempty"` - Op string `json:"op,omitempty"` // "install" | "delete" | "upgrade" + Element string `json:"element,omitempty"` + Op string `json:"op,omitempty"` // "install" | "delete" | "upgrade" + ConfigRevision string `json:"config_revision,omitempty"` } // SubjectCacheInvalidateCollection returns the NATS subject for collection cache invalidation. diff --git a/core/services/modeladmin/config.go b/core/services/modeladmin/config.go index f4fc53d977d4..23de357aa9e2 100644 --- a/core/services/modeladmin/config.go +++ b/core/services/modeladmin/config.go @@ -14,23 +14,38 @@ import ( "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/config/meta" "github.com/mudler/LocalAI/core/gallery" - "github.com/mudler/LocalAI/pkg/model" "github.com/mudler/LocalAI/pkg/utils" + "github.com/mudler/xlog" ) +// ModelRevisionLifecycle applies a persisted model configuration generation to +// the distributed registry. Implementations quarantine stale replicas before +// attempting cleanup. +type ModelRevisionLifecycle interface { + ApplyConfigRevisions(ctx context.Context, transitions []ModelRevisionTransition) (pendingCleanup int, err error) +} + +// ModelRevisionTransition describes one authoritative model identity. Related +// identities, such as both sides of a rename, are published atomically. +type ModelRevisionTransition = config.ModelConfigRevisionTransition + // ConfigService groups operations that read or mutate an installed model's // configuration on disk. It keeps the side-effect surface (loader reload, // model shutdown) explicit so callers know what gets touched. type ConfigService struct { Loader *config.ModelConfigLoader AppConfig *config.ApplicationConfig + Lifecycle ModelRevisionLifecycle } // NewConfigService returns a ConfigService bound to the supplied loader and -// app config. The loader and the system state in AppConfig are mandatory; the -// model loader is required only by EditYAML and ToggleState (for Shutdown). -func NewConfigService(loader *config.ModelConfigLoader, appConfig *config.ApplicationConfig) *ConfigService { - return &ConfigService{Loader: loader, AppConfig: appConfig} +// app config. The loader and the system state in AppConfig are mandatory. +func NewConfigService(loader *config.ModelConfigLoader, appConfig *config.ApplicationConfig, lifecycle ...ModelRevisionLifecycle) *ConfigService { + svc := &ConfigService{Loader: loader, AppConfig: appConfig} + if len(lifecycle) > 0 { + svc.Lifecycle = lifecycle[0] + } + return svc } // ConfigView is the on-disk YAML plus the parsed JSON view, returned by GetConfig. @@ -44,11 +59,19 @@ type ConfigView struct { // EditResult is what EditYAML returns to its caller. type EditResult struct { - Filename string - Renamed bool - OldName string - NewName string - Config config.ModelConfig + Filename string + Renamed bool + OldName string + NewName string + Config config.ModelConfig + ConfigRevision string + PendingCleanup int +} + +type PatchResult struct { + config.ModelConfig + ConfigRevision string + PendingCleanup int } // modelsPath is shorthand for the configured models directory. @@ -89,7 +112,17 @@ func (s *ConfigService) GetConfig(_ context.Context, name string) (*ConfigView, // config — which has SetDefaults applied and would persist runtime defaults // like top_p/temperature/mirostat), deep-merge the patch, validate, write, // reload, preload (preload errors are non-fatal — log only). -func (s *ConfigService) PatchConfig(_ context.Context, name string, patch map[string]any) (*config.ModelConfig, error) { +func (s *ConfigService) PatchConfig(ctx context.Context, name string, patch map[string]any) (*PatchResult, error) { + var result *PatchResult + err := s.Loader.WithModelConfigMutation(func() error { + var err error + result, err = s.patchConfig(ctx, name, patch) + return err + }) + return result, err +} + +func (s *ConfigService) patchConfig(ctx context.Context, name string, patch map[string]any) (*PatchResult, error) { if name == "" { return nil, ErrNameRequired } @@ -100,6 +133,9 @@ func (s *ConfigService) PatchConfig(_ context.Context, name string, patch map[st if !exists { return nil, ErrNotFound } + if patchedName, ok := patch["name"].(string); ok && patchedName != name { + return nil, fmt.Errorf("%w: PATCH cannot rename model %q to %q; use the model edit endpoint", ErrInvalidConfig, name, patchedName) + } configPath := cfg.GetModelConfigFile() if err := utils.VerifyPath(configPath, s.modelsPath()); err != nil { return nil, fmt.Errorf("%w: %v", ErrPathNotTrusted, err) @@ -133,15 +169,31 @@ func (s *ConfigService) PatchConfig(_ context.Context, name string, patch map[st if err := s.Loader.ValidateAliasTarget(&updated); err != nil { return nil, fmt.Errorf("%w: %v", ErrInvalidConfig, err) } - if err := writeFileAtomic(configPath, yamlData, 0644); err != nil { - return nil, fmt.Errorf("write config file: %w", err) - } - if err := s.Loader.LoadModelConfigsFromPath(s.modelsPath(), s.AppConfig.ToConfigLoaderOptions()...); err != nil { - return nil, fmt.Errorf("reload configs: %w", err) - } - // Preload is best-effort — a failure here doesn't undo the patch. - _ = s.Loader.Preload(s.modelsPath()) - return &updated, nil + var result *PatchResult + err = s.withMutationRollback([]string{configPath}, func() error { + if err := writeFileAtomic(configPath, yamlData, 0644); err != nil { + return fmt.Errorf("write config file: %w", err) + } + if err := s.Loader.LoadModelConfigsFromPath(s.modelsPath(), s.AppConfig.ToConfigLoaderOptions()...); err != nil { + return fmt.Errorf("reload configs: %w", err) + } + loaded, ok := s.Loader.GetModelConfig(updated.Name) + if !ok { + return fmt.Errorf("reload configs: model %q missing", updated.Name) + } + revision, err := config.ModelConfigRevision(&loaded) + if err != nil { + return fmt.Errorf("compute config revision: %w", err) + } + _ = s.Loader.Preload(s.modelsPath()) + pending, err := s.applyRevision(ctx, name, updated.Name, revision, updated.IsDisabled()) + if err != nil { + return err + } + result = &PatchResult{ModelConfig: updated, ConfigRevision: revision, PendingCleanup: pending} + return nil + }) + return result, err } // mapLeafFieldPaths returns the set of dotted config paths whose schema type is @@ -194,9 +246,18 @@ func patchMerge(dst, src map[string]any, mapLeaves map[string]struct{}, prefix s } // EditYAML replaces the YAML for an installed model, with optional rename -// support. ml may be nil; when set, EditYAML calls ml.ShutdownModel(oldName) -// after a successful write so the next inference picks up the new config. -func (s *ConfigService) EditYAML(_ context.Context, name string, body []byte, ml *model.ModelLoader) (*EditResult, error) { +// support, and applies the resulting semantic revision after reload. +func (s *ConfigService) EditYAML(ctx context.Context, name string, body []byte) (*EditResult, error) { + var result *EditResult + err := s.Loader.WithModelConfigMutation(func() error { + var err error + result, err = s.editYAML(ctx, name, body) + return err + }) + return result, err +} + +func (s *ConfigService) editYAML(ctx context.Context, name string, body []byte) (*EditResult, error) { if name == "" { return nil, ErrNameRequired } @@ -229,6 +290,7 @@ func (s *ConfigService) EditYAML(_ context.Context, name string, body []byte, ml } renamed := req.Name != name + paths := []string{configPath} if renamed { if strings.ContainsRune(req.Name, os.PathSeparator) || strings.Contains(req.Name, "/") || strings.Contains(req.Name, "\\") { return nil, ErrPathSeparator @@ -237,6 +299,7 @@ func (s *ConfigService) EditYAML(_ context.Context, name string, body []byte, ml return nil, fmt.Errorf("%w: %q", ErrConflict, req.Name) } newConfigPath := filepath.Join(modelsPath, req.Name+".yaml") + paths = append(paths, newConfigPath, filepath.Join(modelsPath, gallery.GalleryFileName(name)), filepath.Join(modelsPath, gallery.GalleryFileName(req.Name))) if err := utils.VerifyPath(newConfigPath, modelsPath); err != nil { return nil, fmt.Errorf("%w: %v", ErrPathNotTrusted, err) } @@ -245,46 +308,87 @@ func (s *ConfigService) EditYAML(_ context.Context, name string, body []byte, ml } else if !errors.Is(err, os.ErrNotExist) { return nil, fmt.Errorf("stat new config: %w", err) } - if err := writeFileAtomic(newConfigPath, body, 0644); err != nil { - return nil, fmt.Errorf("write new config: %w", err) + } + + var result *EditResult + err := s.withMutationRollback(paths, func() error { + if renamed { + newConfigPath := filepath.Join(modelsPath, req.Name+".yaml") + if err := writeFileAtomic(newConfigPath, body, 0644); err != nil { + return fmt.Errorf("write new config: %w", err) + } + if configPath != newConfigPath { + if err := os.Remove(configPath); err != nil && !errors.Is(err, os.ErrNotExist) { + return fmt.Errorf("remove old config: %w", err) + } + } + // Move the gallery metadata file so the delete flow can still find it. + oldGalleryPath := filepath.Join(modelsPath, gallery.GalleryFileName(name)) + newGalleryPath := filepath.Join(modelsPath, gallery.GalleryFileName(req.Name)) + if _, err := os.Stat(oldGalleryPath); err == nil { + if err := os.Rename(oldGalleryPath, newGalleryPath); err != nil { + return fmt.Errorf("rename gallery metadata: %w", err) + } + } + // Drop the stale in-memory entry before reload so we don't surface + // both names between scan steps. + s.Loader.RemoveModelConfig(name) + configPath = newConfigPath + } else { + if err := writeFileAtomic(configPath, body, 0644); err != nil { + return fmt.Errorf("write config: %w", err) + } + } + + if err := s.Loader.LoadModelConfigsFromPath(modelsPath, s.AppConfig.ToConfigLoaderOptions()...); err != nil { + return fmt.Errorf("reload configs: %w", err) } - if configPath != newConfigPath { - // Best-effort: a stale old file is cosmetic, not load-bearing. - _ = os.Remove(configPath) + loaded, ok := s.Loader.GetModelConfig(req.Name) + if !ok { + return fmt.Errorf("reload configs: model %q missing", req.Name) } - // Move the gallery metadata file so the delete flow can still find it. - oldGalleryPath := filepath.Join(modelsPath, gallery.GalleryFileName(name)) - newGalleryPath := filepath.Join(modelsPath, gallery.GalleryFileName(req.Name)) - if _, err := os.Stat(oldGalleryPath); err == nil { - _ = os.Rename(oldGalleryPath, newGalleryPath) + revision, err := config.ModelConfigRevision(&loaded) + if err != nil { + return fmt.Errorf("compute config revision: %w", err) } - // Drop the stale in-memory entry before reload so we don't surface - // both names between scan steps. - s.Loader.RemoveModelConfig(name) - configPath = newConfigPath - } else { - if err := writeFileAtomic(configPath, body, 0644); err != nil { - return nil, fmt.Errorf("write config: %w", err) + if err := s.Loader.Preload(modelsPath); err != nil { + return fmt.Errorf("preload after edit: %w", err) } - } + pending, err := s.applyRevision(ctx, name, req.Name, revision, req.IsDisabled()) + if err != nil { + return err + } + result = &EditResult{ + Filename: configPath, + Renamed: renamed, + OldName: name, + NewName: req.Name, + Config: req, + ConfigRevision: revision, + PendingCleanup: pending, + } + return nil + }) + return result, err +} - if err := s.Loader.LoadModelConfigsFromPath(modelsPath, s.AppConfig.ToConfigLoaderOptions()...); err != nil { - return nil, fmt.Errorf("reload configs: %w", err) - } - // Best-effort shutdown: the config is already written; if shutdown fails - // the caller can manually reload. The shutdown uses the OLD name because - // that's what the running instance was started with. - if ml != nil { - _ = ml.ShutdownModel(name) - } - if err := s.Loader.Preload(modelsPath); err != nil { - return nil, fmt.Errorf("preload after edit: %w", err) - } - return &EditResult{ - Filename: configPath, - Renamed: renamed, - OldName: name, - NewName: req.Name, - Config: req, - }, nil +func (s *ConfigService) applyRevision(ctx context.Context, oldName, newName, revision string, disabled bool) (int, error) { + if s.Lifecycle == nil { + return 0, nil + } + transitions := []ModelRevisionTransition{{ModelName: newName, ConfigRevision: revision, Disabled: disabled}} + if oldName != newName { + transitions = []ModelRevisionTransition{ + {ModelName: oldName, ConfigRevision: DeletedModelConfigRevision(oldName), Disabled: true}, + {ModelName: newName, ConfigRevision: revision, Disabled: disabled}, + } + } + pending, err := s.Lifecycle.ApplyConfigRevisions(ctx, transitions) + if err != nil { + return pending, fmt.Errorf("apply config revision: %w", err) + } + if pending > 0 { + xlog.Warn("Model configuration saved with cleanup pending", "model", newName, "configRevision", revision, "pendingCleanup", pending) + } + return pending, nil } diff --git a/core/services/modeladmin/config_test.go b/core/services/modeladmin/config_test.go index 36569c19b0d0..29d35a12e4a1 100644 --- a/core/services/modeladmin/config_test.go +++ b/core/services/modeladmin/config_test.go @@ -2,6 +2,7 @@ package modeladmin import ( "context" + "errors" "os" "path/filepath" @@ -10,9 +11,40 @@ import ( "gopkg.in/yaml.v3" "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/gallery" + "github.com/mudler/LocalAI/pkg/modelartifacts" "github.com/mudler/LocalAI/pkg/system" ) +type failingConfigMaterializer struct{ err error } + +func (f *failingConfigMaterializer) Ensure(context.Context, string, modelartifacts.Spec) (modelartifacts.Result, error) { + return modelartifacts.Result{}, f.err +} + +type fakeRevisionLifecycle struct { + calls []revisionLifecycleCall + batches [][]ModelRevisionTransition + pending int + err error +} + +type revisionLifecycleCall struct { + oldName, newName, revision string + disabled bool +} + +func (f *fakeRevisionLifecycle) ApplyConfigRevisions(_ context.Context, transitions []ModelRevisionTransition) (int, error) { + f.batches = append(f.batches, append([]ModelRevisionTransition(nil), transitions...)) + for _, transition := range transitions { + f.calls = append(f.calls, revisionLifecycleCall{ + oldName: transition.ModelName, newName: transition.ModelName, + revision: transition.ConfigRevision, disabled: transition.Disabled, + }) + } + return f.pending, f.err +} + // newTestService stands up a ConfigService backed by a tmp dir so the file IO // is real but isolated. The model loader is loaded against the same tmp path // so GetModelConfig works. @@ -48,6 +80,38 @@ var _ = Describe("ConfigService", func() { ctx = context.Background() }) + It("rejects a symlink when snapshotting a mutation", func() { + target := filepath.Join(dir, "target.yaml") + Expect(os.WriteFile(target, []byte("name: target\n"), 0o600)).To(Succeed()) + link := filepath.Join(dir, "link.yaml") + Expect(os.Symlink(target, link)).To(Succeed()) + + called := false + Expect(svc.withMutationRollback([]string{link}, func() error { + called = true + return nil + })).ToNot(Succeed()) + Expect(called).To(BeFalse()) + }) + + It("removes a symlink created at a previously absent rollback destination", func() { + target := filepath.Join(dir, "target.yaml") + Expect(os.WriteFile(target, []byte("unchanged"), 0o600)).To(Succeed()) + destination := filepath.Join(dir, "new.yaml") + mutationErr := errors.New("mutation failed") + + err := svc.withMutationRollback([]string{destination}, func() error { + Expect(os.Symlink(target, destination)).To(Succeed()) + return mutationErr + }) + Expect(err).To(MatchError(mutationErr)) + _, statErr := os.Lstat(destination) + Expect(statErr).To(MatchError(os.ErrNotExist)) + data, readErr := os.ReadFile(target) + Expect(readErr).NotTo(HaveOccurred()) + Expect(data).To(Equal([]byte("unchanged"))) + }) + Describe("GetConfig", func() { It("round-trips YAML from disk and exposes the parsed JSON", func() { writeModelYAML(svc, dir, "qwen", map[string]any{"backend": "llama-cpp", "context_size": 4096}) @@ -70,6 +134,109 @@ var _ = Describe("ConfigService", func() { }) Describe("PatchConfig", func() { + It("rejects a patched name change before mutating any state", func() { + lifecycle := &fakeRevisionLifecycle{} + svc.Lifecycle = lifecycle + writeModelYAML(svc, dir, "qwen", map[string]any{"backend": "llama-cpp", "context_size": 4096}) + path := filepath.Join(dir, "qwen.yaml") + before, err := os.ReadFile(path) + Expect(err).ToNot(HaveOccurred()) + + _, err = svc.PatchConfig(ctx, "qwen", map[string]any{"name": "renamed", "context_size": 8192}) + Expect(err).To(MatchError(ContainSubstring("cannot rename"))) + Expect(errors.Is(err, ErrInvalidConfig)).To(BeTrue()) + after, err := os.ReadFile(path) + Expect(err).ToNot(HaveOccurred()) + Expect(after).To(Equal(before)) + Expect(filepath.Join(dir, "renamed.yaml")).NotTo(BeAnExistingFile()) + loaded, ok := svc.Loader.GetModelConfig("qwen") + Expect(ok).To(BeTrue()) + Expect(loaded.ContextSize).To(HaveValue(Equal(4096))) + _, renamed := svc.Loader.GetModelConfig("renamed") + Expect(renamed).To(BeFalse()) + Expect(lifecycle.calls).To(BeEmpty()) + }) + + It("accepts an omitted or unchanged patched name", func() { + writeModelYAML(svc, dir, "qwen", map[string]any{"backend": "llama-cpp", "context_size": 4096}) + + withoutName, err := svc.PatchConfig(ctx, "qwen", map[string]any{"context_size": 8192}) + Expect(err).ToNot(HaveOccurred()) + Expect(withoutName.Name).To(Equal("qwen")) + withSameName, err := svc.PatchConfig(ctx, "qwen", map[string]any{"name": "qwen", "context_size": 10000}) + Expect(err).ToNot(HaveOccurred()) + Expect(withSameName.Name).To(Equal("qwen")) + }) + + It("serializes lifecycle publication across local service instances", func() { + writeModelYAML(svc, dir, "qwen", map[string]any{"backend": "llama-cpp", "context_size": 4096}) + lifecycle := newBlockingRevisionLifecycle() + first := NewConfigService(svc.Loader, svc.AppConfig, lifecycle) + second := NewConfigService(svc.Loader, svc.AppConfig, lifecycle) + firstDone := make(chan error, 1) + secondDone := make(chan error, 1) + + go func() { + _, err := first.PatchConfig(ctx, "qwen", map[string]any{"context_size": 8192}) + firstDone <- err + }() + Eventually(lifecycle.entered).Should(Receive()) + go func() { + _, err := second.PatchConfig(ctx, "qwen", map[string]any{"context_size": 10000}) + secondDone <- err + }() + Consistently(secondDone).ShouldNot(Receive()) + Expect(readMap(filepath.Join(dir, "qwen.yaml"))).To(HaveKeyWithValue("context_size", 8192)) + + close(lifecycle.release) + Eventually(firstDone).Should(Receive(Succeed())) + Eventually(secondDone).Should(Receive(Succeed())) + loaded, ok := svc.Loader.GetModelConfig("qwen") + Expect(ok).To(BeTrue()) + Expect(loaded.ContextSize).To(HaveValue(Equal(10000))) + Expect(readMap(filepath.Join(dir, "qwen.yaml"))).To(HaveKeyWithValue("context_size", 10000)) + }) + + It("restores disk and loader when revision publication fails", func() { + svc.Lifecycle = &fakeRevisionLifecycle{err: errors.New("registry unavailable")} + writeModelYAML(svc, dir, "qwen", map[string]any{"backend": "llama-cpp", "context_size": 4096}) + + _, err := svc.PatchConfig(ctx, "qwen", map[string]any{"context_size": 8192}) + Expect(err).To(MatchError(ContainSubstring("registry unavailable"))) + Expect(readMap(filepath.Join(dir, "qwen.yaml"))).To(HaveKeyWithValue("context_size", 4096)) + loaded, ok := svc.Loader.GetModelConfig("qwen") + Expect(ok).To(BeTrue()) + Expect(loaded.ContextSize).To(HaveValue(Equal(4096))) + restarted := config.NewModelConfigLoader(dir) + Expect(restarted.LoadModelConfigsFromPath(dir)).To(Succeed()) + reloaded, ok := restarted.GetModelConfig("qwen") + Expect(ok).To(BeTrue()) + Expect(reloaded.ContextSize).To(HaveValue(Equal(4096))) + }) + It("applies the persisted semantic revision and reports pending cleanup", func() { + lifecycle := &fakeRevisionLifecycle{pending: 2} + svc.Lifecycle = lifecycle + writeModelYAML(svc, dir, "qwen", map[string]any{"backend": "llama-cpp", "context_size": 4096}) + + updated, err := svc.PatchConfig(ctx, "qwen", map[string]any{"context_size": 8192}) + Expect(err).ToNot(HaveOccurred()) + Expect(updated.ConfigRevision).ToNot(BeEmpty()) + Expect(updated.PendingCleanup).To(Equal(2)) + Expect(lifecycle.calls).To(ConsistOf(revisionLifecycleCall{ + oldName: "qwen", newName: "qwen", revision: updated.ConfigRevision, + })) + }) + + It("keeps a durable patch successful when cleanup remains pending", func() { + lifecycle := &fakeRevisionLifecycle{pending: 1} + svc.Lifecycle = lifecycle + writeModelYAML(svc, dir, "qwen", map[string]any{"backend": "llama-cpp", "context_size": 4096}) + + updated, err := svc.PatchConfig(ctx, "qwen", map[string]any{"context_size": 8192}) + Expect(err).ToNot(HaveOccurred()) + Expect(updated.PendingCleanup).To(Equal(1)) + Expect(readMap(filepath.Join(dir, "qwen.yaml"))).To(HaveKeyWithValue("context_size", 8192)) + }) It("deep-merges the patch and preserves untouched siblings", func() { writeModelYAML(svc, dir, "qwen", map[string]any{ "backend": "llama-cpp", @@ -168,11 +335,103 @@ var _ = Describe("ConfigService", func() { }) Describe("EditYAML", func() { + It("does not publish an in-place revision when preload preparation fails", func() { + materializer := &failingConfigMaterializer{err: errors.New("artifact unavailable")} + svc.Loader = config.NewModelConfigLoader(dir, config.WithArtifactMaterializer(materializer)) + lifecycle := &fakeRevisionLifecycle{} + svc.Lifecycle = lifecycle + writeModelYAML(svc, dir, "qwen", map[string]any{"backend": "llama-cpp", "context_size": 4096}) + + body := []byte("name: qwen\nbackend: llama-cpp\ncontext_size: 8192\nartifacts:\n - name: model\n target: model\n source: {type: huggingface, repo: owner/repo}\n") + _, err := svc.EditYAML(ctx, "qwen", body) + Expect(err).To(MatchError(ContainSubstring("artifact unavailable"))) + Expect(lifecycle.calls).To(BeEmpty()) + Expect(readMap(filepath.Join(dir, "qwen.yaml"))).To(HaveKeyWithValue("context_size", 4096)) + loaded, ok := svc.Loader.GetModelConfig("qwen") + Expect(ok).To(BeTrue()) + Expect(loaded.ContextSize).To(HaveValue(Equal(4096))) + restarted := config.NewModelConfigLoader(dir) + Expect(restarted.LoadModelConfigsFromPath(dir)).To(Succeed()) + reloaded, ok := restarted.GetModelConfig("qwen") + Expect(ok).To(BeTrue()) + Expect(reloaded.ContextSize).To(HaveValue(Equal(4096))) + }) + + It("does not publish a rename revision when preload preparation fails", func() { + materializer := &failingConfigMaterializer{err: errors.New("artifact unavailable")} + svc.Loader = config.NewModelConfigLoader(dir, config.WithArtifactMaterializer(materializer)) + lifecycle := &fakeRevisionLifecycle{} + svc.Lifecycle = lifecycle + writeModelYAML(svc, dir, "old", map[string]any{"backend": "llama-cpp", "context_size": 4096}) + + body := []byte("name: new\nbackend: llama-cpp\ncontext_size: 8192\nartifacts:\n - name: model\n target: model\n source: {type: huggingface, repo: owner/repo}\n") + _, err := svc.EditYAML(ctx, "old", body) + Expect(err).To(MatchError(ContainSubstring("artifact unavailable"))) + Expect(lifecycle.calls).To(BeEmpty()) + Expect(filepath.Join(dir, "old.yaml")).To(BeAnExistingFile()) + Expect(filepath.Join(dir, "new.yaml")).NotTo(BeAnExistingFile()) + _, oldOK := svc.Loader.GetModelConfig("old") + _, newOK := svc.Loader.GetModelConfig("new") + Expect(oldOK).To(BeTrue()) + Expect(newOK).To(BeFalse()) + restarted := config.NewModelConfigLoader(dir) + Expect(restarted.LoadModelConfigsFromPath(dir)).To(Succeed()) + _, oldOK = restarted.GetModelConfig("old") + _, newOK = restarted.GetModelConfig("new") + Expect(oldOK).To(BeTrue()) + Expect(newOK).To(BeFalse()) + }) + + It("restores an in-place edit when revision publication fails", func() { + svc.Lifecycle = &fakeRevisionLifecycle{err: errors.New("registry unavailable")} + writeModelYAML(svc, dir, "qwen", map[string]any{"backend": "llama-cpp", "context_size": 4096}) + + _, err := svc.EditYAML(ctx, "qwen", []byte("name: qwen\nbackend: llama-cpp\ncontext_size: 8192\n")) + Expect(err).To(MatchError(ContainSubstring("registry unavailable"))) + Expect(readMap(filepath.Join(dir, "qwen.yaml"))).To(HaveKeyWithValue("context_size", 4096)) + loaded, ok := svc.Loader.GetModelConfig("qwen") + Expect(ok).To(BeTrue()) + Expect(loaded.ContextSize).To(HaveValue(Equal(4096))) + }) + + It("restores both identities and gallery metadata when rename publication fails", func() { + svc.Lifecycle = &fakeRevisionLifecycle{err: errors.New("registry unavailable")} + writeModelYAML(svc, dir, "old", map[string]any{"backend": "llama-cpp", "context_size": 4096}) + oldGallery := filepath.Join(dir, gallery.GalleryFileName("old")) + Expect(os.WriteFile(oldGallery, []byte("metadata"), 0644)).To(Succeed()) + + _, err := svc.EditYAML(ctx, "old", []byte("name: new\nbackend: llama-cpp\ncontext_size: 8192\n")) + Expect(err).To(MatchError(ContainSubstring("registry unavailable"))) + Expect(filepath.Join(dir, "old.yaml")).To(BeAnExistingFile()) + Expect(filepath.Join(dir, "new.yaml")).NotTo(BeAnExistingFile()) + Expect(oldGallery).To(BeAnExistingFile()) + Expect(filepath.Join(dir, gallery.GalleryFileName("new"))).NotTo(BeAnExistingFile()) + _, oldOK := svc.Loader.GetModelConfig("old") + _, newOK := svc.Loader.GetModelConfig("new") + Expect(oldOK).To(BeTrue()) + Expect(newOK).To(BeFalse()) + }) + It("applies both rename identities in one revision lifecycle batch", func() { + lifecycle := &fakeRevisionLifecycle{pending: 1} + svc.Lifecycle = lifecycle + writeModelYAML(svc, dir, "old", map[string]any{"backend": "llama-cpp"}) + body := []byte("name: new\nbackend: llama-cpp\ncontext_size: 8192\n") + + result, err := svc.EditYAML(ctx, "old", body) + Expect(err).ToNot(HaveOccurred()) + Expect(result.ConfigRevision).ToNot(BeEmpty()) + Expect(result.PendingCleanup).To(Equal(1)) + Expect(lifecycle.batches).To(HaveLen(1)) + Expect(lifecycle.calls).To(Equal([]revisionLifecycleCall{ + {oldName: "old", newName: "old", revision: DeletedModelConfigRevision("old"), disabled: true}, + {oldName: "new", newName: "new", revision: result.ConfigRevision}, + })) + }) It("renames the on-disk file and reindexes the loader", func() { writeModelYAML(svc, dir, "old-name", map[string]any{"backend": "llama-cpp"}) body := []byte("name: new-name\nbackend: llama-cpp\n") - result, err := svc.EditYAML(ctx, "old-name", body, nil) + result, err := svc.EditYAML(ctx, "old-name", body) Expect(err).ToNot(HaveOccurred()) Expect(result.Renamed).To(BeTrue()) Expect(result.OldName).To(Equal("old-name")) @@ -194,7 +453,7 @@ var _ = Describe("ConfigService", func() { writeModelYAML(svc, dir, "beta", map[string]any{"backend": "llama-cpp"}) body := []byte("name: beta\nbackend: llama-cpp\n") - _, err := svc.EditYAML(ctx, "alpha", body, nil) + _, err := svc.EditYAML(ctx, "alpha", body) Expect(err).To(MatchError(ErrConflict)) }) @@ -202,13 +461,13 @@ var _ = Describe("ConfigService", func() { writeModelYAML(svc, dir, "alpha", map[string]any{"backend": "llama-cpp"}) body := []byte("name: ../escape\nbackend: llama-cpp\n") - _, err := svc.EditYAML(ctx, "alpha", body, nil) + _, err := svc.EditYAML(ctx, "alpha", body) Expect(err).To(MatchError(ErrPathSeparator)) }) It("returns ErrEmptyBody when the body is nil", func() { writeModelYAML(svc, dir, "alpha", map[string]any{"backend": "llama-cpp"}) - _, err := svc.EditYAML(ctx, "alpha", nil, nil) + _, err := svc.EditYAML(ctx, "alpha", nil) Expect(err).To(MatchError(ErrEmptyBody)) }) @@ -216,7 +475,7 @@ var _ = Describe("ConfigService", func() { writeModelYAML(svc, dir, "base", map[string]any{"backend": "llama-cpp"}) body := []byte("name: base\nalias: ghost\n") - _, err := svc.EditYAML(ctx, "base", body, nil) + _, err := svc.EditYAML(ctx, "base", body) Expect(err).To(MatchError(ErrInvalidConfig)) Expect(err.Error()).To(ContainSubstring("ghost")) }) @@ -226,7 +485,7 @@ var _ = Describe("ConfigService", func() { writeModelYAML(svc, dir, "target", map[string]any{"backend": "llama-cpp"}) body := []byte("name: base\nalias: target\n") - _, err := svc.EditYAML(ctx, "base", body, nil) + _, err := svc.EditYAML(ctx, "base", body) Expect(err).ToNot(HaveOccurred()) }) }) diff --git a/core/services/modeladmin/lifecycle.go b/core/services/modeladmin/lifecycle.go new file mode 100644 index 000000000000..fa1c4676ec11 --- /dev/null +++ b/core/services/modeladmin/lifecycle.go @@ -0,0 +1,73 @@ +package modeladmin + +import ( + "context" + "fmt" + + "github.com/mudler/LocalAI/core/services/nodes" +) + +type revisionRegistry interface { + AdvanceModelConfigRevisions(ctx context.Context, transitions []nodes.ModelConfigRevisionTransition) ([]nodes.NodeModel, error) +} + +type revisionCleanup interface { + Cleanup(ctx context.Context, replicas []nodes.NodeModel, force bool) int +} + +type localModelShutdown interface { + ShutdownModel(modelName string) error +} + +type LocalModelRevisionLifecycle struct{ loader localModelShutdown } + +func NewLocalModelRevisionLifecycle(loader localModelShutdown) *LocalModelRevisionLifecycle { + if loader == nil { + return nil + } + return &LocalModelRevisionLifecycle{loader: loader} +} + +func (s *LocalModelRevisionLifecycle) ApplyConfigRevisions(_ context.Context, transitions []ModelRevisionTransition) (int, error) { + pending := 0 + seen := make(map[string]struct{}, len(transitions)) + for _, transition := range transitions { + if _, exists := seen[transition.ModelName]; exists { + continue + } + seen[transition.ModelName] = struct{}{} + if err := s.loader.ShutdownModel(transition.ModelName); err != nil { + pending++ + } + } + return pending, nil +} + +// DistributedModelRevisionLifecycle makes the registry transition authoritative +// before any worker network call. Failed exact stops remain durable unloading +// rows and are retried by ModelCleanupService.Run. +type DistributedModelRevisionLifecycle struct { + registry revisionRegistry + cleanup revisionCleanup +} + +func NewDistributedModelRevisionLifecycle(registry revisionRegistry, cleanup revisionCleanup) *DistributedModelRevisionLifecycle { + if registry == nil || cleanup == nil { + return nil + } + return &DistributedModelRevisionLifecycle{registry: registry, cleanup: cleanup} +} + +func (s *DistributedModelRevisionLifecycle) ApplyConfigRevisions(ctx context.Context, transitions []ModelRevisionTransition) (int, error) { + registryTransitions := make([]nodes.ModelConfigRevisionTransition, 0, len(transitions)) + for _, transition := range transitions { + registryTransitions = append(registryTransitions, nodes.ModelConfigRevisionTransition{ + ModelName: transition.ModelName, ConfigRevision: transition.ConfigRevision, + }) + } + quarantined, err := s.registry.AdvanceModelConfigRevisions(ctx, registryTransitions) + if err != nil { + return 0, fmt.Errorf("advance model config revisions: %w", err) + } + return s.cleanup.Cleanup(ctx, quarantined, false), nil +} diff --git a/core/services/modeladmin/lifecycle_test.go b/core/services/modeladmin/lifecycle_test.go new file mode 100644 index 000000000000..f5fa6e138f5c --- /dev/null +++ b/core/services/modeladmin/lifecycle_test.go @@ -0,0 +1,89 @@ +package modeladmin + +import ( + "context" + "errors" + + "github.com/mudler/LocalAI/core/services/nodes" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +type lifecycleRegistry struct { + models map[string][]nodes.NodeModel + calls []string + err error +} + +func (r *lifecycleRegistry) AdvanceModelConfigRevisions(_ context.Context, transitions []nodes.ModelConfigRevisionTransition) ([]nodes.NodeModel, error) { + for _, transition := range transitions { + r.calls = append(r.calls, transition.ModelName+":"+transition.ConfigRevision) + } + if r.err != nil { + return nil, r.err + } + var quarantined []nodes.NodeModel + for _, transition := range transitions { + quarantined = append(quarantined, r.models[transition.ModelName]...) + } + return quarantined, nil +} + +type lifecycleCleanup struct { + registry *lifecycleRegistry + seen []nodes.NodeModel + pending int +} + +func (c *lifecycleCleanup) Cleanup(_ context.Context, replicas []nodes.NodeModel, _ bool) int { + Expect(c.registry.calls).ToNot(BeEmpty(), "quarantine must precede worker cleanup") + c.seen = append(c.seen, replicas...) + return c.pending +} + +var _ = Describe("DistributedModelRevisionLifecycle", func() { + It("advances the registry before cleanup and reports incomplete exact stops", func() { + registry := &lifecycleRegistry{models: map[string][]nodes.NodeModel{ + "model": {{ID: "stale", ModelName: "model", State: "unloading"}}, + }} + cleanup := &lifecycleCleanup{registry: registry, pending: 1} + lifecycle := NewDistributedModelRevisionLifecycle(registry, cleanup) + + pending, err := lifecycle.ApplyConfigRevisions(context.Background(), []ModelRevisionTransition{{ModelName: "model", ConfigRevision: "rev-new"}}) + Expect(err).ToNot(HaveOccurred()) + Expect(pending).To(Equal(1)) + Expect(registry.calls).To(Equal([]string{"model:rev-new"})) + Expect(cleanup.seen).To(HaveLen(1)) + }) + + It("quarantines the old identity and establishes the renamed identity", func() { + registry := &lifecycleRegistry{models: map[string][]nodes.NodeModel{ + "old": {{ID: "old-replica", ModelName: "old"}}, + }} + cleanup := &lifecycleCleanup{registry: registry} + lifecycle := NewDistributedModelRevisionLifecycle(registry, cleanup) + + _, err := lifecycle.ApplyConfigRevisions(context.Background(), []ModelRevisionTransition{ + {ModelName: "old", ConfigRevision: DeletedModelConfigRevision("old"), Disabled: true}, + {ModelName: "new", ConfigRevision: "rev-renamed"}, + }) + Expect(err).ToNot(HaveOccurred()) + Expect(registry.calls).To(Equal([]string{"old:" + DeletedModelConfigRevision("old"), "new:rev-renamed"})) + Expect(cleanup.seen).To(ConsistOf(nodes.NodeModel{ID: "old-replica", ModelName: "old"})) + }) + + It("does not clean up a partially advanced rename when the atomic transition fails", func() { + registry := &lifecycleRegistry{err: errors.New("injected rename transition failure")} + cleanup := &lifecycleCleanup{registry: registry} + lifecycle := NewDistributedModelRevisionLifecycle(registry, cleanup) + + pending, err := lifecycle.ApplyConfigRevisions(context.Background(), []ModelRevisionTransition{ + {ModelName: "old", ConfigRevision: DeletedModelConfigRevision("old"), Disabled: true}, + {ModelName: "new", ConfigRevision: "rev-renamed"}, + }) + Expect(err).To(MatchError(ContainSubstring("injected rename transition failure"))) + Expect(pending).To(BeZero()) + Expect(registry.calls).To(Equal([]string{"old:" + DeletedModelConfigRevision("old"), "new:rev-renamed"})) + Expect(cleanup.seen).To(BeEmpty()) + }) +}) diff --git a/core/services/modeladmin/mutation.go b/core/services/modeladmin/mutation.go new file mode 100644 index 000000000000..9f458ae79e54 --- /dev/null +++ b/core/services/modeladmin/mutation.go @@ -0,0 +1,76 @@ +package modeladmin + +import ( + "errors" + "fmt" + "os" + "path/filepath" + "strings" + + "github.com/mudler/LocalAI/pkg/safefile" +) + +type savedMutationFile struct { + path string + data []byte + mode os.FileMode + exists bool +} + +func (s *ConfigService) withMutationRollback(paths []string, mutate func() error) error { + configs := s.Loader.GetAllModelsConfigs() + files := make([]savedMutationFile, 0, len(paths)) + seen := map[string]struct{}{} + for _, path := range paths { + if _, ok := seen[path]; ok { + continue + } + seen[path] = struct{}{} + name, err := directMutationEntry(s.modelsPath(), path) + if err != nil { + return fmt.Errorf("snapshot config mutation: %w", err) + } + file := savedMutationFile{path: path} + file.data, file.mode, err = safefile.ReadRegularAt(s.modelsPath(), name) + if err == nil { + file.exists = true + } + if err != nil && !errors.Is(err, os.ErrNotExist) { + return fmt.Errorf("snapshot config mutation: %w", err) + } + files = append(files, file) + } + + if err := mutate(); err != nil { + var restoreErr error + for _, file := range files { + if file.exists { + restoreErr = errors.Join(restoreErr, writeFileAtomic(file.path, file.data, file.mode)) + } else if removeErr := os.Remove(file.path); removeErr != nil && !errors.Is(removeErr, os.ErrNotExist) { + restoreErr = errors.Join(restoreErr, removeErr) + } + } + s.Loader.ReplaceModelConfigs(configs) + if restoreErr != nil { + return errors.Join(err, fmt.Errorf("restore prior model configuration: %w", restoreErr)) + } + return err + } + return nil +} + +func directMutationEntry(modelsPath, path string) (string, error) { + root, err := filepath.Abs(modelsPath) + if err != nil { + return "", err + } + candidate, err := filepath.Abs(path) + if err != nil { + return "", err + } + rel, err := filepath.Rel(root, candidate) + if err != nil || rel == "." || rel == ".." || filepath.IsAbs(rel) || strings.HasPrefix(rel, ".."+string(filepath.Separator)) || filepath.Dir(candidate) != root { + return "", fmt.Errorf("config path %q is not a direct entry of the configured models directory", path) + } + return filepath.Base(candidate), nil +} diff --git a/core/services/modeladmin/remote_sync.go b/core/services/modeladmin/remote_sync.go index 5acf5bf9ac7b..844239a8829d 100644 --- a/core/services/modeladmin/remote_sync.go +++ b/core/services/modeladmin/remote_sync.go @@ -1,53 +1,119 @@ package modeladmin import ( + "context" + "crypto/sha256" + "fmt" + "sort" + "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/services/messaging" - "github.com/mudler/LocalAI/pkg/model" - - "github.com/mudler/xlog" ) -// opDelete is the CacheInvalidateEvent.Op value the gallery delete path and the -// admin delete endpoint use; a delete must prune (a reload-from-path cannot). -const opDelete = "delete" - // ApplyRemoteChange refreshes this replica's in-memory model state from a peer // replica's model-config change broadcast (messaging.CacheInvalidateEvent on // SubjectCacheInvalidateModels). It is the subscriber-side counterpart to // GalleryService.BroadcastModelsChanged. // -// The op matters because LoadModelConfigsFromPath is additive: it loads every -// YAML on disk into the loader but never removes an entry whose file is gone. -// So a delete cannot be propagated by a plain reload - the deleted element must -// be explicitly pruned. Specifically: -// -// - op == "delete" with a named element: prune that element from the loader. -// - otherwise: reload all configs from disk (picks up creates and edits). +// The event is only a wake-up signal. Its operation and revision may be stale +// or reordered, so named changes are always reconciled against the current +// shared filesystem state. // -// In both cases, when an element is named, any running instance on this replica -// is shut down (best-effort) so the next request rebuilds it from the new -// config instead of serving the stale one - mirroring what the originating -// replica does on a local edit/delete. -// -// ml may be nil (no running instances to shut down). modelsPath and opts are -// forwarded to LoadModelConfigsFromPath. -func ApplyRemoteChange(cl *config.ModelConfigLoader, ml *model.ModelLoader, modelsPath string, evt messaging.CacheInvalidateEvent, opts ...config.ConfigLoaderOption) error { - if evt.Op == opDelete && evt.Element != "" { - cl.RemoveModelConfig(evt.Element) - } else if err := cl.LoadModelConfigsFromPath(modelsPath, opts...); err != nil { +// Revision-aware events apply the same idempotent lifecycle transition as the +// originating frontend. modelsPath and opts are forwarded to +// LoadModelConfigsFromPath. +func ApplyRemoteChange(ctx context.Context, cl *config.ModelConfigLoader, modelsPath string, evt messaging.CacheInvalidateEvent, lifecycle ModelRevisionLifecycle, opts ...config.ConfigLoaderOption) error { + return cl.WithModelConfigMutation(func() error { + return applyRemoteChange(ctx, cl, modelsPath, evt, lifecycle, opts...) + }) +} + +func applyRemoteChange(ctx context.Context, cl *config.ModelConfigLoader, modelsPath string, evt messaging.CacheInvalidateEvent, lifecycle ModelRevisionLifecycle, opts ...config.ConfigLoaderOption) error { + authoritative := config.NewModelConfigLoader(modelsPath) + if err := authoritative.LoadModelConfigsFromPathStrict(modelsPath, opts...); err != nil { + return err + } + current := configsByName(cl.GetAllModelsConfigs()) + snapshotConfigs := authoritative.GetAllModelsConfigs() + snapshot := configsByName(snapshotConfigs) + changed, err := changedConfigNames(current, snapshot, evt.Element) + if err != nil { return err } - // Drop any running instance of the affected model so the next request - // rebuilds it from the refreshed config instead of serving the stale one. - // Best-effort: the model may not be loaded on this replica, which surfaces - // as a benign error here. - if ml != nil && evt.Element != "" { - if err := ml.ShutdownModel(evt.Element); err != nil { - xlog.Debug("ApplyRemoteChange: could not shut down model instance (likely not loaded)", - "model", evt.Element, "error", err) + if lifecycle != nil { + transitions := make([]ModelRevisionTransition, 0, len(changed)) + for _, name := range changed { + cfg, exists := snapshot[name] + revision := DeletedModelConfigRevision(name) + disabled := true + if exists { + var err error + revision, err = config.ModelConfigRevision(&cfg) + if err != nil { + return fmt.Errorf("compute authoritative model config revision for %q: %w", name, err) + } + disabled = cfg.IsDisabled() + } + transitions = append(transitions, ModelRevisionTransition{ModelName: name, ConfigRevision: revision, Disabled: disabled}) + } + if len(transitions) > 0 { + if _, err := lifecycle.ApplyConfigRevisions(ctx, transitions); err != nil { + return err + } } } + cl.ReplaceModelConfigs(snapshotConfigs) return nil } + +func configsByName(configs []config.ModelConfig) map[string]config.ModelConfig { + result := make(map[string]config.ModelConfig, len(configs)) + for _, cfg := range configs { + result[cfg.Name] = cfg + } + return result +} + +func changedConfigNames(current, snapshot map[string]config.ModelConfig, named string) ([]string, error) { + changed := map[string]struct{}{} + for name, cfg := range snapshot { + previous, exists := current[name] + if !exists { + changed[name] = struct{}{} + continue + } + previousRevision, err := config.ModelConfigRevision(&previous) + if err != nil { + return nil, fmt.Errorf("compute current model config revision for %q: %w", name, err) + } + revision, err := config.ModelConfigRevision(&cfg) + if err != nil { + return nil, fmt.Errorf("compute authoritative model config revision for %q: %w", name, err) + } + if previousRevision != revision { + changed[name] = struct{}{} + } + } + for name := range current { + if _, exists := snapshot[name]; !exists { + changed[name] = struct{}{} + } + } + if named != "" { + changed[named] = struct{}{} + } + names := make([]string, 0, len(changed)) + for name := range changed { + names = append(names, name) + } + sort.Strings(names) + return names, nil +} + +// DeletedModelConfigRevision is a stable tombstone generation for an absent +// model. It lets every frontend derive the same authoritative state regardless +// of which reordered cache-invalidation event woke it up. +func DeletedModelConfigRevision(modelName string) string { + return fmt.Sprintf("%x", sha256.Sum256([]byte("deleted\x00"+modelName))) +} diff --git a/core/services/modeladmin/remote_sync_test.go b/core/services/modeladmin/remote_sync_test.go index a01cb7b56b88..32289429fa3a 100644 --- a/core/services/modeladmin/remote_sync_test.go +++ b/core/services/modeladmin/remote_sync_test.go @@ -1,8 +1,11 @@ package modeladmin import ( + "context" + "errors" "os" "path/filepath" + "sync" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" @@ -10,6 +13,7 @@ import ( "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/pkg/system" ) var _ = Describe("ApplyRemoteChange", func() { @@ -37,15 +41,89 @@ var _ = Describe("ApplyRemoteChange", func() { _, ok := loader.GetModelConfig("peer-alias") Expect(ok).To(BeFalse(), "precondition: not yet in memory") - err := ApplyRemoteChange(loader, nil, dir, messaging.CacheInvalidateEvent{ + err := ApplyRemoteChange(context.Background(), loader, dir, messaging.CacheInvalidateEvent{ Element: "peer-alias", Op: "install", - }) + }, nil) Expect(err).ToNot(HaveOccurred()) _, ok = loader.GetModelConfig("peer-alias") Expect(ok).To(BeTrue(), "install event must reload the new config from disk") }) + It("idempotently reconciles the authoritative revision instead of the event revision", func() { + writeYAML("peer-alias", map[string]any{"alias": "qwen"}) + lifecycle := &fakeRevisionLifecycle{} + evt := messaging.CacheInvalidateEvent{Element: "peer-alias", Op: "install", ConfigRevision: "stale-event-revision"} + + Expect(ApplyRemoteChange(context.Background(), loader, dir, evt, lifecycle)).To(Succeed()) + Expect(ApplyRemoteChange(context.Background(), loader, dir, evt, lifecycle)).To(Succeed()) + Expect(lifecycle.calls).To(HaveLen(2)) + loaded, ok := loader.GetModelConfig("peer-alias") + Expect(ok).To(BeTrue()) + revision, err := config.ModelConfigRevision(&loaded) + Expect(err).ToNot(HaveOccurred()) + Expect(lifecycle.calls[0].revision).To(Equal(revision)) + Expect(lifecycle.calls[1].revision).To(Equal(revision)) + }) + + It("uses the authoritative installed config for reordered install events", func() { + writeYAML("peer-alias", map[string]any{"backend": "llama-cpp", "context_size": 8192}) + lifecycle := &fakeRevisionLifecycle{} + Expect(ApplyRemoteChange(context.Background(), loader, dir, messaging.CacheInvalidateEvent{ + Element: "peer-alias", Op: "install", ConfigRevision: "old", + }, lifecycle)).To(Succeed()) + + writeYAML("peer-alias", map[string]any{"backend": "llama-cpp", "context_size": 10000}) + Expect(ApplyRemoteChange(context.Background(), loader, dir, messaging.CacheInvalidateEvent{ + Element: "peer-alias", Op: "install", ConfigRevision: "new-event-arrived-first", + }, lifecycle)).To(Succeed()) + Expect(ApplyRemoteChange(context.Background(), loader, dir, messaging.CacheInvalidateEvent{ + Element: "peer-alias", Op: "install", ConfigRevision: "old-event-arrived-late", + }, lifecycle)).To(Succeed()) + + loaded, ok := loader.GetModelConfig("peer-alias") + Expect(ok).To(BeTrue()) + Expect(loaded.ContextSize).To(HaveValue(Equal(10000))) + revision, err := config.ModelConfigRevision(&loaded) + Expect(err).ToNot(HaveOccurred()) + Expect(lifecycle.calls).To(HaveLen(3)) + Expect(lifecycle.calls[1].revision).To(Equal(revision)) + Expect(lifecycle.calls[2].revision).To(Equal(revision)) + }) + + It("does not let a delayed delete prune a reinstalled config", func() { + writeYAML("reinstalled", map[string]any{"backend": "llama-cpp", "context_size": 10000}) + lifecycle := &fakeRevisionLifecycle{} + Expect(ApplyRemoteChange(context.Background(), loader, dir, messaging.CacheInvalidateEvent{ + Element: "reinstalled", Op: "delete", ConfigRevision: "obsolete-delete", + }, lifecycle)).To(Succeed()) + + loaded, ok := loader.GetModelConfig("reinstalled") + Expect(ok).To(BeTrue()) + revision, err := config.ModelConfigRevision(&loaded) + Expect(err).ToNot(HaveOccurred()) + Expect(lifecycle.calls).To(HaveLen(1)) + Expect(lifecycle.calls[0].revision).To(Equal(revision)) + Expect(lifecycle.calls[0].disabled).To(BeFalse()) + }) + + It("does not let a delayed install resurrect an authoritative delete", func() { + writeYAML("deleted", map[string]any{"backend": "llama-cpp"}) + Expect(loader.LoadModelConfigsFromPath(dir)).To(Succeed()) + Expect(os.Remove(filepath.Join(dir, "deleted.yaml"))).To(Succeed()) + lifecycle := &fakeRevisionLifecycle{} + + Expect(ApplyRemoteChange(context.Background(), loader, dir, messaging.CacheInvalidateEvent{ + Element: "deleted", Op: "install", ConfigRevision: "obsolete-install", + }, lifecycle)).To(Succeed()) + + _, ok := loader.GetModelConfig("deleted") + Expect(ok).To(BeFalse()) + Expect(lifecycle.calls).To(HaveLen(1)) + Expect(lifecycle.calls[0].disabled).To(BeTrue()) + Expect(lifecycle.calls[0].revision).To(Equal(DeletedModelConfigRevision("deleted"))) + }) + It("prunes a peer-deleted config that a reload-from-path cannot drop", func() { // Model is present in memory (loaded earlier) but its file is now gone // from the shared dir. LoadModelConfigsFromPath is additive, so only an @@ -56,9 +134,9 @@ var _ = Describe("ApplyRemoteChange", func() { Expect(ok).To(BeTrue(), "precondition: in memory") Expect(os.Remove(filepath.Join(dir, "doomed.yaml"))).To(Succeed()) - err := ApplyRemoteChange(loader, nil, dir, messaging.CacheInvalidateEvent{ + err := ApplyRemoteChange(context.Background(), loader, dir, messaging.CacheInvalidateEvent{ Element: "doomed", Op: "delete", - }) + }, nil) Expect(err).ToNot(HaveOccurred()) _, ok = loader.GetModelConfig("doomed") @@ -69,7 +147,7 @@ var _ = Describe("ApplyRemoteChange", func() { writeYAML("m1", map[string]any{"alias": "qwen"}) writeYAML("m2", map[string]any{"alias": "qwen"}) - err := ApplyRemoteChange(loader, nil, dir, messaging.CacheInvalidateEvent{}) + err := ApplyRemoteChange(context.Background(), loader, dir, messaging.CacheInvalidateEvent{}, nil) Expect(err).ToNot(HaveOccurred()) _, ok1 := loader.GetModelConfig("m1") @@ -78,6 +156,119 @@ var _ = Describe("ApplyRemoteChange", func() { Expect(ok2).To(BeTrue()) }) + It("authoritatively reconciles changed and deleted configs when no element is named", func() { + writeYAML("changed", map[string]any{"backend": "llama-cpp", "context_size": 8192}) + writeYAML("deleted", map[string]any{"backend": "llama-cpp"}) + Expect(loader.LoadModelConfigsFromPath(dir)).To(Succeed()) + + writeYAML("changed", map[string]any{"backend": "llama-cpp", "context_size": 10000}) + Expect(os.Remove(filepath.Join(dir, "deleted.yaml"))).To(Succeed()) + lifecycle := &fakeRevisionLifecycle{} + + Expect(ApplyRemoteChange(context.Background(), loader, dir, messaging.CacheInvalidateEvent{}, lifecycle)).To(Succeed()) + + loaded, ok := loader.GetModelConfig("changed") + Expect(ok).To(BeTrue()) + Expect(loaded.ContextSize).To(HaveValue(Equal(10000))) + _, ok = loader.GetModelConfig("deleted") + Expect(ok).To(BeFalse()) + changedRevision, err := config.ModelConfigRevision(&loaded) + Expect(err).ToNot(HaveOccurred()) + Expect(lifecycle.calls).To(ConsistOf( + revisionLifecycleCall{oldName: "changed", newName: "changed", revision: changedRevision}, + revisionLifecycleCall{oldName: "deleted", newName: "deleted", revision: DeletedModelConfigRevision("deleted"), disabled: true}, + )) + Expect(lifecycle.batches).To(HaveLen(1)) + Expect(lifecycle.batches[0]).To(HaveLen(2)) + }) + + It("keeps the complete live snapshot unchanged when a batched transition fails", func() { + writeYAML("changed", map[string]any{"backend": "llama-cpp", "context_size": 8192}) + writeYAML("deleted", map[string]any{"backend": "llama-cpp"}) + Expect(loader.LoadModelConfigsFromPath(dir)).To(Succeed()) + + writeYAML("changed", map[string]any{"backend": "llama-cpp", "context_size": 10000}) + Expect(os.Remove(filepath.Join(dir, "deleted.yaml"))).To(Succeed()) + lifecycle := &fakeRevisionLifecycle{err: errors.New("injected second transition failure")} + + Expect(ApplyRemoteChange(context.Background(), loader, dir, messaging.CacheInvalidateEvent{}, lifecycle)).To( + MatchError(ContainSubstring("injected second transition failure")), + ) + Expect(lifecycle.batches).To(HaveLen(1)) + Expect(lifecycle.batches[0]).To(HaveLen(2)) + loaded, ok := loader.GetModelConfig("changed") + Expect(ok).To(BeTrue()) + Expect(loaded.ContextSize).To(HaveValue(Equal(8192))) + _, ok = loader.GetModelConfig("deleted") + Expect(ok).To(BeTrue()) + }) + + It("serializes authoritative reads through lifecycle publication", func() { + writeYAML("ordered", map[string]any{"backend": "llama-cpp", "context_size": 8192}) + lifecycle := newBlockingRevisionLifecycle() + firstDone := make(chan error, 1) + secondDone := make(chan error, 1) + + go func() { + firstDone <- ApplyRemoteChange(context.Background(), loader, dir, messaging.CacheInvalidateEvent{Element: "ordered"}, lifecycle) + }() + Eventually(lifecycle.entered).Should(Receive()) + + writeYAML("ordered", map[string]any{"backend": "llama-cpp", "context_size": 10000}) + go func() { + secondDone <- ApplyRemoteChange(context.Background(), loader, dir, messaging.CacheInvalidateEvent{}, lifecycle) + }() + Consistently(lifecycle.entered).ShouldNot(Receive()) + + close(lifecycle.release) + Eventually(firstDone).Should(Receive(Succeed())) + Eventually(secondDone).Should(Receive(Succeed())) + Eventually(lifecycle.entered).Should(Receive()) + + loaded, ok := loader.GetModelConfig("ordered") + Expect(ok).To(BeTrue()) + Expect(loaded.ContextSize).To(HaveValue(Equal(10000))) + revision, err := config.ModelConfigRevision(&loaded) + Expect(err).ToNot(HaveOccurred()) + Expect(lifecycle.revisions()).To(HaveLen(2)) + Expect(lifecycle.revisions()[1]).To(Equal(revision)) + }) + + It("serializes peer publication before a newer local edit", func() { + writeYAML("ordered", map[string]any{"backend": "llama-cpp", "context_size": 8192}) + lifecycle := newBlockingRevisionLifecycle() + appConfig := &config.ApplicationConfig{SystemState: &system.SystemState{Model: system.Model{ModelsPath: dir}}} + svc := NewConfigService(loader, appConfig, lifecycle) + peerDone := make(chan error, 1) + localDone := make(chan error, 1) + + go func() { + peerDone <- ApplyRemoteChange(context.Background(), loader, dir, messaging.CacheInvalidateEvent{Element: "ordered"}, lifecycle) + }() + Eventually(lifecycle.entered).Should(Receive()) + + go func() { + _, err := svc.EditYAML(context.Background(), "ordered", []byte("name: ordered\nbackend: llama-cpp\ncontext_size: 10000\n")) + localDone <- err + }() + Consistently(localDone).ShouldNot(Receive()) + Expect(readMap(filepath.Join(dir, "ordered.yaml"))).To(HaveKeyWithValue("context_size", 8192)) + + close(lifecycle.release) + Eventually(peerDone).Should(Receive(Succeed())) + Eventually(localDone).Should(Receive(Succeed())) + Eventually(lifecycle.entered).Should(Receive()) + + loaded, ok := loader.GetModelConfig("ordered") + Expect(ok).To(BeTrue()) + Expect(loaded.ContextSize).To(HaveValue(Equal(10000))) + Expect(readMap(filepath.Join(dir, "ordered.yaml"))).To(HaveKeyWithValue("context_size", 10000)) + revision, err := config.ModelConfigRevision(&loaded) + Expect(err).ToNot(HaveOccurred()) + Expect(lifecycle.revisions()).To(HaveLen(2)) + Expect(lifecycle.revisions()[1]).To(Equal(revision)) + }) + It("loads a peer-persisted artifact binding without materializing", func() { const relative = ".artifacts/huggingface/0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef/snapshot" writeYAML("peer-managed", map[string]any{ @@ -93,9 +284,9 @@ var _ = Describe("ApplyRemoteChange", func() { }}, "parameters": map[string]any{"model": "owner/repo"}, }) - Expect(ApplyRemoteChange(loader, nil, dir, messaging.CacheInvalidateEvent{ + Expect(ApplyRemoteChange(context.Background(), loader, dir, messaging.CacheInvalidateEvent{ Element: "peer-managed", Op: "install", - })).To(Succeed()) + }, nil)).To(Succeed()) loaded, found := loader.GetModelConfig("peer-managed") Expect(found).To(BeTrue()) Expect(loaded.Model).To(Equal("owner/repo")) @@ -104,3 +295,32 @@ var _ = Describe("ApplyRemoteChange", func() { Expect(loaded.Artifacts[0].Resolved.CacheKey).To(HaveLen(64)) }) }) + +type blockingRevisionLifecycle struct { + mu sync.Mutex + calls []string + entered chan struct{} + release chan struct{} + once sync.Once +} + +func newBlockingRevisionLifecycle() *blockingRevisionLifecycle { + return &blockingRevisionLifecycle{entered: make(chan struct{}, 2), release: make(chan struct{})} +} + +func (l *blockingRevisionLifecycle) ApplyConfigRevisions(_ context.Context, transitions []ModelRevisionTransition) (int, error) { + l.mu.Lock() + for _, transition := range transitions { + l.calls = append(l.calls, transition.ConfigRevision) + } + l.mu.Unlock() + l.entered <- struct{}{} + l.once.Do(func() { <-l.release }) + return 0, nil +} + +func (l *blockingRevisionLifecycle) revisions() []string { + l.mu.Lock() + defer l.mu.Unlock() + return append([]string(nil), l.calls...) +} diff --git a/core/services/modeladmin/state.go b/core/services/modeladmin/state.go index ed23f1c8f795..5d87f6675409 100644 --- a/core/services/modeladmin/state.go +++ b/core/services/modeladmin/state.go @@ -7,23 +7,35 @@ import ( "gopkg.in/yaml.v3" - "github.com/mudler/LocalAI/pkg/model" + "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/pkg/utils" ) // ToggleResult is shared by ToggleState and TogglePinned. type ToggleResult struct { - Filename string - Action Action + Filename string + Action Action + ConfigRevision string + PendingCleanup int } // ToggleState enables or disables an installed model. action must be -// ActionEnable or ActionDisable. When ml is non-nil and the action is -// ActionDisable, ToggleState calls ml.ShutdownModel — best-effort. +// ActionEnable or ActionDisable. The revision lifecycle quarantines existing +// replicas before cleanup when the state changes. // // The on-disk YAML is mutated as a generic map so unrelated fields are // preserved verbatim; we only set or remove the `disabled` key. -func (s *ConfigService) ToggleState(_ context.Context, name string, action Action, ml *model.ModelLoader) (*ToggleResult, error) { +func (s *ConfigService) ToggleState(ctx context.Context, name string, action Action) (*ToggleResult, error) { + var result *ToggleResult + err := s.Loader.WithModelConfigMutation(func() error { + var err error + result, err = s.toggleState(ctx, name, action) + return err + }) + return result, err +} + +func (s *ConfigService) toggleState(ctx context.Context, name string, action Action) (*ToggleResult, error) { if name == "" { return nil, ErrNameRequired } @@ -41,17 +53,30 @@ func (s *ConfigService) ToggleState(_ context.Context, name string, action Actio if err := utils.VerifyPath(configPath, s.modelsPath()); err != nil { return nil, fmt.Errorf("%w: %v", ErrPathNotTrusted, err) } - if err := mutateYAMLBoolFlag(configPath, "disabled", action == ActionDisable); err != nil { - return nil, err - } - if err := s.Loader.LoadModelConfigsFromPath(s.modelsPath(), s.AppConfig.ToConfigLoaderOptions()...); err != nil { - return nil, fmt.Errorf("reload configs: %w", err) - } - if action == ActionDisable && ml != nil { - // Best-effort: the YAML is saved; shutdown is a courtesy. - _ = ml.ShutdownModel(name) - } - return &ToggleResult{Filename: configPath, Action: action}, nil + var result *ToggleResult + err := s.withMutationRollback([]string{configPath}, func() error { + if err := mutateYAMLBoolFlag(configPath, "disabled", action == ActionDisable); err != nil { + return err + } + if err := s.Loader.LoadModelConfigsFromPath(s.modelsPath(), s.AppConfig.ToConfigLoaderOptions()...); err != nil { + return fmt.Errorf("reload configs: %w", err) + } + loaded, ok := s.Loader.GetModelConfig(name) + if !ok { + return fmt.Errorf("reload configs: model %q missing", name) + } + revision, err := config.ModelConfigRevision(&loaded) + if err != nil { + return fmt.Errorf("compute config revision: %w", err) + } + pending, err := s.applyRevision(ctx, name, name, revision, action == ActionDisable) + if err != nil { + return err + } + result = &ToggleResult{Filename: configPath, Action: action, ConfigRevision: revision, PendingCleanup: pending} + return nil + }) + return result, err } // mutateYAMLBoolFlag is a small helper shared by ToggleState and diff --git a/core/services/modeladmin/state_test.go b/core/services/modeladmin/state_test.go index 954e0723690d..a424826ce28d 100644 --- a/core/services/modeladmin/state_test.go +++ b/core/services/modeladmin/state_test.go @@ -2,6 +2,7 @@ package modeladmin import ( "context" + "errors" "os" "path/filepath" @@ -35,17 +36,43 @@ var _ = Describe("ConfigService.ToggleState", func() { It("disables a model by writing disabled: true", func() { writeModelYAML(svc, dir, "qwen", map[string]any{"backend": "llama-cpp"}) - _, err := svc.ToggleState(ctx, "qwen", ActionDisable, nil) + _, err := svc.ToggleState(ctx, "qwen", ActionDisable) Expect(err).ToNot(HaveOccurred()) got := readMap(filepath.Join(dir, "qwen.yaml")) Expect(got).To(HaveKeyWithValue("disabled", true)) }) + It("applies disable through the revision lifecycle", func() { + lifecycle := &fakeRevisionLifecycle{pending: 3} + svc.Lifecycle = lifecycle + writeModelYAML(svc, dir, "qwen", map[string]any{"backend": "llama-cpp"}) + + result, err := svc.ToggleState(ctx, "qwen", ActionDisable) + Expect(err).ToNot(HaveOccurred()) + Expect(result.ConfigRevision).ToNot(BeEmpty()) + Expect(result.PendingCleanup).To(Equal(3)) + Expect(lifecycle.calls).To(ConsistOf(revisionLifecycleCall{ + oldName: "qwen", newName: "qwen", revision: result.ConfigRevision, disabled: true, + })) + }) + + It("restores disk and loader when state publication fails", func() { + svc.Lifecycle = &fakeRevisionLifecycle{err: errors.New("registry unavailable")} + writeModelYAML(svc, dir, "qwen", map[string]any{"backend": "llama-cpp"}) + + _, err := svc.ToggleState(ctx, "qwen", ActionDisable) + Expect(err).To(MatchError(ContainSubstring("registry unavailable"))) + Expect(readMap(filepath.Join(dir, "qwen.yaml"))).NotTo(HaveKey("disabled")) + loaded, ok := svc.Loader.GetModelConfig("qwen") + Expect(ok).To(BeTrue()) + Expect(loaded.IsDisabled()).To(BeFalse()) + }) + It("enables a model by removing the disabled key entirely", func() { writeModelYAML(svc, dir, "qwen", map[string]any{"backend": "llama-cpp", "disabled": true}) - _, err := svc.ToggleState(ctx, "qwen", ActionEnable, nil) + _, err := svc.ToggleState(ctx, "qwen", ActionEnable) Expect(err).ToNot(HaveOccurred()) got := readMap(filepath.Join(dir, "qwen.yaml")) @@ -54,12 +81,12 @@ var _ = Describe("ConfigService.ToggleState", func() { It("rejects unknown actions with ErrBadAction", func() { writeModelYAML(svc, dir, "qwen", map[string]any{"backend": "llama-cpp"}) - _, err := svc.ToggleState(ctx, "qwen", Action("noop"), nil) + _, err := svc.ToggleState(ctx, "qwen", Action("noop")) Expect(err).To(MatchError(ErrBadAction)) }) It("returns ErrNotFound for an unknown model", func() { - _, err := svc.ToggleState(ctx, "ghost", ActionDisable, nil) + _, err := svc.ToggleState(ctx, "ghost", ActionDisable) Expect(err).To(MatchError(ErrNotFound)) }) }) diff --git a/core/services/nodes/disk_headroom_test.go b/core/services/nodes/disk_headroom_test.go index 217fb71fd44d..f14abd8292d2 100644 --- a/core/services/nodes/disk_headroom_test.go +++ b/core/services/nodes/disk_headroom_test.go @@ -155,8 +155,9 @@ var _ = Describe("scheduling a model onto a cluster without disk headroom", func }) route := func(modelFile string) error { - _, err := router.Route(context.Background(), "longcat-video-avatar-1.5", "models/big.gguf", "llama-cpp", + _, err := router.Route(context.Background(), "longcat-video-avatar-1.5", "models/big.gguf", "llama-cpp", "", &pb.ModelOptions{Model: "models/big.gguf", ModelFile: modelFile}, false) + return err } diff --git a/core/services/nodes/interfaces.go b/core/services/nodes/interfaces.go index 1be5f20e84a9..c204752de19b 100644 --- a/core/services/nodes/interfaces.go +++ b/core/services/nodes/interfaces.go @@ -4,9 +4,20 @@ import ( "context" "time" + "github.com/mudler/LocalAI/core/services/messaging" grpc "github.com/mudler/LocalAI/pkg/grpc" ) +type ExactModelStopper interface { + StopModelReplica(ctx context.Context, nodeID string, replica NodeModel, force bool) (messaging.ModelStopReply, error) +} + +type ModelCleanupRegistry interface { + ClaimModelCleanupRetries(ctx context.Context, now, leaseUntil time.Time, limit int) ([]NodeModel, error) + RecordModelCleanupFailure(ctx context.Context, nodeID, modelName string, replicaIndex int, cleanupErr string, nextRetry time.Time) error + RemoveClaimedModelCleanup(ctx context.Context, replica NodeModel) (bool, error) +} + // ModelRouter is used by SmartRouter for routing decisions and model lifecycle. type ModelRouter interface { FindAndLockNodeWithModel(ctx context.Context, modelName string, candidateNodeIDs []string, pref *RoutePreference) (*BackendNode, *NodeModel, error) @@ -16,9 +27,19 @@ type ModelRouter interface { RemoveAllNodeModelReplicas(ctx context.Context, nodeID, modelName string) error TouchNodeModel(ctx context.Context, nodeID, modelName string, replicaIndex int) SetNodeModel(ctx context.Context, nodeID, modelName string, replicaIndex int, state, address string, initialInFlight int) error + SetNodeModelRevision(ctx context.Context, nodeID, modelName string, replicaIndex int, state, address string, initialInFlight int, revision, effectiveOptionsHash string) error SetNodeModelLoadInfo(ctx context.Context, nodeID, modelName string, replicaIndex int, backendType string, optsBlob []byte) error + SetNodeModelLoadInfoRevision(ctx context.Context, nodeID, modelName string, replicaIndex int, backendType, revision string, optsBlob []byte) error UpsertModelLoadInfo(ctx context.Context, modelName, backendType string, optsBlob []byte) error + UpsertModelLoadInfoRevision(ctx context.Context, modelName, backendType, revision string, optsBlob []byte) error GetModelLoadInfo(ctx context.Context, modelName string) (backendType string, optsBlob []byte, err error) + GetModelLoadInfoRevision(ctx context.Context, modelName string) (backendType, revision string, optsBlob []byte, err error) + AdvanceModelConfigRevision(ctx context.Context, modelName, revision string) ([]NodeModel, error) + EstablishModelConfigRevision(ctx context.Context, modelName, revision string) error + GetModelConfigRevision(ctx context.Context, modelName string) (string, error) + GetNodeModel(ctx context.Context, nodeID, modelName string, replicaIndex int) (*NodeModel, error) + RecordModelCleanupFailure(ctx context.Context, nodeID, modelName string, replicaIndex int, cleanupErr string, nextRetry time.Time) error + ListModelCleanupRetries(ctx context.Context, now time.Time, limit int) ([]NodeModel, error) NextFreeReplicaIndex(ctx context.Context, nodeID, modelName string, maxSlots int) (int, error) CountReplicasOnNode(ctx context.Context, nodeID, modelName string) (int, error) FindNodeWithVRAM(ctx context.Context, minBytes uint64) (*BackendNode, error) diff --git a/core/services/nodes/local_stub_invalidator_test.go b/core/services/nodes/local_stub_invalidator_test.go index a8f591417887..00ed820dc6c2 100644 --- a/core/services/nodes/local_stub_invalidator_test.go +++ b/core/services/nodes/local_stub_invalidator_test.go @@ -9,8 +9,8 @@ import ( . "github.com/onsi/gomega" "gorm.io/gorm" - "github.com/mudler/LocalAI/pkg/model" "github.com/mudler/LocalAI/core/services/testutil" + "github.com/mudler/LocalAI/pkg/model" ) // In distributed mode the frontend keeps an in-process stub for every model it diff --git a/core/services/nodes/model_cleanup.go b/core/services/nodes/model_cleanup.go new file mode 100644 index 000000000000..cff11226d69d --- /dev/null +++ b/core/services/nodes/model_cleanup.go @@ -0,0 +1,113 @@ +package nodes + +import ( + "context" + "errors" + "strings" + "time" + + "github.com/mudler/xlog" +) + +const ( + modelCleanupInterval = time.Second + // Exact stops are bounded to ten seconds. Claiming one row for two minutes + // keeps ownership durable even under scheduler stalls and avoids a batch's + // later rows losing their lease while earlier stops run. + modelCleanupLease = 2 * time.Minute + modelCleanupBatch = 1 + modelCleanupMaxDelay = 5 * time.Minute +) + +type ModelCleanupService struct { + registry ModelCleanupRegistry + stopper ExactModelStopper + now func() time.Time +} + +func NewModelCleanupService(registry ModelCleanupRegistry, stopper ExactModelStopper) *ModelCleanupService { + return &ModelCleanupService{registry: registry, stopper: stopper, now: time.Now} +} + +func (s *ModelCleanupService) Cleanup(ctx context.Context, replicas []NodeModel, force bool) int { + pending := 0 + for _, replica := range replicas { + reply, err := s.stopper.StopModelReplica(ctx, replica.NodeID, replica, force) + if err == nil && reply.Terminated { + removed, removeErr := s.registry.RemoveClaimedModelCleanup(ctx, replica) + if removeErr != nil { + xlog.Warn("Removing terminated model replica failed", "nodeID", replica.NodeID, "model", replica.ModelName, "replica", replica.ReplicaIndex, "error", removeErr) + pending++ + } else if !removed { + pending++ + } + continue + } + pending++ + + cleanupErr := conciseCleanupError(err, reply.Error) + nextRetry := s.now().Add(modelCleanupBackoff(replica.CleanupAttempts)) + if recordErr := s.registry.RecordModelCleanupFailure(ctx, replica.NodeID, replica.ModelName, replica.ReplicaIndex, cleanupErr, nextRetry); recordErr != nil { + xlog.Warn("Recording model cleanup retry failed", "nodeID", replica.NodeID, "model", replica.ModelName, "replica", replica.ReplicaIndex, "error", recordErr) + } + } + return pending +} + +func (s *ModelCleanupService) Run(ctx context.Context) { + ticker := time.NewTicker(modelCleanupInterval) + defer ticker.Stop() + for { + s.runOnce(ctx) + select { + case <-ctx.Done(): + return + case <-ticker.C: + } + } +} + +func (s *ModelCleanupService) runOnce(ctx context.Context) { + now := s.now() + replicas, err := s.registry.ClaimModelCleanupRetries(ctx, now, now.Add(modelCleanupLease), modelCleanupBatch) + if err != nil { + if !errors.Is(err, context.Canceled) { + xlog.Warn("Claiming model cleanup retries failed", "error", err) + } + return + } + s.Cleanup(ctx, replicas, false) +} + +func modelCleanupBackoff(attempts int) time.Duration { + if attempts < 0 { + attempts = 0 + } + if attempts > 8 { + attempts = 8 + } + delay := time.Second * time.Duration(1< modelCleanupMaxDelay { + return modelCleanupMaxDelay + } + return delay +} + +func conciseCleanupError(err error, replyError string) string { + message := strings.TrimSpace(replyError) + if err != nil { + message = err.Error() + } + if i := strings.LastIndex(message, ": "); i >= 0 { + message = message[i+2:] + } + message = strings.TrimSpace(message) + if message == "" { + return "termination not confirmed" + } + const max = 240 + if len(message) > max { + return message[:max] + } + return message +} diff --git a/core/services/nodes/model_cleanup_test.go b/core/services/nodes/model_cleanup_test.go new file mode 100644 index 000000000000..c09bbef4d3d0 --- /dev/null +++ b/core/services/nodes/model_cleanup_test.go @@ -0,0 +1,203 @@ +package nodes + +import ( + "context" + "errors" + "sync" + "time" + + "github.com/mudler/LocalAI/core/services/messaging" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +type fakeCleanupRegistry struct { + mu sync.Mutex + due []NodeModel + claimed bool + removed []modelReplicaRef + failures []string + next []time.Time +} + +func (f *fakeCleanupRegistry) ClaimModelCleanupRetries(_ context.Context, _ time.Time, _ time.Time, _ int) ([]NodeModel, error) { + f.mu.Lock() + defer f.mu.Unlock() + if f.claimed { + return nil, nil + } + f.claimed = true + return append([]NodeModel(nil), f.due...), nil +} + +func (f *fakeCleanupRegistry) RemoveClaimedModelCleanup(_ context.Context, claimed NodeModel) (bool, error) { + f.mu.Lock() + defer f.mu.Unlock() + f.removed = append(f.removed, modelReplicaRef{claimed.NodeID, claimed.ModelName, claimed.ReplicaIndex}) + return true, nil +} + +func (f *fakeCleanupRegistry) RecordModelCleanupFailure(_ context.Context, _, _ string, _ int, cleanupErr string, next time.Time) error { + f.mu.Lock() + defer f.mu.Unlock() + f.failures = append(f.failures, cleanupErr) + f.next = append(f.next, next) + return nil +} + +type fakeExactStopper struct { + mu sync.Mutex + replies []messaging.ModelStopReply + errs []error + calls []NodeModel + block chan struct{} +} + +type leasingCleanupRegistry struct { + mu sync.Mutex + row NodeModel + leaseUntil time.Time +} + +func (f *leasingCleanupRegistry) ClaimModelCleanupRetries(_ context.Context, now, leaseUntil time.Time, _ int) ([]NodeModel, error) { + f.mu.Lock() + defer f.mu.Unlock() + if !f.leaseUntil.IsZero() && f.leaseUntil.After(now) { + return nil, nil + } + f.leaseUntil = leaseUntil + return []NodeModel{f.row}, nil +} + +func (f *leasingCleanupRegistry) RemoveClaimedModelCleanup(_ context.Context, _ NodeModel) (bool, error) { + return true, nil +} + +func (f *leasingCleanupRegistry) RecordModelCleanupFailure(_ context.Context, _, _ string, _ int, _ string, _ time.Time) error { + return nil +} + +type blockingExactStopper struct { + mu sync.Mutex + entered chan struct{} + release chan struct{} + calls int +} + +func (f *blockingExactStopper) StopModelReplica(_ context.Context, _ string, _ NodeModel, _ bool) (messaging.ModelStopReply, error) { + f.mu.Lock() + f.calls++ + if f.calls == 1 { + close(f.entered) + } + f.mu.Unlock() + <-f.release + return messaging.ModelStopReply{Matched: true, Terminated: true}, nil +} + +func (f *fakeExactStopper) StopModelReplica(_ context.Context, _ string, replica NodeModel, _ bool) (messaging.ModelStopReply, error) { + if f.block != nil { + <-f.block + } + f.mu.Lock() + defer f.mu.Unlock() + i := len(f.calls) + f.calls = append(f.calls, replica) + var reply messaging.ModelStopReply + var err error + if i < len(f.replies) { + reply = f.replies[i] + } + if i < len(f.errs) { + err = f.errs[i] + } + return reply, err +} + +var _ = Describe("ModelCleanupService", func() { + var now time.Time + BeforeEach(func() { now = time.Date(2026, 8, 21, 12, 0, 0, 0, time.UTC) }) + + It("deletes only replicas whose termination is confirmed", func() { + registry := &fakeCleanupRegistry{} + stopper := &fakeExactStopper{replies: []messaging.ModelStopReply{{Matched: true, Terminated: true}}} + service := NewModelCleanupService(registry, stopper) + service.now = func() time.Time { return now } + service.Cleanup(context.Background(), []NodeModel{{NodeID: "n1", ModelName: "m", ReplicaIndex: 3}}, false) + Expect(registry.removed).To(Equal([]modelReplicaRef{{"n1", "m", 3}})) + Expect(registry.failures).To(BeEmpty()) + }) + + It("treats exact process absence as idempotent success", func() { + registry := &fakeCleanupRegistry{} + stopper := &fakeExactStopper{replies: []messaging.ModelStopReply{{Matched: false, Terminated: true}}} + service := NewModelCleanupService(registry, stopper) + service.Cleanup(context.Background(), []NodeModel{{NodeID: "n1", ModelName: "m"}}, false) + Expect(registry.removed).To(HaveLen(1)) + }) + + It("keeps and backs off a replica when no worker responds", func() { + registry := &fakeCleanupRegistry{} + stopper := &fakeExactStopper{errs: []error{errors.New("NATS request: no responders available")}} + service := NewModelCleanupService(registry, stopper) + service.now = func() time.Time { return now } + service.Cleanup(context.Background(), []NodeModel{{NodeID: "n1", ModelName: "m", CleanupAttempts: 2}}, false) + Expect(registry.removed).To(BeEmpty()) + Expect(registry.failures).To(Equal([]string{"no responders available"})) + Expect(registry.next[0]).To(Equal(now.Add(4 * time.Second))) + }) + + It("retries transient failures and later removes the row", func() { + registry := &fakeCleanupRegistry{} + stopper := &fakeExactStopper{errs: []error{errors.New("timeout"), nil}, replies: []messaging.ModelStopReply{{}, {Matched: true, Terminated: true}}} + service := NewModelCleanupService(registry, stopper) + r := NodeModel{NodeID: "n1", ModelName: "m"} + service.Cleanup(context.Background(), []NodeModel{r}, false) + service.Cleanup(context.Background(), []NodeModel{r}, false) + Expect(registry.failures).To(HaveLen(1)) + Expect(registry.removed).To(HaveLen(1)) + }) + + It("records a negative reply and tolerates a concurrent row deletion", func() { + registry := &fakeCleanupRegistry{} + stopper := &fakeExactStopper{replies: []messaging.ModelStopReply{{Matched: true, Terminated: false, Error: "address mismatch"}}} + service := NewModelCleanupService(registry, stopper) + service.Cleanup(context.Background(), []NodeModel{{NodeID: "n1", ModelName: "m"}}, false) + Expect(registry.failures).To(Equal([]string{"address mismatch"})) + }) + + It("leases due work so two runners do not own the same replica", func() { + registry := &fakeCleanupRegistry{due: []NodeModel{{NodeID: "n1", ModelName: "m"}}} + stopper := &fakeExactStopper{replies: []messaging.ModelStopReply{{Matched: true, Terminated: true}}} + a := NewModelCleanupService(registry, stopper) + b := NewModelCleanupService(registry, stopper) + a.runOnce(context.Background()) + b.runOnce(context.Background()) + Expect(stopper.calls).To(HaveLen(1)) + }) + + It("keeps single ownership while a slow stop advances past the old lease boundary", func() { + clock := now + registry := &leasingCleanupRegistry{row: NodeModel{ID: "claimed-row", NodeID: "n1", ModelName: "m", State: "unloading"}} + stopper := &blockingExactStopper{entered: make(chan struct{}), release: make(chan struct{})} + a := NewModelCleanupService(registry, stopper) + b := NewModelCleanupService(registry, stopper) + a.now = func() time.Time { return clock } + b.now = func() time.Time { return clock } + + done := make(chan struct{}) + go func() { + defer close(done) + a.runOnce(context.Background()) + }() + Eventually(stopper.entered).Should(BeClosed()) + clock = clock.Add(31 * time.Second) + b.runOnce(context.Background()) + + stopper.mu.Lock() + Expect(stopper.calls).To(Equal(1)) + stopper.mu.Unlock() + close(stopper.release) + Eventually(done).Should(BeClosed()) + }) +}) diff --git a/core/services/nodes/model_router.go b/core/services/nodes/model_router.go index 2f87bc5e6053..2d29fe528a92 100644 --- a/core/services/nodes/model_router.go +++ b/core/services/nodes/model_router.go @@ -38,7 +38,7 @@ func NewModelRouterAdapter(router *SmartRouter) *ModelRouterAdapter { // It delegates to SmartRouter.Route() and returns a Model that wraps the // remote gRPC client with file staging if configured. func (a *ModelRouterAdapter) Route(ctx context.Context, backend, modelID, modelName, modelFile string, - opts *pb.ModelOptions, parallel bool) (*model.Model, error) { + configRevision string, opts *pb.ModelOptions, parallel bool) (*model.Model, error) { backendType := backend @@ -51,7 +51,7 @@ func (a *ModelRouterAdapter) Route(ctx context.Context, backend, modelID, modelN // Route to a remote node (SmartRouter handles model pre-staging via FileStager) // Pass modelID so the DB tracks models by their logical ID, not the file path - result, err := a.router.Route(ctx, modelID, modelName, backendType, opts, parallel) + result, err := a.router.Route(ctx, modelID, modelName, backendType, configRevision, opts, parallel) if err != nil { return nil, fmt.Errorf("routing model %s: %w", modelName, err) } diff --git a/core/services/nodes/model_router_test.go b/core/services/nodes/model_router_test.go index b31e9dc064bb..43002006a1b0 100644 --- a/core/services/nodes/model_router_test.go +++ b/core/services/nodes/model_router_test.go @@ -4,11 +4,13 @@ import ( "context" "fmt" "sync" + "time" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" pb "github.com/mudler/LocalAI/pkg/grpc/proto" + "gorm.io/gorm" ) // --- fakeModelRouterForSmartRouter implements ModelRouter --- @@ -54,15 +56,46 @@ func (f *fakeModelRouterForSmartRouter) TouchNodeModel(_ context.Context, _, _ s func (f *fakeModelRouterForSmartRouter) SetNodeModel(_ context.Context, _, _ string, _ int, _, _ string, _ int) error { return nil } +func (f *fakeModelRouterForSmartRouter) SetNodeModelRevision(ctx context.Context, nodeID, modelName string, replicaIndex int, state, address string, initialInFlight int, _, _ string) error { + return f.SetNodeModel(ctx, nodeID, modelName, replicaIndex, state, address, initialInFlight) +} func (f *fakeModelRouterForSmartRouter) SetNodeModelLoadInfo(_ context.Context, _, _ string, _ int, _ string, _ []byte) error { return nil } +func (f *fakeModelRouterForSmartRouter) SetNodeModelLoadInfoRevision(ctx context.Context, nodeID, modelName string, replicaIndex int, backendType, _ string, optsBlob []byte) error { + return f.SetNodeModelLoadInfo(ctx, nodeID, modelName, replicaIndex, backendType, optsBlob) +} func (f *fakeModelRouterForSmartRouter) UpsertModelLoadInfo(_ context.Context, _, _ string, _ []byte) error { return nil } +func (f *fakeModelRouterForSmartRouter) UpsertModelLoadInfoRevision(ctx context.Context, modelName, backendType, _ string, optsBlob []byte) error { + return f.UpsertModelLoadInfo(ctx, modelName, backendType, optsBlob) +} func (f *fakeModelRouterForSmartRouter) GetModelLoadInfo(_ context.Context, _ string) (string, []byte, error) { return "", nil, fmt.Errorf("not found") } +func (f *fakeModelRouterForSmartRouter) GetModelLoadInfoRevision(ctx context.Context, modelName string) (string, string, []byte, error) { + backend, blob, err := f.GetModelLoadInfo(ctx, modelName) + return backend, "", blob, err +} +func (f *fakeModelRouterForSmartRouter) AdvanceModelConfigRevision(_ context.Context, _, _ string) ([]NodeModel, error) { + return nil, nil +} +func (f *fakeModelRouterForSmartRouter) EstablishModelConfigRevision(_ context.Context, _, _ string) error { + return nil +} +func (f *fakeModelRouterForSmartRouter) GetModelConfigRevision(_ context.Context, _ string) (string, error) { + return "", gorm.ErrRecordNotFound +} +func (f *fakeModelRouterForSmartRouter) GetNodeModel(_ context.Context, nodeID, modelName string, replicaIndex int) (*NodeModel, error) { + return &NodeModel{NodeID: nodeID, ModelName: modelName, ReplicaIndex: replicaIndex}, nil +} +func (f *fakeModelRouterForSmartRouter) RecordModelCleanupFailure(_ context.Context, _, _ string, _ int, _ string, _ time.Time) error { + return nil +} +func (f *fakeModelRouterForSmartRouter) ListModelCleanupRetries(_ context.Context, _ time.Time, _ int) ([]NodeModel, error) { + return nil, nil +} func (f *fakeModelRouterForSmartRouter) NextFreeReplicaIndex(_ context.Context, _, _ string, _ int) (int, error) { return 0, nil } @@ -186,7 +219,7 @@ var _ = Describe("ModelRouterAdapter", func() { adapter := NewModelRouterAdapter(sr) opts := &pb.ModelOptions{Model: "test-model"} - m, err := adapter.Route(context.Background(), "llama-cpp", "test-model", "test-model", "model.gguf", opts, false) + m, err := adapter.Route(context.Background(), "llama-cpp", "test-model", "test-model", "model.gguf", "", opts, false) Expect(err).NotTo(HaveOccurred()) Expect(m).NotTo(BeNil()) diff --git a/core/services/nodes/reconciler.go b/core/services/nodes/reconciler.go index 38fd97cc2f37..f7fefc8356c8 100644 --- a/core/services/nodes/reconciler.go +++ b/core/services/nodes/reconciler.go @@ -450,7 +450,7 @@ const probeFailuresBeforeReap = 3 func (rc *ReplicaReconciler) probeLoadedModels(ctx context.Context) { var stale []NodeModel cutoff := time.Now().Add(-rc.probeStaleAfter) - err := rc.registry.db.WithContext(ctx). + err := currentModelRevision(rc.registry.db.WithContext(ctx)). Joins("JOIN backend_nodes ON backend_nodes.id = node_models.node_id"). Where("node_models.state = ? AND backend_nodes.status = ? AND node_models.updated_at < ? AND node_models.address != ''", "loaded", StatusHealthy, cutoff). @@ -532,7 +532,7 @@ const inFlightLeakConfirmations = 2 func (rc *ReplicaReconciler) sweepLeakedInFlight(ctx context.Context) { var suspects []NodeModel cutoff := time.Now().Add(-inFlightLeakIdleAfter) - err := rc.registry.db.WithContext(ctx). + err := currentModelRevision(rc.registry.db.WithContext(ctx)). Joins("JOIN backend_nodes ON backend_nodes.id = node_models.node_id"). Where("node_models.state = ? AND backend_nodes.status = ? AND node_models.in_flight > 0 AND node_models.last_used < ? AND node_models.address != ''", "loaded", StatusHealthy, cutoff). @@ -632,7 +632,7 @@ func (rc *ReplicaReconciler) reconcileNodeProcesses(ctx context.Context) { var stale []NodeModel cutoff := time.Now().Add(-rc.probeStaleAfter) - err := rc.registry.db.WithContext(ctx). + err := currentModelRevision(rc.registry.db.WithContext(ctx)). Joins("JOIN backend_nodes ON backend_nodes.id = node_models.node_id"). Where("node_models.state = ? AND backend_nodes.status = ? AND backend_nodes.node_type = ? AND node_models.updated_at < ?", "loaded", StatusHealthy, NodeTypeBackend, cutoff). @@ -1067,8 +1067,8 @@ func (rc *ReplicaReconciler) scaleDownIdle(ctx context.Context, cfg ModelSchedul // and matching the worker supervisor's port-recycling behavior. cutoff := time.Now().Add(-rc.scaleDownDelay) var idleModels []NodeModel - rc.registry.db.WithContext(ctx). - Where("model_name = ? AND state = ? AND in_flight = 0 AND last_used < ?", + currentModelRevision(rc.registry.db.WithContext(ctx)). + Where("node_models.model_name = ? AND node_models.state = ? AND node_models.in_flight = 0 AND node_models.last_used < ?", cfg.ModelName, "loaded", cutoff). Order("replica_index DESC, last_used ASC"). Find(&idleModels) @@ -1097,8 +1097,8 @@ func (rc *ReplicaReconciler) scaleDownIdle(ctx context.Context, cfg ModelSchedul // allReplicasBusy returns true if all loaded replicas of a model have in-flight requests. func (rc *ReplicaReconciler) allReplicasBusy(ctx context.Context, modelName string) bool { var idleCount int64 - rc.registry.db.WithContext(ctx).Model(&NodeModel{}). - Where("model_name = ? AND state = ? AND in_flight = 0", modelName, "loaded"). + currentModelRevision(rc.registry.db.WithContext(ctx).Model(&NodeModel{})). + Where("node_models.model_name = ? AND node_models.state = ? AND node_models.in_flight = 0", modelName, "loaded"). Count(&idleCount) return idleCount == 0 } diff --git a/core/services/nodes/registry.go b/core/services/nodes/registry.go index 1a59e361f2e5..d6fd59662ff0 100644 --- a/core/services/nodes/registry.go +++ b/core/services/nodes/registry.go @@ -123,19 +123,24 @@ const ( // gRPC Address (each replica is a separate worker process on its own port), // and its own InFlight counter. type NodeModel struct { - ID string `gorm:"primaryKey;size:36" json:"id"` - NodeID string `gorm:"index;size:36" json:"node_id"` - ModelName string `gorm:"index;size:255" json:"model_name"` - ReplicaIndex int `gorm:"column:replica_index;default:0;index" json:"replica_index"` - Address string `gorm:"size:255" json:"address"` // gRPC address for this replica's backend process - State string `gorm:"size:32;default:idle" json:"state"` // staging, loading, loaded, unloading, idle - InFlight int `json:"in_flight"` // number of active requests on this replica - LastUsed time.Time `json:"last_used"` - LoadingBy string `gorm:"size:36" json:"loading_by,omitempty"` // frontend ID that triggered loading - BackendType string `gorm:"size:128" json:"backend_type,omitempty"` // e.g. "llama-cpp"; used by reconciler to replicate loads - ModelOptsBlob []byte `gorm:"type:bytea" json:"-"` // serialized pb.ModelOptions for replica scale-ups - CreatedAt time.Time `json:"created_at"` - UpdatedAt time.Time `json:"updated_at"` + ID string `gorm:"primaryKey;size:36" json:"id"` + NodeID string `gorm:"index;size:36" json:"node_id"` + ModelName string `gorm:"index;size:255" json:"model_name"` + ReplicaIndex int `gorm:"column:replica_index;default:0;index" json:"replica_index"` + Address string `gorm:"size:255" json:"address"` // gRPC address for this replica's backend process + State string `gorm:"size:32;default:idle" json:"state"` // staging, loading, loaded, unloading, idle + InFlight int `json:"in_flight"` // number of active requests on this replica + LastUsed time.Time `json:"last_used"` + LoadingBy string `gorm:"size:36" json:"loading_by,omitempty"` // frontend ID that triggered loading + BackendType string `gorm:"size:128" json:"backend_type,omitempty"` // e.g. "llama-cpp"; used by reconciler to replicate loads + ModelOptsBlob []byte `gorm:"type:bytea" json:"-"` // serialized pb.ModelOptions for replica scale-ups + ConfigRevision string `gorm:"column:config_revision;size:255" json:"config_revision,omitempty"` + EffectiveOptionsHash string `gorm:"column:effective_options_hash;size:128" json:"effective_options_hash,omitempty"` + CleanupError string `gorm:"column:cleanup_error;type:text" json:"cleanup_error,omitempty"` + CleanupAttempts int `gorm:"column:cleanup_attempts;default:0" json:"cleanup_attempts,omitempty"` + CleanupNextRetryAt *time.Time `gorm:"column:cleanup_next_retry_at" json:"cleanup_next_retry_at,omitempty"` + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` } // ModelLoadInfo is per-model load metadata kept independently of NodeModel rows @@ -156,13 +161,24 @@ type NodeModel struct { // That is identical to the per-NodeModel-row semantics today; if a stronger // guarantee is needed in the future, the row carries UpdatedAt for ordering. type ModelLoadInfo struct { - ModelName string `gorm:"primaryKey;size:255" json:"model_name"` - BackendType string `gorm:"size:128" json:"backend_type"` - ModelOptsBlob []byte `gorm:"type:bytea" json:"-"` - CreatedAt time.Time `json:"created_at"` - UpdatedAt time.Time `json:"updated_at"` + ModelName string `gorm:"primaryKey;size:255" json:"model_name"` + BackendType string `gorm:"size:128" json:"backend_type"` + ModelOptsBlob []byte `gorm:"type:bytea" json:"-"` + ConfigRevision string `gorm:"column:config_revision;size:255" json:"config_revision,omitempty"` + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` } +// ModelConfigState is the controller's current configuration generation for a model. +type ModelConfigState struct { + ModelName string `gorm:"primaryKey;size:255" json:"model_name"` + ConfigRevision string `gorm:"column:config_revision;size:255;check:model_config_states_revision_nonempty,config_revision <> ''" json:"config_revision"` + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` +} + +var ErrStaleModelConfigRevision = errors.New("stale model config revision") + // NodeLabel is a key-value label on a node (like K8s labels). type NodeLabel struct { ID string `gorm:"primaryKey;size:36" json:"id"` @@ -383,7 +399,7 @@ func (r *NodeRegistry) nodeModelNames(ctx context.Context, db *gorm.DB, nodeID s // when multiple instances (frontend + workers) start at the same time. func NewNodeRegistry(db *gorm.DB) (*NodeRegistry, error) { if err := advisorylock.WithLockCtx(context.Background(), db, advisorylock.KeySchemaMigrate, func() error { - return db.AutoMigrate(&BackendNode{}, &NodeModel{}, &NodeLabel{}, &ModelSchedulingConfig{}, &PendingBackendOp{}, &ModelLoadInfo{}, &ModelLoadJob{}) + return db.AutoMigrate(&BackendNode{}, &NodeModel{}, &NodeLabel{}, &ModelSchedulingConfig{}, &PendingBackendOp{}, &ModelLoadInfo{}, &ModelLoadJob{}, &ModelConfigState{}) }); err != nil { return nil, fmt.Errorf("migrating node tables: %w", err) } @@ -658,13 +674,14 @@ func (r *NodeRegistry) MarkOffline(ctx context.Context, nodeID string) error { func (r *NodeRegistry) FindNodeWithVRAM(ctx context.Context, minBytes uint64) (*BackendNode, error) { db := r.db.WithContext(ctx) - loadedModels := db.Model(&NodeModel{}). + loadedModels := currentModelRevision(db.Model(&NodeModel{})). Select("node_id"). - Where("state = ?", "loaded"). + Where("node_models.state = ?", "loaded"). Group("node_id") - subquery := db.Model(&NodeModel{}). + subquery := currentModelRevision(db.Model(&NodeModel{})). Select("node_id, COALESCE(SUM(in_flight), 0) as total_inflight"). + Where("node_models.state = ?", "loaded"). Group("node_id") // Try idle nodes with enough effectively-free VRAM first, prefer the one @@ -928,16 +945,16 @@ func (r *NodeRegistry) GetWithExtras(ctx context.Context, nodeID string) (*NodeW } var modelCount int64 - if err := r.db.WithContext(ctx).Model(&NodeModel{}). - Where("node_id = ? AND state = ?", nodeID, "loaded"). + if err := currentModelRevision(r.db.WithContext(ctx).Model(&NodeModel{})). + Where("node_models.node_id = ? AND node_models.state = ?", nodeID, "loaded"). Count(&modelCount).Error; err != nil { xlog.Warn("GetWithExtras: failed to get model count", "node", nodeID, "error", err) } var inFlight struct{ Total int } - if err := r.db.WithContext(ctx).Model(&NodeModel{}). + if err := currentModelRevision(r.db.WithContext(ctx).Model(&NodeModel{})). Select("COALESCE(SUM(in_flight), 0) as total"). - Where("node_id = ? AND state IN ?", nodeID, []string{"loaded", "unloading"}). + Where("node_models.node_id = ? AND node_models.state = ?", nodeID, "loaded"). Scan(&inFlight).Error; err != nil { xlog.Warn("GetWithExtras: failed to get in-flight count", "node", nodeID, "error", err) } @@ -1028,26 +1045,60 @@ func (r *NodeRegistry) FindStaleNodes(ctx context.Context, threshold time.Durati // replicaIndex identifies which slot on the node this replica occupies // (0..MaxReplicasPerModel-1). Pass 0 for single-replica scheduling. func (r *NodeRegistry) SetNodeModel(ctx context.Context, nodeID, modelName string, replicaIndex int, state, address string, initialInFlight int) error { + revision, _ := r.GetModelConfigRevision(ctx, modelName) + return r.setNodeModelRevision(ctx, nodeID, modelName, replicaIndex, state, address, initialInFlight, revision, "", false) +} + +func (r *NodeRegistry) SetNodeModelRevision(ctx context.Context, nodeID, modelName string, replicaIndex int, state, address string, initialInFlight int, revision, effectiveOptionsHash string) error { + return r.setNodeModelRevision(ctx, nodeID, modelName, replicaIndex, state, address, initialInFlight, revision, effectiveOptionsHash, true) +} + +func (r *NodeRegistry) setNodeModelRevision(ctx context.Context, nodeID, modelName string, replicaIndex int, state, address string, initialInFlight int, revision, effectiveOptionsHash string, revisionRequired bool) error { + if err := validateRevisionWrite(modelName, revision, revisionRequired); err != nil { + return err + } now := time.Now() // Use Attrs for creation-only fields (ID) and Assign for update-only fields. // Attrs is applied only when creating a new record. Assign is applied on // both create and update. This prevents overwriting the primary key on // subsequent calls for the same (node, model, replica_index). - var nm NodeModel - result := r.db.WithContext(ctx).Where("node_id = ? AND model_name = ? AND replica_index = ?", nodeID, modelName, replicaIndex). - Attrs(NodeModel{ID: uuid.New().String(), NodeID: nodeID, ModelName: modelName, ReplicaIndex: replicaIndex}). - Assign(map[string]any{"address": address, "state": state, "last_used": now, "in_flight": initialInFlight}). - FirstOrCreate(&nm) - return result.Error + return r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + if err := requireCurrentRevision(tx, modelName, revision); err != nil { + return err + } + var nm NodeModel + return tx.Where("node_id = ? AND model_name = ? AND replica_index = ?", nodeID, modelName, replicaIndex). + Attrs(NodeModel{ID: uuid.New().String(), NodeID: nodeID, ModelName: modelName, ReplicaIndex: replicaIndex}). + Assign(map[string]any{"address": address, "state": state, "last_used": now, "in_flight": initialInFlight, + "config_revision": revision, "effective_options_hash": effectiveOptionsHash}). + FirstOrCreate(&nm).Error + }) } // SetNodeModelLoadInfo stores the backend type and serialized model options on // an existing NodeModel record. This metadata is used by the reconciler to // replicate model loads during scale-up. func (r *NodeRegistry) SetNodeModelLoadInfo(ctx context.Context, nodeID, modelName string, replicaIndex int, backendType string, optsBlob []byte) error { - return r.db.WithContext(ctx).Model(&NodeModel{}). - Where("node_id = ? AND model_name = ? AND replica_index = ?", nodeID, modelName, replicaIndex). - Updates(map[string]any{"backend_type": backendType, "model_opts_blob": optsBlob}).Error + revision, _ := r.GetModelConfigRevision(ctx, modelName) + return r.setNodeModelLoadInfoRevision(ctx, nodeID, modelName, replicaIndex, backendType, revision, optsBlob, false) +} + +func (r *NodeRegistry) SetNodeModelLoadInfoRevision(ctx context.Context, nodeID, modelName string, replicaIndex int, backendType, revision string, optsBlob []byte) error { + return r.setNodeModelLoadInfoRevision(ctx, nodeID, modelName, replicaIndex, backendType, revision, optsBlob, true) +} + +func (r *NodeRegistry) setNodeModelLoadInfoRevision(ctx context.Context, nodeID, modelName string, replicaIndex int, backendType, revision string, optsBlob []byte, revisionRequired bool) error { + if err := validateRevisionWrite(modelName, revision, revisionRequired); err != nil { + return err + } + return r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + if err := requireCurrentRevision(tx, modelName, revision); err != nil { + return err + } + return tx.Model(&NodeModel{}). + Where("node_id = ? AND model_name = ? AND replica_index = ?", nodeID, modelName, replicaIndex). + Updates(map[string]any{"backend_type": backendType, "model_opts_blob": optsBlob, "config_revision": revision}).Error + }) } // UpsertModelLoadInfo records or replaces the per-model load info in the @@ -1061,25 +1112,75 @@ func (r *NodeRegistry) SetNodeModelLoadInfo(ctx context.Context, nodeID, modelNa // opts converge on whichever transaction committed last; that matches the // existing per-replica blob semantics today. func (r *NodeRegistry) UpsertModelLoadInfo(ctx context.Context, modelName, backendType string, optsBlob []byte) error { - if modelName == "" { - return fmt.Errorf("model name is required") + revision, _ := r.GetModelConfigRevision(ctx, modelName) + return r.upsertModelLoadInfoRevision(ctx, modelName, backendType, revision, optsBlob, false) +} + +func (r *NodeRegistry) UpsertModelLoadInfoRevision(ctx context.Context, modelName, backendType, revision string, optsBlob []byte) error { + return r.upsertModelLoadInfoRevision(ctx, modelName, backendType, revision, optsBlob, true) +} + +func (r *NodeRegistry) upsertModelLoadInfoRevision(ctx context.Context, modelName, backendType, revision string, optsBlob []byte, revisionRequired bool) error { + if err := validateRevisionWrite(modelName, revision, revisionRequired); err != nil { + return err } now := time.Now() rec := ModelLoadInfo{ - ModelName: modelName, - BackendType: backendType, - ModelOptsBlob: optsBlob, - CreatedAt: now, - UpdatedAt: now, + ModelName: modelName, + BackendType: backendType, + ModelOptsBlob: optsBlob, + ConfigRevision: revision, + CreatedAt: now, + UpdatedAt: now, } - return r.db.WithContext(ctx).Clauses(clause.OnConflict{ - Columns: []clause.Column{{Name: "model_name"}}, - DoUpdates: clause.Assignments(map[string]any{ - "backend_type": backendType, - "model_opts_blob": optsBlob, - "updated_at": now, - }), - }).Create(&rec).Error + return r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + if err := requireCurrentRevision(tx, modelName, revision); err != nil { + return err + } + return tx.Clauses(clause.OnConflict{ + Columns: []clause.Column{{Name: "model_name"}}, + DoUpdates: clause.Assignments(map[string]any{ + "backend_type": backendType, + "model_opts_blob": optsBlob, + "config_revision": revision, + "updated_at": now, + }), + }).Create(&rec).Error + }) +} + +func requireCurrentRevision(tx *gorm.DB, modelName, revision string) error { + var state ModelConfigState + err := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("model_name = ?", modelName).First(&state).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil + } + if err != nil { + return err + } + if state.ConfigRevision != revision { + return ErrStaleModelConfigRevision + } + return nil +} + +func validateRevisionWrite(modelName, revision string, revisionRequired bool) error { + if modelName == "" { + return fmt.Errorf("model name is required") + } + if revisionRequired && revision == "" { + return fmt.Errorf("config revision is required") + } + return nil +} + +// currentModelRevision limits node_models queries to rows that are safe to +// publish. Legacy rows remain eligible until a current state exists; once it +// does, only an exact revision match is eligible. The correlated subquery +// deliberately avoids adding a JOIN, so callers that already join +// node_models cannot generate duplicate table aliases on PostgreSQL. +func currentModelRevision(db *gorm.DB) *gorm.DB { + return db.Where("NOT EXISTS (SELECT 1 FROM model_config_states WHERE model_config_states.model_name = node_models.model_name) OR node_models.config_revision = (SELECT config_revision FROM model_config_states WHERE model_config_states.model_name = node_models.model_name)") } // GetModelLoadInfo retrieves the stored backend type and serialized model @@ -1089,23 +1190,172 @@ func (r *NodeRegistry) UpsertModelLoadInfo(ctx context.Context, modelName, backe // UpsertModelLoadInfo (rolling-upgrade transition). Returns // gorm.ErrRecordNotFound when neither source has an entry. func (r *NodeRegistry) GetModelLoadInfo(ctx context.Context, modelName string) (backendType string, optsBlob []byte, err error) { + backendType, _, optsBlob, err = r.GetModelLoadInfoRevision(ctx, modelName) + return backendType, optsBlob, err +} + +func (r *NodeRegistry) GetModelLoadInfoRevision(ctx context.Context, modelName string) (backendType, revision string, optsBlob []byte, err error) { var info ModelLoadInfo - err = r.db.WithContext(ctx).Where("model_name = ?", modelName).First(&info).Error + err = r.db.WithContext(ctx). + Where("model_load_infos.model_name = ?", modelName). + Where("NOT EXISTS (SELECT 1 FROM model_config_states WHERE model_config_states.model_name = model_load_infos.model_name) OR (model_load_infos.config_revision <> '' AND model_load_infos.config_revision = (SELECT config_revision FROM model_config_states WHERE model_config_states.model_name = model_load_infos.model_name))"). + First(&info).Error if err == nil { - return info.BackendType, info.ModelOptsBlob, nil + return info.BackendType, info.ConfigRevision, info.ModelOptsBlob, nil } if !errors.Is(err, gorm.ErrRecordNotFound) { - return "", nil, err + return "", "", nil, err } var nm NodeModel err = r.db.WithContext(ctx). Where("model_name = ? AND state = ? AND model_opts_blob IS NOT NULL", modelName, "loaded"). + Where("NOT EXISTS (SELECT 1 FROM model_config_states WHERE model_config_states.model_name = node_models.model_name) OR node_models.config_revision = (SELECT config_revision FROM model_config_states WHERE model_config_states.model_name = node_models.model_name)"). First(&nm).Error if err != nil { - return "", nil, err + return "", "", nil, err + } + return nm.BackendType, nm.ConfigRevision, nm.ModelOptsBlob, nil +} + +func (r *NodeRegistry) GetModelConfigRevision(ctx context.Context, modelName string) (string, error) { + var state ModelConfigState + err := r.db.WithContext(ctx).Where("model_name = ?", modelName).First(&state).Error + if err != nil { + return "", err } - return nm.BackendType, nm.ModelOptsBlob, nil + return state.ConfigRevision, nil +} + +// EstablishModelConfigRevision creates the initial current revision without +// ever replacing one. Inference requests use this operation so a late request +// carrying an older config cannot roll controller state backward. +func (r *NodeRegistry) EstablishModelConfigRevision(ctx context.Context, modelName, revision string) error { + if err := validateRevisionWrite(modelName, revision, true); err != nil { + return err + } + now := time.Now() + return r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + state := ModelConfigState{ModelName: modelName, ConfigRevision: revision, CreatedAt: now, UpdatedAt: now} + if err := tx.Clauses(clause.OnConflict{Columns: []clause.Column{{Name: "model_name"}}, DoNothing: true}).Create(&state).Error; err != nil { + return err + } + return requireCurrentRevision(tx, modelName, revision) + }) +} + +type ModelConfigRevisionTransition struct { + ModelName string + ConfigRevision string +} + +func (r *NodeRegistry) AdvanceModelConfigRevision(ctx context.Context, modelName, revision string) ([]NodeModel, error) { + return r.AdvanceModelConfigRevisions(ctx, []ModelConfigRevisionTransition{{ModelName: modelName, ConfigRevision: revision}}) +} + +// AdvanceModelConfigRevisions publishes one or more related configuration +// identities in a single transaction. Renames use this boundary so the old +// identity cannot advance when establishing the new identity fails. +func (r *NodeRegistry) AdvanceModelConfigRevisions(ctx context.Context, transitions []ModelConfigRevisionTransition) ([]NodeModel, error) { + if len(transitions) == 0 { + return nil, errors.New("at least one model config revision transition is required") + } + for _, transition := range transitions { + if err := validateRevisionWrite(transition.ModelName, transition.ConfigRevision, true); err != nil { + return nil, err + } + } + var quarantined []NodeModel + err := r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + for _, transition := range transitions { + now := time.Now() + state := ModelConfigState{ModelName: transition.ModelName, ConfigRevision: transition.ConfigRevision, CreatedAt: now, UpdatedAt: now} + if err := tx.Clauses(clause.OnConflict{Columns: []clause.Column{{Name: "model_name"}}, DoUpdates: clause.Assignments(map[string]any{"config_revision": transition.ConfigRevision, "updated_at": now})}).Create(&state).Error; err != nil { + return err + } + staleReplica := "model_name = ? AND state IN ? AND (config_revision IS NULL OR config_revision = '' OR config_revision <> ?)" + activeStates := []string{"loaded", "loading", "staging"} + var transitionQuarantined []NodeModel + if err := tx.Where(staleReplica, transition.ModelName, activeStates, transition.ConfigRevision).Find(&transitionQuarantined).Error; err != nil { + return err + } + if err := tx.Model(&NodeModel{}).Where(staleReplica, transition.ModelName, activeStates, transition.ConfigRevision).Updates(map[string]any{"state": "unloading", "cleanup_error": "", "cleanup_attempts": 0, "cleanup_next_retry_at": nil}).Error; err != nil { + return err + } + for i := range transitionQuarantined { + transitionQuarantined[i].State = "unloading" + transitionQuarantined[i].CleanupError = "" + transitionQuarantined[i].CleanupAttempts = 0 + transitionQuarantined[i].CleanupNextRetryAt = nil + } + quarantined = append(quarantined, transitionQuarantined...) + if err := tx.Where("model_name = ? AND (config_revision IS NULL OR config_revision <> ?)", transition.ModelName, transition.ConfigRevision).Delete(&ModelLoadInfo{}).Error; err != nil { + return err + } + } + return nil + }) + if err != nil { + return nil, err + } + return quarantined, nil +} + +func (r *NodeRegistry) RecordModelCleanupFailure(ctx context.Context, nodeID, modelName string, replicaIndex int, cleanupErr string, nextRetry time.Time) error { + return r.db.WithContext(ctx).Model(&NodeModel{}).Where("node_id = ? AND model_name = ? AND replica_index = ? AND state = ?", nodeID, modelName, replicaIndex, "unloading").Updates(map[string]any{"cleanup_error": cleanupErr, "cleanup_attempts": gorm.Expr("cleanup_attempts + 1"), "cleanup_next_retry_at": nextRetry}).Error +} + +func (r *NodeRegistry) ListModelCleanupRetries(ctx context.Context, now time.Time, limit int) ([]NodeModel, error) { + var models []NodeModel + q := r.db.WithContext(ctx).Where("state = ? AND cleanup_next_retry_at IS NOT NULL AND cleanup_next_retry_at <= ?", "unloading", now).Order("cleanup_next_retry_at ASC") + if limit > 0 { + q = q.Limit(limit) + } + err := q.Find(&models).Error + return models, err +} + +// ClaimModelCleanupRetries leases due quarantine rows in one transaction. The +// row locks prevent two frontends from sending the same exact-stop request, +// while SKIP LOCKED lets each frontend take different work without waiting. +func (r *NodeRegistry) ClaimModelCleanupRetries(ctx context.Context, now, leaseUntil time.Time, limit int) ([]NodeModel, error) { + var models []NodeModel + err := r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + q := tx.Clauses(clause.Locking{Strength: "UPDATE", Options: "SKIP LOCKED"}). + Where("state = ? AND (cleanup_next_retry_at IS NULL OR cleanup_next_retry_at <= ?)", "unloading", now). + Order("cleanup_next_retry_at ASC") + if limit > 0 { + q = q.Limit(limit) + } + if err := q.Find(&models).Error; err != nil || len(models) == 0 { + return err + } + ids := make([]string, len(models)) + for i := range models { + ids[i] = models[i].ID + } + return tx.Model(&NodeModel{}).Where("id IN ?", ids).Update("cleanup_next_retry_at", leaseUntil).Error + }) + return models, err +} + +// RemoveClaimedModelCleanup deletes only the exact quarantine row that was +// stopped. A worker may re-register a replacement in the same logical slot +// while the stop request is in flight; matching the immutable row identity and +// stop inputs prevents cleanup from deleting that replacement. +func (r *NodeRegistry) RemoveClaimedModelCleanup(ctx context.Context, replica NodeModel) (bool, error) { + result := r.db.WithContext(ctx). + Where("id = ? AND node_id = ? AND model_name = ? AND replica_index = ? AND state = ? AND address = ? AND config_revision = ?", + replica.ID, replica.NodeID, replica.ModelName, replica.ReplicaIndex, "unloading", replica.Address, replica.ConfigRevision). + Delete(&NodeModel{}) + if result.Error != nil { + return false, result.Error + } + if result.RowsAffected == 0 { + return false, nil + } + r.fireReplicaRemoved(replica.ModelName, replica.NodeID, replica.ReplicaIndex) + return true, nil } // RemoveNodeModel removes a single replica of a model from a node. @@ -1142,6 +1392,7 @@ func (r *NodeRegistry) FindNodesWithModel(ctx context.Context, modelName string) if err := r.db.WithContext(ctx).Joins("JOIN node_models ON node_models.node_id = backend_nodes.id"). Where("node_models.model_name = ? AND node_models.state = ? AND backend_nodes.status = ?", modelName, "loaded", StatusHealthy). + Where("NOT EXISTS (SELECT 1 FROM model_config_states WHERE model_config_states.model_name = node_models.model_name) OR node_models.config_revision = (SELECT config_revision FROM model_config_states WHERE model_config_states.model_name = node_models.model_name)"). Order("node_models.in_flight ASC"). Find(&nodes).Error; err != nil { return nil, fmt.Errorf("finding nodes with model %s: %w", modelName, err) @@ -1186,6 +1437,22 @@ func (r *NodeRegistry) FindAndLockNodeWithModel(ctx context.Context, modelName s var node BackendNode err := r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + // Lock the current model revision before selecting a replica. Revision + // advancement takes this lock before quarantining replica rows too, so + // an edit and a route claim have one ordering: either this transaction + // reserves a replica before the edit, or it observes the new revision + // and cannot reserve the old replica. A revision subquery alone is not + // sufficient under READ COMMITTED because the state can change between + // SELECT and the in-flight increment. + var currentState ModelConfigState + hasCurrentRevision := false + if err := tx.Clauses(clause.Locking{Strength: "SHARE"}). + Where("model_name = ?", modelName).First(¤tState).Error; err == nil { + hasCurrentRevision = true + } else if !errors.Is(err, gorm.ErrRecordNotFound) { + return err + } + // Mirror of PickBestReplica's policy (see replicapicker.go): // 1. in_flight ASC — least busy replica. // 2. last_used ASC — round-robin between equally-loaded replicas. @@ -1208,6 +1475,9 @@ func (r *NodeRegistry) FindAndLockNodeWithModel(ctx context.Context, modelName s Joins("JOIN backend_nodes ON backend_nodes.id = node_models.node_id"). Where("node_models.model_name = ? AND node_models.state = ? AND backend_nodes.status = ?", modelName, "loaded", StatusHealthy) + if hasCurrentRevision { + base = base.Where("node_models.config_revision = ?", currentState.ConfigRevision) + } if len(candidateNodeIDs) > 0 { base = base.Where("node_models.node_id IN ?", candidateNodeIDs) } @@ -1269,7 +1539,7 @@ func (r *NodeRegistry) LoadedReplicaStats(ctx context.Context, modelName string, LastUsed time.Time AvailableVRAM uint64 } - q := r.db.WithContext(ctx).Model(&NodeModel{}). + q := currentModelRevision(r.db.WithContext(ctx).Model(&NodeModel{})). Joins("JOIN backend_nodes ON backend_nodes.id = node_models.node_id"). Where("node_models.model_name = ? AND node_models.state = ? AND backend_nodes.status = ?", modelName, "loaded", StatusHealthy) @@ -1321,7 +1591,8 @@ func (r *NodeRegistry) GetNodeModel(ctx context.Context, nodeID, modelName strin func (r *NodeRegistry) CountReplicasOnNode(ctx context.Context, nodeID, modelName string) (int, error) { var count int64 if err := r.db.WithContext(ctx).Model(&NodeModel{}). - Where("node_id = ? AND model_name = ?", nodeID, modelName). + Where("node_id = ? AND model_name = ? AND state <> ?", nodeID, modelName, "unloading"). + Where("NOT EXISTS (SELECT 1 FROM model_config_states WHERE model_config_states.model_name = node_models.model_name) OR node_models.config_revision = (SELECT config_revision FROM model_config_states WHERE model_config_states.model_name = node_models.model_name)"). Count(&count).Error; err != nil { return 0, err } @@ -1344,8 +1615,8 @@ func (r *NodeRegistry) NextFreeReplicaIndex(ctx context.Context, nodeID, modelNa return 0, ErrNoFreeSlot } var taken []int - if err := r.db.WithContext(ctx).Model(&NodeModel{}). - Where("node_id = ? AND model_name = ?", nodeID, modelName). + if err := currentModelRevision(r.db.WithContext(ctx).Model(&NodeModel{})). + Where("node_models.node_id = ? AND node_models.model_name = ? AND node_models.state <> ?", nodeID, modelName, "unloading"). Pluck("replica_index", &taken).Error; err != nil { return 0, err } @@ -1368,8 +1639,9 @@ func (r *NodeRegistry) FindLeastLoadedNode(ctx context.Context) (*BackendNode, e var node BackendNode query := db.Where("status = ? AND node_type = ?", StatusHealthy, NodeTypeBackend) // Order by total in-flight across all models on the node - subquery := db.Model(&NodeModel{}). + subquery := currentModelRevision(db.Model(&NodeModel{})). Select("node_id, COALESCE(SUM(in_flight), 0) as total_inflight"). + Where("node_models.state = ?", "loaded"). Group("node_id") err := query.Joins("LEFT JOIN (?) AS load ON load.node_id = backend_nodes.id", subquery). @@ -1387,9 +1659,9 @@ func (r *NodeRegistry) FindIdleNode(ctx context.Context) (*BackendNode, error) { db := r.db.WithContext(ctx) var node BackendNode - loadedModels := db.Model(&NodeModel{}). + loadedModels := currentModelRevision(db.Model(&NodeModel{})). Select("node_id"). - Where("state = ?", "loaded"). + Where("node_models.state = ?", "loaded"). Group("node_id") err := db.Where("status = ? AND node_type = ? AND id NOT IN (?)", StatusHealthy, NodeTypeBackend, loadedModels). Order("available_vram DESC"). @@ -1447,6 +1719,7 @@ func (r *NodeRegistry) ListAllLoadedModels(ctx context.Context) ([]NodeModel, er var models []NodeModel err := r.db.WithContext(ctx).Joins("JOIN backend_nodes ON backend_nodes.id = node_models.node_id"). Where("node_models.state = ? AND backend_nodes.status = ?", "loaded", StatusHealthy). + Where("NOT EXISTS (SELECT 1 FROM model_config_states WHERE model_config_states.model_name = node_models.model_name) OR node_models.config_revision = (SELECT config_revision FROM model_config_states WHERE model_config_states.model_name = node_models.model_name)"). Find(&models).Error if err != nil { return nil, fmt.Errorf("listing all loaded models: %w", err) @@ -1467,7 +1740,7 @@ func (r *NodeRegistry) FindNodeForModel(ctx context.Context, modelName string) ( // FindLRUModel returns the least-recently-used model on a node. func (r *NodeRegistry) FindLRUModel(ctx context.Context, nodeID string) (*NodeModel, error) { var nm NodeModel - err := r.db.WithContext(ctx).Where("node_id = ? AND state = ? AND in_flight = 0", nodeID, "loaded"). + err := currentModelRevision(r.db.WithContext(ctx)).Where("node_models.node_id = ? AND node_models.state = ? AND node_models.in_flight = 0", nodeID, "loaded"). Order("last_used ASC").First(&nm).Error if err != nil { return nil, fmt.Errorf("finding LRU model on node %s: %w", nodeID, err) @@ -1480,7 +1753,7 @@ func (r *NodeRegistry) FindLRUModel(ctx context.Context, nodeID string) (*NodeMo // Used by the router for preemptive eviction when no node has free VRAM. func (r *NodeRegistry) FindGlobalLRUModelWithZeroInFlight(ctx context.Context) (*NodeModel, error) { var nm NodeModel - err := r.db.WithContext(ctx).Joins("JOIN backend_nodes ON backend_nodes.id = node_models.node_id"). + err := currentModelRevision(r.db.WithContext(ctx)).Joins("JOIN backend_nodes ON backend_nodes.id = node_models.node_id"). Where("node_models.state = ? AND node_models.in_flight = 0 AND backend_nodes.status = ? AND backend_nodes.node_type = ?", "loaded", StatusHealthy, NodeTypeBackend). Order("node_models.last_used ASC"). @@ -1600,13 +1873,14 @@ func (r *NodeRegistry) FindNodesBySelector(ctx context.Context, selector map[str func (r *NodeRegistry) FindNodeWithVRAMFromSet(ctx context.Context, minBytes uint64, nodeIDs []string) (*BackendNode, error) { db := r.db.WithContext(ctx) - loadedModels := db.Model(&NodeModel{}). + loadedModels := currentModelRevision(db.Model(&NodeModel{})). Select("node_id"). - Where("state = ?", "loaded"). + Where("node_models.state = ?", "loaded"). Group("node_id") - subquery := db.Model(&NodeModel{}). + subquery := currentModelRevision(db.Model(&NodeModel{})). Select("node_id, COALESCE(SUM(in_flight), 0) as total_inflight"). + Where("node_models.state = ?", "loaded"). Group("node_id") // Try idle nodes with enough effectively-free VRAM first. @@ -1636,9 +1910,9 @@ func (r *NodeRegistry) FindIdleNodeFromSet(ctx context.Context, nodeIDs []string db := r.db.WithContext(ctx) var node BackendNode - loadedModels := db.Model(&NodeModel{}). + loadedModels := currentModelRevision(db.Model(&NodeModel{})). Select("node_id"). - Where("state = ?", "loaded"). + Where("node_models.state = ?", "loaded"). Group("node_id") err := db.Where("status = ? AND node_type = ? AND id NOT IN (?) AND id IN ?", StatusHealthy, NodeTypeBackend, loadedModels, nodeIDs). Order("available_vram DESC"). @@ -1656,8 +1930,9 @@ func (r *NodeRegistry) FindLeastLoadedNodeFromSet(ctx context.Context, nodeIDs [ var node BackendNode query := db.Where("status = ? AND node_type = ? AND backend_nodes.id IN ?", StatusHealthy, NodeTypeBackend, nodeIDs) // Order by total in-flight across all models on the node - subquery := db.Model(&NodeModel{}). + subquery := currentModelRevision(db.Model(&NodeModel{})). Select("node_id, COALESCE(SUM(in_flight), 0) as total_inflight"). + Where("node_models.state = ?", "loaded"). Group("node_id") err := query.Joins("LEFT JOIN (?) AS load ON load.node_id = backend_nodes.id", subquery). @@ -1737,7 +2012,9 @@ func (r *NodeRegistry) DeleteModelScheduling(ctx context.Context, modelName stri // CountLoadedReplicas returns the number of loaded replicas for a model. func (r *NodeRegistry) CountLoadedReplicas(ctx context.Context, modelName string) (int64, error) { var count int64 - err := r.db.WithContext(ctx).Model(&NodeModel{}).Where("model_name = ? AND state = ?", modelName, "loaded").Count(&count).Error + err := currentModelRevision(r.db.WithContext(ctx).Model(&NodeModel{})). + Where("node_models.model_name = ? AND node_models.state = ?", modelName, "loaded"). + Count(&count).Error return count, err } @@ -1758,9 +2035,9 @@ func (r *NodeRegistry) FindNodesWithFreeSlot(ctx context.Context, modelName stri // Subquery: per-node count of loaded+loading replicas of this model. // We count any non-removed row (state != deleted) so a load in progress // counts against the cap and a second concurrent scale-up can't overshoot. - subq := r.db.Model(&NodeModel{}). + subq := currentModelRevision(r.db.Model(&NodeModel{})). Select("node_id, COUNT(*) as cnt"). - Where("model_name = ?", modelName). + Where("node_models.model_name = ? AND node_models.state <> ?", modelName, "unloading"). Group("node_id") var out []BackendNode @@ -1787,9 +2064,9 @@ func (r *NodeRegistry) ClusterCapacityForModel(ctx context.Context, modelName st if len(candidateNodeIDs) > 0 { q = q.Where("id IN ?", candidateNodeIDs) } - subq := r.db.Model(&NodeModel{}). + subq := currentModelRevision(r.db.Model(&NodeModel{})). Select("node_id, COUNT(*) as cnt"). - Where("model_name = ?", modelName). + Where("node_models.model_name = ? AND node_models.state <> ?", modelName, "unloading"). Group("node_id") var nodes []struct { @@ -1986,9 +2263,9 @@ func (r *NodeRegistry) ListWithExtras(ctx context.Context) ([]NodeWithExtras, er Count int } var counts []modelCount - if err := r.db.WithContext(ctx).Model(&NodeModel{}). + if err := currentModelRevision(r.db.WithContext(ctx).Model(&NodeModel{})). Select("node_id, COUNT(*) as count"). - Where("state = ?", "loaded"). + Where("node_models.state = ?", "loaded"). Group("node_id"). Find(&counts).Error; err != nil { xlog.Warn("ListWithExtras: failed to get model counts", "error", err) @@ -2005,9 +2282,9 @@ func (r *NodeRegistry) ListWithExtras(ctx context.Context) ([]NodeWithExtras, er Total int } var inFlights []inFlightCount - if err := r.db.WithContext(ctx).Model(&NodeModel{}). + if err := currentModelRevision(r.db.WithContext(ctx).Model(&NodeModel{})). Select("node_id, COALESCE(SUM(in_flight), 0) as total"). - Where("state IN ?", []string{"loaded", "unloading"}). + Where("node_models.state = ?", "loaded"). Group("node_id"). Find(&inFlights).Error; err != nil { xlog.Warn("ListWithExtras: failed to get in-flight counts", "error", err) diff --git a/core/services/nodes/registry_test.go b/core/services/nodes/registry_test.go index a9d4b059aaa1..c240f2f015ef 100644 --- a/core/services/nodes/registry_test.go +++ b/core/services/nodes/registry_test.go @@ -2,6 +2,7 @@ package nodes import ( "context" + "errors" "runtime" "time" @@ -1520,6 +1521,254 @@ var _ = Describe("NodeRegistry", func() { Expect(err).To(HaveOccurred()) }) }) + + Describe("model config revisions", func() { + It("rejects empty model names and required revisions without writing state", func() { + ctx := context.Background() + node := makeNode("revision-validation", "10.0.2.9:50051", 8_000_000_000) + Expect(registry.Register(ctx, node, true)).To(Succeed()) + + _, err := registry.AdvanceModelConfigRevision(ctx, "", "rev-1") + Expect(err).To(HaveOccurred()) + _, err = registry.AdvanceModelConfigRevision(ctx, "model", "") + Expect(err).To(HaveOccurred()) + Expect(registry.SetNodeModelRevision(ctx, node.ID, "model", 0, "loaded", node.Address, 0, "", "hash")).To(HaveOccurred()) + Expect(registry.SetNodeModelLoadInfoRevision(ctx, node.ID, "model", 0, "llama-cpp", "", []byte("opts"))).To(HaveOccurred()) + Expect(registry.UpsertModelLoadInfoRevision(ctx, "model", "llama-cpp", "", []byte("opts"))).To(HaveOccurred()) + _, err = registry.GetModelConfigRevision(ctx, "model") + Expect(err).To(MatchError(gorm.ErrRecordNotFound)) + }) + + It("returns no quarantined rows when advancing the revision rolls back", func() { + ctx := context.Background() + node := makeNode("revision-rollback", "10.0.2.10:50051", 8_000_000_000) + Expect(registry.Register(ctx, node, true)).To(Succeed()) + Expect(registry.AdvanceModelConfigRevision(ctx, "rollback-model", "rev-1")).To(BeEmpty()) + Expect(registry.SetNodeModelRevision(ctx, node.ID, "rollback-model", 0, "loaded", node.Address, 0, "rev-1", "hash-1")).To(Succeed()) + Expect(registry.UpsertModelLoadInfoRevision(ctx, "rollback-model", "llama-cpp", "rev-1", []byte("opts-1"))).To(Succeed()) + + callbackName := "test:fail-model-load-info-delete" + Expect(db.Callback().Delete().Before("gorm:delete").Register(callbackName, func(tx *gorm.DB) { + if tx.Statement.Table == "model_load_infos" { + _ = tx.AddError(errors.New("injected delete failure")) + } + })).To(Succeed()) + DeferCleanup(func() { Expect(db.Callback().Delete().Remove(callbackName)).To(Succeed()) }) + + quarantined, err := registry.AdvanceModelConfigRevision(ctx, "rollback-model", "rev-2") + Expect(err).To(MatchError("injected delete failure")) + Expect(quarantined).To(BeEmpty(), "rolled-back rows must never escape as cleanup work") + Expect(registry.GetModelConfigRevision(ctx, "rollback-model")).To(Equal("rev-1")) + persisted, err := registry.GetNodeModel(ctx, node.ID, "rollback-model", 0) + Expect(err).ToNot(HaveOccurred()) + Expect(persisted.State).To(Equal("loaded")) + }) + + It("rolls back both rename identities when the second transition fails", func() { + ctx := context.Background() + node := makeNode("revision-rename-rollback", "10.0.2.15:50051", 8_000_000_000) + Expect(registry.Register(ctx, node, true)).To(Succeed()) + Expect(registry.AdvanceModelConfigRevision(ctx, "old-name", "rev-1")).To(BeEmpty()) + Expect(registry.SetNodeModelRevision(ctx, node.ID, "old-name", 0, "loaded", node.Address, 0, "rev-1", "hash-1")).To(Succeed()) + + callbackName := "test:fail-second-rename-revision" + Expect(db.Callback().Create().Before("gorm:create").Register(callbackName, func(tx *gorm.DB) { + if state, ok := tx.Statement.Dest.(*ModelConfigState); ok && state.ModelName == "new-name" { + _ = tx.AddError(errors.New("injected second transition failure")) + } + })).To(Succeed()) + DeferCleanup(func() { Expect(db.Callback().Create().Remove(callbackName)).To(Succeed()) }) + + quarantined, err := registry.AdvanceModelConfigRevisions(ctx, []ModelConfigRevisionTransition{ + {ModelName: "old-name", ConfigRevision: "rev-2"}, + {ModelName: "new-name", ConfigRevision: "rev-2"}, + }) + Expect(err).To(MatchError("injected second transition failure")) + Expect(quarantined).To(BeEmpty()) + Expect(registry.GetModelConfigRevision(ctx, "old-name")).To(Equal("rev-1")) + _, err = registry.GetModelConfigRevision(ctx, "new-name") + Expect(err).To(MatchError(gorm.ErrRecordNotFound)) + persisted, err := registry.GetNodeModel(ctx, node.ID, "old-name", 0) + Expect(err).ToNot(HaveOccurred()) + Expect(persisted.State).To(Equal("loaded")) + }) + + It("excludes empty, mismatched, and unloading replicas from routing and statistics", func() { + ctx := context.Background() + matching := makeNode("revision-current", "10.0.2.11:50051", 8_000_000_000) + empty := makeNode("revision-empty", "10.0.2.12:50051", 8_000_000_000) + mismatched := makeNode("revision-mismatch", "10.0.2.13:50051", 8_000_000_000) + unloading := makeNode("revision-unloading", "10.0.2.14:50051", 8_000_000_000) + for _, node := range []*BackendNode{matching, empty, mismatched, unloading} { + Expect(registry.Register(ctx, node, true)).To(Succeed()) + } + Expect(registry.AdvanceModelConfigRevision(ctx, "filtered-model", "rev-2")).To(BeEmpty()) + Expect(registry.SetNodeModelRevision(ctx, matching.ID, "filtered-model", 0, "loaded", matching.Address, 3, "rev-2", "hash-2")).To(Succeed()) + Expect(db.Create(&NodeModel{ID: "empty-revision", NodeID: empty.ID, ModelName: "filtered-model", State: "loaded", InFlight: 5}).Error).ToNot(HaveOccurred()) + Expect(registry.SetNodeModelRevision(ctx, mismatched.ID, "filtered-model", 0, "loaded", mismatched.Address, 7, "rev-1", "hash-1")).To(MatchError(ErrStaleModelConfigRevision)) + Expect(db.Create(&NodeModel{ID: "mismatched-revision", NodeID: mismatched.ID, ModelName: "filtered-model", State: "loaded", InFlight: 7, ConfigRevision: "rev-1"}).Error).ToNot(HaveOccurred()) + Expect(registry.SetNodeModelRevision(ctx, unloading.ID, "filtered-model", 0, "unloading", unloading.Address, 11, "rev-2", "hash-2")).To(Succeed()) + + picked, _, err := registry.FindAndLockNodeWithModel(ctx, "filtered-model", nil, nil) + Expect(err).ToNot(HaveOccurred()) + Expect(picked.ID).To(Equal(matching.ID)) + + stats, err := registry.LoadedReplicaStats(ctx, "filtered-model", nil) + Expect(err).ToNot(HaveOccurred()) + Expect(stats).To(HaveLen(1)) + Expect(stats[0].NodeID).To(Equal(matching.ID)) + + count, err := registry.CountLoadedReplicas(ctx, "filtered-model") + Expect(err).ToNot(HaveOccurred()) + Expect(count).To(Equal(int64(1))) + + nodes, err := registry.ListWithExtras(ctx) + Expect(err).ToNot(HaveOccurred()) + byID := make(map[string]NodeWithExtras, len(nodes)) + for _, node := range nodes { + byID[node.ID] = node + } + Expect(byID[matching.ID].ModelCount).To(Equal(1)) + Expect(byID[matching.ID].InFlightCount).To(Equal(4)) + for _, node := range []*BackendNode{empty, mismatched, unloading} { + Expect(byID[node.ID].ModelCount).To(BeZero()) + Expect(byID[node.ID].InFlightCount).To(BeZero()) + } + + // Scheduling aggregates must treat all three ineligible rows as absent. + Expect(db.Model(&BackendNode{}).Where("id = ?", empty.ID).Update("available_vram", 30_000_000_000).Error).ToNot(HaveOccurred()) + Expect(db.Model(&BackendNode{}).Where("id = ?", mismatched.ID).Update("available_vram", 20_000_000_000).Error).ToNot(HaveOccurred()) + Expect(db.Model(&BackendNode{}).Where("id = ?", unloading.ID).Update("available_vram", 10_000_000_000).Error).ToNot(HaveOccurred()) + candidateIDs := []string{matching.ID, empty.ID, mismatched.ID, unloading.ID} + for _, find := range []func() (*BackendNode, error){ + func() (*BackendNode, error) { return registry.FindNodeWithVRAM(ctx, 0) }, + func() (*BackendNode, error) { return registry.FindNodeWithVRAMFromSet(ctx, 0, candidateIDs) }, + func() (*BackendNode, error) { return registry.FindIdleNode(ctx) }, + func() (*BackendNode, error) { return registry.FindIdleNodeFromSet(ctx, candidateIDs) }, + func() (*BackendNode, error) { return registry.FindLeastLoadedNode(ctx) }, + func() (*BackendNode, error) { return registry.FindLeastLoadedNodeFromSet(ctx, candidateIDs) }, + } { + found, err := find() + Expect(err).ToNot(HaveOccurred()) + Expect(found.ID).To(Equal(empty.ID)) + } + + free, err := registry.FindNodesWithFreeSlot(ctx, "filtered-model", candidateIDs) + Expect(err).ToNot(HaveOccurred()) + Expect(free).To(ConsistOf( + HaveField("ID", empty.ID), HaveField("ID", mismatched.ID), HaveField("ID", unloading.ID), + )) + capacity, err := registry.ClusterCapacityForModel(ctx, "filtered-model", candidateIDs) + Expect(err).ToNot(HaveOccurred()) + Expect(capacity).To(Equal(3)) + }) + + It("atomically advances and quarantines stale active replicas", func() { + ctx := context.Background() + node := makeNode("revision-node", "10.0.2.1:50051", 8_000_000_000) + Expect(registry.Register(ctx, node, true)).To(Succeed()) + Expect(registry.AdvanceModelConfigRevision(ctx, "revision-model", "rev-1")).To(BeEmpty()) + Expect(registry.SetNodeModelRevision(ctx, node.ID, "revision-model", 0, "loaded", node.Address, 0, "rev-1", "hash-1")).To(Succeed()) + Expect(registry.UpsertModelLoadInfoRevision(ctx, "revision-model", "llama-cpp", "rev-1", []byte("opts-1"))).To(Succeed()) + + quarantined, err := registry.AdvanceModelConfigRevision(ctx, "revision-model", "rev-2") + Expect(err).ToNot(HaveOccurred()) + Expect(quarantined).To(HaveLen(1)) + Expect(quarantined[0].State).To(Equal("unloading")) + Expect(quarantined[0].CleanupAttempts).To(Equal(0)) + Expect(quarantined[0].CleanupError).To(BeEmpty()) + + revision, err := registry.GetModelConfigRevision(ctx, "revision-model") + Expect(err).ToNot(HaveOccurred()) + Expect(revision).To(Equal("rev-2")) + _, _, _, err = registry.GetModelLoadInfoRevision(ctx, "revision-model") + Expect(err).To(MatchError(gorm.ErrRecordNotFound)) + }) + + It("rejects stale load-info writes without changing matching replay info", func() { + ctx := context.Background() + Expect(registry.AdvanceModelConfigRevision(ctx, "stale-model", "rev-2")).To(BeEmpty()) + Expect(registry.UpsertModelLoadInfoRevision(ctx, "stale-model", "llama-cpp", "rev-2", []byte("good"))).To(Succeed()) + Expect(registry.UpsertModelLoadInfoRevision(ctx, "stale-model", "vllm", "rev-1", []byte("stale"))).To(MatchError(ErrStaleModelConfigRevision)) + + backend, revision, blob, err := registry.GetModelLoadInfoRevision(ctx, "stale-model") + Expect(err).ToNot(HaveOccurred()) + Expect(backend).To(Equal("llama-cpp")) + Expect(revision).To(Equal("rev-2")) + Expect(blob).To(Equal([]byte("good"))) + }) + + It("never replays mismatched or empty load info once current state exists", func() { + ctx := context.Background() + Expect(registry.AdvanceModelConfigRevision(ctx, "replay-filter", "rev-2")).To(BeEmpty()) + for _, revision := range []string{"", "rev-1"} { + Expect(db.Where("model_name = ?", "replay-filter").Delete(&ModelLoadInfo{}).Error).ToNot(HaveOccurred()) + Expect(db.Create(&ModelLoadInfo{ModelName: "replay-filter", BackendType: "llama-cpp", ConfigRevision: revision, ModelOptsBlob: []byte("stale")}).Error).ToNot(HaveOccurred()) + _, _, _, err := registry.GetModelLoadInfoRevision(ctx, "replay-filter") + Expect(err).To(MatchError(gorm.ErrRecordNotFound)) + } + }) + + It("records cleanup failures and lists only due unloading retries", func() { + ctx := context.Background() + node := makeNode("cleanup-node", "10.0.2.2:50051", 8_000_000_000) + Expect(registry.Register(ctx, node, true)).To(Succeed()) + Expect(registry.SetNodeModelRevision(ctx, node.ID, "cleanup-model", 0, "unloading", node.Address, 0, "rev-1", "hash")).To(Succeed()) + now := time.Now() + Expect(registry.RecordModelCleanupFailure(ctx, node.ID, "cleanup-model", 0, "worker unavailable", now.Add(-time.Second))).To(Succeed()) + + retries, err := registry.ListModelCleanupRetries(ctx, now, 10) + Expect(err).ToNot(HaveOccurred()) + Expect(retries).To(HaveLen(1)) + Expect(retries[0].CleanupAttempts).To(Equal(1)) + Expect(retries[0].CleanupError).To(Equal("worker unavailable")) + }) + + It("preserves a replacement registered in the same slot after cleanup was claimed", func() { + ctx := context.Background() + node := makeNode("cleanup-replacement", "10.0.2.20:50051", 8_000_000_000) + Expect(registry.Register(ctx, node, true)).To(Succeed()) + Expect(registry.SetNodeModelRevision(ctx, node.ID, "cleanup-race", 0, "unloading", "10.0.2.20:6001", 0, "rev-old", "hash-old")).To(Succeed()) + + claimed, err := registry.ClaimModelCleanupRetries(ctx, time.Now(), time.Now().Add(time.Minute), 1) + Expect(err).ToNot(HaveOccurred()) + Expect(claimed).To(HaveLen(1)) + + Expect(db.Where("id = ?", claimed[0].ID).Delete(&NodeModel{}).Error).To(Succeed()) + Expect(registry.SetNodeModelRevision(ctx, node.ID, "cleanup-race", 0, "loaded", "10.0.2.20:7001", 0, "rev-new", "hash-new")).To(Succeed()) + + deleted, err := registry.RemoveClaimedModelCleanup(ctx, claimed[0]) + Expect(err).ToNot(HaveOccurred()) + Expect(deleted).To(BeFalse()) + models, err := registry.GetNodeModels(ctx, node.ID) + Expect(err).ToNot(HaveOccurred()) + Expect(models).To(ConsistOf(And( + HaveField("ConfigRevision", "rev-new"), + HaveField("Address", "10.0.2.20:7001"), + HaveField("State", "loaded"), + ))) + }) + + It("preserves revision state and replay info across worker re-registration", func() { + ctx := context.Background() + node := makeNode("revision-reregister", "10.0.2.3:50051", 8_000_000_000) + Expect(registry.Register(ctx, node, true)).To(Succeed()) + Expect(registry.AdvanceModelConfigRevision(ctx, "replay-model", "rev-1")).To(BeEmpty()) + Expect(registry.UpsertModelLoadInfoRevision(ctx, "replay-model", "llama-cpp", "rev-1", []byte("opts"))).To(Succeed()) + Expect(registry.SetNodeModelRevision(ctx, node.ID, "replay-model", 0, "loaded", node.Address, 0, "rev-1", "hash")).To(Succeed()) + + restarted := makeNode("revision-reregister", "10.0.2.3:50052", 8_000_000_000) + Expect(registry.Register(ctx, restarted, true)).To(Succeed()) + models, err := registry.GetNodeModels(ctx, node.ID) + Expect(err).ToNot(HaveOccurred()) + Expect(models).To(BeEmpty()) + Expect(registry.GetModelConfigRevision(ctx, "replay-model")).To(Equal("rev-1")) + _, revision, blob, err := registry.GetModelLoadInfoRevision(ctx, "replay-model") + Expect(err).ToNot(HaveOccurred()) + Expect(revision).To(Equal("rev-1")) + Expect(blob).To(Equal([]byte("opts"))) + }) + }) }) var _ = Describe("ModelScheduling spread + seeding", func() { diff --git a/core/services/nodes/revision_eligibility_test.go b/core/services/nodes/revision_eligibility_test.go new file mode 100644 index 000000000000..96ea5010b704 --- /dev/null +++ b/core/services/nodes/revision_eligibility_test.go @@ -0,0 +1,280 @@ +package nodes + +import ( + "context" + "runtime" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/gorm" + + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/testutil" + pb "github.com/mudler/LocalAI/pkg/grpc/proto" +) + +var _ = Describe("revision eligibility consumers", func() { + var ( + ctx context.Context + db *gorm.DB + registry *NodeRegistry + nodes map[string]*BackendNode + ) + + const modelName = "revision-matrix" + + BeforeEach(func() { + if runtime.GOOS == "darwin" { + Skip("testcontainers requires Docker, not available on macOS CI") + } + ctx = context.Background() + db = testutil.SetupTestDB() + var err error + registry, err = NewNodeRegistry(db) + Expect(err).NotTo(HaveOccurred()) + Expect(registry.AdvanceModelConfigRevision(ctx, modelName, "current")).To(BeEmpty()) + + nodes = map[string]*BackendNode{} + for i, kind := range []string{"current", "empty", "mismatch", "unloading"} { + node := &BackendNode{Name: "revision-" + kind, NodeType: NodeTypeBackend, Address: "10.0.0." + string(rune('1'+i)) + ":50051", AvailableVRAM: uint64(100 + i)} + Expect(registry.Register(ctx, node, true)).To(Succeed()) + nodes[kind] = node + state := "loaded" + revision := kind + if kind == "current" { + revision = "current" + } + if kind == "empty" { + revision = "" + } + if kind == "unloading" { + state = "unloading" + revision = "current" + } + Expect(db.Create(&NodeModel{ + ID: kind, NodeID: node.ID, ModelName: modelName, ReplicaIndex: i, + Address: kind, State: state, ConfigRevision: revision, LastUsed: time.Now().Add(time.Duration(i) * time.Minute), + UpdatedAt: time.Now().Add(-time.Hour), + }).Error).To(Succeed()) + } + }) + + It("establishes the first request revision without allowing a later request to roll it back", func() { + const freshModel = "first-request-revision" + Expect(registry.EstablishModelConfigRevision(ctx, freshModel, "new")).To(Succeed()) + Expect(registry.EstablishModelConfigRevision(ctx, freshModel, "old")).To(MatchError(ErrStaleModelConfigRevision)) + revision, err := registry.GetModelConfigRevision(ctx, freshModel) + Expect(err).NotTo(HaveOccurred()) + Expect(revision).To(Equal("new")) + }) + + It("rejects mixed old-context and old-parallel requests before placement", func() { + router := NewSmartRouter(registry, SmartRouterOptions{}) + for _, opts := range []*pb.ModelOptions{ + {ContextSize: 8192}, + {ContextSize: 100000, Options: []string{"parallel:4"}}, + } { + _, err := router.Route(ctx, modelName, "models/revision.gguf", "llama-cpp", "old", opts, false) + Expect(err).To(MatchError(ContainSubstring("stale model config revision"))) + } + revision, getErr := registry.GetModelConfigRevision(ctx, modelName) + Expect(getErr).NotTo(HaveOccurred()) + Expect(revision).To(Equal("current")) + }) + + DescribeTable("excludes empty, mismatched, and unloading rows after current state exists", + func(query func() []string) { + Expect(query()).To(ConsistOf("current")) + }, + Entry("FindNodesWithModel", func() []string { + got, err := registry.FindNodesWithModel(ctx, modelName) + Expect(err).NotTo(HaveOccurred()) + out := make([]string, 0, len(got)) + for _, node := range got { + out = append(out, node.Name[len("revision-"):]) + } + return out + }), + Entry("ListAllLoadedModels", func() []string { + got, err := registry.ListAllLoadedModels(ctx) + Expect(err).NotTo(HaveOccurred()) + out := make([]string, 0, len(got)) + for _, row := range got { + out = append(out, row.ID) + } + return out + }), + Entry("FindLRUModel", func() []string { + Expect(db.Model(&NodeModel{}).Where("id IN ?", []string{"empty", "mismatch", "unloading"}).Update("node_id", nodes["current"].ID).Error).To(Succeed()) + row, err := registry.FindLRUModel(ctx, nodes["current"].ID) + Expect(err).NotTo(HaveOccurred()) + return []string{row.ID} + }), + Entry("FindGlobalLRUModelWithZeroInFlight", func() []string { + row, err := registry.FindGlobalLRUModelWithZeroInFlight(ctx) + Expect(err).NotTo(HaveOccurred()) + return []string{row.ID} + }), + ) + + It("applies eligibility to replica counts and slot allocation", func() { + // Put stale rows into slots 1 and 2 on the same node. They must not + // consume capacity once a current state exists. + Expect(db.Model(&NodeModel{}).Where("id IN ?", []string{"empty", "mismatch"}).Update("node_id", nodes["current"].ID).Error).To(Succeed()) + count, err := registry.CountReplicasOnNode(ctx, nodes["current"].ID, modelName) + Expect(err).NotTo(HaveOccurred()) + Expect(count).To(Equal(1)) + + idx, err := registry.NextFreeReplicaIndex(ctx, nodes["current"].ID, modelName, 4) + Expect(err).NotTo(HaveOccurred()) + Expect(idx).To(Equal(1)) + }) + + It("GetWithExtras counts only current loaded rows in both per-node queries", func() { + Expect(db.Model(&NodeModel{}).Where("model_name = ?", modelName). + Updates(map[string]any{"node_id": nodes["current"].ID, "in_flight": 7}).Error).To(Succeed()) + + got, err := registry.GetWithExtras(ctx, nodes["current"].ID) + Expect(err).NotTo(HaveOccurred()) + Expect(got.ModelCount).To(Equal(1)) + Expect(got.InFlightCount).To(Equal(7)) + }) + + It("scaleDownIdle selects a current idle row but not stale, empty, or unloading rows", func() { + Expect(db.Model(&NodeModel{}).Where("model_name = ?", modelName). + Updates(map[string]any{"node_id": nodes["current"].ID, "last_used": time.Now().Add(-2 * time.Hour)}).Error).To(Succeed()) + Expect(db.Model(&NodeModel{}).Where("id = ?", "empty").Update("replica_index", 8).Error).To(Succeed()) + Expect(db.Model(&NodeModel{}).Where("id = ?", "mismatch").Update("replica_index", 9).Error).To(Succeed()) + Expect(db.Create(&NodeModel{ + ID: "current-extra", NodeID: nodes["current"].ID, ModelName: modelName, + ReplicaIndex: 4, Address: "current-extra", State: "loaded", + ConfigRevision: "current", LastUsed: time.Now().Add(-time.Hour), + }).Error).To(Succeed()) + + unloader := &fakeUnloader{} + rc := NewReplicaReconciler(ReplicaReconcilerOptions{ + Registry: registry, DB: db, Unloader: unloader, ScaleDownDelay: time.Minute, + }) + rc.scaleDownIdle(ctx, ModelSchedulingConfig{ModelName: modelName}, 2, 1) + + Expect(unloader.unloadCalls).To(ConsistOf(nodes["current"].ID + ":" + modelName)) + var remaining []string + Expect(db.Model(&NodeModel{}).Where("model_name = ?", modelName).Order("id").Pluck("id", &remaining).Error).To(Succeed()) + Expect(remaining).To(ConsistOf("current", "empty", "mismatch", "unloading")) + }) + + It("reconciler busy checks ignore stale idle replicas", func() { + Expect(db.Model(&NodeModel{}).Where("id = ?", "current").Update("in_flight", 1).Error).To(Succeed()) + Expect(db.Model(&NodeModel{}).Where("id IN ?", []string{"empty", "mismatch"}).Update("node_id", nodes["current"].ID).Error).To(Succeed()) + rc := NewReplicaReconciler(ReplicaReconcilerOptions{Registry: registry, DB: db}) + Expect(rc.allReplicasBusy(ctx, modelName)).To(BeTrue()) + }) + + DescribeTable("limits reconciler state queries to eligible rows", + func(run func(*recordingEligibilityProber, *recordingEligibilityLister)) { + prober := &recordingEligibilityProber{} + lister := &recordingEligibilityLister{} + run(prober, lister) + Expect(prober.addresses).NotTo(ContainElements("empty", "mismatch", "unloading")) + Expect(lister.nodeIDs).NotTo(ContainElements(nodes["empty"].ID, nodes["mismatch"].ID, nodes["unloading"].ID)) + }, + Entry("probeLoadedModels", func(prober *recordingEligibilityProber, _ *recordingEligibilityLister) { + rc := NewReplicaReconciler(ReplicaReconcilerOptions{Registry: registry, DB: db, Prober: prober, ProbeStaleAfter: time.Minute}) + rc.probeLoadedModels(ctx) + Expect(prober.addresses).To(ConsistOf("current")) + }), + Entry("sweepLeakedInFlight", func(prober *recordingEligibilityProber, _ *recordingEligibilityLister) { + Expect(db.Model(&NodeModel{}).Where("model_name = ?", modelName).Updates(map[string]any{"in_flight": 1, "last_used": time.Now().Add(-2 * inFlightLeakIdleAfter)}).Error).To(Succeed()) + rc := NewReplicaReconciler(ReplicaReconcilerOptions{Registry: registry, DB: db, Prober: prober}) + rc.sweepLeakedInFlight(ctx) + Expect(prober.addresses).To(ConsistOf("current")) + }), + Entry("reconcileNodeProcesses", func(_ *recordingEligibilityProber, lister *recordingEligibilityLister) { + lister.running = map[string][]messaging.RunningModelInfo{nodes["current"].ID: {{ModelID: modelName, ReplicaIndex: 0}}} + rc := NewReplicaReconciler(ReplicaReconcilerOptions{Registry: registry, DB: db, ProcessLister: lister, ProbeStaleAfter: time.Minute}) + rc.reconcileNodeProcesses(ctx) + Expect(lister.nodeIDs).To(ConsistOf(nodes["current"].ID)) + }), + ) + + It("router eviction minimum-replica count ignores stale rows", func() { + // A scheduling row makes the first OR branch false. With one current + // replica at a minimum of one, the stale rows must not inflate the + // revision-filtered count and make any row evictable. + Expect(db.Create(&ModelSchedulingConfig{ModelName: modelName, MinReplicas: 1, MaxReplicas: 2}).Error).To(Succeed()) + shortCtx, cancel := context.WithTimeout(ctx, 100*time.Millisecond) + defer cancel() + router := NewSmartRouter(registry, SmartRouterOptions{DB: db, Unloader: &fakeUnloader{}}) + _, err := router.evictLRUAndFreeNode(shortCtx) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("context cancelled")) + + // Once a second current replica exists, the count is genuinely above + // the minimum and the oldest eligible current row may be selected. + Expect(db.Create(&NodeModel{ + ID: "current-extra", NodeID: nodes["current"].ID, ModelName: modelName, + ReplicaIndex: 4, Address: "current-extra", State: "loaded", + ConfigRevision: "current", LastUsed: time.Now().Add(time.Minute), + }).Error).To(Succeed()) + unloader := &fakeUnloader{} + router = NewSmartRouter(registry, SmartRouterOptions{DB: db, Unloader: unloader}) + node, err := router.evictLRUAndFreeNode(ctx) + Expect(err).NotTo(HaveOccurred()) + Expect(node.ID).To(Equal(nodes["current"].ID)) + Expect(unloader.unloadCalls).To(ConsistOf(nodes["current"].ID + ":" + modelName)) + for _, id := range []string{"empty", "mismatch", "unloading"} { + var count int64 + Expect(db.Model(&NodeModel{}).Where("id = ?", id).Count(&count).Error).To(Succeed()) + Expect(count).To(Equal(int64(1)), id+" must not be evicted") + } + }) + + It("ordinary worker reaping preserves current state and matching replay info", func() { + Expect(registry.UpsertModelLoadInfoRevision(ctx, modelName, "llama-cpp", "current", []byte("opts"))).To(Succeed()) + lister := &recordingEligibilityLister{} + rc := NewReplicaReconciler(ReplicaReconcilerOptions{Registry: registry, DB: db, ProcessLister: lister, ProbeStaleAfter: time.Minute}) + for range workerMissesBeforeReap { + rc.reconcileNodeProcesses(ctx) + Expect(db.Model(&NodeModel{}).Where("id = ?", "current").Update("updated_at", time.Now().Add(-time.Hour)).Error).To(Succeed()) + } + revision, err := registry.GetModelConfigRevision(ctx, modelName) + Expect(err).NotTo(HaveOccurred()) + Expect(revision).To(Equal("current")) + backend, revision, opts, err := registry.GetModelLoadInfoRevision(ctx, modelName) + Expect(err).NotTo(HaveOccurred()) + Expect(backend).To(Equal("llama-cpp")) + Expect(revision).To(Equal("current")) + Expect(opts).To(Equal([]byte("opts"))) + }) + + It("health/offline reaping preserves current state and matching replay info", func() { + Expect(registry.UpsertModelLoadInfoRevision(ctx, modelName, "llama-cpp", "current", []byte("opts"))).To(Succeed()) + Expect(registry.MarkOffline(ctx, nodes["current"].ID)).To(Succeed()) + revision, err := registry.GetModelConfigRevision(ctx, modelName) + Expect(err).NotTo(HaveOccurred()) + Expect(revision).To(Equal("current")) + backend, revision, opts, err := registry.GetModelLoadInfoRevision(ctx, modelName) + Expect(err).NotTo(HaveOccurred()) + Expect(backend).To(Equal("llama-cpp")) + Expect(revision).To(Equal("current")) + Expect(opts).To(Equal([]byte("opts"))) + }) +}) + +type recordingEligibilityProber struct{ addresses []string } + +func (p *recordingEligibilityProber) Probe(_ context.Context, address string) ProbeOutcome { + p.addresses = append(p.addresses, address) + return ProbeAlive +} + +type recordingEligibilityLister struct { + nodeIDs []string + running map[string][]messaging.RunningModelInfo +} + +func (l *recordingEligibilityLister) ListRunningModels(nodeID string) (*messaging.ModelsRunningReply, error) { + l.nodeIDs = append(l.nodeIDs, nodeID) + return &messaging.ModelsRunningReply{Models: l.running[nodeID]}, nil +} diff --git a/core/services/nodes/router.go b/core/services/nodes/router.go index dbcd23bb8eb7..465b2d44ec9a 100644 --- a/core/services/nodes/router.go +++ b/core/services/nodes/router.go @@ -39,7 +39,10 @@ var companionSuffixes = map[string][]string{ // SmartRouterOptions holds all dependencies for constructing a SmartRouter. // Passing them at construction time eliminates data races from post-creation setters. type SmartRouterOptions struct { - Unloader NodeCommandSender + Unloader NodeCommandSender + // ModelCleanup performs acknowledged exact-process cleanup when a load + // finishes after its configuration revision became stale. + ModelCleanup *ModelCleanupService FileStager FileStager GalleriesJSON string AuthToken string @@ -150,7 +153,8 @@ func ModelLoadCeilingFor(installTimeout, loadTimeout time.Duration) time.Duratio // It uses the ModelRouter interface (backed by NodeRegistry in production) for routing decisions. type SmartRouter struct { registry ModelRouter - unloader NodeCommandSender // optional, for NATS-driven load/unload + unloader NodeCommandSender // optional, for NATS-driven load/unload + modelCleanup *ModelCleanupService fileStager FileStager // optional, for distributed file transfer galleriesJSON string // backend gallery config for dynamic installation clientFactory BackendClientFactory // creates gRPC backend clients @@ -233,6 +237,7 @@ func NewSmartRouter(registry ModelRouter, opts SmartRouterOptions) *SmartRouter return &SmartRouter{ registry: registry, unloader: opts.Unloader, + modelCleanup: opts.ModelCleanup, fileStager: opts.FileStager, galleriesJSON: opts.GalleriesJSON, clientFactory: factory, @@ -323,7 +328,7 @@ func applyNodeHardwareDefaults(opts *pb.ModelOptions, node *BackendNode, backend // scheduleNewModel allocates the replica index internally so the worker's // processKey, port, and the registry row all agree. func (r *SmartRouter) scheduleAndLoad(ctx context.Context, backendType, trackingKey, modelName string, - modelOpts *pb.ModelOptions, parallel bool, initialInFlight int) (*scheduleLoadResult, error) { + configRevision string, modelOpts *pb.ModelOptions, parallel bool, initialInFlight int) (*scheduleLoadResult, error) { node, backendAddr, replicaIndex, err := r.scheduleNewModel(ctx, backendType, trackingKey, modelOpts) if err != nil { @@ -341,8 +346,8 @@ func (r *SmartRouter) scheduleAndLoad(ctx context.Context, backendType, tracking // nothing happening. The row also reserves the replica slot against // concurrent schedulers. Removed on any failure below so a dead load does // not leave a phantom replica. - if err := r.registry.SetNodeModel(ctx, node.ID, trackingKey, replicaIndex, "staging", backendAddr, 0); err != nil { - xlog.Warn("Failed to record staging state", "node", node.Name, "model", trackingKey, "replica", replicaIndex, "error", err) + if err := r.setNodeModelState(ctx, node.ID, trackingKey, replicaIndex, "staging", backendAddr, 0, configRevision, ""); err != nil { + return nil, fmt.Errorf("recording staging state: %w", err) } reportLoadPhase(ctx, LoadJobStateStaging, node, replicaIndex) lifecycleSettled := false @@ -351,6 +356,14 @@ func (r *SmartRouter) scheduleAndLoad(ctx context.Context, backendType, tracking return } cleanupCtx := context.WithoutCancel(ctx) + // An edit may have quarantined this row while staging/loading was in + // flight. Its cleanup intent is durable and must not be erased by the + // ordinary failed-load cleanup path. + if configRevision != "" { + if current, err := r.registry.GetModelConfigRevision(cleanupCtx, trackingKey); err == nil && current != configRevision { + return + } + } if err := r.registry.RemoveNodeModel(cleanupCtx, node.ID, trackingKey, replicaIndex); err != nil { xlog.Warn("Failed to clear lifecycle row after failed load", "node", node.Name, "model", trackingKey, "replica", replicaIndex, "error", err) } @@ -371,6 +384,14 @@ func (r *SmartRouter) scheduleAndLoad(ctx context.Context, backendType, tracking } loadOpts = staged } + effectiveOptionsHash := "" + if loadOpts != nil { + var err error + effectiveOptionsHash, err = config.EffectiveModelOptionsHash(loadOpts) + if err != nil { + return nil, fmt.Errorf("hashing effective model options: %w", err) + } + } client := r.buildClientForAddr(node, backendAddr, parallel) @@ -380,8 +401,8 @@ func (r *SmartRouter) scheduleAndLoad(ctx context.Context, backendType, tracking "payloadBytes", payloadBytes, "loadBudget", loadTimeout) // Staging is done; the checkpoint load on the worker begins. - if err := r.registry.SetNodeModel(ctx, node.ID, trackingKey, replicaIndex, "loading", backendAddr, 0); err != nil { - xlog.Warn("Failed to record loading state", "node", node.Name, "model", trackingKey, "replica", replicaIndex, "error", err) + if err := r.setNodeModelState(ctx, node.ID, trackingKey, replicaIndex, "loading", backendAddr, 0, configRevision, effectiveOptionsHash); err != nil { + return nil, fmt.Errorf("recording loading state: %w", err) } reportLoadPhase(ctx, LoadJobStateLoading, node, replicaIndex) @@ -425,10 +446,14 @@ func (r *SmartRouter) scheduleAndLoad(ctx context.Context, backendType, tracking // Record the model as loaded on this node (specific replica slot). From // here the row is authoritative; the failure-cleanup defer must not touch it. - lifecycleSettled = true - if err := r.registry.SetNodeModel(ctx, node.ID, trackingKey, replicaIndex, "loaded", backendAddr, initialInFlight); err != nil { - xlog.Warn("Failed to record model on node", "node", node.Name, "model", trackingKey, "replica", replicaIndex, "error", err) + if err := r.setNodeModelState(ctx, node.ID, trackingKey, replicaIndex, "loaded", backendAddr, initialInFlight, configRevision, effectiveOptionsHash); err != nil { + if errors.Is(err, ErrStaleModelConfigRevision) { + lifecycleSettled = true + r.cleanupStaleLoad(ctx, node, trackingKey, replicaIndex, backendAddr, configRevision, effectiveOptionsHash) + } + return nil, fmt.Errorf("publishing loaded model: %w", err) } + lifecycleSettled = true // Store load metadata for future replica scale-ups by the reconciler. // Writes both per-replica (NodeModel.model_opts_blob) for backward compat @@ -436,10 +461,10 @@ func (r *SmartRouter) scheduleAndLoad(ctx context.Context, backendType, tracking // every replica row has been removed (Bug-1). if modelOpts != nil { if optsBlob, marshalErr := proto.Marshal(modelOpts); marshalErr == nil { - if storeErr := r.registry.SetNodeModelLoadInfo(ctx, node.ID, trackingKey, replicaIndex, backendType, optsBlob); storeErr != nil { + if storeErr := r.setNodeModelLoadInfo(ctx, node.ID, trackingKey, replicaIndex, backendType, configRevision, optsBlob); storeErr != nil { xlog.Warn("Failed to store model load info", "node", node.Name, "model", trackingKey, "replica", replicaIndex, "error", storeErr) } - if storeErr := r.registry.UpsertModelLoadInfo(ctx, trackingKey, backendType, optsBlob); storeErr != nil { + if storeErr := r.upsertModelLoadInfo(ctx, trackingKey, backendType, configRevision, optsBlob); storeErr != nil { xlog.Warn("Failed to upsert per-model load info", "model", trackingKey, "error", storeErr) } } @@ -448,6 +473,44 @@ func (r *SmartRouter) scheduleAndLoad(ctx context.Context, backendType, tracking return &scheduleLoadResult{Node: node, Client: client, BackendAddr: backendAddr, ReplicaIndex: replicaIndex}, nil } +func (r *SmartRouter) cleanupStaleLoad(ctx context.Context, node *BackendNode, modelName string, replicaIndex int, address, revision, hash string) { + if r.modelCleanup == nil { + xlog.Warn("Stale model load requires exact cleanup", "node", node.Name, "model", modelName, "replica", replicaIndex) + return + } + replica, err := r.registry.GetNodeModel(context.WithoutCancel(ctx), node.ID, modelName, replicaIndex) + if err != nil { + replica = &NodeModel{NodeID: node.ID, ModelName: modelName, ReplicaIndex: replicaIndex, Address: address, State: "unloading", ConfigRevision: revision, EffectiveOptionsHash: hash} + } + r.modelCleanup.Cleanup(context.WithoutCancel(ctx), []NodeModel{*replica}, false) +} + +func (r *SmartRouter) setNodeModelState(ctx context.Context, nodeID, modelName string, replicaIndex int, state, address string, initialInFlight int, revision, hash string) error { + if revision == "" { + if _, err := r.registry.GetModelConfigRevision(ctx, modelName); err == nil { + return ErrStaleModelConfigRevision + } else if !errors.Is(err, gorm.ErrRecordNotFound) { + return err + } + return r.registry.SetNodeModel(ctx, nodeID, modelName, replicaIndex, state, address, initialInFlight) + } + return r.registry.SetNodeModelRevision(ctx, nodeID, modelName, replicaIndex, state, address, initialInFlight, revision, hash) +} + +func (r *SmartRouter) setNodeModelLoadInfo(ctx context.Context, nodeID, modelName string, replicaIndex int, backendType, revision string, blob []byte) error { + if revision == "" { + return r.registry.SetNodeModelLoadInfo(ctx, nodeID, modelName, replicaIndex, backendType, blob) + } + return r.registry.SetNodeModelLoadInfoRevision(ctx, nodeID, modelName, replicaIndex, backendType, revision, blob) +} + +func (r *SmartRouter) upsertModelLoadInfo(ctx context.Context, modelName, backendType, revision string, blob []byte) error { + if revision == "" { + return r.registry.UpsertModelLoadInfo(ctx, modelName, backendType, blob) + } + return r.registry.UpsertModelLoadInfoRevision(ctx, modelName, backendType, revision, blob) +} + // loadAbandonedOnWorker reports whether a failed remote LoadModel left the // worker process still running the load. // @@ -498,7 +561,7 @@ func (r *SmartRouter) reapAbandonedLoad(node *BackendNode, trackingKey string, r // full load sequence (stage files, LoadModel, SetNodeModel) on a new node. func (r *SmartRouter) ScheduleAndLoadModel(ctx context.Context, modelName string, candidateNodeIDs []string) (*BackendNode, error) { // Get load info from an existing replica (stored when Route() first loaded the model) - backendType, optsBlob, err := r.registry.GetModelLoadInfo(ctx, modelName) + backendType, revision, optsBlob, err := r.registry.GetModelLoadInfoRevision(ctx, modelName) if err != nil { // No replica has ever been loaded for this model, so we have no // backend type or model options to replicate. The previous fallback @@ -517,7 +580,7 @@ func (r *SmartRouter) ScheduleAndLoadModel(ctx context.Context, modelName string // 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) + result, err := r.scheduleAndLoad(ctx, backendType, modelName, modelName, revision, &modelOpts, false, 0) if err != nil { return nil, err } @@ -542,12 +605,24 @@ type RouteResult struct { // modelID is the logical model identifier used for DB tracking (e.g. "qwen_qwen3.5-0.8b"). // modelName is the model file path used for gRPC LoadModel (e.g. "llama-cpp/models/Qwen_...gguf"). // When modelID is empty, modelName is used for both purposes (backward compat). -func (r *SmartRouter) Route(ctx context.Context, modelID, modelName, backendType string, modelOpts *pb.ModelOptions, parallel bool) (*RouteResult, error) { +func (r *SmartRouter) Route(ctx context.Context, modelID, modelName, backendType, configRevision string, modelOpts *pb.ModelOptions, parallel bool) (*RouteResult, error) { // Use modelID for DB tracking; fall back to modelName if empty trackingKey := modelID if trackingKey == "" { trackingKey = modelName } + if configRevision != "" { + if err := r.registry.EstablishModelConfigRevision(ctx, trackingKey, configRevision); err != nil { + return nil, fmt.Errorf("establishing config revision for %s: %w", trackingKey, err) + } + } else if _, err := r.registry.GetModelConfigRevision(ctx, trackingKey); err == nil { + return nil, fmt.Errorf("routing %s without a config revision: %w", trackingKey, ErrStaleModelConfigRevision) + } else if !errors.Is(err, gorm.ErrRecordNotFound) { + return nil, fmt.Errorf("reading config revision for %s: %w", trackingKey, err) + } + if modelOpts != nil { + modelOpts = proto.Clone(modelOpts).(*pb.ModelOptions) + } // Fetch the model's scheduling config once: it is immutable for the life of // this request, and resolveSelectorCandidates, buildPreference, and @@ -576,6 +651,7 @@ func (r *SmartRouter) Route(ctx context.Context, modelID, modelName, backendType trackingKey: trackingKey, modelName: modelName, backendType: backendType, + configRevision: configRevision, modelOpts: modelOpts, parallel: parallel, sched: sched, @@ -616,6 +692,7 @@ type routeAttempt struct { trackingKey string modelName string backendType string + configRevision string modelOpts *pb.ModelOptions parallel bool sched *ModelSchedulingConfig @@ -681,7 +758,7 @@ func (r *SmartRouter) tryWarmPath(ctx context.Context, att *routeAttempt) *Route // the replica it landed on. initialInFlight reserves the slot for the calling // request; the job runner passes 0 because it is loading on nobody's behalf. func (r *SmartRouter) coldLoad(ctx context.Context, att *routeAttempt, initialInFlight int) (*RouteResult, error) { - result, err := r.scheduleAndLoad(ctx, att.backendType, att.trackingKey, att.modelName, att.modelOpts, att.parallel, initialInFlight) + result, err := r.scheduleAndLoad(ctx, att.backendType, att.trackingKey, att.modelName, att.configRevision, att.modelOpts, att.parallel, initialInFlight) if err != nil { return nil, err } @@ -1904,12 +1981,14 @@ func (r *SmartRouter) evictLRUAndFreeNode(ctx context.Context) (*BackendNode, er var lru NodeModel err := r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { // Lock the row so no other frontend can evict the same model - if err := tx.Clauses(clause.Locking{Strength: "UPDATE"}). + if err := currentModelRevision(tx.Clauses(clause.Locking{Strength: "UPDATE"})). Joins("JOIN backend_nodes ON backend_nodes.id = node_models.node_id"). Where(`node_models.in_flight = 0 AND node_models.state = ? AND backend_nodes.status = ? AND ( NOT EXISTS (SELECT 1 FROM model_scheduling_configs sc WHERE sc.model_name = node_models.model_name AND (sc.min_replicas > 0 OR sc.max_replicas > 0)) - OR (SELECT COUNT(*) FROM node_models nm2 WHERE nm2.model_name = node_models.model_name AND nm2.state = 'loaded') + OR (SELECT COUNT(*) FROM node_models nm2 WHERE nm2.model_name = node_models.model_name AND nm2.state = 'loaded' + AND (NOT EXISTS (SELECT 1 FROM model_config_states mcs2 WHERE mcs2.model_name = nm2.model_name) + OR nm2.config_revision = (SELECT mcs3.config_revision FROM model_config_states mcs3 WHERE mcs3.model_name = nm2.model_name))) > COALESCE((SELECT sc2.min_replicas FROM model_scheduling_configs sc2 WHERE sc2.model_name = node_models.model_name), 1) )`, "loaded", StatusHealthy). Order("node_models.last_used ASC"). diff --git a/core/services/nodes/router_load_budget_test.go b/core/services/nodes/router_load_budget_test.go index 44779694a3b8..921b19905d4a 100644 --- a/core/services/nodes/router_load_budget_test.go +++ b/core/services/nodes/router_load_budget_test.go @@ -114,8 +114,9 @@ var _ = Describe("size-derived remote LoadModel budget", func() { }) routeFile := func(router *SmartRouter, modelFile string) { - _, err := router.Route(context.Background(), "big-model", "models/big.gguf", "llama-cpp", + _, err := router.Route(context.Background(), "big-model", "models/big.gguf", "llama-cpp", "", &pb.ModelOptions{Model: "models/big.gguf", ModelFile: modelFile}, false) + Expect(err).ToNot(HaveOccurred()) } diff --git a/core/services/nodes/router_load_job_test.go b/core/services/nodes/router_load_job_test.go index f67e1399b5fc..65d1ef939c35 100644 --- a/core/services/nodes/router_load_job_test.go +++ b/core/services/nodes/router_load_job_test.go @@ -74,8 +74,9 @@ var _ = Describe("Route cold-load jobs", func() { first := make(chan error, 1) go func() { defer GinkgoRecover() - _, err := router.Route(context.Background(), "big-model", "models/big.gguf", "llama-cpp", + _, err := router.Route(context.Background(), "big-model", "models/big.gguf", "llama-cpp", "", &pb.ModelOptions{Model: "models/big.gguf"}, false) + first <- err }() @@ -88,8 +89,9 @@ var _ = Describe("Route cold-load jobs", func() { second := make(chan error, 1) go func() { defer GinkgoRecover() - _, err := router.Route(context.Background(), "big-model", "models/big.gguf", "llama-cpp", + _, err := router.Route(context.Background(), "big-model", "models/big.gguf", "llama-cpp", "", &pb.ModelOptions{Model: "models/big.gguf"}, false) + second <- err }() @@ -128,8 +130,9 @@ var _ = Describe("Route cold-load jobs", func() { first := make(chan error, 1) go func() { defer GinkgoRecover() - _, err := router.Route(context.Background(), "doomed", "models/doomed.gguf", "llama-cpp", + _, err := router.Route(context.Background(), "doomed", "models/doomed.gguf", "llama-cpp", "", &pb.ModelOptions{Model: "models/doomed.gguf"}, false) + first <- err }() Eventually(func() *ModelLoadJob { @@ -140,8 +143,9 @@ var _ = Describe("Route cold-load jobs", func() { second := make(chan error, 1) go func() { defer GinkgoRecover() - _, err := router.Route(context.Background(), "doomed", "models/doomed.gguf", "llama-cpp", + _, err := router.Route(context.Background(), "doomed", "models/doomed.gguf", "llama-cpp", "", &pb.ModelOptions{Model: "models/doomed.gguf"}, false) + second <- err }() @@ -165,8 +169,9 @@ var _ = Describe("Route cold-load jobs", func() { go func() { defer GinkgoRecover() - _, _ = router.Route(context.Background(), "detached", "models/detached.gguf", "llama-cpp", + _, _ = router.Route(context.Background(), "detached", "models/detached.gguf", "llama-cpp", "", &pb.ModelOptions{Model: "models/detached.gguf"}, false) + }() Eventually(func() *ModelLoadJob { job, _ := registry.GetLoadJob(context.Background(), "detached") @@ -177,8 +182,9 @@ var _ = Describe("Route cold-load jobs", func() { waiter := make(chan error, 1) go func() { defer GinkgoRecover() - _, err := router.Route(ctx, "detached", "models/detached.gguf", "llama-cpp", + _, err := router.Route(ctx, "detached", "models/detached.gguf", "llama-cpp", "", &pb.ModelOptions{Model: "models/detached.gguf"}, false) + waiter <- err }() time.Sleep(100 * time.Millisecond) @@ -207,8 +213,9 @@ var _ = Describe("Route cold-load jobs", func() { }) start := time.Now() - _, err := router.Route(context.Background(), "slow-model", "models/slow.gguf", "llama-cpp", + _, err := router.Route(context.Background(), "slow-model", "models/slow.gguf", "llama-cpp", "", &pb.ModelOptions{Model: "models/slow.gguf"}, false) + Expect(err).To(HaveOccurred()) Expect(time.Since(start)).To(BeNumerically("<", 10*time.Second)) @@ -239,8 +246,9 @@ var _ = Describe("Route cold-load jobs", func() { done := make(chan error, 1) go func() { defer GinkgoRecover() - _, err := router.Route(context.Background(), "patient", "models/patient.gguf", "llama-cpp", + _, err := router.Route(context.Background(), "patient", "models/patient.gguf", "llama-cpp", "", &pb.ModelOptions{Model: "models/patient.gguf"}, false) + done <- err }() @@ -259,8 +267,9 @@ var _ = Describe("Route cold-load jobs", func() { router := newRouter() go func() { defer GinkgoRecover() - _, _ = router.Route(context.Background(), "beating", "models/beating.gguf", "llama-cpp", + _, _ = router.Route(context.Background(), "beating", "models/beating.gguf", "llama-cpp", "", &pb.ModelOptions{Model: "models/beating.gguf"}, false) + }() var first *ModelLoadJob diff --git a/core/services/nodes/router_load_timeout_test.go b/core/services/nodes/router_load_timeout_test.go index d6ba09fa0b51..295dc35d81fd 100644 --- a/core/services/nodes/router_load_timeout_test.go +++ b/core/services/nodes/router_load_timeout_test.go @@ -72,8 +72,9 @@ var _ = Describe("remote LoadModel deadline", func() { }) route := func(router *SmartRouter) { - _, err := router.Route(context.Background(), "big-model", "models/big.gguf", "llama-cpp", + _, err := router.Route(context.Background(), "big-model", "models/big.gguf", "llama-cpp", "", &pb.ModelOptions{Model: "models/big.gguf"}, false) + Expect(err).ToNot(HaveOccurred()) } diff --git a/core/services/nodes/router_reap_load_test.go b/core/services/nodes/router_reap_load_test.go index 4c8e7cf8bb45..67376c06f535 100644 --- a/core/services/nodes/router_reap_load_test.go +++ b/core/services/nodes/router_reap_load_test.go @@ -80,8 +80,9 @@ var _ = Describe("reaping an abandoned remote load", func() { ModelLoadTimeout: time.Minute, ModelLoadCeiling: time.Hour, }) - _, err := router.Route(context.Background(), "big-model", "models/big.gguf", "llama-cpp", + _, err := router.Route(context.Background(), "big-model", "models/big.gguf", "llama-cpp", "", &pb.ModelOptions{Model: "models/big.gguf"}, false) + return err } diff --git a/core/services/nodes/router_revision_lifecycle_test.go b/core/services/nodes/router_revision_lifecycle_test.go new file mode 100644 index 000000000000..b9b78f790436 --- /dev/null +++ b/core/services/nodes/router_revision_lifecycle_test.go @@ -0,0 +1,287 @@ +package nodes + +import ( + "context" + "errors" + "runtime" + "sync" + "time" + + corebackend "github.com/mudler/LocalAI/core/backend" + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/testutil" + pb "github.com/mudler/LocalAI/pkg/grpc/proto" + "github.com/mudler/LocalAI/pkg/model" + "github.com/mudler/LocalAI/pkg/system" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "google.golang.org/protobuf/proto" + "gorm.io/gorm" + "gorm.io/gorm/clause" +) + +type recordingRevisionStopper struct { + mu sync.Mutex + replicas []NodeModel + err error +} + +func (s *recordingRevisionStopper) StopModelReplica(_ context.Context, _ string, replica NodeModel, _ bool) (messaging.ModelStopReply, error) { + s.mu.Lock() + s.replicas = append(s.replicas, replica) + s.mu.Unlock() + return messaging.ModelStopReply{}, s.err +} + +var _ = Describe("revision-bound load publication", func() { + var ( + ctx context.Context + db *gorm.DB + registry *NodeRegistry + node *BackendNode + backend *stubBackend + unloader *fakeUnloader + ) + + BeforeEach(func() { + if runtime.GOOS == "darwin" { + Skip("testcontainers requires Docker, not available on macOS CI") + } + ctx = context.Background() + db = testutil.SetupTestDB() + var err error + registry, err = NewNodeRegistry(db) + Expect(err).NotTo(HaveOccurred()) + node = &BackendNode{Name: "revision-worker", NodeType: NodeTypeBackend, Address: "10.0.0.1:50051", TotalVRAM: 64_000_000_000, AvailableVRAM: 64_000_000_000} + Expect(registry.Register(ctx, node, true)).To(Succeed()) + backend = &stubBackend{healthResult: true, loadResult: &pb.Result{Success: true}} + unloader = &fakeUnloader{installReply: &messaging.BackendInstallReply{Success: true, Address: "10.0.0.1:9001"}} + }) + + It("quarantines and exactly stops a load that finishes after its revision changes", func() { + const modelName = "edited-while-loading" + Expect(registry.EstablishModelConfigRevision(ctx, modelName, "rev-old")).To(Succeed()) + + entered := make(chan struct{}) + release := make(chan struct{}) + backend.loadHook = func(*pb.ModelOptions) { + close(entered) + <-release + } + stopper := &recordingRevisionStopper{err: errors.New("worker temporarily unreachable")} + router := NewSmartRouter(registry, SmartRouterOptions{ + Unloader: unloader, + ClientFactory: &stubClientFactory{client: backend}, + ModelCleanup: NewModelCleanupService(registry, stopper), + DB: db, + }) + + done := make(chan error, 1) + go func() { + defer GinkgoRecover() + _, err := router.Route(ctx, modelName, "models/edited.gguf", "llama-cpp", "rev-old", &pb.ModelOptions{ContextSize: 8192}, false) + done <- err + }() + Eventually(entered, 5*time.Second).Should(BeClosed()) + + quarantined, err := registry.AdvanceModelConfigRevision(ctx, modelName, "rev-new") + Expect(err).NotTo(HaveOccurred()) + Expect(quarantined).To(HaveLen(1)) + close(release) + + var routeErr error + Eventually(done, 10*time.Second).Should(Receive(&routeErr)) + Expect(routeErr).To(MatchError(ContainSubstring("stale model config revision"))) + Eventually(func() int { + stopper.mu.Lock() + defer stopper.mu.Unlock() + return len(stopper.replicas) + }).Should(Equal(1)) + + var rows []NodeModel + Expect(db.Where("model_name = ?", modelName).Find(&rows).Error).To(Succeed()) + Expect(rows).To(HaveLen(1)) + Expect(rows[0].State).To(Equal("unloading")) + Expect(rows[0].ConfigRevision).To(Equal("rev-old")) + Expect(rows[0].CleanupAttempts).To(Equal(1)) + Expect(rows[0].CleanupNextRetryAt).NotTo(BeNil()) + var replayCount int64 + Expect(db.Model(&ModelLoadInfo{}).Where("model_name = ?", modelName).Count(&replayCount).Error).To(Succeed()) + Expect(replayCount).To(BeZero()) + revision, err := registry.GetModelConfigRevision(ctx, modelName) + Expect(err).NotTo(HaveOccurred()) + Expect(revision).To(Equal("rev-new")) + job, err := registry.GetLoadJob(ctx, modelName) + Expect(err).NotTo(HaveOccurred()) + Expect(job).NotTo(BeNil()) + Expect(job.State).To(Equal(LoadJobStateFailed)) + Expect(job.LastError).To(ContainSubstring("stale model config revision")) + }) + + It("rechecks the revision transactionally when claiming a loaded replica", func() { + const modelName = "claim-race" + Expect(registry.EstablishModelConfigRevision(ctx, modelName, "rev-old")).To(Succeed()) + Expect(registry.SetNodeModelRevision(ctx, node.ID, modelName, 0, "loaded", "10.0.0.1:9001", 0, "rev-old", "hash-old")).To(Succeed()) + + edit := db.Begin() + Expect(edit.Error).NotTo(HaveOccurred()) + var state ModelConfigState + Expect(edit.Clauses(clause.Locking{Strength: "UPDATE"}).Where("model_name = ?", modelName).First(&state).Error).To(Succeed()) + Expect(edit.Model(&ModelConfigState{}).Where("model_name = ?", modelName).Update("config_revision", "rev-new").Error).To(Succeed()) + Expect(edit.Model(&NodeModel{}).Where("model_name = ?", modelName).Update("state", "unloading").Error).To(Succeed()) + + type claimResult struct { + nm *NodeModel + err error + } + claimed := make(chan claimResult, 1) + go func() { + defer GinkgoRecover() + _, nm, err := registry.FindAndLockNodeWithModel(ctx, modelName, nil, nil) + claimed <- claimResult{nm: nm, err: err} + }() + Consistently(claimed, 200*time.Millisecond).ShouldNot(Receive()) + Expect(edit.Commit().Error).To(Succeed()) + + var result claimResult + Eventually(claimed, 5*time.Second).Should(Receive(&result)) + Expect(result.err).To(MatchError(gorm.ErrRecordNotFound)) + Expect(result.nm).To(BeNil()) + var row NodeModel + Expect(db.Where("model_name = ?", modelName).First(&row).Error).To(Succeed()) + Expect(row.InFlight).To(BeZero()) + Expect(row.State).To(Equal("unloading")) + }) + + It("loads changed context and parallel options as a new revision", func() { + const modelName = "changed-options" + stopper := &recordingRevisionStopper{} + router := NewSmartRouter(registry, SmartRouterOptions{ + Unloader: unloader, + ClientFactory: &stubClientFactory{client: backend}, + ModelCleanup: NewModelCleanupService(registry, stopper), + }) + + first, err := router.Route(ctx, modelName, "models/changed.gguf", "llama-cpp", "rev-8k", &pb.ModelOptions{ContextSize: 8192}, false) + Expect(err).NotTo(HaveOccurred()) + first.Release() + quarantined, err := registry.AdvanceModelConfigRevision(ctx, modelName, "rev-100k") + Expect(err).NotTo(HaveOccurred()) + NewModelCleanupService(registry, stopper).Cleanup(ctx, quarantined, false) + + second, err := router.Route(ctx, modelName, "models/changed.gguf", "llama-cpp", "rev-100k", &pb.ModelOptions{ContextSize: 100000, Options: []string{"parallel:4"}}, true) + Expect(err).NotTo(HaveOccurred()) + second.Release() + backend.mu.Lock() + loads := append([]*pb.ModelOptions(nil), backend.loadOpts...) + backend.mu.Unlock() + Expect(loads).To(HaveLen(2)) + Expect(loads[0].ContextSize).To(Equal(int32(8192))) + Expect(loads[1].ContextSize).To(Equal(int32(100000))) + Expect(loads[1].Options).To(ContainElement("parallel:4")) + var loaded NodeModel + Expect(db.Where("model_name = ? AND state = ?", modelName, "loaded").First(&loaded).Error).To(Succeed()) + Expect(loaded.ConfigRevision).To(Equal("rev-100k")) + }) + + It("recovers min replicas only from matching-revision replay information", func() { + const modelName = "matching-replay" + Expect(registry.EstablishModelConfigRevision(ctx, modelName, "rev-current")).To(Succeed()) + current, err := proto.Marshal(&pb.ModelOptions{ContextSize: 100000, Options: []string{"parallel:4"}}) + Expect(err).NotTo(HaveOccurred()) + Expect(registry.UpsertModelLoadInfoRevision(ctx, modelName, "llama-cpp", "rev-current", current)).To(Succeed()) + Expect(registry.SetModelScheduling(ctx, &ModelSchedulingConfig{ModelName: modelName, MinReplicas: 1, MaxReplicas: 1})).To(Succeed()) + + router := NewSmartRouter(registry, SmartRouterOptions{Unloader: unloader, ClientFactory: &stubClientFactory{client: backend}}) + rc := NewReplicaReconciler(ReplicaReconcilerOptions{Registry: registry, Scheduler: router, DB: db}) + rc.reconcileModel(ctx, ModelSchedulingConfig{ModelName: modelName, MinReplicas: 1, MaxReplicas: 1}) + + var loaded NodeModel + Expect(db.Where("model_name = ? AND state = ?", modelName, "loaded").First(&loaded).Error).To(Succeed()) + Expect(loaded.ConfigRevision).To(Equal("rev-current")) + backend.mu.Lock() + defer backend.mu.Unlock() + Expect(backend.loadOpts).To(HaveLen(1)) + Expect(backend.loadOpts[0].ContextSize).To(Equal(int32(100000))) + }) + + It("hashes each replica's post-default effective options on heterogeneous nodes", func() { + const modelName = "heterogeneous-options" + Expect(db.Model(&BackendNode{}).Where("id = ?", node.ID).Updates(map[string]any{ + "gpu_vendor": "NVIDIA", "gpu_compute_capability": "12.0", "max_replicas_per_model": 1, + }).Error).To(Succeed()) + secondNode := &BackendNode{ + Name: "hopper-worker", NodeType: NodeTypeBackend, Address: "10.0.0.2:50051", + GPUVendor: "NVIDIA", GPUComputeCapability: "9.0", TotalVRAM: 16_000_000_000, + AvailableVRAM: 16_000_000_000, MaxReplicasPerModel: 1, + } + Expect(registry.Register(ctx, secondNode, true)).To(Succeed()) + router := NewSmartRouter(registry, SmartRouterOptions{Unloader: unloader, ClientFactory: &stubClientFactory{client: backend}}) + + first, err := router.Route(ctx, modelName, "models/heterogeneous.gguf", "llama-cpp", "rev-one", &pb.ModelOptions{ContextSize: 8192, NBatch: 512}, false) + Expect(err).NotTo(HaveOccurred()) + first.Release() + _, err = router.ScheduleAndLoadModel(ctx, modelName, nil) + Expect(err).NotTo(HaveOccurred()) + + var replicas []NodeModel + Expect(db.Where("model_name = ? AND state = ?", modelName, "loaded").Order("node_id").Find(&replicas).Error).To(Succeed()) + Expect(replicas).To(HaveLen(2)) + Expect(replicas[0].ConfigRevision).To(Equal("rev-one")) + Expect(replicas[1].ConfigRevision).To(Equal("rev-one")) + Expect(replicas[0].EffectiveOptionsHash).NotTo(BeEmpty()) + Expect(replicas[1].EffectiveOptionsHash).NotTo(BeEmpty()) + Expect(replicas[0].EffectiveOptionsHash).NotTo(Equal(replicas[1].EffectiveOptionsHash)) + + backend.mu.Lock() + loads := append([]*pb.ModelOptions(nil), backend.loadOpts...) + backend.mu.Unlock() + Expect(loads).To(HaveLen(2)) + Expect([]int32{loads[0].NBatch, loads[1].NBatch}).To(ConsistOf(int32(2048), int32(512))) + }) + + It("carries one immutable revision from backend options through the loader, adapter, and durable attempt", func() { + contextSize := 10000 + cfg := config.ModelConfig{ + Name: "full-revision-flow", + Backend: "llama-cpp", + LLMConfig: config.LLMConfig{ContextSize: &contextSize}, + } + cfg.Model = "models/full-flow.gguf" + expectedRevision, err := config.ModelConfigRevision(&cfg) + Expect(err).NotTo(HaveOccurred()) + + router := NewSmartRouter(registry, SmartRouterOptions{ + Unloader: unloader, + ClientFactory: &stubClientFactory{client: backend}, + DB: db, + }) + adapter := NewModelRouterAdapter(router) + state := &system.SystemState{} + loader := model.NewModelLoader(state) + loader.SetModelRouter(adapter.AsModelRouter()) + appCfg := &config.ApplicationConfig{Context: ctx, SystemState: state} + options := corebackend.ModelOptions(cfg, appCfg) + + // Mutating the source config after ModelOptions resolved it must not alter + // the revision captured by the durable load attempt. + *cfg.ContextSize = 8192 + client, err := loader.Load(options...) + Expect(err).NotTo(HaveOccurred()) + Expect(client).NotTo(BeNil()) + + revision, err := registry.GetModelConfigRevision(ctx, "full-revision-flow") + Expect(err).NotTo(HaveOccurred()) + Expect(revision).To(Equal(expectedRevision)) + var loaded NodeModel + Expect(db.Where("model_name = ? AND state = ?", "full-revision-flow", "loaded").First(&loaded).Error).To(Succeed()) + Expect(loaded.ConfigRevision).To(Equal(expectedRevision)) + _, replayRevision, replay, err := registry.GetModelLoadInfoRevision(ctx, "full-revision-flow") + Expect(err).NotTo(HaveOccurred()) + Expect(replayRevision).To(Equal(expectedRevision)) + var replayOpts pb.ModelOptions + Expect(proto.Unmarshal(replay, &replayOpts)).To(Succeed()) + Expect(replayOpts.ContextSize).To(Equal(int32(10000))) + }) +}) diff --git a/core/services/nodes/router_staging_context_test.go b/core/services/nodes/router_staging_context_test.go index 6fa892689929..f0b07a7a53e8 100644 --- a/core/services/nodes/router_staging_context_test.go +++ b/core/services/nodes/router_staging_context_test.go @@ -68,7 +68,7 @@ var _ = Describe("Route cold-load staging context", func() { stager.cancelRequest = cancel defer cancel() - result, err := router.Route(ctx, "big-model", filepath.Join("models", "big.gguf"), "llama-cpp", + result, err := router.Route(ctx, "big-model", filepath.Join("models", "big.gguf"), "llama-cpp", "", &pb.ModelOptions{Model: "big.gguf", ModelFile: modelFile}, false) Expect(err).ToNot(HaveOccurred()) diff --git a/core/services/nodes/router_staging_deadline_test.go b/core/services/nodes/router_staging_deadline_test.go index 8b62b6ffbc67..35d1d2ae568c 100644 --- a/core/services/nodes/router_staging_deadline_test.go +++ b/core/services/nodes/router_staging_deadline_test.go @@ -107,8 +107,10 @@ var _ = Describe("cold-load staging deadline", func() { modelFile := filepath.Join(modelDir, "big.gguf") Expect(os.WriteFile(modelFile, []byte("weights"), 0o644)).To(Succeed()) _, err := router.Route(context.Background(), "longcat-video-avatar-1.5", - filepath.Join("models", "big.gguf"), "llama-cpp", + filepath.Join("models", "big.gguf"), "llama-cpp", "", + &pb.ModelOptions{Model: "big.gguf", ModelFile: modelFile}, false) + return err } diff --git a/core/services/nodes/router_test.go b/core/services/nodes/router_test.go index 6e1d86d47929..96db9b93fcb9 100644 --- a/core/services/nodes/router_test.go +++ b/core/services/nodes/router_test.go @@ -18,6 +18,7 @@ import ( grpc "github.com/mudler/LocalAI/pkg/grpc" pb "github.com/mudler/LocalAI/pkg/grpc/proto" ggrpc "google.golang.org/grpc" + "google.golang.org/protobuf/proto" "gorm.io/gorm" ) @@ -261,18 +262,49 @@ func (f *fakeModelRouter) SetNodeModel(_ context.Context, nodeID, modelName stri f.setCalls = append(f.setCalls, fmt.Sprintf("%s:%s:%s:%s", nodeID, modelName, state, address)) return nil } +func (f *fakeModelRouter) SetNodeModelRevision(ctx context.Context, nodeID, modelName string, replicaIndex int, state, address string, initialInFlight int, _, _ string) error { + return f.SetNodeModel(ctx, nodeID, modelName, replicaIndex, state, address, initialInFlight) +} func (f *fakeModelRouter) SetNodeModelLoadInfo(_ context.Context, _, _ string, _ int, _ string, _ []byte) error { return nil } +func (f *fakeModelRouter) SetNodeModelLoadInfoRevision(ctx context.Context, nodeID, modelName string, replicaIndex int, backendType, _ string, optsBlob []byte) error { + return f.SetNodeModelLoadInfo(ctx, nodeID, modelName, replicaIndex, backendType, optsBlob) +} func (f *fakeModelRouter) UpsertModelLoadInfo(_ context.Context, _, _ string, _ []byte) error { return nil } +func (f *fakeModelRouter) UpsertModelLoadInfoRevision(ctx context.Context, modelName, backendType, _ string, optsBlob []byte) error { + return f.UpsertModelLoadInfo(ctx, modelName, backendType, optsBlob) +} func (f *fakeModelRouter) GetModelLoadInfo(_ context.Context, _ string) (string, []byte, error) { return "", nil, fmt.Errorf("not found") } +func (f *fakeModelRouter) GetModelLoadInfoRevision(ctx context.Context, modelName string) (string, string, []byte, error) { + backend, blob, err := f.GetModelLoadInfo(ctx, modelName) + return backend, "", blob, err +} +func (f *fakeModelRouter) AdvanceModelConfigRevision(_ context.Context, _, _ string) ([]NodeModel, error) { + return nil, nil +} +func (f *fakeModelRouter) EstablishModelConfigRevision(_ context.Context, _, _ string) error { + return nil +} +func (f *fakeModelRouter) GetModelConfigRevision(_ context.Context, _ string) (string, error) { + return "", gorm.ErrRecordNotFound +} +func (f *fakeModelRouter) GetNodeModel(_ context.Context, nodeID, modelName string, replicaIndex int) (*NodeModel, error) { + return &NodeModel{NodeID: nodeID, ModelName: modelName, ReplicaIndex: replicaIndex}, nil +} +func (f *fakeModelRouter) RecordModelCleanupFailure(_ context.Context, _, _ string, _ int, _ string, _ time.Time) error { + return nil +} +func (f *fakeModelRouter) ListModelCleanupRetries(_ context.Context, _ time.Time, _ int) ([]NodeModel, error) { + return nil, nil +} func (f *fakeModelRouter) NextFreeReplicaIndex(_ context.Context, _, _ string, _ int) (int, error) { return 0, nil @@ -386,13 +418,23 @@ type stubBackend struct { healthErr error loadResult *pb.Result loadErr error + loadHook func(*pb.ModelOptions) + loadOpts []*pb.ModelOptions + mu sync.Mutex } func (f *stubBackend) HealthCheck(_ context.Context) (bool, error) { return f.healthResult, f.healthErr } -func (f *stubBackend) LoadModel(_ context.Context, _ *pb.ModelOptions, _ ...ggrpc.CallOption) (*pb.Result, error) { +func (f *stubBackend) LoadModel(_ context.Context, opts *pb.ModelOptions, _ ...ggrpc.CallOption) (*pb.Result, error) { + cloned := proto.Clone(opts).(*pb.ModelOptions) + f.mu.Lock() + f.loadOpts = append(f.loadOpts, cloned) + f.mu.Unlock() + if f.loadHook != nil { + f.loadHook(cloned) + } return f.loadResult, f.loadErr } @@ -531,7 +573,7 @@ var _ = Describe("SmartRouter", func() { ClientFactory: factory, }) - result, err := router.Route(context.Background(), "my-model", "models/my-model.gguf", "llama-cpp", nil, false) + result, err := router.Route(context.Background(), "my-model", "models/my-model.gguf", "llama-cpp", "", nil, false) Expect(err).ToNot(HaveOccurred()) Expect(result).ToNot(BeNil()) Expect(result.Node.ID).To(Equal("n1")) @@ -569,7 +611,7 @@ var _ = Describe("SmartRouter", func() { ClientFactory: factory, }) - result, err := router.Route(context.Background(), "some-model", "models/some-model.gguf", "llama-cpp", nil, false) + result, err := router.Route(context.Background(), "some-model", "models/some-model.gguf", "llama-cpp", "", nil, false) Expect(err).ToNot(HaveOccurred()) Expect(result).ToNot(BeNil()) Expect(result.Node.ID).To(Equal("n2")) @@ -598,7 +640,7 @@ var _ = Describe("SmartRouter", func() { // DB is nil — no advisory lock }) - result, err := router.Route(context.Background(), "new-model", "models/new.gguf", "llama-cpp", nil, false) + result, err := router.Route(context.Background(), "new-model", "models/new.gguf", "llama-cpp", "", nil, false) Expect(err).ToNot(HaveOccurred()) Expect(result.Node.ID).To(Equal("n3")) }) @@ -629,8 +671,10 @@ var _ = Describe("SmartRouter", func() { go func() { defer GinkgoRecover() _, err := router.Route(context.Background(), "wedged-model", - "models/wedged.gguf", "llama-cpp", + "models/wedged.gguf", "llama-cpp", "", + &pb.ModelOptions{Model: "models/wedged.gguf"}, false) + done <- err }() @@ -693,7 +737,7 @@ var _ = Describe("SmartRouter", func() { idleNode := &BackendNode{ID: "idle-vram", Name: "idle", Address: "10.0.0.11:50051"} reg.findIdleNode = idleNode - result, err := router.Route(context.Background(), "m1", "models/m1.gguf", "llama-cpp", &pb.ModelOptions{}, false) + result, err := router.Route(context.Background(), "m1", "models/m1.gguf", "llama-cpp", "", &pb.ModelOptions{}, false) Expect(err).ToNot(HaveOccurred()) Expect(result.Node.ID).To(Equal("idle-vram")) }) @@ -708,7 +752,7 @@ var _ = Describe("SmartRouter", func() { ClientFactory: factory, }) - result, err := router.Route(context.Background(), "m2", "models/m2.gguf", "llama-cpp", nil, false) + result, err := router.Route(context.Background(), "m2", "models/m2.gguf", "llama-cpp", "", nil, false) Expect(err).ToNot(HaveOccurred()) Expect(result.Node.ID).To(Equal("idle-1")) }) @@ -724,7 +768,7 @@ var _ = Describe("SmartRouter", func() { ClientFactory: factory, }) - result, err := router.Route(context.Background(), "m3", "models/m3.gguf", "llama-cpp", nil, false) + result, err := router.Route(context.Background(), "m3", "models/m3.gguf", "llama-cpp", "", nil, false) Expect(err).ToNot(HaveOccurred()) Expect(result.Node.ID).To(Equal("ll-1")) }) @@ -740,7 +784,7 @@ var _ = Describe("SmartRouter", func() { // DB is nil — evictLRUAndFreeNode will fail because r.db is nil }) - _, err := router.Route(context.Background(), "m4", "models/m4.gguf", "llama-cpp", nil, false) + _, err := router.Route(context.Background(), "m4", "models/m4.gguf", "llama-cpp", "", nil, false) Expect(err).To(HaveOccurred()) Expect(err.Error()).To(ContainSubstring("no available nodes")) }) @@ -843,7 +887,7 @@ var _ = Describe("SmartRouter", func() { ClientFactory: factory, }) - result, err := router.Route(context.Background(), "selector-model", "models/selector.gguf", "llama-cpp", nil, false) + result, err := router.Route(context.Background(), "selector-model", "models/selector.gguf", "llama-cpp", "", nil, false) Expect(err).ToNot(HaveOccurred()) Expect(result).ToNot(BeNil()) Expect(result.Node.ID).To(Equal("gpu-1")) @@ -862,7 +906,7 @@ var _ = Describe("SmartRouter", func() { ClientFactory: factory, }) - _, err := router.Route(context.Background(), "no-match-model", "models/nomatch.gguf", "llama-cpp", nil, false) + _, err := router.Route(context.Background(), "no-match-model", "models/nomatch.gguf", "llama-cpp", "", nil, false) Expect(err).To(HaveOccurred()) Expect(err.Error()).To(ContainSubstring("no healthy nodes match selector")) }) @@ -877,7 +921,7 @@ var _ = Describe("SmartRouter", func() { ClientFactory: factory, }) - result, err := router.Route(context.Background(), "regular-model", "models/regular.gguf", "llama-cpp", nil, false) + result, err := router.Route(context.Background(), "regular-model", "models/regular.gguf", "llama-cpp", "", nil, false) Expect(err).ToNot(HaveOccurred()) Expect(result).ToNot(BeNil()) Expect(result.Node.ID).To(Equal("regular-1")) @@ -924,7 +968,7 @@ var _ = Describe("SmartRouter", func() { ClientFactory: factory, }) - result, err := router.Route(context.Background(), "sel-model", "models/sel.gguf", "llama-cpp", nil, false) + result, err := router.Route(context.Background(), "sel-model", "models/sel.gguf", "llama-cpp", "", nil, false) Expect(err).ToNot(HaveOccurred()) Expect(result).ToNot(BeNil()) // Should have fallen through to the new node @@ -1305,7 +1349,7 @@ var _ = Describe("SmartRouter prefix-cache routing", func() { router := NewSmartRouter(reg, SmartRouterOptions{Unloader: unloader, ClientFactory: factory}) ctx := distributedhdr.WithPrefixChain(context.Background(), []uint64{1, 2, 3}) - _, err := router.Route(ctx, "m", "models/m.gguf", "llama-cpp", nil, false) + _, err := router.Route(ctx, "m", "models/m.gguf", "llama-cpp", "", nil, false) Expect(err).ToNot(HaveOccurred()) Expect(reg.findAndLockPrefs).ToNot(BeEmpty()) @@ -1327,7 +1371,7 @@ var _ = Describe("SmartRouter prefix-cache routing", func() { }) ctx := distributedhdr.WithPrefixChain(context.Background(), []uint64{1, 2, 3}) - _, err := router.Route(ctx, "m", "models/m.gguf", "llama-cpp", nil, false) + _, err := router.Route(ctx, "m", "models/m.gguf", "llama-cpp", "", nil, false) Expect(err).ToNot(HaveOccurred()) Expect(prov.decideCalls).To(BeNumerically(">=", 1)) @@ -1352,7 +1396,7 @@ var _ = Describe("SmartRouter prefix-cache routing", func() { }) ctx := distributedhdr.WithPrefixChain(context.Background(), []uint64{7, 8, 9}) - _, err := router.Route(ctx, "m", "models/m.gguf", "llama-cpp", nil, false) + _, err := router.Route(ctx, "m", "models/m.gguf", "llama-cpp", "", nil, false) Expect(err).ToNot(HaveOccurred()) // First request landed on X (cold placement on the only candidate) // and observed the prefix there. @@ -1362,7 +1406,7 @@ var _ = Describe("SmartRouter prefix-cache routing", func() { // Second request, same chain: X is now the warm-cache hot match, so // the preference must point at it. - _, err = router.Route(ctx, "m", "models/m.gguf", "llama-cpp", nil, false) + _, err = router.Route(ctx, "m", "models/m.gguf", "llama-cpp", "", nil, false) Expect(err).ToNot(HaveOccurred()) last := reg.findAndLockPrefs[len(reg.findAndLockPrefs)-1] Expect(last).ToNot(BeNil()) @@ -1402,7 +1446,7 @@ var _ = Describe("SmartRouter prefix-cache routing", func() { }) ctx := distributedhdr.WithPrefixChain(context.Background(), []uint64{1, 2, 3}) - _, err := router.Route(ctx, "m", "models/m.gguf", "llama-cpp", nil, false) + _, err := router.Route(ctx, "m", "models/m.gguf", "llama-cpp", "", nil, false) Expect(err).ToNot(HaveOccurred()) pref := reg.findAndLockPrefs[0] @@ -1422,7 +1466,7 @@ var _ = Describe("SmartRouter prefix-cache routing", func() { PrefixConfig: prefixcache.DefaultConfig(), }) - _, err := router.Route(context.Background(), "m", "models/m.gguf", "llama-cpp", nil, false) + _, err := router.Route(context.Background(), "m", "models/m.gguf", "llama-cpp", "", nil, false) Expect(err).ToNot(HaveOccurred()) Expect(prov.decideCalls).To(Equal(0)) Expect(prov.observed).To(BeEmpty()) @@ -1441,7 +1485,7 @@ var _ = Describe("SmartRouter prefix-cache routing", func() { }) ctx := distributedhdr.WithPrefixChain(context.Background(), []uint64{1, 2, 3}) - _, err := router.Route(ctx, "m", "models/m.gguf", "llama-cpp", nil, false) + _, err := router.Route(ctx, "m", "models/m.gguf", "llama-cpp", "", nil, false) Expect(err).ToNot(HaveOccurred()) Expect(prov.decideCalls).To(Equal(0)) Expect(prov.observed).To(BeEmpty()) @@ -1487,7 +1531,7 @@ var _ = Describe("SmartRouter prefix-cache routing", func() { }) ctx := distributedhdr.WithPrefixChain(context.Background(), []uint64{1, 2, 3}) - _, err := router.Route(ctx, "m", "models/m.gguf", "llama-cpp", nil, false) + _, err := router.Route(ctx, "m", "models/m.gguf", "llama-cpp", "", nil, false) Expect(err).ToNot(HaveOccurred()) Expect(pressure.Count("m", time.Now())).To(BeNumerically(">", 0), @@ -1511,7 +1555,7 @@ var _ = Describe("SmartRouter prefix-cache routing", func() { }) ctx := distributedhdr.WithPrefixChain(context.Background(), []uint64{1, 2, 3}) - _, err := router.Route(ctx, "m", "models/m.gguf", "llama-cpp", nil, false) + _, err := router.Route(ctx, "m", "models/m.gguf", "llama-cpp", "", nil, false) Expect(err).ToNot(HaveOccurred()) Expect(pressure.Count("m", time.Now())).To(Equal(0), @@ -1535,7 +1579,7 @@ var _ = Describe("SmartRouter prefix-cache routing", func() { }) ctx := distributedhdr.WithPrefixChain(context.Background(), []uint64{1, 2, 3}) - _, err := router.Route(ctx, "m", "models/m.gguf", "llama-cpp", nil, false) + _, err := router.Route(ctx, "m", "models/m.gguf", "llama-cpp", "", nil, false) Expect(err).ToNot(HaveOccurred()) Expect(pressure.Count("m", time.Now())).To(Equal(0), @@ -1556,7 +1600,7 @@ var _ = Describe("SmartRouter prefix-cache routing", func() { ctx := distributedhdr.WithPrefixChain(context.Background(), []uint64{5, 6}) // Warm the cache: X now holds the prefix. - _, err := router.Route(ctx, "m", "models/m.gguf", "llama-cpp", nil, false) + _, err := router.Route(ctx, "m", "models/m.gguf", "llama-cpp", "", nil, false) Expect(err).ToNot(HaveOccurred()) Expect(idx.Decide("m", []uint64{5, 6}, []prefixcache.ReplicaKey{{NodeID: "X", Replica: 0}}, time.Now()).Hot).To(Equal(prefixcache.ReplicaKey{NodeID: "X", Replica: 0})) diff --git a/core/services/nodes/unloader.go b/core/services/nodes/unloader.go index decfa2b12dbb..8d47d71a622d 100644 --- a/core/services/nodes/unloader.go +++ b/core/services/nodes/unloader.go @@ -81,8 +81,48 @@ var ( _ model.RemoteModelUnloader = (*RemoteUnloaderAdapter)(nil) _ model.RemoteModelContextUnloader = (*RemoteUnloaderAdapter)(nil) _ model.RemoteModelPresenceChecker = (*RemoteUnloaderAdapter)(nil) + _ ExactModelStopper = (*RemoteUnloaderAdapter)(nil) ) +const exactModelStopTimeout = 10 * time.Second + +// StopModelReplica stops only the process represented by replica. Configuration +// cleanup intentionally has no backend.stop fallback: an old worker that does +// not understand this request leaves the quarantine row for a later retry. +func (a *RemoteUnloaderAdapter) StopModelReplica(ctx context.Context, nodeID string, replica NodeModel, force bool) (messaging.ModelStopReply, error) { + if ctx == nil { + ctx = context.Background() + } + ctx, cancel := context.WithTimeout(ctx, exactModelStopTimeout) + defer cancel() + + type result struct { + reply *messaging.ModelStopReply + err error + } + done := make(chan result, 1) + go func() { + reply, err := messaging.RequestJSON[messaging.ModelStopRequest, messaging.ModelStopReply](a.nats, messaging.SubjectNodeModelStop(nodeID), messaging.ModelStopRequest{ + ModelName: replica.ModelName, + ProcessKey: model.BackendProcessKey(replica.ModelName, replica.ReplicaIndex), + ExpectedAddress: replica.Address, + Force: force, + ConfigRevision: replica.ConfigRevision, + }, exactModelStopTimeout) + done <- result{reply: reply, err: err} + }() + + select { + case <-ctx.Done(): + return messaging.ModelStopReply{}, ctx.Err() + case result := <-done: + if result.err != nil { + return messaging.ModelStopReply{}, result.err + } + return *result.reply, nil + } +} + // UnloadRemoteModel finds the node(s) hosting the given model and tells them // to stop their backend process via NATS backend.stop event. // The worker process handles a bounded Free() followed by process termination; diff --git a/core/services/nodes/unloader_test.go b/core/services/nodes/unloader_test.go index f95169eb83e7..8e51aca6cd75 100644 --- a/core/services/nodes/unloader_test.go +++ b/core/services/nodes/unloader_test.go @@ -247,6 +247,26 @@ var _ = Describe("RemoteUnloaderAdapter", func() { }) }) + Describe("StopModelReplica", func() { + It("requests an acknowledged stop for the exact process", func() { + mc.requestReply, _ = json.Marshal(messaging.ModelStopReply{Matched: true, Terminated: true, ProcessKey: "llama#2"}) + replica := NodeModel{ModelName: "llama", ReplicaIndex: 2, Address: "127.0.0.1:5002", ConfigRevision: "rev-1"} + + reply, err := adapter.StopModelReplica(context.Background(), "node-1", replica, true) + Expect(err).NotTo(HaveOccurred()) + Expect(reply.Terminated).To(BeTrue()) + Expect(mc.requestCalls).To(HaveLen(1)) + Expect(mc.requestCalls[0].Subject).To(Equal(messaging.SubjectNodeModelStop("node-1"))) + Expect(mc.requestCalls[0].Timeout).To(BeNumerically(">", 0)) + + var request messaging.ModelStopRequest + Expect(json.Unmarshal(mc.requestCalls[0].Data, &request)).To(Succeed()) + Expect(request).To(Equal(messaging.ModelStopRequest{ + ModelName: "llama", ProcessKey: "llama#2", ExpectedAddress: "127.0.0.1:5002", Force: true, ConfigRevision: "rev-1", + })) + }) + }) + Describe("StopNode", func() { It("publishes to correct subject", func() { Expect(adapter.StopNode("node-abc")).To(Succeed()) diff --git a/core/services/worker/lifecycle.go b/core/services/worker/lifecycle.go index c80e00ea0a02..0c30c01f3b2a 100644 --- a/core/services/worker/lifecycle.go +++ b/core/services/worker/lifecycle.go @@ -42,6 +42,9 @@ func (s *backendSupervisor) subscribeLifecycleEvents() error { if _, err := s.nats.SubscribeReply(messaging.SubjectNodeModelUnload(s.nodeID), s.handleModelUnload); err != nil { return fmt.Errorf("subscribing to model unload events: %w", err) } + if _, err := s.nats.SubscribeReply(messaging.SubjectNodeModelStop(s.nodeID), s.handleModelStop); err != nil { + return fmt.Errorf("subscribing to model stop events: %w", err) + } if _, err := s.nats.SubscribeReply(messaging.SubjectNodeModelDelete(s.nodeID), s.handleModelDelete); err != nil { return fmt.Errorf("subscribing to model delete events: %w", err) } @@ -51,6 +54,15 @@ func (s *backendSupervisor) subscribeLifecycleEvents() error { return nil } +func (s *backendSupervisor) handleModelStop(data []byte, reply func([]byte)) { + var req messaging.ModelStopRequest + if err := json.Unmarshal(data, &req); err != nil { + replyJSON(reply, messaging.ModelStopReply{Error: fmt.Sprintf("invalid request: %v", err)}) + return + } + replyJSON(reply, s.stopModelExact(req)) +} + // handleBackendInstall is the NATS callback for backend.install — install // backend (idempotent: skips download if binary exists on disk) + start gRPC // process (request-reply). diff --git a/core/services/worker/model_stop_test.go b/core/services/worker/model_stop_test.go new file mode 100644 index 000000000000..345b61a30976 --- /dev/null +++ b/core/services/worker/model_stop_test.go @@ -0,0 +1,131 @@ +package worker + +import ( + "context" + "encoding/json" + "errors" + "net" + "sync/atomic" + + "github.com/mudler/LocalAI/core/services/messaging" + process "github.com/mudler/go-processmanager" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + gogrpc "google.golang.org/grpc" + + pb "github.com/mudler/LocalAI/pkg/grpc/proto" +) + +type modelStopBackend struct { + pb.UnimplementedBackendServer + freeCalls atomic.Int32 + freeErr error +} + +func (b *modelStopBackend) Free(context.Context, *pb.HealthMessage) (*pb.Result, error) { + b.freeCalls.Add(1) + return &pb.Result{Success: b.freeErr == nil}, b.freeErr +} + +func startModelStopBackend(backend *modelStopBackend) (string, int, func()) { + lis, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).NotTo(HaveOccurred()) + server := gogrpc.NewServer() + pb.RegisterBackendServer(server, backend) + go func() { _ = server.Serve(lis) }() + return lis.Addr().String(), lis.Addr().(*net.TCPAddr).Port, server.Stop +} + +func startModelStopProcess() *process.Process { + proc := process.New(process.WithTemporaryStateDir(), process.WithName("/bin/sleep"), process.WithArgs("300")) + Expect(proc.Run()).To(Succeed()) + return proc +} + +func requestModelStop(s *backendSupervisor, req messaging.ModelStopRequest) messaging.ModelStopReply { + data, err := json.Marshal(req) + Expect(err).NotTo(HaveOccurred()) + var response []byte + s.handleModelStop(data, func(data []byte) { response = append([]byte(nil), data...) }) + var reply messaging.ModelStopReply + Expect(json.Unmarshal(response, &reply)).To(Succeed()) + return reply +} + +var _ = Describe("Acknowledged exact model stop", func() { + It("stops only the exact process key and releases its port before replying", func() { + backend := &modelStopBackend{} + addr, port, stopServer := startModelStopBackend(backend) + defer stopServer() + proc := startModelStopProcess() + other := &backendProcess{addr: "127.0.0.1:59999", port: 59999} + s := &backendSupervisor{cfg: &Config{}, processes: map[string]*backendProcess{ + "model#0": {proc: proc, addr: addr, port: port}, + "model#1": other, + }} + + reply := requestModelStop(s, messaging.ModelStopRequest{ModelName: "model", ProcessKey: "model#0", ExpectedAddress: addr}) + + Expect(reply).To(Equal(messaging.ModelStopReply{Matched: true, Freed: true, Terminated: true, ProcessKey: "model#0", Address: addr})) + Expect(backend.freeCalls.Load()).To(Equal(int32(1))) + Expect(s.processes).To(HaveKeyWithValue("model#1", other)) + Expect(s.processes).NotTo(HaveKey("model#0")) + Expect(quarantinedPortNumbers(s)).To(ConsistOf(port)) + Expect(proc.Done()).To(BeClosed()) + }) + + It("rejects an address mismatch without stopping anything", func() { + proc := startModelStopProcess() + defer func() { + if pidAlive(proc.CurrentPID()) { + _ = proc.Stop() + } + }() + s := &backendSupervisor{cfg: &Config{}, processes: map[string]*backendProcess{"model#0": {proc: proc, addr: "127.0.0.1:50051", port: 50051}}} + + reply := requestModelStop(s, messaging.ModelStopRequest{ProcessKey: "model#0", ExpectedAddress: "127.0.0.1:50052"}) + + Expect(reply.Matched).To(BeTrue()) + Expect(reply.Terminated).To(BeFalse()) + Expect(reply.Error).To(ContainSubstring("address mismatch")) + Expect(s.processes).To(HaveKey("model#0")) + Expect(pidAlive(proc.CurrentPID())).To(BeTrue()) + }) + + It("treats an absent exact process key as idempotently terminated", func() { + s := &backendSupervisor{cfg: &Config{}, processes: map[string]*backendProcess{}} + reply := requestModelStop(s, messaging.ModelStopRequest{ProcessKey: "missing#0", ExpectedAddress: "127.0.0.1:50051"}) + Expect(reply).To(Equal(messaging.ModelStopReply{Matched: false, Terminated: true, ProcessKey: "missing#0"})) + }) + + It("reports Free failure but still terminates the process", func() { + backend := &modelStopBackend{freeErr: errors.New("free failed")} + addr, port, stopServer := startModelStopBackend(backend) + defer stopServer() + proc := startModelStopProcess() + s := &backendSupervisor{cfg: &Config{}, processes: map[string]*backendProcess{"model#0": {proc: proc, addr: addr, port: port}}} + + reply := requestModelStop(s, messaging.ModelStopRequest{ProcessKey: "model#0", ExpectedAddress: addr}) + + Expect(reply.Matched).To(BeTrue()) + Expect(reply.Freed).To(BeFalse()) + Expect(reply.Terminated).To(BeTrue()) + Expect(reply.Error).To(ContainSubstring("free failed")) + Expect(s.processes).NotTo(HaveKey("model#0")) + }) + + It("skips Free when forced", func() { + backend := &modelStopBackend{} + addr, port, stopServer := startModelStopBackend(backend) + defer stopServer() + proc := startModelStopProcess() + s := &backendSupervisor{cfg: &Config{}, processes: map[string]*backendProcess{"model#0": {proc: proc, addr: addr, port: port}}} + + reply := requestModelStop(s, messaging.ModelStopRequest{ProcessKey: "model#0", ExpectedAddress: addr, Force: true}) + + Expect(reply.Matched).To(BeTrue()) + Expect(reply.Freed).To(BeFalse()) + Expect(reply.Terminated).To(BeTrue()) + Expect(backend.freeCalls.Load()).To(BeZero()) + }) +}) diff --git a/core/services/worker/supervisor.go b/core/services/worker/supervisor.go index 1e01cf44160b..cf95e8b63aaa 100644 --- a/core/services/worker/supervisor.go +++ b/core/services/worker/supervisor.go @@ -843,6 +843,62 @@ func (s *backendSupervisor) stopBackendExact(key string, force bool) error { return s.finishBackendStop(key, bp, stopErr) } +// stopModelExact implements the acknowledged controller-to-worker stop path. +// The address check and stopping reservation are one critical section so a +// stale controller request can never stop a replacement under the same key. +func (s *backendSupervisor) stopModelExact(req messaging.ModelStopRequest) messaging.ModelStopReply { + reply := messaging.ModelStopReply{ProcessKey: req.ProcessKey} + + s.mu.Lock() + bp, ok := s.processes[req.ProcessKey] + if !ok || bp.proc == nil { + s.mu.Unlock() + reply.Terminated = true + return reply + } + reply.Matched = true + reply.Address = bp.addr + if bp.addr != req.ExpectedAddress { + s.mu.Unlock() + reply.Error = fmt.Sprintf("address mismatch for process %s: recorded %q, expected %q", req.ProcessKey, bp.addr, req.ExpectedAddress) + return reply + } + if bp.stopping { + s.mu.Unlock() + reply.Error = fmt.Sprintf("process %s is already stopping", req.ProcessKey) + return reply + } + bp.stopping = true + s.mu.Unlock() + + if !req.Force { + client := grpc.NewClientWithToken(bp.addr, false, nil, false, s.cfg.RegistrationToken) + freeCtx, cancel := context.WithTimeout(context.Background(), workerBackendFreeTimeout) + freeErr := client.Free(freeCtx) + cancel() + if freeErr != nil { + reply.Error = fmt.Sprintf("freeing process %s: %v", req.ProcessKey, freeErr) + } else { + reply.Freed = true + } + } + + stopErr := bp.proc.Stop() + if stopErr == nil { + <-bp.proc.Done() + } + if err := s.finishBackendStop(req.ProcessKey, bp, stopErr); err != nil { + if reply.Error != "" { + reply.Error += "; " + err.Error() + } else { + reply.Error = err.Error() + } + return reply + } + reply.Terminated = true + return reply +} + // beginBackendStop reserves both the process entry and its port while network // cleanup and process termination run without the supervisor mutex. func (s *backendSupervisor) beginBackendStop(key string) *backendProcess { diff --git a/docs/content/features/distributed-mode.md b/docs/content/features/distributed-mode.md index bf5a872619fc..a15fef132755 100644 --- a/docs/content/features/distributed-mode.md +++ b/docs/content/features/distributed-mode.md @@ -486,6 +486,37 @@ Used by the WebUI and admin API consumers. Requires admin authentication. The **Nodes** page in the React WebUI provides a visual overview of all registered workers, their statuses, and loaded models. The page opens with a one-line **cluster pulse** summarising node health and an **attention callout** that surfaces nodes needing action (for example pending approvals). Below that, a roster of **node panels** lists each worker with its inline model chips (no expand click needed), filtered by an **All / Backend / Agent** segmented control. Selecting a panel opens a dedicated **node detail page** at `/app/nodes/:id` with per-node metrics, models, and backend actions. Model scheduling lives on its own **Scheduling** page (separate nav item), not as a tab on the Nodes page. +### Model configuration revisions + +Distributed mode assigns a `config_revision` to each validated model configuration. It hashes the persisted semantic configuration, including fields such as `context_size` and parallel settings. YAML formatting, comments, and map order do not change it. + +The first request for a model establishes its current revision and replay information. The replica reconciler uses only replay information that matches the current revision. This lets `min_replicas` recover after an ordinary worker failure without restoring an old configuration. + +When you save a valid model edit, LocalAI makes replicas from the old revision ineligible immediately. New requests cannot route to those replicas. This rule applies to raw YAML edits, structured patches, renames, disabled models, and changes from another frontend. + +The edit response includes these fields: + +- `config_revision` identifies the saved semantic configuration. +- `pending_cleanup` counts old replicas that still need cleanup when the response returns. + +LocalAI sends an acknowledged stop request for each exact backend process. If a worker or NATS is unreachable, LocalAI keeps the replica in the `unloading` state and retries with durable backoff. The saved edit remains successful while cleanup is pending. + +Workers must support the exact model-stop protocol. Upgrade all workers before you rely on revision cleanup. An older worker cannot acknowledge the request, so its stale replica remains `unloading` until cleanup succeeds or the worker re-registers. + +Worker re-registration removes stale live-replica rows, but it preserves the current model revision and matching replay information. A temporary worker outage therefore does not make an old revision routable. The reconciler can restore the current revision after the worker becomes healthy. + +The responses from `GET /api/node/:id/models` and `GET /api/nodes/:id/models` include these replica fields: + +| Field | Meaning | +|-------|---------| +| `config_revision` | Hash of the persisted semantic model configuration that created the replica. Routable replicas match the current revision. | +| `effective_options_hash` | Hash of the final node-specific load options after defaults and file staging have been applied. Different hashes can be valid on heterogeneous workers when `config_revision` matches. | +| `state` | Replica lifecycle state, such as `staging`, `loading`, `loaded`, or `unloading`. Only eligible `loaded` replicas receive requests. | +| `cleanup_error` | Last exact-stop error. This field appears while cleanup is pending. | +| `cleanup_next_retry_at` | Time of the next durable cleanup attempt. This field appears after a failed attempt. | + +`model.unload` releases model memory inside a running backend. It does not replace the exact process stop that configuration cleanup requires. The `backend.stop` operation remains an administrative backend operation. + ### Per-node VRAM budget Each worker advertises its detected VRAM, and the SmartRouter uses that number when picking a node with enough free memory. You can cap the VRAM a node offers for placement so it never gets scheduled beyond a chosen limit, leaving headroom for other workloads on that machine. @@ -980,6 +1011,15 @@ Notes: **Backend not installing:** - Check the worker logs for `backend.install` events +**Requests still report an old context size or another old load option:** +- Query `/api/nodes/:id/models` for every worker that hosts the model. +- Confirm that every routable replica has `state: loaded` and the same current `config_revision`. +- Treat a different `effective_options_hash` as diagnostic information. Node-specific defaults can cause valid differences. +- Check `cleanup_error` and `cleanup_next_retry_at` on replicas in the `unloading` state. +- Check connectivity to the worker and NATS when cleanup reports a timeout or no responder. +- Upgrade the worker when it does not support the exact model-stop request. +- Stop and restart the stale backend only as an operational recovery action. LocalAI keeps it non-routable while durable cleanup is pending. + **Port conflicts on workers:** - Each model gets its own gRPC process on an incrementing port (50051, 50052, ...) - The HTTP file transfer server runs on the base port - 1 (default: 50050) diff --git a/docs/superpowers/specs/2026-08-21-distributed-model-config-revisions-design.md b/docs/superpowers/specs/2026-08-21-distributed-model-config-revisions-design.md new file mode 100644 index 000000000000..6d1de2e63af9 --- /dev/null +++ b/docs/superpowers/specs/2026-08-21-distributed-model-config-revisions-design.md @@ -0,0 +1,357 @@ +# Distributed Model Configuration Revisions + +## Problem + +Editing a model configuration in a distributed LocalAI deployment can leave the +cluster serving different effective configurations for the same logical model. +The frontend reloads the edited YAML and asks workers to stop the model, but the +existing `backend.stop` message is fire-and-forget. The frontend therefore +removes routing state without knowing whether the worker process stopped. + +Separately, the replica reconciler persists `ModelLoadInfo` independently of +live `NodeModel` rows. This is necessary for restoring `min_replicas` after a +worker failure, but the persisted options currently have no relationship to a +specific revision of the model configuration. After an edit, the reconciler can +restore a replica from options captured before the edit. + +The observed result was one replica serving a context near 100K while another +served the default 8K context and default parallelism. Requests behaved +differently depending on which replica the router selected. The problem is not +specific to `context_size`: any load-time model option can be stale. + +## Goals + +- Make all routable replicas of a logical model belong to the current model + configuration revision. +- Prevent the reconciler and late load jobs from restoring options belonging to + an older revision. +- Remove a model from routing before attempting distributed cleanup. +- Confirm that the exact worker process exited before deleting its registry + row. +- Recover safely when a worker or NATS is temporarily unreachable. +- Apply the same lifecycle to raw YAML edits, structured configuration patches, + renames, disabling, and changes received from peer frontends. +- Preserve the existing ability to restore `min_replicas` after ordinary + worker or backend failure when the model configuration has not changed. +- Expose enough state to diagnose why two replicas have different effective + options. + +## Non-goals + +- Requiring identical hardware-derived options on heterogeneous workers. +- Changing `model.unload`, which remains a memory-release operation. +- Replacing backend administration operations such as backend upgrade, delete, + or stop-all. +- Automatically upgrading workers that do not support the new stop protocol. +- Making arbitrary out-of-band filesystem edits transactional across multiple + machines. Such edits are detected when the model configuration loader next + refreshes the model. + +## Configuration identity + +Each validated model configuration has a `config_revision`. The revision is a +SHA-256 digest of a canonical semantic representation of the validated model +configuration. Formatting, YAML comments, and map ordering do not affect the +revision. Load-time request overrides and node-specific hardware tuning are not +part of this digest. + +Canonicalization must use the typed, validated configuration rather than raw +YAML bytes. The canonical representation includes every field that can affect +model loading or serving. Fields used only to locate the source file or report +runtime status are excluded. The canonical encoder must produce stable field +and map ordering and must distinguish absent values where absence has different +semantics from an explicit zero value. + +The revision is carried with the model options from configuration loading into +the distributed router. It is also persisted in: + +- `ModelConfigState`, keyed by logical model name, as the currently accepted + revision; +- `ModelLoadInfo`, alongside the serialized `pb.ModelOptions` used for future + reconciliation; +- `NodeModel`, identifying the revision used for that live replica. + +Each `NodeModel` also records an `effective_options_hash`, computed from the +fully materialized `pb.ModelOptions` after node-specific hardware defaults and +file-path staging rewrites. This hash is diagnostic only. Two replicas may have +different effective hashes and remain compatible when they share the same +configuration revision. + +Rows created by older versions have an empty revision. They remain usable until +the model's first revision-aware configuration mutation. Once a current +revision is recorded, empty-revision rows are stale and cannot be routed. + +## Registry invariants + +The database is the coordination boundary shared by frontend replicas. + +1. At most one current configuration revision exists per logical model name. +2. A `NodeModel` is routable only when it is in the loaded state and its + `config_revision` equals the current `ModelConfigState` revision. +3. A `ModelLoadInfo` row is reconcilable only when its revision equals the + current `ModelConfigState` revision. +4. A load job may publish `NodeModel` or `ModelLoadInfo` state only when its + captured revision still equals the current revision. +5. Advancing the current revision and quarantining prior-revision replica rows + happen in one database transaction. + +The load-info upsert becomes compare-and-set rather than unconditional +last-write-wins. If the load's revision is no longer current, the upsert returns +a typed stale-revision error. The load is then abandoned and its worker process +is stopped through the exact stop protocol. A late load can therefore neither +be routed nor overwrite current reconciliation options. + +Normal worker death does not change `ModelConfigState` or delete matching +`ModelLoadInfo`; this preserves restart recovery. A configuration mutation +advances `ModelConfigState` and invalidates older load information. + +## Configuration mutation lifecycle + +All model configuration mutation entry points use one model administration +lifecycle service. The structured PATCH endpoint must no longer bypass local +shutdown behavior. + +For an edit that keeps the same logical model name, the service: + +1. Validates and persists the new configuration. +2. Reloads it and computes its semantic revision. +3. In one transaction, records the new current revision, marks every replica + from another or empty revision as `unloading`, and removes or supersedes old + `ModelLoadInfo`. +4. Broadcasts the revision-aware invalidation to peer frontends. +5. Starts cleanup for each quarantined replica using exact `model.stop`. +6. Deletes a replica row only after confirmed process termination or confirmed + absence of that exact process. + +Marking rows `unloading` precedes network calls. A worker that cannot be reached +therefore cannot continue receiving inference traffic through LocalAI even if +its old backend process is still alive. + +The configuration save is durable even if cleanup is incomplete. The endpoint +must not report that saving failed after the new file and revision have +committed. Its response reports that cleanup is pending, and the condition is +also logged and exposed through the existing model/node lifecycle status +surfaces. Subsequent retries finish cleanup. + +For rename, the old identity is quarantined and stopped under its old name. The +new identity receives its own current revision. Old load information is not +copied to the new name. Disable performs the same quarantine and cleanup but +does not permit fresh loads while disabled. Delete follows the existing file +deletion lifecycle after exact process cleanup. + +Peer invalidation events carry the logical model name, operation, and new +revision. Applying an event is idempotent. A peer that already observes that +revision refreshes its in-memory configuration but does not create a second +cleanup generation. + +When the existing configuration watcher detects an out-of-band file change, it +computes the revision after validation and submits the same lifecycle +transition. A parse or validation failure leaves the last accepted revision +current and does not quarantine its replicas. This does not make filesystem +writes atomic, but it ensures a successfully observed external edit cannot +silently bypass revision-aware routing. + +## Exact worker process stop + +A new request/reply NATS operation, `model.stop`, is separate from the existing +ambiguous `backend.stop` operation. + +The request contains: + +```text +model_name +process_key +expected_address +force +config_revision +``` + +`process_key` is the exact supervisor key, including replica index. The +controller derives it from the registry row rather than asking the worker to +resolve a bare backend or model name. `expected_address` prevents a stale row +from stopping an unrelated process after port reuse. `config_revision` is +included for auditability; process key and expected address are the worker-side +identity checks because workers do not own the configuration database. + +The reply contains: + +```text +matched +freed +terminated +process_key +address +error +``` + +The worker verifies that both process key and address identify the same +supervised process. A mismatched address is an error and never stops anything. +An absent process is a successful idempotent outcome with `matched=false` and +`terminated=true` because there is no process left to clean up. + +For a graceful request, the worker performs bounded gRPC `Free()` and then +terminates the supervised process. A `Free()` failure is recorded but does not +prevent termination. A forced request skips `Free()`. The worker replies only +after the process has exited and its supervisor bookkeeping and port ownership +have been updated. + +The existing operations retain their meanings: + +- `model.unload` calls gRPC `Free()` without promising process termination; +- `backend.stop` remains an administration and compatibility operation whose + identifier may be a backend name; +- `model.stop` is the only operation used to confirm configuration-generation + cleanup for an exact replica. + +Sending both `model.unload` and `model.stop` is unnecessary because graceful +`model.stop` already performs bounded `Free()` before termination. + +## Unreachable workers and retry + +An `unloading` replica is never routable. Failed `model.stop` attempts retain +the row with its last error, attempt count, and next retry time. A bounded, +backoff-based cleanup loop retries exact stops. Retries are idempotent and are +claimed through the database so multiple frontend replicas do not concurrently +own the same attempt. + +The existing recovery paths remain backstops: + +- Worker re-registration clears all `NodeModel` rows for that node because a + restarted worker has no surviving supervised backend processes. +- The per-model health monitor removes rows after consecutive unreachable + backend probes. +- Node offline handling prevents scheduling onto a worker with stale + heartbeats. + +Cleanup-row removal through any of these paths fires the existing replica +removal hooks. It does not restore stale `ModelLoadInfo` because only the +current revision is eligible for reconciliation. + +If a worker keeps heartbeating but does not support `model.stop`, the row stays +quarantined and the error clearly identifies an incompatible worker version. +The system favors temporary unavailability over silently serving an obsolete +configuration. Restarting or upgrading that worker lets re-registration or a +subsequent retry complete cleanup. + +## Reconciliation and loading + +The reconciler reads the current revision and matching `ModelLoadInfo` in one +consistent operation. If no matching load information exists, it does not use +an older blob. It records a diagnostic explaining that the model must first be +loaded under its current revision. + +The next inference request builds options from the current configuration, +captures its revision, and performs the normal install, staging, and load +sequence. On success, it transactionally records the replica and current +`ModelLoadInfo`. The reconciler may then restore additional `min_replicas` +using that revision. + +Every scheduling and routing decision rechecks revision eligibility when it +claims a replica. A replica selected immediately before a concurrent edit must +fail the claim after the edit advances the current revision. Existing in-flight +requests may finish; no new request is assigned to the old replica. Graceful +cleanup waits for bounded `Free()` behavior and then terminates it. + +## API and observability + +Model and node lifecycle responses should expose, where replica details are +already returned: + +- current model `config_revision`; +- replica `config_revision`; +- `effective_options_hash`; +- lifecycle state, including `unloading`; +- pending cleanup error and retry time. + +Logs for routing, reconciliation, load completion, stale-load rejection, and +cleanup include model name, replica index, node ID, and abbreviated revision. +No serialized model options or request content is added to logs. + +The Web UI does not require a new workflow. After saving, it may show that the +configuration is saved while one or more old replicas are still being cleaned +up. User-facing distributed-model documentation explains this state and the +requirement to upgrade workers that lack acknowledged `model.stop` support. + +## Rolling upgrades + +Database migrations add nullable revision and cleanup columns so old binaries +can continue reading existing rows. New frontends treat missing revisions as +legacy state according to the compatibility rule above. + +The new NATS subject avoids changing the semantics of `backend.stop` for old +workers. A new frontend receiving no responder for `model.stop` leaves the +replica quarantined and reports the compatibility problem. It must not fall +back to fire-and-forget `backend.stop`, because doing so would recreate the +original false-success failure. + +Deployments should upgrade workers before or together with frontends. Mixed +frontend versions are tolerated at the database level, but old frontends do +not enforce revision-aware routing. Documentation must state that strict +cross-replica consistency is guaranteed only after all frontend replicas run +the revision-aware version. + +## Testing + +All Go tests use Ginkgo and Gomega. + +### Registry tests + +- Advancing a revision and quarantining old replicas is atomic. +- Only loaded replicas matching the current revision are returned for routing. +- Empty legacy revisions become stale after a revision-aware mutation. +- Load-info compare-and-set rejects a late old-revision write. +- Matching load information survives ordinary replica removal and worker + failure. +- Re-registration removes quarantined rows without changing current revision or + matching load information. + +### Router and reconciler tests + +- Given one 8K old-revision replica and one 100K current-revision replica, every + new request routes to the current revision. +- Changing `parallel` produces the same revision transition behavior as changing + `context_size`. +- The reconciler never loads from stale `ModelLoadInfo`. +- A late durable load job cannot publish a stale replica or overwrite current + load information. +- A request racing a configuration edit cannot claim the old generation. +- Heterogeneous effective option hashes remain routable when their + configuration revision matches. + +### Worker protocol tests + +- Exact process key and address stop the intended process and wait for exit. +- An address mismatch stops nothing. +- An already-absent process returns idempotent success. +- Graceful stop attempts bounded `Free()` and still terminates after a failure. +- Forced stop skips `Free()`. +- Replica port ownership and quarantine are updated before replying. + +### Lifecycle tests + +- Raw YAML edit, structured PATCH, rename, disable, and peer application all + advance or apply the expected revision and quarantine old replicas. +- A successful stop deletes the matching row. +- A timeout leaves a non-routable `unloading` row with retry state. +- Retry eventually deletes the row after the worker recovers. +- A worker without `model.stop` support produces a visible compatibility error + and never triggers fire-and-forget fallback. +- Partial cleanup does not roll back an already persisted configuration edit. + +### Live distributed regression + +An integration scenario loads a model on two workers, edits context and +parallel settings, and verifies that no request is routed to an old revision. +After cleanup and reload, every replica reports the current revision. The test +also disconnects one worker during the edit, verifies its replica is +quarantined, reconnects it, and verifies retry or re-registration removes the +stale row. + +## Documentation impact + +The implementation updates the distributed model lifecycle documentation under +`docs/content/` in the same change. It documents revision consistency, +quarantined cleanup state, rolling-upgrade requirements, and why an edited model +may wait for its first request before `min_replicas` can be restored. + +No configuration key or public inference API changes are introduced. diff --git a/pkg/mcp/localaitools/inproc/client.go b/pkg/mcp/localaitools/inproc/client.go index 1d80c30dffc2..48e5af0cb20a 100644 --- a/pkg/mcp/localaitools/inproc/client.go +++ b/pkg/mcp/localaitools/inproc/client.go @@ -663,7 +663,7 @@ func (c *Client) VRAMEstimate(ctx context.Context, req localaitools.VRAMEstimate // ---- State ---- func (c *Client) ToggleModelState(ctx context.Context, name string, action modeladmin.Action) error { - _, err := c.modelAdmin.ToggleState(ctx, name, action, c.ModelLoader) + _, err := c.modelAdmin.ToggleState(ctx, name, action) return err } diff --git a/pkg/model/initializers.go b/pkg/model/initializers.go index 2462973e67f6..e73d45535e07 100644 --- a/pkg/model/initializers.go +++ b/pkg/model/initializers.go @@ -68,7 +68,7 @@ func (ml *ModelLoader) grpcModel(backend string, o *Options) func(string, string ml.mu.Unlock() if router != nil { xlog.Info("Routing model to remote node via ModelRouter", "modelID", modelID, "backend", backend) - return router(o.context, backend, modelID, modelName, modelFile, o.gRPCOptions, o.parallelRequests) + return router(o.context, backend, modelID, modelName, modelFile, o.configRevision, o.gRPCOptions, o.parallelRequests) } uri := ml.GetAllExternalBackends(o)[backend] diff --git a/pkg/model/loader.go b/pkg/model/loader.go index 77b7a0ed65bd..322b11e36c96 100644 --- a/pkg/model/loader.go +++ b/pkg/model/loader.go @@ -59,7 +59,7 @@ type RemoteModelPresenceChecker interface { // instead of starting a local process. When set on the ModelLoader, // grpcModel() will delegate to this function before attempting local loading. type ModelRouter func(ctx context.Context, backend, modelID, modelName, modelFile string, - opts *pb.ModelOptions, parallel bool) (*Model, error) + configRevision string, opts *pb.ModelOptions, parallel bool) (*Model, error) // BackendLoadEvent describes one actual backend load attempt: a backend // process spawn (or remote-address attach) followed by its LoadModel RPC. diff --git a/pkg/model/loader_options.go b/pkg/model/loader_options.go index 9552c3c67830..b6db0a47d5cd 100644 --- a/pkg/model/loader_options.go +++ b/pkg/model/loader_options.go @@ -7,11 +7,12 @@ import ( ) type Options struct { - backendString string - model string - modelFile string - modelID string - context context.Context + backendString string + model string + modelFile string + modelID string + configRevision string + context context.Context gRPCOptions *pb.ModelOptions @@ -27,6 +28,13 @@ type Options struct { modelSizeBytes int64 } +// WithConfigRevision binds a load to the semantic revision of the resolved +// model configuration. The value is copied into Options and therefore remains +// stable even if the configuration loader refreshes while a load is running. +func WithConfigRevision(revision string) Option { + return func(o *Options) { o.configRevision = revision } +} + type Option func(*Options) var EnableParallelRequests = func(o *Options) { diff --git a/pkg/model/loader_options_revision_test.go b/pkg/model/loader_options_revision_test.go new file mode 100644 index 000000000000..b254b9c14ddd --- /dev/null +++ b/pkg/model/loader_options_revision_test.go @@ -0,0 +1,13 @@ +package model + +import ( + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("config revision options", func() { + It("captures the revision supplied by the resolved model config", func() { + opts := NewOptions(WithConfigRevision("revision-a")) + Expect(opts.configRevision).To(Equal("revision-a")) + }) +}) diff --git a/pkg/safefile/read_fifo_unix_test.go b/pkg/safefile/read_fifo_unix_test.go new file mode 100644 index 000000000000..14a18268c736 --- /dev/null +++ b/pkg/safefile/read_fifo_unix_test.go @@ -0,0 +1,22 @@ +//go:build aix || darwin || dragonfly || freebsd || linux || netbsd || openbsd || solaris + +package safefile_test + +import ( + "path/filepath" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "golang.org/x/sys/unix" + + "github.com/mudler/LocalAI/pkg/safefile" +) + +var _ = Describe("ReadRegularAt Unix special files", func() { + It("rejects a FIFO without blocking for a writer", func() { + dir := GinkgoT().TempDir() + Expect(unix.Mkfifo(filepath.Join(dir, "model.yaml"), 0o600)).To(Succeed()) + _, _, err := safefile.ReadRegularAt(dir, "model.yaml") + Expect(err).To(MatchError(ContainSubstring("not a regular file"))) + }) +}) diff --git a/pkg/safefile/read_other.go b/pkg/safefile/read_other.go new file mode 100644 index 000000000000..6120b959fd54 --- /dev/null +++ b/pkg/safefile/read_other.go @@ -0,0 +1,43 @@ +//go:build !aix && !darwin && !dragonfly && !freebsd && !linux && !netbsd && !openbsd && !solaris + +package safefile + +import ( + "fmt" + "io" + "os" + "path/filepath" +) + +// ReadRegularAt is the portable fallback for platforms without openat. +func ReadRegularAt(dir, name string) ([]byte, os.FileMode, error) { + if name == "" || filepath.Base(name) != name { + return nil, 0, fmt.Errorf("%q is not a direct directory entry", name) + } + root, err := os.OpenRoot(dir) + if err != nil { + return nil, 0, err + } + defer func() { _ = root.Close() }() + info, err := root.Lstat(name) + if err != nil { + return nil, 0, err + } + if !info.Mode().IsRegular() { + return nil, 0, fmt.Errorf("%q is not a regular file", name) + } + file, err := root.Open(name) + if err != nil { + return nil, 0, err + } + defer func() { _ = file.Close() }() + openedInfo, err := file.Stat() + if err != nil { + return nil, 0, err + } + if !openedInfo.Mode().IsRegular() || !os.SameFile(info, openedInfo) { + return nil, 0, fmt.Errorf("%q changed while opening", name) + } + data, err := io.ReadAll(file) + return data, openedInfo.Mode().Perm(), err +} diff --git a/pkg/safefile/read_test.go b/pkg/safefile/read_test.go new file mode 100644 index 000000000000..a6992a358d37 --- /dev/null +++ b/pkg/safefile/read_test.go @@ -0,0 +1,38 @@ +package safefile_test + +import ( + "os" + "path/filepath" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/pkg/safefile" +) + +var _ = Describe("ReadRegularAt", func() { + It("reads a regular direct child", func() { + dir := GinkgoT().TempDir() + Expect(os.WriteFile(filepath.Join(dir, "model.yaml"), []byte("safe"), 0o640)).To(Succeed()) + data, mode, err := safefile.ReadRegularAt(dir, "model.yaml") + Expect(err).NotTo(HaveOccurred()) + Expect(data).To(Equal([]byte("safe"))) + Expect(mode).To(Equal(os.FileMode(0o640))) + }) + + It("rejects symbolic links", func() { + dir := GinkgoT().TempDir() + target := filepath.Join(dir, "target") + Expect(os.WriteFile(target, []byte("secret"), 0o600)).To(Succeed()) + Expect(os.Symlink(target, filepath.Join(dir, "model.yaml"))).To(Succeed()) + _, _, err := safefile.ReadRegularAt(dir, "model.yaml") + Expect(err).To(HaveOccurred()) + }) + + It("rejects non-regular files without reading them", func() { + dir := GinkgoT().TempDir() + Expect(os.Mkdir(filepath.Join(dir, "model.yaml"), 0o700)).To(Succeed()) + _, _, err := safefile.ReadRegularAt(dir, "model.yaml") + Expect(err).To(MatchError(ContainSubstring("not a regular file"))) + }) +}) diff --git a/pkg/safefile/read_unix.go b/pkg/safefile/read_unix.go new file mode 100644 index 000000000000..72671024030b --- /dev/null +++ b/pkg/safefile/read_unix.go @@ -0,0 +1,49 @@ +//go:build aix || darwin || dragonfly || freebsd || linux || netbsd || openbsd || solaris + +package safefile + +import ( + "fmt" + "io" + "os" + "path/filepath" + + "golang.org/x/sys/unix" +) + +// ReadRegularAt reads a direct child of dir without following symbolic links. +// Opening relative to a held directory descriptor closes the check/open race. +func ReadRegularAt(dir, name string) ([]byte, os.FileMode, error) { + if name == "" || filepath.Base(name) != name { + return nil, 0, fmt.Errorf("%q is not a direct directory entry", name) + } + dirFD, err := unix.Open(dir, unix.O_RDONLY|unix.O_DIRECTORY|unix.O_CLOEXEC, 0) + if err != nil { + return nil, 0, err + } + defer func() { _ = unix.Close(dirFD) }() + + fd, err := unix.Openat(dirFD, name, unix.O_RDONLY|unix.O_CLOEXEC|unix.O_NOFOLLOW|unix.O_NONBLOCK, 0) + if err != nil { + return nil, 0, err + } + file := os.NewFile(uintptr(fd), name) + if file == nil { + _ = unix.Close(fd) + return nil, 0, fmt.Errorf("open %q: invalid file descriptor", name) + } + defer func() { _ = file.Close() }() + + info, err := file.Stat() + if err != nil { + return nil, 0, err + } + if !info.Mode().IsRegular() { + return nil, 0, fmt.Errorf("%q is not a regular file", name) + } + data, err := io.ReadAll(file) + if err != nil { + return nil, 0, err + } + return data, info.Mode().Perm(), nil +} diff --git a/pkg/safefile/suite_test.go b/pkg/safefile/suite_test.go new file mode 100644 index 000000000000..5cdfca9ff5c6 --- /dev/null +++ b/pkg/safefile/suite_test.go @@ -0,0 +1,13 @@ +package safefile_test + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestSafeFile(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Safe file Suite") +} diff --git a/tests/e2e/distributed/distributed_full_flow_test.go b/tests/e2e/distributed/distributed_full_flow_test.go index 1f36a9625997..5eb9ff44281d 100644 --- a/tests/e2e/distributed/distributed_full_flow_test.go +++ b/tests/e2e/distributed/distributed_full_flow_test.go @@ -284,7 +284,7 @@ var _ = Describe("Full Distributed Inference Flow", Label("Distributed"), func() router := newTestSmartRouter(registry) // The model is not loaded yet, so Route will pick the node and call LoadModel - result, err := router.Route(ctx, "", "test-model", "llama-cpp", &pb.ModelOptions{ + result, err := router.Route(ctx, "", "test-model", "llama-cpp", "", &pb.ModelOptions{ Model: "test-model", }, false) Expect(err).ToNot(HaveOccurred()) @@ -343,7 +343,7 @@ var _ = Describe("Full Distributed Inference Flow", Label("Distributed"), func() // Route should pick node-2 (least loaded) thanks to ORDER BY in_flight ASC router := newTestSmartRouter(registry) - result, err := router.Route(ctx, "", "test-model", "llama-cpp", nil, false) + result, err := router.Route(ctx, "", "test-model", "llama-cpp", "", nil, false) Expect(err).ToNot(HaveOccurred()) Expect(result.Node.Name).To(Equal("node-light")) result.Release() @@ -362,7 +362,7 @@ var _ = Describe("Full Distributed Inference Flow", Label("Distributed"), func() // Route should pick this node and call LoadModel on it router := newTestSmartRouter(registry) - result, err := router.Route(ctx, "", "new-model", "llama-cpp", &pb.ModelOptions{ + result, err := router.Route(ctx, "", "new-model", "llama-cpp", "", &pb.ModelOptions{ Model: "new-model", }, false) Expect(err).ToNot(HaveOccurred()) @@ -427,7 +427,7 @@ var _ = Describe("Full Distributed Inference Flow", Label("Distributed"), func() adapter := nodes.NewModelRouterAdapter(router) // Call adapter.Route() (same signature ModelLoader uses) - m, err := adapter.Route(ctx, "llama-cpp", "test-model-id", "test-model", "", + m, err := adapter.Route(ctx, "llama-cpp", "test-model-id", "test-model", "", "", &pb.ModelOptions{Model: "test-model"}, false) Expect(err).ToNot(HaveOccurred()) Expect(m).ToNot(BeNil()) @@ -489,7 +489,7 @@ var _ = Describe("Full Distributed Inference Flow", Label("Distributed"), func() router := newTestSmartRouter(registry, nodes.SmartRouterOptions{FileStager: stager}) // Route with ModelOptions that have file paths — SmartRouter should stage them - result, err := router.Route(ctx, "", "staged-model", "llama-cpp", &pb.ModelOptions{ + result, err := router.Route(ctx, "", "staged-model", "llama-cpp", "", &pb.ModelOptions{ Model: "staged-model", ModelFile: modelPath, MMProj: mmprojPath, @@ -562,7 +562,7 @@ var _ = Describe("Full Distributed Inference Flow", Label("Distributed"), func() modelPath := filepath.Join(modelDir, "vision.gguf") Expect(os.WriteFile(modelPath, []byte("vision model data"), 0644)).To(Succeed()) - result, err := router.Route(ctx, "", "vision-model", "llama-cpp", &pb.ModelOptions{ + result, err := router.Route(ctx, "", "vision-model", "llama-cpp", "", &pb.ModelOptions{ Model: "vision-model", ModelFile: modelPath, }, false) @@ -660,7 +660,7 @@ var _ = Describe("Full Distributed Inference Flow", Label("Distributed"), func() router := newTestSmartRouter(registry, nodes.SmartRouterOptions{FileStager: stager}) - result, err := router.Route(ctx, "", modelName, backendType, &pb.ModelOptions{ + result, err := router.Route(ctx, "", modelName, backendType, "", &pb.ModelOptions{ Model: modelName, }, false) Expect(err).ToNot(HaveOccurred()) @@ -889,7 +889,7 @@ var _ = Describe("Full Distributed Inference Flow", Label("Distributed"), func() router := newTestSmartRouter(registry, nodes.SmartRouterOptions{FileStager: stager}) // Route with ModelFile pointing to the .onnx file (triggers model staging) - result, err := router.Route(ctx, "voice-it-paola-medium", "it-paola-medium.onnx", "piper", &pb.ModelOptions{ + result, err := router.Route(ctx, "voice-it-paola-medium", "it-paola-medium.onnx", "piper", "", &pb.ModelOptions{ Model: "it-paola-medium.onnx", ModelFile: modelFile, }, false) @@ -973,7 +973,7 @@ var _ = Describe("Full Distributed Inference Flow", Label("Distributed"), func() router := newTestSmartRouter(registry, nodes.SmartRouterOptions{FileStager: stager}) // Route with ModelFile pointing to the .onnx file - result, err := router.Route(ctx, "piper-companion-test", "my-model.onnx", "piper", &pb.ModelOptions{ + result, err := router.Route(ctx, "piper-companion-test", "my-model.onnx", "piper", "", &pb.ModelOptions{ Model: "my-model.onnx", ModelFile: modelFile, }, false) diff --git a/tests/e2e/distributed/model_config_revision_test.go b/tests/e2e/distributed/model_config_revision_test.go new file mode 100644 index 000000000000..548a35eb4b93 --- /dev/null +++ b/tests/e2e/distributed/model_config_revision_test.go @@ -0,0 +1,162 @@ +package distributed_test + +import ( + "context" + "errors" + "sync" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/nodes" + pb "github.com/mudler/LocalAI/pkg/grpc/proto" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + pgdriver "gorm.io/driver/postgres" + "gorm.io/gorm" + "gorm.io/gorm/logger" +) + +type revisionCleanupStopper struct { + mu sync.Mutex + unreachable string + stopped []nodes.NodeModel +} + +func (s *revisionCleanupStopper) StopModelReplica(_ context.Context, nodeID string, replica nodes.NodeModel, _ bool) (messaging.ModelStopReply, error) { + s.mu.Lock() + defer s.mu.Unlock() + s.stopped = append(s.stopped, replica) + if nodeID == s.unreachable { + return messaging.ModelStopReply{}, errors.New("worker unreachable") + } + return messaging.ModelStopReply{ + Matched: true, + Terminated: true, + ProcessKey: replica.ModelName, + Address: replica.Address, + }, nil +} + +var _ = Describe("distributed model configuration revisions", Label("Distributed"), func() { + It("quarantines old replicas cluster-wide and converges after an unreachable worker re-registers", func() { + infra := SetupInfra("localai_model_revision_test") + ctx := infra.Ctx + + db, err := gorm.Open(pgdriver.Open(infra.PGURL), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)}) + Expect(err).NotTo(HaveOccurred()) + frontendA, err := nodes.NewNodeRegistry(db) + Expect(err).NotTo(HaveOccurred()) + // A separate registry instance represents another frontend sharing the + // PostgreSQL coordination boundary. + frontendB, err := nodes.NewNodeRegistry(db.Session(&gorm.Session{NewDB: true})) + Expect(err).NotTo(HaveOccurred()) + + workerA := &nodes.BackendNode{Name: "revision-worker-a", Address: "10.0.0.1:9001"} + workerB := &nodes.BackendNode{Name: "revision-worker-b", Address: "10.0.0.2:9002"} + Expect(frontendA.Register(ctx, workerA, true)).To(Succeed()) + Expect(frontendA.Register(ctx, workerB, true)).To(Succeed()) + + const ( + model = "ornith" + oldRevision = "revision-8k" + newRevision = "revision-100k-parallel-4" + ) + oldOptions := &pb.ModelOptions{ContextSize: 8192, Options: []string{"parallel:1"}} + newOptionsA := &pb.ModelOptions{ContextSize: 100000, Options: []string{"parallel:4", "gpu-layers:80"}} + newOptionsB := &pb.ModelOptions{ContextSize: 100000, Options: []string{"parallel:4", "gpu-layers:60"}} + oldHash, err := config.EffectiveModelOptionsHash(oldOptions) + Expect(err).NotTo(HaveOccurred()) + newHashA, err := config.EffectiveModelOptionsHash(newOptionsA) + Expect(err).NotTo(HaveOccurred()) + newHashB, err := config.EffectiveModelOptionsHash(newOptionsB) + Expect(err).NotTo(HaveOccurred()) + + Expect(frontendA.EstablishModelConfigRevision(ctx, model, oldRevision)).To(Succeed()) + Expect(frontendA.SetNodeModelRevision(ctx, workerA.ID, model, 0, "loaded", workerA.Address, 0, oldRevision, oldHash)).To(Succeed()) + Expect(frontendA.SetNodeModelRevision(ctx, workerB.ID, model, 0, "loaded", workerB.Address, 0, oldRevision, oldHash)).To(Succeed()) + Expect(frontendA.UpsertModelLoadInfoRevision(ctx, model, "llama-cpp", oldRevision, []byte("8k-parallel-1"))).To(Succeed()) + + // Before the edit, either frontend can claim an old-generation replica. + _, claimed, err := frontendB.FindAndLockNodeWithModel(ctx, model, nil, nil) + Expect(err).NotTo(HaveOccurred()) + Expect(claimed.ConfigRevision).To(Equal(oldRevision)) + Expect(frontendB.DecrementInFlight(ctx, claimed.NodeID, model, claimed.ReplicaIndex)).To(Succeed()) + + quarantined, err := frontendA.AdvanceModelConfigRevision(ctx, model, newRevision) + Expect(err).NotTo(HaveOccurred()) + Expect(quarantined).To(HaveLen(2)) + + // The other frontend immediately observes the transaction: no stale or + // unloading replica can receive a request, and a late durable completion + // cannot restore old replay options. + _, claimed, err = frontendB.FindAndLockNodeWithModel(ctx, model, nil, nil) + Expect(err).To(MatchError(gorm.ErrRecordNotFound)) + Expect(claimed).To(BeNil()) + Expect(frontendB.UpsertModelLoadInfoRevision(ctx, model, "llama-cpp", oldRevision, []byte("late-old-load"))).To(MatchError(nodes.ErrStaleModelConfigRevision)) + + // A backend load that began before the edit may finish after the new + // revision is current. Publishing that late completion must fail at the + // shared registry boundary and must not revive the quarantined replica. + Expect(frontendB.SetNodeModelRevision(ctx, workerA.ID, model, 0, "loaded", workerA.Address, 0, oldRevision, oldHash)).To(MatchError(nodes.ErrStaleModelConfigRevision)) + modelsAfterLatePublication, err := frontendA.GetNodeModels(ctx, workerA.ID) + Expect(err).NotTo(HaveOccurred()) + Expect(modelsAfterLatePublication).NotTo(ContainElement(And( + HaveField("ConfigRevision", oldRevision), + HaveField("State", "loaded"), + ))) + + stopper := &revisionCleanupStopper{unreachable: workerB.ID} + pending := nodes.NewModelCleanupService(frontendA, stopper).Cleanup(ctx, quarantined, false) + Expect(pending).To(Equal(1)) + modelsA, err := frontendB.GetNodeModels(ctx, workerA.ID) + Expect(err).NotTo(HaveOccurred()) + Expect(modelsA).To(BeEmpty(), "confirmed exact stop removes only its claimed row") + modelsB, err := frontendB.GetNodeModels(ctx, workerB.ID) + Expect(err).NotTo(HaveOccurred()) + Expect(modelsB).To(HaveLen(1)) + Expect(modelsB[0].State).To(Equal("unloading")) + Expect(modelsB[0].CleanupError).To(ContainSubstring("unreachable")) + Expect(modelsB[0].CleanupNextRetryAt).NotTo(BeNil()) + + stopper.mu.Lock() + stopped := append([]nodes.NodeModel(nil), stopper.stopped...) + stopper.mu.Unlock() + Expect(stopped).To(HaveLen(2)) + Expect(stopped).To(ConsistOf( + HaveField("ConfigRevision", oldRevision), + HaveField("ConfigRevision", oldRevision), + )) + + // Re-registration is the recovery path for the unreachable worker. It + // clears its quarantined process row without rolling back current state. + restartedB := &nodes.BackendNode{Name: workerB.Name, Address: "10.0.0.2:9012"} + Expect(frontendB.Register(ctx, restartedB, true)).To(Succeed()) + modelsB, err = frontendA.GetNodeModels(ctx, restartedB.ID) + Expect(err).NotTo(HaveOccurred()) + Expect(modelsB).To(BeEmpty()) + current, err := frontendB.GetModelConfigRevision(ctx, model) + Expect(err).NotTo(HaveOccurred()) + Expect(current).To(Equal(newRevision)) + + // The model can now converge on the new context and parallel setting. + // Hardware-specific effective options may differ without changing the + // semantic configuration revision. + Expect(frontendA.SetNodeModelRevision(ctx, workerA.ID, model, 0, "loaded", workerA.Address, 0, newRevision, newHashA)).To(Succeed()) + Expect(frontendB.SetNodeModelRevision(ctx, restartedB.ID, model, 0, "loaded", restartedB.Address, 0, newRevision, newHashB)).To(Succeed()) + Expect(frontendA.UpsertModelLoadInfoRevision(ctx, model, "llama-cpp", newRevision, []byte("100k-parallel-4"))).To(Succeed()) + + for range 8 { + _, replica, claimErr := frontendB.FindAndLockNodeWithModel(ctx, model, nil, nil) + Expect(claimErr).NotTo(HaveOccurred()) + Expect(replica.ConfigRevision).To(Equal(newRevision)) + Expect(replica.EffectiveOptionsHash).To(Or(Equal(newHashA), Equal(newHashB))) + Expect(frontendB.DecrementInFlight(ctx, replica.NodeID, model, replica.ReplicaIndex)).To(Succeed()) + } + + backend, replayRevision, replay, err := frontendB.GetModelLoadInfoRevision(ctx, model) + Expect(err).NotTo(HaveOccurred()) + Expect(backend).To(Equal("llama-cpp")) + Expect(replayRevision).To(Equal(newRevision)) + Expect(string(replay)).To(Equal("100k-parallel-4")) + }) +}) diff --git a/tests/e2e/distributed/prefix_cache_routing_test.go b/tests/e2e/distributed/prefix_cache_routing_test.go index 542cd29e0120..9b1e3c117718 100644 --- a/tests/e2e/distributed/prefix_cache_routing_test.go +++ b/tests/e2e/distributed/prefix_cache_routing_test.go @@ -78,7 +78,7 @@ var _ = Describe("Prefix-cache aware routing", Label("Distributed"), func() { routeAndSettle := func(chain []uint64) string { GinkgoHelper() ctx := distributedhdr.WithPrefixChain(context.Background(), chain) - result, err := router.Route(ctx, model, model, "llama-cpp", + result, err := router.Route(ctx, model, model, "llama-cpp", "", &pb.ModelOptions{ModelFile: model}, false) Expect(err).ToNot(HaveOccurred()) Expect(result).ToNot(BeNil()) diff --git a/tests/e2e/distributed/registry_extra_test.go b/tests/e2e/distributed/registry_extra_test.go index 91a84bd67e3b..8c346fa727e8 100644 --- a/tests/e2e/distributed/registry_extra_test.go +++ b/tests/e2e/distributed/registry_extra_test.go @@ -104,6 +104,23 @@ var _ = Describe("NodeRegistry extra methods", Label("Distributed"), func() { }) }) + Context("revision-aware PostgreSQL queries", func() { + It("routes a current replica without adding node_models twice", func() { + ctx := context.Background() + node := &nodes.BackendNode{Name: "revision-query-node", Address: "revision-query:5000"} + Expect(registry.Register(ctx, node, true)).To(Succeed()) + quarantined, err := registry.AdvanceModelConfigRevision(ctx, "revision-query-model", "rev-1") + Expect(err).ToNot(HaveOccurred()) + Expect(quarantined).To(BeEmpty()) + Expect(registry.SetNodeModelRevision(ctx, node.ID, "revision-query-model", 0, "loaded", node.Address, 0, "rev-1", "hash-1")).To(Succeed()) + + found, replica, err := registry.FindAndLockNodeWithModel(ctx, "revision-query-model", nil, nil) + Expect(err).ToNot(HaveOccurred()) + Expect(found.ID).To(Equal(node.ID)) + Expect(replica.ConfigRevision).To(Equal("rev-1")) + }) + }) + Context("FindNodeForModel", func() { It("returns (node, true) when model is loaded on healthy node", func() { node := &nodes.BackendNode{ diff --git a/tests/e2e/distributed/router_tracking_test.go b/tests/e2e/distributed/router_tracking_test.go index fe7cced3def1..9691b31b02b7 100644 --- a/tests/e2e/distributed/router_tracking_test.go +++ b/tests/e2e/distributed/router_tracking_test.go @@ -92,7 +92,7 @@ var _ = Describe("SmartRouter trackingKey", Label("Distributed"), func() { }) It("records model under modelID when modelID is provided", func() { - result, err := router.Route(infra.Ctx, "my-model-id", "path/to/model.gguf", "llama-cpp", + result, err := router.Route(infra.Ctx, "my-model-id", "path/to/model.gguf", "llama-cpp", "", &pb.ModelOptions{ModelFile: "path/to/model.gguf"}, false) Expect(err).ToNot(HaveOccurred()) defer result.Release() @@ -105,7 +105,7 @@ var _ = Describe("SmartRouter trackingKey", Label("Distributed"), func() { }) It("records model under modelName when modelID is empty (backward compat)", func() { - result, err := router.Route(infra.Ctx, "", "legacy/model.bin", "llama-cpp", + result, err := router.Route(infra.Ctx, "", "legacy/model.bin", "llama-cpp", "", &pb.ModelOptions{ModelFile: "legacy/model.bin"}, false) Expect(err).ToNot(HaveOccurred()) defer result.Release() @@ -117,7 +117,7 @@ var _ = Describe("SmartRouter trackingKey", Label("Distributed"), func() { }) It("FindNodesWithModel(modelID) finds node; FindNodesWithModel(modelName) does not", func() { - result, err := router.Route(infra.Ctx, "distinct-id", "distinct/path.gguf", "llama-cpp", + result, err := router.Route(infra.Ctx, "distinct-id", "distinct/path.gguf", "llama-cpp", "", &pb.ModelOptions{ModelFile: "distinct/path.gguf"}, false) Expect(err).ToNot(HaveOccurred()) defer result.Release() @@ -135,7 +135,7 @@ var _ = Describe("SmartRouter trackingKey", Label("Distributed"), func() { It("InFlight tracking increments and decrements via registry", func() { // Route to establish model record - result, err := router.Route(infra.Ctx, "release-model", "release/path.gguf", "llama-cpp", + result, err := router.Route(infra.Ctx, "release-model", "release/path.gguf", "llama-cpp", "", &pb.ModelOptions{ModelFile: "release/path.gguf"}, false) Expect(err).ToNot(HaveOccurred()) defer result.Release() @@ -178,7 +178,7 @@ var _ = Describe("SmartRouter trackingKey", Label("Distributed"), func() { It("clears stale model record when node is unreachable", func() { // First route to establish the model record - result, err := router.Route(infra.Ctx, "stale-check", "stale/path.gguf", "llama-cpp", + result, err := router.Route(infra.Ctx, "stale-check", "stale/path.gguf", "llama-cpp", "", &pb.ModelOptions{ModelFile: "stale/path.gguf"}, false) Expect(err).ToNot(HaveOccurred()) result.Release() @@ -195,7 +195,7 @@ var _ = Describe("SmartRouter trackingKey", Label("Distributed"), func() { // Route again — should detect unreachable node and clear stale record // (it will fall through to FindLeastLoadedNode + backend.install which succeeds, // but the LoadModel gRPC call will fail since the server is down) - _, err = router.Route(infra.Ctx, "stale-check", "stale/path.gguf", "llama-cpp", + _, err = router.Route(infra.Ctx, "stale-check", "stale/path.gguf", "llama-cpp", "", &pb.ModelOptions{ModelFile: "stale/path.gguf"}, false) // Expect an error since the only node is down (LoadModel fails) Expect(err).To(HaveOccurred())