fix(distributed): stage backend companion assets to remote nodes (#10330)

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 <mudler@localai.io>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
LocalAI [bot]
2026-06-14 16:42:59 +02:00
committed by GitHub
parent 4d3d54d61b
commit e5c95e0449
3 changed files with 165 additions and 13 deletions

View File

@@ -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.

View File

@@ -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))
})
})

View File

@@ -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