internal/router: priority queues so batch jobs yield to interactive requests

The GPU is a size-1 resource, so a single long job monopolises the box for its
whole duration and every interactive request queues behind it. Callers can now
declare intent with an X-LlamaSwap-Priority header and the serial scheduler
dispatches by score instead of by arrival.

- X-LlamaSwap-Priority: signed integer, 0 default, absent/unparseable means 0.
  interactive/normal/batch aliases resolve to +100/0/-100. Values are not
  clamped: the caller composes band and any per-user offset itself.
- serial dispatch score = priority + swap affinity + aging. Bands sit 100 apart
  so a small caller offset orders work inside a band without crossing one;
  aging is unbounded so low-priority work cannot starve.
- routing.scheduler.settings.serial.{agingDivisor,swapAffinityBonus}, defaulting
  to 60s/point and +10. swapAffinityBonus is capped at 99 so it can never
  promote a request into the next band.
- fifo adds the header to its per-model priority, so the header is not silently
  ignored under that scheduler.
- /metrics exports per-band queue depth, oldest wait and dispatch counts, plus
  counters for how often aging or swap affinity changed the pick. Each request
  records its priority, band, score and queue wait in the activity log.

Note swapAffinityBonus defaults to 10, so equal-priority requests for the
already-loaded model now run before older requests that need a swap. Set it to
0 for the previous strict arrival order.

