mirror of
https://github.com/mudler/LocalAI.git
synced 2026-09-12 22:33:54 -04:00
feat(distributed): key scheduling rules by a model alias (#11771)
Node placement and replica rules could only name a model, so an operator who pinned "llama3" to the GPU tier had to rewrite the rule whenever a different model took over that job. An alias already gives a stable name for whichever model serves it, and a rule on that name makes it a deployment slot: repoint the alias and the placement follows. A rule keeps the name the operator chose. Reads resolve that name through the config loader to the model the rule governs, so the reconciler counts, schedules and trims replicas of the target, and the router finds an alias-keyed rule from the target it is already routing. An alias that resolves to nothing governs nothing loadable, so the reconciler skips it and the write paths refuse it. A replica is shared by every name that resolves to it, so only one rule can decide where it runs. The REST and MCP write paths reject a rule whose target another rule already governs. A pair that arrives some other way, such as a seed file or an alias repointed onto a model that already has a rule, resolves in favour of the rule named after the model itself and then the oldest, and the rest are listed as shadowed. The eviction guard is the exception: it matches rules to replicas in raw SQL inside a locking transaction and cannot resolve an alias. It reads a stored target that the reconciler refreshes each tick, and falls back to the rule's own name when that target is empty. Assisted-by: Claude:claude-opus-5 golangci-lint eslint Signed-off-by: Ettore Di Giacinto <mudler@localai.io> Co-authored-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
1 parent
f9f4d2751f
commit
80e3240f2d
24 files changed
+1291
-34
No files matched your search
@@ -162,6 +162,15 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade
|
||||
}
|
||||
xlog.Info("Node registry initialized")
|
||||
|
||||
// Let scheduling rules be keyed by a model alias. The registry resolves a
|
||||
// rule's name through the config loader to find the model it governs, so an
|
||||
// operator can pin placement to a stable name like "production" and have it
|
||||
// follow the alias when the alias is repointed. Wired before the seed below
|
||||
// and before the reconciler starts, so the first tick already resolves.
|
||||
if configLoader != nil {
|
||||
registry.SetAliasResolver(configLoader)
|
||||
}
|
||||
|
||||
// Seed declarative per-model scheduling config (LOCALAI_MODEL_SCHEDULING /
|
||||
// LOCALAI_MODEL_SCHEDULING_CONFIG). Authoritative: overwrites matching models
|
||||
// on every boot. Runs before the reconciler starts so the first tick already
|
||||
|
||||
@@ -441,6 +441,26 @@ func (bcl *ModelConfigLoader) ResolveAlias(cfg *ModelConfig) (*ModelConfig, bool
|
||||
return &target, true, nil
|
||||
}
|
||||
|
||||
// ResolveAliasName maps a model name to the name of the model that actually
|
||||
// serves it: an alias resolves to its target, anything else resolves to
|
||||
// itself. The second return reports whether name was an alias.
|
||||
//
|
||||
// Unlike ResolveAlias this never errors. A name with no config (a rule may be
|
||||
// authored before the model is installed), a dangling alias, and a chained
|
||||
// alias all resolve to themselves, so callers keep a usable name that simply
|
||||
// has no model behind it rather than silently governing a different model.
|
||||
func (bcl *ModelConfigLoader) ResolveAliasName(name string) (string, bool) {
|
||||
cfg, exists := bcl.GetModelConfig(name)
|
||||
if !exists || !cfg.IsAlias() {
|
||||
return name, false
|
||||
}
|
||||
target, exists := bcl.GetModelConfig(cfg.Alias)
|
||||
if !exists || target.IsAlias() {
|
||||
return name, true
|
||||
}
|
||||
return target.Name, true
|
||||
}
|
||||
|
||||
// ValidateAliasTarget checks an alias config's target at create/swap time:
|
||||
// the target must exist, must not be an alias, and must not be disabled.
|
||||
// Returns nil for non-alias configs.
|
||||
|
||||
@@ -314,3 +314,57 @@ var _ = Describe("ModelConfigLoader alias resolution", func() {
|
||||
Expect(loader.ValidateAliasTarget(&bad)).To(MatchError(ContainSubstring("itself an alias")))
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("ModelConfigLoader ResolveAliasName", func() {
|
||||
var loader *ModelConfigLoader
|
||||
|
||||
BeforeEach(func() {
|
||||
loader = NewModelConfigLoader("")
|
||||
loader.configs["real"] = ModelConfig{Name: "real", Backend: "llama-cpp"}
|
||||
loader.configs["production"] = ModelConfig{Name: "production", Alias: "real"}
|
||||
loader.configs["chain"] = ModelConfig{Name: "chain", Alias: "production"}
|
||||
loader.configs["dangling"] = ModelConfig{Name: "dangling", Alias: "nope"}
|
||||
})
|
||||
|
||||
It("maps an alias name to the model that actually serves it", func() {
|
||||
target, isAlias := loader.ResolveAliasName("production")
|
||||
Expect(isAlias).To(BeTrue())
|
||||
Expect(target).To(Equal("real"))
|
||||
})
|
||||
|
||||
It("maps a real model name to itself", func() {
|
||||
target, isAlias := loader.ResolveAliasName("real")
|
||||
Expect(isAlias).To(BeFalse())
|
||||
Expect(target).To(Equal("real"))
|
||||
})
|
||||
|
||||
// A rule may be authored for a model that is not installed yet (pre-staging
|
||||
// placement before standing up a node), so an unknown name must resolve to
|
||||
// itself rather than to the empty string.
|
||||
It("maps an unknown name to itself", func() {
|
||||
target, isAlias := loader.ResolveAliasName("not-installed-yet")
|
||||
Expect(isAlias).To(BeFalse())
|
||||
Expect(target).To(Equal("not-installed-yet"))
|
||||
})
|
||||
|
||||
// A broken alias has no model behind it. Resolving to itself keeps the
|
||||
// caller on a name that simply has no replicas, instead of silently
|
||||
// governing some other model.
|
||||
It("maps a dangling alias to itself", func() {
|
||||
target, isAlias := loader.ResolveAliasName("dangling")
|
||||
Expect(isAlias).To(BeTrue())
|
||||
Expect(target).To(Equal("dangling"))
|
||||
})
|
||||
|
||||
It("maps a chained alias to itself rather than following the chain", func() {
|
||||
target, isAlias := loader.ResolveAliasName("chain")
|
||||
Expect(isAlias).To(BeTrue())
|
||||
Expect(target).To(Equal("chain"))
|
||||
})
|
||||
|
||||
It("maps the empty name to itself", func() {
|
||||
target, isAlias := loader.ResolveAliasName("")
|
||||
Expect(isAlias).To(BeFalse())
|
||||
Expect(target).To(BeEmpty())
|
||||
})
|
||||
})
|
||||
@@ -1218,6 +1218,20 @@ func SetSchedulingEndpoint(registry *nodes.NodeRegistry) echo.HandlerFunc {
|
||||
return c.JSON(http.StatusBadRequest, nodeError(http.StatusBadRequest, err.Error()))
|
||||
}
|
||||
|
||||
// A rule may be keyed by an alias, in which case it governs whatever
|
||||
// that alias currently points at. Reject an alias that resolves to
|
||||
// nothing, and reject a second rule for a model some other rule already
|
||||
// governs, so the operator hears about the clash instead of silently
|
||||
// writing a rule that never takes effect.
|
||||
target, err := registry.ValidateSchedulingTarget(ctx, req.ModelName)
|
||||
if err != nil {
|
||||
status := http.StatusBadRequest
|
||||
if errors.Is(err, nodes.ErrSchedulingConflict) {
|
||||
status = http.StatusConflict
|
||||
}
|
||||
return c.JSON(status, nodeError(status, err.Error()))
|
||||
}
|
||||
|
||||
// Serialize node selector to JSON
|
||||
var selectorJSON string
|
||||
if len(req.NodeSelector) > 0 {
|
||||
@@ -1230,6 +1244,7 @@ func SetSchedulingEndpoint(registry *nodes.NodeRegistry) echo.HandlerFunc {
|
||||
|
||||
config := &nodes.ModelSchedulingConfig{
|
||||
ModelName: req.ModelName,
|
||||
TargetModel: target,
|
||||
NodeSelector: selectorJSON,
|
||||
MinReplicas: req.MinReplicas,
|
||||
MaxReplicas: req.MaxReplicas,
|
||||
|
||||
@@ -0,0 +1,125 @@
|
||||
package localai
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
|
||||
"github.com/labstack/echo/v4"
|
||||
"github.com/mudler/LocalAI/core/services/nodes"
|
||||
"github.com/mudler/LocalAI/core/services/testutil"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
// aliasResolverStub maps alias names to targets in place of a config loader.
|
||||
type aliasResolverStub struct{ aliases map[string]string }
|
||||
|
||||
func (s *aliasResolverStub) ResolveAliasName(name string) (string, bool) {
|
||||
target, ok := s.aliases[name]
|
||||
if !ok {
|
||||
return name, false
|
||||
}
|
||||
return target, true
|
||||
}
|
||||
|
||||
var _ = Describe("Scheduling endpoints with model aliases", func() {
|
||||
var (
|
||||
registry *nodes.NodeRegistry
|
||||
resolver *aliasResolverStub
|
||||
)
|
||||
|
||||
BeforeEach(func() {
|
||||
db := testutil.SetupTestDB()
|
||||
var err error
|
||||
registry, err = nodes.NewNodeRegistry(db)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
resolver = &aliasResolverStub{aliases: map[string]string{"production": "qwen3"}}
|
||||
registry.SetAliasResolver(resolver)
|
||||
})
|
||||
|
||||
post := func(body string) *httptest.ResponseRecorder {
|
||||
e := echo.New()
|
||||
req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(body))
|
||||
req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON)
|
||||
rec := httptest.NewRecorder()
|
||||
c := e.NewContext(req, rec)
|
||||
ExpectWithOffset(1, SetSchedulingEndpoint(registry)(c)).To(Succeed())
|
||||
return rec
|
||||
}
|
||||
|
||||
It("accepts a rule keyed by an alias and reports the model it governs", func() {
|
||||
rec := post(`{"model_name":"production","min_replicas":2,"node_selector":{"tier":"gpu"}}`)
|
||||
Expect(rec.Code).To(Equal(http.StatusOK))
|
||||
|
||||
var resp map[string]any
|
||||
Expect(json.Unmarshal(rec.Body.Bytes(), &resp)).To(Succeed())
|
||||
Expect(resp["model_name"]).To(Equal("production"))
|
||||
Expect(resp["target_model"]).To(Equal("qwen3"))
|
||||
})
|
||||
|
||||
It("rejects a second rule for a model an alias rule already governs", func() {
|
||||
Expect(post(`{"model_name":"production","min_replicas":2}`).Code).To(Equal(http.StatusOK))
|
||||
|
||||
rec := post(`{"model_name":"qwen3","min_replicas":1}`)
|
||||
Expect(rec.Code).To(Equal(http.StatusConflict))
|
||||
Expect(rec.Body.String()).To(ContainSubstring("production"))
|
||||
})
|
||||
|
||||
It("rejects an alias rule for a model that already has its own rule", func() {
|
||||
Expect(post(`{"model_name":"qwen3","min_replicas":1}`).Code).To(Equal(http.StatusOK))
|
||||
|
||||
rec := post(`{"model_name":"production","min_replicas":2}`)
|
||||
Expect(rec.Code).To(Equal(http.StatusConflict))
|
||||
Expect(rec.Body.String()).To(ContainSubstring("qwen3"))
|
||||
})
|
||||
|
||||
It("still allows editing a rule in place", func() {
|
||||
Expect(post(`{"model_name":"production","min_replicas":2}`).Code).To(Equal(http.StatusOK))
|
||||
|
||||
rec := post(`{"model_name":"production","min_replicas":4}`)
|
||||
Expect(rec.Code).To(Equal(http.StatusOK))
|
||||
|
||||
stored, err := registry.GetModelScheduling(context.Background(), "production")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(stored.MinReplicas).To(Equal(4))
|
||||
})
|
||||
|
||||
It("rejects a rule keyed by an alias that does not resolve", func() {
|
||||
resolver.aliases["orphan"] = "orphan"
|
||||
|
||||
rec := post(`{"model_name":"orphan","min_replicas":1}`)
|
||||
Expect(rec.Code).To(Equal(http.StatusBadRequest))
|
||||
Expect(rec.Body.String()).To(ContainSubstring("does not resolve"))
|
||||
})
|
||||
|
||||
It("still accepts a rule for a model that is not installed yet", func() {
|
||||
rec := post(`{"model_name":"not-installed-yet","min_replicas":1}`)
|
||||
Expect(rec.Code).To(Equal(http.StatusOK))
|
||||
})
|
||||
|
||||
It("labels a rule that another rule shadows when listing", func() {
|
||||
// A seed file or a repointed alias can leave two rules on one model,
|
||||
// which the write path above rejects but cannot retract.
|
||||
Expect(registry.SetModelScheduling(context.Background(), &nodes.ModelSchedulingConfig{ModelName: "production", MinReplicas: 2})).To(Succeed())
|
||||
Expect(registry.SetModelScheduling(context.Background(), &nodes.ModelSchedulingConfig{ModelName: "qwen3", MinReplicas: 1})).To(Succeed())
|
||||
|
||||
e := echo.New()
|
||||
req := httptest.NewRequest(http.MethodGet, "/", nil)
|
||||
rec := httptest.NewRecorder()
|
||||
c := e.NewContext(req, rec)
|
||||
Expect(ListSchedulingEndpoint(registry)(c)).To(Succeed())
|
||||
|
||||
var listed []map[string]any
|
||||
Expect(json.Unmarshal(rec.Body.Bytes(), &listed)).To(Succeed())
|
||||
byName := map[string]map[string]any{}
|
||||
for _, item := range listed {
|
||||
byName[item["model_name"].(string)] = item
|
||||
}
|
||||
Expect(byName["qwen3"]["shadowed"]).To(BeNil())
|
||||
Expect(byName["production"]["shadowed"]).To(Equal(true))
|
||||
})
|
||||
})
|
||||
@@ -176,6 +176,75 @@ test.describe('Scheduling page', () => {
|
||||
await expect(page.getByLabel('Node selector').getByText('gpu.vendor=nvidia', { exact: true })).toBeVisible()
|
||||
})
|
||||
|
||||
// A rule may be keyed by an alias, in which case it governs whichever model
|
||||
// the alias points at. The page has to say which model that is, because the
|
||||
// rule's own name no longer tells you.
|
||||
test.describe('rules keyed by a model alias', () => {
|
||||
const aliasRule = {
|
||||
model_name: 'production',
|
||||
target_model: 'llama-3.3',
|
||||
model_is_alias: true,
|
||||
node_selector: { tier: 'gpu' },
|
||||
min_replicas: 2,
|
||||
max_replicas: 4,
|
||||
}
|
||||
|
||||
async function mockAliases(page, aliases = [{ name: 'production', target: 'llama-3.3' }]) {
|
||||
await page.route('**/api/aliases', route => route.fulfill({
|
||||
status: 200,
|
||||
contentType: 'application/json',
|
||||
body: JSON.stringify(aliases),
|
||||
}))
|
||||
await page.route('**/api/models/capabilities', route => route.fulfill({
|
||||
status: 200,
|
||||
contentType: 'application/json',
|
||||
body: JSON.stringify({ object: 'list', data: [{ id: 'llama-3.3' }, { id: 'production' }] }),
|
||||
}))
|
||||
}
|
||||
|
||||
test('names the model an alias rule governs', async ({ page }) => {
|
||||
await mockScheduling(page, { rules: [aliasRule] })
|
||||
await mockAliases(page)
|
||||
await page.goto('/app/scheduling')
|
||||
|
||||
await expect(page.getByText('production')).toBeVisible()
|
||||
await expect(page.locator('.scheduling-rule-target')).toHaveText(/llama-3\.3/)
|
||||
})
|
||||
|
||||
test('marks a rule another rule already governs as shadowed', async ({ page }) => {
|
||||
await mockScheduling(page, { rules: [{ ...aliasRule, shadowed: true }, rule] })
|
||||
await mockAliases(page)
|
||||
await page.goto('/app/scheduling')
|
||||
|
||||
await expect(page.locator('.scheduling-rule-shadowed')).toHaveCount(1)
|
||||
await expect(page.locator('.scheduling-rule-shadowed')).toContainText('Shadowed')
|
||||
})
|
||||
|
||||
test('flags an alias rule that no longer resolves', async ({ page }) => {
|
||||
await mockScheduling(page, {
|
||||
rules: [{ model_name: 'orphan', target_model: 'orphan', model_is_alias: true, min_replicas: 1 }],
|
||||
})
|
||||
await mockAliases(page, [])
|
||||
await page.goto('/app/scheduling')
|
||||
|
||||
await expect(page.locator('.scheduling-rule-target--broken')).toBeVisible()
|
||||
})
|
||||
|
||||
test('offers aliases in the model picker, tagged with their target', async ({ page }) => {
|
||||
await mockScheduling(page)
|
||||
await mockAliases(page)
|
||||
await page.goto('/app/scheduling')
|
||||
await page.getByRole('button', { name: 'Add Scheduling Rule' }).click()
|
||||
|
||||
const picker = page.locator('.searchable-model-select input')
|
||||
await picker.click()
|
||||
await expect(page.locator('.sms-hint')).toHaveText('alias of llama-3.3')
|
||||
|
||||
await page.getByRole('option', { name: /production/ }).click()
|
||||
await expect(page.getByText(/production is an alias for llama-3\.3/)).toBeVisible()
|
||||
})
|
||||
})
|
||||
|
||||
test('keeps rule actions reachable on a narrow viewport', async ({ page }) => {
|
||||
await page.setViewportSize({ width: 390, height: 844 })
|
||||
await mockScheduling(page, { nodeList: nodes.slice(0, 2) })
|
||||
|
||||
@@ -2781,6 +2781,31 @@ select.input {
|
||||
justify-content: flex-end;
|
||||
}
|
||||
|
||||
/* Second line of a rule's Model cell: the model an alias-keyed rule currently
|
||||
governs. The cell itself is bold, so the weight is reset here rather than
|
||||
inherited. */
|
||||
.scheduling-rule-target {
|
||||
font-weight: 400;
|
||||
}
|
||||
|
||||
.scheduling-rule-target--broken {
|
||||
font-weight: 400;
|
||||
color: var(--color-warning);
|
||||
}
|
||||
|
||||
/* Status pill for a rule another rule already governs, so it has no effect.
|
||||
Mirrors the unsatisfiable pill's shape. */
|
||||
.scheduling-rule-shadowed {
|
||||
display: inline-block;
|
||||
font-size: var(--text-xs);
|
||||
padding: 2px 8px;
|
||||
border-radius: var(--radius-sm);
|
||||
font-weight: 600;
|
||||
background: var(--color-bg-tertiary);
|
||||
border: 1px solid var(--color-warning);
|
||||
color: var(--color-warning);
|
||||
}
|
||||
|
||||
@media (max-width: 640px) {
|
||||
.scheduling-rule-actions {
|
||||
width: 100%;
|
||||
|
||||
@@ -7,7 +7,10 @@ import { useModels } from '../hooks/useModels'
|
||||
// query isn't treated as a chosen value. After a commit the field is cleared,
|
||||
// matching the add-and-clear flow. Default false keeps the as-you-type
|
||||
// behaviour single-value editors rely on.
|
||||
export default function SearchableModelSelect({ value, onChange, capability, placeholder = 'Type or select a model...', style, commitOnly = false }) {
|
||||
// hints: optional { [modelId]: string } shown as muted text beside an entry and
|
||||
// searchable along with the name. Used to mark aliases with the model they
|
||||
// point at, so a picker that lists both can tell them apart.
|
||||
export default function SearchableModelSelect({ value, onChange, capability, placeholder = 'Type or select a model...', style, commitOnly = false, hints = {} }) {
|
||||
const { models, loading } = useModels(capability)
|
||||
const [query, setQuery] = useState('')
|
||||
const [open, setOpen] = useState(false)
|
||||
@@ -29,8 +32,10 @@ export default function SearchableModelSelect({ value, onChange, capability, pla
|
||||
return () => document.removeEventListener('mousedown', handler)
|
||||
}, [])
|
||||
|
||||
const needle = query.toLowerCase()
|
||||
const filtered = models.filter(m =>
|
||||
m.id.toLowerCase().includes(query.toLowerCase())
|
||||
m.id.toLowerCase().includes(needle) ||
|
||||
(hints[m.id] || '').toLowerCase().includes(needle)
|
||||
)
|
||||
|
||||
// Which item Enter will select — matches SearchableSelect behavior
|
||||
@@ -126,6 +131,11 @@ export default function SearchableModelSelect({ value, onChange, capability, pla
|
||||
color: var(--color-primary);
|
||||
font-weight: 600;
|
||||
}
|
||||
.sms-hint {
|
||||
color: var(--color-text-muted);
|
||||
font-size: 0.75rem;
|
||||
flex-shrink: 0;
|
||||
}
|
||||
.sms-empty {
|
||||
padding: 8px 10px;
|
||||
font-size: 0.8125rem;
|
||||
@@ -172,6 +182,9 @@ export default function SearchableModelSelect({ value, onChange, capability, pla
|
||||
}}
|
||||
>
|
||||
<span style={{ flex: 1, overflow: 'hidden', textOverflow: 'ellipsis', whiteSpace: 'nowrap' }}>{m.id}</span>
|
||||
{hints[m.id] && (
|
||||
<span className="sms-hint">{hints[m.id]}</span>
|
||||
)}
|
||||
{isEnterTarget && (
|
||||
<span style={{ color: 'var(--color-text-muted)', fontSize: '0.75rem', flexShrink: 0 }}>↵</span>
|
||||
)}
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
import { useState, useEffect, useCallback } from 'react'
|
||||
import { useOutletContext } from 'react-router-dom'
|
||||
import { useTranslation } from 'react-i18next'
|
||||
import { nodesApi } from '../utils/api'
|
||||
import { nodesApi, modelsApi } from '../utils/api'
|
||||
import PageHeader from '../components/PageHeader'
|
||||
import ConfirmDialog from '../components/ConfirmDialog'
|
||||
import ResponsiveTable from '../components/ResponsiveTable'
|
||||
@@ -66,7 +66,7 @@ function configMode(config) {
|
||||
return 'placement'
|
||||
}
|
||||
|
||||
function SchedulingForm({ initialConfig, onSave, onCancel, labels }) {
|
||||
function SchedulingForm({ initialConfig, onSave, onCancel, labels, aliases }) {
|
||||
const [mode, setMode] = useState(() => configMode(initialConfig))
|
||||
const [modelName, setModelName] = useState(initialConfig?.model_name || '')
|
||||
// Selector is now a chip-builder map instead of a comma-separated string.
|
||||
@@ -84,6 +84,10 @@ function SchedulingForm({ initialConfig, onSave, onCancel, labels }) {
|
||||
const [minPrefixMatch, setMinPrefixMatch] = useState(initialConfig?.min_prefix_match ?? 0)
|
||||
|
||||
const hasSelector = Object.keys(selector).length > 0
|
||||
// Aliases are listed in the picker alongside models, tagged with the model
|
||||
// they resolve to so the two are distinguishable in one flat list.
|
||||
const aliasHints = Object.fromEntries(Object.entries(aliases || {}).map(([name, target]) => [name, `alias of ${target}`]))
|
||||
const aliasTarget = (aliases || {})[modelName]
|
||||
|
||||
const isValid = () => {
|
||||
if (!modelName) return false
|
||||
@@ -159,9 +163,20 @@ function SchedulingForm({ initialConfig, onSave, onCancel, labels }) {
|
||||
<SearchableModelSelect
|
||||
value={modelName}
|
||||
onChange={setModelName}
|
||||
placeholder="Type to search models, or paste a name..."
|
||||
placeholder="Type to search models or aliases, or paste a name..."
|
||||
hints={aliasHints}
|
||||
/>
|
||||
)}
|
||||
{/* An alias is a stable name for whichever model currently serves it,
|
||||
so a rule on one is a rule on a slot rather than on a model. Say
|
||||
so at the point of choosing, because the consequence (repointing
|
||||
the alias carries the rule along) is not visible anywhere else. */}
|
||||
{aliasTarget && (
|
||||
<span className="text-meta d-block mt-xs">
|
||||
<i className="fas fa-link icon-before" aria-hidden="true" />
|
||||
{modelName} is an alias for {aliasTarget}. This rule applies to whichever model the alias points at, and follows it if you repoint it.
|
||||
</span>
|
||||
)}
|
||||
</div>
|
||||
|
||||
<div>
|
||||
@@ -301,6 +316,21 @@ export default function Scheduling() {
|
||||
// is swallowed rather than surfaced: the field still commits whatever is
|
||||
// typed into it.
|
||||
const [labels, setLabels] = useState(() => labelIndex([]))
|
||||
// name -> target for every configured alias. Feeds the picker so aliases are
|
||||
// listed as schedulable names. Failing to load costs the annotation and
|
||||
// nothing else: an alias typed by hand still resolves server-side.
|
||||
const [aliases, setAliases] = useState({})
|
||||
|
||||
useEffect(() => {
|
||||
let cancelled = false
|
||||
modelsApi.listAliases()
|
||||
.then(data => {
|
||||
if (cancelled || !Array.isArray(data)) return
|
||||
setAliases(Object.fromEntries(data.map(a => [a.name, a.target])))
|
||||
})
|
||||
.catch(() => {})
|
||||
return () => { cancelled = true }
|
||||
}, [])
|
||||
|
||||
useEffect(() => {
|
||||
let cancelled = false
|
||||
@@ -356,6 +386,7 @@ export default function Scheduling() {
|
||||
onSave={handleSave}
|
||||
onCancel={() => setFormState(null)}
|
||||
labels={labels}
|
||||
aliases={aliases}
|
||||
/>
|
||||
)}
|
||||
{schedulingConfigs.length === 0 && !formState ? (
|
||||
@@ -387,9 +418,26 @@ export default function Scheduling() {
|
||||
// of the model silently failing to scale.
|
||||
const unsatisfiableUntil = cfg.unsatisfiable_until ? new Date(cfg.unsatisfiable_until) : null
|
||||
const isUnsatisfiable = unsatisfiableUntil && unsatisfiableUntil.getTime() > Date.now()
|
||||
// A rule keyed by an alias names a slot, so the model it
|
||||
// currently governs is worth showing next to it.
|
||||
const governs = cfg.target_model && cfg.target_model !== cfg.model_name ? cfg.target_model : null
|
||||
const danglingAlias = cfg.model_is_alias && !governs
|
||||
return (
|
||||
<tr key={cfg.id || cfg.model_name}>
|
||||
<td style={{ fontWeight: 600, fontSize: '0.875rem' }}>{cfg.model_name}</td>
|
||||
<td style={{ fontWeight: 600, fontSize: '0.875rem' }}>
|
||||
{cfg.model_name}
|
||||
{governs && (
|
||||
<div className="text-meta scheduling-rule-target">
|
||||
<i className="fas fa-arrow-right icon-before" aria-hidden="true" />
|
||||
{governs}
|
||||
</div>
|
||||
)}
|
||||
{danglingAlias && (
|
||||
<div className="text-meta scheduling-rule-target--broken">
|
||||
alias points at nothing
|
||||
</div>
|
||||
)}
|
||||
</td>
|
||||
<td>
|
||||
<span style={{
|
||||
display: 'inline-block', fontSize: '0.75rem', padding: '2px 8px', borderRadius: "var(--radius-sm)",
|
||||
@@ -436,7 +484,15 @@ export default function Scheduling() {
|
||||
) : '-'}
|
||||
</td>
|
||||
<td>
|
||||
{isUnsatisfiable ? (
|
||||
{cfg.shadowed ? (
|
||||
<span
|
||||
className="scheduling-rule-shadowed"
|
||||
title="Another rule already governs the same model, so this one has no effect. Placement decides where a single shared load runs, so only one rule per model can apply."
|
||||
>
|
||||
<i className="fas fa-eye-slash icon-before" />
|
||||
Shadowed
|
||||
</span>
|
||||
) : isUnsatisfiable ? (
|
||||
<span
|
||||
title={`Reconciler couldn't satisfy this rule (capacity exhausted). Will retry by ${unsatisfiableUntil.toLocaleString()}, or sooner on a node lifecycle change.`}
|
||||
style={{
|
||||
|
||||
@@ -0,0 +1,230 @@
|
||||
package nodes
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
|
||||
"github.com/mudler/xlog"
|
||||
)
|
||||
|
||||
// AliasResolver maps a model name to the name of the model that actually
|
||||
// serves it: an alias resolves to its target, anything else to itself. The
|
||||
// second return reports whether the name was an alias.
|
||||
//
|
||||
// core/config.ModelConfigLoader implements this. It is an interface here so
|
||||
// the registry stays testable without building a full config loader.
|
||||
type AliasResolver interface {
|
||||
ResolveAliasName(name string) (string, bool)
|
||||
}
|
||||
|
||||
// SetAliasResolver installs the resolver used to map a scheduling rule's model
|
||||
// name onto the model the rule actually governs. Called once at startup before
|
||||
// serving. Leaving it unset makes every rule govern its own name, which is the
|
||||
// behaviour from before rules could be keyed by an alias.
|
||||
func (r *NodeRegistry) SetAliasResolver(resolver AliasResolver) {
|
||||
r.aliasResolver.Store(&resolver)
|
||||
}
|
||||
|
||||
// resolveAlias maps a name through the installed resolver, or returns it
|
||||
// unchanged when no resolver is wired.
|
||||
func (r *NodeRegistry) resolveAlias(name string) (string, bool) {
|
||||
p := r.aliasResolver.Load()
|
||||
if p == nil || *p == nil {
|
||||
return name, false
|
||||
}
|
||||
return (*p).ResolveAliasName(name)
|
||||
}
|
||||
|
||||
// applyTarget fills in the rule's derived TargetModel. Every read path runs a
|
||||
// rule through this so callers can tell the rule's key (ModelName, the name
|
||||
// the operator chose) apart from the model it governs (TargetModel).
|
||||
func (r *NodeRegistry) applyTarget(cfg *ModelSchedulingConfig) {
|
||||
if cfg == nil {
|
||||
return
|
||||
}
|
||||
cfg.TargetModel, cfg.ModelIsAlias = r.resolveAlias(cfg.ModelName)
|
||||
}
|
||||
|
||||
// GetGoverningScheduling returns the rule that governs a physical model: the
|
||||
// rule keyed by the model's own name when one exists, otherwise the rule of an
|
||||
// alias that resolves to it. Returns nil when no rule governs the model.
|
||||
//
|
||||
// The reverse lookup exists because rules stay keyed by the name the operator
|
||||
// chose, so that an alias rule survives repointing the alias, while the router
|
||||
// only ever sees the resolved model name (request middleware resolves the
|
||||
// alias long before routing).
|
||||
func (r *NodeRegistry) GetGoverningScheduling(ctx context.Context, modelName string) (*ModelSchedulingConfig, error) {
|
||||
if direct, err := r.GetModelScheduling(ctx, modelName); err != nil || direct != nil {
|
||||
return direct, err
|
||||
}
|
||||
rule, err := r.aliasRuleFor(ctx, modelName, "")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return rule, nil
|
||||
}
|
||||
|
||||
// aliasRuleFor returns the oldest alias-keyed rule resolving to targetModel,
|
||||
// skipping the rule named by exclude. Returns nil when no alias rule resolves
|
||||
// there.
|
||||
//
|
||||
// Several names can resolve to one model, but placement governs a single
|
||||
// shared load, so exactly one rule can win. Ordering by creation time (then by
|
||||
// name, so rules written in the same transaction still order the same way)
|
||||
// makes every frontend and every reconciler tick pick the same rule.
|
||||
//
|
||||
// This scans the rule table rather than filtering in SQL, because the mapping
|
||||
// from a rule's name to the model it governs lives in the config loader, not
|
||||
// in the database. The scan is bounded by the number of scheduling rules an
|
||||
// operator has written, not by the number of models in the cluster.
|
||||
func (r *NodeRegistry) aliasRuleFor(ctx context.Context, targetModel, exclude string) (*ModelSchedulingConfig, error) {
|
||||
if targetModel == "" {
|
||||
return nil, nil
|
||||
}
|
||||
var configs []ModelSchedulingConfig
|
||||
if err := r.db.WithContext(ctx).Order("created_at ASC, model_name ASC").Find(&configs).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for i := range configs {
|
||||
cfg := &configs[i]
|
||||
// A rule keyed by the target's own name is the direct rule, not an
|
||||
// alias rule; callers handle that case with a precedence of its own.
|
||||
if cfg.ModelName == targetModel || cfg.ModelName == exclude {
|
||||
continue
|
||||
}
|
||||
r.applyTarget(cfg)
|
||||
if cfg.Target() == targetModel {
|
||||
return cfg, nil
|
||||
}
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// SchedulingConflict reports the name of an existing rule that already governs
|
||||
// targetModel, or "" when the target is free. exclude is the rule being
|
||||
// created or edited, which never conflicts with itself.
|
||||
//
|
||||
// Placement governs one shared load, so two rules resolving to the same model
|
||||
// would each claim to decide where it runs. Write paths use this to reject the
|
||||
// second one instead of leaving the outcome to the tiebreak in
|
||||
// GetGoverningScheduling.
|
||||
func (r *NodeRegistry) SchedulingConflict(ctx context.Context, targetModel, exclude string) (string, error) {
|
||||
if targetModel == "" {
|
||||
return "", nil
|
||||
}
|
||||
if targetModel != exclude {
|
||||
direct, err := r.GetModelScheduling(ctx, targetModel)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if direct != nil {
|
||||
return direct.ModelName, nil
|
||||
}
|
||||
}
|
||||
rule, err := r.aliasRuleFor(ctx, targetModel, exclude)
|
||||
if err != nil || rule == nil {
|
||||
return "", err
|
||||
}
|
||||
return rule.ModelName, nil
|
||||
}
|
||||
|
||||
// ResolveRuleTarget returns the model a rule keyed by ruleName would govern,
|
||||
// and whether ruleName is an alias. A name that is an alias but comes back
|
||||
// unchanged is one that does not resolve.
|
||||
func (r *NodeRegistry) ResolveRuleTarget(ruleName string) (string, bool) {
|
||||
return r.resolveAlias(ruleName)
|
||||
}
|
||||
|
||||
// markShadowed flags every rule that resolves to a model some other rule
|
||||
// already governs. Placement decides where one shared load runs, so only one
|
||||
// rule per target can take effect.
|
||||
//
|
||||
// The precedence matches GetGoverningScheduling: a rule keyed by the target's
|
||||
// own name wins, otherwise the oldest rule does. configs is modified in place.
|
||||
func markShadowed(configs []ModelSchedulingConfig) {
|
||||
governing := make(map[string]int, len(configs))
|
||||
for i := range configs {
|
||||
target := configs[i].Target()
|
||||
best, seen := governing[target]
|
||||
if !seen {
|
||||
governing[target] = i
|
||||
continue
|
||||
}
|
||||
if rulePrecedes(configs[i], configs[best], target) {
|
||||
governing[target] = i
|
||||
}
|
||||
}
|
||||
for i := range configs {
|
||||
configs[i].Shadowed = governing[configs[i].Target()] != i
|
||||
}
|
||||
}
|
||||
|
||||
// rulePrecedes reports whether rule a governs target instead of rule b.
|
||||
func rulePrecedes(a, b ModelSchedulingConfig, target string) bool {
|
||||
if (a.ModelName == target) != (b.ModelName == target) {
|
||||
return a.ModelName == target
|
||||
}
|
||||
if !a.CreatedAt.Equal(b.CreatedAt) {
|
||||
return a.CreatedAt.Before(b.CreatedAt)
|
||||
}
|
||||
return a.ModelName < b.ModelName
|
||||
}
|
||||
|
||||
// ErrSchedulingConflict is returned when a rule would govern a model that
|
||||
// another rule already governs. Callers map it onto a conflict status.
|
||||
var ErrSchedulingConflict = errors.New("model already has a scheduling rule")
|
||||
|
||||
// ValidateSchedulingTarget checks that a rule keyed by ruleName can be written,
|
||||
// and returns the model it will govern.
|
||||
//
|
||||
// It rejects two cases. An alias that does not resolve governs nothing
|
||||
// loadable, so a rule on it would sit inert forever. And a model that another
|
||||
// rule already governs cannot take a second one, because placement decides
|
||||
// where a single shared load runs: two rules would each claim to decide, and
|
||||
// only one could win.
|
||||
func (r *NodeRegistry) ValidateSchedulingTarget(ctx context.Context, ruleName string) (string, error) {
|
||||
target, isAlias := r.resolveAlias(ruleName)
|
||||
if isAlias && target == ruleName {
|
||||
return "", fmt.Errorf("%q is an alias that does not resolve to a model: point it at an existing model before giving it a scheduling rule", ruleName)
|
||||
}
|
||||
conflict, err := r.SchedulingConflict(ctx, target, ruleName)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if conflict != "" {
|
||||
return "", fmt.Errorf("%w: rule %q already governs model %q, so edit or delete that rule instead of adding a second one", ErrSchedulingConflict, conflict, target)
|
||||
}
|
||||
return target, nil
|
||||
}
|
||||
|
||||
// RefreshSchedulingTargets rewrites each rule's stored target_model to match
|
||||
// the current alias mapping, and returns the number of rows it changed.
|
||||
//
|
||||
// Go callers resolve aliases live and never read the stored copy. It exists for
|
||||
// the eviction guard, which matches rules to loaded replicas in raw SQL inside
|
||||
// a locking transaction and so cannot resolve an alias itself. Repointing an
|
||||
// alias therefore reaches that guard one reconciler tick later, which is early
|
||||
// enough: until then the guard protects the previous target, and the reconciler
|
||||
// is already reloading the new one.
|
||||
func (r *NodeRegistry) RefreshSchedulingTargets(ctx context.Context) error {
|
||||
var configs []ModelSchedulingConfig
|
||||
if err := r.db.WithContext(ctx).Find(&configs).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
for i := range configs {
|
||||
stored := configs[i].TargetModel
|
||||
live, _ := r.resolveAlias(configs[i].ModelName)
|
||||
if stored == live {
|
||||
continue
|
||||
}
|
||||
if err := r.db.WithContext(ctx).Model(&ModelSchedulingConfig{}).
|
||||
Where("id = ?", configs[i].ID).
|
||||
Update("target_model", live).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
xlog.Info("Scheduling rule now governs a different model",
|
||||
"rule", configs[i].ModelName, "was", stored, "now", live)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,196 @@
|
||||
package nodes
|
||||
|
||||
import (
|
||||
"context"
|
||||
"runtime"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/testutil"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// fakeAliasResolver maps alias names to targets from a plain map, standing in
|
||||
// for the config loader so these specs don't need a model directory.
|
||||
type fakeAliasResolver struct{ aliases map[string]string }
|
||||
|
||||
func (f *fakeAliasResolver) ResolveAliasName(name string) (string, bool) {
|
||||
target, ok := f.aliases[name]
|
||||
if !ok {
|
||||
return name, false
|
||||
}
|
||||
return target, true
|
||||
}
|
||||
|
||||
var _ = Describe("Alias-keyed scheduling rules", func() {
|
||||
var (
|
||||
db *gorm.DB
|
||||
registry *NodeRegistry
|
||||
resolver *fakeAliasResolver
|
||||
)
|
||||
|
||||
BeforeEach(func() {
|
||||
if runtime.GOOS == "darwin" {
|
||||
Skip("testcontainers requires Docker, not available on macOS CI")
|
||||
}
|
||||
db = testutil.SetupTestDB()
|
||||
var err error
|
||||
registry, err = NewNodeRegistry(db)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
resolver = &fakeAliasResolver{aliases: map[string]string{"production": "qwen3"}}
|
||||
registry.SetAliasResolver(resolver)
|
||||
})
|
||||
|
||||
set := func(cfg *ModelSchedulingConfig) {
|
||||
ExpectWithOffset(1, registry.SetModelScheduling(context.Background(), cfg)).To(Succeed())
|
||||
}
|
||||
|
||||
Describe("resolving a rule to the model it governs", func() {
|
||||
It("reports the alias target as a rule's target model", func() {
|
||||
set(&ModelSchedulingConfig{ModelName: "production", MinReplicas: 2})
|
||||
|
||||
got, err := registry.GetModelScheduling(context.Background(), "production")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(got).ToNot(BeNil())
|
||||
// The rule keeps the operator's name; only what it governs resolves.
|
||||
Expect(got.ModelName).To(Equal("production"))
|
||||
Expect(got.Target()).To(Equal("qwen3"))
|
||||
})
|
||||
|
||||
It("reports a plain model rule as governing itself", func() {
|
||||
set(&ModelSchedulingConfig{ModelName: "qwen3", MinReplicas: 1})
|
||||
|
||||
got, err := registry.GetModelScheduling(context.Background(), "qwen3")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(got.Target()).To(Equal("qwen3"))
|
||||
})
|
||||
|
||||
It("resolves targets when listing every rule", func() {
|
||||
set(&ModelSchedulingConfig{ModelName: "production", MinReplicas: 2})
|
||||
|
||||
configs, err := registry.ListModelSchedulings(context.Background())
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(configs).To(HaveLen(1))
|
||||
Expect(configs[0].Target()).To(Equal("qwen3"))
|
||||
})
|
||||
|
||||
It("resolves targets when listing auto-scaling rules", func() {
|
||||
set(&ModelSchedulingConfig{ModelName: "production", MinReplicas: 2})
|
||||
|
||||
configs, err := registry.ListAutoScalingConfigs(context.Background())
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(configs).To(HaveLen(1))
|
||||
Expect(configs[0].Target()).To(Equal("qwen3"))
|
||||
})
|
||||
|
||||
It("governs the new target after the alias is repointed", func() {
|
||||
set(&ModelSchedulingConfig{ModelName: "production", MinReplicas: 2})
|
||||
resolver.aliases["production"] = "llama4"
|
||||
|
||||
got, err := registry.GetModelScheduling(context.Background(), "production")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
// The rule did not move: its settings now apply to llama4.
|
||||
Expect(got.ModelName).To(Equal("production"))
|
||||
Expect(got.Target()).To(Equal("llama4"))
|
||||
Expect(got.MinReplicas).To(Equal(2))
|
||||
})
|
||||
|
||||
It("governs its own name when no resolver is installed", func() {
|
||||
plain, err := NewNodeRegistry(db)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(plain.SetModelScheduling(context.Background(), &ModelSchedulingConfig{ModelName: "production", MinReplicas: 2})).To(Succeed())
|
||||
|
||||
got, err := plain.GetModelScheduling(context.Background(), "production")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(got.Target()).To(Equal("production"))
|
||||
})
|
||||
})
|
||||
|
||||
Describe("finding the rule that governs a loaded model", func() {
|
||||
It("finds an alias rule from the model the alias resolves to", func() {
|
||||
set(&ModelSchedulingConfig{ModelName: "production", MinReplicas: 2, NodeSelector: `{"tier":"gpu"}`})
|
||||
|
||||
got, err := registry.GetGoverningScheduling(context.Background(), "qwen3")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(got).ToNot(BeNil())
|
||||
Expect(got.ModelName).To(Equal("production"))
|
||||
Expect(got.NodeSelector).To(Equal(`{"tier":"gpu"}`))
|
||||
})
|
||||
|
||||
It("returns nil when nothing governs the model", func() {
|
||||
got, err := registry.GetGoverningScheduling(context.Background(), "qwen3")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(got).To(BeNil())
|
||||
})
|
||||
|
||||
It("prefers a rule on the model itself over an alias rule", func() {
|
||||
set(&ModelSchedulingConfig{ModelName: "production", MinReplicas: 2})
|
||||
set(&ModelSchedulingConfig{ModelName: "qwen3", MinReplicas: 7})
|
||||
|
||||
got, err := registry.GetGoverningScheduling(context.Background(), "qwen3")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(got.ModelName).To(Equal("qwen3"))
|
||||
Expect(got.MinReplicas).To(Equal(7))
|
||||
})
|
||||
|
||||
// Two aliases onto one model is a conflict the write paths reject, but
|
||||
// a config-file edit can still produce it. Whichever rule wins, it must
|
||||
// be the same one on every frontend and every tick.
|
||||
It("breaks a two-alias tie deterministically on the older rule", func() {
|
||||
resolver.aliases["staging"] = "qwen3"
|
||||
set(&ModelSchedulingConfig{ModelName: "production", MinReplicas: 2})
|
||||
set(&ModelSchedulingConfig{ModelName: "staging", MinReplicas: 5})
|
||||
|
||||
got, err := registry.GetGoverningScheduling(context.Background(), "qwen3")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(got.ModelName).To(Equal("production"))
|
||||
})
|
||||
|
||||
It("stops governing the old target once the alias is repointed", func() {
|
||||
set(&ModelSchedulingConfig{ModelName: "production", MinReplicas: 2})
|
||||
resolver.aliases["production"] = "llama4"
|
||||
|
||||
gone, err := registry.GetGoverningScheduling(context.Background(), "qwen3")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(gone).To(BeNil())
|
||||
|
||||
moved, err := registry.GetGoverningScheduling(context.Background(), "llama4")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(moved).ToNot(BeNil())
|
||||
Expect(moved.ModelName).To(Equal("production"))
|
||||
})
|
||||
|
||||
It("ignores an alias rule whose target is a different model", func() {
|
||||
set(&ModelSchedulingConfig{ModelName: "production", MinReplicas: 2})
|
||||
|
||||
got, err := registry.GetGoverningScheduling(context.Background(), "some-other-model")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(got).To(BeNil())
|
||||
})
|
||||
})
|
||||
|
||||
Describe("rules that already resolve to a model", func() {
|
||||
It("reports the rule already governing a target", func() {
|
||||
set(&ModelSchedulingConfig{ModelName: "production", MinReplicas: 2})
|
||||
|
||||
conflict, err := registry.SchedulingConflict(context.Background(), "qwen3", "")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(conflict).To(Equal("production"))
|
||||
})
|
||||
|
||||
It("does not report the rule being edited as its own conflict", func() {
|
||||
set(&ModelSchedulingConfig{ModelName: "production", MinReplicas: 2})
|
||||
|
||||
conflict, err := registry.SchedulingConflict(context.Background(), "qwen3", "production")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(conflict).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("reports no conflict when the target is free", func() {
|
||||
conflict, err := registry.SchedulingConflict(context.Background(), "qwen3", "")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(conflict).To(BeEmpty())
|
||||
})
|
||||
})
|
||||
})
|
||||
@@ -49,6 +49,7 @@ type ModelRouter interface {
|
||||
FindLRUModel(ctx context.Context, nodeID string) (*NodeModel, error)
|
||||
Get(ctx context.Context, nodeID string) (*BackendNode, error)
|
||||
GetModelScheduling(ctx context.Context, modelName string) (*ModelSchedulingConfig, error)
|
||||
GetGoverningScheduling(ctx context.Context, modelName string) (*ModelSchedulingConfig, error)
|
||||
FindNodesBySelector(ctx context.Context, selector map[string]string) ([]BackendNode, error)
|
||||
FindNodesWithFreeSlot(ctx context.Context, modelName string, candidateNodeIDs []string) ([]BackendNode, error)
|
||||
NarrowByDiskHeadroom(ctx context.Context, candidateNodeIDs []string, required uint64) ([]string, error)
|
||||
|
||||
@@ -128,6 +128,9 @@ func (f *fakeModelRouterForSmartRouter) Get(_ context.Context, nodeID string) (*
|
||||
func (f *fakeModelRouterForSmartRouter) GetModelScheduling(_ context.Context, _ string) (*ModelSchedulingConfig, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (f *fakeModelRouterForSmartRouter) GetGoverningScheduling(_ context.Context, _ string) (*ModelSchedulingConfig, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (f *fakeModelRouterForSmartRouter) FindNodesBySelector(_ context.Context, _ map[string]string) ([]BackendNode, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
@@ -783,6 +783,12 @@ func (rc *ReplicaReconciler) pruneProbeFailures(seen map[string]struct{}) {
|
||||
}
|
||||
|
||||
func (rc *ReplicaReconciler) reconcile(ctx context.Context) {
|
||||
// Keep each rule's stored target in step with the alias mapping. Only the
|
||||
// eviction guard reads that column, and it cannot resolve aliases itself.
|
||||
if err := rc.registry.RefreshSchedulingTargets(ctx); err != nil {
|
||||
xlog.Warn("Reconciler: failed to refresh scheduling targets", "error", err)
|
||||
}
|
||||
|
||||
configs, err := rc.registry.ListAutoScalingConfigs(ctx)
|
||||
if err != nil {
|
||||
xlog.Warn("Reconciler: failed to list auto-scaling configs", "error", err)
|
||||
@@ -834,7 +840,26 @@ func (rc *ReplicaReconciler) candidateNodeIDsForSelector(ctx context.Context, cf
|
||||
return ids, true
|
||||
}
|
||||
|
||||
// reconcileModel brings one scheduling rule's replica count in line with what
|
||||
// the rule asks for.
|
||||
//
|
||||
// A rule is keyed by the name the operator chose, which may be an alias, while
|
||||
// the model it governs is cfg.Target(). The two are used deliberately: anything
|
||||
// that touches a loaded replica (counting, capacity, scheduling, eviction,
|
||||
// cache pressure) goes through the target, and the rule's own bookkeeping
|
||||
// columns (the unsatisfiable counter and cooldown) stay keyed by cfg.ModelName.
|
||||
func (rc *ReplicaReconciler) reconcileModel(ctx context.Context, cfg ModelSchedulingConfig) {
|
||||
// An alias that resolves to itself is one that no longer resolves at all:
|
||||
// its target was removed, or it was pointed at another alias. Scheduling it
|
||||
// would ask a worker to load a pure redirect, which has no backend and no
|
||||
// model file behind it, so leave the rule alone until the alias is fixed.
|
||||
if cfg.ModelIsAlias && cfg.Target() == cfg.ModelName {
|
||||
xlog.Warn("Reconciler: scheduling rule is keyed by an alias that does not resolve; skipping",
|
||||
"rule", cfg.ModelName)
|
||||
return
|
||||
}
|
||||
target := cfg.Target()
|
||||
|
||||
// spread_all: derive a dynamic replica target equal to the number of nodes
|
||||
// currently matching the selector (all healthy backend nodes when the
|
||||
// selector is empty). Feeding it through Min==Max==target reuses every
|
||||
@@ -864,9 +889,9 @@ func (rc *ReplicaReconciler) reconcileModel(ctx context.Context, cfg ModelSchedu
|
||||
return
|
||||
}
|
||||
|
||||
current, err := rc.registry.CountLoadedReplicas(ctx, cfg.ModelName)
|
||||
current, err := rc.registry.CountLoadedReplicas(ctx, target)
|
||||
if err != nil {
|
||||
xlog.Warn("Reconciler: failed to count replicas", "model", cfg.ModelName, "error", err)
|
||||
xlog.Warn("Reconciler: failed to count replicas", "model", target, "error", err)
|
||||
return
|
||||
}
|
||||
|
||||
@@ -877,14 +902,14 @@ func (rc *ReplicaReconciler) reconcileModel(ctx context.Context, cfg ModelSchedu
|
||||
if cfg.MinReplicas > 0 && int(current) < cfg.MinReplicas {
|
||||
candidateNodeIDs, selectorMatched := rc.candidateNodeIDsForSelector(ctx, cfg)
|
||||
if !selectorMatched {
|
||||
xlog.Warn("Reconciler: no nodes match selector", "model", cfg.ModelName, "selector", cfg.NodeSelector)
|
||||
xlog.Warn("Reconciler: no nodes match selector", "model", target, "selector", cfg.NodeSelector)
|
||||
rc.markCapacityProblem(ctx, cfg.ModelName, "no nodes match selector")
|
||||
return
|
||||
}
|
||||
|
||||
capacity, capErr := rc.registry.ClusterCapacityForModel(ctx, cfg.ModelName, candidateNodeIDs)
|
||||
capacity, capErr := rc.registry.ClusterCapacityForModel(ctx, target, candidateNodeIDs)
|
||||
if capErr != nil {
|
||||
xlog.Warn("Reconciler: failed to compute cluster capacity", "model", cfg.ModelName, "error", capErr)
|
||||
xlog.Warn("Reconciler: failed to compute cluster capacity", "model", target, "error", capErr)
|
||||
return
|
||||
}
|
||||
|
||||
@@ -898,11 +923,11 @@ func (rc *ReplicaReconciler) reconcileModel(ctx context.Context, cfg ModelSchedu
|
||||
}
|
||||
// Cap to actual capacity so we don't try harder than possible.
|
||||
if needed > capacity {
|
||||
xlog.Info("Reconciler: capping scale-up at cluster capacity", "model", cfg.ModelName,
|
||||
xlog.Info("Reconciler: capping scale-up at cluster capacity", "model", target,
|
||||
"need", needed, "capacity", capacity)
|
||||
needed = capacity
|
||||
}
|
||||
xlog.Info("Reconciler: scaling up to meet minimum", "model", cfg.ModelName,
|
||||
xlog.Info("Reconciler: scaling up to meet minimum", "model", target,
|
||||
"current", current, "min", cfg.MinReplicas, "adding", needed)
|
||||
if rc.scaleUp(ctx, cfg, needed) {
|
||||
// A real (or partial) scale-up clears the hysteresis so a future
|
||||
@@ -925,19 +950,19 @@ func (rc *ReplicaReconciler) reconcileModel(ctx context.Context, cfg ModelSchedu
|
||||
|
||||
// 2. Auto-scale up if all replicas are busy
|
||||
if current > 0 && (cfg.MaxReplicas == 0 || int(current) < cfg.MaxReplicas) {
|
||||
if rc.allReplicasBusy(ctx, cfg.ModelName) {
|
||||
if rc.allReplicasBusy(ctx, target) {
|
||||
candidateNodeIDs, selectorMatched := rc.candidateNodeIDsForSelector(ctx, cfg)
|
||||
if !selectorMatched {
|
||||
return
|
||||
}
|
||||
capacity, capErr := rc.registry.ClusterCapacityForModel(ctx, cfg.ModelName, candidateNodeIDs)
|
||||
capacity, capErr := rc.registry.ClusterCapacityForModel(ctx, target, candidateNodeIDs)
|
||||
if capErr != nil || capacity == 0 {
|
||||
// All busy AND no slot available — burst load above capacity.
|
||||
// Don't enter cooldown for this case (it's transient demand,
|
||||
// not a misconfig); the next tick will retry naturally.
|
||||
return
|
||||
}
|
||||
xlog.Info("Reconciler: all replicas busy, scaling up", "model", cfg.ModelName,
|
||||
xlog.Info("Reconciler: all replicas busy, scaling up", "model", target,
|
||||
"current", current)
|
||||
// Only mark the tick as having scaled up if a replica was actually
|
||||
// added. On a failed scaleUp, leave scaledUp false so the pressure
|
||||
@@ -958,13 +983,13 @@ func (rc *ReplicaReconciler) reconcileModel(ctx context.Context, cfg ModelSchedu
|
||||
// Skipped when the busy-burst path already scaled up this tick: at most
|
||||
// one scaleUp(+1) per tick (see scaledUp above).
|
||||
if !scaledUp && rc.pressure != nil && current > 0 && (cfg.MaxReplicas == 0 || int(current) < cfg.MaxReplicas) {
|
||||
if pressureCount := rc.pressure.Count(cfg.ModelName, time.Now()); pressureCount >= rc.pressureThreshold {
|
||||
if pressureCount := rc.pressure.Count(target, time.Now()); pressureCount >= rc.pressureThreshold {
|
||||
candidateNodeIDs, selectorMatched := rc.candidateNodeIDsForSelector(ctx, cfg)
|
||||
if selectorMatched {
|
||||
capacity, capErr := rc.registry.ClusterCapacityForModel(ctx, cfg.ModelName, candidateNodeIDs)
|
||||
capacity, capErr := rc.registry.ClusterCapacityForModel(ctx, target, candidateNodeIDs)
|
||||
if capErr == nil && capacity > 0 {
|
||||
xlog.Info("Reconciler: prefix-cache forced-disturb pressure, scaling up",
|
||||
"model", cfg.ModelName, "current", current,
|
||||
"model", target, "current", current,
|
||||
"pressure", pressureCount,
|
||||
"threshold", rc.pressureThreshold)
|
||||
if rc.scaleUp(ctx, cfg, 1) {
|
||||
@@ -979,7 +1004,7 @@ func (rc *ReplicaReconciler) reconcileModel(ctx context.Context, cfg ModelSchedu
|
||||
// we preserve the signal so the next tick retries off
|
||||
// the same accumulated pressure instead of having to
|
||||
// re-accumulate a full window from scratch.
|
||||
rc.pressure.Reset(cfg.ModelName)
|
||||
rc.pressure.Reset(target)
|
||||
}
|
||||
}
|
||||
// No capacity: transient demand, not a misconfig - let the next
|
||||
@@ -1046,14 +1071,14 @@ func (rc *ReplicaReconciler) scaleUp(ctx context.Context, cfg ModelSchedulingCon
|
||||
|
||||
scheduled := 0
|
||||
for i := 0; i < count; i++ {
|
||||
node, err := rc.scheduler.ScheduleAndLoadModel(ctx, cfg.ModelName, candidateNodeIDs)
|
||||
node, err := rc.scheduler.ScheduleAndLoadModel(ctx, cfg.Target(), candidateNodeIDs)
|
||||
if err != nil {
|
||||
xlog.Warn("Reconciler: failed to scale up replica", "model", cfg.ModelName,
|
||||
xlog.Warn("Reconciler: failed to scale up replica", "model", cfg.Target(),
|
||||
"attempt", i+1, "error", err)
|
||||
break // stop trying on first failure
|
||||
}
|
||||
scheduled++
|
||||
xlog.Info("Reconciler: scaled up replica", "model", cfg.ModelName, "node", node.Name)
|
||||
xlog.Info("Reconciler: scaled up replica", "model", cfg.Target(), "node", node.Name)
|
||||
}
|
||||
return scheduled > 0
|
||||
}
|
||||
@@ -1073,7 +1098,7 @@ func (rc *ReplicaReconciler) scaleDownIdle(ctx context.Context, cfg ModelSchedul
|
||||
var idleModels []NodeModel
|
||||
currentModelRevision(rc.registry.db.WithContext(ctx)).
|
||||
Where("node_models.model_name = ? AND node_models.state = ? AND node_models.in_flight = 0 AND node_models.last_used < ?",
|
||||
cfg.ModelName, "loaded", cutoff).
|
||||
cfg.Target(), "loaded", cutoff).
|
||||
Order("replica_index DESC, last_used ASC").
|
||||
Find(&idleModels)
|
||||
|
||||
@@ -1093,7 +1118,7 @@ func (rc *ReplicaReconciler) scaleDownIdle(ctx context.Context, cfg ModelSchedul
|
||||
if err := rc.unloader.UnloadModelOnNode(nm.NodeID, nm.ModelName); err != nil {
|
||||
xlog.Warn("Reconciler: unload failed (model already removed from registry)", "error", err)
|
||||
}
|
||||
xlog.Info("Reconciler: scaled down idle replica", "model", cfg.ModelName, "node", nm.NodeID, "replica", nm.ReplicaIndex)
|
||||
xlog.Info("Reconciler: scaled down idle replica", "model", cfg.Target(), "node", nm.NodeID, "replica", nm.ReplicaIndex)
|
||||
removed++
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,133 @@
|
||||
package nodes
|
||||
|
||||
import (
|
||||
"context"
|
||||
"runtime"
|
||||
"time"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/testutil"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
var _ = Describe("ReplicaReconciler with alias-keyed rules", func() {
|
||||
var (
|
||||
db *gorm.DB
|
||||
registry *NodeRegistry
|
||||
resolver *fakeAliasResolver
|
||||
)
|
||||
|
||||
BeforeEach(func() {
|
||||
if runtime.GOOS == "darwin" {
|
||||
Skip("testcontainers requires Docker, not available on macOS CI")
|
||||
}
|
||||
db = testutil.SetupTestDB()
|
||||
var err error
|
||||
registry, err = NewNodeRegistry(db)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
resolver = &fakeAliasResolver{aliases: map[string]string{"production": "qwen3"}}
|
||||
registry.SetAliasResolver(resolver)
|
||||
})
|
||||
|
||||
registerNode := func(name, address string) *BackendNode {
|
||||
node := &BackendNode{
|
||||
Name: name,
|
||||
NodeType: NodeTypeBackend,
|
||||
Address: address,
|
||||
MaxReplicasPerModel: 4,
|
||||
}
|
||||
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
||||
return node
|
||||
}
|
||||
|
||||
setRule := func(cfg *ModelSchedulingConfig) ModelSchedulingConfig {
|
||||
ExpectWithOffset(1, registry.SetModelScheduling(context.Background(), cfg)).To(Succeed())
|
||||
return mustGetSched(registry, cfg.ModelName)
|
||||
}
|
||||
|
||||
It("loads the model the alias points at, not the alias itself", func() {
|
||||
node := registerNode("alias-n1", "10.9.0.1:50051")
|
||||
rule := setRule(&ModelSchedulingConfig{ModelName: "production", MinReplicas: 1, MaxReplicas: 2})
|
||||
|
||||
scheduler := &fakeScheduler{scheduleNode: node}
|
||||
reconciler := NewReplicaReconciler(ReplicaReconcilerOptions{Registry: registry, Scheduler: scheduler, DB: db})
|
||||
|
||||
reconciler.reconcileModel(context.Background(), rule)
|
||||
|
||||
Expect(scheduler.scheduleCalls).To(HaveLen(1))
|
||||
Expect(scheduler.scheduleCalls[0].modelName).To(Equal("qwen3"))
|
||||
})
|
||||
|
||||
It("counts the target's replicas when deciding whether the floor is met", func() {
|
||||
node := registerNode("alias-n2", "10.9.0.2:50051")
|
||||
Expect(registry.SetNodeModel(context.Background(), node.ID, "qwen3", 0, "loaded", "", 0)).To(Succeed())
|
||||
rule := setRule(&ModelSchedulingConfig{ModelName: "production", MinReplicas: 1, MaxReplicas: 2})
|
||||
|
||||
scheduler := &fakeScheduler{scheduleNode: node}
|
||||
reconciler := NewReplicaReconciler(ReplicaReconcilerOptions{Registry: registry, Scheduler: scheduler, DB: db})
|
||||
|
||||
reconciler.reconcileModel(context.Background(), rule)
|
||||
|
||||
// The floor is already met by the target's replica. Counting against
|
||||
// the alias name instead would see zero and load a redundant replica.
|
||||
Expect(scheduler.scheduleCalls).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("scales down idle replicas of the target", func() {
|
||||
n1 := registerNode("alias-n3", "10.9.0.3:50051")
|
||||
n2 := registerNode("alias-n4", "10.9.0.4:50051")
|
||||
past := time.Now().Add(-10 * time.Minute)
|
||||
for _, n := range []*BackendNode{n1, n2} {
|
||||
Expect(registry.SetNodeModel(context.Background(), n.ID, "qwen3", 0, "loaded", "", 0)).To(Succeed())
|
||||
db.Model(&NodeModel{}).Where("node_id = ? AND model_name = ?", n.ID, "qwen3").Update("last_used", past)
|
||||
}
|
||||
rule := setRule(&ModelSchedulingConfig{ModelName: "production", MinReplicas: 1, MaxReplicas: 4})
|
||||
|
||||
unloader := &fakeUnloader{}
|
||||
reconciler := NewReplicaReconciler(ReplicaReconcilerOptions{
|
||||
Registry: registry, Unloader: unloader, DB: db, ScaleDownDelay: time.Minute,
|
||||
})
|
||||
|
||||
reconciler.reconcileModel(context.Background(), rule)
|
||||
|
||||
remaining, err := registry.CountLoadedReplicas(context.Background(), "qwen3")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(remaining).To(BeNumerically("==", 1))
|
||||
})
|
||||
|
||||
It("skips a rule whose alias no longer resolves", func() {
|
||||
registerNode("alias-n5", "10.9.0.5:50051")
|
||||
resolver.aliases["orphan"] = "orphan" // target removed: resolves to itself
|
||||
rule := setRule(&ModelSchedulingConfig{ModelName: "orphan", MinReplicas: 1})
|
||||
|
||||
scheduler := &fakeScheduler{}
|
||||
reconciler := NewReplicaReconciler(ReplicaReconcilerOptions{Registry: registry, Scheduler: scheduler, DB: db})
|
||||
|
||||
reconciler.reconcileModel(context.Background(), rule)
|
||||
|
||||
// Loading the alias name would ask a worker to start a pure redirect
|
||||
// that has no backend and no model file behind it.
|
||||
Expect(scheduler.scheduleCalls).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("records unsatisfiable capacity against the rule, not the target", func() {
|
||||
registerNode("alias-n6", "10.9.0.6:50051")
|
||||
rule := setRule(&ModelSchedulingConfig{ModelName: "production", MinReplicas: 1, NodeSelector: `{"tier":"absent"}`})
|
||||
|
||||
reconciler := NewReplicaReconciler(ReplicaReconcilerOptions{Registry: registry, Scheduler: &fakeScheduler{}, DB: db})
|
||||
for i := 0; i < unsatisfiableTickThreshold; i++ {
|
||||
reconciler.reconcileModel(context.Background(), rule)
|
||||
}
|
||||
|
||||
stored, err := registry.GetModelScheduling(context.Background(), "production")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(stored.UnsatisfiableUntil).ToNot(BeNil())
|
||||
|
||||
// The bookkeeping belongs to the rule row; the target has no rule.
|
||||
targetRule, err := registry.GetModelScheduling(context.Background(), "qwen3")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(targetRule).To(BeNil())
|
||||
})
|
||||
})
|
||||
@@ -225,6 +225,42 @@ type ModelSchedulingConfig struct {
|
||||
UnsatisfiableTicks int `gorm:"column:unsatisfiable_ticks;default:0" json:"unsatisfiable_ticks"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
|
||||
// TargetModel is the model this rule actually governs: ModelName itself,
|
||||
// or, when ModelName is an alias, the model that alias points at. Callers
|
||||
// must use Target() for anything that touches a loaded replica (counting,
|
||||
// capacity, scheduling, eviction) and ModelName for anything that touches
|
||||
// this rule's own row.
|
||||
//
|
||||
// Every read re-derives it from the live alias mapping, so Go callers never
|
||||
// see a stale value. It is stored as well, purely so the eviction guard in
|
||||
// evictLRUAndFreeNodeFrom can match a rule to a loaded replica inside its
|
||||
// locking transaction: that check is raw SQL and cannot resolve an alias.
|
||||
// RefreshSchedulingTargets rewrites the stored copy on every reconciler
|
||||
// tick, so repointing an alias reaches the guard within a tick.
|
||||
TargetModel string `gorm:"column:target_model;size:255" json:"target_model,omitempty"`
|
||||
// ModelIsAlias reports whether ModelName is an alias rather than a model.
|
||||
// Also derived on every read. An alias whose TargetModel equals ModelName
|
||||
// is one that could not be resolved (its target is gone, or points at
|
||||
// another alias): it governs nothing loadable.
|
||||
ModelIsAlias bool `gorm:"-" json:"model_is_alias,omitempty"`
|
||||
// Shadowed reports that another rule already governs this rule's target, so
|
||||
// this one has no effect. Set only by ListModelSchedulings, which sees every
|
||||
// rule at once. Write paths reject creating such a pair, but one can still
|
||||
// arrive from a seed file or from repointing an alias onto a model that
|
||||
// already has a rule, and an inert rule the operator cannot see is worse
|
||||
// than one that is labelled.
|
||||
Shadowed bool `gorm:"-" json:"shadowed,omitempty"`
|
||||
}
|
||||
|
||||
// Target returns the model this rule governs. It falls back to ModelName when
|
||||
// the rule was built by hand rather than read through the registry, so a rule
|
||||
// that was never alias-resolved still governs itself.
|
||||
func (c ModelSchedulingConfig) Target() string {
|
||||
if c.TargetModel != "" {
|
||||
return c.TargetModel
|
||||
}
|
||||
return c.ModelName
|
||||
}
|
||||
|
||||
// NodeWithExtras extends BackendNode with computed fields for list views.
|
||||
@@ -323,6 +359,13 @@ type NodeRegistry struct {
|
||||
// Stored in an atomic.Pointer to an immutable slice so the startup wiring
|
||||
// (append) and request / reconcile handling (fire) are race-free.
|
||||
replicaRemovedHooks atomic.Pointer[[]func(modelName, nodeID string, replicaIndex int)]
|
||||
|
||||
// aliasResolver maps a scheduling rule's model name onto the model it
|
||||
// governs, so a rule can be keyed by an alias. Installed once at startup
|
||||
// (see SetAliasResolver); nil means every rule governs its own name.
|
||||
// Held in an atomic.Pointer for the same reason as the hooks above: the
|
||||
// startup wiring writes it while request handling reads it.
|
||||
aliasResolver atomic.Pointer[AliasResolver]
|
||||
}
|
||||
|
||||
// AddReplicaRemovedHook registers a callback invoked after a replica row for
|
||||
@@ -404,6 +447,14 @@ func NewNodeRegistry(db *gorm.DB) (*NodeRegistry, error) {
|
||||
return nil, fmt.Errorf("migrating node tables: %w", err)
|
||||
}
|
||||
|
||||
// Rules written before scheduling rules could be keyed by an alias have no
|
||||
// stored target. They are all direct rules, so their target is their own
|
||||
// name, and the eviction guard needs the column filled in to match them.
|
||||
_ = advisorylock.WithLockCtx(context.Background(), db, advisorylock.KeySchemaMigrate, func() error {
|
||||
return db.Exec(`UPDATE model_scheduling_configs SET target_model = model_name
|
||||
WHERE target_model IS NULL OR target_model = ''`).Error
|
||||
})
|
||||
|
||||
// One-shot cleanup of queue rows that can never drain: ops targeted at
|
||||
// agent workers (wrong subscription set), at non-existent nodes, or with
|
||||
// an empty backend name. The guard in enqueueAndDrainBackendOp prevents
|
||||
@@ -1968,13 +2019,14 @@ func (r *NodeRegistry) SetModelScheduling(ctx context.Context, config *ModelSche
|
||||
if config.ID == "" {
|
||||
config.ID = uuid.New().String()
|
||||
}
|
||||
r.applyTarget(config)
|
||||
return r.db.WithContext(ctx).
|
||||
Clauses(clause.OnConflict{
|
||||
Columns: []clause.Column{{Name: "model_name"}},
|
||||
DoUpdates: clause.AssignmentColumns([]string{
|
||||
"node_selector", "min_replicas", "max_replicas", "spread_all",
|
||||
"route_policy", "balance_abs_threshold", "balance_rel_threshold", "min_prefix_match",
|
||||
"updated_at",
|
||||
"target_model", "updated_at",
|
||||
}),
|
||||
}).
|
||||
Create(config).Error
|
||||
@@ -2004,6 +2056,7 @@ func (r *NodeRegistry) GetModelScheduling(ctx context.Context, modelName string)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
r.applyTarget(&config)
|
||||
return &config, nil
|
||||
}
|
||||
|
||||
@@ -2011,6 +2064,10 @@ func (r *NodeRegistry) GetModelScheduling(ctx context.Context, modelName string)
|
||||
func (r *NodeRegistry) ListModelSchedulings(ctx context.Context) ([]ModelSchedulingConfig, error) {
|
||||
var configs []ModelSchedulingConfig
|
||||
err := r.db.WithContext(ctx).Order("model_name ASC").Find(&configs).Error
|
||||
for i := range configs {
|
||||
r.applyTarget(&configs[i])
|
||||
}
|
||||
markShadowed(configs)
|
||||
return configs, err
|
||||
}
|
||||
|
||||
@@ -2018,6 +2075,9 @@ func (r *NodeRegistry) ListModelSchedulings(ctx context.Context) ([]ModelSchedul
|
||||
func (r *NodeRegistry) ListAutoScalingConfigs(ctx context.Context) ([]ModelSchedulingConfig, error) {
|
||||
var configs []ModelSchedulingConfig
|
||||
err := r.db.WithContext(ctx).Where("min_replicas > 0 OR max_replicas > 0 OR spread_all = ?", true).Find(&configs).Error
|
||||
for i := range configs {
|
||||
r.applyTarget(&configs[i])
|
||||
}
|
||||
return configs, err
|
||||
}
|
||||
|
||||
|
||||
@@ -629,7 +629,11 @@ func (r *SmartRouter) Route(ctx context.Context, modelID, modelName, backendType
|
||||
// nodeMatchesScheduling all read it. Fetching once gives a consistent
|
||||
// snapshot and avoids three DB round-trips for one row. nil sched means
|
||||
// "no scheduling constraints", same as before.
|
||||
sched, _ := r.registry.GetModelScheduling(ctx, trackingKey)
|
||||
// GetGoverningScheduling, not GetModelScheduling: a rule may be keyed by an
|
||||
// alias of this model. Request middleware resolves an alias to its target
|
||||
// long before routing, so by here trackingKey is always the target's name
|
||||
// and the alias's rule can only be found by resolving the other way.
|
||||
sched, _ := r.registry.GetGoverningScheduling(ctx, trackingKey)
|
||||
|
||||
// Resolve the model's NodeSelector once so cached-replica lookup and the
|
||||
// new-load scheduler agree on the candidate set. Without this, a cached
|
||||
@@ -1047,7 +1051,7 @@ func (r *SmartRouter) scheduleNewModel(ctx context.Context, backendType, modelID
|
||||
// Check for scheduling constraints (node selector). If a selector is set,
|
||||
// we restrict the candidate pool to matching nodes; otherwise nil means
|
||||
// "any healthy node".
|
||||
sched, _ := r.registry.GetModelScheduling(ctx, modelID)
|
||||
sched, _ := r.registry.GetGoverningScheduling(ctx, modelID)
|
||||
candidateNodeIDs, err := r.resolveSelectorCandidates(ctx, modelID, sched)
|
||||
if err != nil {
|
||||
return nil, "", 0, err
|
||||
@@ -2006,16 +2010,27 @@ func (r *SmartRouter) evictLRUAndFreeNodeFrom(ctx context.Context, candidateNode
|
||||
for attempt := range maxEvictionRetries {
|
||||
var lru NodeModel
|
||||
err := r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
// Lock the row so no other frontend can evict the same model
|
||||
// Lock the row so no other frontend can evict the same model.
|
||||
//
|
||||
// The replica-floor guard matches a rule to a replica through
|
||||
// sc.target_model, so a rule keyed by an alias protects the model
|
||||
// the alias points at. It falls back to sc.model_name when the
|
||||
// stored target is empty, which keeps a rule inserted by some path
|
||||
// that never resolved it protecting itself rather than nothing.
|
||||
//
|
||||
// target_model is not unique (two names can resolve to one model),
|
||||
// so the floor is MAX over the matching rules: only one of them
|
||||
// governs, but over-protecting costs a retry while under-protecting
|
||||
// evicts below a floor the reconciler then has to rebuild.
|
||||
q := currentModelRevision(tx.Clauses(clause.Locking{Strength: "UPDATE"})).
|
||||
Joins("JOIN backend_nodes ON backend_nodes.id = node_models.node_id").
|
||||
Where(`node_models.in_flight = 0 AND node_models.state = ? AND backend_nodes.status = ?
|
||||
AND (
|
||||
NOT EXISTS (SELECT 1 FROM model_scheduling_configs sc WHERE sc.model_name = node_models.model_name AND (sc.min_replicas > 0 OR sc.max_replicas > 0))
|
||||
NOT EXISTS (SELECT 1 FROM model_scheduling_configs sc WHERE COALESCE(NULLIF(sc.target_model, ''), sc.model_name) = node_models.model_name AND (sc.min_replicas > 0 OR sc.max_replicas > 0))
|
||||
OR (SELECT COUNT(*) FROM node_models nm2 WHERE nm2.model_name = node_models.model_name AND nm2.state = 'loaded'
|
||||
AND (NOT EXISTS (SELECT 1 FROM model_config_states mcs2 WHERE mcs2.model_name = nm2.model_name)
|
||||
OR nm2.config_revision = (SELECT mcs3.config_revision FROM model_config_states mcs3 WHERE mcs3.model_name = nm2.model_name)))
|
||||
> COALESCE((SELECT sc2.min_replicas FROM model_scheduling_configs sc2 WHERE sc2.model_name = node_models.model_name), 1)
|
||||
> COALESCE((SELECT MAX(sc2.min_replicas) FROM model_scheduling_configs sc2 WHERE COALESCE(NULLIF(sc2.target_model, ''), sc2.model_name) = node_models.model_name), 1)
|
||||
)`, "loaded", StatusHealthy)
|
||||
if len(candidateNodeIDs) > 0 {
|
||||
q = q.Where("node_models.node_id IN ?", candidateNodeIDs)
|
||||
|
||||
@@ -0,0 +1,104 @@
|
||||
package nodes
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"runtime"
|
||||
"time"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
"gorm.io/gorm"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/testutil"
|
||||
)
|
||||
|
||||
// A replica floor protects a model from LRU eviction. The guard that enforces
|
||||
// it matches a scheduling rule against the loaded model's name, so a rule keyed
|
||||
// by an alias has to reach the model the alias points at. Without that, the
|
||||
// router evicts under the floor and the reconciler reloads on its next tick,
|
||||
// which is the replica flapping this floor exists to prevent.
|
||||
var _ = Describe("Eviction against an alias-keyed replica floor", func() {
|
||||
var (
|
||||
db *gorm.DB
|
||||
registry *NodeRegistry
|
||||
router *SmartRouter
|
||||
ctx context.Context
|
||||
)
|
||||
|
||||
BeforeEach(func() {
|
||||
if runtime.GOOS == "darwin" {
|
||||
Skip("testcontainers requires Docker, not available on macOS CI")
|
||||
}
|
||||
db = testutil.SetupTestDB()
|
||||
var err error
|
||||
registry, err = NewNodeRegistry(db)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
registry.SetAliasResolver(&fakeAliasResolver{aliases: map[string]string{"production": "qwen3"}})
|
||||
router = NewSmartRouter(registry, SmartRouterOptions{DB: db})
|
||||
ctx = context.Background()
|
||||
})
|
||||
|
||||
register := func(name string) *BackendNode {
|
||||
node := &BackendNode{Name: name, NodeType: NodeTypeBackend, Address: name + ":50051"}
|
||||
Expect(registry.Register(ctx, node, true)).To(Succeed())
|
||||
fetched, err := registry.GetByName(ctx, name)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
return fetched
|
||||
}
|
||||
|
||||
rowID := 0
|
||||
seedLoaded := func(node *BackendNode, model string, idleFor time.Duration) {
|
||||
rowID++
|
||||
Expect(db.Create(&NodeModel{
|
||||
ID: fmt.Sprintf("alias-row-%d", rowID), NodeID: node.ID, ModelName: model,
|
||||
Address: node.Address, State: "loaded", InFlight: 0,
|
||||
LastUsed: time.Now().Add(-idleFor), UpdatedAt: time.Now(),
|
||||
}).Error).To(Succeed())
|
||||
}
|
||||
|
||||
rowExists := func(model string) bool {
|
||||
var n int64
|
||||
Expect(db.Model(&NodeModel{}).Where("model_name = ?", model).Count(&n).Error).To(Succeed())
|
||||
return n > 0
|
||||
}
|
||||
|
||||
It("protects the target of an alias-keyed rule at its floor", func() {
|
||||
node := register("floor-node")
|
||||
Expect(registry.SetModelScheduling(ctx, &ModelSchedulingConfig{ModelName: "production", MinReplicas: 1})).To(Succeed())
|
||||
seedLoaded(node, "qwen3", time.Hour)
|
||||
|
||||
_, err := router.evictLRUAndFreeNodeFrom(ctx, nil)
|
||||
|
||||
Expect(err).To(HaveOccurred(), "the only candidate sits at its replica floor")
|
||||
Expect(rowExists("qwen3")).To(BeTrue())
|
||||
})
|
||||
|
||||
It("still evicts the target above its floor", func() {
|
||||
n1 := register("floor-node-a")
|
||||
n2 := register("floor-node-b")
|
||||
Expect(registry.SetModelScheduling(ctx, &ModelSchedulingConfig{ModelName: "production", MinReplicas: 1})).To(Succeed())
|
||||
seedLoaded(n1, "qwen3", 2*time.Hour)
|
||||
seedLoaded(n2, "qwen3", time.Hour)
|
||||
|
||||
_, err := router.evictLRUAndFreeNodeFrom(ctx, nil)
|
||||
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
var remaining int64
|
||||
Expect(db.Model(&NodeModel{}).Where("model_name = ?", "qwen3").Count(&remaining).Error).To(Succeed())
|
||||
Expect(remaining).To(BeNumerically("==", 1))
|
||||
})
|
||||
|
||||
It("stops protecting the old target once the alias is repointed", func() {
|
||||
node := register("floor-node-c")
|
||||
Expect(registry.SetModelScheduling(ctx, &ModelSchedulingConfig{ModelName: "production", MinReplicas: 1})).To(Succeed())
|
||||
seedLoaded(node, "qwen3", time.Hour)
|
||||
registry.SetAliasResolver(&fakeAliasResolver{aliases: map[string]string{"production": "llama4"}})
|
||||
Expect(registry.RefreshSchedulingTargets(ctx)).To(Succeed())
|
||||
|
||||
_, err := router.evictLRUAndFreeNodeFrom(ctx, nil)
|
||||
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(rowExists("qwen3")).To(BeFalse())
|
||||
})
|
||||
})
|
||||
@@ -99,7 +99,12 @@ type fakeModelRouter struct {
|
||||
|
||||
// GetModelScheduling returns
|
||||
getModelScheduling *ModelSchedulingConfig
|
||||
getModelSchedErr error
|
||||
// getGoverningScheduling stands in for a rule keyed by an alias of the
|
||||
// routed model: the model has no rule under its own name, but one governs
|
||||
// it all the same. Falls back to getModelScheduling so specs that set up a
|
||||
// single direct rule need no change.
|
||||
getGoverningScheduling *ModelSchedulingConfig
|
||||
getModelSchedErr error
|
||||
|
||||
// FindNodesBySelector returns
|
||||
findBySelectorNodes []BackendNode
|
||||
@@ -347,6 +352,13 @@ func (f *fakeModelRouter) GetModelScheduling(_ context.Context, _ string) (*Mode
|
||||
return f.getModelScheduling, f.getModelSchedErr
|
||||
}
|
||||
|
||||
func (f *fakeModelRouter) GetGoverningScheduling(_ context.Context, _ string) (*ModelSchedulingConfig, error) {
|
||||
if f.getGoverningScheduling != nil {
|
||||
return f.getGoverningScheduling, f.getModelSchedErr
|
||||
}
|
||||
return f.getModelScheduling, f.getModelSchedErr
|
||||
}
|
||||
|
||||
func (f *fakeModelRouter) FindNodesBySelector(_ context.Context, _ map[string]string) ([]BackendNode, error) {
|
||||
return f.findBySelectorNodes, f.findBySelectorErr
|
||||
}
|
||||
@@ -922,6 +934,30 @@ var _ = Describe("SmartRouter", func() {
|
||||
Expect(result.Node.ID).To(Equal("gpu-1"))
|
||||
})
|
||||
|
||||
It("applies a rule keyed by an alias of the routed model", func() {
|
||||
// No rule under the model's own name: the rule the operator wrote
|
||||
// is keyed "production", an alias that resolves to this model. Its
|
||||
// selector matches nothing, so honouring it is the only way to
|
||||
// reach the selector error — ignoring it would route successfully.
|
||||
reg.getModelScheduling = nil
|
||||
reg.getGoverningScheduling = &ModelSchedulingConfig{
|
||||
ModelName: "production",
|
||||
TargetModel: "aliased-model",
|
||||
NodeSelector: `{"gpu.vendor":"tpu"}`,
|
||||
}
|
||||
reg.findBySelectorNodes = nil
|
||||
reg.findIdleNode = &BackendNode{ID: "cpu-1", Name: "cpu-node", Address: "10.0.0.52:50051"}
|
||||
|
||||
router := NewSmartRouter(reg, SmartRouterOptions{
|
||||
Unloader: unloader,
|
||||
ClientFactory: factory,
|
||||
})
|
||||
|
||||
_, err := router.Route(context.Background(), "aliased-model", "models/aliased.gguf", "llama-cpp", "", nil, false)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("no healthy nodes match selector"))
|
||||
})
|
||||
|
||||
It("returns error when no nodes match selector", func() {
|
||||
reg.getModelScheduling = &ModelSchedulingConfig{
|
||||
ModelName: "no-match-model",
|
||||
|
||||
@@ -914,6 +914,40 @@ All fields are optional and composable:
|
||||
- Replicas only: auto-scale across all nodes
|
||||
- Both: auto-scale on matching nodes only
|
||||
|
||||
### Scheduling a model alias
|
||||
|
||||
`model_name` accepts a [model alias](/features/model-aliases/) as well as a
|
||||
model. A rule keyed by an alias governs whatever model that alias currently
|
||||
points at, and keeps governing it after you repoint the alias:
|
||||
|
||||
```bash
|
||||
# "production" is an alias for llama3
|
||||
curl -X POST http://frontend:8080/api/nodes/scheduling \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"model_name": "production", "node_selector": {"tier": "gpu"}, "min_replicas": 2}'
|
||||
|
||||
# Repoint the alias at a new model: the rule follows, llama4 now runs
|
||||
# two replicas on the GPU tier and llama3 falls back to on-demand placement.
|
||||
```
|
||||
|
||||
This makes an alias a stable deployment slot: the placement policy belongs to
|
||||
the slot, and the model filling it can change without rewriting the rule. The
|
||||
WebUI lists aliases in the model picker on the **Scheduling** page, tagged with
|
||||
the model each one resolves to.
|
||||
|
||||
Two constraints follow from replicas being shared. A single load of `llama3`
|
||||
serves both `production` and any request that names `llama3` directly, so only
|
||||
one rule can decide where it runs: a rule whose target is already governed by
|
||||
another rule is rejected with `409 Conflict` naming the rule that has it. And a
|
||||
rule keyed by an alias that resolves to nothing (its target was deleted, or it
|
||||
points at another alias) is rejected, since it would govern nothing loadable.
|
||||
|
||||
A rule can still end up inert if the pair is created some other way, for example
|
||||
by a declarative seed or by repointing an alias onto a model that already has a
|
||||
rule. The rule that governs is the one keyed by the model's own name, or failing
|
||||
that the oldest one; the rest are listed as **Shadowed** in the WebUI and carry
|
||||
`"shadowed": true` in `GET /api/nodes/scheduling`.
|
||||
|
||||
### Declarative per-model scheduling (unattended installs)
|
||||
|
||||
In distributed mode you can declare per-model scheduling at startup, instead of
|
||||
|
||||
@@ -74,6 +74,23 @@ is fully supported and is the live-swap path: the alias config has no backend of
|
||||
its own, so swapping its target stays a valid pure redirect.
|
||||
{{% /notice %}}
|
||||
|
||||
## Aliases as deployment slots (distributed mode)
|
||||
|
||||
In [distributed mode]({{%relref "features/distributed-mode" %}}) an alias can
|
||||
carry a scheduling rule. `POST /api/nodes/scheduling` accepts an alias for
|
||||
`model_name`, and the rule then governs whatever model the alias points at:
|
||||
|
||||
```bash
|
||||
curl -X POST http://frontend:8080/api/nodes/scheduling \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"model_name": "production", "node_selector": {"tier": "gpu"}, "min_replicas": 2}'
|
||||
```
|
||||
|
||||
Re-point `production` and the placement policy follows it, so the alias behaves
|
||||
as a stable slot whose contents you can swap. Because a replica is shared by
|
||||
every name that resolves to it, only one rule may govern a given model at a
|
||||
time. See [Scheduling a model alias]({{%relref "features/distributed-mode" %}}#scheduling-a-model-alias).
|
||||
|
||||
## Limits
|
||||
|
||||
Aliases are a static 1:1 redirect. For classifier-based or load-balanced
|
||||
|
||||
@@ -128,13 +128,19 @@ type ModelSchedulingConfig struct {
|
||||
BalanceAbsThreshold int `json:"balance_abs_threshold,omitempty"`
|
||||
BalanceRelThreshold float64 `json:"balance_rel_threshold,omitempty"`
|
||||
MinPrefixMatch float64 `json:"min_prefix_match,omitempty"`
|
||||
// TargetModel is the model the rule governs. It differs from ModelName when
|
||||
// the rule is keyed by an alias, in which case it follows the alias.
|
||||
TargetModel string `json:"target_model,omitempty"`
|
||||
// Shadowed reports that another rule already governs TargetModel, leaving
|
||||
// this one with no effect.
|
||||
Shadowed bool `json:"shadowed,omitempty"`
|
||||
}
|
||||
|
||||
// SetSchedulingRequest is the input for set_scheduling. It mirrors
|
||||
// /api/nodes/scheduling so standalone MCP and REST callers preserve the same
|
||||
// PATCH-style semantics for the optional prefix-cache routing fields.
|
||||
type SetSchedulingRequest struct {
|
||||
ModelName string `json:"model_name" jsonschema:"Installed model name whose distributed scheduling rule should be created or updated."`
|
||||
ModelName string `json:"model_name" jsonschema:"Installed model name, or model alias, whose distributed scheduling rule should be created or updated. A rule keyed by an alias follows that alias to whatever model it currently points at."`
|
||||
NodeSelector map[string]string `json:"node_selector,omitempty" jsonschema:"Optional node-label selector. Empty means any healthy backend node."`
|
||||
MinReplicas int `json:"min_replicas" jsonschema:"Minimum desired replicas. Mutually exclusive with spread_all."`
|
||||
MaxReplicas int `json:"max_replicas" jsonschema:"Maximum desired replicas. Must be >= min_replicas when non-zero. Mutually exclusive with spread_all."`
|
||||
|
||||
@@ -600,6 +600,14 @@ func (c *Client) SetScheduling(ctx context.Context, req localaitools.SetScheduli
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Same alias rules as POST /api/nodes/scheduling: a rule may be keyed by an
|
||||
// alias and then follows it, an alias that resolves to nothing is refused,
|
||||
// and a model already governed by another rule cannot take a second one.
|
||||
target, err := c.NodeRegistry.ValidateSchedulingTarget(ctx, req.ModelName)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var selectorJSON string
|
||||
if len(req.NodeSelector) > 0 {
|
||||
b, err := json.Marshal(req.NodeSelector)
|
||||
@@ -610,6 +618,7 @@ func (c *Client) SetScheduling(ctx context.Context, req localaitools.SetScheduli
|
||||
}
|
||||
config := &nodes.ModelSchedulingConfig{
|
||||
ModelName: req.ModelName,
|
||||
TargetModel: target,
|
||||
NodeSelector: selectorJSON,
|
||||
MinReplicas: req.MinReplicas,
|
||||
MaxReplicas: req.MaxReplicas,
|
||||
|
||||
@@ -13,6 +13,8 @@ func SchedulingConfigFromNode(config nodes.ModelSchedulingConfig) ModelSchedulin
|
||||
BalanceAbsThreshold: config.BalanceAbsThreshold,
|
||||
BalanceRelThreshold: config.BalanceRelThreshold,
|
||||
MinPrefixMatch: config.MinPrefixMatch,
|
||||
TargetModel: config.Target(),
|
||||
Shadowed: config.Shadowed,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in new issue
Block a user