fix(proxy): propagate model env and accept direct keys (#12940)

* feat(proxy): accept direct upstream API keys

Support direct credentials for models created in the UI. Keep env and
file references as mutually exclusive alternatives. Use model env for
frontend probes and omit credentials from structured config logs.

Assisted-by: nib:gpt-6.1-sol
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* feat(distributed): isolate model environments in worker processes

Forward complete model environment snapshots through backend installation.
Restart worker children when overrides change or are removed, and include
keyed snapshot identities in frontend install coalescing.

Cover wire compatibility, routing, real worker children, and option copying;
document rollout requirements and control-transport security.

Assisted-by: nib:gpt-6.1-sol
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* fix(worker): serialize installs by model replica

Different backend names can target the same process slot. Hold a slot
lock through startup, readiness, and operation attachment so concurrent
installs cannot reuse another configuration or an unready child.

Assisted-by: nib:gpt-6.1-sol
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* test(nodes): update endpoint stub for model env

Keep the endpoint command stub aligned with the environment-aware install
interface. Assert interface conformance so future changes fail at the
stub declaration rather than at endpoint call sites.

Assisted-by: nib:gpt-6.1-sol
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

---------

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
Co-authored-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
mudler-agentandEttore Di Giacinto authored and GitHub committed 2026-10-27 10:27:45 +01:00
1 parent 2f5ec65909
commit 3a02f29edc
33 files changed
+820 -59

No files matched your search

+3 -3
View File
@@ -544,9 +544,8 @@ message ModelOptions {
// ProxyOptions configures the cloud-proxy backend. UpstreamURL and
// Mode are always meaningful; Provider only matters in translate mode.
// The two api_key_* fields are mutually exclusive and resolved by the
// backend at LoadModel — core forwards the references rather than the
// plaintext key.
// The api_key, api_key_env and api_key_file fields are mutually exclusive
// and resolved by the backend at LoadModel. api_key is plaintext admin config.
message ProxyOptions {
string upstream_url = 1;
string mode = 2;
@@ -561,6 +560,7 @@ message ProxyOptions {
// Cuts input cost on repeated/agentic calls (cache read = 0.1x). Only
// meaningful for mode=translate + provider=anthropic; ignored otherwise.
bool cache_prompt = 8;
string api_key = 9;
}
message Result {
@@ -0,0 +1,48 @@
// SPDX-License-Identifier: MIT
package main
import (
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
"os"
"path/filepath"
)
var _ = Describe("Proxy credential sources", func() {
It("resolves direct, absent, env and file credentials", func() {
key, err := resolveAPIKey("direct-canary", "", "")
Expect(err).NotTo(HaveOccurred())
Expect(key).To(Equal("direct-canary"))
key, err = resolveAPIKey("", "", "")
Expect(err).NotTo(HaveOccurred())
Expect(key).To(BeEmpty())
DeferCleanup(os.Unsetenv, "LOCALAI_TEST_PROXY_CREDENTIAL")
Expect(os.Setenv("LOCALAI_TEST_PROXY_CREDENTIAL", "env-canary")).To(Succeed())
key, err = resolveAPIKey("", "LOCALAI_TEST_PROXY_CREDENTIAL", "")
Expect(err).NotTo(HaveOccurred())
Expect(key).To(Equal("env-canary"))
dir, err := os.MkdirTemp("", "proxy-credentials-")
Expect(err).NotTo(HaveOccurred())
DeferCleanup(os.RemoveAll, dir)
path := filepath.Join(dir, "key")
Expect(os.WriteFile(path, []byte("file-canary\n"), 0600)).To(Succeed())
key, err = resolveAPIKey("", "", path)
Expect(err).NotTo(HaveOccurred())
Expect(key).To(Equal("file-canary"))
_, err = resolveAPIKey("", "", path+"-missing")
Expect(err).To(HaveOccurred())
})
It("rejects every conflicting source without exposing the key", func() {
for _, args := range [][3]string{{"secret-canary", "ENV", ""}, {"secret-canary", "", "/missing"}, {"", "ENV", "/missing"}, {"secret-canary", "ENV", "/missing"}} {
_, err := resolveAPIKey(args[0], args[1], args[2])
Expect(err).To(HaveOccurred())
Expect(err.Error()).NotTo(ContainSubstring("secret-canary"))
}
})
It("loads direct credentials from protobuf options", func() {
proxy := NewCloudProxy()
Expect(proxy.Load(&pb.ModelOptions{Model: "alias", Proxy: &pb.ProxyOptions{UpstreamUrl: "http://example.invalid", ApiKey: "direct-canary"}})).To(Succeed())
Expect(proxy.cfg.Load().apiKey).To(Equal("direct-canary"))
})
})
+15 -3
View File
@@ -95,7 +95,7 @@ func (c *CloudProxy) Load(opts *pb.ModelOptions) error {
return fmt.Errorf("cloud-proxy: unknown mode %q", mode)
}
key, err := resolveAPIKey(po.GetApiKeyEnv(), po.GetApiKeyFile())
key, err := resolveAPIKey(po.GetApiKey(), po.GetApiKeyEnv(), po.GetApiKeyFile())
if err != nil {
return err
}
@@ -120,8 +120,20 @@ func (c *CloudProxy) Load(opts *pb.ModelOptions) error {
// resolveAPIKey mirrors config.ProxyConfig.ResolveAPIKey. Duplicated
// (a few lines) rather than importing core/config from a backend
// binary — keeps backends independent of core's package layout.
// Mutual-exclusion is enforced upstream in core/config.Validate.
func resolveAPIKey(envName, filePath string) (string, error) {
// Also validate here for callers that bypass core config validation.
func resolveAPIKey(directKey, envName, filePath string) (string, error) {
sources := 0
for _, value := range []string{directKey, envName, filePath} {
if value != "" {
sources++
}
}
if sources > 1 {
return "", fmt.Errorf("cloud-proxy: api_key, api_key_env and api_key_file are mutually exclusive")
}
if directKey != "" {
return directKey, nil
}
if envName != "" {
v := os.Getenv(envName)
if v == "" {
@@ -0,0 +1,48 @@
// SPDX-License-Identifier: MIT
package main
import (
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
"os"
"path/filepath"
)
var _ = Describe("Proxy credential sources", func() {
It("resolves direct, absent, env and file credentials", func() {
key, err := resolveAPIKey("direct-canary", "", "")
Expect(err).NotTo(HaveOccurred())
Expect(key).To(Equal("direct-canary"))
key, err = resolveAPIKey("", "", "")
Expect(err).NotTo(HaveOccurred())
Expect(key).To(BeEmpty())
DeferCleanup(os.Unsetenv, "LOCALAI_TEST_PROXY_CREDENTIAL")
Expect(os.Setenv("LOCALAI_TEST_PROXY_CREDENTIAL", "env-canary")).To(Succeed())
key, err = resolveAPIKey("", "LOCALAI_TEST_PROXY_CREDENTIAL", "")
Expect(err).NotTo(HaveOccurred())
Expect(key).To(Equal("env-canary"))
dir, err := os.MkdirTemp("", "proxy-credentials-")
Expect(err).NotTo(HaveOccurred())
DeferCleanup(os.RemoveAll, dir)
path := filepath.Join(dir, "key")
Expect(os.WriteFile(path, []byte("file-canary\n"), 0600)).To(Succeed())
key, err = resolveAPIKey("", "", path)
Expect(err).NotTo(HaveOccurred())
Expect(key).To(Equal("file-canary"))
_, err = resolveAPIKey("", "", path+"-missing")
Expect(err).To(HaveOccurred())
})
It("rejects every conflicting source without exposing the key", func() {
for _, args := range [][3]string{{"secret-canary", "ENV", ""}, {"secret-canary", "", "/missing"}, {"", "ENV", "/missing"}, {"secret-canary", "ENV", "/missing"}} {
_, err := resolveAPIKey(args[0], args[1], args[2])
Expect(err).To(HaveOccurred())
Expect(err.Error()).NotTo(ContainSubstring("secret-canary"))
}
})
It("loads direct credentials from protobuf options", func() {
proxy := NewLocalAIProxy()
Expect(proxy.Load(&pb.ModelOptions{Model: "alias", Proxy: &pb.ProxyOptions{UpstreamUrl: "http://example.invalid", ApiKey: "direct-canary"}})).To(Succeed())
Expect(proxy.cfg.Load().apiKey).To(Equal("direct-canary"))
})
})
+14 -2
View File
@@ -92,7 +92,7 @@ func (p *LocalAIProxy) Load(opts *pb.ModelOptions) error {
"mode", po.GetMode(), "provider", po.GetProvider())
}
key, err := resolveAPIKey(po.GetApiKeyEnv(), po.GetApiKeyFile())
key, err := resolveAPIKey(po.GetApiKey(), po.GetApiKeyEnv(), po.GetApiKeyFile())
if err != nil {
return err
}
@@ -151,7 +151,19 @@ func (p *LocalAIProxy) model(req string) string {
// resolveAPIKey mirrors config.ProxyConfig.ResolveAPIKey (and cloud-proxy's
// copy). Duplicated so the backend binary does not depend on core's layout.
func resolveAPIKey(envName, filePath string) (string, error) {
func resolveAPIKey(directKey, envName, filePath string) (string, error) {
sources := 0
for _, value := range []string{directKey, envName, filePath} {
if value != "" {
sources++
}
}
if sources > 1 {
return "", fmt.Errorf("localai-proxy: api_key, api_key_env and api_key_file are mutually exclusive")
}
if directKey != "" {
return directKey, nil
}
if envName != "" {
v := os.Getenv(envName)
if v == "" {
+1
View File
@@ -547,6 +547,7 @@ func grpcModelOpts(c config.ModelConfig, modelPath string) *pb.ModelOptions {
UpstreamUrl: c.Proxy.UpstreamURL,
Mode: c.Proxy.Mode,
Provider: c.Proxy.Provider,
ApiKey: c.Proxy.APIKey,
ApiKeyEnv: c.Proxy.APIKeyEnv,
ApiKeyFile: c.Proxy.APIKeyFile,
UpstreamModel: c.Proxy.UpstreamModel,
+10
View File
@@ -469,3 +469,13 @@ var _ = Describe("grpcModelOpts option forwarding", func() {
Expect(grpcModelOpts(cfg, "/models").Options).To(Equal(cfg.Options))
})
})
var _ = Describe("Proxy credential options", func() {
It("forwards the direct credential and preserves wire field 9", func() {
threads := 1
cfg := config.ModelConfig{Threads: &threads, Backend: "cloud-proxy", Proxy: config.ProxyConfig{APIKey: "direct-canary", UpstreamURL: "http://example.invalid"}}
opts := grpcModelOpts(cfg, "")
Expect(opts.Proxy.GetApiKey()).To(Equal("direct-canary"))
Expect(int(opts.Proxy.ProtoReflect().Descriptor().Fields().ByName("api_key").Number())).To(Equal(9))
})
})
+7
View File
@@ -1204,6 +1204,13 @@ func DefaultRegistry() map[string]FieldMetaOverride {
Component: "input",
Order: 210,
},
"proxy.api_key": {
Section: "proxy",
Label: "Proxy API Key",
Description: "Plaintext administrator credential stored in model YAML/JSON. Use only one of api_key, api_key_env or api_key_file. Masking is not encryption.",
Component: "password",
Order: 211,
},
"proxy.api_key_env": {
Section: "proxy",
Label: "Proxy API Key Env Var",
+6
View File
@@ -97,3 +97,9 @@ var _ = Describe("backend-only upscale options", func() {
}
})
})
var _ = Describe("Proxy credential metadata", func() {
It("masks the direct key with a password input", func() {
Expect(meta.DefaultRegistry()["proxy.api_key"].Component).To(Equal("password"))
})
})
+38 -10
View File
@@ -281,14 +281,18 @@ type ProxyConfig struct {
// wire format there is whatever the client sent.
Provider string `yaml:"provider,omitempty" json:"provider,omitempty"`
// APIKey is a plaintext administrator-managed upstream credential.
// Mutually exclusive with APIKeyEnv and APIKeyFile.
APIKey string `yaml:"api_key,omitempty" json:"api_key,omitempty"`
// APIKeyEnv names the environment variable holding the upstream
// API key. Mutually exclusive with APIKeyFile. Both empty is
// API key. Mutually exclusive with APIKey and APIKeyFile. All empty is
// allowed (no-auth upstreams).
APIKeyEnv string `yaml:"api_key_env,omitempty" json:"api_key_env,omitempty"`
// APIKeyFile is a path to a file whose contents are the upstream
// API key. Trailing whitespace is trimmed. Mutually exclusive
// with APIKeyEnv. The integration point for K8s secret mounts,
// with APIKey and APIKeyEnv. The integration point for K8s secret mounts,
// Vault agent files, and similar external-secret workflows.
APIKeyFile string `yaml:"api_key_file,omitempty" json:"api_key_file,omitempty"`
@@ -329,12 +333,40 @@ const (
ProxyProviderAnthropic = "anthropic"
)
// ResolveAPIKey returns the upstream key from api_key_env or api_key_file, or
// LogValue uses an allowlist so future credential fields cannot leak through config logs.
func (c ModelConfig) LogValue() slog.Value {
return slog.GroupValue(slog.String("name", c.Name), slog.String("backend", c.Backend))
}
// LogValue deliberately omits credentials and potentially credential-bearing URLs.
func (p ProxyConfig) LogValue() slog.Value {
return slog.GroupValue(slog.String("mode", p.Mode), slog.String("provider", p.Provider))
}
func (p ProxyConfig) validateCredentialSources() error {
sources := 0
for _, value := range []string{p.APIKey, p.APIKeyEnv, p.APIKeyFile} {
if value != "" {
sources++
}
}
if sources > 1 {
return fmt.Errorf("proxy: api_key, api_key_env and api_key_file are mutually exclusive")
}
return nil
}
// ResolveAPIKey returns the upstream key from api_key, api_key_env or api_key_file, or
// "" when neither is set. Mirrored (not imported, to keep backends independent
// of core's package layout) by resolveAPIKey in backend/go/cloud-proxy/proxy.go
// — keep the two in sync, empty-value handling included.
// and backend/go/localai-proxy/proxy.go — keep all copies in sync.
func (p ProxyConfig) ResolveAPIKey(envLookup func(string) string) (string, error) {
if err := p.validateCredentialSources(); err != nil {
return "", err
}
switch {
case p.APIKey != "":
return p.APIKey, nil
case p.APIKeyEnv != "":
var v string
if envLookup != nil {
@@ -2058,12 +2090,8 @@ func (c *ModelConfig) Validate() (bool, error) {
}
}
// Cloud-proxy: at most one of api_key_env / api_key_file may be
// set. Both empty means no Authorization header (no-auth upstream
// or a development passthrough). The mode field accepts the empty
// string (defaults to passthrough), "passthrough", or "translate".
if c.Proxy.APIKeyEnv != "" && c.Proxy.APIKeyFile != "" {
return false, fmt.Errorf("proxy: api_key_env and api_key_file are mutually exclusive")
if err := c.Proxy.validateCredentialSources(); err != nil {
return false, err
}
switch c.Proxy.Mode {
case "", ProxyModePassthrough, ProxyModeTranslate:
+44
View File
@@ -0,0 +1,44 @@
// SPDX-License-Identifier: MIT
package config
import (
"bytes"
"encoding/json"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
"gopkg.in/yaml.v3"
"log/slog"
)
var _ = Describe("Proxy credentials", func() {
It("preserves and resolves a direct admin credential", func() {
var p ProxyConfig
Expect(json.Unmarshal([]byte(`{"api_key":"fake-direct-key"}`), &p)).To(Succeed())
key, err := p.ResolveAPIKey(nil)
Expect(err).NotTo(HaveOccurred())
Expect(key).To(Equal("fake-direct-key"))
encoded, err := json.Marshal(p)
Expect(err).NotTo(HaveOccurred())
Expect(string(encoded)).To(ContainSubstring(`"api_key":"fake-direct-key"`))
})
It("rejects direct credentials combined with another source", func() {
var c ModelConfig
Expect(json.Unmarshal([]byte(`{"proxy":{"api_key":"fake-key","api_key_env":"OTHER"}}`), &c)).To(Succeed())
_, err := c.Validate()
Expect(err).To(HaveOccurred())
})
})
var _ = Describe("Proxy credential logging", func() {
It("keeps secrets out of structured config logs without changing storage", func() {
var output bytes.Buffer
logger := slog.New(slog.NewJSONHandler(&output, nil))
c := ModelConfig{Proxy: ProxyConfig{APIKey: "fake-direct-secret"}, Environment: map[string]string{"KEY": "fake-env-secret"}}
logger.Info("config", "config", c, "proxy", c.Proxy)
Expect(output.String()).NotTo(ContainSubstring("fake-direct-secret"))
Expect(output.String()).NotTo(ContainSubstring("fake-env-secret"))
data, err := yaml.Marshal(c)
Expect(err).NotTo(HaveOccurred())
Expect(string(data)).To(ContainSubstring("fake-direct-secret"))
})
})
@@ -22,7 +22,9 @@ type stubNodeCommandSender struct {
listBackendsCalled bool
}
func (s *stubNodeCommandSender) InstallBackend(_, _, _, _, _, _, _ string, _ int, _ string, _ func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendInstallReply, error) {
var _ nodes.NodeCommandSender = (*stubNodeCommandSender)(nil)
func (s *stubNodeCommandSender) InstallBackend(_, _, _, _, _, _, _ string, _ int, _ string, _ func(workerctl.BackendInstallProgressEvent), _ ...map[string]string) (*workerctl.BackendInstallReply, error) {
return &workerctl.BackendInstallReply{}, nil
}
@@ -46,7 +48,7 @@ func (s *stubNodeCommandSender) UnloadModelOnNode(_, _ string) error { return ni
func (s *stubNodeCommandSender) PingNode(_ string) error { return nil }
// The load operation verbs. This stub never reaches them.
func (s *stubNodeCommandSender) InstallBackendOp(_, _, _, _ string, _ int, _, _ string, _ time.Duration, _ func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendInstallReply, error) {
func (s *stubNodeCommandSender) InstallBackendOp(_, _, _, _ string, _ int, _, _ string, _ time.Duration, _ func(workerctl.BackendInstallProgressEvent), _ ...map[string]string) (*workerctl.BackendInstallReply, error) {
return &workerctl.BackendInstallReply{}, nil
}
@@ -498,7 +498,8 @@ export default function ConfigFieldRenderer({ field, value, onChange, onRemove,
// Default: text input
return (
<SettingRow label={<FieldLabel field={field} />} description={description}>
<input className="input cfr-select" value={value ?? ''}
<input className="input cfr-select" type={component === 'password' ? 'password' : 'text'}
autoComplete={component === 'password' ? 'new-password' : undefined} value={value ?? ''}
onChange={e => handleChange(e.target.value)}
placeholder={field.placeholder}
/>
@@ -0,0 +1,36 @@
// SPDX-License-Identifier: MIT
import test from 'node:test'
import assert from 'node:assert/strict'
import { readFileSync } from 'node:fs'
import * as React from 'react'
import { renderToStaticMarkup } from 'react-dom/server'
import { transformWithOxc } from 'vite'
test('config password input preserves the real value on rerender and edits', async () => {
// Exercise the renderer's actual default-input JSX without loading unrelated
// editors (CodeMirror and model selectors require a browser).
const source = readFileSync(new URL('../../src/components/ConfigFieldRenderer.jsx', import.meta.url), 'utf8')
const input = source.slice(source.lastIndexOf('<input'), source.lastIndexOf('/>') + 2)
const { code } = await transformWithOxc(`return (${input})`, 'field.jsx', { jsx: { runtime: 'classic' } })
const renderInput = new Function('React', 'component', 'value', 'handleChange', 'field', code)
const secret = 'fake-secret<&"'
const values = { 'proxy.api_key': secret, name: 'before' }
const render = () => renderInput(React, 'password', values['proxy.api_key'], value => {
values['proxy.api_key'] = value
}, {})
const first = render()
assert.equal(first.props.type, 'password')
assert.equal(first.props.autoComplete, 'new-password')
assert.equal(first.props.value, secret)
assert.match(renderToStaticMarkup(first), /type="password"/)
values.name = 'after'
assert.equal(render().props.value, secret)
assert.equal(JSON.parse(JSON.stringify(values))['proxy.api_key'], secret)
render().props.onChange({ target: { value: 'replacement-secret' } })
assert.equal(render().props.value, 'replacement-secret')
render().props.onChange({ target: { value: '' } })
assert.equal(render().props.value, '')
assert.equal(renderInput(React, 'password', undefined, () => {}, {}).props.value, '')
assert.equal(renderInput(React, 'input', 'plain', () => {}, {}).props.type, 'text')
})
@@ -0,0 +1,34 @@
// SPDX-License-Identifier: MIT
package failover
import (
"github.com/mudler/LocalAI/core/config"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
"net/http"
)
var _ = Describe("Model env credentials", func() {
It("authorizes with a direct credential or frontend fallback", func() {
p := &DefaultProber{EnvLookup: func(string) string { return "frontend-key" }}
req, err := http.NewRequest("GET", "http://example.invalid", nil)
Expect(err).NotTo(HaveOccurred())
c := config.ModelConfig{Proxy: config.ProxyConfig{APIKey: "direct-key"}}
Expect(p.authorize(req, c)).To(Succeed())
Expect(req.Header.Get("Authorization")).To(Equal("Bearer direct-key"))
c.Proxy = config.ProxyConfig{APIKeyEnv: "KEY"}
Expect(p.authorize(req, c)).To(Succeed())
Expect(req.Header.Get("Authorization")).To(Equal("Bearer frontend-key"))
})
It("uses model env before the frontend env, including an explicit empty value", func() {
p := &DefaultProber{EnvLookup: func(string) string { return "frontend-key" }}
c := config.ModelConfig{Environment: map[string]string{"ANTHROPIC_API_KEY": "model-key"}}
c.Proxy.APIKeyEnv = "ANTHROPIC_API_KEY"
req, err := http.NewRequest("GET", "http://example.invalid", nil)
Expect(err).NotTo(HaveOccurred())
Expect(p.authorize(req, c)).To(Succeed())
Expect(req.Header.Get("Authorization")).To(Equal("Bearer model-key"))
c.Environment["ANTHROPIC_API_KEY"] = ""
Expect(p.authorize(req, c)).NotTo(Succeed())
})
})
+11 -2
View File
@@ -110,12 +110,21 @@ func PrepareTarget(cfg *config.ModelConfig) {
}
func (p *DefaultProber) authorize(req *http.Request, cfg config.ModelConfig) error {
if cfg.Proxy.APIKeyEnv != "" && p.EnvLookup == nil {
_, hasModelEnv := cfg.Environment[cfg.Proxy.APIKeyEnv]
if cfg.Proxy.APIKeyEnv != "" && !hasModelEnv && p.EnvLookup == nil {
if _, warned := p.unresolvableEnvWarned.LoadOrStore(cfg.Proxy.APIKeyEnv, struct{}{}); !warned {
xlog.Warn("failover: remote target has api_key_env set but no credential lookup is configured; liveness and requests will fail authorization", "target", cfg.Name, "api_key_env", cfg.Proxy.APIKeyEnv)
}
}
key, err := cfg.Proxy.ResolveAPIKey(p.EnvLookup)
key, err := cfg.Proxy.ResolveAPIKey(func(name string) string {
if value, ok := cfg.Environment[name]; ok {
return value
}
if p.EnvLookup != nil {
return p.EnvLookup(name)
}
return ""
})
if err != nil || key == "" {
return err
}
+60
View File
@@ -0,0 +1,60 @@
// SPDX-License-Identifier: MIT
package nodes
import (
"context"
"github.com/mudler/LocalAI/core/services/workerctl"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
"time"
)
var _ = Describe("Router install environment", func() {
It("keeps changed snapshots out of an in-flight identical slot", func() {
started := make(chan struct{}, 2)
release := make(chan struct{})
defer close(release)
unloader := &fakeUnloader{installReply: &workerctl.BackendInstallReply{Success: true, Address: "127.0.0.1:1234"}, installHook: func() { started <- struct{}{}; <-release }}
router := NewSmartRouter(&fakeModelRouter{}, SmartRouterOptions{Unloader: unloader})
done := make(chan error, 2)
for _, value := range []string{"first", "second"} {
go func(value string) {
_, err := router.installBackendOnNode(context.Background(), &BackendNode{ID: "node"}, "backend", "model", 0, map[string]string{"TOKEN": value})
done <- err
}(value)
}
Eventually(started).Should(Receive())
Eventually(started).Should(Receive())
// Both leaders entered before either could finish, so they cannot have shared
// one install reply (and therefore one child's credentials).
go func() { release <- struct{}{}; release <- struct{}{} }()
Eventually(done).Should(Receive(BeNil()))
Eventually(done).Should(Receive(BeNil()))
Expect(unloader.installCalls).To(HaveLen(2))
Expect([]string{unloader.installCalls[0].env["TOKEN"], unloader.installCalls[1].env["TOKEN"]}).To(ConsistOf("first", "second"))
})
It("coalesces identical snapshots without placing credentials in the identity", func() {
a := map[string]string{"TOKEN": "test-value", "EMPTY": ""}
b := map[string]string{"EMPTY": "", "TOKEN": "test-value"}
Expect(installEnvIdentity(a)).To(Equal(installEnvIdentity(b)))
Expect(installEnvIdentity(a)).NotTo(ContainSubstring("test-value"))
Expect(installEnvIdentity(nil)).NotTo(Equal(installEnvIdentity(map[string]string{})))
})
It("forwards complete and explicitly empty snapshots through the operation-aware fake", func() {
fake := &fakeOpWorker{fakeUnloader: &fakeUnloader{}}
for _, env := range []map[string]string{{"TOKEN": "test-value"}, {}} {
_, err := fake.InstallBackendOp("node", "backend", "model", "", 0, "", "operation", time.Minute, nil, env)
Expect(err).NotTo(HaveOccurred())
Expect(fake.installs[len(fake.installs)-1].EnvVars).To(Equal(env))
Expect(fake.installCalls[len(fake.installCalls)-1].env).To(Equal(env))
}
})
It("copies snapshots and preserves absence separately from removal", func() {
Expect(installEnvSnapshot(nil)).To(BeNil())
Expect(installEnvSnapshot([]map[string]string{nil})).NotTo(BeNil())
env := map[string]string{"TOKEN": "test-value"}
snapshot := installEnvSnapshot([]map[string]string{env})
env["TOKEN"] = "changed"
Expect(snapshot["TOKEN"]).To(Equal("test-value"))
})
})
+1 -1
View File
@@ -88,7 +88,7 @@ type ReplicaUnloader interface {
// LoadOperationInstaller starts a backend as a load operation the worker bounds.
type LoadOperationInstaller interface {
InstallBackendOp(nodeID, backendType, modelID, galleriesJSON string, replicaIndex int, opID, operationID string, deadline time.Duration, onProgress func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendInstallReply, error)
InstallBackendOp(nodeID, backendType, modelID, galleriesJSON string, replicaIndex int, opID, operationID string, deadline time.Duration, onProgress func(workerctl.BackendInstallProgressEvent), env ...map[string]string) (*workerctl.BackendInstallReply, error)
}
// LoadOperationStopper is the only path that kills remote load work. It stops
+3 -3
View File
@@ -50,11 +50,11 @@ func (w *fakeOpWorker) StopModelReplica(_ context.Context, _ string, replica Nod
return workerctl.ModelStopReply{Matched: true, Terminated: true}, nil
}
func (w *fakeOpWorker) InstallBackendOp(nodeID, backend, modelID, galleries string, replica int, opID, operationID string, deadline time.Duration, progress func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendInstallReply, error) {
func (w *fakeOpWorker) InstallBackendOp(nodeID, backend, modelID, galleries string, replica int, opID, operationID string, deadline time.Duration, progress func(workerctl.BackendInstallProgressEvent), env ...map[string]string) (*workerctl.BackendInstallReply, error) {
w.mu.Lock()
w.installs = append(w.installs, workerctl.BackendInstallRequest{ModelID: modelID, ReplicaIndex: int32(replica), OperationID: operationID, DeadlineMs: deadline.Milliseconds()})
w.installs = append(w.installs, workerctl.BackendInstallRequest{EnvVars: installEnvSnapshot(env), ModelID: modelID, ReplicaIndex: int32(replica), OperationID: operationID, DeadlineMs: deadline.Milliseconds()})
w.mu.Unlock()
reply, err := w.fakeUnloader.InstallBackend(nodeID, backend, modelID, galleries, "", "", "", replica, opID, progress)
reply, err := w.fakeUnloader.InstallBackend(nodeID, backend, modelID, galleries, "", "", "", replica, opID, progress, env...)
if reply != nil {
copy := *reply
copy.ProcessInstance, copy.ReportsOperations = "", false
+26 -5
View File
@@ -2,6 +2,9 @@ package nodes
import (
"context"
"crypto/hmac"
"crypto/rand"
"crypto/sha256"
"encoding/json"
"errors"
"fmt"
@@ -1275,7 +1278,7 @@ func (r *SmartRouter) scheduleNewModel(ctx context.Context, backendType, modelID
// Send backend.install — the worker installs the backend if needed and
// starts the gRPC process bound to a port for this (model, replica) slot.
addr, installErr := r.installBackendOnNode(ctx, node, backendType, modelID, replicaIdx)
addr, installErr := r.installBackendOnNode(ctx, node, backendType, modelID, replicaIdx, modelOpts.GetEnvVars())
if installErr != nil {
// Roll back the reservation explicitly so the column is accurate
// before the next heartbeat. Best-effort.
@@ -1388,13 +1391,14 @@ func (r *SmartRouter) estimateModelVRAM(ctx context.Context, opts *pb.ModelOptio
// Routine load: the worker's fast-path "already running → return current
// address" is correct here. Upgrades go through
// DistributedBackendManager.UpgradeBackend on the backend.upgrade subject.
func (r *SmartRouter) installBackendOnNode(ctx context.Context, node *BackendNode, backendType, modelID string, replicaIndex int) (string, error) {
func (r *SmartRouter) installBackendOnNode(ctx context.Context, node *BackendNode, backendType, modelID string, replicaIndex int, env ...map[string]string) (string, error) {
if r.unloader == nil {
return "", fmt.Errorf("no NATS connection for backend installation")
}
snapshot := installEnvSnapshot(env)
reportLoadPhase(ctx, LoadJobStateInstalling, node, replicaIndex)
key := fmt.Sprintf("%s|%s|%s|%d", node.ID, backendType, modelID, replicaIndex)
key := fmt.Sprintf("%s|%s|%s|%d|%s", node.ID, backendType, modelID, replicaIndex, installEnvIdentity(snapshot))
// DoChan rather than Do so this wait honors ctx cancellation. InstallBackend
// blocks for its full NATS deadline (15m by default) when a worker accepts
// the request but never replies (e.g. it died mid-install). Without ctx
@@ -1410,9 +1414,9 @@ func (r *SmartRouter) installBackendOnNode(ctx context.Context, node *BackendNod
var reply *workerctl.BackendInstallReply
var err error
if owned {
reply, err = r.unloader.InstallBackendOp(node.ID, backendType, modelID, r.galleriesJSON, replicaIndex, "", ref.Generation, r.loadOperationDeadline(), nil)
reply, err = r.unloader.InstallBackendOp(node.ID, backendType, modelID, r.galleriesJSON, replicaIndex, "", ref.Generation, r.loadOperationDeadline(), nil, snapshot)
} else {
reply, err = r.unloader.InstallBackend(node.ID, backendType, modelID, r.galleriesJSON, "", "", "", replicaIndex, "", nil)
reply, err = r.unloader.InstallBackend(node.ID, backendType, modelID, r.galleriesJSON, "", "", "", replicaIndex, "", nil, snapshot)
}
if err != nil {
return "", err
@@ -2262,3 +2266,20 @@ func (r *SmartRouter) loadOperationDeadline() time.Duration {
}
return modelLoadAbsoluteMax
}
// Keyed identities let identical installs coalesce without retaining credentials
// in singleflight keys or exposing a dictionary-testable credential hash.
var installEnvKey = func() []byte {
key := make([]byte, 32)
if _, err := rand.Read(key); err != nil {
panic(err)
}
return key
}()
func installEnvIdentity(env map[string]string) string {
body, _ := json.Marshal(env)
mac := hmac.New(sha256.New, installEnvKey)
mac.Write(body)
return fmt.Sprintf("%x", mac.Sum(nil))
}
+6 -5
View File
@@ -573,6 +573,7 @@ type fakeUnloader struct {
// every call so tests can verify both presence and shape (e.g. that backend
// is non-empty).
type installCall struct {
env map[string]string
nodeID string
backend string
modelID string
@@ -585,7 +586,7 @@ type upgradeCall struct {
replica int
}
func (f *fakeUnloader) InstallBackend(nodeID, backend, modelID, _, _, _, _ string, replica int, _ string, _ func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendInstallReply, error) {
func (f *fakeUnloader) InstallBackend(nodeID, backend, modelID, _, _, _, _ string, replica int, _ string, _ func(workerctl.BackendInstallProgressEvent), env ...map[string]string) (*workerctl.BackendInstallReply, error) {
// installHook intentionally runs OUTSIDE the mutex: the hook may block
// on a channel and we don't want to serialize concurrent callers,
// which would defeat the singleflight-overlap test.
@@ -593,15 +594,15 @@ func (f *fakeUnloader) InstallBackend(nodeID, backend, modelID, _, _, _, _ strin
f.installHook()
}
f.mu.Lock()
f.installCalls = append(f.installCalls, installCall{nodeID, backend, modelID, replica})
f.installCalls = append(f.installCalls, installCall{nodeID: nodeID, backend: backend, modelID: modelID, replica: replica, env: installEnvSnapshot(env)})
f.mu.Unlock()
return f.installReply, f.installErr
}
// The load operation verbs of the carrier seam. The default fake is a worker
// that names operations and acknowledges every stop.
func (f *fakeUnloader) InstallBackendOp(nodeID, backend, modelID, galleries string, replica int, opID, _ string, _ time.Duration, progress func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendInstallReply, error) {
reply, err := f.InstallBackend(nodeID, backend, modelID, galleries, "", "", "", replica, opID, progress)
func (f *fakeUnloader) InstallBackendOp(nodeID, backend, modelID, galleries string, replica int, opID, _ string, _ time.Duration, progress func(workerctl.BackendInstallProgressEvent), env ...map[string]string) (*workerctl.BackendInstallReply, error) {
reply, err := f.InstallBackend(nodeID, backend, modelID, galleries, "", "", "", replica, opID, progress, env...)
if reply != nil {
withOps := *reply
withOps.ReportsOperations = true
@@ -1406,7 +1407,7 @@ var _ = Describe("SmartRouter", func() {
done := make(chan error, 5)
for i := 0; i < 5; i++ {
go func() {
_, err := router.installBackendOnNode(context.Background(), node, "llama-cpp", "my-model", 0)
_, err := router.installBackendOnNode(context.Background(), node, "llama-cpp", "my-model", 0, map[string]string{"TOKEN": "same-snapshot"})
done <- err
}()
}
+18 -2
View File
@@ -31,7 +31,7 @@ import (
// UpgradeBackend returns ErrNoRoute on an old worker that does not serve
// backend.upgrade, and the caller falls back to the legacy install.
type NodeCommandSender interface {
InstallBackend(nodeID, backendType, modelID, galleriesJSON, uri, name, alias string, replicaIndex int, opID string, onProgress func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendInstallReply, error)
InstallBackend(nodeID, backendType, modelID, galleriesJSON, uri, name, alias string, replicaIndex int, opID string, onProgress func(workerctl.BackendInstallProgressEvent), env ...map[string]string) (*workerctl.BackendInstallReply, error)
UpgradeBackend(nodeID, backendType, galleriesJSON, uri, name, alias string, replicaIndex int, opID string, onProgress func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendUpgradeReply, error)
DeleteBackend(nodeID, backendName string) (*workerctl.BackendDeleteReply, error)
ListBackends(nodeID string) (*workerctl.BackendListReply, error)
@@ -158,8 +158,9 @@ const operationControlTimeout = 5 * time.Second
// InstallBackendOp is InstallBackend for a load operation: the request carries
// the operation id and the longest the load may run, so the worker can bound it.
func (a *RemoteUnloaderAdapter) InstallBackendOp(nodeID, backendType, modelID, galleriesJSON string, replicaIndex int, opID, operationID string, deadline time.Duration, onProgress func(workerctl.BackendInstallProgressEvent)) (*workerctl.BackendInstallReply, error) {
func (a *RemoteUnloaderAdapter) InstallBackendOp(nodeID, backendType, modelID, galleriesJSON string, replicaIndex int, opID, operationID string, deadline time.Duration, onProgress func(workerctl.BackendInstallProgressEvent), env ...map[string]string) (*workerctl.BackendInstallReply, error) {
return a.installBackend(nodeID, workerctl.BackendInstallRequest{
EnvVars: installEnvSnapshot(env),
Backend: backendType,
ModelID: modelID,
BackendGalleries: galleriesJSON,
@@ -257,8 +258,10 @@ func (a *RemoteUnloaderAdapter) InstallBackend(
replicaIndex int,
opID string,
onProgress func(workerctl.BackendInstallProgressEvent),
env ...map[string]string,
) (*workerctl.BackendInstallReply, error) {
return a.installBackend(nodeID, workerctl.BackendInstallRequest{
EnvVars: installEnvSnapshot(env),
Backend: backendType,
ModelID: modelID,
BackendGalleries: galleriesJSON,
@@ -672,3 +675,16 @@ func (a *RemoteUnloaderAdapter) StopNode(nodeID string) error {
subject := messaging.SubjectNodeStop(nodeID)
return a.nats.Publish(subject, nil)
}
// installEnvSnapshot preserves nil for legacy callers and an explicit empty
// snapshot for new callers that remove their last model override.
func installEnvSnapshot(env []map[string]string) map[string]string {
if len(env) == 0 {
return nil
}
result := make(map[string]string, len(env[0]))
for k, v := range env[0] {
result[k] = v
}
return result
}
+187
View File
@@ -0,0 +1,187 @@
// SPDX-License-Identifier: MIT
package worker
import (
"context"
"fmt"
"net"
"os"
"path/filepath"
"strconv"
"strings"
"time"
"github.com/mudler/LocalAI/core/services/workerctl"
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
"github.com/mudler/LocalAI/pkg/model"
"github.com/mudler/LocalAI/pkg/system"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
"google.golang.org/grpc"
)
// The real worker starts this test executable as a tiny gRPC backend. Only the
// child's explicit environment enables helper mode; the test runner is untouched.
func init() {
if os.Getenv("LOCALAI_ENV_TEST_CHILD") != "1" {
return
}
var addr string
for i, arg := range os.Args {
if arg == "--addr" && i+1 < len(os.Args) {
addr = os.Args[i+1]
}
}
listener, err := net.Listen("tcp", addr)
if err != nil {
os.Exit(2)
}
value, present := os.LookupEnv("LOCALAI_ENV_TEST_VALUE")
if err := os.WriteFile(os.Getenv("LOCALAI_ENV_TEST_OUTPUT"), []byte(fmt.Sprintf("%t:%s", present, value)), 0600); err != nil {
os.Exit(3)
}
// A filesystem gate lets the parent hold a real child in startup while
// concurrent install requests arrive, without mutating the runner's env.
if gate := os.Getenv("LOCALAI_ENV_TEST_GATE"); gate != "" {
for {
if _, err := os.Stat(gate); err == nil {
break
}
time.Sleep(10 * time.Millisecond)
}
}
server := grpc.NewServer()
pb.RegisterBackendServer(server, envTestBackend{})
if server.Serve(listener) != nil {
os.Exit(4)
}
os.Exit(0)
}
type envTestBackend struct{ pb.UnimplementedBackendServer }
func (envTestBackend) Health(context.Context, *pb.HealthMessage) (*pb.Reply, error) {
return &pb.Reply{Message: []byte("OK")}, nil
}
func (envTestBackend) Free(context.Context, *pb.HealthMessage) (*pb.Result, error) {
return &pb.Result{Success: true}, nil
}
var _ = Describe("Worker model environment", func() {
It("serializes concurrent serveInstall requests for the exact slot across backend names", func() {
root := GinkgoT().TempDir()
exe, err := os.Executable()
Expect(err).NotTo(HaveOccurred())
for _, name := range []string{"first", "second", "sibling"} {
Expect(os.WriteFile(filepath.Join(root, name), []byte("#!/bin/sh\nexec "+strconv.Quote(exe)+" \"$@\"\n"), 0700)).To(Succeed())
}
state, err := system.GetSystemState(system.WithBackendPath(root), system.WithBackendSystemPath(root), system.WithModelPath(root))
Expect(err).NotTo(HaveOccurred())
sup := &backendSupervisor{cfg: &Config{BackendsPath: root}, ml: model.NewModelLoader(state), processes: map[string]*backendProcess{}, minPort: 32000, maxPort: 42000, nextPort: 32000}
gate := filepath.Join(root, "ready")
firstDone, secondDone, siblingDone := make(chan workerctl.BackendInstallReply, 1), make(chan workerctl.BackendInstallReply, 1), make(chan workerctl.BackendInstallReply, 1)
DeferCleanup(func() {
_ = os.WriteFile(gate, nil, 0600)
// Drain all handlers before inspecting/stopping their children.
Eventually(firstDone, "40s").Should(Receive())
Eventually(secondDone, "40s").Should(Receive())
Eventually(siblingDone, "40s").Should(Receive())
for key := range sup.processes {
Expect(sup.stopBackendExact(key, true)).To(Succeed())
}
})
request := func(backend string, replica int) workerctl.BackendInstallRequest {
return workerctl.BackendInstallRequest{ModelID: "model", Backend: backend, ReplicaIndex: int32(replica), OperationID: backend,
EnvVars: map[string]string{"LOCALAI_ENV_TEST_CHILD": "1", "LOCALAI_ENV_TEST_OUTPUT": filepath.Join(root, backend+"-env"), "LOCALAI_ENV_TEST_VALUE": backend}}
}
first := request("first", 0)
first.EnvVars["LOCALAI_ENV_TEST_GATE"] = gate
go func() { firstDone <- sup.serveInstall(context.Background(), first, nil) }()
Eventually(func() bool { _, err := os.Stat(filepath.Join(root, "first-env")); return err == nil }, "10s").Should(BeTrue())
sup.mu.Lock()
original := sup.processes["model#0"]
sup.mu.Unlock()
go func() { secondDone <- sup.serveInstall(context.Background(), request("second", 0), nil) }()
go func() { siblingDone <- sup.serveInstall(context.Background(), request("sibling", 1), nil) }()
// A different replica must reach readiness while replica 0 is gated.
var sibling workerctl.BackendInstallReply
Eventually(siblingDone, "10s").Should(Receive(&sibling))
siblingDone <- sibling
Expect(sibling.Success).To(BeTrue(), sibling.Error)
Consistently(func() bool {
sup.mu.Lock()
defer sup.mu.Unlock()
return sup.processes["model#0"] == original && original.proc.IsAlive() && original.operationID == ""
}, "500ms").Should(BeTrue(), "the second install must not replace or attach to the unready first child")
Consistently(secondDone, "200ms").ShouldNot(Receive())
Expect(os.WriteFile(gate, nil, 0600)).To(Succeed())
var a, b workerctl.BackendInstallReply
Eventually(firstDone, "10s").Should(Receive(&a))
firstDone <- a
Eventually(secondDone, "10s").Should(Receive(&b))
secondDone <- b
Expect(a.Success).To(BeTrue(), a.Error)
Expect(b.Success).To(BeTrue(), b.Error)
Expect(a.ProcessInstance).To(Equal(original.instance))
Expect(b.ProcessInstance).NotTo(Equal(a.ProcessInstance))
body, err := os.ReadFile(filepath.Join(root, "second-env"))
Expect(err).NotTo(HaveOccurred())
Expect(string(body)).To(Equal("true:second"))
sup.mu.Lock()
defer sup.mu.Unlock()
Expect(sup.processes["model#0"].backendName).To(Equal("second"))
Expect(sup.processes["model#0"].operationID).To(Equal("second"))
Expect(sup.operations["second"].instance).To(Equal(b.ProcessInstance))
Expect(sup.processes["model#1"].instance).To(Equal(sibling.ProcessInstance))
})
It("passes isolated snapshots to real children and restarts on replacement and removal", func() {
root := GinkgoT().TempDir()
exe, err := os.Executable()
Expect(err).NotTo(HaveOccurred())
backend := filepath.Join(root, "backend")
Expect(os.WriteFile(backend, []byte("#!/bin/sh\nexec "+strconv.Quote(exe)+" \"$@\"\n"), 0700)).To(Succeed())
state, err := system.GetSystemState(system.WithBackendPath(root), system.WithBackendSystemPath(root), system.WithModelPath(root))
Expect(err).NotTo(HaveOccurred())
sup := &backendSupervisor{cfg: &Config{}, ml: model.NewModelLoader(state), processes: map[string]*backendProcess{}, minPort: 32000, maxPort: 42000, nextPort: 32000}
DeferCleanup(func() {
for key := range sup.processes {
Expect(sup.stopBackendExact(key, true)).To(Succeed())
}
})
snapshot := func(slot, value string) map[string]string {
return map[string]string{"LOCALAI_ENV_TEST_CHILD": "1", "LOCALAI_ENV_TEST_OUTPUT": filepath.Join(root, slot), "LOCALAI_ENV_TEST_VALUE": value}
}
first := snapshot("a", "first")
_, err = sup.startBackend("a#0", "test", backend, first)
Expect(err).NotTo(HaveOccurred())
original := sup.processes["a#0"]
first["LOCALAI_ENV_TEST_VALUE"] = "caller mutation"
Expect(original.envVars["LOCALAI_ENV_TEST_VALUE"]).To(Equal("first"))
_, err = sup.startBackend("a#0", "test", backend, snapshot("a", "first"))
Expect(err).NotTo(HaveOccurred())
Expect(sup.processes["a#0"]).To(BeIdenticalTo(original))
_, err = sup.startBackend("b#0", "test", backend, snapshot("b", "second"))
Expect(err).NotTo(HaveOccurred())
body, err := os.ReadFile(filepath.Join(root, "b"))
Expect(err).NotTo(HaveOccurred())
Expect(string(body)).To(Equal("true:second"))
_, err = sup.startBackend("a#0", "test", backend, snapshot("a", "changed"))
Expect(err).NotTo(HaveOccurred())
Expect(sup.processes["a#0"]).NotTo(BeIdenticalTo(original))
removed := snapshot("a", "")
delete(removed, "LOCALAI_ENV_TEST_VALUE")
_, err = sup.startBackend("a#0", "test", backend, removed)
Expect(err).NotTo(HaveOccurred())
body, err = os.ReadFile(filepath.Join(root, "a"))
Expect(err).NotTo(HaveOccurred())
Expect(string(body)).To(Equal("false:"))
body, err = os.ReadFile(filepath.Join(root, "b"))
Expect(err).NotTo(HaveOccurred())
Expect(string(body)).To(Equal("true:second"))
Expect(strings.Join(os.Environ(), "\n")).NotTo(ContainSubstring("LOCALAI_ENV_TEST_CHILD=1"))
Expect(sup.processMatchesEnv("b#0", map[string]string{})).To(BeFalse())
Expect(sup.processMatchesEnv("b#0", nil)).To(BeTrue())
})
})
+23 -3
View File
@@ -38,7 +38,8 @@ func buildProcessKey(modelID, backend string, replicaIndex int) string {
// upgrade path; force=false is the routine load path.
//
// The caller is responsible for holding s.lockBackend(req.Backend) for
// the duration of the call so the gallery directory isn't raced.
// the duration of the call so the gallery directory isn't raced. Production
// serveInstall also holds lockInstallSlot through readiness and attachment.
//
// 1. If already running for this (model, replica) slot AND force is false,
// return existing address (the fast path used by routine load events that
@@ -73,7 +74,7 @@ func (s *backendSupervisor) installBackend(req workerctl.BackendInstallRequest,
// running" again, and the cluster loops on a dead replica forever.
if addr := s.getAddr(processKey); addr != "" {
switch {
case !s.processMatchesBackend(processKey, req.Backend):
case !s.processMatchesBackend(processKey, req.Backend) || !s.processMatchesEnv(processKey, req.EnvVars):
// The slot is held by a process that must not serve this
// backend: either it was started from a DIFFERENT backend
// (this model's previous backend was deleted, or superseded by
@@ -169,7 +170,7 @@ func (s *backendSupervisor) installBackend(req workerctl.BackendInstallRequest,
// Start the gRPC process on a new port (keyed by model, not just backend).
// req.Backend is recorded on the process so a later backend.delete/stop can
// find it by backend name.
return s.startBackend(processKey, req.Backend, backendPath)
return s.startBackend(processKey, req.Backend, backendPath, req.EnvVars)
}
// upgradeBackend stops every running process for `backend`, force-reinstalls
@@ -276,3 +277,22 @@ func (s *backendSupervisor) lockBackend(name string) func() {
m.Lock()
return m.Unlock
}
// lockInstallSlot protects one complete serveInstall transaction, including
// readiness and operation attachment. It is acquired after lockBackend. Stops
// (including upgrade's multi-slot stops) must not acquire it: they can interrupt
// startup, and must not wait on an install holding an artifact lock they need.
func (s *backendSupervisor) lockInstallSlot(key string) func() {
s.mu.Lock()
if s.installSlotLocks == nil {
s.installSlotLocks = make(map[string]*sync.Mutex)
}
m := s.installSlotLocks[key]
if m == nil {
m = &sync.Mutex{}
s.installSlotLocks[key] = m
}
s.mu.Unlock()
m.Lock()
return m.Unlock
}
+5
View File
@@ -127,6 +127,11 @@ func (s *backendSupervisor) serveInstall(_ context.Context, req workerctl.Backen
xlog.Info("Received NATS backend.install event")
release := s.lockBackend(req.Backend)
defer release()
// A model/replica process key does not include the backend name. Keep
// competing installs on that exact slot serialized through readiness and
// operation attachment, not just the gallery artifact download.
releaseSlot := s.lockInstallSlot(buildProcessKey(req.ModelID, req.Backend, int(req.ReplicaIndex)))
defer releaseSlot()
downloadCb, flush := s.downloadProgress(req.OpID, req.Backend, progress)
defer flush()
+32 -5
View File
@@ -27,9 +27,10 @@ import (
// backendProcess represents a single gRPC backend process.
type backendProcess struct {
proc *process.Process
addr string // gRPC address (host:port)
port int
envVars map[string]string
proc *process.Process
addr string // gRPC address (host:port)
port int
// serving marks a process that has answered a gRPC health check and is
// therefore actually listening on addr. It is the opening bracket of the
// lifecycle that stopping closes.
@@ -160,6 +161,10 @@ type backendSupervisor struct {
// download path doesn't race itself on the same directory.
backendLocks map[string]*sync.Mutex
// installSlotLocks serializes installs for a process key, even when their
// backend names (and therefore artifact locks) differ. Worker-local only.
installSlotLocks map[string]*sync.Mutex
// operations are the loads the watchdog bounds, by operation id. Guarded by
// mu. See operations.go.
operations map[string]*loadOperation
@@ -497,7 +502,16 @@ func (s *backendSupervisor) releasePort(port int) {
// `backend` is the process key (modelID#replica); `backendName` is the gallery
// backend it was started from, recorded so delete/stop can find this process
// by backend name later.
func (s *backendSupervisor) startBackend(backend, backendName, backendPath string) (string, error) {
func (s *backendSupervisor) startBackend(backend, backendName, backendPath string, env ...map[string]string) (string, error) {
var snapshot map[string]string
if len(env) > 0 {
snapshot = maps.Clone(env[0])
}
if !s.processMatchesEnv(backend, snapshot) {
if err := s.stopBackendExact(backend, false); err != nil {
return "", err
}
}
// Mirrors the working directory pkg/model.startProcess hands the child.
backendDir := filepath.Dir(backendPath)
@@ -537,7 +551,7 @@ func (s *backendSupervisor) startBackend(backend, backendName, backendPath strin
bindAddr := fmt.Sprintf("0.0.0.0:%d", port)
clientAddr := fmt.Sprintf("127.0.0.1:%d", port)
proc, err := s.ml.StartProcess(backendPath, backend, bindAddr, nil)
proc, err := s.ml.StartProcess(backendPath, backend, bindAddr, snapshot)
if err != nil {
s.releasePortForKey(backend, port)
s.mu.Unlock()
@@ -557,6 +571,7 @@ func (s *backendSupervisor) startBackend(backend, backendName, backendPath strin
}
s.processes[backend] = &backendProcess{
envVars: maps.Clone(snapshot),
proc: proc,
addr: clientAddr,
port: port,
@@ -1167,3 +1182,15 @@ func (s *backendSupervisor) LoadedBackendAddresses() []string {
}
return addrs
}
// An absent snapshot retains old-controller reuse semantics. An explicit
// empty snapshot must still restart a slot that carries removed overrides.
func (s *backendSupervisor) processMatchesEnv(key string, env map[string]string) bool {
if env == nil {
return true
}
s.mu.Lock()
defer s.mu.Unlock()
bp := s.processes[key]
return bp == nil || maps.Equal(bp.envVars, env)
}
+35 -3
View File
@@ -1,10 +1,15 @@
package workerctl
import "encoding/json"
// BackendInstallRequest is the payload for a backend.install control request.
type BackendInstallRequest struct {
Backend string `json:"backend"`
ModelID string `json:"model_id,omitempty"`
BackendGalleries string `json:"backend_galleries,omitempty"`
// EnvVars is a complete per-process snapshot. nil means a legacy caller;
// a non-nil empty map explicitly removes previous model overrides.
EnvVars map[string]string `json:"-"`
Backend string `json:"backend"`
ModelID string `json:"model_id,omitempty"`
BackendGalleries string `json:"backend_galleries,omitempty"`
// URI is set for external installs (OCI image, URL, or path). When non-empty
// the worker routes to InstallExternalBackend instead of the gallery lookup.
URI string `json:"uri,omitempty"`
@@ -188,3 +193,30 @@ type BackendDeleteReply struct {
// conclusion against an older worker.
ReportsStoppedProcesses bool `json:"reports_stopped_processes,omitempty"`
}
// MarshalJSON omits legacy snapshots without collapsing an explicit empty map.
func (r BackendInstallRequest) MarshalJSON() ([]byte, error) {
type wire BackendInstallRequest
var env *map[string]string
if r.EnvVars != nil {
env = &r.EnvVars
}
return json.Marshal(struct {
wire
EnvVars *map[string]string `json:"env_vars,omitempty"`
}{wire(r), env})
}
func (r *BackendInstallRequest) UnmarshalJSON(data []byte) error {
type wire BackendInstallRequest
var decoded struct {
wire
EnvVars map[string]string `json:"env_vars"`
}
if err := json.Unmarshal(data, &decoded); err != nil {
return err
}
*r = BackendInstallRequest(decoded.wire)
r.EnvVars = decoded.EnvVars
return nil
}
+29
View File
@@ -0,0 +1,29 @@
// SPDX-License-Identifier: MIT
package workerctl
import (
"encoding/json"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
)
var _ = Describe("Install environment wire compatibility", func() {
It("preserves absent legacy bytes and explicit empty snapshots", func() {
for _, body := range []string{`{"backend":"test"}`, `{"backend":"test","env_vars":{}}`, `{"backend":"test","env_vars":{"EMPTY":"","TOKEN":"test-value"}}`} {
var req BackendInstallRequest
Expect(json.Unmarshal([]byte(body), &req)).To(Succeed())
encoded, err := json.Marshal(req)
Expect(err).NotTo(HaveOccurred())
Expect(string(encoded)).To(Equal(body))
}
})
It("lets an old receiver ignore the additive snapshot", func() {
body, err := json.Marshal(BackendInstallRequest{Backend: "test", EnvVars: map[string]string{"TOKEN": "test-value"}})
Expect(err).NotTo(HaveOccurred())
var legacy struct {
Backend string `json:"backend"`
}
Expect(json.Unmarshal(body, &legacy)).To(Succeed())
Expect(legacy.Backend).To(Equal("test"))
})
})
+27
View File
@@ -1522,3 +1522,30 @@ Further enhancements, surfaced from a survey of SGLang, vLLM production-stack, R
- **Prefill/decode disaggregation routing** ([#10068](https://github.com/mudler/LocalAI/issues/10068)): route prefill and decode to separate pools with KV transfer.
- **Per-user fairness (VTC)** ([#10069](https://github.com/mudler/LocalAI/issues/10069)): balance per-user token usage against pod load.
- **Minor tuning + MCP parity** ([#10070](https://github.com/mudler/LocalAI/issues/10070)): per-model TTL override, probabilistic LRU updates, and MCP scheduling-config tool parity.
## Per-model worker environment
A model's `env` map is sent with each backend installation as a complete
per-process snapshot. The worker applies it only to that model replica's child
process, on top of the worker's inherited environment; it never changes the
frontend or worker daemon environment. Models sharing a backend executable can
therefore use different credentials or runtime settings without sharing them.
An empty string is an explicit value, not an instruction to remove a variable.
When a slot is installed again with changed overrides (including removal of the
last override), the worker restarts its child rather than reusing a process with
stale settings. Removing an override restores the worker's inherited value, if
one exists. Unload/reload a running model after changing its environment; this
is not live environment mutation of an already loaded child.
This requires updated frontends **and** workers. Legacy install requests omit
the snapshot and retain legacy reuse behavior; older workers ignore the added
field and cannot provide per-model environment isolation. Upgrade workers before
relying on this feature. Environment snapshots may contain secrets and travel
over the control transport: restrict NATS access and enable its authentication
and TLS for untrusted networks. They are not secret-store references.
Install coalescing is per frontend instance and includes a keyed identity of
the snapshot, so concurrent installs with different environments do not share an
install reply. The worker owns process-local snapshots; they are rebuilt from
model configuration on each install, not persisted as a new cluster state store.
+24 -5
View File
@@ -63,9 +63,28 @@ against - and two modes:
`proxy.provider` selects the auth scheme and (in translate mode) the wire
format. Supported values: `openai`, `anthropic`.
If the upstream requires an API key, configure either an environment variable
(`api_key_env`) or a file (`api_key_file`). The key never appears in the config
file or the admin UI. If the upstream requires no API key, omit both fields.
If the upstream requires an API key, configure one of `api_key`, `api_key_env`
(an environment variable name), or `api_key_file` (a file path). If the upstream
requires no API key, omit all three fields.
`api_key` is a plaintext administrator-managed credential, supported by both
`cloud-proxy` and `localai-proxy`:
```yaml
proxy:
api_key: your-upstream-key
```
The model editor masks this input but preserves its actual value on unrelated
saves. Masking is not encryption or vault storage: model YAML/JSON, admin
configuration responses, exports, and backups can contain the plaintext key,
just as they can contain model `env` values. Restrict access and protect files
and backups. Use `api_key_env` or `api_key_file` for external secret management.
Direct keys require updated core and proxy backend binaries. In distributed
installations, update workers and their proxy backends before frontends; older
binaries can ignore this additive field and send unauthenticated requests.
This does not add distributed model-environment propagation.
### OpenAI passthrough
@@ -146,7 +165,7 @@ proxy:
### Upstreams without an API key
For an OpenAI-compatible upstream that accepts requests without authentication,
omit both `api_key_env` and `api_key_file`:
omit `api_key`, `api_key_env`, and `api_key_file`:
```yaml
name: internal-chat-proxy
@@ -239,7 +258,7 @@ proxy:
```
The file is read at backend load time and trimmed of surrounding whitespace.
`api_key_env` and `api_key_file` are mutually exclusive.
`api_key`, `api_key_env`, and `api_key_file` are mutually exclusive.
## Combining with the intelligent router
+18
View File
@@ -0,0 +1,18 @@
// SPDX-License-Identifier: MIT
package model
import (
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
)
var _ = Describe("Model env isolation", func() {
It("copies the caller map for each options instance", func() {
env := map[string]string{"KEY": "original"}
option := WithEnvVars(env)
first, second := NewOptions(option), NewOptions(option)
env["KEY"] = "caller mutation"
first.envVars["KEY"] = "first mutation"
Expect(second.envVars["KEY"]).To(Equal("original"))
})
})
+2 -3
View File
@@ -75,7 +75,7 @@ const (
func (ml *ModelLoader) grpcModel(backend string, o *Options) func(string, string, string) (*Model, error) {
return func(modelID, modelName, modelFile string) (*Model, error) {
xlog.Debug("Loading Model with gRPC", "modelID", modelID, "file", modelFile, "backend", backend, "options", *o)
xlog.Debug("Loading Model with gRPC", "modelID", modelID, "file", modelFile, "backend", backend)
// Distributed mode: delegate to the model router if set. No load
// event is emitted here: this branch runs per inference request and
@@ -180,7 +180,6 @@ func (ml *ModelLoader) spawnGRPCModel(backend, uri string, o *Options, modelID,
}
xlog.Debug("Wait for the service to start up")
xlog.Debug("Options", "options", o.gRPCOptions)
// Wait for the service to start up
ready := false
@@ -212,7 +211,7 @@ func (ml *ModelLoader) spawnGRPCModel(backend, uri string, o *Options, modelID,
options.ModelFile = modelFile
options.ModelPath = ml.ModelPath
xlog.Debug("GRPC: Loading model with options", "options", options)
xlog.Debug("GRPC: Loading model", "modelID", modelID, "backend", backend)
res, err := client.GRPC(o.parallelRequests, ml.wd).LoadModel(o.context, options)
if err != nil {
+3 -1
View File
@@ -2,6 +2,7 @@ package model
import (
"context"
"maps"
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
)
@@ -116,8 +117,9 @@ func WithModelSizeBytes(bytes int64) Option {
}
func WithEnvVars(envVars map[string]string) Option {
snapshot := maps.Clone(envVars)
return func(o *Options) {
o.envVars = envVars
o.envVars = maps.Clone(snapshot)
}
}