llmgw/internal/router/scheduler/pickers.go

package scheduler

import (
	"math"
	"math/rand/v2"

	"gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint"
)

// RandSource is the pickers' randomness source; tests can pin it.
type RandSource interface {
	Float64() float64
	IntN(n int) int
}

type globalRand struct{}

func (globalRand) Float64() float64 { return rand.Float64() }
func (globalRand) IntN(n int) int   { return rand.IntN(n) }

// MaxScore picks the highest-scoring candidate.
//
// Deterministic and easy to explain, but prone to herding: with equally scored
// replicas, consecutive requests all choose the same one. Prefer P2C once there
// are many replicas.
func MaxScore() Picker { return maxScorePicker{} }

type maxScorePicker struct{}

func (maxScorePicker) Name() string { return "max-score" }

func (maxScorePicker) Pick(_ *Ctx, _ []endpoint.Snapshot, scores []float64) int {
	best, bestIdx := math.Inf(-1), 0
	for i, s := range scores {
		if s > best {
			best, bestIdx = s, i
		}
	}
	return bestIdx
}

// P2C samples two candidates at random and takes the better one
// (power of two choices).
//
// It breaks herding: even when two independent requests see the same "best"
// replica, they may draw different pairs and land elsewhere. This is the
// default picker once there are many replicas.
func P2C() Picker { return &p2cPicker{rng: globalRand{}} }

// P2CWith returns a P2C picker with a specific randomness source, for tests.
func P2CWith(r RandSource) Picker { return &p2cPicker{rng: r} }

type p2cPicker struct{ rng RandSource }

func (*p2cPicker) Name() string { return "p2c" }

func (p *p2cPicker) Pick(_ *Ctx, _ []endpoint.Snapshot, scores []float64) int {
	n := len(scores)
	if n == 0 {
		return -1
	}
	if n == 1 {
		return 0
	}
	a := p.rng.IntN(n)
	b := p.rng.IntN(n - 1)
	if b >= a {
		b++ // ensure it differs from a
	}
	if scores[a] >= scores[b] {
		return a
	}
	return b
}

// Softmax treats scores as logits and samples from them (Dynamo's
// router_temperature).
//
// A temperature of 0 degenerates to argmax. A small temperature breaks the
// oscillation that synchronized decisions cause among equally scored replicas;
// a large one approaches uniform random.
func Softmax(temperature float64) Picker {
	return &softmaxPicker{temp: temperature, rng: globalRand{}}
}

// SoftmaxWith returns a Softmax picker with a specific randomness source, for tests.
func SoftmaxWith(temperature float64, r RandSource) Picker {
	return &softmaxPicker{temp: temperature, rng: r}
}

type softmaxPicker struct {
	temp float64
	rng  RandSource
}

func (*softmaxPicker) Name() string { return "softmax" }

func (p *softmaxPicker) Pick(c *Ctx, eps []endpoint.Snapshot, scores []float64) int {
	if p.temp <= 0 {
		return maxScorePicker{}.Pick(c, eps, scores)
	}
	n := len(scores)
	if n == 0 {
		return -1
	}

	// Subtract the maximum to avoid overflow.
	maxS := math.Inf(-1)
	for _, s := range scores {
		maxS = max(maxS, s)
	}
	weights := make([]float64, n)
	var sum float64
	for i, s := range scores {
		w := math.Exp((s - maxS) / p.temp)
		weights[i] = w
		sum += w
	}
	if sum <= 0 || math.IsInf(sum, 0) || math.IsNaN(sum) {
		return maxScorePicker{}.Pick(c, eps, scores)
	}

	r := p.rng.Float64() * sum
	for i, w := range weights {
		r -= w
		if r <= 0 {
			return i
		}
	}
	return n - 1
}

// RoundRobin cycles through candidates. For debugging and simple setups.
func RoundRobin() Picker { return &rrPicker{} }

type rrPicker struct{ n int }

func (*rrPicker) Name() string { return "round-robin" }

func (p *rrPicker) Pick(_ *Ctx, _ []endpoint.Snapshot, scores []float64) int {
	if len(scores) == 0 {
		return -1
	}
	i := p.n % len(scores)
	p.n++
	return i
}