Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
87 changes: 74 additions & 13 deletions core/services/nodes/router.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
}
Expand Down Expand Up @@ -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
}
info, err := os.Stat(absPath)
if err != nil {
continue
}
if _, err := os.Stat(absPath); os.IsNotExist(err) {

// 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.
Expand Down
77 changes: 77 additions & 0 deletions core/services/nodes/router_optionstage_test.go
Original file line number Diff line number Diff line change
@@ -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))
})
})
14 changes: 14 additions & 0 deletions gallery/sherpa-onnx-tts.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Loading