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
2026-08-07 02:12:08 -04:00
parent 2f0bb2556e
commit 0358fe321e
22 changed files with 1515 additions and 67 deletions
@@ -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