Merge pull request #10 from feat/request-priority
Build CUDA image (fork) / build (push) Successful in 1m25s
Build CUDA image (fork) / build (push) Successful in 1m25s
internal/router: priority queues so batch jobs yield to interactive requests Callers declare intent with X-LlamaSwap-Priority (signed int, 0 default, interactive/normal/batch aliases at +100/0/-100). serial dispatches by score = priority + swap affinity + aging; fifo adds the header to the model's configured priority so it is not silently ignored there. Closes #9. Consumer side: steve/mort#1576.
This commit was merged in pull request #10.
This commit is contained in:
@@ -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
@@ -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
@@ -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
|
||||||
|
|||||||
@@ -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.
|
||||||
@@ -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"`
|
||||||
|
|||||||
@@ -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)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -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 {
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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{})
|
||||||
|
|||||||
@@ -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)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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.
|
||||||
|
|||||||
@@ -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 {
|
||||||
|
|||||||
@@ -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,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -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
@@ -293,16 +293,26 @@ 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) {
|
||||||
w.WriteHeader(http.StatusOK)
|
w.WriteHeader(http.StatusOK)
|
||||||
|
|||||||
@@ -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, ""))
|
||||||
|
|
||||||
|
|||||||
@@ -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)
|
||||||
|
}
|
||||||
@@ -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())
|
||||||
|
|||||||
@@ -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
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -0,0 +1,89 @@
|
|||||||
|
package shared
|
||||||
|
|
||||||
|
import (
|
||||||
|
"net/http"
|
||||||
|
"strconv"
|
||||||
|
"strings"
|
||||||
|
)
|
||||||
|
|
||||||
|
// PriorityHeader carries the caller's scheduling priority. Priority is a
|
||||||
|
// property of the caller's intent, not of the model — the same model serves
|
||||||
|
// both an interactive request and a batch job — so it travels as a header
|
||||||
|
// rather than as per-model configuration, and leaves the OpenAI-compatible
|
||||||
|
// request body untouched.
|
||||||
|
const PriorityHeader = "X-LlamaSwap-Priority"
|
||||||
|
|
||||||
|
// Priority band anchors. The number is the interface: callers may send any
|
||||||
|
// signed integer, and these names are conveniences that resolve to a value.
|
||||||
|
//
|
||||||
|
// Bands are spaced 100 apart so a consumer can add a small per-caller offset
|
||||||
|
// (a subscription tier, say) on top of a band without ever promoting a request
|
||||||
|
// across one: a "max member" batch job at -98 still loses to every normal
|
||||||
|
// request at 0.
|
||||||
|
const (
|
||||||
|
PriorityInteractive = 100 // a human is waiting
|
||||||
|
PriorityNormal = 0 // the default
|
||||||
|
PriorityBatch = -100
|
||||||
|
)
|
||||||
|
|
||||||
|
// bandHalfWidth is how far a priority may sit from a band anchor and still be
|
||||||
|
// reported as that band. It is half the 100-point band spacing, so every
|
||||||
|
// integer belongs to exactly one band.
|
||||||
|
const bandHalfWidth = 50
|
||||||
|
|
||||||
|
// Band names used for metrics labels, ordered most to least urgent.
|
||||||
|
const (
|
||||||
|
BandInteractive = "interactive"
|
||||||
|
BandNormal = "normal"
|
||||||
|
BandBatch = "batch"
|
||||||
|
)
|
||||||
|
|
||||||
|
// PriorityBands is every band name, in descending order of urgency. Metrics
|
||||||
|
// emit a series per band regardless of whether any request currently occupies
|
||||||
|
// it, so the label set stays stable.
|
||||||
|
var PriorityBands = []string{BandInteractive, BandNormal, BandBatch}
|
||||||
|
|
||||||
|
var priorityAliases = map[string]int{
|
||||||
|
"interactive": PriorityInteractive,
|
||||||
|
"normal": PriorityNormal,
|
||||||
|
"batch": PriorityBatch,
|
||||||
|
}
|
||||||
|
|
||||||
|
// ParsePriority resolves a PriorityHeader value to a signed integer. A numeric
|
||||||
|
// value is used as-is; a recognised alias resolves to its anchor. Anything
|
||||||
|
// absent or unparseable is PriorityNormal, so a malformed header degrades to
|
||||||
|
// the default rather than failing the request.
|
||||||
|
//
|
||||||
|
// Values are deliberately not clamped: the caller composes band, tier offset
|
||||||
|
// and anything else it wants into one number, and llama-swap honours it.
|
||||||
|
func ParsePriority(value string) int {
|
||||||
|
value = strings.TrimSpace(value)
|
||||||
|
if value == "" {
|
||||||
|
return PriorityNormal
|
||||||
|
}
|
||||||
|
if n, err := strconv.Atoi(value); err == nil {
|
||||||
|
return n
|
||||||
|
}
|
||||||
|
if n, ok := priorityAliases[strings.ToLower(value)]; ok {
|
||||||
|
return n
|
||||||
|
}
|
||||||
|
return PriorityNormal
|
||||||
|
}
|
||||||
|
|
||||||
|
// RequestPriority reads PriorityHeader off r and resolves it. See ParsePriority.
|
||||||
|
func RequestPriority(r *http.Request) int {
|
||||||
|
return ParsePriority(r.Header.Get(PriorityHeader))
|
||||||
|
}
|
||||||
|
|
||||||
|
// PriorityBand buckets a priority into the band it belongs to, for metrics
|
||||||
|
// labels. Values beyond the anchors saturate: +1000 is still "interactive".
|
||||||
|
func PriorityBand(priority int) string {
|
||||||
|
switch {
|
||||||
|
case priority >= PriorityInteractive-bandHalfWidth:
|
||||||
|
return BandInteractive
|
||||||
|
case priority <= PriorityBatch+bandHalfWidth:
|
||||||
|
return BandBatch
|
||||||
|
default:
|
||||||
|
return BandNormal
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,118 @@
|
|||||||
|
package shared
|
||||||
|
|
||||||
|
import (
|
||||||
|
"net/http"
|
||||||
|
"net/http/httptest"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/mostlygeek/llama-swap/internal/config"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestShared_ParsePriority(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
value string
|
||||||
|
want int
|
||||||
|
}{
|
||||||
|
{"absent", "", PriorityNormal},
|
||||||
|
{"zero", "0", 0},
|
||||||
|
{"positive", "100", 100},
|
||||||
|
{"negative", "-100", -100},
|
||||||
|
{"explicit plus", "+42", 42},
|
||||||
|
{"surrounding space", " 100 ", 100},
|
||||||
|
{"large value is not clamped", "100000", 100000},
|
||||||
|
{"alias interactive", "interactive", PriorityInteractive},
|
||||||
|
{"alias normal", "normal", PriorityNormal},
|
||||||
|
{"alias batch", "batch", PriorityBatch},
|
||||||
|
{"alias is case insensitive", "Batch", PriorityBatch},
|
||||||
|
{"tier offset on a band", "-98", -98},
|
||||||
|
{"unknown alias falls back to normal", "urgent", PriorityNormal},
|
||||||
|
{"garbage falls back to normal", "!!", PriorityNormal},
|
||||||
|
{"float falls back to normal", "1.5", PriorityNormal},
|
||||||
|
}
|
||||||
|
for _, tc := range tests {
|
||||||
|
t.Run(tc.name, func(t *testing.T) {
|
||||||
|
if got := ParsePriority(tc.value); got != tc.want {
|
||||||
|
t.Errorf("ParsePriority(%q)=%d want %d", tc.value, got, tc.want)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestShared_RequestPriority(t *testing.T) {
|
||||||
|
r := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
||||||
|
if got := RequestPriority(r); got != PriorityNormal {
|
||||||
|
t.Errorf("no header: got %d want %d", got, PriorityNormal)
|
||||||
|
}
|
||||||
|
|
||||||
|
r.Header.Set(PriorityHeader, "batch")
|
||||||
|
if got := RequestPriority(r); got != PriorityBatch {
|
||||||
|
t.Errorf("batch header: got %d want %d", got, PriorityBatch)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestShared_PriorityBand pins the bucketing the metrics labels rely on. The
|
||||||
|
// important cases are the tier offsets: a max member's batch job at -98 must
|
||||||
|
// still report as batch, not normal.
|
||||||
|
func TestShared_PriorityBand(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
priority int
|
||||||
|
want string
|
||||||
|
}{
|
||||||
|
{PriorityInteractive, BandInteractive},
|
||||||
|
{PriorityInteractive + 2, BandInteractive}, // max tier, interactive
|
||||||
|
{1000, BandInteractive}, // saturates
|
||||||
|
{50, BandInteractive}, // lower edge
|
||||||
|
{49, BandNormal},
|
||||||
|
{PriorityNormal, BandNormal},
|
||||||
|
{2, BandNormal}, // max tier, normal
|
||||||
|
{-49, BandNormal},
|
||||||
|
{-50, BandBatch}, // upper edge
|
||||||
|
{PriorityBatch + 2, BandBatch},
|
||||||
|
{PriorityBatch, BandBatch},
|
||||||
|
{-1000, BandBatch}, // saturates
|
||||||
|
}
|
||||||
|
for _, tc := range tests {
|
||||||
|
if got := PriorityBand(tc.priority); got != tc.want {
|
||||||
|
t.Errorf("PriorityBand(%d)=%q want %q", tc.priority, got, tc.want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestShared_FetchContext_Priority verifies the header reaches the request
|
||||||
|
// context, which is what carries it to the scheduler. Both entry points into
|
||||||
|
// FetchContext are covered: the normal body-parsed path and /upstream/<model>.
|
||||||
|
func TestShared_FetchContext_Priority(t *testing.T) {
|
||||||
|
cfg := config.Config{Models: map[string]config.ModelConfig{"m1": {}}}
|
||||||
|
|
||||||
|
cases := []struct {
|
||||||
|
name string
|
||||||
|
path string
|
||||||
|
header string
|
||||||
|
want int
|
||||||
|
}{
|
||||||
|
{"body path, no header", "/v1/chat/completions", "", PriorityNormal},
|
||||||
|
{"body path, alias", "/v1/chat/completions", "interactive", PriorityInteractive},
|
||||||
|
{"body path, number", "/v1/chat/completions", "-98", -98},
|
||||||
|
{"upstream path", "/upstream/m1/v1/chat/completions", "batch", PriorityBatch},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, c := range cases {
|
||||||
|
t.Run(c.name, func(t *testing.T) {
|
||||||
|
r := httptest.NewRequest(http.MethodPost, c.path, strings.NewReader(`{"model":"m1"}`))
|
||||||
|
r.Header.Set("Content-Type", "application/json")
|
||||||
|
if c.header != "" {
|
||||||
|
r.Header.Set(PriorityHeader, c.header)
|
||||||
|
}
|
||||||
|
|
||||||
|
data, err := FetchContext(r, cfg)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("FetchContext: %v", err)
|
||||||
|
}
|
||||||
|
if data.Priority != c.want {
|
||||||
|
t.Errorf("Priority=%d want %d", data.Priority, c.want)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user