fixes #9

Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
Claude-Session: https://claude.ai/code/session_01WUyhZBgv8BBCC5MduX88gE
This commit is contained in:
2026-08-07 02:12:08 -04:00
co-authored by Claude Opus 5
parent 2f0bb2556e
commit 0358fe321e
22 changed files with 1515 additions and 67 deletions
+3
View File
@@ -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 - `matrix` to run concurrent models with a custom swap logic DSL
- `hooks` to run things on startup - `hooks` to run things on startup
- `macros` reusable snippets - `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 - Model customization
- `ttl` to automatically unload models - `ttl` to automatically unload models
- `aliases` to use familiar model names (e.g., "gpt-4o-mini") - `aliases` to use familiar model names (e.g., "gpt-4o-mini")
+22 -2
View File
@@ -605,17 +605,37 @@
"fifo" "fifo"
], ],
"default": "serial", "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": { "settings": {
"type": "object", "type": "object",
"properties": { "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": { "fifo": {
"type": "object", "type": "object",
"properties": { "properties": {
"priority": { "priority": {
"type": "object", "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": { "additionalProperties": {
"type": "integer" "type": "integer"
} }
+44 -5
View File
@@ -559,23 +559,62 @@ routing:
# scheduler: how queued requests are ordered and run. # scheduler: how queued requests are ordered and run.
# - optional, default on this fork: "serial" # - optional, default on this fork: "serial"
# - valid values: # - valid values:
# - "serial": strict one-model-at-a-time. Requests run in exact arrival # - "serial": strict one-model-at-a-time. Only one request runs at a time;
# order; only one request runs at a time; switching to a different model # switching to a different model evicts every other running model first so
# evicts every other running model first so a single model occupies memory # a single model occupies memory at a time. This ignores group/matrix
# at a time. This ignores group/matrix co-residency entirely. The "fifo" # co-residency entirely. The "fifo" settings below do not apply.
# settings below (priority) do not apply.
# - "fifo": throughput-oriented. Same-model requests are batched to reduce # - "fifo": throughput-oriented. Same-model requests are batched to reduce
# swaps and a model serves up to its concurrencyLimit in parallel; models # swaps and a model serves up to its concurrencyLimit in parallel; models
# in non-exclusive groups can run concurrently. Requests may be reordered. # 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: scheduler:
use: serial use: serial
settings: 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 settings only apply when use: fifo
fifo: fifo:
# priority: a dictionary of model ID -> priority # priority: a dictionary of model ID -> priority
# - optional, default: empty dictionary # - optional, default: empty dictionary
# - models default to priority 0 # - models default to priority 0
# - higher priority requests are serviced first in the queue # - 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: priority:
A: 10 A: 10
B: 5 B: 5
+151
View File
@@ -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.
+54
View File
@@ -168,12 +168,66 @@ type SchedulerConfig struct {
type SchedulerSettings struct { type SchedulerSettings struct {
Fifo FifoConfig `yaml:"fifo"` Fifo FifoConfig `yaml:"fifo"`
Serial SerialConfig `yaml:"serial"`
} }
type FifoConfig struct { type FifoConfig struct {
Priority map[string]int `yaml:"priority"` // model ID -> priority, default 0 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 { type RouterConfig struct {
Use string `yaml:"use"` // "group" (default) | "matrix" Use string `yaml:"use"` // "group" (default) | "matrix"
Settings RouterSettings `yaml:"settings"` Settings RouterSettings `yaml:"settings"`
+74
View File
@@ -1720,3 +1720,77 @@ routing:
require.NoError(t, err) require.NoError(t, err)
assert.Equal(t, 5, cfg.Routing.Scheduler.Settings.Fifo.Priority["gemma"]) 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)
})
}
}
+10
View File
@@ -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) 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 // Clean up hooks preload
if len(config.Hooks.OnStartup.Preload) > 0 { if len(config.Hooks.OnStartup.Preload) > 0 {
+12
View File
@@ -338,6 +338,17 @@ func (b *baseRouter) Handles(model string) bool {
return ok 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) { func (b *baseRouter) ProcessLogger(modelID string) (*logmon.Monitor, bool) {
if p, ok := b.processes[modelID]; ok { if p, ok := b.processes[modelID]; ok {
return p.Logger(), true return p.Logger(), true
@@ -421,6 +432,7 @@ func (b *baseRouter) ServeHTTP(w http.ResponseWriter, req *http.Request) {
hr := scheduler.HandlerReq{ hr := scheduler.HandlerReq{
Model: data.ModelID, Model: data.ModelID,
Priority: data.Priority,
Ctx: req.Context(), Ctx: req.Context(),
// Unbuffered: a successful send on Respond proves the waiter is // Unbuffered: a successful send on Respond proves the waiter is
// alive and consuming. grant() relies on this to avoid handing a // alive and consuming. grant() relies on this to avoid handing a
+72
View File
@@ -5,6 +5,7 @@ import (
"io" "io"
"net/http" "net/http"
"net/http/httptest" "net/http/httptest"
"sync"
"testing" "testing"
"time" "time"
@@ -12,6 +13,7 @@ import (
"github.com/mostlygeek/llama-swap/internal/logmon" "github.com/mostlygeek/llama-swap/internal/logmon"
"github.com/mostlygeek/llama-swap/internal/process" "github.com/mostlygeek/llama-swap/internal/process"
"github.com/mostlygeek/llama-swap/internal/router/scheduler" "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 // 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) { func TestBaseRouter_ModelNotFound(t *testing.T) {
a := newFakeProcess("a") a := newFakeProcess("a")
b := newTestBase(t, map[string]process.Process{"a": a}, &stubPlanner{}) b := newTestBase(t, map[string]process.Process{"a": a}, &stubPlanner{})
+5
View File
@@ -6,6 +6,7 @@ import (
"github.com/mostlygeek/llama-swap/internal/logmon" "github.com/mostlygeek/llama-swap/internal/logmon"
"github.com/mostlygeek/llama-swap/internal/process" "github.com/mostlygeek/llama-swap/internal/process"
"github.com/mostlygeek/llama-swap/internal/router/scheduler"
"github.com/mostlygeek/llama-swap/internal/shared" "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 // modelID must be a real (non-alias) config key. Returns false when the
// model is not known to this router. // model is not known to this router.
ProcessLogger(modelID string) (*logmon.Monitor, bool) 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)
} }
+13 -6
View File
@@ -280,7 +280,7 @@ func (s *FIFO) grantHandler(req HandlerReq, modelID string) {
return 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) 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) 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 // 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 // first queued item whose priority is strictly lower, so higher-priority
// are serviced first while equal-priority requests keep their arrival (FIFO) // requests are serviced first while equal-priority requests keep their arrival
// order. Priorities come from the FifoConfig; unlisted models default to 0. // (FIFO) order.
func (s *FIFO) enqueue(req HandlerReq) { func (s *FIFO) enqueue(req HandlerReq) {
p := s.cfg.Priority[req.Model] p := s.priorityOf(req)
i := len(s.queued) i := len(s.queued)
for j, q := range s.queued { for j, q := range s.queued {
if s.cfg.Priority[q.Model] < p { if s.priorityOf(q) < p {
i = j i = j
break break
} }
+28
View File
@@ -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 // TestFIFO_OnCancel_QueuedRequest verifies that cancelling a queued request
// prevents drainQueue from ever starting a model load for it. Without OnCancel // 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. // the dead request would sit in the queue until a drain triggers a wasted swap.
+40 -2
View File
@@ -95,7 +95,7 @@ type Effects interface {
// New returns a Scheduler selected by conf.Routing.Scheduler.Use, configured from // 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" // conf and bound to the given planner and effects. Supported values are "fifo"
// (throughput-oriented, batches same-model requests) and "serial" (strict // (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 // The deployment default is applied by config loading (LoadConfig sets Use to
// "serial" when unset). The "" fallback here is the library default and remains // "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 return NewFIFO(name, logger, planner, conf.Routing.Scheduler.Settings.Fifo, conf.Models, eff), nil
case "serial": case "serial":
// Serial ignores the group planner: it always evicts every other model. // 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: default:
return nil, fmt.Errorf("unsupported scheduler type: %q", use) return nil, fmt.Errorf("unsupported scheduler type: %q", use)
} }
@@ -119,11 +119,49 @@ func New(conf config.Config, name string, logger *logmon.Monitor, planner Swappe
// HandlerReq is one in-flight ServeHTTP request waiting for a routing decision. // HandlerReq is one in-flight ServeHTTP request waiting for a routing decision.
type HandlerReq struct { 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 Ctx context.Context
Respond chan HandlerResp Respond chan HandlerResp
PositionCh chan int 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 // HandlerResp is the routing decision returned to a HandlerReq's caller: either
// a handler to serve with, or an error. // a handler to serve with, or an error.
type HandlerResp struct { type HandlerResp struct {
+319 -36
View File
@@ -3,40 +3,101 @@ package scheduler
import ( import (
"fmt" "fmt"
"sort" "sort"
"strconv"
"sync"
"time" "time"
"github.com/mostlygeek/llama-swap/internal/config"
"github.com/mostlygeek/llama-swap/internal/logmon" "github.com/mostlygeek/llama-swap/internal/logmon"
"github.com/mostlygeek/llama-swap/internal/process" "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 // Serial is a strict one-model-at-a-time scheduler for a size-1 resource: at
// or batches: requests run in exact arrival order and at most one request runs at // most one request runs at any instant, and when the next request targets a
// any instant. When the next request targets a model other than the one loaded, // model other than the one loaded, every other running model is evicted and the
// every other running model is evicted and the target is loaded before it runs, // target is loaded before it runs. A single model occupies memory at a time, at
// so a single model occupies memory at a time — at the cost of throughput. // 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 ignores group/eviction policy entirely: it always evicts every other // Serial ignores group/eviction policy entirely: it always evicts every other
// running model, regardless of how groups are configured. That is what makes the // running model, regardless of how groups are configured. That is what makes
// single-model guarantee a property of the scheduler rather than of the config. // 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 // # Dispatch order
// internal locking is needed. //
// 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 { type Serial struct {
name string name string
logger *logmon.Monitor logger *logmon.Monitor
effects Effects effects Effects
// queued holds requests in strict arrival order. It is never reordered. // agingDivisor is how long a request must wait to gain one priority
queued []HandlerReq // 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), // active is the one request currently being processed (loading or serving),
// or nil when idle. phase is meaningful only while active != nil. // or nil when idle. phase is meaningful only while active != nil.
active *HandlerReq active *HandlerReq
phase serialPhase 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. // 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 // NewSerial builds a Serial scheduler. It takes no Swapper: eviction is always
// "stop every other running model", so the group planner is not consulted. // "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{ return &Serial{
name: name, name: name,
logger: logger, logger: logger,
effects: eff, 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, // OnRequest validates the model and appends the request to the queue, stamping
// then tries to start the next job. Unknown models fail immediately. // 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) { func (s *Serial) OnRequest(req HandlerReq) {
if _, ok := s.effects.ModelState(req.Model); !ok { if _, ok := s.effects.ModelState(req.Model); !ok {
s.logger.Debugf("%s: model %s not handled by this router", s.name, req.Model) s.logger.Debugf("%s: model %s not handled by this router", s.name, req.Model)
s.effects.GrantError(req, ErrModelNotFound) s.effects.GrantError(req, ErrModelNotFound)
return return
} }
s.queued = append(s.queued, req) s.queued = append(s.queued, queuedReq{req: req, enqueued: s.now()})
broadcastQueuePositions(s.queued) s.queueChanged()
s.startNext() s.startNext()
} }
// startNext begins processing the head of the queue when nothing is active. It // startNext begins processing the highest scoring waiting request when nothing
// fast-paths a request whose model is already the sole loaded-and-ready process; // is active. It fast-paths a request whose model is already the sole
// otherwise it launches a swap that evicts every other running model first. The // loaded-and-ready process; otherwise it launches a swap that evicts every other
// loop skips over requests for models that vanished (e.g. a config reload) and // running model first. The loop skips over requests for models that vanished
// requests whose caller disconnected before they could be served. // (e.g. a config reload) and requests whose caller disconnected before they
// could be served.
func (s *Serial) startNext() { func (s *Serial) startNext() {
if s.active != nil { if s.active != nil {
return // a job is already loading or serving return // a job is already loading or serving
} }
for len(s.queued) > 0 { for len(s.queued) > 0 {
req := s.queued[0] idx, score := s.pick()
s.queued = s.queued[1:] q := s.queued[idx]
broadcastQueuePositions(s.queued) s.queued = append(s.queued[:idx], s.queued[idx+1:]...)
s.queueChanged()
req := q.req
state, ok := s.effects.ModelState(req.Model) state, ok := s.effects.ModelState(req.Model)
if !ok { if !ok {
s.effects.GrantError(req, ErrModelNotFound) s.effects.GrantError(req, ErrModelNotFound)
continue continue
} }
waited := s.now().Sub(q.enqueued)
s.annotate(req, score, waited)
r := req r := req
s.active = &r s.active = &r
evict := s.otherRunning(req.Model) evict := s.otherRunning(req.Model)
if state == process.StateReady && len(evict) == 0 { if state == process.StateReady && len(evict) == 0 {
// Already loaded and the only model running — serve immediately. // 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() { if s.serve() {
s.recordDispatch(req.Priority)
return return
} }
continue // caller gone; pick the next request 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.phase = phaseSwapping
s.recordDispatch(req.Priority)
s.effects.StartSwap(req.Model, evict) s.effects.StartSwap(req.Model, evict)
return 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 // 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 // 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. // 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] kept := s.queued[:0]
removed := false removed := false
for _, q := range s.queued { for _, q := range s.queued {
if q.Respond == req.Respond { if q.req.Respond == req.Respond {
removed = true removed = true
continue continue
} }
@@ -177,7 +354,7 @@ func (s *Serial) OnCancel(req HandlerReq) {
s.queued = kept s.queued = kept
if removed { if removed {
s.logger.Debugf("%s: cancelled request for model %s pruned from queue", s.name, req.Model) 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 { if len(s.queued) > 0 {
kept := s.queued[:0] kept := s.queued[:0]
for _, q := range s.queued { for _, q := range s.queued {
if targetSet[q.Model] { if targetSet[q.req.Model] {
s.effects.GrantError(q, unloadErr) s.effects.GrantError(q.req, unloadErr)
continue continue
} }
kept = append(kept, q) kept = append(kept, q)
} }
s.queued = kept s.queued = kept
broadcastQueuePositions(s.queued) s.queueChanged()
} }
s.effects.StopProcesses(timeout, targets) s.effects.StopProcesses(timeout, targets)
@@ -234,9 +411,10 @@ func (s *Serial) OnShutdown(err error) {
s.phase = phaseIdle s.phase = phaseIdle
} }
for _, q := range s.queued { for _, q := range s.queued {
s.effects.GrantError(q, err) s.effects.GrantError(q.req, err)
} }
s.queued = nil s.queued = nil
s.queueChanged()
} }
// otherRunning returns every running model except target, sorted for // otherRunning returns every running model except target, sorted for
@@ -251,3 +429,108 @@ func (s *Serial) otherRunning(target string) []string {
sort.Strings(out) sort.Strings(out)
return 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,
}
}
+345 -5
View File
@@ -1,13 +1,16 @@
package scheduler package scheduler
import ( import (
"context"
"errors" "errors"
"io" "io"
"testing" "testing"
"time" "time"
"github.com/mostlygeek/llama-swap/internal/config"
"github.com/mostlygeek/llama-swap/internal/logmon" "github.com/mostlygeek/llama-swap/internal/logmon"
"github.com/mostlygeek/llama-swap/internal/process" "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 // 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 // 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. // 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 { 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. // 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, // TestSerial_StrictArrivalOrder covers the scoring function's degenerate case:
// qwen36 execute in EXACTLY that order with evictions between each model switch, // with equal priorities and swap affinity disabled, qwen36, qwen35, sdxl,
// including reloading qwen36 at the end even though it ran first. // 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) { func TestSerial_StrictArrivalOrder(t *testing.T) {
eff := newFakeEffects() eff := newFakeEffects()
for _, m := range []string{"qwen36", "qwen35", "sdxl"} { for _, m := range []string{"qwen36", "qwen35", "sdxl"} {
eff.states[m] = process.StateStopped eff.states[m] = process.StateStopped
} }
s := newSerial(eff) s := newSerialCfg(eff, config.DefaultAgingDivisor, 0)
for _, m := range []string{"qwen36", "qwen35", "sdxl", "qwen36"} { for _, m := range []string{"qwen36", "qwen35", "sdxl", "qwen36"} {
s.OnRequest(req(m)) s.OnRequest(req(m))
@@ -243,6 +277,312 @@ func sameOrder(a, b []string) bool {
return true 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) { func TestSerial_SwapError_FailsCallerAndAdvances(t *testing.T) {
eff := newFakeEffects() eff := newFakeEffects()
eff.states["a"] = process.StateStopped eff.states["a"] = process.StateStopped
+13 -3
View File
@@ -293,15 +293,25 @@ func (s *Server) startPreload() {
}() }()
} }
// handleMetrics serves Prometheus-format performance metrics. Returns 503 when // handleMetrics serves Prometheus-format metrics: system/GPU performance plus
// performance monitoring is disabled. // 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) { 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.WriteHeader(http.StatusServiceUnavailable)
w.Write([]byte("# performance monitor not available\n")) w.Write([]byte("# performance monitor not available\n"))
return return
} }
w.Header().Set("Content-Type", "text/plain; version=0.0.4; charset=utf-8")
if s.perf != nil {
s.perf.MetricsHandler().ServeHTTP(w, r) s.perf.MetricsHandler().ServeHTTP(w, r)
}
if hasQueueStats {
writeSchedulerMetrics(w, queueStats)
}
} }
func handleHealth(w http.ResponseWriter, r *http.Request) { func handleHealth(w http.ResponseWriter, r *http.Request) {
+42
View File
@@ -8,10 +8,12 @@ import (
"regexp" "regexp"
"strings" "strings"
"testing" "testing"
"time"
"github.com/mostlygeek/llama-swap/internal/config" "github.com/mostlygeek/llama-swap/internal/config"
"github.com/mostlygeek/llama-swap/internal/logmon" "github.com/mostlygeek/llama-swap/internal/logmon"
"github.com/mostlygeek/llama-swap/internal/process" "github.com/mostlygeek/llama-swap/internal/process"
"github.com/mostlygeek/llama-swap/internal/router/scheduler"
"github.com/mostlygeek/llama-swap/internal/shared" "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) { 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, "")) s := newTestServer(newStubRouter(nil, ""), newStubRouter(nil, ""))
w := httptest.NewRecorder() 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) { func TestServer_Redirects(t *testing.T) {
s := newTestServer(newStubRouter(nil, ""), newStubRouter(nil, "")) s := newTestServer(newStubRouter(nil, ""), newStubRouter(nil, ""))
+39
View File
@@ -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)
}
+9
View File
@@ -18,6 +18,7 @@ import (
"github.com/mostlygeek/llama-swap/internal/logmon" "github.com/mostlygeek/llama-swap/internal/logmon"
"github.com/mostlygeek/llama-swap/internal/process" "github.com/mostlygeek/llama-swap/internal/process"
"github.com/mostlygeek/llama-swap/internal/router" "github.com/mostlygeek/llama-swap/internal/router"
"github.com/mostlygeek/llama-swap/internal/router/scheduler"
"github.com/mostlygeek/llama-swap/internal/shared" "github.com/mostlygeek/llama-swap/internal/shared"
) )
@@ -29,6 +30,7 @@ type stubRouter struct {
running map[string]process.ProcessState running map[string]process.ProcessState
unloadCalls atomic.Int32 unloadCalls atomic.Int32
loggers map[string]*logmon.Monitor loggers map[string]*logmon.Monitor
queueStats *scheduler.QueueStats
} }
func newStubRouter(models []string, response string) *stubRouter { func newStubRouter(models []string, response string) *stubRouter {
@@ -57,6 +59,13 @@ func (s *stubRouter) ProcessLogger(modelID string) (*logmon.Monitor, bool) {
return nil, false 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. // newTestServer wires a Server with stub routers and a built mux.
func newTestServer(local router.LocalRouter, peer router.Router) *Server { func newTestServer(local router.LocalRouter, peer router.Router) *Server {
ctx, cancel := context.WithCancel(context.Background()) ctx, cancel := context.WithCancel(context.Background())
+5
View File
@@ -26,6 +26,9 @@ type ReqContextData struct {
ModelID string ModelID string
Streaming bool Streaming bool
SendLoadingState 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 // Metadata is a request-scoped key/value bag that handlers may mutate
// while processing. The metrics middleware copies it into ActivityLogEntry. // while processing. The metrics middleware copies it into ActivityLogEntry.
Metadata map[string]string 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 { if mc, ok := cfg.Models[realName]; ok {
data.SendLoadingState = mc.SendLoadingState != nil && *mc.SendLoadingState data.SendLoadingState = mc.SendLoadingState != nil && *mc.SendLoadingState
} }
data.Priority = RequestPriority(r)
*r = *r.WithContext(SetContext(r.Context(), data)) *r = *r.WithContext(SetContext(r.Context(), data))
return data, nil return data, nil
} }
@@ -137,6 +141,7 @@ func extractUpstreamContext(r *http.Request, cfg config.Config) (ReqContextData,
ApiKey: ExtractAPIKey(r), ApiKey: ExtractAPIKey(r),
Streaming: r.URL.Query().Get("stream") == "true", Streaming: r.URL.Query().Get("stream") == "true",
SendLoadingState: sendLoadingState(cfg, realName), SendLoadingState: sendLoadingState(cfg, realName),
Priority: RequestPriority(r),
Metadata: make(map[string]string), Metadata: make(map[string]string),
}, true }, true
} }
+89
View File
@@ -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
}
}
+118
View File
@@ -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)
}
})
}
}