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
@@ -26,6 +26,9 @@ type ReqContextData struct {
 	ModelID          string
 	Streaming        bool
 	SendLoadingState bool
+	// Priority is the caller's scheduling priority, read from PriorityHeader.
+	// Higher is more urgent; 0 (PriorityNormal) is the default. See priority.go.
+	Priority int
 	// Metadata is a request-scoped key/value bag that handlers may mutate
 	// while processing. The metrics middleware copies it into ActivityLogEntry.
 	Metadata map[string]string
@@ -118,6 +121,7 @@ func FetchContext(r *http.Request, cfg config.Config) (ReqContextData, error) {
 		if mc, ok := cfg.Models[realName]; ok {
 			data.SendLoadingState = mc.SendLoadingState != nil && *mc.SendLoadingState
 		}
+		data.Priority = RequestPriority(r)
 		*r = *r.WithContext(SetContext(r.Context(), data))
 		return data, nil
 	}
@@ -137,6 +141,7 @@ func extractUpstreamContext(r *http.Request, cfg config.Config) (ReqContextData,
 		ApiKey:           ExtractAPIKey(r),
 		Streaming:        r.URL.Query().Get("stream") == "true",
 		SendLoadingState: sendLoadingState(cfg, realName),
+		Priority:         RequestPriority(r),
 		Metadata:         make(map[string]string),
 	}, true
 }
