From e5c95e044980ba04a64d5cb42807150d39b83381 Mon Sep 17 00:00:00 2001 From: "LocalAI [bot]" <139863280+localai-bot@users.noreply.github.com> Date: Sun, 14 Jun 2026 16:42:59 +0200 Subject: [PATCH] fix(distributed): stage backend companion assets to remote nodes (#10330) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A model whose ModelFile is a single file (e.g. sherpa-onnx VITS/piper: the .onnx) failed to load on remote worker nodes because the sibling assets the backend resolves from the model dir — tokens.txt, lexicon.txt, the espeak-ng-data / dict directories, Kokoro's voices.bin — were never staged. Only the declared ModelFile was shipped, so the worker hit "failed to create sherpa-onnx TTS engine" and TTS produced no audio. Lean on the existing option-path staging instead of hardcoding filenames: - stageGenericOptions now also resolves an option value relative to the model's own directory (not just the frontend models dir), so a shared config can declare companions with bare names regardless of whether Model includes a subdirectory; and it expands directory-valued options (e.g. espeak-ng-data) file-by-file rather than handing a directory fd to the stager. - gallery/sherpa-onnx-tts.yaml declares the companion assets as option paths (tokens, lexicon, espeak-ng-data, voices.bin, dict, per-lang lexicons). The backend ignores these keys and keeps resolving siblings from the model dir; they exist only so distributed staging ships them. Absent files are skipped. Adds router_optionstage_test.go covering file + directory companion staging via the model-dir fallback. Co-authored-by: Ettore Di Giacinto Co-authored-by: Claude Opus 4.8 (1M context) --- core/services/nodes/router.go | 87 ++++++++++++++++--- .../services/nodes/router_optionstage_test.go | 77 ++++++++++++++++ gallery/sherpa-onnx-tts.yaml | 14 +++ 3 files changed, 165 insertions(+), 13 deletions(-) create mode 100644 core/services/nodes/router_optionstage_test.go diff --git a/core/services/nodes/router.go b/core/services/nodes/router.go index f28200314..e5ce52306 100644 --- a/core/services/nodes/router.go +++ b/core/services/nodes/router.go @@ -908,6 +908,17 @@ func (r *SmartRouter) stageModelFiles(ctx context.Context, node *BackendNode, op frontendModelsDir = filepath.Clean(strings.TrimSuffix(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. + // sherpa-onnx's tokens.txt / espeak-ng-data) live beside the model, so option + // values are resolved relative to this dir as well as frontendModelsDir — + // letting a shared config declare them with bare names regardless of whether + // Model includes a subdirectory. + localModelDir := "" + if opts.ModelFile != "" { + localModelDir = filepath.Dir(opts.ModelFile) + } + // keyMapper generates storage keys namespaced under trackingKey, preserving // subdirectory structure relative to frontendModelsDir. This ensures: // 1. All files for a model land in one directory on the worker for clean deletion @@ -1079,8 +1090,8 @@ func (r *SmartRouter) stageModelFiles(ctx context.Context, node *BackendNode, op // Stage file paths referenced in generic Options (key:value pairs where values // are file paths). Options stay as relative paths — backends resolve them via ModelPath. - r.stageGenericOptions(ctx, node, opts.Options, frontendModelsDir, keyMapper.Key) - r.stageGenericOptions(ctx, node, opts.Overrides, frontendModelsDir, keyMapper.Key) + r.stageGenericOptions(ctx, node, opts.Options, frontendModelsDir, localModelDir, keyMapper.Key) + r.stageGenericOptions(ctx, node, opts.Overrides, frontendModelsDir, localModelDir, keyMapper.Key) return opts, nil } @@ -1196,36 +1207,86 @@ func (r *SmartRouter) stageCompanionFiles(ctx context.Context, node *BackendNode } // stageGenericOptions iterates key:value option strings and stages any values -// that resolve to existing files relative to the frontend models directory. -// Option values are NOT rewritten — backends resolve them via ModelPath. -// keyFn generates the namespaced storage key for each file path. -func (r *SmartRouter) stageGenericOptions(ctx context.Context, node *BackendNode, options []string, frontendModelsDir string, keyFn func(string) string) { +// that resolve to existing files relative to the frontend models directory or +// the model's own directory. Option values are NOT rewritten — backends resolve +// them via ModelPath. keyFn generates the namespaced storage key for each file. +func (r *SmartRouter) stageGenericOptions(ctx context.Context, node *BackendNode, options []string, frontendModelsDir, modelDir string, keyFn func(string) string) { for _, opt := range options { optKey, val, ok := strings.Cut(opt, ":") if !ok || val == "" { continue } - // Check if value is an existing file path (absolute or relative to frontend models dir) - absPath := val - if !filepath.IsAbs(val) && frontendModelsDir != "" { - absPath = filepath.Join(frontendModelsDir, val) + // Resolve the value to an existing path: absolute as-is, otherwise + // relative to frontendModelsDir first, then the model's own directory + // (where backends like sherpa-onnx keep companion assets such as + // tokens.txt and espeak-ng-data). + absPath, ok := resolveOptionPath(val, frontendModelsDir, modelDir) + if !ok { + continue } - if _, err := os.Stat(absPath); os.IsNotExist(err) { + info, err := os.Stat(absPath) + if err != nil { + continue + } + + // A directory option value (e.g. sherpa-onnx's espeak-ng-data) is staged + // file-by-file so the whole tree is recreated beside the model on the + // worker; a single file is staged directly. Values are never rewritten — + // backends resolve relative paths via ModelPath. + if err == nil && info.IsDir() { + r.stageOptionDir(ctx, node, absPath, keyFn) + xlog.Debug("Staged option directory", "option", optKey, "localPath", absPath) continue } - // Stage the file to the worker using the namespaced key key := keyFn(absPath) if _, err := r.fileStager.EnsureRemote(ctx, node.ID, absPath, key); err != nil { xlog.Warn("Failed to stage option file, skipping", "option", opt, "path", absPath, "error", err) continue } - // Leave option value unchanged — backend resolves relative paths via ModelPath xlog.Debug("Staged option file", "option", optKey, "localPath", absPath) } } +// resolveOptionPath finds an existing local path for an option value: an +// absolute path as-is, otherwise relative to frontendModelsDir, then to the +// model's own directory. Returns false when none exists. +func resolveOptionPath(val, frontendModelsDir, modelDir string) (string, bool) { + if filepath.IsAbs(val) { + if _, err := os.Stat(val); err == nil { + return val, true + } + return "", false + } + for _, base := range []string{frontendModelsDir, modelDir} { + if base == "" { + continue + } + p := filepath.Join(base, val) + if _, err := os.Stat(p); err == nil { + return p, true + } + } + return "", false +} + +// stageOptionDir stages every regular file under an option-declared directory +// (e.g. sherpa-onnx's espeak-ng-data) using the structure-preserving key, so the +// tree is recreated beside the model on the worker. Per-file errors are logged +// and skipped; the option value itself is not rewritten. +func (r *SmartRouter) stageOptionDir(ctx context.Context, node *BackendNode, dir string, keyFn func(string) string) { + _ = filepath.WalkDir(dir, func(path string, d fs.DirEntry, walkErr error) error { + if walkErr != nil || d.IsDir() { + return nil + } + if _, err := r.fileStager.EnsureRemote(ctx, node.ID, path, keyFn(path)); err != nil { + xlog.Warn("Failed to stage option directory file, skipping", "path", path, "error", err) + } + return nil + }) +} + // probeHealth checks whether a backend process on the given node/addr is alive // via a gRPC health check with a 2-second timeout. The client is closed after // the check. diff --git a/core/services/nodes/router_optionstage_test.go b/core/services/nodes/router_optionstage_test.go new file mode 100644 index 000000000..42ebb602e --- /dev/null +++ b/core/services/nodes/router_optionstage_test.go @@ -0,0 +1,77 @@ +package nodes + +import ( + "context" + "os" + "path/filepath" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + pb "github.com/mudler/LocalAI/pkg/grpc/proto" +) + +// These tests cover staging of companion assets declared as option file paths +// (the "vae_path:..." convention). Backends like sherpa-onnx keep a single-file +// ModelFile (the .onnx) but resolve sibling assets — tokens.txt and the +// espeak-ng-data directory — relative to the model dir. Those siblings must be +// shipped to remote workers too, including directory-valued options expanded +// file-by-file. +var _ = Describe("stageGenericOptions companion assets", func() { + var ( + stager *fakeFileStager + router *SmartRouter + node *BackendNode + tmp string + ) + + BeforeEach(func() { + stager = &fakeFileStager{} + router = &SmartRouter{ + fileStager: stager, + stagingTracker: NewStagingTracker(), + } + node = &BackendNode{ID: "node-1", Name: "node-1", Address: "10.0.0.1:50051"} + tmp = GinkgoT().TempDir() + }) + + It("stages option-declared sibling files and expands directory options", func() { + modelRel := "vits-piper-it_IT-paola-medium" + modelDir := filepath.Join(tmp, "models", modelRel) + dataDir := filepath.Join(modelDir, "espeak-ng-data") + Expect(os.MkdirAll(filepath.Join(dataDir, "lang"), 0o755)).To(Succeed()) + + onnx := filepath.Join(modelDir, "it_IT-paola-medium.onnx") + tokens := filepath.Join(modelDir, "tokens.txt") + phontab := filepath.Join(dataDir, "phontab") + langIt := filepath.Join(dataDir, "lang", "it") + for _, f := range []string{onnx, tokens, phontab, langIt} { + Expect(os.WriteFile(f, []byte("x"), 0o644)).To(Succeed()) + } + + opts := &pb.ModelOptions{ + Model: filepath.Join(modelRel, "it_IT-paola-medium.onnx"), + ModelFile: onnx, + // Bare names: not found under the models root (Model includes a + // subdir), so they must resolve relative to the model's own dir. + Options: []string{ + "tts.noise_scale=0.667", // not a path; ignored by staging + "tokens:tokens.txt", + "data_dir:espeak-ng-data", + }, + } + + _, err := router.stageModelFiles(context.Background(), node, opts, "track-key") + Expect(err).ToNot(HaveOccurred()) + + staged := make([]string, 0, len(stager.ensureCalls)) + for _, c := range stager.ensureCalls { + staged = append(staged, c.localPath) + } + // The .onnx (ModelFile), the tokens.txt file option, and every file under + // the espeak-ng-data directory option are staged; the directory path + // itself is never handed to the stager. + Expect(staged).To(ContainElements(onnx, tokens, phontab, langIt)) + Expect(staged).ToNot(ContainElement(dataDir)) + }) +}) diff --git a/gallery/sherpa-onnx-tts.yaml b/gallery/sherpa-onnx-tts.yaml index e6bdb1f4b..84dc25484 100644 --- a/gallery/sherpa-onnx-tts.yaml +++ b/gallery/sherpa-onnx-tts.yaml @@ -12,3 +12,17 @@ config_file: | # Speech rate multiplier. Applied at every TTS / TTSStream call # since the TTSRequest proto has no speed field. - tts.speed=1.0 + # Companion assets that sherpa-onnx TTS voices load from beside the .onnx + # (tokens, lexicons, espeak-ng phonemization data, Kokoro voices bank / jieba + # dict). Declared as option paths so distributed inference stages them to + # remote worker nodes too; the backend ignores these keys and resolves the + # files relative to the model dir. Bare names resolve against the model's own + # directory; any that a given voice doesn't ship are skipped during staging. + - tokens:tokens.txt + - lexicon:lexicon.txt + - data_dir:espeak-ng-data + - voices:voices.bin + - dict_dir:dict + - lexicon_us:lexicon-us-en.txt + - lexicon_gb:lexicon-gb-en.txt + - lexicon_zh:lexicon-zh.txt