internal/router: priority queues so batch jobs yield to interactive requests
The GPU is a size-1 resource, so a single long job monopolises the box for its
whole duration and every interactive request queues behind it. Callers can now
declare intent with an X-LlamaSwap-Priority header and the serial scheduler
dispatches by score instead of by arrival.
- X-LlamaSwap-Priority: signed integer, 0 default, absent/unparseable means 0.
interactive/normal/batch aliases resolve to +100/0/-100. Values are not
clamped: the caller composes band and any per-user offset itself.
- serial dispatch score = priority + swap affinity + aging. Bands sit 100 apart
so a small caller offset orders work inside a band without crossing one;
aging is unbounded so low-priority work cannot starve.
- routing.scheduler.settings.serial.{agingDivisor,swapAffinityBonus}, defaulting
to 60s/point and +10. swapAffinityBonus is capped at 99 so it can never
promote a request into the next band.
- fifo adds the header to its per-model priority, so the header is not silently
ignored under that scheduler.
- /metrics exports per-band queue depth, oldest wait and dispatch counts, plus
counters for how often aging or swap affinity changed the pick. Each request
records its priority, band, score and queue wait in the activity log.
Note swapAffinityBonus defaults to 10, so equal-priority requests for the
already-loaded model now run before older requests that need a swap. Set it to
0 for the previous strict arrival order.
fixes #9
Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
Claude-Session: https://claude.ai/code/session_01WUyhZBgv8BBCC5MduX88gE
This commit is contained in:
+14
-2
@@ -338,6 +338,17 @@ func (b *baseRouter) Handles(model string) bool {
|
||||
return ok
|
||||
}
|
||||
|
||||
// SchedulerStats returns a snapshot of the scheduler's queue. ok is false when
|
||||
// the configured scheduler does not report stats (only "serial" does today), in
|
||||
// which case callers should omit the metrics rather than emit zeroes.
|
||||
func (b *baseRouter) SchedulerStats() (scheduler.QueueStats, bool) {
|
||||
reporter, ok := b.schedule.(scheduler.StatsReporter)
|
||||
if !ok {
|
||||
return scheduler.QueueStats{}, false
|
||||
}
|
||||
return reporter.QueueStats(), true
|
||||
}
|
||||
|
||||
func (b *baseRouter) ProcessLogger(modelID string) (*logmon.Monitor, bool) {
|
||||
if p, ok := b.processes[modelID]; ok {
|
||||
return p.Logger(), true
|
||||
@@ -420,8 +431,9 @@ func (b *baseRouter) ServeHTTP(w http.ResponseWriter, req *http.Request) {
|
||||
}
|
||||
|
||||
hr := scheduler.HandlerReq{
|
||||
Model: data.ModelID,
|
||||
Ctx: req.Context(),
|
||||
Model: data.ModelID,
|
||||
Priority: data.Priority,
|
||||
Ctx: req.Context(),
|
||||
// Unbuffered: a successful send on Respond proves the waiter is
|
||||
// alive and consuming. grant() relies on this to avoid handing a
|
||||
// handleFunc to a cancelled waiter and leaking the inFlight count.
|
||||
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@@ -12,6 +13,7 @@ import (
|
||||
"github.com/mostlygeek/llama-swap/internal/logmon"
|
||||
"github.com/mostlygeek/llama-swap/internal/process"
|
||||
"github.com/mostlygeek/llama-swap/internal/router/scheduler"
|
||||
"github.com/mostlygeek/llama-swap/internal/shared"
|
||||
)
|
||||
|
||||
// These tests cover baseRouter's own machinery — the run loop, process
|
||||
@@ -218,6 +220,76 @@ func TestBaseRouter_ContextCancel(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestBaseRouter_SerialPriorityHeader is the end-to-end path for issue #9: the
|
||||
// X-LlamaSwap-Priority header on an HTTP request reaches the serial scheduler
|
||||
// and changes which queued request runs next.
|
||||
//
|
||||
// "hold" occupies the single slot. batch and urgent queue behind it in that
|
||||
// order; when the slot frees, urgent must go first despite arriving later.
|
||||
// urgent blocks inside its handler, so observing urgent's handler entry proves
|
||||
// the scheduler chose it — if batch had won, batch's (unblocked) handler would
|
||||
// already have run by then.
|
||||
func TestBaseRouter_SerialPriorityHeader(t *testing.T) {
|
||||
hold := newFakeProcess("hold")
|
||||
hold.autoReady = true
|
||||
hold.serveBlock = make(chan struct{})
|
||||
batch := newFakeProcess("batch")
|
||||
batch.autoReady = true
|
||||
urgent := newFakeProcess("urgent")
|
||||
urgent.autoReady = true
|
||||
urgent.serveBlock = make(chan struct{})
|
||||
|
||||
conf := config.Config{HealthCheckTimeout: 5}
|
||||
conf.Routing.Scheduler.Use = "serial"
|
||||
b, err := newBaseRouter("test", conf, map[string]process.Process{
|
||||
"hold": hold, "batch": batch, "urgent": urgent,
|
||||
}, logmon.NewWriter(io.Discard), &stubPlanner{})
|
||||
if err != nil {
|
||||
t.Fatalf("newBaseRouter: %v", err)
|
||||
}
|
||||
b.testProcessed = make(chan struct{}, 64)
|
||||
go b.run()
|
||||
|
||||
releaseHold := sync.OnceFunc(func() { close(hold.serveBlock) })
|
||||
releaseUrgent := sync.OnceFunc(func() { close(urgent.serveBlock) })
|
||||
t.Cleanup(func() {
|
||||
releaseHold()
|
||||
releaseUrgent()
|
||||
if !b.shuttingDown.Load() {
|
||||
_ = b.Shutdown(time.Second)
|
||||
}
|
||||
})
|
||||
|
||||
serve := func(model, priority string) {
|
||||
r := newRequest(model)
|
||||
if priority != "" {
|
||||
r.Header.Set(shared.PriorityHeader, priority)
|
||||
}
|
||||
go b.ServeHTTP(httptest.NewRecorder(), r)
|
||||
}
|
||||
|
||||
serve("hold", "")
|
||||
waitProcessed(t, b.testProcessed, 2) // OnRequest, then the swap completing
|
||||
<-hold.serveStarted
|
||||
|
||||
// Both queue behind hold; batch arrives first.
|
||||
serve("batch", "batch")
|
||||
waitProcessed(t, b.testProcessed, 1)
|
||||
serve("urgent", "interactive")
|
||||
waitProcessed(t, b.testProcessed, 1)
|
||||
|
||||
releaseHold()
|
||||
|
||||
select {
|
||||
case <-urgent.serveStarted:
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("interactive request never started")
|
||||
}
|
||||
if got := batch.serveCalls.Load(); got != 0 {
|
||||
t.Errorf("batch serveCalls=%d want 0 — the batch job ran before the interactive request", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBaseRouter_ModelNotFound(t *testing.T) {
|
||||
a := newFakeProcess("a")
|
||||
b := newTestBase(t, map[string]process.Process{"a": a}, &stubPlanner{})
|
||||
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
|
||||
"github.com/mostlygeek/llama-swap/internal/logmon"
|
||||
"github.com/mostlygeek/llama-swap/internal/process"
|
||||
"github.com/mostlygeek/llama-swap/internal/router/scheduler"
|
||||
"github.com/mostlygeek/llama-swap/internal/shared"
|
||||
)
|
||||
|
||||
@@ -49,4 +50,8 @@ type LocalRouter interface {
|
||||
// modelID must be a real (non-alias) config key. Returns false when the
|
||||
// model is not known to this router.
|
||||
ProcessLogger(modelID string) (*logmon.Monitor, bool)
|
||||
|
||||
// SchedulerStats returns a snapshot of the scheduler's queue for metrics.
|
||||
// ok is false when the configured scheduler does not report stats.
|
||||
SchedulerStats() (scheduler.QueueStats, bool)
|
||||
}
|
||||
|
||||
@@ -280,7 +280,7 @@ func (s *FIFO) grantHandler(req HandlerReq, modelID string) {
|
||||
return
|
||||
}
|
||||
|
||||
if err := shared.SetReqData(req.Ctx, "fifo_priority", strconv.Itoa(s.cfg.Priority[req.Model])); err != nil {
|
||||
if err := shared.SetReqData(req.Ctx, "fifo_priority", strconv.Itoa(s.priorityOf(req))); err != nil {
|
||||
s.logger.Debugf("failed to set fifo_priority metadata: %v", err)
|
||||
}
|
||||
|
||||
@@ -311,15 +311,22 @@ func (s *FIFO) startSwap(initial HandlerReq, evict, running []string) {
|
||||
s.effects.StartSwap(initial.Model, evict)
|
||||
}
|
||||
|
||||
// priorityOf is a request's effective queue priority: the caller's
|
||||
// X-LlamaSwap-Priority plus the model's configured priority. Both default to 0,
|
||||
// so a deployment that uses neither keeps plain FIFO order.
|
||||
func (s *FIFO) priorityOf(req HandlerReq) int {
|
||||
return req.Priority + s.cfg.Priority[req.Model]
|
||||
}
|
||||
|
||||
// enqueue inserts req into the queue in priority order: it goes just before the
|
||||
// first queued item whose priority is strictly lower, so higher-priority models
|
||||
// are serviced first while equal-priority requests keep their arrival (FIFO)
|
||||
// order. Priorities come from the FifoConfig; unlisted models default to 0.
|
||||
// first queued item whose priority is strictly lower, so higher-priority
|
||||
// requests are serviced first while equal-priority requests keep their arrival
|
||||
// (FIFO) order.
|
||||
func (s *FIFO) enqueue(req HandlerReq) {
|
||||
p := s.cfg.Priority[req.Model]
|
||||
p := s.priorityOf(req)
|
||||
i := len(s.queued)
|
||||
for j, q := range s.queued {
|
||||
if s.cfg.Priority[q.Model] < p {
|
||||
if s.priorityOf(q) < p {
|
||||
i = j
|
||||
break
|
||||
}
|
||||
|
||||
@@ -571,6 +571,34 @@ func TestFIFO_PriorityQueueOrder(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestFIFO_RequestPriorityOrder verifies the caller's X-LlamaSwap-Priority is
|
||||
// added to the model's configured priority, so a batch caller for a
|
||||
// high-priority model still queues behind an interactive caller for a
|
||||
// low-priority one.
|
||||
func TestFIFO_RequestPriorityOrder(t *testing.T) {
|
||||
eff := newFakeEffects()
|
||||
for _, m := range []string{"z", "hot", "cold"} {
|
||||
eff.states[m] = process.StateStopped
|
||||
}
|
||||
planner := &stubPlanner{evict: map[string][]string{"z": {"hot", "cold"}}}
|
||||
cfg := config.FifoConfig{Priority: map[string]int{"hot": 10}}
|
||||
s := NewFIFO("test", logmon.NewWriter(io.Discard), planner, cfg, nil, eff)
|
||||
|
||||
s.OnRequest(req("z")) // StartSwap(z, [hot, cold]) — everything else queues
|
||||
|
||||
s.OnRequest(HandlerReq{Model: "hot", Priority: shared.PriorityBatch}) // 10 - 100 = -90
|
||||
s.OnRequest(HandlerReq{Model: "cold", Priority: shared.PriorityInteractive}) // 0 + 100 = 100
|
||||
|
||||
got := make([]string, len(s.queued))
|
||||
for i, q := range s.queued {
|
||||
got[i] = q.Model
|
||||
}
|
||||
want := []string{"cold", "hot"}
|
||||
if len(got) != len(want) || got[0] != want[0] || got[1] != want[1] {
|
||||
t.Fatalf("queue=%v want %v", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
// TestFIFO_OnCancel_QueuedRequest verifies that cancelling a queued request
|
||||
// prevents drainQueue from ever starting a model load for it. Without OnCancel
|
||||
// the dead request would sit in the queue until a drain triggers a wasted swap.
|
||||
|
||||
@@ -95,7 +95,7 @@ type Effects interface {
|
||||
// New returns a Scheduler selected by conf.Routing.Scheduler.Use, configured from
|
||||
// conf and bound to the given planner and effects. Supported values are "fifo"
|
||||
// (throughput-oriented, batches same-model requests) and "serial" (strict
|
||||
// one-model-at-a-time, exact arrival order).
|
||||
// one-model-at-a-time, highest-priority-first).
|
||||
//
|
||||
// The deployment default is applied by config loading (LoadConfig sets Use to
|
||||
// "serial" when unset). The "" fallback here is the library default and remains
|
||||
@@ -110,7 +110,7 @@ func New(conf config.Config, name string, logger *logmon.Monitor, planner Swappe
|
||||
return NewFIFO(name, logger, planner, conf.Routing.Scheduler.Settings.Fifo, conf.Models, eff), nil
|
||||
case "serial":
|
||||
// Serial ignores the group planner: it always evicts every other model.
|
||||
return NewSerial(name, logger, eff), nil
|
||||
return NewSerial(name, logger, conf.Routing.Scheduler.Settings.Serial, eff), nil
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported scheduler type: %q", use)
|
||||
}
|
||||
@@ -118,12 +118,50 @@ func New(conf config.Config, name string, logger *logmon.Monitor, planner Swappe
|
||||
|
||||
// HandlerReq is one in-flight ServeHTTP request waiting for a routing decision.
|
||||
type HandlerReq struct {
|
||||
Model string
|
||||
Model string
|
||||
// Priority is the caller's scheduling priority from the
|
||||
// X-LlamaSwap-Priority header. Higher is more urgent, 0 is normal, and
|
||||
// negative is batch work. Schedulers use it to order their queues.
|
||||
Priority int
|
||||
Ctx context.Context
|
||||
Respond chan HandlerResp
|
||||
PositionCh chan int
|
||||
}
|
||||
|
||||
// BandStats is the queue state for one priority band.
|
||||
type BandStats struct {
|
||||
// Depth is how many requests are waiting in this band.
|
||||
Depth int
|
||||
// OldestWait is how long the longest-waiting request in this band has
|
||||
// been queued, or 0 when the band is empty.
|
||||
OldestWait time.Duration
|
||||
// Dispatched counts requests in this band handed off since start.
|
||||
Dispatched uint64
|
||||
}
|
||||
|
||||
// QueueStats is a point-in-time snapshot of a scheduler's queue, for metrics.
|
||||
// The reorder counters answer "did the scoring function actually change
|
||||
// anything?" — without them, tuning aging and swap affinity is guesswork.
|
||||
type QueueStats struct {
|
||||
// Bands is keyed by shared.PriorityBand and always carries an entry for
|
||||
// every band in shared.PriorityBands, so metric series stay stable.
|
||||
Bands map[string]BandStats
|
||||
// AgingReorders counts dispatches where the aging term changed which
|
||||
// request was picked.
|
||||
AgingReorders uint64
|
||||
// AffinityReorders counts dispatches where the swap-affinity term changed
|
||||
// which request was picked.
|
||||
AffinityReorders uint64
|
||||
}
|
||||
|
||||
// StatsReporter is implemented by schedulers that can report queue statistics.
|
||||
// It is optional: callers type-assert and skip schedulers that do not provide
|
||||
// it. Unlike the Scheduler methods, QueueStats is called from arbitrary
|
||||
// goroutines (the /metrics handler) and must be safe for concurrent use.
|
||||
type StatsReporter interface {
|
||||
QueueStats() QueueStats
|
||||
}
|
||||
|
||||
// HandlerResp is the routing decision returned to a HandlerReq's caller: either
|
||||
// a handler to serve with, or an error.
|
||||
type HandlerResp struct {
|
||||
|
||||
@@ -3,40 +3,101 @@ package scheduler
|
||||
import (
|
||||
"fmt"
|
||||
"sort"
|
||||
"strconv"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/mostlygeek/llama-swap/internal/config"
|
||||
"github.com/mostlygeek/llama-swap/internal/logmon"
|
||||
"github.com/mostlygeek/llama-swap/internal/process"
|
||||
"github.com/mostlygeek/llama-swap/internal/shared"
|
||||
)
|
||||
|
||||
// Serial is a strict one-model-at-a-time scheduler. Unlike FIFO it never reorders
|
||||
// or batches: requests run in exact arrival order and at most one request runs at
|
||||
// any instant. When the next request targets a model other than the one loaded,
|
||||
// every other running model is evicted and the target is loaded before it runs,
|
||||
// so a single model occupies memory at a time — at the cost of throughput.
|
||||
//
|
||||
// Example: A B C A is served as A B C A. The final A reloads its model even
|
||||
// though it ran first, because B and C displaced it in between. (FIFO, by
|
||||
// contrast, would batch the two A requests: A A B C.)
|
||||
// Serial is a strict one-model-at-a-time scheduler for a size-1 resource: at
|
||||
// most one request runs at any instant, and when the next request targets a
|
||||
// model other than the one loaded, every other running model is evicted and the
|
||||
// target is loaded before it runs. A single model occupies memory at a time, at
|
||||
// the cost of throughput.
|
||||
//
|
||||
// Serial ignores group/eviction policy entirely: it always evicts every other
|
||||
// running model, regardless of how groups are configured. That is what makes the
|
||||
// single-model guarantee a property of the scheduler rather than of the config.
|
||||
// running model, regardless of how groups are configured. That is what makes
|
||||
// the single-model guarantee a property of the scheduler rather than of the
|
||||
// config.
|
||||
//
|
||||
// Like FIFO, every method runs on the router's single run-loop goroutine, so no
|
||||
// internal locking is needed.
|
||||
// # Dispatch order
|
||||
//
|
||||
// Scheduling is non-preemptive — a running job is never interrupted, because a
|
||||
// generation cannot be stopped mid-sample without discarding the work. Priority
|
||||
// applies at dispatch time: whenever the slot frees, the highest scoring
|
||||
// waiting request goes next, where
|
||||
//
|
||||
// score = request priority + swap affinity + aging
|
||||
//
|
||||
// - request priority is the caller's X-LlamaSwap-Priority (0 by default).
|
||||
// Bands sit 100 apart: +100 interactive, 0 normal, -100 batch.
|
||||
// - swap affinity is a bounded bonus for a request that can run without a
|
||||
// model swap. Bounded well below the band spacing, so it only breaks
|
||||
// near-ties and can never promote batch work above interactive.
|
||||
// - aging is waited_seconds / agingDivisor, and is unbounded on purpose: it
|
||||
// is the only term allowed to cross a band, which is what stops a batch job
|
||||
// from starving behind an endless stream of normal traffic.
|
||||
//
|
||||
// Equal scores keep arrival order, so with default priorities and swap affinity
|
||||
// disabled this degrades to exact FIFO.
|
||||
//
|
||||
// The consequence of being non-preemptive is worth stating plainly: an
|
||||
// interactive request can still wait up to one full job duration. Priority
|
||||
// alone does not fix that — the complementary lever is client-side, where a
|
||||
// batch producer submits one unit at a time and re-queues so gaps are frequent.
|
||||
//
|
||||
// Like FIFO, every Scheduler method runs on the router's single run-loop
|
||||
// goroutine, so no internal locking is needed for queue state. QueueStats is
|
||||
// the exception: it is read from the /metrics handler, so the stats mirror it
|
||||
// reads is guarded by a mutex.
|
||||
type Serial struct {
|
||||
name string
|
||||
logger *logmon.Monitor
|
||||
effects Effects
|
||||
|
||||
// queued holds requests in strict arrival order. It is never reordered.
|
||||
queued []HandlerReq
|
||||
// agingDivisor is how long a request must wait to gain one priority
|
||||
// point. Zero disables aging.
|
||||
agingDivisor time.Duration
|
||||
// swapAffinityBonus is added to a request that needs no model swap. Zero
|
||||
// disables the term.
|
||||
swapAffinityBonus int
|
||||
|
||||
// now is the clock, injectable so tests can drive aging deterministically.
|
||||
now func() time.Time
|
||||
|
||||
// queued holds waiting requests in arrival order. Dispatch picks by score
|
||||
// rather than popping the head, but the slice itself is never reordered so
|
||||
// index order remains arrival order — which is what makes equal scores
|
||||
// break in favour of the earlier request.
|
||||
queued []queuedReq
|
||||
|
||||
// active is the one request currently being processed (loading or serving),
|
||||
// or nil when idle. phase is meaningful only while active != nil.
|
||||
active *HandlerReq
|
||||
phase serialPhase
|
||||
|
||||
mu sync.Mutex
|
||||
stats serialStats
|
||||
}
|
||||
|
||||
// queuedReq is one waiting request plus the arrival time aging is measured from.
|
||||
type queuedReq struct {
|
||||
req HandlerReq
|
||||
enqueued time.Time
|
||||
}
|
||||
|
||||
// serialStats mirrors queue state and dispatch counters for QueueStats. It is
|
||||
// written on the run loop and read by the metrics handler, both under mu.
|
||||
type serialStats struct {
|
||||
depth map[string]int
|
||||
oldest map[string]time.Time
|
||||
dispatched map[string]uint64
|
||||
agingReorders uint64
|
||||
affinityReorders uint64
|
||||
}
|
||||
|
||||
// serialPhase is the lifecycle stage of the active request.
|
||||
@@ -50,67 +111,183 @@ const (
|
||||
|
||||
// NewSerial builds a Serial scheduler. It takes no Swapper: eviction is always
|
||||
// "stop every other running model", so the group planner is not consulted.
|
||||
func NewSerial(name string, logger *logmon.Monitor, eff Effects) *Serial {
|
||||
func NewSerial(name string, logger *logmon.Monitor, cfg config.SerialConfig, eff Effects) *Serial {
|
||||
return &Serial{
|
||||
name: name,
|
||||
logger: logger,
|
||||
effects: eff,
|
||||
name: name,
|
||||
logger: logger,
|
||||
effects: eff,
|
||||
agingDivisor: time.Duration(cfg.GetAgingDivisor()) * time.Second,
|
||||
swapAffinityBonus: cfg.GetSwapAffinityBonus(),
|
||||
now: time.Now,
|
||||
stats: serialStats{
|
||||
depth: make(map[string]int),
|
||||
oldest: make(map[string]time.Time),
|
||||
dispatched: make(map[string]uint64),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// OnRequest validates the model and appends the request to the tail of the queue,
|
||||
// then tries to start the next job. Unknown models fail immediately.
|
||||
// OnRequest validates the model and appends the request to the queue, stamping
|
||||
// its arrival time so aging can be measured, then tries to start the next job.
|
||||
// Unknown models fail immediately.
|
||||
func (s *Serial) OnRequest(req HandlerReq) {
|
||||
if _, ok := s.effects.ModelState(req.Model); !ok {
|
||||
s.logger.Debugf("%s: model %s not handled by this router", s.name, req.Model)
|
||||
s.effects.GrantError(req, ErrModelNotFound)
|
||||
return
|
||||
}
|
||||
s.queued = append(s.queued, req)
|
||||
broadcastQueuePositions(s.queued)
|
||||
s.queued = append(s.queued, queuedReq{req: req, enqueued: s.now()})
|
||||
s.queueChanged()
|
||||
s.startNext()
|
||||
}
|
||||
|
||||
// startNext begins processing the head of the queue when nothing is active. It
|
||||
// fast-paths a request whose model is already the sole loaded-and-ready process;
|
||||
// otherwise it launches a swap that evicts every other running model first. The
|
||||
// loop skips over requests for models that vanished (e.g. a config reload) and
|
||||
// requests whose caller disconnected before they could be served.
|
||||
// startNext begins processing the highest scoring waiting request when nothing
|
||||
// is active. It fast-paths a request whose model is already the sole
|
||||
// loaded-and-ready process; otherwise it launches a swap that evicts every other
|
||||
// running model first. The loop skips over requests for models that vanished
|
||||
// (e.g. a config reload) and requests whose caller disconnected before they
|
||||
// could be served.
|
||||
func (s *Serial) startNext() {
|
||||
if s.active != nil {
|
||||
return // a job is already loading or serving
|
||||
}
|
||||
for len(s.queued) > 0 {
|
||||
req := s.queued[0]
|
||||
s.queued = s.queued[1:]
|
||||
broadcastQueuePositions(s.queued)
|
||||
idx, score := s.pick()
|
||||
q := s.queued[idx]
|
||||
s.queued = append(s.queued[:idx], s.queued[idx+1:]...)
|
||||
s.queueChanged()
|
||||
|
||||
req := q.req
|
||||
state, ok := s.effects.ModelState(req.Model)
|
||||
if !ok {
|
||||
s.effects.GrantError(req, ErrModelNotFound)
|
||||
continue
|
||||
}
|
||||
|
||||
waited := s.now().Sub(q.enqueued)
|
||||
s.annotate(req, score, waited)
|
||||
|
||||
r := req
|
||||
s.active = &r
|
||||
|
||||
evict := s.otherRunning(req.Model)
|
||||
if state == process.StateReady && len(evict) == 0 {
|
||||
// Already loaded and the only model running — serve immediately.
|
||||
s.logger.Debugf("%s: serving model %s (already loaded)", s.name, req.Model)
|
||||
s.logger.Debugf("%s: serving model %s (already loaded, priority %d, score %d, waited %s)",
|
||||
s.name, req.Model, req.Priority, score, waited.Round(time.Millisecond))
|
||||
if s.serve() {
|
||||
s.recordDispatch(req.Priority)
|
||||
return
|
||||
}
|
||||
continue // caller gone; pick the next request
|
||||
}
|
||||
|
||||
s.logger.Debugf("%s: swapping to model %s, evicting %v", s.name, req.Model, evict)
|
||||
s.logger.Debugf("%s: swapping to model %s (priority %d, score %d, waited %s), evicting %v",
|
||||
s.name, req.Model, req.Priority, score, waited.Round(time.Millisecond), evict)
|
||||
s.phase = phaseSwapping
|
||||
s.recordDispatch(req.Priority)
|
||||
s.effects.StartSwap(req.Model, evict)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// pick returns the index of the highest scoring queued request and its score.
|
||||
// It also accounts for whether the aging and swap-affinity terms changed the
|
||||
// outcome: each is recomputed with that term removed, and a different winner
|
||||
// means the term reordered this dispatch. The queue is never empty here.
|
||||
func (s *Serial) pick() (int, int) {
|
||||
now := s.now()
|
||||
resident, _ := s.noSwapModel()
|
||||
|
||||
idx, score := s.best(now, resident, true, true)
|
||||
if len(s.queued) > 1 {
|
||||
if other, _ := s.best(now, resident, false, true); other != idx {
|
||||
s.countReorder(true)
|
||||
}
|
||||
if other, _ := s.best(now, resident, true, false); other != idx {
|
||||
s.countReorder(false)
|
||||
}
|
||||
}
|
||||
return idx, score
|
||||
}
|
||||
|
||||
// best returns the index and score of the highest scoring queued request, with
|
||||
// the aging and swap-affinity terms individually switchable so pick can measure
|
||||
// their effect. Ties keep the earliest arrival because the queue is in arrival
|
||||
// order and the comparison is strictly greater-than.
|
||||
func (s *Serial) best(now time.Time, resident string, withAging, withAffinity bool) (int, int) {
|
||||
bestIdx, bestScore := 0, 0
|
||||
for i, q := range s.queued {
|
||||
score := q.req.Priority
|
||||
if withAffinity && resident != "" && q.req.Model == resident {
|
||||
score += s.swapAffinityBonus
|
||||
}
|
||||
if withAging {
|
||||
score += s.aging(q, now)
|
||||
}
|
||||
if i == 0 || score > bestScore {
|
||||
bestIdx, bestScore = i, score
|
||||
}
|
||||
}
|
||||
return bestIdx, bestScore
|
||||
}
|
||||
|
||||
// aging converts how long a request has waited into priority points. It is
|
||||
// unbounded: a request that waits long enough eventually outranks anything.
|
||||
func (s *Serial) aging(q queuedReq, now time.Time) int {
|
||||
if s.agingDivisor <= 0 {
|
||||
return 0
|
||||
}
|
||||
waited := now.Sub(q.enqueued)
|
||||
if waited <= 0 {
|
||||
return 0
|
||||
}
|
||||
return int(waited / s.agingDivisor)
|
||||
}
|
||||
|
||||
// noSwapModel returns the model that can be served without any swap: the sole
|
||||
// running process, when it is ready. Anything else — nothing loaded, a process
|
||||
// still loading, or more than one running — means every request costs a swap,
|
||||
// so no request earns the affinity bonus.
|
||||
func (s *Serial) noSwapModel() (string, bool) {
|
||||
if s.swapAffinityBonus == 0 {
|
||||
return "", false
|
||||
}
|
||||
running := s.effects.RunningModels()
|
||||
if len(running) != 1 {
|
||||
return "", false
|
||||
}
|
||||
for id, state := range running {
|
||||
if state == process.StateReady {
|
||||
return id, true
|
||||
}
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
|
||||
// annotate records the dispatch decision on the request's context so it reaches
|
||||
// the activity log. This is how "why did my request take 20 minutes" gets
|
||||
// answered after the fact.
|
||||
func (s *Serial) annotate(req HandlerReq, score int, waited time.Duration) {
|
||||
if req.Ctx == nil {
|
||||
return
|
||||
}
|
||||
fields := [...]struct{ key, value string }{
|
||||
{"serial_priority", strconv.Itoa(req.Priority)},
|
||||
{"serial_band", shared.PriorityBand(req.Priority)},
|
||||
{"serial_score", strconv.Itoa(score)},
|
||||
{"serial_queue_wait_ms", strconv.FormatInt(waited.Milliseconds(), 10)},
|
||||
}
|
||||
for _, f := range fields {
|
||||
if err := shared.SetReqData(req.Ctx, f.key, f.value); err != nil {
|
||||
// Every key shares one context, so the first failure means none
|
||||
// of them will land.
|
||||
s.logger.Debugf("%s: failed to set %s metadata: %v", s.name, f.key, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// serve hands the active request its tracked handler. It returns true when the
|
||||
// request is now serving (await OnServeDone); false when the caller had already
|
||||
// disconnected, in which case active is cleared so the next job can start.
|
||||
@@ -168,7 +345,7 @@ func (s *Serial) OnCancel(req HandlerReq) {
|
||||
kept := s.queued[:0]
|
||||
removed := false
|
||||
for _, q := range s.queued {
|
||||
if q.Respond == req.Respond {
|
||||
if q.req.Respond == req.Respond {
|
||||
removed = true
|
||||
continue
|
||||
}
|
||||
@@ -177,7 +354,7 @@ func (s *Serial) OnCancel(req HandlerReq) {
|
||||
s.queued = kept
|
||||
if removed {
|
||||
s.logger.Debugf("%s: cancelled request for model %s pruned from queue", s.name, req.Model)
|
||||
broadcastQueuePositions(s.queued)
|
||||
s.queueChanged()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -205,14 +382,14 @@ func (s *Serial) OnUnload(targets []string, timeout time.Duration) {
|
||||
if len(s.queued) > 0 {
|
||||
kept := s.queued[:0]
|
||||
for _, q := range s.queued {
|
||||
if targetSet[q.Model] {
|
||||
s.effects.GrantError(q, unloadErr)
|
||||
if targetSet[q.req.Model] {
|
||||
s.effects.GrantError(q.req, unloadErr)
|
||||
continue
|
||||
}
|
||||
kept = append(kept, q)
|
||||
}
|
||||
s.queued = kept
|
||||
broadcastQueuePositions(s.queued)
|
||||
s.queueChanged()
|
||||
}
|
||||
|
||||
s.effects.StopProcesses(timeout, targets)
|
||||
@@ -234,9 +411,10 @@ func (s *Serial) OnShutdown(err error) {
|
||||
s.phase = phaseIdle
|
||||
}
|
||||
for _, q := range s.queued {
|
||||
s.effects.GrantError(q, err)
|
||||
s.effects.GrantError(q.req, err)
|
||||
}
|
||||
s.queued = nil
|
||||
s.queueChanged()
|
||||
}
|
||||
|
||||
// otherRunning returns every running model except target, sorted for
|
||||
@@ -251,3 +429,108 @@ func (s *Serial) otherRunning(target string) []string {
|
||||
sort.Strings(out)
|
||||
return out
|
||||
}
|
||||
|
||||
// queueChanged is called after every mutation of s.queued. It refreshes the
|
||||
// stats mirror and tells each waiter its new position.
|
||||
func (s *Serial) queueChanged() {
|
||||
s.syncStats()
|
||||
s.broadcastPositions()
|
||||
}
|
||||
|
||||
// broadcastPositions sends each waiter its 1-indexed place in dispatch order
|
||||
// rather than arrival order, so a queued caller sees the position priority
|
||||
// actually earned it. The ranking is a snapshot: aging keeps moving, so a
|
||||
// position only holds until the next queue change.
|
||||
func (s *Serial) broadcastPositions() {
|
||||
if len(s.queued) == 0 {
|
||||
return
|
||||
}
|
||||
now := s.now()
|
||||
resident, _ := s.noSwapModel()
|
||||
|
||||
scores := make([]int, len(s.queued))
|
||||
order := make([]int, len(s.queued))
|
||||
for i, q := range s.queued {
|
||||
score := q.req.Priority
|
||||
if resident != "" && q.req.Model == resident {
|
||||
score += s.swapAffinityBonus
|
||||
}
|
||||
scores[i] = score + s.aging(q, now)
|
||||
order[i] = i
|
||||
}
|
||||
// Stable so equal scores keep arrival order, matching best().
|
||||
sort.SliceStable(order, func(a, b int) bool { return scores[order[a]] > scores[order[b]] })
|
||||
|
||||
ranked := make([]HandlerReq, len(order))
|
||||
for rank, i := range order {
|
||||
ranked[rank] = s.queued[i].req
|
||||
}
|
||||
broadcastQueuePositions(ranked)
|
||||
}
|
||||
|
||||
// syncStats refreshes the per-band depth and oldest-arrival mirror that
|
||||
// QueueStats reads. Enqueue times are stored rather than durations so the wait
|
||||
// is computed fresh at read time instead of going stale between queue changes.
|
||||
func (s *Serial) syncStats() {
|
||||
depth := make(map[string]int, len(shared.PriorityBands))
|
||||
oldest := make(map[string]time.Time, len(shared.PriorityBands))
|
||||
for _, q := range s.queued {
|
||||
band := shared.PriorityBand(q.req.Priority)
|
||||
depth[band]++
|
||||
if t, ok := oldest[band]; !ok || q.enqueued.Before(t) {
|
||||
oldest[band] = q.enqueued
|
||||
}
|
||||
}
|
||||
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.stats.depth = depth
|
||||
s.stats.oldest = oldest
|
||||
}
|
||||
|
||||
// recordDispatch counts one request handed off, bucketed by band.
|
||||
func (s *Serial) recordDispatch(priority int) {
|
||||
band := shared.PriorityBand(priority)
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.stats.dispatched[band]++
|
||||
}
|
||||
|
||||
// countReorder records that aging (or swap affinity) changed which request a
|
||||
// dispatch picked.
|
||||
func (s *Serial) countReorder(aging bool) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if aging {
|
||||
s.stats.agingReorders++
|
||||
} else {
|
||||
s.stats.affinityReorders++
|
||||
}
|
||||
}
|
||||
|
||||
// QueueStats implements StatsReporter. It is safe to call from any goroutine.
|
||||
func (s *Serial) QueueStats() QueueStats {
|
||||
now := s.now()
|
||||
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
bands := make(map[string]BandStats, len(shared.PriorityBands))
|
||||
for _, band := range shared.PriorityBands {
|
||||
bs := BandStats{
|
||||
Depth: s.stats.depth[band],
|
||||
Dispatched: s.stats.dispatched[band],
|
||||
}
|
||||
if t, ok := s.stats.oldest[band]; ok {
|
||||
if wait := now.Sub(t); wait > 0 {
|
||||
bs.OldestWait = wait
|
||||
}
|
||||
}
|
||||
bands[band] = bs
|
||||
}
|
||||
return QueueStats{
|
||||
Bands: bands,
|
||||
AgingReorders: s.stats.agingReorders,
|
||||
AffinityReorders: s.stats.affinityReorders,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,13 +1,16 @@
|
||||
package scheduler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/mostlygeek/llama-swap/internal/config"
|
||||
"github.com/mostlygeek/llama-swap/internal/logmon"
|
||||
"github.com/mostlygeek/llama-swap/internal/process"
|
||||
"github.com/mostlygeek/llama-swap/internal/shared"
|
||||
)
|
||||
|
||||
// Serial methods all run on the router's single run-loop goroutine, so these
|
||||
@@ -15,8 +18,36 @@ import (
|
||||
// req/reqCh helpers from fifo_test.go. A load completes via OnSwapDone and a
|
||||
// served request finishes via OnServeDone — the events the run loop delivers.
|
||||
|
||||
// newSerial builds a Serial with the production defaults: aging at one point
|
||||
// per minute, swap affinity at +10.
|
||||
func newSerial(eff Effects) *Serial {
|
||||
return NewSerial("test", logmon.NewWriter(io.Discard), eff)
|
||||
return NewSerial("test", logmon.NewWriter(io.Discard), config.SerialConfig{}, eff)
|
||||
}
|
||||
|
||||
// newSerialCfg builds a Serial with explicit scoring settings. Passing 0 for
|
||||
// either term disables it.
|
||||
func newSerialCfg(eff Effects, agingDivisor, swapAffinityBonus int) *Serial {
|
||||
return NewSerial("test", logmon.NewWriter(io.Discard), config.SerialConfig{
|
||||
AgingDivisor: &agingDivisor,
|
||||
SwapAffinityBonus: &swapAffinityBonus,
|
||||
}, eff)
|
||||
}
|
||||
|
||||
// fakeClock replaces a Serial's clock with one the test advances by hand, so
|
||||
// aging is exercised without sleeping.
|
||||
type fakeClock struct{ t time.Time }
|
||||
|
||||
func newFakeClock(s *Serial) *fakeClock {
|
||||
c := &fakeClock{t: time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)}
|
||||
s.now = func() time.Time { return c.t }
|
||||
return c
|
||||
}
|
||||
|
||||
func (c *fakeClock) advance(d time.Duration) { c.t = c.t.Add(d) }
|
||||
|
||||
// reqP is a HandlerReq carrying an explicit caller priority.
|
||||
func reqP(model string, priority int) HandlerReq {
|
||||
return HandlerReq{Model: model, Priority: priority}
|
||||
}
|
||||
|
||||
// lastStart returns the most recent StartSwap record.
|
||||
@@ -181,15 +212,18 @@ func TestSerial_SameModelConsecutive_NoReload(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestSerial_StrictArrivalOrder is the core guarantee: qwen36, qwen35, sdxl,
|
||||
// qwen36 execute in EXACTLY that order with evictions between each model switch,
|
||||
// including reloading qwen36 at the end even though it ran first.
|
||||
// TestSerial_StrictArrivalOrder covers the scoring function's degenerate case:
|
||||
// with equal priorities and swap affinity disabled, qwen36, qwen35, sdxl,
|
||||
// qwen36 execute in EXACTLY that order with evictions between each model
|
||||
// switch, including reloading qwen36 at the end even though it ran first.
|
||||
// Aging cannot reorder them either — they all arrive at the same instant, so
|
||||
// they age at the same rate.
|
||||
func TestSerial_StrictArrivalOrder(t *testing.T) {
|
||||
eff := newFakeEffects()
|
||||
for _, m := range []string{"qwen36", "qwen35", "sdxl"} {
|
||||
eff.states[m] = process.StateStopped
|
||||
}
|
||||
s := newSerial(eff)
|
||||
s := newSerialCfg(eff, config.DefaultAgingDivisor, 0)
|
||||
|
||||
for _, m := range []string{"qwen36", "qwen35", "sdxl", "qwen36"} {
|
||||
s.OnRequest(req(m))
|
||||
@@ -243,6 +277,312 @@ func sameOrder(a, b []string) bool {
|
||||
return true
|
||||
}
|
||||
|
||||
// stepModel completes the current load+serve for model and returns control to
|
||||
// the scheduler so it dispatches the next queued request.
|
||||
func stepModel(t *testing.T, s *Serial, eff *fakeEffects, model string) {
|
||||
t.Helper()
|
||||
if len(eff.starts) > 0 {
|
||||
if st := eff.starts[len(eff.starts)-1]; st.model == model {
|
||||
for _, e := range st.evict {
|
||||
eff.states[e] = process.StateStopped
|
||||
}
|
||||
eff.states[model] = process.StateReady
|
||||
s.OnSwapDone(SwapDone{ModelID: model})
|
||||
}
|
||||
}
|
||||
s.OnServeDone(ServeDoneEvent{ModelID: model})
|
||||
}
|
||||
|
||||
// TestSerial_HigherPriorityDispatchesFirst is the point of the whole exercise:
|
||||
// a batch job already queued must yield to an interactive request that arrives
|
||||
// later, because dispatch order is by score, not arrival.
|
||||
func TestSerial_HigherPriorityDispatchesFirst(t *testing.T) {
|
||||
eff := newFakeEffects()
|
||||
for _, m := range []string{"running", "batch", "interactive"} {
|
||||
eff.states[m] = process.StateStopped
|
||||
}
|
||||
s := newSerial(eff)
|
||||
|
||||
s.OnRequest(reqP("running", shared.PriorityNormal)) // dispatches immediately
|
||||
s.OnRequest(reqP("batch", shared.PriorityBatch)) // queued first...
|
||||
s.OnRequest(reqP("interactive", shared.PriorityInteractive)) // ...but this jumps it
|
||||
|
||||
stepModel(t, s, eff, "running")
|
||||
|
||||
if got := eff.startsFor("interactive"); got != 1 {
|
||||
t.Fatalf("StartSwap(interactive)=%d want 1 (must overtake the queued batch job)", got)
|
||||
}
|
||||
if got := eff.startsFor("batch"); got != 0 {
|
||||
t.Fatalf("StartSwap(batch)=%d want 0 (still waiting behind interactive)", got)
|
||||
}
|
||||
|
||||
stepModel(t, s, eff, "interactive")
|
||||
if got := eff.startsFor("batch"); got != 1 {
|
||||
t.Fatalf("StartSwap(batch)=%d want 1 once nothing outranks it", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSerial_TierOffsetBreaksTieWithinBand verifies the composition rule the
|
||||
// design depends on: a small per-caller offset orders requests inside a band
|
||||
// and never promotes one across a band. A "max member" batch job at -98 beats
|
||||
// other batch work but still loses to every normal request at 0.
|
||||
func TestSerial_TierOffsetBreaksTieWithinBand(t *testing.T) {
|
||||
eff := newFakeEffects()
|
||||
for _, m := range []string{"running", "free", "max", "normal"} {
|
||||
eff.states[m] = process.StateStopped
|
||||
}
|
||||
s := newSerialCfg(eff, 0, 0) // isolate the priority term
|
||||
|
||||
s.OnRequest(reqP("running", 0))
|
||||
s.OnRequest(reqP("free", shared.PriorityBatch)) // -100
|
||||
s.OnRequest(reqP("max", shared.PriorityBatch+2)) // -98, max tier
|
||||
s.OnRequest(reqP("normal", shared.PriorityNormal)) // 0
|
||||
|
||||
stepModel(t, s, eff, "running")
|
||||
if got := eff.startsFor("normal"); got != 1 {
|
||||
t.Fatalf("StartSwap(normal)=%d want 1 (a tier bonus must not cross a band)", got)
|
||||
}
|
||||
|
||||
stepModel(t, s, eff, "normal")
|
||||
if got := eff.startsFor("max"); got != 1 {
|
||||
t.Fatalf("StartSwap(max)=%d want 1 (max tier outranks free within the batch band)", got)
|
||||
}
|
||||
|
||||
stepModel(t, s, eff, "max")
|
||||
if got := eff.startsFor("free"); got != 1 {
|
||||
t.Fatalf("StartSwap(free)=%d want 1 (last)", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSerial_AgingPreventsStarvation verifies the one term allowed to cross
|
||||
// bands: a batch job that has waited long enough eventually beats an
|
||||
// interactive request that arrived just now.
|
||||
func TestSerial_AgingPreventsStarvation(t *testing.T) {
|
||||
eff := newFakeEffects()
|
||||
for _, m := range []string{"running", "batch", "interactive"} {
|
||||
eff.states[m] = process.StateStopped
|
||||
}
|
||||
s := newSerialCfg(eff, 60, 0) // one point per minute, no affinity
|
||||
clock := newFakeClock(s)
|
||||
|
||||
s.OnRequest(reqP("running", shared.PriorityNormal))
|
||||
s.OnRequest(reqP("batch", shared.PriorityBatch)) // -100
|
||||
|
||||
// The batch job waits long enough to gain 201 points: -100 + 201 = 101,
|
||||
// just past a fresh interactive request at +100.
|
||||
clock.advance(201 * time.Minute)
|
||||
s.OnRequest(reqP("interactive", shared.PriorityInteractive))
|
||||
|
||||
stepModel(t, s, eff, "running")
|
||||
if got := eff.startsFor("batch"); got != 1 {
|
||||
t.Fatalf("StartSwap(batch)=%d want 1 (aging must eventually beat a fresh interactive request)", got)
|
||||
}
|
||||
if got := eff.startsFor("interactive"); got != 0 {
|
||||
t.Fatalf("StartSwap(interactive)=%d want 0 (outranked by the aged batch job)", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSerial_AgingCannotCrossBandTooEarly is the other half of aging: a batch
|
||||
// job that has waited only a little still loses to interactive traffic.
|
||||
func TestSerial_AgingCannotCrossBandTooEarly(t *testing.T) {
|
||||
eff := newFakeEffects()
|
||||
for _, m := range []string{"running", "batch", "interactive"} {
|
||||
eff.states[m] = process.StateStopped
|
||||
}
|
||||
s := newSerialCfg(eff, 60, 0)
|
||||
clock := newFakeClock(s)
|
||||
|
||||
s.OnRequest(reqP("running", shared.PriorityNormal))
|
||||
s.OnRequest(reqP("batch", shared.PriorityBatch))
|
||||
clock.advance(30 * time.Minute) // -100 + 30 = -70
|
||||
s.OnRequest(reqP("interactive", shared.PriorityInteractive))
|
||||
|
||||
stepModel(t, s, eff, "running")
|
||||
if got := eff.startsFor("interactive"); got != 1 {
|
||||
t.Fatalf("StartSwap(interactive)=%d want 1 (30 minutes is not enough aging)", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSerial_SwapAffinity_PrefersResidentModel verifies the bounded bonus keeps
|
||||
// the model cache from thrashing: with equal priorities, a queued request for
|
||||
// the already-loaded model runs before earlier requests that would each force a
|
||||
// cold load.
|
||||
func TestSerial_SwapAffinity_PrefersResidentModel(t *testing.T) {
|
||||
eff := newFakeEffects()
|
||||
for _, m := range []string{"qwen36", "qwen35", "sdxl"} {
|
||||
eff.states[m] = process.StateStopped
|
||||
}
|
||||
s := newSerial(eff) // affinity +10 by default
|
||||
|
||||
for _, m := range []string{"qwen36", "qwen35", "sdxl", "qwen36"} {
|
||||
s.OnRequest(req(m))
|
||||
}
|
||||
|
||||
stepModel(t, s, eff, "qwen36") // qwen36 now resident and ready
|
||||
if got := eff.served("qwen36"); got != 2 {
|
||||
t.Fatalf("served(qwen36)=%d want 2 (the trailing qwen36 should skip the swap queue)", got)
|
||||
}
|
||||
if len(eff.starts) != 1 {
|
||||
t.Fatalf("starts=%+v want only the initial qwen36 load (no reload)", eff.starts)
|
||||
}
|
||||
|
||||
stepModel(t, s, eff, "qwen36")
|
||||
if got := eff.startsFor("qwen35"); got != 1 {
|
||||
t.Fatalf("StartSwap(qwen35)=%d want 1 once no resident request remains", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSerial_SwapAffinity_CannotCrossBand verifies the bonus is bounded: an
|
||||
// interactive request that needs a cold load still beats a batch job for the
|
||||
// already-resident model.
|
||||
func TestSerial_SwapAffinity_CannotCrossBand(t *testing.T) {
|
||||
eff := newFakeEffects()
|
||||
eff.states["resident"] = process.StateReady
|
||||
eff.states["cold"] = process.StateStopped
|
||||
s := newSerialCfg(eff, 0, 99) // the largest bonus config permits
|
||||
|
||||
s.OnRequest(reqP("resident", shared.PriorityNormal)) // dispatches immediately
|
||||
s.OnRequest(reqP("resident", shared.PriorityBatch)) // queued, would get +99
|
||||
s.OnRequest(reqP("cold", shared.PriorityInteractive))
|
||||
|
||||
s.OnServeDone(ServeDoneEvent{ModelID: "resident"})
|
||||
if got := eff.startsFor("cold"); got != 1 {
|
||||
t.Fatalf("StartSwap(cold)=%d want 1 (-100+99 must still lose to +100)", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSerial_DispatchSetsMetadata verifies the dispatch decision is recorded on
|
||||
// the request context, which is what puts it in the activity log.
|
||||
func TestSerial_DispatchSetsMetadata(t *testing.T) {
|
||||
eff := newFakeEffects()
|
||||
eff.states["hold"] = process.StateReady
|
||||
eff.states["a"] = process.StateStopped
|
||||
s := newSerialCfg(eff, 60, 0)
|
||||
clock := newFakeClock(s)
|
||||
|
||||
s.OnRequest(req("hold")) // serves immediately, occupying the slot
|
||||
|
||||
ctx := shared.SetContext(context.Background(), shared.ReqContextData{ModelID: "a", Metadata: make(map[string]string)})
|
||||
s.OnRequest(HandlerReq{Model: "a", Priority: shared.PriorityBatch, Ctx: ctx})
|
||||
|
||||
clock.advance(2 * time.Minute) // -100 + 2 aging points
|
||||
s.OnServeDone(ServeDoneEvent{ModelID: "hold"})
|
||||
|
||||
// a is dispatched (and annotated) here; finish its load so it is granted.
|
||||
eff.states["hold"] = process.StateStopped
|
||||
eff.states["a"] = process.StateReady
|
||||
s.OnSwapDone(SwapDone{ModelID: "a"})
|
||||
|
||||
data, ok := shared.ReadContext(eff.lastServeReq.Ctx)
|
||||
if !ok {
|
||||
t.Fatal("context data missing from granted request")
|
||||
}
|
||||
want := map[string]string{
|
||||
"serial_priority": "-100",
|
||||
"serial_band": shared.BandBatch,
|
||||
"serial_score": "-98",
|
||||
"serial_queue_wait_ms": "120000",
|
||||
}
|
||||
for k, v := range want {
|
||||
if got := data.Metadata[k]; got != v {
|
||||
t.Errorf("%s = %q, want %q", k, got, v)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestSerial_QueueStats checks the numbers the /metrics endpoint reports:
|
||||
// per-band depth and wait while queued, dispatch counts after the fact, and the
|
||||
// reorder counters that say whether the scoring terms changed anything.
|
||||
func TestSerial_QueueStats(t *testing.T) {
|
||||
eff := newFakeEffects()
|
||||
for _, m := range []string{"running", "batch", "interactive"} {
|
||||
eff.states[m] = process.StateStopped
|
||||
}
|
||||
s := newSerialCfg(eff, 60, 0)
|
||||
clock := newFakeClock(s)
|
||||
|
||||
s.OnRequest(reqP("running", shared.PriorityNormal)) // dispatched, not queued
|
||||
s.OnRequest(reqP("batch", shared.PriorityBatch))
|
||||
clock.advance(5 * time.Minute)
|
||||
s.OnRequest(reqP("interactive", shared.PriorityInteractive))
|
||||
|
||||
stats := s.QueueStats()
|
||||
if got := stats.Bands[shared.BandBatch].Depth; got != 1 {
|
||||
t.Errorf("batch depth=%d want 1", got)
|
||||
}
|
||||
if got := stats.Bands[shared.BandBatch].OldestWait; got != 5*time.Minute {
|
||||
t.Errorf("batch oldest wait=%s want 5m", got)
|
||||
}
|
||||
if got := stats.Bands[shared.BandInteractive].Depth; got != 1 {
|
||||
t.Errorf("interactive depth=%d want 1", got)
|
||||
}
|
||||
if got := stats.Bands[shared.BandNormal].Depth; got != 0 {
|
||||
t.Errorf("normal depth=%d want 0 (it dispatched immediately)", got)
|
||||
}
|
||||
if got := stats.Bands[shared.BandNormal].Dispatched; got != 1 {
|
||||
t.Errorf("normal dispatched=%d want 1", got)
|
||||
}
|
||||
|
||||
// Dispatching interactive over the older batch job is priority doing its
|
||||
// job, not aging: with only 5 minutes of aging the winner is unchanged.
|
||||
stepModel(t, s, eff, "running")
|
||||
stats = s.QueueStats()
|
||||
if got := stats.Bands[shared.BandInteractive].Dispatched; got != 1 {
|
||||
t.Errorf("interactive dispatched=%d want 1", got)
|
||||
}
|
||||
if stats.AgingReorders != 0 {
|
||||
t.Errorf("aging reorders=%d want 0", stats.AgingReorders)
|
||||
}
|
||||
if stats.AffinityReorders != 0 {
|
||||
t.Errorf("affinity reorders=%d want 0 (affinity disabled)", stats.AffinityReorders)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSerial_QueueStats_CountsReorders verifies each scoring term is credited
|
||||
// when it actually changes the dispatch decision.
|
||||
func TestSerial_QueueStats_CountsReorders(t *testing.T) {
|
||||
t.Run("aging", func(t *testing.T) {
|
||||
eff := newFakeEffects()
|
||||
for _, m := range []string{"running", "old", "new"} {
|
||||
eff.states[m] = process.StateStopped
|
||||
}
|
||||
s := newSerialCfg(eff, 60, 0)
|
||||
clock := newFakeClock(s)
|
||||
|
||||
s.OnRequest(reqP("running", 0))
|
||||
s.OnRequest(reqP("old", -10))
|
||||
clock.advance(30 * time.Minute) // old: -10 + 30 = 20
|
||||
s.OnRequest(reqP("new", 0)) // new: 0
|
||||
|
||||
stepModel(t, s, eff, "running")
|
||||
if got := eff.startsFor("old"); got != 1 {
|
||||
t.Fatalf("StartSwap(old)=%d want 1 (aging promoted it)", got)
|
||||
}
|
||||
if got := s.QueueStats().AgingReorders; got != 1 {
|
||||
t.Errorf("aging reorders=%d want 1", got)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("swap affinity", func(t *testing.T) {
|
||||
eff := newFakeEffects()
|
||||
eff.states["resident"] = process.StateReady
|
||||
eff.states["cold"] = process.StateStopped
|
||||
s := newSerialCfg(eff, 0, 10)
|
||||
|
||||
s.OnRequest(req("resident")) // dispatches immediately
|
||||
s.OnRequest(reqP("cold", 5)) // higher priority, but needs a swap
|
||||
s.OnRequest(reqP("resident", 0)) // +10 affinity beats it
|
||||
s.OnServeDone(ServeDoneEvent{ModelID: "resident"})
|
||||
|
||||
if got := eff.served("resident"); got != 2 {
|
||||
t.Fatalf("served(resident)=%d want 2 (affinity outranked the cold request)", got)
|
||||
}
|
||||
if got := s.QueueStats().AffinityReorders; got != 1 {
|
||||
t.Errorf("affinity reorders=%d want 1", got)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestSerial_SwapError_FailsCallerAndAdvances(t *testing.T) {
|
||||
eff := newFakeEffects()
|
||||
eff.states["a"] = process.StateStopped
|
||||
|
||||
Reference in New Issue
Block a user