@@ -0,0 +1,89 @@
+package shared
+
+import (
+	"net/http"
+	"strconv"
+	"strings"
+)
+
+// PriorityHeader carries the caller's scheduling priority. Priority is a
+// property of the caller's intent, not of the model — the same model serves
+// both an interactive request and a batch job — so it travels as a header
+// rather than as per-model configuration, and leaves the OpenAI-compatible
+// request body untouched.
+const PriorityHeader = "X-LlamaSwap-Priority"
+
+// Priority band anchors. The number is the interface: callers may send any
+// signed integer, and these names are conveniences that resolve to a value.
+//
+// Bands are spaced 100 apart so a consumer can add a small per-caller offset
+// (a subscription tier, say) on top of a band without ever promoting a request
+// across one: a "max member" batch job at -98 still loses to every normal
+// request at 0.
+const (
+	PriorityInteractive = 100 // a human is waiting
+	PriorityNormal      = 0   // the default
+	PriorityBatch       = -100
+)
+
+// bandHalfWidth is how far a priority may sit from a band anchor and still be
+// reported as that band. It is half the 100-point band spacing, so every
+// integer belongs to exactly one band.
+const bandHalfWidth = 50
+
+// Band names used for metrics labels, ordered most to least urgent.
+const (
+	BandInteractive = "interactive"
+	BandNormal      = "normal"
+	BandBatch       = "batch"
+)
+
+// PriorityBands is every band name, in descending order of urgency. Metrics
+// emit a series per band regardless of whether any request currently occupies
+// it, so the label set stays stable.
+var PriorityBands = []string{BandInteractive, BandNormal, BandBatch}
+
+var priorityAliases = map[string]int{
+	"interactive": PriorityInteractive,
+	"normal":      PriorityNormal,
+	"batch":       PriorityBatch,
+}
+
+// ParsePriority resolves a PriorityHeader value to a signed integer. A numeric
+// value is used as-is; a recognised alias resolves to its anchor. Anything
+// absent or unparseable is PriorityNormal, so a malformed header degrades to
+// the default rather than failing the request.
+//
+// Values are deliberately not clamped: the caller composes band, tier offset
+// and anything else it wants into one number, and llama-swap honours it.
+func ParsePriority(value string) int {
+	value = strings.TrimSpace(value)
+	if value == "" {
+		return PriorityNormal
+	}
+	if n, err := strconv.Atoi(value); err == nil {
+		return n
+	}
+	if n, ok := priorityAliases[strings.ToLower(value)]; ok {
+		return n
+	}
+	return PriorityNormal
+}
+
+// RequestPriority reads PriorityHeader off r and resolves it. See ParsePriority.
+func RequestPriority(r *http.Request) int {
+	return ParsePriority(r.Header.Get(PriorityHeader))
+}
+
+// PriorityBand buckets a priority into the band it belongs to, for metrics
+// labels. Values beyond the anchors saturate: +1000 is still "interactive".
+func PriorityBand(priority int) string {
+	switch {
+	case priority >= PriorityInteractive-bandHalfWidth:
+		return BandInteractive
+	case priority <= PriorityBatch+bandHalfWidth:
+		return BandBatch
+	default:
+		return BandNormal
+	}
+}
@@ -0,0 +1,118 @@
+package shared
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"github.com/mostlygeek/llama-swap/internal/config"
+)
+
+func TestShared_ParsePriority(t *testing.T) {
+	tests := []struct {
+		name  string
+		value string
+		want  int
+	}{
+		{"absent", "", PriorityNormal},
+		{"zero", "0", 0},
+		{"positive", "100", 100},
+		{"negative", "-100", -100},
+		{"explicit plus", "+42", 42},
+		{"surrounding space", "  100  ", 100},
+		{"large value is not clamped", "100000", 100000},
+		{"alias interactive", "interactive", PriorityInteractive},
+		{"alias normal", "normal", PriorityNormal},
+		{"alias batch", "batch", PriorityBatch},
+		{"alias is case insensitive", "Batch", PriorityBatch},
+		{"tier offset on a band", "-98", -98},
+		{"unknown alias falls back to normal", "urgent", PriorityNormal},
+		{"garbage falls back to normal", "!!", PriorityNormal},
+		{"float falls back to normal", "1.5", PriorityNormal},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			if got := ParsePriority(tc.value); got != tc.want {
+				t.Errorf("ParsePriority(%q)=%d want %d", tc.value, got, tc.want)
+			}
+		})
+	}
+}
+
+func TestShared_RequestPriority(t *testing.T) {
+	r := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
+	if got := RequestPriority(r); got != PriorityNormal {
+		t.Errorf("no header: got %d want %d", got, PriorityNormal)
+	}
+
+	r.Header.Set(PriorityHeader, "batch")
+	if got := RequestPriority(r); got != PriorityBatch {
+		t.Errorf("batch header: got %d want %d", got, PriorityBatch)
+	}
+}
+
+// TestShared_PriorityBand pins the bucketing the metrics labels rely on. The
+// important cases are the tier offsets: a max member's batch job at -98 must
+// still report as batch, not normal.
+func TestShared_PriorityBand(t *testing.T) {
+	tests := []struct {
+		priority int
+		want     string
+	}{
+		{PriorityInteractive, BandInteractive},
+		{PriorityInteractive + 2, BandInteractive}, // max tier, interactive
+		{1000, BandInteractive},                    // saturates
+		{50, BandInteractive},                      // lower edge
+		{49, BandNormal},
+		{PriorityNormal, BandNormal},
+		{2, BandNormal}, // max tier, normal
+		{-49, BandNormal},
+		{-50, BandBatch}, // upper edge
+		{PriorityBatch + 2, BandBatch},
+		{PriorityBatch, BandBatch},
+		{-1000, BandBatch}, // saturates
+	}
+	for _, tc := range tests {
+		if got := PriorityBand(tc.priority); got != tc.want {
+			t.Errorf("PriorityBand(%d)=%q want %q", tc.priority, got, tc.want)
+		}
+	}
+}
+
+// TestShared_FetchContext_Priority verifies the header reaches the request
+// context, which is what carries it to the scheduler. Both entry points into
+// FetchContext are covered: the normal body-parsed path and /upstream/<model>.
+func TestShared_FetchContext_Priority(t *testing.T) {
+	cfg := config.Config{Models: map[string]config.ModelConfig{"m1": {}}}
+
+	cases := []struct {
+		name   string
+		path   string
+		header string
+		want   int
+	}{
+		{"body path, no header", "/v1/chat/completions", "", PriorityNormal},
+		{"body path, alias", "/v1/chat/completions", "interactive", PriorityInteractive},
+		{"body path, number", "/v1/chat/completions", "-98", -98},
+		{"upstream path", "/upstream/m1/v1/chat/completions", "batch", PriorityBatch},
+	}
+
+	for _, c := range cases {
+		t.Run(c.name, func(t *testing.T) {
+			r := httptest.NewRequest(http.MethodPost, c.path, strings.NewReader(`{"model":"m1"}`))
+			r.Header.Set("Content-Type", "application/json")
+			if c.header != "" {
+				r.Header.Set(PriorityHeader, c.header)
+			}
+
+			data, err := FetchContext(r, cfg)
+			if err != nil {
+				t.Fatalf("FetchContext: %v", err)
+			}
+			if data.Priority != c.want {
+				t.Errorf("Priority=%d want %d", data.Priority, c.want)
+			}
+		})
+	}
+}