diff --git a/core/config/model_artifact_binding.go b/core/config/model_artifact_binding.go index 416ad3e73c19..752d8158bb35 100644 --- a/core/config/model_artifact_binding.go +++ b/core/config/model_artifact_binding.go @@ -10,7 +10,15 @@ import ( "github.com/mudler/LocalAI/pkg/modelartifacts" ) -func persistArtifactBinding(fileName, modelName string, result modelartifacts.Result) error { +// persistArtifactBinding writes the resolved artifact set back into a model's +// config document. It replaces the whole `artifacts:` list, so the caller must +// pass EVERY artifact the model declares — the primary and all companions — not +// just the one that triggered the write. Persisting only the primary silently +// dropped companions from disk, and on the next controller restart the reloaded +// config had no companion at all: withCompanionArtifactOptions then synthesized +// no companion option and a remote backend fell back to fetching the companion +// repo itself, failing the load (the distributed longcat-video base_model bug). +func persistArtifactBinding(fileName, modelName string, artifacts []modelartifacts.Spec) error { data, err := os.ReadFile(fileName) if err != nil { return err @@ -24,7 +32,7 @@ func persistArtifactBinding(fileName, modelName string, result modelartifacts.Re return err } artifactValue := &yaml.Node{} - encoded, err := yaml.Marshal([]modelartifacts.Spec{result.Spec}) + encoded, err := yaml.Marshal(artifacts) if err != nil { return err } diff --git a/core/config/model_artifact_binding_test.go b/core/config/model_artifact_binding_test.go index c61228d04e65..8e3985008688 100644 --- a/core/config/model_artifact_binding_test.go +++ b/core/config/model_artifact_binding_test.go @@ -25,19 +25,16 @@ var _ = Describe("artifact binding persistence", func() { sibling_only: true parameters: {model: sibling.gguf} `), 0644)).To(Succeed()) - result := modelartifacts.Result{ - RelativePath: ".artifacts/huggingface/0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef/snapshot", - Spec: modelartifacts.Spec{ - Name: "model", Target: "model", - Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo", Revision: "main"}, - Resolved: &modelartifacts.Resolved{ - Endpoint: "https://huggingface.co", - Revision: "0123456789abcdef0123456789abcdef01234567", - CacheKey: "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", - }, + primary := modelartifacts.Spec{ + Name: "model", Target: "model", + Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo", Revision: "main"}, + Resolved: &modelartifacts.Resolved{ + Endpoint: "https://huggingface.co", + Revision: "0123456789abcdef0123456789abcdef01234567", + CacheKey: "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", }, } - Expect(persistArtifactBinding(fileName, "managed", result)).To(Succeed()) + Expect(persistArtifactBinding(fileName, "managed", []modelartifacts.Spec{primary})).To(Succeed()) updated, err := os.ReadFile(fileName) Expect(err).NotTo(HaveOccurred()) Expect(string(updated)).To(ContainSubstring("name: sibling")) @@ -47,4 +44,48 @@ var _ = Describe("artifact binding persistence", func() { Expect(string(updated)).To(ContainSubstring("cache_key: 0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef")) Expect(string(updated)).To(ContainSubstring("revision: 0123456789abcdef0123456789abcdef01234567")) }) + + It("writes back every artifact it is given, primary and companion", func() { + // The binding replaces the whole artifacts list, so a companion is only + // retained if it is passed in. Dropping it here is what lost companions + // on a controller restart (the distributed longcat-video base_model bug). + fileName := filepath.Join(GinkgoT().TempDir(), "models.yaml") + Expect(os.WriteFile(fileName, []byte(` +- name: avatar + backend: longcat-video + artifacts: + - name: model + target: model + source: {type: huggingface, repo: owner/avatar} + - name: base_model + target: companion + source: {type: huggingface, repo: owner/base} + parameters: {model: owner/avatar} +`), 0644)).To(Succeed()) + primaryKey := "1111111111111111111111111111111111111111111111111111111111111111" + companionKey := "2222222222222222222222222222222222222222222222222222222222222222" + resolved := func(repo, key string) modelartifacts.Spec { + return modelartifacts.Spec{ + Name: "x", Target: "companion", + Source: modelartifacts.Source{Type: "huggingface", Repo: repo, Revision: "main"}, + Resolved: &modelartifacts.Resolved{ + Endpoint: "https://huggingface.co", + Revision: "0123456789abcdef0123456789abcdef01234567", + CacheKey: key, + }, + } + } + primary := resolved("owner/avatar", primaryKey) + primary.Name, primary.Target = "model", "model" + companion := resolved("owner/base", companionKey) + companion.Name = "base_model" + + Expect(persistArtifactBinding(fileName, "avatar", []modelartifacts.Spec{primary, companion})).To(Succeed()) + + updated, err := os.ReadFile(fileName) + Expect(err).NotTo(HaveOccurred()) + Expect(string(updated)).To(ContainSubstring("name: base_model")) + Expect(string(updated)).To(ContainSubstring(primaryKey)) + Expect(string(updated)).To(ContainSubstring(companionKey)) + }) }) diff --git a/core/config/model_config_loader.go b/core/config/model_config_loader.go index bd9e93865035..be1fc2695c52 100644 --- a/core/config/model_config_loader.go +++ b/core/config/model_config_loader.go @@ -435,12 +435,15 @@ func (bcl *ModelConfigLoader) PreloadWithContext(ctx context.Context, modelPath bcl.Unlock() continue } - if artifactResult != nil && bindingNeedsPersistence(current, *artifactResult) && current.modelConfigFile != "" { + // Persist the WHOLE resolved artifact set (primary + every companion), + // not just the primary result: writing back only the primary dropped + // companions from disk and lost them on the next restart. + if artifactResult != nil && bindingNeedsPersistence(current, updated.Artifacts) && current.modelConfigFile != "" { modelartifacts.ReportProgress(ctx, modelartifacts.ProgressEvent{ Phase: modelartifacts.PhasePersisting, Artifact: artifactResult.Spec.Name, }) - if err := persistArtifactBinding(current.modelConfigFile, current.Name, *artifactResult); err != nil { + if err := persistArtifactBinding(current.modelConfigFile, current.Name, updated.Artifacts); err != nil { bcl.Unlock() return err } @@ -549,8 +552,14 @@ func (bcl *ModelConfigLoader) preloadOne( return updated, artifactResult, nil } -func bindingNeedsPersistence(current ModelConfig, result modelartifacts.Result) bool { - return len(current.Artifacts) == 0 || !reflect.DeepEqual(current.Artifacts[0], result.Spec) +// bindingNeedsPersistence reports whether the freshly resolved artifact set +// differs from what is currently on the config, and so has to be written back. +// It compares the WHOLE set, not just the primary: a companion that resolved +// for the first time (or changed) must trigger a write even when the primary is +// unchanged, or its resolved state would never reach disk and would be lost on +// the next restart. +func bindingNeedsPersistence(current ModelConfig, resolved []modelartifacts.Spec) bool { + return !reflect.DeepEqual(current.Artifacts, resolved) } func (bcl *ModelConfigLoader) displayPreloadedModel(config ModelConfig) { diff --git a/core/config/model_config_loader_companion_test.go b/core/config/model_config_loader_companion_test.go index 028d7516c8c6..a5b4243a0df4 100644 --- a/core/config/model_config_loader_companion_test.go +++ b/core/config/model_config_loader_companion_test.go @@ -111,6 +111,39 @@ parameters: {model: meituan-longcat/LongCat-Video-Avatar-1.5} Expect(loaded.ModelFileName()).To(ContainSubstring(loaded.Artifacts[0].Resolved.CacheKey)) }) + It("keeps every resolved artifact in the persisted file across a reload", func() { + // Regression for the distributed longcat-video companion loss: a + // controller resolves the primary and companion in memory, but if the + // binding it writes back to disk carries only the primary, the companion + // is gone the moment the process restarts and reloads the file. With no + // companion in the config, withCompanionArtifactOptions synthesizes no + // base_model option, so the remote backend falls back to downloading the + // base repo itself and fails ("base_model must point to a LongCat-Video + // checkpoint"). The persisted document, reloaded fresh, must still name + // the companion. + modelsPath := GinkgoT().TempDir() + configPath := filepath.Join(modelsPath, "avatar.yaml") + Expect(os.WriteFile(configPath, []byte(companionConfig), 0644)).To(Succeed()) + + fake := &companionMaterializer{} + loader := NewModelConfigLoader(modelsPath, WithArtifactMaterializer(fake)) + Expect(loader.LoadModelConfigsFromPath(modelsPath)).To(Succeed()) + Expect(loader.PreloadWithContext(context.Background(), modelsPath)).To(Succeed()) + + // A fresh loader models the restart: it only ever sees what was written + // back to disk, never the in-memory state the first loader held. + reloaded := NewModelConfigLoader(modelsPath, WithArtifactMaterializer(&companionMaterializer{})) + Expect(reloaded.LoadModelConfigsFromPath(modelsPath)).To(Succeed()) + + persisted, found := reloaded.GetModelConfig("avatar") + Expect(found).To(BeTrue()) + Expect(persisted.Artifacts).To(HaveLen(2)) + Expect(persisted.Artifacts[1].Name).To(Equal("base_model")) + Expect(persisted.Artifacts[1].Target).To(Equal(modelartifacts.TargetCompanion)) + Expect(persisted.Artifacts[1].Resolved).ToNot(BeNil()) + Expect(persisted.Artifacts[1].Resolved.CacheKey).ToNot(BeEmpty()) + }) + It("fails the load when an explicitly declared companion cannot be acquired", func() { // Explicit artifacts are all-or-nothing: a config that names a companion // is asserting the backend needs it, so silently loading without it