diff --git a/README.md b/README.md index 09d4256e..9eeed097 100644 --- a/README.md +++ b/README.md @@ -213,6 +213,9 @@ Almost all configuration settings are optional and can be added one step at a ti - `matrix` to run concurrent models with a custom swap logic DSL - `hooks` to run things on startup - `macros` reusable snippets + - `routing.scheduler` to control how queued requests are ordered — see + [request priority](docs/request-priority.md) for the `X-LlamaSwap-Priority` + header that lets batch jobs yield to interactive requests - Model customization - `ttl` to automatically unload models - `aliases` to use familiar model names (e.g., "gpt-4o-mini") diff --git a/config-schema.json b/config-schema.json index d0dc2da0..1865a67e 100644 --- a/config-schema.json +++ b/config-schema.json @@ -605,17 +605,37 @@ "fifo" ], "default": "serial", - "description": "Scheduler to use. 'serial' (default on this fork): strict one-model-at-a-time, requests run in exact arrival order, switching models evicts every other model first. 'fifo': throughput-oriented, batches same-model requests and allows parallel/co-resident models." + "description": "Scheduler to use. 'serial' (default on this fork): strict one-model-at-a-time, highest-priority request runs next, switching models evicts every other model first. 'fifo': throughput-oriented, batches same-model requests and allows parallel/co-resident models. Both honour the per-request X-LlamaSwap-Priority header." }, "settings": { "type": "object", "properties": { + "serial": { + "type": "object", + "description": "Tunes how the serial scheduler scores waiting requests: score = X-LlamaSwap-Priority + swap affinity + aging. Ignored unless use is 'serial'.", + "properties": { + "agingDivisor": { + "type": "integer", + "minimum": 0, + "default": 60, + "description": "Seconds a request must wait to gain one priority point. 0 disables aging. Unbounded on purpose: aging is the only term allowed to promote a request across a priority band, which is what prevents starvation." + }, + "swapAffinityBonus": { + "type": "integer", + "minimum": 0, + "maximum": 99, + "default": 10, + "description": "Bonus for a request that can be served without swapping models, so strict priority does not force a cold load on every dispatch. Capped below the 100-point band spacing so it can only break near-ties. 0 disables it, giving strict priority-then-arrival order." + } + }, + "additionalProperties": false + }, "fifo": { "type": "object", "properties": { "priority": { "type": "object", - "description": "Per-model priority. Keys are model IDs, values are integers (default 0). Higher values are serviced first.", + "description": "Per-model priority. Keys are model IDs, values are integers (default 0). Higher values are serviced first. Added to the caller's X-LlamaSwap-Priority.", "additionalProperties": { "type": "integer" } diff --git a/config.example.yaml b/config.example.yaml index f495cadd..15410103 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -559,23 +559,62 @@ routing: # scheduler: how queued requests are ordered and run. # - optional, default on this fork: "serial" # - valid values: - # - "serial": strict one-model-at-a-time. Requests run in exact arrival - # order; only one request runs at a time; switching to a different model - # evicts every other running model first so a single model occupies memory - # at a time. This ignores group/matrix co-residency entirely. The "fifo" - # settings below (priority) do not apply. + # - "serial": strict one-model-at-a-time. Only one request runs at a time; + # switching to a different model evicts every other running model first so + # a single model occupies memory at a time. This ignores group/matrix + # co-residency entirely. The "fifo" settings below do not apply. # - "fifo": throughput-oriented. Same-model requests are batched to reduce # swaps and a model serves up to its concurrencyLimit in parallel; models # in non-exclusive groups can run concurrently. Requests may be reordered. + # + # Callers set per-request priority with the X-LlamaSwap-Priority header + # (a signed integer; "interactive" = 100, "normal" = 0, "batch" = -100; + # absent or unparseable means 0). Both schedulers honour it. Scheduling is + # non-preemptive: priority applies when a job finishes, so an interactive + # request can still wait up to one full job duration. scheduler: use: serial settings: + # serial settings only apply when use: serial + # + # At each dispatch the serial scheduler runs the highest scoring waiting + # request, where: + # + # score = X-LlamaSwap-Priority + swap affinity + aging + # + # Priority bands sit 100 apart, which leaves room for a caller to add a + # small per-user offset (a subscription tier, say) without crossing a + # band: a "max member" batch job at -98 beats other batch work but still + # loses to every normal request at 0. + serial: + # agingDivisor: seconds a request must wait to gain one priority point + # - optional, default: 60 (one point per minute) + # - 0 disables aging + # - unbounded on purpose: aging is the only term allowed to promote a + # request across a band, which is what stops low-priority work from + # starving. At 60, a batch job at -100 overtakes normal traffic after + # ~100 minutes of waiting. + agingDivisor: 60 + + # swapAffinityBonus: bonus for a request that needs no model swap + # - optional, default: 10 + # - must be 0..99, so it can only break near-ties and can never promote + # a request into the next band + # - 0 disables it, giving strict priority-then-arrival order + # - cold loads are expensive, so this keeps a run of same-model requests + # together instead of forcing a reload on every dispatch. Note this + # means equal-priority requests are NOT served in strict arrival + # order; set it to 0 if you need that. + swapAffinityBonus: 10 + # fifo settings only apply when use: fifo fifo: # priority: a dictionary of model ID -> priority # - optional, default: empty dictionary # - models default to priority 0 # - higher priority requests are serviced first in the queue + # - added to the caller's X-LlamaSwap-Priority, so a model priority of + # 10 and a header of "batch" (-100) queue at -90 priority: A: 10 B: 5 diff --git a/docs/request-priority.md b/docs/request-priority.md new file mode 100644 index 00000000..c7c47402 --- /dev/null +++ b/docs/request-priority.md @@ -0,0 +1,151 @@ +# Request priority + +A GPU is a size-1 resource: one job at a time. Without priorities, a single long +job monopolises the box for its whole duration and every interactive request +queues behind it. A 16-minute render blocking a chat completion is +uncomfortable; a 12-hour batch render blocking everything for a day is not +workable. + +Callers declare intent with a header, and the scheduler dispatches by score. + +## The header + +``` +X-LlamaSwap-Priority: 0 # normal — the default +X-LlamaSwap-Priority: 100 # interactive, a human is waiting +X-LlamaSwap-Priority: -100 # batch, nobody is watching +``` + +Priority is a **signed integer**; higher is more urgent. An absent or +unparseable header means `0`. Values are not clamped — llama-swap trusts the +caller. + +Named aliases resolve to numbers for convenience, but **the number is the +interface**: callers may send any value. + +| alias | value | +| ------------- | -----: | +| `interactive` | `100` | +| `normal` | `0` | +| `batch` | `-100` | + +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 — which is why it +travels as a header rather than as per-model config. It also leaves the +OpenAI-compatible request body untouched. + +```bash +curl http://localhost:8080/v1/chat/completions \ + -H 'Content-Type: application/json' \ + -H 'X-LlamaSwap-Priority: interactive' \ + -d '{"model":"qwen3","messages":[{"role":"user","content":"hi"}]}' +``` + +## Bands and offsets + +The three anchors sit 100 apart. That spacing is deliberate: it leaves room for +a caller to add a small per-user offset without ever crossing a band. + +| tier | offset | +| ---- | -----: | +| free | `+0` | +| pro | `+1` | +| max | `+2` | + +A max member's batch job is `-100 + 2 = -98`: ahead of pro's `-99` and free's +`-100`, and still **below** any normal request at `0`. A tier bonus breaks ties +*inside* a band and can never promote across one. + +> **Invariant:** band spacing (100) must stay far larger than the largest offset +> a caller adds plus `swapAffinityBonus`. Otherwise a max member's batch work +> could outrank a free member's interactive request. + +## How the serial scheduler dispatches + +Scheduling is **non-preemptive**. A generation cannot be interrupted mid-sample +without discarding the work, so priority applies at *dispatch* time: when a job +finishes, the best-scoring waiting request goes next. + +``` +score = X-LlamaSwap-Priority + swap affinity + aging +``` + +| term | range | crosses bands? | +| --------------------- | ------- | ----------------------- | +| **request priority** | any | — it *is* the band | +| **swap affinity** | `0..99` | **never** | +| **aging** | `0..∞` | **yes, deliberately** | + +Equal scores keep arrival order. + +**Swap affinity** is a bounded bonus for a request that can run without a model +swap. Cold loads are expensive, so this keeps a run of same-model requests +together instead of forcing a reload on every dispatch. Because it is bounded +well below the band spacing, it can only break near-ties — it can never promote +batch work above interactive. + +**Aging** is `waited_seconds / agingDivisor`, and is unbounded on purpose: a +batch job that has waited long enough *should* eventually beat normal traffic, +or it starves. At the default divisor of 60 (one point per minute), a batch job +at `-100` overtakes normal traffic after ~100 minutes of waiting. + +That asymmetry is the design: **bounded offsets express policy, unbounded aging +prevents starvation.** + +### What priority does not fix + +Because dispatch is non-preemptive, an interactive request can still wait **up +to one full job duration**. The complementary lever is client-side: a batch +producer should submit **one unit at a time** and re-queue, so gaps are frequent. +For a long render that means one shot per submission, not one 12-hour chain. + +Priority is also not fair-share. If one caller submits 200 batch jobs they all +sit at the same priority, and llama-swap will work through them in order. Per-user +fairness belongs in the client that submits the work. + +## Configuration + +```yaml +routing: + scheduler: + use: serial + settings: + serial: + # Seconds of waiting per priority point. 0 disables aging. + agingDivisor: 60 + # Bonus for a request needing no model swap. 0..99; 0 disables it. + swapAffinityBonus: 10 +``` + +Setting `swapAffinityBonus: 0` gives strict priority-then-arrival order — with +uniform priorities that is exact FIFO. + +The `fifo` scheduler honours the header too, adding it to the model's configured +priority: a model at `10` with a `batch` header queues at `-90`. + +## Observability + +`GET /metrics` exports the queue in Prometheus text format, labelled by band +(`interactive`, `normal`, `batch`): + +| metric | type | meaning | +| --------------------------------------------- | ------- | ---------------------------------------------------- | +| `llamaswap_scheduler_queue_depth` | gauge | requests waiting | +| `llamaswap_scheduler_queue_oldest_wait_seconds`| gauge | how long the longest-waiting request has waited | +| `llamaswap_scheduler_dispatched_total` | counter | requests dispatched | +| `llamaswap_scheduler_reorders_total` | counter | dispatches where `aging` or `swap_affinity` changed the pick | + +The reorder counters are what make `agingDivisor` and `swapAffinityBonus` +tunable rather than guesswork: they say how often each term actually changed +which request went next. + +Every dispatched request also records its decision in the activity log: + +| key | meaning | +| ---------------------- | ------------------------------------------ | +| `serial_priority` | the caller's priority | +| `serial_band` | which band it fell in | +| `serial_score` | the score it won with | +| `serial_queue_wait_ms` | how long it waited before dispatch | + +That is what answers "why did my request take 20 minutes" after the fact. diff --git a/internal/config/config.go b/internal/config/config.go index 090e3512..7b100b41 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -167,13 +167,67 @@ type SchedulerConfig struct { } type SchedulerSettings struct { - Fifo FifoConfig `yaml:"fifo"` + Fifo FifoConfig `yaml:"fifo"` + Serial SerialConfig `yaml:"serial"` } type FifoConfig struct { Priority map[string]int `yaml:"priority"` // model ID -> priority, default 0 } +// Serial scheduler scoring defaults. See SerialConfig. +const ( + // DefaultAgingDivisor gives a waiting request one priority point per + // minute, so a batch job at -100 overtakes normal traffic after ~100 + // minutes rather than starving behind it forever. + DefaultAgingDivisor = 60 + + // DefaultSwapAffinityBonus is large enough to break near-ties in favour + // of the already-loaded model (cold loads cost 2-3x) and far too small to + // promote a request across a 100-point priority band. + DefaultSwapAffinityBonus = 10 +) + +// SerialConfig tunes how the serial scheduler scores queued requests. At each +// dispatch it picks the highest scoring waiting request, where +// +// score = request priority + swap affinity + aging +// +// Both terms are pointers so an explicit 0 (disable) is distinguishable from +// "unset" (use the default); read them through the accessors below. +type SerialConfig struct { + // AgingDivisor is how many seconds a request must wait to gain one + // priority point. Aging is unbounded on purpose — it is the only term + // allowed to promote a request across a band, which is what stops a + // low-priority job from starving. 0 disables aging. + AgingDivisor *int `yaml:"agingDivisor"` + + // SwapAffinityBonus is added to a request that can be served without + // swapping models, so strict priority does not force a cold load on every + // dispatch. Bounded well below the 100-point band spacing so it can only + // break near-ties, never promote batch work above interactive. 0 disables + // it, restoring strict priority-then-arrival order. + SwapAffinityBonus *int `yaml:"swapAffinityBonus"` +} + +// GetAgingDivisor returns the configured aging divisor in seconds, or +// DefaultAgingDivisor when unset. +func (c SerialConfig) GetAgingDivisor() int { + if c.AgingDivisor == nil { + return DefaultAgingDivisor + } + return *c.AgingDivisor +} + +// GetSwapAffinityBonus returns the configured swap affinity bonus, or +// DefaultSwapAffinityBonus when unset. +func (c SerialConfig) GetSwapAffinityBonus() int { + if c.SwapAffinityBonus == nil { + return DefaultSwapAffinityBonus + } + return *c.SwapAffinityBonus +} + type RouterConfig struct { Use string `yaml:"use"` // "group" (default) | "matrix" Settings RouterSettings `yaml:"settings"` diff --git a/internal/config/config_test.go b/internal/config/config_test.go index eef80b96..c4b883b7 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -1720,3 +1720,77 @@ routing: require.NoError(t, err) assert.Equal(t, 5, cfg.Routing.Scheduler.Settings.Fifo.Priority["gemma"]) } + +func TestConfig_Routing_SerialDefaults(t *testing.T) { + cfg, err := LoadConfigFromReader(strings.NewReader(twoModels)) + require.NoError(t, err) + + serial := cfg.Routing.Scheduler.Settings.Serial + assert.Nil(t, serial.AgingDivisor, "unset should stay nil so the default applies") + assert.Equal(t, DefaultAgingDivisor, serial.GetAgingDivisor()) + assert.Equal(t, DefaultSwapAffinityBonus, serial.GetSwapAffinityBonus()) +} + +// TestConfig_Routing_SerialExplicitZero verifies an explicit 0 disables a term +// rather than falling back to the default — the reason both fields are pointers. +func TestConfig_Routing_SerialExplicitZero(t *testing.T) { + yaml := twoModels + ` +routing: + scheduler: + use: serial + settings: + serial: + agingDivisor: 0 + swapAffinityBonus: 0 +` + cfg, err := LoadConfigFromReader(strings.NewReader(yaml)) + require.NoError(t, err) + + serial := cfg.Routing.Scheduler.Settings.Serial + assert.Equal(t, 0, serial.GetAgingDivisor()) + assert.Equal(t, 0, serial.GetSwapAffinityBonus()) +} + +func TestConfig_Routing_SerialSettings(t *testing.T) { + yaml := twoModels + ` +routing: + scheduler: + use: serial + settings: + serial: + agingDivisor: 30 + swapAffinityBonus: 5 +` + cfg, err := LoadConfigFromReader(strings.NewReader(yaml)) + require.NoError(t, err) + + serial := cfg.Routing.Scheduler.Settings.Serial + assert.Equal(t, 30, serial.GetAgingDivisor()) + assert.Equal(t, 5, serial.GetSwapAffinityBonus()) +} + +func TestConfig_Routing_SerialInvalidSettings(t *testing.T) { + cases := []struct { + name string + settings string + wantErr string + }{ + {"negative aging divisor", "agingDivisor: -1", "agingDivisor"}, + {"negative affinity bonus", "swapAffinityBonus: -1", "swapAffinityBonus"}, + // 100 would let the bonus promote a request into the next band. + {"affinity bonus reaches a band", "swapAffinityBonus: 100", "swapAffinityBonus"}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + yaml := twoModels + ` +routing: + scheduler: + settings: + serial: + ` + c.settings + "\n" + _, err := LoadConfigFromReader(strings.NewReader(yaml)) + require.Error(t, err) + assert.Contains(t, err.Error(), c.wantErr) + }) + } +} diff --git a/internal/config/load.go b/internal/config/load.go index a56358e3..cbdfe952 100644 --- a/internal/config/load.go +++ b/internal/config/load.go @@ -374,6 +374,16 @@ func LoadConfigFromReader(r io.Reader) (Config, error) { return Config{}, fmt.Errorf("routing.scheduler.settings.fifo.priority references unknown model %q", modelID) } } + serialCfg := config.Routing.Scheduler.Settings.Serial + if d := serialCfg.AgingDivisor; d != nil && *d < 0 { + return Config{}, fmt.Errorf("routing.scheduler.settings.serial.agingDivisor: must be >= 0, got %d (0 disables aging)", *d) + } + // Capped below the 100-point band spacing so the bonus can only break + // near-ties. A bonus that could reach the next band would let a batch job + // on the loaded model outrank an interactive request that needs a swap. + if b := serialCfg.SwapAffinityBonus; b != nil && (*b < 0 || *b >= 100) { + return Config{}, fmt.Errorf("routing.scheduler.settings.serial.swapAffinityBonus: must be between 0 and 99, got %d (priority bands are 100 apart and the bonus must never cross one)", *b) + } // Clean up hooks preload if len(config.Hooks.OnStartup.Preload) > 0 { diff --git a/internal/router/base.go b/internal/router/base.go index 7f7232ae..c88fa973 100644 --- a/internal/router/base.go +++ b/internal/router/base.go @@ -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. diff --git a/internal/router/base_test.go b/internal/router/base_test.go index 0c8d2fab..8bc61aea 100644 --- a/internal/router/base_test.go +++ b/internal/router/base_test.go @@ -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{}) diff --git a/internal/router/router.go b/internal/router/router.go index 561c5fc8..51801323 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -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) } diff --git a/internal/router/scheduler/fifo.go b/internal/router/scheduler/fifo.go index 44cd4d7d..dcdd55bb 100644 --- a/internal/router/scheduler/fifo.go +++ b/internal/router/scheduler/fifo.go @@ -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 } diff --git a/internal/router/scheduler/fifo_test.go b/internal/router/scheduler/fifo_test.go index 34b39749..d94b7c4a 100644 --- a/internal/router/scheduler/fifo_test.go +++ b/internal/router/scheduler/fifo_test.go @@ -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. diff --git a/internal/router/scheduler/scheduler.go b/internal/router/scheduler/scheduler.go index 86439ede..9b7c4584 100644 --- a/internal/router/scheduler/scheduler.go +++ b/internal/router/scheduler/scheduler.go @@ -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 { diff --git a/internal/router/scheduler/serial.go b/internal/router/scheduler/serial.go index 3b49ce8a..e019e0c7 100644 --- a/internal/router/scheduler/serial.go +++ b/internal/router/scheduler/serial.go @@ -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, + } +} diff --git a/internal/router/scheduler/serial_test.go b/internal/router/scheduler/serial_test.go index 60de4e43..e54a0098 100644 --- a/internal/router/scheduler/serial_test.go +++ b/internal/router/scheduler/serial_test.go @@ -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 diff --git a/internal/server/api.go b/internal/server/api.go index b1b4f107..c68a3483 100644 --- a/internal/server/api.go +++ b/internal/server/api.go @@ -293,15 +293,25 @@ func (s *Server) startPreload() { }() } -// handleMetrics serves Prometheus-format performance metrics. Returns 503 when -// performance monitoring is disabled. +// handleMetrics serves Prometheus-format metrics: system/GPU performance plus +// the scheduler's request queue. Returns 503 only when neither source is +// available — performance monitoring disabled and a scheduler that does not +// report queue stats. func (s *Server) handleMetrics(w http.ResponseWriter, r *http.Request) { - if s.perf == nil { + queueStats, hasQueueStats := s.local.SchedulerStats() + if s.perf == nil && !hasQueueStats { w.WriteHeader(http.StatusServiceUnavailable) w.Write([]byte("# performance monitor not available\n")) return } - s.perf.MetricsHandler().ServeHTTP(w, r) + + w.Header().Set("Content-Type", "text/plain; version=0.0.4; charset=utf-8") + if s.perf != nil { + s.perf.MetricsHandler().ServeHTTP(w, r) + } + if hasQueueStats { + writeSchedulerMetrics(w, queueStats) + } } func handleHealth(w http.ResponseWriter, r *http.Request) { diff --git a/internal/server/api_test.go b/internal/server/api_test.go index 715a480f..20df000f 100644 --- a/internal/server/api_test.go +++ b/internal/server/api_test.go @@ -8,10 +8,12 @@ import ( "regexp" "strings" "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/router/scheduler" "github.com/mostlygeek/llama-swap/internal/shared" ) @@ -301,6 +303,8 @@ func TestServer_HandleUpstream_MetricsSkipsGET(t *testing.T) { } func TestServer_HandleMetrics_Unavailable(t *testing.T) { + // No perf monitor and a scheduler that reports no queue stats: nothing to + // serve. s := newTestServer(newStubRouter(nil, ""), newStubRouter(nil, "")) w := httptest.NewRecorder() @@ -310,6 +314,44 @@ func TestServer_HandleMetrics_Unavailable(t *testing.T) { } } +// TestServer_HandleMetrics_SchedulerQueue verifies the scheduler queue is +// exported even when performance monitoring is off, and that every band gets a +// series so the label set stays stable while the queue drains. +func TestServer_HandleMetrics_SchedulerQueue(t *testing.T) { + local := newStubRouter(nil, "") + local.queueStats = &scheduler.QueueStats{ + Bands: map[string]scheduler.BandStats{ + shared.BandInteractive: {Depth: 2, OldestWait: 3 * time.Second, Dispatched: 41}, + shared.BandBatch: {Depth: 1, Dispatched: 7}, + }, + AgingReorders: 5, + AffinityReorders: 9, + } + s := newTestServer(local, newStubRouter(nil, "")) + + w := httptest.NewRecorder() + s.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/metrics", nil)) + if w.Code != http.StatusOK { + t.Fatalf("status = %d, want 200", w.Code) + } + + body := w.Body.String() + for _, want := range []string{ + `llamaswap_scheduler_queue_depth{band="interactive"} 2`, + `llamaswap_scheduler_queue_depth{band="normal"} 0`, + `llamaswap_scheduler_queue_depth{band="batch"} 1`, + `llamaswap_scheduler_queue_oldest_wait_seconds{band="interactive"} 3`, + `llamaswap_scheduler_dispatched_total{band="interactive"} 41`, + `llamaswap_scheduler_dispatched_total{band="batch"} 7`, + `llamaswap_scheduler_reorders_total{cause="aging"} 5`, + `llamaswap_scheduler_reorders_total{cause="swap_affinity"} 9`, + } { + if !strings.Contains(body, want) { + t.Errorf("missing %q in:\n%s", want, body) + } + } +} + func TestServer_Redirects(t *testing.T) { s := newTestServer(newStubRouter(nil, ""), newStubRouter(nil, "")) diff --git a/internal/server/metrics_scheduler.go b/internal/server/metrics_scheduler.go new file mode 100644 index 00000000..dd0fa1e7 --- /dev/null +++ b/internal/server/metrics_scheduler.go @@ -0,0 +1,39 @@ +package server + +import ( + "fmt" + "io" + + "github.com/mostlygeek/llama-swap/internal/router/scheduler" + "github.com/mostlygeek/llama-swap/internal/shared" +) + +// writeSchedulerMetrics emits the request queue in Prometheus text format. +// +// The reorder counters are the ones that matter for tuning: they say how often +// aging or swap affinity actually changed which request went next. Without +// them, choosing agingDivisor and swapAffinityBonus is guesswork. +func writeSchedulerMetrics(w io.Writer, stats scheduler.QueueStats) { + fmt.Fprintf(w, "# HELP llamaswap_scheduler_queue_depth Requests waiting in the scheduler queue, by priority band\n") + fmt.Fprintf(w, "# TYPE llamaswap_scheduler_queue_depth gauge\n") + for _, band := range shared.PriorityBands { + fmt.Fprintf(w, "llamaswap_scheduler_queue_depth{band=%q} %d\n", band, stats.Bands[band].Depth) + } + + fmt.Fprintf(w, "# HELP llamaswap_scheduler_queue_oldest_wait_seconds How long the longest-waiting queued request has waited, by priority band\n") + fmt.Fprintf(w, "# TYPE llamaswap_scheduler_queue_oldest_wait_seconds gauge\n") + for _, band := range shared.PriorityBands { + fmt.Fprintf(w, "llamaswap_scheduler_queue_oldest_wait_seconds{band=%q} %g\n", band, stats.Bands[band].OldestWait.Seconds()) + } + + fmt.Fprintf(w, "# HELP llamaswap_scheduler_dispatched_total Requests dispatched by the scheduler, by priority band\n") + fmt.Fprintf(w, "# TYPE llamaswap_scheduler_dispatched_total counter\n") + for _, band := range shared.PriorityBands { + fmt.Fprintf(w, "llamaswap_scheduler_dispatched_total{band=%q} %d\n", band, stats.Bands[band].Dispatched) + } + + fmt.Fprintf(w, "# HELP llamaswap_scheduler_reorders_total Dispatches where a scoring term changed which request was picked\n") + fmt.Fprintf(w, "# TYPE llamaswap_scheduler_reorders_total counter\n") + fmt.Fprintf(w, "llamaswap_scheduler_reorders_total{cause=\"aging\"} %d\n", stats.AgingReorders) + fmt.Fprintf(w, "llamaswap_scheduler_reorders_total{cause=\"swap_affinity\"} %d\n", stats.AffinityReorders) +} diff --git a/internal/server/server_test.go b/internal/server/server_test.go index 964fcb81..20130950 100644 --- a/internal/server/server_test.go +++ b/internal/server/server_test.go @@ -18,6 +18,7 @@ import ( "github.com/mostlygeek/llama-swap/internal/logmon" "github.com/mostlygeek/llama-swap/internal/process" "github.com/mostlygeek/llama-swap/internal/router" + "github.com/mostlygeek/llama-swap/internal/router/scheduler" "github.com/mostlygeek/llama-swap/internal/shared" ) @@ -29,6 +30,7 @@ type stubRouter struct { running map[string]process.ProcessState unloadCalls atomic.Int32 loggers map[string]*logmon.Monitor + queueStats *scheduler.QueueStats } func newStubRouter(models []string, response string) *stubRouter { @@ -57,6 +59,13 @@ func (s *stubRouter) ProcessLogger(modelID string) (*logmon.Monitor, bool) { return nil, false } +func (s *stubRouter) SchedulerStats() (scheduler.QueueStats, bool) { + if s.queueStats == nil { + return scheduler.QueueStats{}, false + } + return *s.queueStats, true +} + // newTestServer wires a Server with stub routers and a built mux. func newTestServer(local router.LocalRouter, peer router.Router) *Server { ctx, cancel := context.WithCancel(context.Background()) diff --git a/internal/shared/http.go b/internal/shared/http.go index ea0d3dfe..ae36ea7a 100644 --- a/internal/shared/http.go +++ b/internal/shared/http.go @@ -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 } diff --git a/internal/shared/priority.go b/internal/shared/priority.go new file mode 100644 index 00000000..4bc9d9c8 --- /dev/null +++ b/internal/shared/priority.go @@ -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 + } +} diff --git a/internal/shared/priority_test.go b/internal/shared/priority_test.go new file mode 100644 index 00000000..db0e4099 --- /dev/null +++ b/internal/shared/priority_test.go @@ -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/. +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) + } + }) + } +}