mirror of
https://github.com/mudler/LocalAI.git
synced 2026-07-30 01:48:06 -04:00
feat(nodes): add pluggable routing pipeline
Route prefix-cache replica selection through composable filters, weighted scorers, and a replaceable picker. Preserve the existing load guard and deterministic selection policy while exposing per-model scorer weights through scheduling configuration. Assisted-by: Codex:gpt-5 [go-vet]
This commit is contained in:
@@ -1117,21 +1117,22 @@ func GetSchedulingEndpoint(registry *nodes.NodeRegistry) echo.HandlerFunc {
|
||||
|
||||
// SetSchedulingRequest is the request body for creating/updating a scheduling config.
|
||||
//
|
||||
// The four prefix-cache fields are POINTERS so an omitted field is
|
||||
// The prefix-cache fields are POINTERS so an omitted field is
|
||||
// distinguishable from an explicit zero. On update, an omitted prefix-cache
|
||||
// field preserves the model's previously-configured value instead of resetting
|
||||
// it (see SetSchedulingEndpoint's PATCH-style merge). ModelName, NodeSelector,
|
||||
// MinReplicas, MaxReplicas and SpreadAll keep their full-replace PUT semantics.
|
||||
type SetSchedulingRequest struct {
|
||||
ModelName string `json:"model_name"`
|
||||
NodeSelector map[string]string `json:"node_selector,omitempty"`
|
||||
MinReplicas int `json:"min_replicas"`
|
||||
MaxReplicas int `json:"max_replicas"`
|
||||
SpreadAll bool `json:"spread_all,omitempty"`
|
||||
RoutePolicy *string `json:"route_policy,omitempty"`
|
||||
BalanceAbsThreshold *int `json:"balance_abs_threshold,omitempty"`
|
||||
BalanceRelThreshold *float64 `json:"balance_rel_threshold,omitempty"`
|
||||
MinPrefixMatch *float64 `json:"min_prefix_match,omitempty"`
|
||||
ModelName string `json:"model_name"`
|
||||
NodeSelector map[string]string `json:"node_selector,omitempty"`
|
||||
MinReplicas int `json:"min_replicas"`
|
||||
MaxReplicas int `json:"max_replicas"`
|
||||
SpreadAll bool `json:"spread_all,omitempty"`
|
||||
RoutePolicy *string `json:"route_policy,omitempty"`
|
||||
BalanceAbsThreshold *int `json:"balance_abs_threshold,omitempty"`
|
||||
BalanceRelThreshold *float64 `json:"balance_rel_threshold,omitempty"`
|
||||
MinPrefixMatch *float64 `json:"min_prefix_match,omitempty"`
|
||||
ScorerWeights *map[string]float64 `json:"scorer_weights,omitempty"`
|
||||
}
|
||||
|
||||
// validateSchedulingRequest enforces the invariants of a scheduling config.
|
||||
@@ -1159,6 +1160,11 @@ func validateSchedulingRequest(req SetSchedulingRequest, routePolicy string, abs
|
||||
if err := prefixcache.ValidateThresholds(routePolicy, absThr, relThr, minMatch); err != nil {
|
||||
return err
|
||||
}
|
||||
if req.ScorerWeights != nil {
|
||||
if err := prefixcache.ValidateScorerWeights(*req.ScorerWeights); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -1166,7 +1172,7 @@ func validateSchedulingRequest(req SetSchedulingRequest, routePolicy string, abs
|
||||
//
|
||||
// The registry upsert full-replaces all columns, so a request that omits the
|
||||
// prefix-cache fields would otherwise wipe a model's previously-configured
|
||||
// routing settings. To avoid that footgun the four prefix-cache fields are
|
||||
// routing settings. To avoid that footgun the prefix-cache fields are
|
||||
// merged PATCH-style: a non-nil request pointer wins; a nil one preserves the
|
||||
// existing config's value (or the zero default when no config exists yet). The
|
||||
// non-prefix fields keep their full-replace PUT behavior.
|
||||
@@ -1195,11 +1201,13 @@ func SetSchedulingEndpoint(registry *nodes.NodeRegistry) echo.HandlerFunc {
|
||||
absThr := 0
|
||||
relThr := 0.0
|
||||
minMatch := 0.0
|
||||
var scorerWeights map[string]float64
|
||||
if existing != nil {
|
||||
routePolicy = existing.RoutePolicy
|
||||
absThr = existing.BalanceAbsThreshold
|
||||
relThr = existing.BalanceRelThreshold
|
||||
minMatch = existing.MinPrefixMatch
|
||||
scorerWeights = existing.ScorerWeights
|
||||
}
|
||||
if req.RoutePolicy != nil {
|
||||
routePolicy = *req.RoutePolicy
|
||||
@@ -1213,6 +1221,9 @@ func SetSchedulingEndpoint(registry *nodes.NodeRegistry) echo.HandlerFunc {
|
||||
if req.MinPrefixMatch != nil {
|
||||
minMatch = *req.MinPrefixMatch
|
||||
}
|
||||
if req.ScorerWeights != nil {
|
||||
scorerWeights = *req.ScorerWeights
|
||||
}
|
||||
|
||||
if err := validateSchedulingRequest(req, routePolicy, absThr, relThr, minMatch); err != nil {
|
||||
return c.JSON(http.StatusBadRequest, nodeError(http.StatusBadRequest, err.Error()))
|
||||
@@ -1238,6 +1249,7 @@ func SetSchedulingEndpoint(registry *nodes.NodeRegistry) echo.HandlerFunc {
|
||||
BalanceAbsThreshold: absThr,
|
||||
BalanceRelThreshold: relThr,
|
||||
MinPrefixMatch: minMatch,
|
||||
ScorerWeights: scorerWeights,
|
||||
}
|
||||
if err := registry.SetModelScheduling(ctx, config); err != nil {
|
||||
return c.JSON(http.StatusInternalServerError, nodeError(http.StatusInternalServerError, "failed to set scheduling config"))
|
||||
|
||||
@@ -25,6 +25,9 @@ type Config struct {
|
||||
// saturated and scale up (subject to MaxReplicas and capacity). Default 1,
|
||||
// i.e. any sustained forced-disturb.
|
||||
PressureScaleThreshold int
|
||||
// ScorerWeights enables and weights named routing signals. Missing names
|
||||
// use their default weight; zero disables that scorer.
|
||||
ScorerWeights map[string]float64
|
||||
}
|
||||
|
||||
func DefaultConfig() Config {
|
||||
@@ -39,6 +42,7 @@ func DefaultConfig() Config {
|
||||
MaxDepth: 64,
|
||||
PressureWindow: time.Minute,
|
||||
PressureScaleThreshold: 1,
|
||||
ScorerWeights: map[string]float64{ScorerPrefixCache: 1},
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
106
core/services/nodes/prefixcache/pipeline.go
Normal file
106
core/services/nodes/prefixcache/pipeline.go
Normal file
@@ -0,0 +1,106 @@
|
||||
package prefixcache
|
||||
|
||||
import "fmt"
|
||||
|
||||
const ScorerPrefixCache = "prefix_cache"
|
||||
|
||||
func ValidateScorerWeights(weights map[string]float64) error {
|
||||
for name, weight := range weights {
|
||||
if name != ScorerPrefixCache {
|
||||
return fmt.Errorf("prefixcache: unknown routing scorer %q", name)
|
||||
}
|
||||
if weight < 0 {
|
||||
return fmt.Errorf("prefixcache: scorer weight for %q must be >= 0", name)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// CandidateFilter removes replicas that are not eligible for a routing
|
||||
// decision. Filters run in declaration order.
|
||||
type CandidateFilter interface {
|
||||
Filter([]Candidate) []Candidate
|
||||
}
|
||||
|
||||
// CandidateScorer returns a normalized score in [0,1]. scored=false means the
|
||||
// scorer has no signal for this candidate and its weight is excluded.
|
||||
type CandidateScorer interface {
|
||||
Score(Candidate) (score float64, scored bool)
|
||||
}
|
||||
|
||||
// WeightedScorer configures one independent routing signal.
|
||||
type WeightedScorer struct {
|
||||
Scorer CandidateScorer
|
||||
Weight float64
|
||||
}
|
||||
|
||||
// ScoredCandidate is the normalized pipeline output passed to a picker.
|
||||
type ScoredCandidate struct {
|
||||
Candidate Candidate
|
||||
Score float64
|
||||
}
|
||||
|
||||
// CandidatePicker owns the final selection policy.
|
||||
type CandidatePicker interface {
|
||||
Pick([]ScoredCandidate) (ReplicaKey, bool)
|
||||
}
|
||||
|
||||
// RoutingPipeline composes eligibility filters and independent scorers. The
|
||||
// highest weighted mean wins; ReplicaKey ordering makes ties deterministic.
|
||||
type RoutingPipeline struct {
|
||||
Filters []CandidateFilter
|
||||
Scorers []WeightedScorer
|
||||
Picker CandidatePicker
|
||||
}
|
||||
|
||||
func (p RoutingPipeline) Pick(candidates []Candidate) (ReplicaKey, bool) {
|
||||
filtered := append([]Candidate(nil), candidates...)
|
||||
for _, filter := range p.Filters {
|
||||
filtered = filter.Filter(filtered)
|
||||
}
|
||||
if len(filtered) == 0 {
|
||||
return ReplicaKey{}, false
|
||||
}
|
||||
|
||||
scored := make([]ScoredCandidate, len(filtered))
|
||||
for i, candidate := range filtered {
|
||||
scored[i] = ScoredCandidate{Candidate: candidate, Score: p.score(candidate)}
|
||||
}
|
||||
if p.Picker != nil {
|
||||
return p.Picker.Pick(scored)
|
||||
}
|
||||
return HighestScorePicker{}.Pick(scored)
|
||||
}
|
||||
|
||||
func (p RoutingPipeline) score(candidate Candidate) float64 {
|
||||
var sum float64
|
||||
for _, weighted := range p.Scorers {
|
||||
if weighted.Scorer == nil || weighted.Weight <= 0 {
|
||||
continue
|
||||
}
|
||||
score, ok := weighted.Scorer.Score(candidate)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
score = max(0, min(1, score))
|
||||
sum += score * weighted.Weight
|
||||
}
|
||||
return sum
|
||||
}
|
||||
|
||||
// HighestScorePicker picks the maximum score and breaks ties by replica key.
|
||||
type HighestScorePicker struct{}
|
||||
|
||||
func (HighestScorePicker) Pick(candidates []ScoredCandidate) (ReplicaKey, bool) {
|
||||
if len(candidates) == 0 {
|
||||
return ReplicaKey{}, false
|
||||
}
|
||||
best := candidates[0]
|
||||
for _, candidate := range candidates[1:] {
|
||||
if candidate.Score > best.Score ||
|
||||
(candidate.Score == best.Score && candidate.Candidate.Key.less(best.Candidate.Key)) {
|
||||
best = candidate
|
||||
}
|
||||
}
|
||||
return best.Candidate.Key, true
|
||||
}
|
||||
149
core/services/nodes/prefixcache/pipeline_test.go
Normal file
149
core/services/nodes/prefixcache/pipeline_test.go
Normal file
@@ -0,0 +1,149 @@
|
||||
package prefixcache_test
|
||||
|
||||
import (
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/nodes/prefixcache"
|
||||
)
|
||||
|
||||
type filterFunc func([]prefixcache.Candidate) []prefixcache.Candidate
|
||||
|
||||
func (f filterFunc) Filter(c []prefixcache.Candidate) []prefixcache.Candidate { return f(c) }
|
||||
|
||||
type scorerFunc func(prefixcache.Candidate) (float64, bool)
|
||||
|
||||
func (f scorerFunc) Score(c prefixcache.Candidate) (float64, bool) { return f(c) }
|
||||
|
||||
type pickerFunc func([]prefixcache.ScoredCandidate) (prefixcache.ReplicaKey, bool)
|
||||
|
||||
func (f pickerFunc) Pick(c []prefixcache.ScoredCandidate) (prefixcache.ReplicaKey, bool) {
|
||||
return f(c)
|
||||
}
|
||||
|
||||
var _ = Describe("RoutingPipeline", func() {
|
||||
cand := func(node string, replica, inflight int) prefixcache.Candidate {
|
||||
return prefixcache.Candidate{Key: rk(node, replica), InFlight: inflight}
|
||||
}
|
||||
|
||||
It("filters ineligible replicas before scoring", func() {
|
||||
pipeline := prefixcache.RoutingPipeline{
|
||||
Filters: []prefixcache.CandidateFilter{filterFunc(func(c []prefixcache.Candidate) []prefixcache.Candidate {
|
||||
return c[1:]
|
||||
})},
|
||||
Scorers: []prefixcache.WeightedScorer{{
|
||||
Weight: 1,
|
||||
Scorer: scorerFunc(func(c prefixcache.Candidate) (float64, bool) {
|
||||
if c.Key.NodeID == "A" {
|
||||
return 1, true
|
||||
}
|
||||
return 0.25, true
|
||||
}),
|
||||
}},
|
||||
}
|
||||
|
||||
got, ok := pipeline.Pick([]prefixcache.Candidate{cand("A", 0, 0), cand("B", 0, 0)})
|
||||
Expect(ok).To(BeTrue())
|
||||
Expect(got).To(Equal(rk("B", 0)))
|
||||
})
|
||||
|
||||
It("combines normalized scorer results using configured weights", func() {
|
||||
pipeline := prefixcache.RoutingPipeline{
|
||||
Scorers: []prefixcache.WeightedScorer{
|
||||
{Weight: 1, Scorer: scorerFunc(func(c prefixcache.Candidate) (float64, bool) {
|
||||
if c.Key.NodeID == "A" {
|
||||
return 1, true
|
||||
}
|
||||
return 0, true
|
||||
})},
|
||||
{Weight: 3, Scorer: scorerFunc(func(c prefixcache.Candidate) (float64, bool) {
|
||||
if c.Key.NodeID == "B" {
|
||||
return 1, true
|
||||
}
|
||||
return 0, true
|
||||
})},
|
||||
},
|
||||
}
|
||||
|
||||
got, ok := pipeline.Pick([]prefixcache.Candidate{cand("A", 0, 0), cand("B", 0, 0)})
|
||||
Expect(ok).To(BeTrue())
|
||||
Expect(got).To(Equal(rk("B", 0)))
|
||||
})
|
||||
|
||||
It("ignores scorers that cannot score a candidate", func() {
|
||||
pipeline := prefixcache.RoutingPipeline{
|
||||
Scorers: []prefixcache.WeightedScorer{
|
||||
{Weight: 100, Scorer: scorerFunc(func(prefixcache.Candidate) (float64, bool) {
|
||||
return 1, false
|
||||
})},
|
||||
{Weight: 1, Scorer: scorerFunc(func(c prefixcache.Candidate) (float64, bool) {
|
||||
if c.Key.NodeID == "B" {
|
||||
return 0.75, true
|
||||
}
|
||||
return 0.25, true
|
||||
})},
|
||||
},
|
||||
}
|
||||
|
||||
got, ok := pipeline.Pick([]prefixcache.Candidate{cand("A", 0, 1), cand("B", 0, 2)})
|
||||
Expect(ok).To(BeTrue())
|
||||
Expect(got).To(Equal(rk("B", 0)))
|
||||
})
|
||||
|
||||
It("uses a weighted sum when signal availability differs by candidate", func() {
|
||||
pipeline := prefixcache.RoutingPipeline{
|
||||
Scorers: []prefixcache.WeightedScorer{
|
||||
{Weight: 1, Scorer: scorerFunc(func(c prefixcache.Candidate) (float64, bool) {
|
||||
return 0.6, c.Key.NodeID == "A"
|
||||
})},
|
||||
{Weight: 1, Scorer: scorerFunc(func(c prefixcache.Candidate) (float64, bool) {
|
||||
return 0.5, c.Key.NodeID == "B"
|
||||
})},
|
||||
{Weight: 1, Scorer: scorerFunc(func(c prefixcache.Candidate) (float64, bool) {
|
||||
return 0.5, c.Key.NodeID == "B"
|
||||
})},
|
||||
},
|
||||
}
|
||||
|
||||
got, ok := pipeline.Pick([]prefixcache.Candidate{cand("A", 0, 0), cand("B", 0, 0)})
|
||||
Expect(ok).To(BeTrue())
|
||||
Expect(got).To(Equal(rk("B", 0)))
|
||||
})
|
||||
|
||||
It("uses replica key ordering as a deterministic score tie-break", func() {
|
||||
pipeline := prefixcache.RoutingPipeline{}
|
||||
got, ok := pipeline.Pick([]prefixcache.Candidate{
|
||||
cand("B", 1, 0), cand("A", 1, 0), cand("A", 0, 0),
|
||||
})
|
||||
Expect(ok).To(BeTrue())
|
||||
Expect(got).To(Equal(rk("A", 0)))
|
||||
})
|
||||
|
||||
It("delegates final selection to the configured picker", func() {
|
||||
pipeline := prefixcache.RoutingPipeline{
|
||||
Picker: pickerFunc(func(candidates []prefixcache.ScoredCandidate) (prefixcache.ReplicaKey, bool) {
|
||||
return candidates[len(candidates)-1].Candidate.Key, true
|
||||
}),
|
||||
}
|
||||
got, ok := pipeline.Pick([]prefixcache.Candidate{cand("A", 0, 0), cand("B", 0, 0)})
|
||||
Expect(ok).To(BeTrue())
|
||||
Expect(got).To(Equal(rk("B", 0)))
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("ValidateScorerWeights", func() {
|
||||
It("accepts the named prefix-cache scorer", func() {
|
||||
Expect(prefixcache.ValidateScorerWeights(map[string]float64{
|
||||
prefixcache.ScorerPrefixCache: 0.5,
|
||||
})).To(Succeed())
|
||||
})
|
||||
|
||||
It("rejects negative and unknown scorer weights", func() {
|
||||
Expect(prefixcache.ValidateScorerWeights(map[string]float64{
|
||||
prefixcache.ScorerPrefixCache: -1,
|
||||
})).To(MatchError(ContainSubstring("must be >= 0")))
|
||||
Expect(prefixcache.ValidateScorerWeights(map[string]float64{
|
||||
"latency": 1,
|
||||
})).To(MatchError(ContainSubstring("unknown routing scorer")))
|
||||
})
|
||||
})
|
||||
@@ -1,5 +1,7 @@
|
||||
package prefixcache
|
||||
|
||||
import "sort"
|
||||
|
||||
// ReplicaKey identifies a specific loaded replica (a backend process). Affinity
|
||||
// is tracked per replica, not per node, because each replica is a separate
|
||||
// process with its own KV cache.
|
||||
@@ -47,47 +49,81 @@ func Select(cands []Candidate, d PrefixDecision, cfg Config) (ReplicaKey, bool)
|
||||
if len(cands) == 0 {
|
||||
return ReplicaKey{}, false
|
||||
}
|
||||
minIF := cands[0].InFlight
|
||||
for _, c := range cands {
|
||||
minIF = min(minIF, c.InFlight)
|
||||
pipeline := RoutingPipeline{
|
||||
Filters: []CandidateFilter{loadGuardFilter{cfg: cfg}},
|
||||
Scorers: []WeightedScorer{{
|
||||
Weight: scorerWeight(cfg.ScorerWeights, ScorerPrefixCache, 1),
|
||||
Scorer: newPrefixPreferenceScorer(cands, d, cfg),
|
||||
}},
|
||||
}
|
||||
eligible := map[ReplicaKey]bool{}
|
||||
for _, c := range cands {
|
||||
withinAbs := c.InFlight <= minIF+cfg.BalanceAbsThreshold
|
||||
// +1 softens the relative guard when minIF==0 so a zero baseline does
|
||||
// not require exact-zero in-flight; the absolute guard governs near 0.
|
||||
withinRel := float64(c.InFlight) <= float64(minIF)*cfg.BalanceRelThreshold+1
|
||||
if withinAbs && withinRel {
|
||||
eligible[c.Key] = true
|
||||
}
|
||||
}
|
||||
// Hot match wins if eligible and strong enough.
|
||||
if d.HasHot && d.MatchRatio >= cfg.MinPrefixMatch && eligible[d.Hot] {
|
||||
return d.Hot, true
|
||||
}
|
||||
// Cold placement: lowest cacheWeight eligible replica.
|
||||
for _, k := range d.ColdOrder {
|
||||
if eligible[k] {
|
||||
return k, true
|
||||
}
|
||||
}
|
||||
// Deterministic eligible fallback: least in-flight, tiebreak NodeID then
|
||||
// Replica. ColdOrder may not cover the eligible set (the caller may pass an
|
||||
// empty ColdOrder), so this guarantees Select still returns the best eligible
|
||||
// replica rather than failing.
|
||||
var best Candidate
|
||||
found := false
|
||||
for _, c := range cands {
|
||||
if !eligible[c.Key] {
|
||||
continue
|
||||
}
|
||||
if !found || c.InFlight < best.InFlight ||
|
||||
(c.InFlight == best.InFlight && c.Key.less(best.Key)) {
|
||||
best, found = c, true
|
||||
}
|
||||
}
|
||||
if found {
|
||||
return best.Key, true
|
||||
}
|
||||
return ReplicaKey{}, false
|
||||
return pipeline.Pick(cands)
|
||||
}
|
||||
|
||||
func scorerWeight(weights map[string]float64, name string, fallback float64) float64 {
|
||||
if weight, ok := weights[name]; ok {
|
||||
return weight
|
||||
}
|
||||
return fallback
|
||||
}
|
||||
|
||||
type loadGuardFilter struct {
|
||||
cfg Config
|
||||
}
|
||||
|
||||
func (f loadGuardFilter) Filter(candidates []Candidate) []Candidate {
|
||||
if len(candidates) == 0 {
|
||||
return nil
|
||||
}
|
||||
minInFlight := candidates[0].InFlight
|
||||
for _, candidate := range candidates[1:] {
|
||||
minInFlight = min(minInFlight, candidate.InFlight)
|
||||
}
|
||||
filtered := make([]Candidate, 0, len(candidates))
|
||||
for _, candidate := range candidates {
|
||||
withinAbs := candidate.InFlight <= minInFlight+f.cfg.BalanceAbsThreshold
|
||||
// +1 softens the relative guard when minInFlight is zero; the absolute
|
||||
// guard remains the controlling signal near zero.
|
||||
withinRel := float64(candidate.InFlight) <= float64(minInFlight)*f.cfg.BalanceRelThreshold+1
|
||||
if withinAbs && withinRel {
|
||||
filtered = append(filtered, candidate)
|
||||
}
|
||||
}
|
||||
return filtered
|
||||
}
|
||||
|
||||
type prefixPreferenceScorer struct {
|
||||
scores map[ReplicaKey]float64
|
||||
}
|
||||
|
||||
func newPrefixPreferenceScorer(candidates []Candidate, decision PrefixDecision, cfg Config) prefixPreferenceScorer {
|
||||
scores := make(map[ReplicaKey]float64, len(candidates))
|
||||
|
||||
// The fallback occupies the bottom quarter of the normalized range. It
|
||||
// preserves the previous least-in-flight, then replica-key ordering.
|
||||
fallback := append([]Candidate(nil), candidates...)
|
||||
sort.Slice(fallback, func(i, j int) bool {
|
||||
if fallback[i].InFlight != fallback[j].InFlight {
|
||||
return fallback[i].InFlight < fallback[j].InFlight
|
||||
}
|
||||
return fallback[i].Key.less(fallback[j].Key)
|
||||
})
|
||||
for i, candidate := range fallback {
|
||||
scores[candidate.Key] = 0.25 * float64(len(fallback)-i) / float64(len(fallback)+1)
|
||||
}
|
||||
|
||||
// Cold placement occupies the middle half, in cache-weight order.
|
||||
for i, key := range decision.ColdOrder {
|
||||
scores[key] = 0.75 - 0.25*float64(i)/float64(len(decision.ColdOrder)+1)
|
||||
}
|
||||
|
||||
// A sufficiently strong hot match remains the highest-scoring signal.
|
||||
if decision.HasHot && decision.MatchRatio >= cfg.MinPrefixMatch {
|
||||
scores[decision.Hot] = 1
|
||||
}
|
||||
return prefixPreferenceScorer{scores: scores}
|
||||
}
|
||||
|
||||
func (s prefixPreferenceScorer) Score(candidate Candidate) (float64, bool) {
|
||||
score, ok := s.scores[candidate.Key]
|
||||
return score, ok
|
||||
}
|
||||
|
||||
@@ -47,6 +47,17 @@ var _ = Describe("Select (filter-then-score)", func() {
|
||||
Expect(got).To(Equal(rk("B", 0))) // cold placement: lowest cacheWeight eligible
|
||||
})
|
||||
|
||||
It("disables prefix scoring when its configured weight is zero", func() {
|
||||
cfg := prefixcache.DefaultConfig()
|
||||
cfg.ScorerWeights[prefixcache.ScorerPrefixCache] = 0
|
||||
cands := []prefixcache.Candidate{cand("B", 0, 0), cand("A", 0, 0)}
|
||||
got, ok := prefixcache.Select(cands, prefixcache.PrefixDecision{
|
||||
Hot: rk("B", 0), HasHot: true, MatchRatio: 1,
|
||||
}, cfg)
|
||||
Expect(ok).To(BeTrue())
|
||||
Expect(got).To(Equal(rk("A", 0)))
|
||||
})
|
||||
|
||||
It("cold-places to lowest-cacheWeight replica within the eligible subset", func() {
|
||||
cands := []prefixcache.Candidate{cand("A", 0, 0), cand("B", 0, 0), cand("C", 0, 9)}
|
||||
got, ok := prefixcache.Select(cands, prefixcache.PrefixDecision{
|
||||
|
||||
@@ -193,10 +193,11 @@ type ModelSchedulingConfig struct {
|
||||
// Prefix-cache-aware routing (epic #10063). RoutePolicy "" means inherit
|
||||
// the cluster-wide default. Thresholds are per-model overrides; 0 means
|
||||
// inherit the global default.
|
||||
RoutePolicy string `gorm:"column:route_policy;size:32" json:"route_policy,omitempty"`
|
||||
BalanceAbsThreshold int `gorm:"column:balance_abs_threshold;default:0" json:"balance_abs_threshold,omitempty"`
|
||||
BalanceRelThreshold float64 `gorm:"column:balance_rel_threshold;default:0" json:"balance_rel_threshold,omitempty"`
|
||||
MinPrefixMatch float64 `gorm:"column:min_prefix_match;default:0" json:"min_prefix_match,omitempty"`
|
||||
RoutePolicy string `gorm:"column:route_policy;size:32" json:"route_policy,omitempty"`
|
||||
BalanceAbsThreshold int `gorm:"column:balance_abs_threshold;default:0" json:"balance_abs_threshold,omitempty"`
|
||||
BalanceRelThreshold float64 `gorm:"column:balance_rel_threshold;default:0" json:"balance_rel_threshold,omitempty"`
|
||||
MinPrefixMatch float64 `gorm:"column:min_prefix_match;default:0" json:"min_prefix_match,omitempty"`
|
||||
ScorerWeights map[string]float64 `gorm:"column:routing_scorer_weights;serializer:json" json:"scorer_weights,omitempty"`
|
||||
// UnsatisfiableUntil is set by the reconciler when no candidate node has
|
||||
// free capacity for this model; while in the future, the reconciler skips
|
||||
// scale-up attempts for this model. Cleared on cluster events that could
|
||||
@@ -1642,6 +1643,7 @@ func (r *NodeRegistry) SetModelScheduling(ctx context.Context, config *ModelSche
|
||||
DoUpdates: clause.AssignmentColumns([]string{
|
||||
"node_selector", "min_replicas", "max_replicas", "spread_all",
|
||||
"route_policy", "balance_abs_threshold", "balance_rel_threshold", "min_prefix_match",
|
||||
"routing_scorer_weights",
|
||||
"updated_at",
|
||||
}),
|
||||
}).
|
||||
|
||||
@@ -743,6 +743,9 @@ func (r *SmartRouter) buildPreference(ctx context.Context, modelID string, candi
|
||||
if sched.MinPrefixMatch > 0 {
|
||||
cfg.MinPrefixMatch = sched.MinPrefixMatch
|
||||
}
|
||||
if sched.ScorerWeights != nil {
|
||||
cfg.ScorerWeights = sched.ScorerWeights
|
||||
}
|
||||
}
|
||||
if policy != prefixcache.RoutePolicyPrefixCache {
|
||||
return nil, nil
|
||||
|
||||
@@ -70,10 +70,11 @@ type SeedSchedulingEntry struct {
|
||||
Replicas *ReplicasSpec `json:"replicas,omitempty" yaml:"replicas,omitempty"`
|
||||
SpreadAll bool `json:"spread_all,omitempty" yaml:"spread_all,omitempty"`
|
||||
|
||||
RoutePolicy string `json:"route_policy,omitempty" yaml:"route_policy,omitempty"`
|
||||
BalanceAbsThreshold int `json:"balance_abs_threshold,omitempty" yaml:"balance_abs_threshold,omitempty"`
|
||||
BalanceRelThreshold float64 `json:"balance_rel_threshold,omitempty" yaml:"balance_rel_threshold,omitempty"`
|
||||
MinPrefixMatch float64 `json:"min_prefix_match,omitempty" yaml:"min_prefix_match,omitempty"`
|
||||
RoutePolicy string `json:"route_policy,omitempty" yaml:"route_policy,omitempty"`
|
||||
BalanceAbsThreshold int `json:"balance_abs_threshold,omitempty" yaml:"balance_abs_threshold,omitempty"`
|
||||
BalanceRelThreshold float64 `json:"balance_rel_threshold,omitempty" yaml:"balance_rel_threshold,omitempty"`
|
||||
MinPrefixMatch float64 `json:"min_prefix_match,omitempty" yaml:"min_prefix_match,omitempty"`
|
||||
ScorerWeights map[string]float64 `json:"scorer_weights,omitempty" yaml:"scorer_weights,omitempty"`
|
||||
}
|
||||
|
||||
// spread reports whether this entry requests spread-to-all-matching-nodes mode,
|
||||
@@ -104,6 +105,9 @@ func ValidateSeedEntry(e SeedSchedulingEntry) error {
|
||||
if err := prefixcache.ValidateThresholds(e.RoutePolicy, e.BalanceAbsThreshold, e.BalanceRelThreshold, e.MinPrefixMatch); err != nil {
|
||||
return fmt.Errorf("%w (model %q)", err, e.ModelName)
|
||||
}
|
||||
if err := prefixcache.ValidateScorerWeights(e.ScorerWeights); err != nil {
|
||||
return fmt.Errorf("%w (model %q)", err, e.ModelName)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -126,6 +130,7 @@ func (e SeedSchedulingEntry) toConfig() (ModelSchedulingConfig, error) {
|
||||
BalanceAbsThreshold: e.BalanceAbsThreshold,
|
||||
BalanceRelThreshold: e.BalanceRelThreshold,
|
||||
MinPrefixMatch: e.MinPrefixMatch,
|
||||
ScorerWeights: e.ScorerWeights,
|
||||
}, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -20,6 +20,16 @@ var _ = Describe("ParseSchedulingSeed", func() {
|
||||
Expect(configs[0].NodeSelector).To(Equal(`{"tier":"gpu"}`))
|
||||
})
|
||||
|
||||
It("parses per-model routing scorer weights", func() {
|
||||
configs, err := ParseSchedulingSeed(
|
||||
`[{"model_name":"chat","route_policy":"prefix_cache","scorer_weights":{"prefix_cache":0.4}}]`,
|
||||
"",
|
||||
)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(configs).To(HaveLen(1))
|
||||
Expect(configs[0].ScorerWeights).To(Equal(map[string]float64{"prefix_cache": 0.4}))
|
||||
})
|
||||
|
||||
It("maps replicas: all to SpreadAll", func() {
|
||||
configs, err := ParseSchedulingSeed(`[{"model_name":"m","replicas":"all"}]`, "")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
@@ -944,6 +944,29 @@ Notes:
|
||||
- Verify the backend gallery configuration is correct
|
||||
- The worker needs network access to download backends from the gallery
|
||||
|
||||
## Routing pipeline
|
||||
|
||||
Loaded replicas are selected through a filter, scorer, and picker pipeline.
|
||||
The initial pipeline applies the load guard as an eligibility filter, scores
|
||||
eligible replicas using prefix-cache affinity and cold-placement order, then
|
||||
picks the highest score with a deterministic node/replica tie-break.
|
||||
|
||||
Per-model scheduling fields configure the initial pipeline:
|
||||
|
||||
- `route_policy` enables `prefix_cache` scoring or selects the
|
||||
`round_robin` floor.
|
||||
- `balance_abs_threshold` and `balance_rel_threshold` configure the load
|
||||
eligibility filter.
|
||||
- `min_prefix_match` controls when prefix affinity contributes the highest
|
||||
score.
|
||||
- `scorer_weights` enables or weights named scorers. The initial scorer is
|
||||
`prefix_cache`; set `scorer_weights: {prefix_cache: 0}` to disable its
|
||||
contribution while retaining the load filter and deterministic picker.
|
||||
|
||||
The pipeline accepts additional independently weighted scorers and alternate
|
||||
pickers without coupling them to `SmartRouter`. This is the extension point for
|
||||
queue depth, precise KV utilization, latency, and fairness signals.
|
||||
|
||||
## Roadmap: Routing and Caching Enhancements
|
||||
|
||||
The scheduling algorithm above is load-based (least in-flight, then least-recently-used). Work is underway to make routing **prefix-cache-aware**: bias each request toward the replica that already holds the relevant KV/prefix cache (multi-turn conversations and shared system prompts), so backends reuse cache instead of recomputing it. The first step is a router-side radix tree of prompt-prefix hashes mapped to nodes, with longest-prefix match, a load guard that preserves round-robin behavior under imbalance, and NATS sync across frontends. It is purely a routing-layer hint (no backend changes) and never routes worse than today's round-robin.
|
||||
@@ -952,7 +975,6 @@ Further enhancements, surfaced from a survey of SGLang, vLLM production-stack, R
|
||||
|
||||
- **Reported/precise KV-event mode** ([#10064](https://github.com/mudler/LocalAI/issues/10064)): subscribe to actual backend KV-cache events for exact residency instead of inferring it from routing history.
|
||||
- **Multi-tier cache-overlap scoring** ([#10065](https://github.com/mudler/LocalAI/issues/10065)): credit GPU/CPU/disk cache tiers separately.
|
||||
- **Pluggable scorer/filter/picker pipeline** ([#10066](https://github.com/mudler/LocalAI/issues/10066)): composable multi-signal routing (cache, queue depth, KV utilization, latency).
|
||||
- **Load-shaping** ([#10067](https://github.com/mudler/LocalAI/issues/10067)): anti-herding (softmax/temperature) and dispatch-time freshness.
|
||||
- **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.
|
||||
|
||||
Reference in New Issue
Block a user