diff --git a/backend/python/longcat-video/backend.py b/backend/python/longcat-video/backend.py index 96e2617452e3..0b54dd71f5c9 100755 --- a/backend/python/longcat-video/backend.py +++ b/backend/python/longcat-video/backend.py @@ -109,6 +109,7 @@ def __init__(self): self.device_index = 0 self.cp_split_hw = None self._dist_store_dir = None + self.model_path = "" def Health(self, request, context): return backend_pb2.Reply(message=b"OK") @@ -118,6 +119,11 @@ def LoadModel(self, request, context): if request.ModelFile and os.path.isdir(request.ModelFile): model = request.ModelFile + # A managed companion snapshot (see base_model) is passed as a path + # relative to the models directory, which is ModelPath here and a + # different directory on a remote worker once staging has rewritten it. + self.model_path = getattr(request, "ModelPath", "") or "" + model_kind = classify_model(model) if model_kind is None: return self._fail( @@ -376,6 +382,26 @@ def _ensure_distributed(self): ) self.cp_split_hw = self.context_parallel_util.get_optimal_split(1) + def _resolve_option_path(self, value): + """Resolve an option that may name a local directory or a HF repo id. + + A managed companion artifact is handed to us relative to the models + directory, so it only becomes a real path once joined with ModelPath; + that indirection is what lets the same config work on a remote worker, + where staging puts the snapshot somewhere else entirely. Anything that + is already an absolute directory, or is not a path at all (a plain repo + id), is passed through untouched for snapshot_download to handle. + """ + if not value: + return value + candidate = normalize_model_source(str(value)) + if os.path.isabs(candidate) or not self.model_path: + return value + joined = os.path.join(self.model_path, candidate) + if os.path.isdir(joined): + return joined + return value + def _resolve_checkpoint(self, model, patterns): source = normalize_model_source(model) if os.path.isdir(source): @@ -440,13 +466,7 @@ def _load_avatar_model(self, model): avatar_patterns.append(f"{model_subfolder}/**") checkpoint = self._resolve_checkpoint(model, avatar_patterns) - base_model = self.options.get("base_model") - if not base_model and os.path.isdir(normalize_model_source(model)): - sibling = os.path.join( - os.path.dirname(normalize_model_source(model)), "LongCat-Video" - ) - if os.path.isdir(sibling): - base_model = sibling + base_model = self._resolve_option_path(self.options.get("base_model")) base_model = base_model or BASE_MODEL_ID if classify_model(str(base_model)) != MODEL_KIND_BASE: raise ValueError("base_model must point to a LongCat-Video checkpoint") diff --git a/core/backend/companion_options_test.go b/core/backend/companion_options_test.go new file mode 100644 index 000000000000..c4b6f3953a66 --- /dev/null +++ b/core/backend/companion_options_test.go @@ -0,0 +1,123 @@ +package backend + +import ( + "strings" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/pkg/modelartifacts" +) + +// A companion snapshot only becomes useful once the backend can find it, and +// its location is a content-addressed cache key that is unknowable until the +// artifact resolves. A static gallery override therefore cannot carry it. The +// path is instead synthesized into ModelOptions.Options at load time, under the +// companion's own name, using the same key:value convention backends already +// parse for options like attention_backend. +var _ = Describe("companion artifact backend options", func() { + const ( + primaryKey = "1111111111111111111111111111111111111111111111111111111111111111" + companionKey = "2222222222222222222222222222222222222222222222222222222222222222" + ) + + resolved := func(key string) *modelartifacts.Resolved { + return &modelartifacts.Resolved{ + Endpoint: "https://huggingface.co", + Revision: "0123456789abcdef0123456789abcdef01234567", + CacheKey: key, + } + } + + threads := 1 + configWithCompanion := func(options ...string) config.ModelConfig { + return config.ModelConfig{ + Backend: "longcat-video", + Options: options, + Threads: &threads, + Artifacts: []modelartifacts.Spec{ + { + Name: modelartifacts.TargetModel, Target: modelartifacts.TargetModel, + Source: modelartifacts.Source{Type: "huggingface", Repo: "meituan-longcat/LongCat-Video-Avatar-1.5", Revision: "main"}, + Resolved: resolved(primaryKey), + }, + { + Name: "base_model", Target: modelartifacts.TargetCompanion, + Source: modelartifacts.Source{Type: "huggingface", Repo: "meituan-longcat/LongCat-Video", Revision: "main"}, + Resolved: resolved(companionKey), + }, + }, + } + } + + optionValue := func(options []string, key string) (string, bool) { + for _, option := range options { + name, value, found := strings.Cut(option, ":") + if found && name == key { + return value, true + } + } + return "", false + } + + It("exposes a resolved companion snapshot as an option named after the artifact", func() { + opts := grpcModelOpts(configWithCompanion("attention_backend:sdpa"), "/models") + + value, found := optionValue(opts.Options, "base_model") + Expect(found).To(BeTrue()) + expected, err := modelartifacts.RelativeSnapshotPath(companionKey) + Expect(err).NotTo(HaveOccurred()) + Expect(value).To(Equal(expected)) + // The companion path must never be confused with the load target. + Expect(value).ToNot(ContainSubstring(primaryKey)) + // Author-supplied options survive untouched. + Expect(opts.Options).To(ContainElement("attention_backend:sdpa")) + }) + + It("keeps the path relative so the worker resolves it under its own ModelPath", func() { + opts := grpcModelOpts(configWithCompanion(), "/models") + + value, found := optionValue(opts.Options, "base_model") + Expect(found).To(BeTrue()) + Expect(strings.HasPrefix(value, "/")).To(BeFalse()) + }) + + It("does not override an explicitly configured option of the same name", func() { + // An operator pinning base_model to a local checkout must win over the + // synthesized value. + opts := grpcModelOpts(configWithCompanion("base_model:/opt/checkouts/LongCat-Video"), "/models") + + Expect(opts.Options).To(ContainElement("base_model:/opt/checkouts/LongCat-Video")) + values := 0 + for _, option := range opts.Options { + if strings.HasPrefix(option, "base_model:") { + values++ + } + } + Expect(values).To(Equal(1)) + }) + + It("synthesizes nothing for a config without companions", func() { + cfg := config.ModelConfig{ + Backend: "transformers", + Options: []string{"attention_backend:sdpa"}, + Threads: &threads, + Artifacts: []modelartifacts.Spec{{ + Name: modelartifacts.TargetModel, Target: modelartifacts.TargetModel, + Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo", Revision: "main"}, + Resolved: resolved(primaryKey), + }}, + } + opts := grpcModelOpts(cfg, "/models") + Expect(opts.Options).To(Equal([]string{"attention_backend:sdpa"})) + }) + + It("skips a companion that has not been resolved yet", func() { + cfg := configWithCompanion() + cfg.Artifacts[1].Resolved = nil + opts := grpcModelOpts(cfg, "/models") + _, found := optionValue(opts.Options, "base_model") + Expect(found).To(BeFalse()) + }) +}) diff --git a/core/backend/options.go b/core/backend/options.go index 664e99af8826..a7cf00dc3dd2 100644 --- a/core/backend/options.go +++ b/core/backend/options.go @@ -7,6 +7,7 @@ import ( "math/rand/v2" "os" "path/filepath" + "slices" "strings" "time" @@ -280,6 +281,48 @@ func EffectiveBatchSize(c config.ModelConfig) int { return DefaultBatchSize } +// withCompanionArtifactOptions surfaces each resolved companion snapshot to the +// backend as ":", reusing the key:value option +// convention backends already parse. +// +// The value is deliberately relative to the models directory and deliberately +// not persisted to the config YAML. It is derived from a content-addressed cache +// key that only exists after the artifact resolves, so a static gallery override +// could not carry it, and a persisted copy would rot the moment a re-resolve +// produced a new key. Staying relative also lets a remote worker resolve it +// under its own ModelPath after staging rewrites the model root. +// +// An option the author set explicitly always wins: pinning a companion to a +// local checkout has to beat the managed snapshot. +func withCompanionArtifactOptions(options []string, artifacts []modelartifacts.Spec) []string { + configured := make(map[string]struct{}, len(options)) + for _, option := range options { + if name, _, found := strings.Cut(option, ":"); found { + configured[name] = struct{}{} + } + } + + // Copy before appending: opts.Options would otherwise share (and could + // reallocate away from) the config's own slice. + combined := slices.Clone(options) + for _, artifact := range artifacts { + if artifact.Target != modelartifacts.TargetCompanion || artifact.Resolved == nil { + continue + } + if _, exists := configured[artifact.Name]; exists { + xlog.Debug("keeping the configured companion option over the managed snapshot", "artifact", artifact.Name) + continue + } + snapshot, err := modelartifacts.RelativeSnapshotPath(artifact.Resolved.CacheKey) + if err != nil { + xlog.Warn("skipping companion artifact with an unusable cache key", "artifact", artifact.Name, "error", err) + continue + } + combined = append(combined, artifact.Name+":"+snapshot) + } + return combined +} + func grpcModelOpts(c config.ModelConfig, modelPath string) *pb.ModelOptions { ctxSize := EffectiveContextSize(c) b := EffectiveBatchSize(c) @@ -370,7 +413,7 @@ func grpcModelOpts(c config.ModelConfig, modelPath string) *pb.ModelOptions { IMG2IMG: c.Diffusers.IMG2IMG, CLIPModel: c.Diffusers.ClipModel, CLIPSubfolder: c.Diffusers.ClipSubFolder, - Options: c.Options, + Options: withCompanionArtifactOptions(c.Options, c.Artifacts), Overrides: c.Overrides, EngineArgs: engineArgsJSON, CLIPSkip: int32(c.Diffusers.ClipSkip), diff --git a/core/config/model_artifact_companion_test.go b/core/config/model_artifact_companion_test.go new file mode 100644 index 000000000000..8bec13c85cd8 --- /dev/null +++ b/core/config/model_artifact_companion_test.go @@ -0,0 +1,113 @@ +package config_test + +import ( + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gopkg.in/yaml.v3" + + "github.com/mudler/LocalAI/core/config" +) + +// A composed pipeline needs more than one snapshot: LongCat-Video-Avatar-1.5 +// loads its own transformer but takes tokenizer, text encoder and VAE from the +// LongCat-Video base repo. The config expresses that as one target: model +// artifact followed by named target: companion artifacts. Exactly one primary, +// and it must be first, because ModelFileName() and every load path resolve the +// load target from Artifacts[0]. +var _ = Describe("multi-artifact config validation", func() { + parse := func(doc string) config.ModelConfig { + var c config.ModelConfig + Expect(yaml.Unmarshal([]byte(doc), &c)).To(Succeed()) + return c + } + + It("accepts a primary followed by a named companion", func() { + c := parse(` +backend: longcat-video +artifacts: + - name: model + target: model + source: {type: huggingface, repo: meituan-longcat/LongCat-Video-Avatar-1.5} + - name: base_model + target: companion + source: {type: huggingface, repo: meituan-longcat/LongCat-Video} +parameters: {model: meituan-longcat/LongCat-Video-Avatar-1.5} +`) + valid, err := c.Validate() + Expect(err).NotTo(HaveOccurred()) + Expect(valid).To(BeTrue()) + }) + + It("rejects a config whose artifacts are all companions", func() { + // The reachable way to end up without a primary. Two primaries cannot + // coexist by construction: both would have to be named "model", so the + // duplicate-name rule catches that case first (covered below). + c := parse(` +backend: longcat-video +artifacts: + - name: base_model + target: companion + source: {type: huggingface, repo: owner/base} +parameters: {model: owner/main} +`) + valid, err := c.Validate() + Expect(valid).To(BeFalse()) + Expect(err).To(MatchError(ContainSubstring("exactly one"))) + }) + + It("rejects a second artifact claiming the primary target", func() { + c := parse(` +backend: longcat-video +artifacts: + - name: model + target: model + source: {type: huggingface, repo: owner/one} + - name: second + target: model + source: {type: huggingface, repo: owner/two} +parameters: {model: owner/one} +`) + valid, err := c.Validate() + Expect(valid).To(BeFalse()) + Expect(err).To(MatchError(ContainSubstring("primary artifact name"))) + }) + + It("rejects a companion declared before the primary", func() { + // Artifacts[0] is the load target everywhere; a companion in that slot + // would silently point the backend at the wrong snapshot. + c := parse(` +backend: longcat-video +artifacts: + - name: base_model + target: companion + source: {type: huggingface, repo: owner/base} + - name: model + target: model + source: {type: huggingface, repo: owner/main} +parameters: {model: owner/main} +`) + valid, err := c.Validate() + Expect(valid).To(BeFalse()) + Expect(err).To(MatchError(ContainSubstring("first"))) + }) + + It("rejects companions sharing a name", func() { + c := parse(` +backend: longcat-video +artifacts: + - name: model + target: model + source: {type: huggingface, repo: owner/main} + - name: base_model + target: companion + source: {type: huggingface, repo: owner/base} + - name: base_model + target: companion + source: {type: huggingface, repo: owner/other} +parameters: {model: owner/main} +`) + valid, err := c.Validate() + Expect(valid).To(BeFalse()) + Expect(err).To(MatchError(ContainSubstring("duplicate"))) + }) +}) diff --git a/core/config/model_artifact_inference.go b/core/config/model_artifact_inference.go index a011a7a6f7cb..b9c90856e611 100644 --- a/core/config/model_artifact_inference.go +++ b/core/config/model_artifact_inference.go @@ -21,6 +21,7 @@ var managedArtifactBackends = map[string]struct{}{ "transformers-musicgen": {}, "mamba": {}, "diffusers": {}, "qwen-asr": {}, "fish-speech": {}, "nemo": {}, "voxcpm": {}, "qwen-tts": {}, "liquid-audio": {}, "vllm": {}, "vllm-omni": {}, "sglang": {}, + "longcat-video": {}, } // IsManagedArtifactBackend reports whether backend consumes a model as a diff --git a/core/config/model_artifact_inference_test.go b/core/config/model_artifact_inference_test.go index 1bc5987ad540..4a6b9a4ce471 100644 --- a/core/config/model_artifact_inference_test.go +++ b/core/config/model_artifact_inference_test.go @@ -42,6 +42,20 @@ var _ = Describe("PrimaryArtifactSpec backend gating", func() { Expect(spec.Source.Repo).To(Equal("owner/repo")) }) + It("infers a managed artifact for longcat-video from a bare repo id", func() { + // longcat-video loads a checkpoint directory (its backend.py takes + // request.ModelFile when os.path.isdir), so it belongs to the same + // class as transformers/vllm/diffusers. Off the allow-list, the + // controller never acquires the weights and the backend re-downloads + // them from HuggingFace inside the remote LoadModel deadline. + c := parse("backend: longcat-video\nparameters: {model: meituan-longcat/LongCat-Video-Avatar-1.5}\n") + spec, inferred, managed, err := c.PrimaryArtifactSpec("/models") + Expect(err).NotTo(HaveOccurred()) + Expect(managed).To(BeTrue()) + Expect(inferred).To(BeTrue()) + Expect(spec.Source.Repo).To(Equal("meituan-longcat/LongCat-Video-Avatar-1.5")) + }) + It("keeps explicit artifacts managed even on a single-file backend", func() { // An explicitly declared artifacts: block is a deliberate choice; // single-file resolution (PrimaryFile) handles the load path. diff --git a/core/config/model_config.go b/core/config/model_config.go index 65bef77cfd54..7ca24be64dc9 100644 --- a/core/config/model_config.go +++ b/core/config/model_config.go @@ -1334,6 +1334,7 @@ func (c *ModelConfig) Validate() (bool, error) { return false, fmt.Errorf("alias model %q cannot declare artifacts", c.Name) } seenArtifacts := make(map[string]struct{}, len(c.Artifacts)) + primaries := 0 for i, artifact := range c.Artifacts { normalized, err := artifact.Normalize() if err != nil { @@ -1343,6 +1344,18 @@ func (c *ModelConfig) Validate() (bool, error) { return false, fmt.Errorf("duplicate artifact name %q", normalized.Name) } seenArtifacts[normalized.Name] = struct{}{} + if normalized.Target == modelartifacts.TargetModel { + primaries++ + // Artifacts[0] is the load target for every consumer (ModelFileName, + // size estimation, staging), so the primary has to occupy that slot + // rather than merely exist somewhere in the list. + if i != 0 { + return false, fmt.Errorf("the primary artifact must be declared first, found it at position %d", i) + } + } + } + if len(c.Artifacts) > 0 && primaries != 1 { + return false, fmt.Errorf("a config with artifacts must declare exactly one %q target, found %d", modelartifacts.TargetModel, primaries) } // An alias is a pure redirect: validate only its own shape here. Target diff --git a/core/config/model_config_loader.go b/core/config/model_config_loader.go index 4423fddadc5c..eea1058d322d 100644 --- a/core/config/model_config_loader.go +++ b/core/config/model_config_loader.go @@ -505,6 +505,24 @@ func (bcl *ModelConfigLoader) preloadOne( } } + // Companions are only ever declared explicitly, so unlike the inferred + // primary above they have no legacy path to fall back to: a config that + // names one is asserting the backend needs it. Failing here keeps the error + // at the acquisition boundary instead of surfacing as a confusing + // missing-weights error inside the backend. + if managedPrimary { + for i := 1; i < len(updated.Artifacts); i++ { + if updated.Artifacts[i].Target != modelartifacts.TargetCompanion { + continue + } + companion, err := bcl.artifactMaterializer.Ensure(ctx, modelPath, updated.Artifacts[i]) + if err != nil { + return ModelConfig{}, nil, fmt.Errorf("materialize companion artifact %q: %w", updated.Artifacts[i].Name, err) + } + updated.Artifacts[i] = companion.Spec + } + } + if !managedPrimary && updated.IsModelURL() { modelFileName := updated.ModelFileName() uri := downloader.URI(updated.Model) diff --git a/core/config/model_config_loader_companion_test.go b/core/config/model_config_loader_companion_test.go new file mode 100644 index 000000000000..028d7516c8c6 --- /dev/null +++ b/core/config/model_config_loader_companion_test.go @@ -0,0 +1,128 @@ +package config + +import ( + "context" + "fmt" + "os" + "path/filepath" + "strings" + "sync" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/pkg/modelartifacts" +) + +// companionMaterializer resolves each requested repo to its own snapshot, so a +// test can tell which artifacts were actually acquired rather than assuming the +// single fixed result the other preload fakes return. +type companionMaterializer struct { + mu sync.Mutex + specs []modelartifacts.Spec + fail string +} + +func (f *companionMaterializer) Ensure(_ context.Context, _ string, spec modelartifacts.Spec) (modelartifacts.Result, error) { + f.mu.Lock() + defer f.mu.Unlock() + f.specs = append(f.specs, spec) + if f.fail != "" && spec.Source.Repo == f.fail { + return modelartifacts.Result{}, fmt.Errorf("materialization refused for %s", spec.Source.Repo) + } + // A distinct, well-formed cache key per repo, so the resulting snapshot + // paths differ and a test can prove the companion is not aliasing the + // primary. + key := strings.Repeat(fmt.Sprintf("%x", len(spec.Source.Repo)%16), 64) + resolved := spec + resolved.Source.Revision = "main" + resolved.Resolved = &modelartifacts.Resolved{ + Endpoint: "https://huggingface.co", + Revision: "0123456789abcdef0123456789abcdef01234567", + CacheKey: key, + } + return modelartifacts.Result{ + Spec: resolved, + RelativePath: filepath.Join(".artifacts", "huggingface", key, "snapshot"), + }, nil +} + +func (f *companionMaterializer) repos() []string { + f.mu.Lock() + defer f.mu.Unlock() + seen := make([]string, 0, len(f.specs)) + for _, spec := range f.specs { + seen = append(seen, spec.Source.Repo) + } + return seen +} + +var _ = Describe("companion artifact materialization", func() { + const companionConfig = ` +name: avatar +backend: longcat-video +artifacts: + - name: model + target: model + source: {type: huggingface, repo: meituan-longcat/LongCat-Video-Avatar-1.5} + - name: base_model + target: companion + source: {type: huggingface, repo: meituan-longcat/LongCat-Video} +parameters: {model: meituan-longcat/LongCat-Video-Avatar-1.5} +` + + It("acquires every declared artifact, not only the primary", func() { + 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()) + + Expect(fake.repos()).To(ConsistOf( + "meituan-longcat/LongCat-Video-Avatar-1.5", + "meituan-longcat/LongCat-Video", + )) + }) + + It("persists the resolved state of every artifact", func() { + 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()) + + loaded, found := loader.GetModelConfig("avatar") + Expect(found).To(BeTrue()) + Expect(loaded.Artifacts).To(HaveLen(2)) + Expect(loaded.Artifacts[0].Target).To(Equal(modelartifacts.TargetModel)) + Expect(loaded.Artifacts[0].Resolved).ToNot(BeNil()) + Expect(loaded.Artifacts[1].Name).To(Equal("base_model")) + Expect(loaded.Artifacts[1].Target).To(Equal(modelartifacts.TargetCompanion)) + Expect(loaded.Artifacts[1].Resolved).ToNot(BeNil()) + // The companion must land in its own snapshot, never alias the primary. + Expect(loaded.Artifacts[1].Resolved.CacheKey).ToNot(Equal(loaded.Artifacts[0].Resolved.CacheKey)) + // The load target still resolves from the primary alone. + Expect(loaded.ModelFileName()).To(ContainSubstring(loaded.Artifacts[0].Resolved.CacheKey)) + }) + + 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 + // would just move the failure somewhere less legible. + modelsPath := GinkgoT().TempDir() + configPath := filepath.Join(modelsPath, "avatar.yaml") + Expect(os.WriteFile(configPath, []byte(companionConfig), 0644)).To(Succeed()) + + fake := &companionMaterializer{fail: "meituan-longcat/LongCat-Video"} + loader := NewModelConfigLoader(modelsPath, WithArtifactMaterializer(fake)) + Expect(loader.LoadModelConfigsFromPath(modelsPath)).To(Succeed()) + err := loader.PreloadWithContext(context.Background(), modelsPath) + Expect(err).To(MatchError(ContainSubstring("LongCat-Video"))) + }) +}) diff --git a/core/gallery/model_artifacts.go b/core/gallery/model_artifacts.go index 07924d9dfc16..78bd88905de7 100644 --- a/core/gallery/model_artifacts.go +++ b/core/gallery/model_artifacts.go @@ -53,6 +53,19 @@ func bindPrimaryArtifact(ctx context.Context, modelsPath string, typed *config.M next = append(next, typed.Artifacts[1:]...) } typed.Artifacts = next + // Companions are always explicit, so there is no legacy path to degrade to: + // failing here leaves any previously installed config untouched rather than + // persisting a half-acquired model that would fail later inside the backend. + for i := 1; i < len(typed.Artifacts); i++ { + if typed.Artifacts[i].Target != modelartifacts.TargetCompanion { + continue + } + companion, err := materializer.Ensure(ctx, modelsPath, typed.Artifacts[i]) + if err != nil { + return false, fmt.Errorf("materialize companion artifact %q: %w", typed.Artifacts[i].Name, err) + } + typed.Artifacts[i] = companion.Spec + } artifactYAML, err := yaml.Marshal(typed.Artifacts) if err != nil { return false, err diff --git a/core/gallery/model_artifacts_companion_test.go b/core/gallery/model_artifacts_companion_test.go new file mode 100644 index 000000000000..41a4cde43683 --- /dev/null +++ b/core/gallery/model_artifacts_companion_test.go @@ -0,0 +1,99 @@ +package gallery_test + +import ( + "context" + "fmt" + "os" + "path/filepath" + "strings" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gopkg.in/yaml.v3" + + "github.com/mudler/LocalAI/core/gallery" + "github.com/mudler/LocalAI/pkg/modelartifacts" + "github.com/mudler/LocalAI/pkg/system" +) + +// perRepoMaterializer resolves each repo to its own snapshot so a test can +// distinguish the primary from its companions in the persisted config. +type perRepoMaterializer struct { + seen []modelartifacts.Spec + fail string +} + +func (f *perRepoMaterializer) Ensure(_ context.Context, _ string, spec modelartifacts.Spec) (modelartifacts.Result, error) { + f.seen = append(f.seen, spec) + if f.fail != "" && spec.Source.Repo == f.fail { + return modelartifacts.Result{}, fmt.Errorf("acquisition failed for %s", spec.Source.Repo) + } + key := strings.Repeat(fmt.Sprintf("%x", len(spec.Source.Repo)%16), 64) + resolved := spec + resolved.Source.Revision = "main" + resolved.Resolved = &modelartifacts.Resolved{ + Endpoint: "https://huggingface.co", + Revision: "0123456789abcdef0123456789abcdef01234567", + CacheKey: key, + } + return modelartifacts.Result{ + Spec: resolved, + RelativePath: filepath.Join(".artifacts", "huggingface", key, "snapshot"), + }, nil +} + +var _ = Describe("gallery companion artifact installation", func() { + const companionConfig = ` +backend: longcat-video +artifacts: + - name: model + target: model + source: {type: huggingface, repo: meituan-longcat/LongCat-Video-Avatar-1.5} + - name: base_model + target: companion + source: {type: huggingface, repo: meituan-longcat/LongCat-Video} +parameters: {model: meituan-longcat/LongCat-Video-Avatar-1.5} +` + + It("acquires and persists every declared artifact at install time", func() { + modelsPath := GinkgoT().TempDir() + state, err := system.GetSystemState(system.WithModelPath(modelsPath)) + Expect(err).NotTo(HaveOccurred()) + fake := &perRepoMaterializer{} + definition := &gallery.ModelConfig{Name: "avatar", ConfigFile: companionConfig} + + installed, err := gallery.InstallModel(context.Background(), state, "avatar", definition, nil, nil, false, + gallery.WithArtifactMaterializer(fake)) + Expect(err).NotTo(HaveOccurred()) + Expect(fake.seen).To(HaveLen(2)) + Expect(installed.Artifacts).To(HaveLen(2)) + Expect(installed.Artifacts[1].Name).To(Equal("base_model")) + Expect(installed.Artifacts[1].Resolved).ToNot(BeNil()) + + data, err := os.ReadFile(filepath.Join(modelsPath, "avatar.yaml")) + Expect(err).NotTo(HaveOccurred()) + var persisted map[string]any + Expect(yaml.Unmarshal(data, &persisted)).To(Succeed()) + artifacts := persisted["artifacts"].([]any) + Expect(artifacts).To(HaveLen(2)) + companion := artifacts[1].(map[string]any) + Expect(companion["name"]).To(Equal("base_model")) + Expect(companion).To(HaveKey("resolved")) + }) + + It("does not install when a companion cannot be acquired", func() { + modelsPath := GinkgoT().TempDir() + state, err := system.GetSystemState(system.WithModelPath(modelsPath)) + Expect(err).NotTo(HaveOccurred()) + configPath := filepath.Join(modelsPath, "avatar.yaml") + Expect(os.WriteFile(configPath, []byte("name: old\n"), 0644)).To(Succeed()) + fake := &perRepoMaterializer{fail: "meituan-longcat/LongCat-Video"} + definition := &gallery.ModelConfig{Name: "avatar", ConfigFile: companionConfig} + + _, err = gallery.InstallModel(context.Background(), state, "avatar", definition, nil, nil, false, + gallery.WithArtifactMaterializer(fake)) + Expect(err).To(MatchError(ContainSubstring("LongCat-Video"))) + // The pre-existing config must survive a failed install. + Expect(os.ReadFile(configPath)).To(Equal([]byte("name: old\n"))) + }) +}) diff --git a/core/services/nodes/router.go b/core/services/nodes/router.go index 01100123401a..4fc5dce4cdaa 100644 --- a/core/services/nodes/router.go +++ b/core/services/nodes/router.go @@ -1030,10 +1030,9 @@ func (r *SmartRouter) stageModelFiles(ctx context.Context, node *BackendNode, op // Derive the frontend models directory from ModelFile and Model. // Example: ModelFile="/models/sd-cpp/models/flux.gguf", Model="sd-cpp/models/flux.gguf" // → frontendModelsDir="/models" - frontendModelsDir := "" - if opts.ModelFile != "" && opts.Model != "" { - frontendModelsDir = filepath.Clean(strings.TrimSuffix(opts.ModelFile, opts.Model)) - } + // A managed artifact anchors on the models root that holds the whole + // .artifacts tree instead, so sibling companion snapshots stay in scope. + frontendModelsDir := ModelsRootForModelFile(opts.ModelFile, opts.Model) // Local model directory, captured before the ModelFile field is rewritten to // its remote path below. Companion assets declared as option paths (e.g. @@ -1101,9 +1100,15 @@ func (r *SmartRouter) stageModelFiles(ctx context.Context, node *BackendNode, op if *f.val == "" { continue } - // Skip non-existent files + // Skip non-existent files. This is legitimate — a backend that takes a + // bare HuggingFace repo id gets an optimistically constructed path that + // was never materialized, and fetches its own weights on the worker — so + // it must not fail the load. But it is warn-level because the same skip + // is what a genuine controller-side acquisition gap looks like, and at + // debug it left the operator with a reassuring "Staging model files" + // line for work that never happened. if _, err := os.Stat(*f.val); os.IsNotExist(err) { - xlog.Debug("Skipping staging for non-existent path", "field", f.name, "path", *f.val) + xlog.Warn("Skipping staging for non-existent path; the worker will have to source this itself", "field", f.name, "path", *f.val, "node", node.Name, "trackingKey", trackingKey) *f.val = "" continue } @@ -1126,8 +1131,8 @@ func (r *SmartRouter) stageModelFiles(ctx context.Context, node *BackendNode, op continue } *f.val = remoteDir - if f.name == "ModelFile" && opts.Model != "" { - opts.ModelPath = DeriveRemoteModelPath(remoteDir, opts.Model) + if f.name == "ModelFile" { + opts.ModelPath = DeriveRemoteModelPath(remoteDir, relativeToModelsDir(frontendModelsDir, localPath, opts.Model)) xlog.Debug("Derived remote ModelPath", "modelPath", opts.ModelPath) } continue @@ -1164,8 +1169,8 @@ func (r *SmartRouter) stageModelFiles(ctx context.Context, node *BackendNode, op // remotePath = "/worker/models/{trackingKey}/sd-cpp/models/flux.gguf" // Model = "sd-cpp/models/flux.gguf" // → ModelPath = "/worker/models/{trackingKey}" - if f.name == "ModelFile" && opts.Model != "" { - opts.ModelPath = DeriveRemoteModelPath(remotePath, opts.Model) + if f.name == "ModelFile" { + opts.ModelPath = DeriveRemoteModelPath(remotePath, relativeToModelsDir(frontendModelsDir, localPath, opts.Model)) xlog.Debug("Derived remote ModelPath", "modelPath", opts.ModelPath) } diff --git a/core/services/nodes/router_companionstage_test.go b/core/services/nodes/router_companionstage_test.go new file mode 100644 index 000000000000..92958aaeb741 --- /dev/null +++ b/core/services/nodes/router_companionstage_test.go @@ -0,0 +1,136 @@ +package nodes + +import ( + "context" + "os" + "path/filepath" + "strings" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/storage" + pb "github.com/mudler/LocalAI/pkg/grpc/proto" +) + +// Staging a managed model has to preserve the models-root-relative layout of +// the content-addressed artifact tree, not just the contents of the primary +// snapshot. +// +// A composed pipeline resolves a companion snapshot through an option holding a +// path relative to the models directory (see withCompanionArtifactOptions). +// That only works if the worker's ModelPath is the staged models ROOT: the +// companion lives in a sibling .artifacts/huggingface//snapshot directory, +// outside the primary snapshot entirely, so anchoring on the primary would put +// it out of reach and collapse its files to bare basenames. +var _ = Describe("stageModelFiles managed artifact trees", func() { + const ( + primaryKey = "1111111111111111111111111111111111111111111111111111111111111111" + companionKey = "2222222222222222222222222222222222222222222222222222222222222222" + ) + + var ( + stager *fakeFileStager + router *SmartRouter + node *BackendNode + modelsDir string + primaryRel string + compRel string + ) + + write := func(root string, files map[string]string) { + for relative, contents := range files { + path := filepath.Join(root, relative) + Expect(os.MkdirAll(filepath.Dir(path), 0o750)).To(Succeed()) + Expect(os.WriteFile(path, []byte(contents), 0o644)).To(Succeed()) + } + } + + BeforeEach(func() { + stager = &fakeFileStager{} + router = &SmartRouter{fileStager: stager, stagingTracker: NewStagingTracker()} + node = &BackendNode{ID: "node-1", Name: "nvidia-thor", Address: "10.0.0.1:50051"} + modelsDir = filepath.Join(GinkgoT().TempDir(), "models") + primaryRel = filepath.Join(".artifacts", "huggingface", primaryKey, "snapshot") + compRel = filepath.Join(".artifacts", "huggingface", companionKey, "snapshot") + + write(filepath.Join(modelsDir, primaryRel), map[string]string{ + "model_index.json": "{}", + "base_model/diffusion.safetensors": "weights", + }) + write(filepath.Join(modelsDir, compRel), map[string]string{ + "model_index.json": "{}", + "vae/vae.safetensors": "vae", + "tokenizer/vocab.json": "{}", + }) + }) + + stagedKeys := func() []string { + keys := make([]string, 0, len(stager.ensureCalls)) + for _, call := range stager.ensureCalls { + keys = append(keys, call.key) + } + return keys + } + + It("stages a companion snapshot referenced by a relative option", func() { + opts := &pb.ModelOptions{ + Model: "meituan-longcat/LongCat-Video-Avatar-1.5", + ModelFile: filepath.Join(modelsDir, primaryRel), + Options: []string{"attention_backend:sdpa", "base_model:" + compRel}, + } + + staged, err := router.stageModelFiles(context.Background(), node, opts, "avatar") + Expect(err).ToNot(HaveOccurred()) + + // Every file of both snapshots ships, each under its models-root-relative + // path so the two trees stay distinguishable on the worker. + Expect(stagedKeys()).To(ConsistOf( + storage.ModelKey(filepath.Join("avatar", primaryRel, "model_index.json")), + storage.ModelKey(filepath.Join("avatar", primaryRel, "base_model/diffusion.safetensors")), + storage.ModelKey(filepath.Join("avatar", compRel, "model_index.json")), + storage.ModelKey(filepath.Join("avatar", compRel, "vae/vae.safetensors")), + storage.ModelKey(filepath.Join("avatar", compRel, "tokenizer/vocab.json")), + )) + + // The option value is never rewritten: it stays relative and resolves + // against the worker's ModelPath. + Expect(staged.Options).To(ContainElement("base_model:" + compRel)) + Expect(staged.Options).To(ContainElement("attention_backend:sdpa")) + }) + + It("points ModelPath at the staged models root, not at the primary snapshot", func() { + opts := &pb.ModelOptions{ + Model: "meituan-longcat/LongCat-Video-Avatar-1.5", + ModelFile: filepath.Join(modelsDir, primaryRel), + Options: []string{"base_model:" + compRel}, + } + + staged, err := router.stageModelFiles(context.Background(), node, opts, "avatar") + Expect(err).ToNot(HaveOccurred()) + + root := filepath.Join("/remote", storage.ModelKey("avatar")) + Expect(staged.ModelPath).To(Equal(root)) + // The load target still addresses the primary snapshot itself. + Expect(staged.ModelFile).To(Equal(filepath.Join(root, primaryRel))) + // And joining ModelPath with the untouched option value has to land on + // the companion the worker actually received. + Expect(filepath.Join(staged.ModelPath, compRel)).To(Equal(filepath.Join(root, compRel))) + }) + + It("keeps a legacy relative model file anchored on the models directory", func() { + // The pre-artifact layout, where Model is a real relative path under the + // models dir, must keep deriving the same root it always did. + legacy := filepath.Join(modelsDir, "sd-cpp", "models") + write(legacy, map[string]string{"flux.gguf": "weights"}) + opts := &pb.ModelOptions{ + Model: filepath.Join("sd-cpp", "models", "flux.gguf"), + ModelFile: filepath.Join(legacy, "flux.gguf"), + } + + staged, err := router.stageModelFiles(context.Background(), node, opts, "flux") + Expect(err).ToNot(HaveOccurred()) + Expect(strings.HasSuffix(staged.ModelPath, storage.ModelKey("flux"))).To(BeTrue()) + Expect(staged.ModelFile).To(Equal(filepath.Join(staged.ModelPath, "sd-cpp", "models", "flux.gguf"))) + }) +}) diff --git a/core/services/nodes/router_dirstage_test.go b/core/services/nodes/router_dirstage_test.go index fef138f48a4a..41fed3bf4c5e 100644 --- a/core/services/nodes/router_dirstage_test.go +++ b/core/services/nodes/router_dirstage_test.go @@ -96,14 +96,22 @@ var _ = Describe("stageModelFiles directory models", func() { expectedKeys := make([]string, 0, len(files)) for relative := range files { expectedPaths = append(expectedPaths, filepath.Join(snapshot, relative)) - expectedKeys = append(expectedKeys, storage.ModelKey(filepath.Join("managed-model", relative))) + // Keys keep the full models-root-relative path, including the + // .artifacts/huggingface//snapshot prefix. This changed when + // companion artifacts arrived: keys used to be relative to the + // snapshot itself, which flattened the prefix away and left two + // snapshots of one model indistinguishable on the worker. + expectedKeys = append(expectedKeys, storage.ModelKey(filepath.Join("managed-model", relativeSnapshot, relative))) } Expect(stagedPaths).To(ConsistOf(expectedPaths)) Expect(stagedKeys).To(ConsistOf(expectedKeys)) - remoteSnapshot := filepath.Join("/remote", storage.ModelKey("managed-model")) + remoteRoot := filepath.Join("/remote", storage.ModelKey("managed-model")) Expect(stagedOpts.Model).To(Equal(logicalModel)) - Expect(stagedOpts.ModelFile).To(Equal(remoteSnapshot)) - Expect(stagedOpts.ModelPath).To(Equal(remoteSnapshot)) + Expect(stagedOpts.ModelFile).To(Equal(filepath.Join(remoteRoot, relativeSnapshot))) + // ModelPath is the staged models ROOT, not the snapshot. It used to be + // the snapshot, which made every relative option resolve inside the + // primary and put a sibling companion snapshot out of reach. + Expect(stagedOpts.ModelPath).To(Equal(remoteRoot)) Expect(opts.Model).To(Equal(logicalModel)) Expect(opts.ModelFile).To(Equal(snapshot)) }) diff --git a/core/services/nodes/router_missingsource_test.go b/core/services/nodes/router_missingsource_test.go new file mode 100644 index 000000000000..401a01170cb7 --- /dev/null +++ b/core/services/nodes/router_missingsource_test.go @@ -0,0 +1,71 @@ +package nodes + +import ( + "bytes" + "context" + "log/slog" + "path/filepath" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + pb "github.com/mudler/LocalAI/pkg/grpc/proto" + "github.com/mudler/xlog" +) + +// A ModelFile that does not exist on the controller is legitimate: every +// backend outside config.managedArtifactBackends that takes a bare HuggingFace +// repo id gets an optimistically constructed path (ModelFileName falls through +// to the raw model reference) that was never materialized, and downloads its +// own weights on the worker. Staging must therefore keep skipping it rather +// than failing. What it must NOT do is skip silently: the operator sees +// "Staging model files for remote node" and no further trace, so a genuine +// acquisition gap (#10910's allow-list omitting a directory-consuming backend) +// looks like a remote LoadModel timeout instead. The skip is logged at warn so +// it survives a default log level. +var _ = Describe("stageModelFiles missing source visibility", func() { + var ( + router *SmartRouter + node *BackendNode + captured *bytes.Buffer + ) + + BeforeEach(func() { + router = &SmartRouter{ + fileStager: &fakeFileStager{}, + stagingTracker: NewStagingTracker(), + } + node = &BackendNode{ID: "node-1", Name: "nvidia-thor", Address: "10.0.0.1:50051"} + + // Capture at warn level so a debug-level emission is filtered out and + // the assertion fails on visibility, not on wording. + captured = &bytes.Buffer{} + handler := slog.NewTextHandler(captured, &slog.HandlerOptions{Level: slog.LevelWarn}) + xlog.SetLogger(xlog.NewLoggerWithHandler(handler, xlog.LogLevelWarn)) + }) + + AfterEach(func() { + // xlog exposes no getter for the package logger, so restore the same + // default the suite entrypoint installs rather than the prior value. + xlog.SetLogger(xlog.NewLogger(xlog.LogLevel("info"), "text")) + }) + + It("warns, and names the field and path, when the model file does not exist locally", func() { + missing := filepath.Join(GinkgoT().TempDir(), "meituan-longcat", "LongCat-Video-Avatar-1.5") + opts := &pb.ModelOptions{ + Model: "meituan-longcat/LongCat-Video-Avatar-1.5", + ModelFile: missing, + } + + staged, err := router.stageModelFiles(context.Background(), node, opts, "longcat-video-avatar-1.5") + Expect(err).ToNot(HaveOccurred()) + + // Behavior is unchanged: the field is still blanked so the worker does + // not receive a controller-only path. + Expect(staged.ModelFile).To(BeEmpty()) + + logged := captured.String() + Expect(logged).To(ContainSubstring("ModelFile")) + Expect(logged).To(ContainSubstring(missing)) + }) +}) diff --git a/core/services/nodes/staging_keys.go b/core/services/nodes/staging_keys.go index 7235ea46d916..489b50c735ae 100644 --- a/core/services/nodes/staging_keys.go +++ b/core/services/nodes/staging_keys.go @@ -29,14 +29,62 @@ func (m *StagingKeyMapper) Key(localPath string) string { } // DeriveRemoteModelPath computes the worker's ModelPath from the staged -// ModelFile remote path and the original Model relative path. -// Returns the ModelPath that makes all relative option values resolve correctly. +// ModelFile remote path and that file's path relative to the frontend models +// directory. Returns the ModelPath that makes all relative option values +// resolve correctly. // // Example: // // remotePath = "/worker/models/my-model/sd-cpp/models/flux.gguf" -// model = "sd-cpp/models/flux.gguf" +// relative = "sd-cpp/models/flux.gguf" // → returns "/worker/models/my-model" -func DeriveRemoteModelPath(remotePath, model string) string { - return filepath.Clean(strings.TrimSuffix(remotePath, model)) +func DeriveRemoteModelPath(remotePath, relative string) string { + return filepath.Clean(strings.TrimSuffix(remotePath, relative)) +} + +// relativeToModelsDir returns the staged file's path relative to the frontend +// models directory, which is the suffix DeriveRemoteModelPath strips to recover +// the worker's models root. It falls back to the Model field for the callers +// that predate a resolvable models directory. +func relativeToModelsDir(frontendModelsDir, localPath, model string) string { + if frontendModelsDir != "" { + if rel, err := filepath.Rel(frontendModelsDir, localPath); err == nil && + !strings.HasPrefix(rel, "..") && rel != "." { + return rel + } + } + return model +} + +// managedArtifactDir is the root of the content-addressed artifact tree, kept +// in sync with the layout modelartifacts.LayoutFor builds. +const managedArtifactDir = ".artifacts" + +// ModelsRootForModelFile finds the frontend models directory a model file sits +// under. +// +// The legacy answer is to strip the Model relative path off the end of +// ModelFile. That silently degrades for a managed artifact, whose ModelFile is +// a content-addressed snapshot directory while Model stays a bare HuggingFace +// repo id: nothing matches, the suffix strip is a no-op, and the "models +// directory" comes out as the snapshot itself. Everything anchored on it then +// narrows to that one snapshot, which is precisely where a sibling companion +// snapshot becomes unreachable and its files collapse to bare basenames. +// +// So when the path runs through the managed artifact tree, anchor on the +// directory that contains that tree instead. +func ModelsRootForModelFile(modelFile, model string) string { + if modelFile == "" { + return "" + } + parts := strings.Split(filepath.Clean(modelFile), string(filepath.Separator)) + for i, part := range parts { + if part == managedArtifactDir && i > 0 { + return filepath.Clean(strings.Join(parts[:i], string(filepath.Separator))) + } + } + if model == "" { + return "" + } + return filepath.Clean(strings.TrimSuffix(modelFile, model)) } diff --git a/gallery/index.yaml b/gallery/index.yaml index 79e7e265a49d..b382a022ccca 100644 --- a/gallery/index.yaml +++ b/gallery/index.yaml @@ -5711,6 +5711,44 @@ options: - attention_backend:sdpa - use_distill:true + # Avatar generation needs two repositories: the avatar checkpoint itself + # plus the tokenizer, text encoder and VAE from the LongCat-Video base + # repo. Declaring the base repo as a companion has LocalAI acquire both up + # front, so a distributed worker is handed staged weights instead of + # downloading them itself inside the load deadline. + # + # allow_patterns keeps this to what the configured options actually load. + # The avatar repo ships both base_model/ and base_model_int8/ and only one + # is ever read, so fetching both would roughly double the download. These + # patterns match use_distill:true with use_int8 left at its default of + # false; enabling use_int8 on this entry additionally requires adding + # base_model_int8/** here. + artifacts: + - name: model + target: model + source: + type: huggingface + repo: meituan-longcat/LongCat-Video-Avatar-1.5 + allow_patterns: + - config.json + - model_index.json + - scheduler/** + - lora/dmd_lora.safetensors + - whisper-large-v3/config.json + - whisper-large-v3/model.safetensors + - whisper-large-v3/preprocessor_config.json + - base_model/** + - name: base_model + target: companion + source: + type: huggingface + repo: meituan-longcat/LongCat-Video + allow_patterns: + - config.json + - model_index.json + - text_encoder/** + - tokenizer/** + - vae/** parameters: model: meituan-longcat/LongCat-Video-Avatar-1.5 - name: vllm-omni-qwen3-omni-30b diff --git a/pkg/modelartifacts/types.go b/pkg/modelartifacts/types.go index be375d9099b2..850ad414b952 100644 --- a/pkg/modelartifacts/types.go +++ b/pkg/modelartifacts/types.go @@ -12,12 +12,18 @@ import ( const ( SourceTypeHuggingFace = "huggingface" TargetModel = "model" - HuggingFaceTokenEnv = "HF_TOKEN" + // TargetCompanion marks a snapshot that supports the primary model rather + // than being it: a composed pipeline may pull its tokenizer, text encoder or + // VAE from a separate repository. Companions are surfaced to the backend as + // named options, never as the load target. + TargetCompanion = "companion" + HuggingFaceTokenEnv = "HF_TOKEN" ) var ( - commitPattern = regexp.MustCompile(`^[0-9a-f]{40}$`) - cacheKeyPattern = regexp.MustCompile(`^[0-9a-f]{64}$`) + commitPattern = regexp.MustCompile(`^[0-9a-f]{40}$`) + cacheKeyPattern = regexp.MustCompile(`^[0-9a-f]{64}$`) + artifactNamePattern = regexp.MustCompile(`^[a-z0-9][a-z0-9_-]{0,63}$`) ) type Spec struct { @@ -50,19 +56,39 @@ type Resolved struct { } func (s Spec) Normalize() (Spec, error) { - if strings.TrimSpace(s.Name) == "" { - s.Name = TargetModel - } if strings.TrimSpace(s.Target) == "" { s.Target = TargetModel } s.Name = strings.TrimSpace(s.Name) s.Target = strings.TrimSpace(s.Target) + // The primary is always named "model": the whole codebase addresses it + // positionally as Artifacts[0], so a second spelling would buy nothing and + // break the pre-companion configs that omit the field entirely. + if s.Target == TargetModel && s.Name == "" { + s.Name = TargetModel + } s.Source.Type = strings.TrimSpace(s.Source.Type) s.Source.TokenEnv = strings.TrimSpace(s.Source.TokenEnv) - if s.Name != TargetModel || s.Target != TargetModel { - return Spec{}, fmt.Errorf("only artifact name/target %q is supported", TargetModel) + switch s.Target { + case TargetModel: + if s.Name != TargetModel { + return Spec{}, fmt.Errorf("the primary artifact name must be %q, got %q", TargetModel, s.Name) + } + case TargetCompanion: + // A companion's name is the option key the backend receives, so it has + // to survive that round trip unambiguously. + if !artifactNamePattern.MatchString(s.Name) { + return Spec{}, fmt.Errorf("companion artifact name %q must match %s", s.Name, artifactNamePattern) + } + if s.Name == TargetModel { + return Spec{}, fmt.Errorf("companion artifact name %q is reserved for the primary", TargetModel) + } + default: + return Spec{}, fmt.Errorf("unsupported artifact target %q", s.Target) + } + if s.Source.Type != SourceTypeHuggingFace { + return Spec{}, fmt.Errorf("unsupported artifact source type %q", s.Source.Type) } if s.Source.Type != SourceTypeHuggingFace { return Spec{}, fmt.Errorf("unsupported artifact source type %q", s.Source.Type) @@ -112,6 +138,12 @@ func (s Spec) Normalize() (Spec, error) { } resolved.PrimaryFile = strings.TrimSpace(resolved.PrimaryFile) if resolved.PrimaryFile != "" { + // PrimaryFile exists to point a single-file backend at the weight + // file inside a snapshot. A companion is never the load target, so + // carrying one would encode an intent nothing acts on. + if s.Target == TargetCompanion { + return Spec{}, fmt.Errorf("companion artifact %q must not declare primary_file", s.Name) + } if err := ValidateRelativeHubPath(resolved.PrimaryFile); err != nil { return Spec{}, err } diff --git a/pkg/modelartifacts/types_test.go b/pkg/modelartifacts/types_test.go index e61fef317fab..eb9ddca66e37 100644 --- a/pkg/modelartifacts/types_test.go +++ b/pkg/modelartifacts/types_test.go @@ -67,4 +67,67 @@ var _ = Describe("artifact configuration", func() { } Expect(spec.Validate()).To(Succeed()) }) + + It("normalizes a named companion artifact", func() { + // Composed pipelines (LongCat-Video-Avatar pulling tokenizer, text + // encoder and VAE from the LongCat-Video base repo) need more than one + // snapshot. A companion is an ordinary HuggingFace snapshot addressed by + // name; the name is what the backend later receives as an option key. + spec, err := (modelartifacts.Spec{ + Name: "base_model", + Target: modelartifacts.TargetCompanion, + Source: modelartifacts.Source{Type: "huggingface", Repo: "meituan-longcat/LongCat-Video"}, + }).Normalize() + Expect(err).NotTo(HaveOccurred()) + Expect(spec.Name).To(Equal("base_model")) + Expect(spec.Target).To(Equal(modelartifacts.TargetCompanion)) + Expect(spec.Source.Revision).To(Equal("main")) + }) + + DescribeTable("rejects malformed companion declarations", + func(spec modelartifacts.Spec, message string) { + Expect(spec.Validate()).To(MatchError(ContainSubstring(message))) + }, + Entry("unknown target", modelartifacts.Spec{Name: "controlnet", Target: "controlnet", Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo"}}, "target"), + Entry("uppercase name", modelartifacts.Spec{Name: "Base_Model", Target: modelartifacts.TargetCompanion, Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo"}}, "name"), + Entry("name with a slash", modelartifacts.Spec{Name: "base/model", Target: modelartifacts.TargetCompanion, Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo"}}, "name"), + Entry("companion declaring a primary file", modelartifacts.Spec{Name: "base_model", Target: modelartifacts.TargetCompanion, Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo"}, Resolved: &modelartifacts.Resolved{Endpoint: "https://huggingface.co", Revision: "0123456789abcdef0123456789abcdef01234567", CacheKey: "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", PrimaryFile: "model.gguf"}}, "primary_file"), + ) + + // The cache key is the on-disk identity of a materialized snapshot + // (.artifacts/huggingface//snapshot). If widening the artifact model + // perturbed it, every already-installed managed model would miss its cache + // and silently re-download on next load. These two specs pin that. + It("keeps the cache key independent of artifact name and target", func() { + source := modelartifacts.Source{Type: "huggingface", Repo: "owner/repo"} + resolved := func() *modelartifacts.Resolved { + return &modelartifacts.Resolved{ + Endpoint: "https://huggingface.co", + Revision: "0123456789abcdef0123456789abcdef01234567", + } + } + primary, err := modelartifacts.CacheKey(modelartifacts.Spec{ + Name: modelartifacts.TargetModel, Target: modelartifacts.TargetModel, + Source: source, Resolved: resolved(), + }) + Expect(err).NotTo(HaveOccurred()) + companion, err := modelartifacts.CacheKey(modelartifacts.Spec{ + Name: "base_model", Target: modelartifacts.TargetCompanion, + Source: source, Resolved: resolved(), + }) + Expect(err).NotTo(HaveOccurred()) + Expect(companion).To(Equal(primary)) + }) + + It("pins the cache key digest for a known primary artifact", func() { + key, err := modelartifacts.CacheKey(modelartifacts.Spec{ + Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo"}, + Resolved: &modelartifacts.Resolved{ + Endpoint: "https://huggingface.co", + Revision: "0123456789abcdef0123456789abcdef01234567", + }, + }) + Expect(err).NotTo(HaveOccurred()) + Expect(key).To(Equal("b79e23e0b9c50af094d627582df30109eff8637438864172d64be07dfc5a98f9")) + }) })