Compare commits

...
19 Commits
Author SHA1 Message Date
steve c87893ba6f Merge pull request #10 from feat/request-priority
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.
2026-08-08 01:23:14 +00:00
steveandClaude Opus 5 0358fe321e internal/router: priority queues so batch jobs yield to interactive requests
The GPU is a size-1 resource, so a single long job monopolises the box for its
whole duration and every interactive request queues behind it. Callers can now
declare intent with an X-LlamaSwap-Priority header and the serial scheduler
dispatches by score instead of by arrival.

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

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

fixes #9

Co-Authored-By: Claude Opus 5 (1M context) <[email protected]>
Claude-Session: https://claude.ai/code/session_01WUyhZBgv8BBCC5MduX88gE
2026-08-07 02:12:08 -04:00
steveandClaude Fable 5 2f0bb2556e ci: push immutable sha- and moving main-cuda- tags alongside the version label
Build CUDA image (fork) / build (push) Successful in 1m9s
Push-triggered builds always defaulted to the v230 label, silently
overwriting it on every main merge — a stack pinned to any other
version label (v231) kept pulling stale images. Every build now also
pushes sha-<short> (immutable pin) and main-cuda-<llamacpp> (explicit
latest-main), so deployments can choose mutable-tracking or exact-pin
deliberately.

Co-Authored-By: Claude Fable 5 <[email protected]>
Claude-Session: https://claude.ai/code/session_01AXQxVhXBw8PwFAtsVrXSmj
2026-07-12 11:25:51 -04:00
steve fdd89d6580 Merge pull request 'feat(ui): Video playground tab + video capability plumbing' (#2) from feat/video-ui into main
Build CUDA image (fork) / build (push) Successful in 2m57s
Reviewed-on: #2
2026-07-12 14:56:35 +00:00
steveandClaude Fable 5 1934efb0b5 feat(ui): Video playground tab + video capability plumbing
New Video tab in the Playground: model selector filtered to
video_generation/image_to_video capabilities, prompt + negative prompt,
size/frames/fps/steps/guidance/seed knobs (0 = model default), optional
conditioning-image upload (image-to-video), elapsed-time spinner for
the minutes-long blocking call, inline <video> playback + download.
Calls POST /v1/videos/sync (multipart); resolution rides as both
width/height and size so either upstream convention honors it.

Config side: "video" joins the valid capability modalities, and
/v1/models maps text->video to video_generation and image->video to
image_to_video, mirroring the image mappings. Declare e.g.:

  videogen-wan22-5b:
    capabilities:
      in: [text, image]
      out: [video]

Co-Authored-By: Claude Fable 5 <[email protected]>
Claude-Session: https://claude.ai/code/session_01AXQxVhXBw8PwFAtsVrXSmj
2026-07-12 10:53:23 -04:00
steveandClaude Fable 5 b83edeb9ef ci: rebuild fork image on main source pushes, not just build-definition edits
Build CUDA image (fork) / build (push) Successful in 4m49s
The merge of the /v1/videos/sync routes produced no image because the
push trigger only watched the workflow file and Containerfile. Now any
Go/UI/deps change on main rebuilds; docs/config-example edits still
don't. (The video-routes image itself was built via manual dispatch,
tag prefix v231.)

Co-Authored-By: Claude Fable 5 <[email protected]>
Claude-Session: https://claude.ai/code/session_01AXQxVhXBw8PwFAtsVrXSmj
2026-07-12 10:24:00 -04:00
steve 7886758053 Merge pull request 'feat(server): dispatch /v1/videos/sync by model' (#1) from feat/video-routes into main
Reviewed-on: #1
2026-07-12 14:13:05 +00:00
steveandClaude Fable 5 d297852b9e fix: scope video dispatch to /v1/videos/sync + capture mask
Review findings: the async POST /v1/videos leg is deliberately NOT
dispatched — its poll/download companions carry no model field to route
by, so registering creation alone would start unretrievable upstream
jobs (a stray OpenAI-videos client now gets a clean 404). And
/v1/videos/sync joins captureFieldsByPath (headers only, both
directions) so enabling captures doesn't buffer conditioning frames or
whole clips through cbor+zstd.

Co-Authored-By: Claude Fable 5 <[email protected]>
Claude-Session: https://claude.ai/code/session_01AXQxVhXBw8PwFAtsVrXSmj
2026-07-12 09:53:33 -04:00
steveandClaude Fable 5 cecb89880e feat(server): dispatch /v1/videos and /v1/videos/sync by model
Registers the OpenAI/vLLM-Omni video generation endpoints as
model-dispatched routes so llama-swap can swap in video upstreams
(e.g. vllm serve <model> --omni, or a ComfyUI shim exposing the same
shape). Extraction is content-type driven, so both multipart form and
JSON bodies resolve the model field.

Co-Authored-By: Claude Fable 5 <[email protected]>
Claude-Session: https://claude.ai/code/session_01AXQxVhXBw8PwFAtsVrXSmj
2026-07-12 09:25:55 -04:00
steveandClaude Opus 4.8 0292c90ca1 ci: copy ui-svelte/.npmrc before npm ci in fork-cuda build
Build CUDA image (fork) / build (push) Successful in 12m49s
npm ci ran without .npmrc (legacy-peer-deps=true), failing on the
tailwind/vite peer dependency conflict. Copy .npmrc with the manifest.

Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
2026-06-28 12:56:21 -04:00
steveandClaude Opus 4.8 617c7dc6b9 ci: add Gitea workflow to build fork CUDA image
Build CUDA image (fork) / build (push) Failing after 2m23s
Add a Gitea Actions workflow and multi-stage Containerfile that build
this fork's llama-swap (serial scheduler + embedded Svelte UI) from
source and layer it on a pinned llama.cpp CUDA server base, then push to
the Gitea container registry as v230-cuda-b9821.

- docker/fork-cuda.Containerfile: node UI -> go build -> cuda runtime,
  runs as root to match the upstream non-suffixed image
- .gitea/workflows/build-cuda-image.yml: workflow_dispatch (version +
  llama.cpp build inputs) and push-on-build-files; logs in with
  REGISTRY_USER/REGISTRY_PASSWORD

Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
2026-06-28 12:48:48 -04:00
steveandClaude Opus 4.8 542b79dacf internal/router/scheduler: add serial scheduler, default on this fork
Validate JSON Schema / validate-schema (push) Successful in 9m53s
Linux CI / run-tests (push) Failing after 15m57s
Windows CI / run-tests (push) Has been cancelled
Add a strict one-model-at-a-time scheduler. Requests run in exact
arrival order; at most one runs at a time; switching to a different
model evicts every other running model first so a single model occupies
memory at a time. Unlike fifo it never reorders or batches same-model
requests, and it ignores group/matrix co-residency entirely, making the
single-model guarantee a property of the scheduler rather than the config.

- new Serial scheduler implementing the Scheduler interface
- register "serial" in scheduler.New; default routing.scheduler.use to
  "serial" at config load (fifo still selectable for upstream behavior)
- update config schema, example config, and config defaults tests

Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
2026-06-28 12:17:32 -04:00
Benson Wong 0a25b3bd31 AGENTS.md: small tweaks 2026-06-25 20:31:48 -07:00
Benson Wong 32bc781326 internal/config,watcher: add -config-dir (#873)
Over time the llama-swap configuration file can get really long and
challenging to work with. The -config-dir flag is used for a directory
of configuration YAML fragments.

These fragments are merged together and into a full configuration and
tested for validity. All previous configuration functionality remains
unchanged.
2026-06-24 20:48:51 -07:00
Benson Wong 316ad63f76 config,server: add upstream.ignorePaths (#869)
Add upstream.ignorePaths config to prevent model swaps for static-asset
requests made through the /upstream/<model>/<path> passthrough endpoint.

- add UpstreamConfig with compiled *regexp.Regexp slice; invalid regex
returns an error at load time
- apply a default pattern matching common static-asset suffixes
(.js/.json/.css/.png/.gif/.jpg/.jpeg/.ico/.txt) when unset
- in handleUpstream, return 409 Conflict when a path matches and the
local model is not already loaded; peer and already-loaded models fall
through to normal dispatch
- update config-schema.json and config.example.yaml

Updates discussion: #868
2026-06-21 13:49:53 -07:00
g2mt e37077a963 feat: hide performance menu item if disabled (#832)
Hide the Performance UI item of the navigation bar if its disabled.
2026-06-21 13:38:29 -07:00
Benson Wong eff9b60434 server: capture failed (non-200) LLM requests (#862)
Store a request/response capture for non-200 responses so failed
requests can be inspected in the activity log's Capture dialog, matching
the existing behavior for successful requests.

- extract storeCapture/decodeResponseBody helpers to share capture logic
between the success and non-200 paths
- record non-200 bodies (decompressed) so error details are viewable
- the activity UI already gates the View button on has_capture, so it
now appears for failed requests with no UI changes
- add tests for capturing failed requests and the disabled-captures case

closes #766
2026-06-20 11:50:35 -07:00
Wojciech 9bcddad91b internal/server,ui: add new Acitivty page column - Drafted (#859)
Add draft metrics to activity log
2026-06-18 20:55:02 -07:00
Benson Wong a15e47922c proxy: meter /upstream requests via metrics middleware (#858)
Wrap /upstream/{upstreamPath...} in the metrics middleware so activity
log entries are recorded for model-dispatched endpoints accessed through
the upstream passthrough.

- Move findModelInPath to shared.FindModelInPath and reuse it in
handleUpstream, the log monitor lookup, and FetchContext.
- Extend FetchContext to resolve the model from /upstream/<model>/...
paths without consuming the request body.
- Add isMetricsRecordPath to limit recording to the model-dispatched
endpoints that produce token usage/timings.
- Add tests for upstream metrics recording and FetchContext upstream
path resolution.

Fixes #855
2026-06-17 17:38:52 -07:00
55 changed files with 5435 additions and 805 deletions
+89
View File
@@ -0,0 +1,89 @@
name: Build CUDA image (fork)
# Builds this fork's llama-swap (serial scheduler + embedded UI) from source and
# layers it on a pinned llama.cpp CUDA server base, then pushes to the Gitea
# container registry, e.g. gitea.stevedudenhoeffer.com/steve/llama-swap:v230-cuda-b9821
#
# Requires repo secrets: REGISTRY_USER, REGISTRY_PASSWORD (push to the registry).
on:
workflow_dispatch:
inputs:
llama_swap_version:
description: "llama-swap version label (image tag prefix)"
required: false
default: "v230"
llamacpp_build:
description: "llama.cpp CUDA server build (base image tag suffix)"
required: false
default: "b9821"
# Any change that lands in the shipped binary or image kicks off a fresh
# build (source, UI, deps, build definition). Docs/config-example edits
# don't.
push:
branches: [main]
paths:
- ".gitea/workflows/build-cuda-image.yml"
- "docker/fork-cuda.Containerfile"
- "**/*.go"
- "go.mod"
- "go.sum"
- "ui-svelte/**"
- "llama-swap.go"
env:
REGISTRY: gitea.stevedudenhoeffer.com
jobs:
build:
runs-on: ubuntu-latest
steps:
- name: Checkout
uses: actions/checkout@v4
- name: Compute image metadata
id: meta
run: |
LS_VER="${{ inputs.llama_swap_version || 'v230' }}"
LCPP="${{ inputs.llamacpp_build || 'b9821' }}"
SHORT_SHA="$(echo '${{ github.sha }}' | cut -c1-7)"
IMAGE="${REGISTRY}/${{ github.repository }}"
# Three tags per build: the version label (mutable, human-facing),
# an explicit moving main-cuda tag, and an immutable sha tag so a
# deployment can always pin the exact build.
{
echo "image=${IMAGE}"
echo "tag=${LS_VER}-cuda-${LCPP}"
echo "tags=${IMAGE}:${LS_VER}-cuda-${LCPP},${IMAGE}:main-cuda-${LCPP},${IMAGE}:sha-${SHORT_SHA}"
echo "base_tag=server-cuda-${LCPP}"
echo "ls_version=${LS_VER}"
echo "build_date=$(date -u +%Y-%m-%dT%H:%M:%SZ)"
} >> "$GITHUB_OUTPUT"
- name: Set up Docker Buildx
uses: docker/setup-buildx-action@v3
- name: Log in to Gitea registry
uses: docker/login-action@v3
with:
registry: ${{ env.REGISTRY }}
username: ${{ secrets.REGISTRY_USER }}
password: ${{ secrets.REGISTRY_PASSWORD }}
- name: Build and push
uses: docker/build-push-action@v6
with:
context: .
file: docker/fork-cuda.Containerfile
push: true
provenance: false
build-args: |
BASE_TAG=${{ steps.meta.outputs.base_tag }}
LS_VERSION=${{ steps.meta.outputs.ls_version }}
GIT_HASH=${{ github.sha }}
BUILD_DATE=${{ steps.meta.outputs.build_date }}
tags: ${{ steps.meta.outputs.tags }}
- name: Summary
run: |
echo "Pushed ${{ steps.meta.outputs.tags }}" >> "$GITHUB_STEP_SUMMARY"
+4 -6
View File
@@ -5,16 +5,14 @@ llama-swap is a light weight, transparent proxy server that provides automatic m
## Tech stack
- golang
- typescript, vite and svelt5 for UI (located in ui/)
- typescript, vite and svelte 5 for UI (located in ui-svelte/)
## Workflow Tasks
- when summarizing changes only include details that require further action
- just say "Done." when there is no further action
- use the github CLI `gh` to create pull requests and work with github
- Rules for creating pull requests:
- keep them short and focused on changes.
- never include a test plan
- keep them short and focused on changes
- skip the test plan
- write the summary using the same style rules as commit message
## Testing
@@ -30,7 +28,7 @@ llama-swap is a light weight, transparent proxy server that provides automatic m
### Commit message example format:
```
proxy: add new feature
internal/server: add new feature
Add new feature that implements functionality X and Y.
+3
View File
@@ -213,6 +213,9 @@ Almost all configuration settings are optional and can be added one step at a ti
- `matrix` to run concurrent models with a custom swap logic DSL
- `hooks` to run things on startup
- `macros` reusable snippets
- `routing.scheduler` to control how queued requests are ordered — see
[request priority](docs/request-priority.md) for the `X-LlamaSwap-Priority`
header that lets batch jobs yield to interactive requests
- Model customization
- `ttl` to automatically unload models
- `aliases` to use familiar model names (e.g., "gpt-4o-mini")
+42 -3
View File
@@ -572,6 +572,24 @@
"default": {},
"description": "A dictionary of remote peers and models they provide. Peers can be another llama-swap or any server that provides the /v1/ generative API endpoints supported by llama-swap."
},
"upstream": {
"type": "object",
"description": "Controls behaviour of the /upstream passthrough endpoint. Recommended to only use in special use cases; leaving it as the default will typically be the best experience.",
"properties": {
"ignorePaths": {
"type": "array",
"items": {
"type": "string"
},
"default": [
".*\\.(js|json|css|png|gif|jpg|jpeg|ico|txt)$"
],
"description": "List of RE2 compatible regular expressions. Any request to a path matching any of the regular expressions will be ignored and not trigger a swap. When not specified, defaults to a pattern matching common static-asset suffixes (.js, .json, .css, .png, .gif, .jpg, .jpeg, .ico, .txt)."
}
},
"additionalProperties": false,
"default": {}
},
"routing": {
"type": "object",
"description": "Canonical routing/scheduling configuration. Alternative to the legacy top-level 'groups'/'matrix' keys; a config must not use both styles.",
@@ -583,20 +601,41 @@
"use": {
"type": "string",
"enum": [
"serial",
"fifo"
],
"default": "fifo",
"description": "Scheduler to use. Only 'fifo' is currently supported."
"default": "serial",
"description": "Scheduler to use. 'serial' (default on this fork): strict one-model-at-a-time, highest-priority request runs next, switching models evicts every other model first. 'fifo': throughput-oriented, batches same-model requests and allows parallel/co-resident models. Both honour the per-request X-LlamaSwap-Priority header."
},
"settings": {
"type": "object",
"properties": {
"serial": {
"type": "object",
"description": "Tunes how the serial scheduler scores waiting requests: score = X-LlamaSwap-Priority + swap affinity + aging. Ignored unless use is 'serial'.",
"properties": {
"agingDivisor": {
"type": "integer",
"minimum": 0,
"default": 60,
"description": "Seconds a request must wait to gain one priority point. 0 disables aging. Unbounded on purpose: aging is the only term allowed to promote a request across a priority band, which is what prevents starvation."
},
"swapAffinityBonus": {
"type": "integer",
"minimum": 0,
"maximum": 99,
"default": 10,
"description": "Bonus for a request that can be served without swapping models, so strict priority does not force a cold load on every dispatch. Capped below the 100-point band spacing so it can only break near-ties. 0 disables it, giving strict priority-then-arrival order."
}
},
"additionalProperties": false
},
"fifo": {
"type": "object",
"properties": {
"priority": {
"type": "object",
"description": "Per-model priority. Keys are model IDs, values are integers (default 0). Higher values are serviced first.",
"description": "Per-model priority. Keys are model IDs, values are integers (default 0). Higher values are serviced first. Added to the caller's X-LlamaSwap-Priority.",
"additionalProperties": {
"type": "integer"
}
+64 -3
View File
@@ -134,6 +134,18 @@ apiKeys:
- "${env.API_KEY_1}"
- "${env.API_KEY_2}"
# upstream: controls behaviour of the /upstream passthrough endpoint
# - optional, default: empty dictionary
# - recommended to only use in special use cases. Leaving it as the
# default will typically be the best experience
upstream:
# ignorePaths: list of RE2 compatible regular expressions
# - default: (see below)
# - any request to a path matching any of the regular expressions
# will be ignored and not trigger a swap
ignorePaths:
- '.*\.(js|json|css|png|gif|jpg|jpeg|ico|txt)$'
# models: a dictionary of model configurations
# - required
# - each key is the model's ID, used in API requests
@@ -544,16 +556,65 @@ routing:
# expands to: [L]
full: "L"
# scheduler: how queued requests are ordered.
# The default and only valid scheduler is "fifo"
# scheduler: how queued requests are ordered and run.
# - optional, default on this fork: "serial"
# - valid values:
# - "serial": strict one-model-at-a-time. Only one request runs at a time;
# switching to a different model evicts every other running model first so
# a single model occupies memory at a time. This ignores group/matrix
# co-residency entirely. The "fifo" settings below do not apply.
# - "fifo": throughput-oriented. Same-model requests are batched to reduce
# swaps and a model serves up to its concurrencyLimit in parallel; models
# in non-exclusive groups can run concurrently. Requests may be reordered.
#
# Callers set per-request priority with the X-LlamaSwap-Priority header
# (a signed integer; "interactive" = 100, "normal" = 0, "batch" = -100;
# absent or unparseable means 0). Both schedulers honour it. Scheduling is
# non-preemptive: priority applies when a job finishes, so an interactive
# request can still wait up to one full job duration.
scheduler:
use: fifo
use: serial
settings:
# serial settings only apply when use: serial
#
# At each dispatch the serial scheduler runs the highest scoring waiting
# request, where:
#
# score = X-LlamaSwap-Priority + swap affinity + aging
#
# Priority bands sit 100 apart, which leaves room for a caller to add a
# small per-user offset (a subscription tier, say) without crossing a
# band: a "max member" batch job at -98 beats other batch work but still
# loses to every normal request at 0.
serial:
# agingDivisor: seconds a request must wait to gain one priority point
# - optional, default: 60 (one point per minute)
# - 0 disables aging
# - unbounded on purpose: aging is the only term allowed to promote a
# request across a band, which is what stops low-priority work from
# starving. At 60, a batch job at -100 overtakes normal traffic after
# ~100 minutes of waiting.
agingDivisor: 60
# swapAffinityBonus: bonus for a request that needs no model swap
# - optional, default: 10
# - must be 0..99, so it can only break near-ties and can never promote
# a request into the next band
# - 0 disables it, giving strict priority-then-arrival order
# - cold loads are expensive, so this keeps a run of same-model requests
# together instead of forcing a reload on every dispatch. Note this
# means equal-priority requests are NOT served in strict arrival
# order; set it to 0 if you need that.
swapAffinityBonus: 10
# fifo settings only apply when use: fifo
fifo:
# priority: a dictionary of model ID -> priority
# - optional, default: empty dictionary
# - models default to priority 0
# - higher priority requests are serviced first in the queue
# - added to the caller's X-LlamaSwap-Priority, so a model priority of
# 10 and a header of "batch" (-100) queue at -90
priority:
A: 10
B: 5
+74
View File
@@ -0,0 +1,74 @@
# Build a CUDA llama-swap image FROM THIS FORK's source (includes the serial
# scheduler) and layer it on a pinned llama.cpp CUDA server base. Produces e.g.:
# gitea.stevedudenhoeffer.com/steve/llama-swap:v230-cuda-b9821
#
# BASE_TAG selects the llama.cpp CUDA runtime + llama-server build, e.g.
# "server-cuda-b9821". The llama-swap binary (with the embedded Svelte UI) is
# compiled from the repo at build time, so no GitHub release is required.
#
# Build context is the repo root:
# docker build -f docker/fork-cuda.Containerfile \
# --build-arg BASE_TAG=server-cuda-b9821 -t llama-swap:v230-cuda-b9821 .
ARG BASE_IMAGE=ghcr.io/ggml-org/llama.cpp
ARG BASE_TAG=server-cuda-b9821
# ---- Stage 1: build the Svelte UI (embedded into the binary) ----
FROM node:22-bookworm-slim AS ui
WORKDIR /src/ui-svelte
# Install deps first for layer caching. .npmrc carries legacy-peer-deps=true,
# which the project relies on (tailwind/vite peer ranges), so copy it before
# npm ci or the strict resolver fails with ERESOLVE.
COPY ui-svelte/package.json ui-svelte/package-lock.json ui-svelte/.npmrc ./
RUN npm ci
COPY ui-svelte/ ./
# `npm run build` is `vite build --emptyOutDir`; vite.config.ts writes to
# ../internal/server/ui_dist, which //go:embed picks up in the next stage.
RUN mkdir -p /src/internal/server && npm run build
# ---- Stage 2: build the llama-swap binary with the embedded UI ----
FROM golang:1.26-bookworm AS build
WORKDIR /src
# Cache modules independently of source churn.
COPY go.mod go.sum ./
RUN go mod download
COPY . .
# Overlay the freshly built UI so //go:embed ui_dist ships the real assets
# instead of the committed placeholder.
COPY --from=ui /src/internal/server/ui_dist/ ./internal/server/ui_dist/
ARG LS_VERSION=v230
ARG GIT_HASH=unknown
ARG BUILD_DATE=unknown
RUN CGO_ENABLED=0 GOOS=linux go build \
-ldflags="-X main.version=${LS_VERSION} -X main.commit=${GIT_HASH} -X main.date=${BUILD_DATE}" \
-o /out/llama-swap .
# ---- Stage 3: runtime image on the pinned llama.cpp CUDA base ----
FROM ${BASE_IMAGE}:${BASE_TAG}
# Run as root by default to match the upstream `vNNN-cuda-bNNNN` (non-suffixed)
# image that ragnaros pulls today: it needs root to reach the mounted docker
# socket for container-backed models (sd-server). Override UID/GID at build time
# for a non-root variant.
ARG UID=0
ARG GID=0
ARG USER_HOME=/root
ENV HOME=$USER_HOME
RUN set -eux; \
if [ "$UID" -ne 0 ]; then \
if [ "$GID" -ne 0 ]; then groupadd --system --gid "$GID" app; fi; \
useradd --system --uid "$UID" --gid "$GID" --home "$USER_HOME" app; \
fi; \
mkdir --parents "$HOME" /app; \
chown --recursive "$UID:$GID" "$HOME" /app
COPY --from=build --chown=$UID:$GID /out/llama-swap /app/llama-swap
COPY --chown=$UID:$GID docker/config.example.yaml /app/config.yaml
USER $UID:$GID
WORKDIR /app
ENV PATH="/app:${PATH}"
HEALTHCHECK CMD curl -f http://localhost:8080/ || exit 1
ENTRYPOINT [ "/app/llama-swap", "-config", "/app/config.yaml" ]
+151
View File
@@ -0,0 +1,151 @@
# Request priority
A GPU is a size-1 resource: one job at a time. Without priorities, a single long
job monopolises the box for its whole duration and every interactive request
queues behind it. A 16-minute render blocking a chat completion is
uncomfortable; a 12-hour batch render blocking everything for a day is not
workable.
Callers declare intent with a header, and the scheduler dispatches by score.
## The header
```
X-LlamaSwap-Priority: 0 # normal — the default
X-LlamaSwap-Priority: 100 # interactive, a human is waiting
X-LlamaSwap-Priority: -100 # batch, nobody is watching
```
Priority is a **signed integer**; higher is more urgent. An absent or
unparseable header means `0`. Values are not clamped — llama-swap trusts the
caller.
Named aliases resolve to numbers for convenience, but **the number is the
interface**: callers may send any value.
| alias | value |
| ------------- | -----: |
| `interactive` | `100` |
| `normal` | `0` |
| `batch` | `-100` |
Priority is a property of the *caller's intent*, not of the model — the same
model serves both an interactive request and a batch job — which is why it
travels as a header rather than as per-model config. It also leaves the
OpenAI-compatible request body untouched.
```bash
curl http://localhost:8080/v1/chat/completions \
-H 'Content-Type: application/json' \
-H 'X-LlamaSwap-Priority: interactive' \
-d '{"model":"qwen3","messages":[{"role":"user","content":"hi"}]}'
```
## Bands and offsets
The three anchors sit 100 apart. That spacing is deliberate: it leaves room for
a caller to add a small per-user offset without ever crossing a band.
| tier | offset |
| ---- | -----: |
| free | `+0` |
| pro | `+1` |
| max | `+2` |
A max member's batch job is `-100 + 2 = -98`: ahead of pro's `-99` and free's
`-100`, and still **below** any normal request at `0`. A tier bonus breaks ties
*inside* a band and can never promote across one.
> **Invariant:** band spacing (100) must stay far larger than the largest offset
> a caller adds plus `swapAffinityBonus`. Otherwise a max member's batch work
> could outrank a free member's interactive request.
## How the serial scheduler dispatches
Scheduling is **non-preemptive**. A generation cannot be interrupted mid-sample
without discarding the work, so priority applies at *dispatch* time: when a job
finishes, the best-scoring waiting request goes next.
```
score = X-LlamaSwap-Priority + swap affinity + aging
```
| term | range | crosses bands? |
| --------------------- | ------- | ----------------------- |
| **request priority** | any | — it *is* the band |
| **swap affinity** | `0..99` | **never** |
| **aging** | `0..∞` | **yes, deliberately** |
Equal scores keep arrival order.
**Swap affinity** is a bounded bonus for a request that can run without a model
swap. Cold loads are expensive, so this keeps a run of same-model requests
together instead of forcing a reload on every dispatch. Because it is bounded
well below the band spacing, it can only break near-ties — it can never promote
batch work above interactive.
**Aging** is `waited_seconds / agingDivisor`, and is unbounded on purpose: a
batch job that has waited long enough *should* eventually beat normal traffic,
or it starves. At the default divisor of 60 (one point per minute), a batch job
at `-100` overtakes normal traffic after ~100 minutes of waiting.
That asymmetry is the design: **bounded offsets express policy, unbounded aging
prevents starvation.**
### What priority does not fix
Because dispatch is non-preemptive, an interactive request can still wait **up
to one full job duration**. The complementary lever is client-side: a batch
producer should submit **one unit at a time** and re-queue, so gaps are frequent.
For a long render that means one shot per submission, not one 12-hour chain.
Priority is also not fair-share. If one caller submits 200 batch jobs they all
sit at the same priority, and llama-swap will work through them in order. Per-user
fairness belongs in the client that submits the work.
## Configuration
```yaml
routing:
scheduler:
use: serial
settings:
serial:
# Seconds of waiting per priority point. 0 disables aging.
agingDivisor: 60
# Bonus for a request needing no model swap. 0..99; 0 disables it.
swapAffinityBonus: 10
```
Setting `swapAffinityBonus: 0` gives strict priority-then-arrival order — with
uniform priorities that is exact FIFO.
The `fifo` scheduler honours the header too, adding it to the model's configured
priority: a model at `10` with a `batch` header queues at `-90`.
## Observability
`GET /metrics` exports the queue in Prometheus text format, labelled by band
(`interactive`, `normal`, `batch`):
| metric | type | meaning |
| --------------------------------------------- | ------- | ---------------------------------------------------- |
| `llamaswap_scheduler_queue_depth` | gauge | requests waiting |
| `llamaswap_scheduler_queue_oldest_wait_seconds`| gauge | how long the longest-waiting request has waited |
| `llamaswap_scheduler_dispatched_total` | counter | requests dispatched |
| `llamaswap_scheduler_reorders_total` | counter | dispatches where `aging` or `swap_affinity` changed the pick |
The reorder counters are what make `agingDivisor` and `swapAffinityBonus`
tunable rather than guesswork: they say how often each term actually changed
which request went next.
Every dispatched request also records its decision in the activity log:
| key | meaning |
| ---------------------- | ------------------------------------------ |
| `serial_priority` | the caller's priority |
| `serial_band` | which band it fell in |
| `serial_score` | the score it won with |
| `serial_queue_wait_ms` | how long it waited before dispatch |
That is what answers "why did my request take 20 minutes" after the fact.
+57
View File
@@ -0,0 +1,57 @@
package config
import (
"fmt"
"runtime"
"strings"
"github.com/billziss-gh/golib/shlex"
)
func SanitizeCommand(cmdStr string) ([]string, error) {
var cleanedLines []string
for _, line := range strings.Split(cmdStr, "\n") {
trimmed := strings.TrimSpace(line)
// Skip comment lines
if strings.HasPrefix(trimmed, "#") {
continue
}
// Handle trailing backslashes by replacing with space
if strings.HasSuffix(trimmed, "\\") {
cleanedLines = append(cleanedLines, strings.TrimSuffix(trimmed, "\\")+" ")
} else {
cleanedLines = append(cleanedLines, line)
}
}
// put it back together
cmdStr = strings.Join(cleanedLines, "\n")
// Split the command into arguments
var args []string
if runtime.GOOS == "windows" {
args = shlex.Windows.Split(cmdStr)
} else {
args = shlex.Posix.Split(cmdStr)
}
// Ensure the command is not empty
if len(args) == 0 {
return nil, fmt.Errorf("empty command")
}
return args, nil
}
func StripComments(cmdStr string) string {
var cleanedLines []string
for _, line := range strings.Split(cmdStr, "\n") {
trimmed := strings.TrimSpace(line)
// Skip comment lines
if strings.HasPrefix(trimmed, "#") {
continue
}
cleanedLines = append(cleanedLines, line)
}
return strings.Join(cleanedLines, "\n")
}
+58 -662
View File
@@ -2,16 +2,9 @@ package config
import (
"fmt"
"io"
"net/url"
"os"
"regexp"
"runtime"
"sort"
"strings"
"time"
"github.com/billziss-gh/golib/shlex"
"gopkg.in/yaml.v3"
)
@@ -85,12 +78,6 @@ type GroupConfig struct {
Members []string `yaml:"members"`
}
var (
macroNameRegex = regexp.MustCompile(`^[a-zA-Z0-9_-]+$`)
macroPatternRegex = regexp.MustCompile(`\$\{([a-zA-Z0-9_-]+)\}`)
envMacroRegex = regexp.MustCompile(`\$\{env\.([a-zA-Z_][a-zA-Z0-9_]*)\}`)
)
// set default values for GroupConfig
func (c *GroupConfig) UnmarshalYAML(unmarshal func(interface{}) error) error {
type rawGroupConfig GroupConfig
@@ -163,6 +150,9 @@ type Config struct {
// support remote peers, see issue #433, #296
Peers PeerDictionaryConfig `yaml:"peers"`
// upstream controls behaviour of the /upstream passthrough endpoint
Upstream UpstreamConfig `yaml:"upstream"`
}
// RoutingConfig is the canonical, normalized routing/scheduling configuration.
@@ -177,13 +167,67 @@ type SchedulerConfig struct {
}
type SchedulerSettings struct {
Fifo FifoConfig `yaml:"fifo"`
Fifo FifoConfig `yaml:"fifo"`
Serial SerialConfig `yaml:"serial"`
}
type FifoConfig struct {
Priority map[string]int `yaml:"priority"` // model ID -> priority, default 0
}
// Serial scheduler scoring defaults. See SerialConfig.
const (
// DefaultAgingDivisor gives a waiting request one priority point per
// minute, so a batch job at -100 overtakes normal traffic after ~100
// minutes rather than starving behind it forever.
DefaultAgingDivisor = 60
// DefaultSwapAffinityBonus is large enough to break near-ties in favour
// of the already-loaded model (cold loads cost 2-3x) and far too small to
// promote a request across a 100-point priority band.
DefaultSwapAffinityBonus = 10
)
// SerialConfig tunes how the serial scheduler scores queued requests. At each
// dispatch it picks the highest scoring waiting request, where
//
// score = request priority + swap affinity + aging
//
// Both terms are pointers so an explicit 0 (disable) is distinguishable from
// "unset" (use the default); read them through the accessors below.
type SerialConfig struct {
// AgingDivisor is how many seconds a request must wait to gain one
// priority point. Aging is unbounded on purpose — it is the only term
// allowed to promote a request across a band, which is what stops a
// low-priority job from starving. 0 disables aging.
AgingDivisor *int `yaml:"agingDivisor"`
// SwapAffinityBonus is added to a request that can be served without
// swapping models, so strict priority does not force a cold load on every
// dispatch. Bounded well below the 100-point band spacing so it can only
// break near-ties, never promote batch work above interactive. 0 disables
// it, restoring strict priority-then-arrival order.
SwapAffinityBonus *int `yaml:"swapAffinityBonus"`
}
// GetAgingDivisor returns the configured aging divisor in seconds, or
// DefaultAgingDivisor when unset.
func (c SerialConfig) GetAgingDivisor() int {
if c.AgingDivisor == nil {
return DefaultAgingDivisor
}
return *c.AgingDivisor
}
// GetSwapAffinityBonus returns the configured swap affinity bonus, or
// DefaultSwapAffinityBonus when unset.
func (c SerialConfig) GetSwapAffinityBonus() int {
if c.SwapAffinityBonus == nil {
return DefaultSwapAffinityBonus
}
return *c.SwapAffinityBonus
}
type RouterConfig struct {
Use string `yaml:"use"` // "group" (default) | "matrix"
Settings RouterSettings `yaml:"settings"`
@@ -221,424 +265,6 @@ func LoadConfig(path string) (Config, error) {
return LoadConfigFromReader(file)
}
func LoadConfigFromReader(r io.Reader) (Config, error) {
data, err := io.ReadAll(r)
if err != nil {
return Config{}, err
}
yamlStr := string(data)
// Phase 1: Substitute all ${env.VAR} macros at string level
// This is safe because env values are simple strings without YAML formatting
yamlStr, err = substituteEnvMacros(yamlStr)
if err != nil {
return Config{}, err
}
// Unmarshal into full Config with defaults
config := Config{
HealthCheckTimeout: 120,
StartPort: 5800,
LogLevel: "info",
LogTimeFormat: "",
LogToStdout: LogToStdoutProxy,
MetricsMaxInMemory: 1000,
CaptureBuffer: 5,
GlobalTTL: 0,
}
if err = yaml.Unmarshal([]byte(yamlStr), &config); err != nil {
return Config{}, err
}
if config.HealthCheckTimeout < 15 {
config.HealthCheckTimeout = 15
}
// Apply defaults for performance config when section is missing
if config.Performance.Every == 0 {
config.Performance.Every = 5 * time.Second
}
if err = config.Performance.Validate(); err != nil {
return Config{}, fmt.Errorf("performance: %w", err)
}
if config.StartPort < 1 {
return Config{}, fmt.Errorf("startPort must be greater than 1")
}
if config.GlobalTTL < 0 {
return Config{}, fmt.Errorf("globalTTL must be >= 0")
}
switch config.LogToStdout {
case LogToStdoutProxy, LogToStdoutUpstream, LogToStdoutBoth, LogToStdoutNone:
default:
return Config{}, fmt.Errorf("logToStdout must be one of: proxy, upstream, both, none")
}
// Populate the aliases map
config.aliases = make(map[string]string)
for modelName, modelConfig := range config.Models {
for _, alias := range modelConfig.Aliases {
if _, found := config.aliases[alias]; found {
return Config{}, fmt.Errorf("duplicate alias %s found in model: %s", alias, modelName)
}
config.aliases[alias] = modelName
}
}
// Validate global macros
for _, macro := range config.Macros {
if err = validateMacro(macro.Name, macro.Value); err != nil {
return Config{}, err
}
}
// Get and sort all model IDs for consistent port assignment
modelIds := make([]string, 0, len(config.Models))
for modelId := range config.Models {
modelIds = append(modelIds, modelId)
}
sort.Strings(modelIds)
nextPort := config.StartPort
for _, modelId := range modelIds {
modelConfig := config.Models[modelId]
modelConfig.HealthCheckTimeout = config.HealthCheckTimeout
// Strip comments from command fields
modelConfig.Cmd = StripComments(modelConfig.Cmd)
modelConfig.CmdStop = StripComments(modelConfig.CmdStop)
// set model TTL to globalTTL it is the default value
if modelConfig.UnloadAfter == MODEL_CONFIG_DEFAULT_TTL {
modelConfig.UnloadAfter = config.GlobalTTL
}
if modelConfig.UnloadAfter < 0 {
return Config{}, fmt.Errorf("model %s: invalid TTL value %d", modelId, modelConfig.UnloadAfter)
}
// Validate model macros
for _, macro := range modelConfig.Macros {
if err = validateMacro(macro.Name, macro.Value); err != nil {
return Config{}, fmt.Errorf("model %s: %s", modelId, err.Error())
}
}
// Build merged macro list: MODEL_ID + global macros + model macros (model overrides global)
mergedMacros := make(MacroList, 0, len(config.Macros)+len(modelConfig.Macros)+1)
mergedMacros = append(mergedMacros, MacroEntry{Name: "MODEL_ID", Value: modelId})
mergedMacros = append(mergedMacros, config.Macros...)
// Add model macros (override globals with same name)
for _, entry := range modelConfig.Macros {
found := false
for i, existing := range mergedMacros {
if existing.Name == entry.Name {
mergedMacros[i] = entry
found = true
break
}
}
if !found {
mergedMacros = append(mergedMacros, entry)
}
}
// Substitute remaining macros in model fields (LIFO order)
for i := len(mergedMacros) - 1; i >= 0; i-- {
entry := mergedMacros[i]
macroSlug := fmt.Sprintf("${%s}", entry.Name)
macroStr := fmt.Sprintf("%v", entry.Value)
modelConfig.Cmd = strings.ReplaceAll(modelConfig.Cmd, macroSlug, macroStr)
modelConfig.CmdStop = strings.ReplaceAll(modelConfig.CmdStop, macroSlug, macroStr)
modelConfig.Proxy = strings.ReplaceAll(modelConfig.Proxy, macroSlug, macroStr)
modelConfig.CheckEndpoint = strings.ReplaceAll(modelConfig.CheckEndpoint, macroSlug, macroStr)
modelConfig.Filters.StripParams = strings.ReplaceAll(modelConfig.Filters.StripParams, macroSlug, macroStr)
modelConfig.Name = strings.ReplaceAll(modelConfig.Name, macroSlug, macroStr)
modelConfig.Description = strings.ReplaceAll(modelConfig.Description, macroSlug, macroStr)
// Substitute macros in SetParamsByID keys and values
if len(modelConfig.Filters.SetParamsByID) > 0 {
newSetParamsByID := make(map[string]map[string]any, len(modelConfig.Filters.SetParamsByID))
for key, paramMap := range modelConfig.Filters.SetParamsByID {
newKey := strings.ReplaceAll(key, macroSlug, macroStr)
newValAny, err := substituteMacroInValue(any(paramMap), entry.Name, entry.Value)
if err != nil {
return Config{}, fmt.Errorf("model %s filters.setParamsByID: %s", modelId, err.Error())
}
newParamMap, ok := newValAny.(map[string]any)
if !ok {
return Config{}, fmt.Errorf("model %s filters.setParamsByID: unexpected type after macro substitution", modelId)
}
newSetParamsByID[newKey] = newParamMap
}
modelConfig.Filters.SetParamsByID = newSetParamsByID
}
// Substitute in metadata (type-preserving)
if len(modelConfig.Metadata) > 0 {
result, err := substituteMacroInValue(modelConfig.Metadata, entry.Name, entry.Value)
if err != nil {
return Config{}, fmt.Errorf("model %s metadata: %s", modelId, err.Error())
}
modelConfig.Metadata = result.(map[string]any)
}
}
// Handle PORT macro - only allocate if cmd uses it
cmdHasPort := strings.Contains(modelConfig.Cmd, "${PORT}")
proxyHasPort := strings.Contains(modelConfig.Proxy, "${PORT}")
if cmdHasPort || proxyHasPort {
if !cmdHasPort && proxyHasPort {
return Config{}, fmt.Errorf("model %s: proxy uses ${PORT} but cmd does not - ${PORT} is only available when used in cmd", modelId)
}
macroSlug := "${PORT}"
macroStr := fmt.Sprintf("%v", nextPort)
modelConfig.Cmd = strings.ReplaceAll(modelConfig.Cmd, macroSlug, macroStr)
modelConfig.CmdStop = strings.ReplaceAll(modelConfig.CmdStop, macroSlug, macroStr)
modelConfig.Proxy = strings.ReplaceAll(modelConfig.Proxy, macroSlug, macroStr)
modelConfig.Name = strings.ReplaceAll(modelConfig.Name, macroSlug, macroStr)
modelConfig.Description = strings.ReplaceAll(modelConfig.Description, macroSlug, macroStr)
if len(modelConfig.Metadata) > 0 {
result, err := substituteMacroInValue(modelConfig.Metadata, "PORT", nextPort)
if err != nil {
return Config{}, fmt.Errorf("model %s metadata: %s", modelId, err.Error())
}
modelConfig.Metadata = result.(map[string]any)
}
nextPort++
}
// Validate no unknown macros remain
fieldMap := map[string]string{
"cmd": modelConfig.Cmd,
"cmdStop": modelConfig.CmdStop,
"proxy": modelConfig.Proxy,
"checkEndpoint": modelConfig.CheckEndpoint,
"filters.stripParams": modelConfig.Filters.StripParams,
"name": modelConfig.Name,
"description": modelConfig.Description,
}
for fieldName, fieldValue := range fieldMap {
matches := macroPatternRegex.FindAllStringSubmatch(fieldValue, -1)
for _, match := range matches {
macroName := match[1]
if macroName == "PID" && fieldName == "cmdStop" {
continue // replaced at runtime
}
if macroName == "PORT" || macroName == "MODEL_ID" {
return Config{}, fmt.Errorf("macro '${%s}' should have been substituted in %s.%s", macroName, modelId, fieldName)
}
return Config{}, fmt.Errorf("unknown macro '${%s}' found in %s.%s", macroName, modelId, fieldName)
}
}
if len(modelConfig.Metadata) > 0 {
if err := validateNestedForUnknownMacros(modelConfig.Metadata, fmt.Sprintf("model %s metadata", modelId)); err != nil {
return Config{}, err
}
}
if err = modelConfig.Capabilities.Validate(); err != nil {
return Config{}, fmt.Errorf("model %s: %w", modelId, err)
}
// Validate SetParamsByID keys and values
for key, paramMap := range modelConfig.Filters.SetParamsByID {
if matches := macroPatternRegex.FindAllStringSubmatch(key, -1); len(matches) > 0 {
return Config{}, fmt.Errorf("unknown macro '${%s}' found in model %s filters.setParamsByID key", matches[0][1], modelId)
}
if err := validateNestedForUnknownMacros(any(paramMap), fmt.Sprintf("model %s filters.setParamsByID[%s]", modelId, key)); err != nil {
return Config{}, err
}
}
// Auto-register setParamsByID keys as aliases (skip the model's own ID)
for key := range modelConfig.Filters.SetParamsByID {
if key == modelId {
continue
}
if _, exists := config.Models[key]; exists {
return Config{}, fmt.Errorf("model %s filters.setParamsByID: key '%s' conflicts with an existing model ID", modelId, key)
}
if existingModel, exists := config.aliases[key]; exists {
if existingModel != modelId {
return Config{}, fmt.Errorf("duplicate alias '%s' in model %s filters.setParamsByID, already used by model %s", key, modelId, existingModel)
}
continue // already registered as explicit alias for this model
}
config.aliases[key] = modelId
modelConfig.Aliases = append(modelConfig.Aliases, key)
}
if _, err := url.Parse(modelConfig.Proxy); err != nil {
return Config{}, fmt.Errorf("model %s: invalid proxy URL: %w", modelId, err)
}
if modelConfig.SendLoadingState == nil {
v := config.SendLoadingState
modelConfig.SendLoadingState = &v
}
config.Models[modelId] = modelConfig
}
// Normalize routing config. The legacy top-level `matrix`/`groups` keys and
// the new `routing.router` block are mutually exclusive: a config may use
// either style, never both.
hasTopLevel := config.Matrix != nil || len(config.Groups) > 0
rtr := config.Routing.Router
hasRouting := rtr.Use != "" || rtr.Settings.Matrix != nil || len(rtr.Settings.Groups) > 0
if hasTopLevel && hasRouting {
return Config{}, fmt.Errorf("config uses both the legacy top-level 'matrix'/'groups' keys and the new 'routing.router' block; please migrate the top-level keys into 'routing.router' and remove them")
}
if !hasTopLevel {
// Both groups and matrix may be defined under routing.router.settings;
// routing.router.use selects which one is active, so there is no conflict.
rs := config.Routing.Router.Settings
switch config.Routing.Router.Use {
case "matrix":
if rs.Matrix == nil {
return Config{}, fmt.Errorf("routing.router.use is 'matrix' but routing.router.settings.matrix is not set")
}
config.Matrix = rs.Matrix
case "group", "":
config.Groups = rs.Groups
default:
return Config{}, fmt.Errorf("routing.router.use: unknown router %q (valid: group, matrix)", config.Routing.Router.Use)
}
}
// groups XOR matrix
if config.Matrix != nil && len(config.Groups) > 0 {
return Config{}, fmt.Errorf("config cannot use both 'groups' and 'matrix'")
}
if config.Matrix != nil {
expandedSets, err := ValidateMatrix(*config.Matrix, config.Models)
if err != nil {
return Config{}, fmt.Errorf("matrix: %w", err)
}
config.Matrix.ExpandedSets = expandedSets
} else {
config = AddDefaultGroupToConfig(config)
// Validate group members
memberUsage := make(map[string]string)
for groupID, groupConfig := range config.Groups {
prevSet := make(map[string]bool)
for _, member := range groupConfig.Members {
if _, found := prevSet[member]; found {
return Config{}, fmt.Errorf("duplicate model member %s found in group: %s", member, groupID)
}
prevSet[member] = true
if existingGroup, exists := memberUsage[member]; exists {
return Config{}, fmt.Errorf("model member %s is used in multiple groups: %s and %s", member, existingGroup, groupID)
}
memberUsage[member] = groupID
}
}
}
// Build the canonical Config.Routing from the effective result. Both legacy
// and new-style configs converge here. The Matrix pointer is shared so
// ExpandedSets stays in one place.
if config.Matrix != nil {
config.Routing.Router.Use = "matrix"
} else {
config.Routing.Router.Use = "group"
}
config.Routing.Router.Settings.Matrix = config.Matrix
config.Routing.Router.Settings.Groups = config.Groups
if config.Routing.Scheduler.Use == "" {
config.Routing.Scheduler.Use = "fifo"
}
if config.Routing.Scheduler.Use != "fifo" {
return Config{}, fmt.Errorf("routing.scheduler.use: unknown scheduler %q (valid: fifo)", config.Routing.Scheduler.Use)
}
for modelID := range config.Routing.Scheduler.Settings.Fifo.Priority {
if _, found := config.RealModelName(modelID); !found {
return Config{}, fmt.Errorf("routing.scheduler.settings.fifo.priority references unknown model %q", modelID)
}
}
// Clean up hooks preload
if len(config.Hooks.OnStartup.Preload) > 0 {
var toPreload []string
for _, modelID := range config.Hooks.OnStartup.Preload {
modelID = strings.TrimSpace(modelID)
if modelID == "" {
continue
}
if real, found := config.RealModelName(modelID); found {
toPreload = append(toPreload, real)
}
}
config.Hooks.OnStartup.Preload = toPreload
}
// Validate API keys (env macros already substituted at string level)
for i, apikey := range config.RequiredAPIKeys {
if apikey == "" {
return Config{}, fmt.Errorf("empty api key found in apiKeys")
}
if strings.Contains(apikey, " ") {
return Config{}, fmt.Errorf("api key cannot contain spaces: `%s`", apikey)
}
config.RequiredAPIKeys[i] = apikey
}
// Process peers with global macro substitution
for peerName, peerConfig := range config.Peers {
// Substitute global macros (LIFO order)
for i := len(config.Macros) - 1; i >= 0; i-- {
entry := config.Macros[i]
macroSlug := fmt.Sprintf("${%s}", entry.Name)
macroStr := fmt.Sprintf("%v", entry.Value)
peerConfig.ApiKey = strings.ReplaceAll(peerConfig.ApiKey, macroSlug, macroStr)
peerConfig.Filters.StripParams = strings.ReplaceAll(peerConfig.Filters.StripParams, macroSlug, macroStr)
// Substitute in setParams (type-preserving)
if len(peerConfig.Filters.SetParams) > 0 {
result, err := substituteMacroInValue(peerConfig.Filters.SetParams, entry.Name, entry.Value)
if err != nil {
return Config{}, fmt.Errorf("peers.%s.filters.setParams: %w", peerName, err)
}
peerConfig.Filters.SetParams = result.(map[string]any)
}
}
// Validate no unknown macros remain
if matches := macroPatternRegex.FindAllStringSubmatch(peerConfig.ApiKey, -1); len(matches) > 0 {
return Config{}, fmt.Errorf("peers.%s.apiKey: unknown macro '${%s}'", peerName, matches[0][1])
}
if matches := macroPatternRegex.FindAllStringSubmatch(peerConfig.Filters.StripParams, -1); len(matches) > 0 {
return Config{}, fmt.Errorf("peers.%s.filters.stripParams: unknown macro '${%s}'", peerName, matches[0][1])
}
if len(peerConfig.Filters.SetParams) > 0 {
if err := validateNestedForUnknownMacros(peerConfig.Filters.SetParams, fmt.Sprintf("peers.%s.filters.setParams", peerName)); err != nil {
return Config{}, err
}
}
config.Peers[peerName] = peerConfig
}
return config, nil
}
// rewrites the yaml to include a default group with any orphaned models
func AddDefaultGroupToConfig(config Config) Config {
@@ -683,233 +309,3 @@ func AddDefaultGroupToConfig(config Config) Config {
return config
}
func SanitizeCommand(cmdStr string) ([]string, error) {
var cleanedLines []string
for _, line := range strings.Split(cmdStr, "\n") {
trimmed := strings.TrimSpace(line)
// Skip comment lines
if strings.HasPrefix(trimmed, "#") {
continue
}
// Handle trailing backslashes by replacing with space
if strings.HasSuffix(trimmed, "\\") {
cleanedLines = append(cleanedLines, strings.TrimSuffix(trimmed, "\\")+" ")
} else {
cleanedLines = append(cleanedLines, line)
}
}
// put it back together
cmdStr = strings.Join(cleanedLines, "\n")
// Split the command into arguments
var args []string
if runtime.GOOS == "windows" {
args = shlex.Windows.Split(cmdStr)
} else {
args = shlex.Posix.Split(cmdStr)
}
// Ensure the command is not empty
if len(args) == 0 {
return nil, fmt.Errorf("empty command")
}
return args, nil
}
func StripComments(cmdStr string) string {
var cleanedLines []string
for _, line := range strings.Split(cmdStr, "\n") {
trimmed := strings.TrimSpace(line)
// Skip comment lines
if strings.HasPrefix(trimmed, "#") {
continue
}
cleanedLines = append(cleanedLines, line)
}
return strings.Join(cleanedLines, "\n")
}
// validateMacro validates macro name and value constraints
func validateMacro(name string, value any) error {
if len(name) >= 64 {
return fmt.Errorf("macro name '%s' exceeds maximum length of 63 characters", name)
}
if !macroNameRegex.MatchString(name) {
return fmt.Errorf("macro name '%s' contains invalid characters, must match pattern ^[a-zA-Z0-9_-]+$", name)
}
// Validate that value is a scalar type
switch v := value.(type) {
case string:
// Check for self-reference
macroSlug := fmt.Sprintf("${%s}", name)
if strings.Contains(v, macroSlug) {
return fmt.Errorf("macro '%s' contains self-reference", name)
}
case int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64, float32, float64, bool:
// These types are allowed
default:
return fmt.Errorf("macro '%s' has invalid type %T, must be a scalar type (string, int, float, or bool)", name, value)
}
switch name {
case "PORT", "MODEL_ID":
return fmt.Errorf("macro name '%s' is reserved", name)
}
return nil
}
// validateNestedForUnknownMacros recursively checks for any remaining macro references in nested structures
func validateNestedForUnknownMacros(value any, context string) error {
switch v := value.(type) {
case string:
matches := macroPatternRegex.FindAllStringSubmatch(v, -1)
for _, match := range matches {
macroName := match[1]
return fmt.Errorf("%s: unknown macro '${%s}'", context, macroName)
}
// Check for unsubstituted env macros
envMatches := envMacroRegex.FindAllStringSubmatch(v, -1)
for _, match := range envMatches {
varName := match[1]
return fmt.Errorf("%s: environment variable '%s' not set", context, varName)
}
return nil
case map[string]any:
for _, val := range v {
if err := validateNestedForUnknownMacros(val, context); err != nil {
return err
}
}
return nil
case []any:
for _, val := range v {
if err := validateNestedForUnknownMacros(val, context); err != nil {
return err
}
}
return nil
default:
// Scalar types don't contain macros
return nil
}
}
// substituteMacroInValue recursively substitutes a single macro in a value structure
// This is called once per macro, allowing LIFO substitution order
func substituteMacroInValue(value any, macroName string, macroValue any) (any, error) {
macroSlug := fmt.Sprintf("${%s}", macroName)
macroStr := fmt.Sprintf("%v", macroValue)
switch v := value.(type) {
case string:
// Check if this is a direct macro substitution
if v == macroSlug {
return macroValue, nil
}
// Handle string interpolation
if strings.Contains(v, macroSlug) {
return strings.ReplaceAll(v, macroSlug, macroStr), nil
}
return v, nil
case map[string]any:
// Recursively process map values
newMap := make(map[string]any)
for key, val := range v {
newVal, err := substituteMacroInValue(val, macroName, macroValue)
if err != nil {
return nil, err
}
newMap[key] = newVal
}
return newMap, nil
case []any:
// Recursively process slice elements
newSlice := make([]any, len(v))
for i, val := range v {
newVal, err := substituteMacroInValue(val, macroName, macroValue)
if err != nil {
return nil, err
}
newSlice[i] = newVal
}
return newSlice, nil
default:
// Return scalar types as-is
return value, nil
}
}
// substituteEnvMacros replaces ${env.VAR_NAME} with environment variable values.
// Returns error if any referenced env var is not set or contains invalid characters.
// Env macros inside YAML comments are ignored by unmarshalling the YAML first
// (which strips comments) and only checking the comment-free version for macros.
func substituteEnvMacros(s string) (string, error) {
// Unmarshal and remarshal to strip YAML comments
var raw any
if err := yaml.Unmarshal([]byte(s), &raw); err != nil {
// If YAML is invalid, fall back to scanning the original string
// so the user gets the env var error rather than a confusing YAML parse error
return substituteEnvMacrosInString(s, s)
}
clean, err := yaml.Marshal(raw)
if err != nil {
return substituteEnvMacrosInString(s, s)
}
return substituteEnvMacrosInString(s, string(clean))
}
// substituteEnvMacrosInString finds ${env.VAR} macros in scanStr and substitutes
// them in target. This separation allows scanning comment-free YAML while
// substituting in the original string.
func substituteEnvMacrosInString(target, scanStr string) (string, error) {
result := target
matches := envMacroRegex.FindAllStringSubmatch(scanStr, -1)
for _, match := range matches {
fullMatch := match[0] // ${env.VAR_NAME}
varName := match[1] // VAR_NAME
value, exists := os.LookupEnv(varName)
if !exists {
return "", fmt.Errorf("environment variable '%s' is not set", varName)
}
// Sanitize the value for safe YAML substitution
value, err := sanitizeEnvValueForYAML(value, varName)
if err != nil {
return "", err
}
result = strings.ReplaceAll(result, fullMatch, value)
}
return result, nil
}
// sanitizeEnvValueForYAML ensures an environment variable value is safe for YAML substitution.
// It rejects values with characters that break YAML structure and escapes quotes/backslashes
// for compatibility with double-quoted YAML strings.
func sanitizeEnvValueForYAML(value, varName string) (string, error) {
// Reject values that would break YAML structure regardless of quoting context
if strings.ContainsAny(value, "\n\r\x00") {
return "", fmt.Errorf("environment variable '%s' contains newlines or null bytes which are not allowed in YAML substitution", varName)
}
// Escape backslashes and double quotes for safe use in double-quoted YAML strings.
// In unquoted contexts, these escapes appear literally (harmless for most use cases).
// In double-quoted contexts, they are interpreted correctly.
value = strings.ReplaceAll(value, `\`, `\\`)
value = strings.ReplaceAll(value, `"`, `\"`)
return value, nil
}
+4 -1
View File
@@ -266,6 +266,9 @@ groups:
"mthree": "model3",
},
Groups: expectedGroups,
Upstream: UpstreamConfig{
IgnorePaths: DefaultUpstreamIgnorePaths(),
},
Routing: RoutingConfig{
Router: RouterConfig{
Use: "group",
@@ -274,7 +277,7 @@ groups:
},
},
Scheduler: SchedulerConfig{
Use: "fifo",
Use: "serial",
},
},
}
+85 -6
View File
@@ -777,22 +777,27 @@ func TestConfig_APIKeys_Invalid(t *testing.T) {
{
name: "blank spaces only",
content: `apiKeys: [" "]`,
expectedErr: "api key cannot contain spaces: ` `",
expectedErr: "apiKeys[0]: api key cannot contain spaces",
},
{
name: "contains leading space",
content: `apiKeys: [" key123"]`,
expectedErr: "api key cannot contain spaces: ` key123`",
expectedErr: "apiKeys[0]: api key cannot contain spaces",
},
{
name: "contains trailing space",
content: `apiKeys: ["key123 "]`,
expectedErr: "api key cannot contain spaces: `key123 `",
expectedErr: "apiKeys[0]: api key cannot contain spaces",
},
{
name: "contains middle space",
content: `apiKeys: ["key 123"]`,
expectedErr: "api key cannot contain spaces: `key 123`",
expectedErr: "apiKeys[0]: api key cannot contain spaces",
},
{
name: "space in second key reports correct index",
content: `apiKeys: ["valid-key", "bad key"]`,
expectedErr: "apiKeys[1]: api key cannot contain spaces",
},
{
name: "empty in list with valid keys",
@@ -1567,7 +1572,7 @@ groups:
assert.Equal(t, "group", cfg.Routing.Router.Use)
// default group injected for orphaned models (none here) still leaves g1
assert.Contains(t, cfg.Routing.Router.Settings.Groups, "g1")
assert.Equal(t, "fifo", cfg.Routing.Scheduler.Use)
assert.Equal(t, "serial", cfg.Routing.Scheduler.Use)
}
func TestConfig_Routing_LegacyTopLevelMatrix(t *testing.T) {
@@ -1626,7 +1631,7 @@ func TestConfig_Routing_DefaultsToGroup(t *testing.T) {
cfg, err := LoadConfigFromReader(strings.NewReader(twoModels))
require.NoError(t, err)
assert.Equal(t, "group", cfg.Routing.Router.Use)
assert.Equal(t, "fifo", cfg.Routing.Scheduler.Use)
assert.Equal(t, "serial", cfg.Routing.Scheduler.Use)
}
func TestConfig_Routing_LegacyAndRoutingConflict(t *testing.T) {
@@ -1715,3 +1720,77 @@ routing:
require.NoError(t, err)
assert.Equal(t, 5, cfg.Routing.Scheduler.Settings.Fifo.Priority["gemma"])
}
func TestConfig_Routing_SerialDefaults(t *testing.T) {
cfg, err := LoadConfigFromReader(strings.NewReader(twoModels))
require.NoError(t, err)
serial := cfg.Routing.Scheduler.Settings.Serial
assert.Nil(t, serial.AgingDivisor, "unset should stay nil so the default applies")
assert.Equal(t, DefaultAgingDivisor, serial.GetAgingDivisor())
assert.Equal(t, DefaultSwapAffinityBonus, serial.GetSwapAffinityBonus())
}
// TestConfig_Routing_SerialExplicitZero verifies an explicit 0 disables a term
// rather than falling back to the default — the reason both fields are pointers.
func TestConfig_Routing_SerialExplicitZero(t *testing.T) {
yaml := twoModels + `
routing:
scheduler:
use: serial
settings:
serial:
agingDivisor: 0
swapAffinityBonus: 0
`
cfg, err := LoadConfigFromReader(strings.NewReader(yaml))
require.NoError(t, err)
serial := cfg.Routing.Scheduler.Settings.Serial
assert.Equal(t, 0, serial.GetAgingDivisor())
assert.Equal(t, 0, serial.GetSwapAffinityBonus())
}
func TestConfig_Routing_SerialSettings(t *testing.T) {
yaml := twoModels + `
routing:
scheduler:
use: serial
settings:
serial:
agingDivisor: 30
swapAffinityBonus: 5
`
cfg, err := LoadConfigFromReader(strings.NewReader(yaml))
require.NoError(t, err)
serial := cfg.Routing.Scheduler.Settings.Serial
assert.Equal(t, 30, serial.GetAgingDivisor())
assert.Equal(t, 5, serial.GetSwapAffinityBonus())
}
func TestConfig_Routing_SerialInvalidSettings(t *testing.T) {
cases := []struct {
name string
settings string
wantErr string
}{
{"negative aging divisor", "agingDivisor: -1", "agingDivisor"},
{"negative affinity bonus", "swapAffinityBonus: -1", "swapAffinityBonus"},
// 100 would let the bonus promote a request into the next band.
{"affinity bonus reaches a band", "swapAffinityBonus: 100", "swapAffinityBonus"},
}
for _, c := range cases {
t.Run(c.name, func(t *testing.T) {
yaml := twoModels + `
routing:
scheduler:
settings:
serial:
` + c.settings + "\n"
_, err := LoadConfigFromReader(strings.NewReader(yaml))
require.Error(t, err)
assert.Contains(t, err.Error(), c.wantErr)
})
}
}
+4 -1
View File
@@ -255,6 +255,9 @@ groups:
"mthree": "model3",
},
Groups: expectedGroups,
Upstream: UpstreamConfig{
IgnorePaths: DefaultUpstreamIgnorePaths(),
},
Routing: RoutingConfig{
Router: RouterConfig{
Use: "group",
@@ -263,7 +266,7 @@ groups:
},
},
Scheduler: SchedulerConfig{
Use: "fifo",
Use: "serial",
},
},
}
+451
View File
@@ -0,0 +1,451 @@
package config
import (
"fmt"
"io"
"net/url"
"sort"
"strings"
"time"
"gopkg.in/yaml.v3"
)
func LoadConfigFromReader(r io.Reader) (Config, error) {
data, err := io.ReadAll(r)
if err != nil {
return Config{}, err
}
yamlStr := string(data)
// Phase 1: Substitute all ${env.VAR} macros at string level
// This is safe because env values are simple strings without YAML formatting
yamlStr, err = substituteEnvMacros(yamlStr)
if err != nil {
return Config{}, err
}
// Unmarshal into full Config with defaults
config := Config{
HealthCheckTimeout: 120,
StartPort: 5800,
LogLevel: "info",
LogTimeFormat: "",
LogToStdout: LogToStdoutProxy,
MetricsMaxInMemory: 1000,
CaptureBuffer: 5,
GlobalTTL: 0,
}
if err = yaml.Unmarshal([]byte(yamlStr), &config); err != nil {
return Config{}, err
}
if config.HealthCheckTimeout < 15 {
config.HealthCheckTimeout = 15
}
// Apply defaults for performance config when section is missing
if config.Performance.Every == 0 {
config.Performance.Every = 5 * time.Second
}
if err = config.Performance.Validate(); err != nil {
return Config{}, fmt.Errorf("performance: %w", err)
}
if config.StartPort < 1 {
return Config{}, fmt.Errorf("startPort must be greater than 1")
}
if config.GlobalTTL < 0 {
return Config{}, fmt.Errorf("globalTTL must be >= 0")
}
// Apply default for upstream.ignorePaths when not specified. The default
// matches common static-asset suffixes so they do not trigger a swap.
if len(config.Upstream.IgnorePaths) == 0 {
config.Upstream.IgnorePaths = DefaultUpstreamIgnorePaths()
}
switch config.LogToStdout {
case LogToStdoutProxy, LogToStdoutUpstream, LogToStdoutBoth, LogToStdoutNone:
default:
return Config{}, fmt.Errorf("logToStdout must be one of: proxy, upstream, both, none")
}
// Populate the aliases map
config.aliases = make(map[string]string)
for modelName, modelConfig := range config.Models {
for _, alias := range modelConfig.Aliases {
if _, found := config.aliases[alias]; found {
return Config{}, fmt.Errorf("duplicate alias %s found in model: %s", alias, modelName)
}
config.aliases[alias] = modelName
}
}
// Validate global macros
for _, macro := range config.Macros {
if err = validateMacro(macro.Name, macro.Value); err != nil {
return Config{}, err
}
}
// Get and sort all model IDs for consistent port assignment
modelIds := make([]string, 0, len(config.Models))
for modelId := range config.Models {
modelIds = append(modelIds, modelId)
}
sort.Strings(modelIds)
nextPort := config.StartPort
for _, modelId := range modelIds {
modelConfig := config.Models[modelId]
modelConfig.HealthCheckTimeout = config.HealthCheckTimeout
// Strip comments from command fields
modelConfig.Cmd = StripComments(modelConfig.Cmd)
modelConfig.CmdStop = StripComments(modelConfig.CmdStop)
// set model TTL to globalTTL it is the default value
if modelConfig.UnloadAfter == MODEL_CONFIG_DEFAULT_TTL {
modelConfig.UnloadAfter = config.GlobalTTL
}
if modelConfig.UnloadAfter < 0 {
return Config{}, fmt.Errorf("model %s: invalid TTL value %d", modelId, modelConfig.UnloadAfter)
}
// Validate model macros
for _, macro := range modelConfig.Macros {
if err = validateMacro(macro.Name, macro.Value); err != nil {
return Config{}, fmt.Errorf("model %s: %s", modelId, err.Error())
}
}
// Build merged macro list: MODEL_ID + global macros + model macros (model overrides global)
mergedMacros := make(MacroList, 0, len(config.Macros)+len(modelConfig.Macros)+1)
mergedMacros = append(mergedMacros, MacroEntry{Name: "MODEL_ID", Value: modelId})
mergedMacros = append(mergedMacros, config.Macros...)
// Add model macros (override globals with same name)
for _, entry := range modelConfig.Macros {
found := false
for i, existing := range mergedMacros {
if existing.Name == entry.Name {
mergedMacros[i] = entry
found = true
break
}
}
if !found {
mergedMacros = append(mergedMacros, entry)
}
}
// Substitute remaining macros in model fields (LIFO order)
for i := len(mergedMacros) - 1; i >= 0; i-- {
entry := mergedMacros[i]
macroSlug := fmt.Sprintf("${%s}", entry.Name)
macroStr := fmt.Sprintf("%v", entry.Value)
modelConfig.Cmd = strings.ReplaceAll(modelConfig.Cmd, macroSlug, macroStr)
modelConfig.CmdStop = strings.ReplaceAll(modelConfig.CmdStop, macroSlug, macroStr)
modelConfig.Proxy = strings.ReplaceAll(modelConfig.Proxy, macroSlug, macroStr)
modelConfig.CheckEndpoint = strings.ReplaceAll(modelConfig.CheckEndpoint, macroSlug, macroStr)
modelConfig.Filters.StripParams = strings.ReplaceAll(modelConfig.Filters.StripParams, macroSlug, macroStr)
modelConfig.Name = strings.ReplaceAll(modelConfig.Name, macroSlug, macroStr)
modelConfig.Description = strings.ReplaceAll(modelConfig.Description, macroSlug, macroStr)
// Substitute macros in SetParamsByID keys and values
if len(modelConfig.Filters.SetParamsByID) > 0 {
newSetParamsByID := make(map[string]map[string]any, len(modelConfig.Filters.SetParamsByID))
for key, paramMap := range modelConfig.Filters.SetParamsByID {
newKey := strings.ReplaceAll(key, macroSlug, macroStr)
newValAny, err := substituteMacroInValue(any(paramMap), entry.Name, entry.Value)
if err != nil {
return Config{}, fmt.Errorf("model %s filters.setParamsByID: %s", modelId, err.Error())
}
newParamMap, ok := newValAny.(map[string]any)
if !ok {
return Config{}, fmt.Errorf("model %s filters.setParamsByID: unexpected type after macro substitution", modelId)
}
newSetParamsByID[newKey] = newParamMap
}
modelConfig.Filters.SetParamsByID = newSetParamsByID
}
// Substitute in metadata (type-preserving)
if len(modelConfig.Metadata) > 0 {
result, err := substituteMacroInValue(modelConfig.Metadata, entry.Name, entry.Value)
if err != nil {
return Config{}, fmt.Errorf("model %s metadata: %s", modelId, err.Error())
}
modelConfig.Metadata = result.(map[string]any)
}
}
// Handle PORT macro - only allocate if cmd uses it
cmdHasPort := strings.Contains(modelConfig.Cmd, "${PORT}")
proxyHasPort := strings.Contains(modelConfig.Proxy, "${PORT}")
if cmdHasPort || proxyHasPort {
if !cmdHasPort && proxyHasPort {
return Config{}, fmt.Errorf("model %s: proxy uses ${PORT} but cmd does not - ${PORT} is only available when used in cmd", modelId)
}
macroSlug := "${PORT}"
macroStr := fmt.Sprintf("%v", nextPort)
modelConfig.Cmd = strings.ReplaceAll(modelConfig.Cmd, macroSlug, macroStr)
modelConfig.CmdStop = strings.ReplaceAll(modelConfig.CmdStop, macroSlug, macroStr)
modelConfig.Proxy = strings.ReplaceAll(modelConfig.Proxy, macroSlug, macroStr)
modelConfig.Name = strings.ReplaceAll(modelConfig.Name, macroSlug, macroStr)
modelConfig.Description = strings.ReplaceAll(modelConfig.Description, macroSlug, macroStr)
if len(modelConfig.Metadata) > 0 {
result, err := substituteMacroInValue(modelConfig.Metadata, "PORT", nextPort)
if err != nil {
return Config{}, fmt.Errorf("model %s metadata: %s", modelId, err.Error())
}
modelConfig.Metadata = result.(map[string]any)
}
nextPort++
}
// Validate no unknown macros remain
fieldMap := map[string]string{
"cmd": modelConfig.Cmd,
"cmdStop": modelConfig.CmdStop,
"proxy": modelConfig.Proxy,
"checkEndpoint": modelConfig.CheckEndpoint,
"filters.stripParams": modelConfig.Filters.StripParams,
"name": modelConfig.Name,
"description": modelConfig.Description,
}
for fieldName, fieldValue := range fieldMap {
matches := macroPatternRegex.FindAllStringSubmatch(fieldValue, -1)
for _, match := range matches {
macroName := match[1]
if macroName == "PID" && fieldName == "cmdStop" {
continue // replaced at runtime
}
if macroName == "PORT" || macroName == "MODEL_ID" {
return Config{}, fmt.Errorf("macro '${%s}' should have been substituted in %s.%s", macroName, modelId, fieldName)
}
return Config{}, fmt.Errorf("unknown macro '${%s}' found in %s.%s", macroName, modelId, fieldName)
}
}
if len(modelConfig.Metadata) > 0 {
if err := validateNestedForUnknownMacros(modelConfig.Metadata, fmt.Sprintf("model %s metadata", modelId)); err != nil {
return Config{}, err
}
}
if err = modelConfig.Capabilities.Validate(); err != nil {
return Config{}, fmt.Errorf("model %s: %w", modelId, err)
}
// Validate SetParamsByID keys and values
for key, paramMap := range modelConfig.Filters.SetParamsByID {
if matches := macroPatternRegex.FindAllStringSubmatch(key, -1); len(matches) > 0 {
return Config{}, fmt.Errorf("unknown macro '${%s}' found in model %s filters.setParamsByID key", matches[0][1], modelId)
}
if err := validateNestedForUnknownMacros(any(paramMap), fmt.Sprintf("model %s filters.setParamsByID[%s]", modelId, key)); err != nil {
return Config{}, err
}
}
// Auto-register setParamsByID keys as aliases (skip the model's own ID)
for key := range modelConfig.Filters.SetParamsByID {
if key == modelId {
continue
}
if _, exists := config.Models[key]; exists {
return Config{}, fmt.Errorf("model %s filters.setParamsByID: key '%s' conflicts with an existing model ID", modelId, key)
}
if existingModel, exists := config.aliases[key]; exists {
if existingModel != modelId {
return Config{}, fmt.Errorf("duplicate alias '%s' in model %s filters.setParamsByID, already used by model %s", key, modelId, existingModel)
}
continue // already registered as explicit alias for this model
}
config.aliases[key] = modelId
modelConfig.Aliases = append(modelConfig.Aliases, key)
}
if _, err := url.Parse(modelConfig.Proxy); err != nil {
return Config{}, fmt.Errorf("model %s: invalid proxy URL: %w", modelId, err)
}
if modelConfig.SendLoadingState == nil {
v := config.SendLoadingState
modelConfig.SendLoadingState = &v
}
config.Models[modelId] = modelConfig
}
// Normalize routing config. The legacy top-level `matrix`/`groups` keys and
// the new `routing.router` block are mutually exclusive: a config may use
// either style, never both.
hasTopLevel := config.Matrix != nil || len(config.Groups) > 0
rtr := config.Routing.Router
hasRouting := rtr.Use != "" || rtr.Settings.Matrix != nil || len(rtr.Settings.Groups) > 0
if hasTopLevel && hasRouting {
return Config{}, fmt.Errorf("config uses both the legacy top-level 'matrix'/'groups' keys and the new 'routing.router' block; please migrate the top-level keys into 'routing.router' and remove them")
}
if !hasTopLevel {
// Both groups and matrix may be defined under routing.router.settings;
// routing.router.use selects which one is active, so there is no conflict.
rs := config.Routing.Router.Settings
switch config.Routing.Router.Use {
case "matrix":
if rs.Matrix == nil {
return Config{}, fmt.Errorf("routing.router.use is 'matrix' but routing.router.settings.matrix is not set")
}
config.Matrix = rs.Matrix
case "group", "":
config.Groups = rs.Groups
default:
return Config{}, fmt.Errorf("routing.router.use: unknown router %q (valid: group, matrix)", config.Routing.Router.Use)
}
}
// groups XOR matrix
if config.Matrix != nil && len(config.Groups) > 0 {
return Config{}, fmt.Errorf("config cannot use both 'groups' and 'matrix'")
}
if config.Matrix != nil {
expandedSets, err := ValidateMatrix(*config.Matrix, config.Models)
if err != nil {
return Config{}, fmt.Errorf("matrix: %w", err)
}
config.Matrix.ExpandedSets = expandedSets
} else {
config = AddDefaultGroupToConfig(config)
// Validate group members
memberUsage := make(map[string]string)
for groupID, groupConfig := range config.Groups {
prevSet := make(map[string]bool)
for _, member := range groupConfig.Members {
if _, found := prevSet[member]; found {
return Config{}, fmt.Errorf("duplicate model member %s found in group: %s", member, groupID)
}
prevSet[member] = true
if existingGroup, exists := memberUsage[member]; exists {
return Config{}, fmt.Errorf("model member %s is used in multiple groups: %s and %s", member, existingGroup, groupID)
}
memberUsage[member] = groupID
}
}
}
// Build the canonical Config.Routing from the effective result. Both legacy
// and new-style configs converge here. The Matrix pointer is shared so
// ExpandedSets stays in one place.
if config.Matrix != nil {
config.Routing.Router.Use = "matrix"
} else {
config.Routing.Router.Use = "group"
}
config.Routing.Router.Settings.Matrix = config.Matrix
config.Routing.Router.Settings.Groups = config.Groups
// This fork defaults to the "serial" scheduler: one model loaded at a time,
// requests served in strict arrival order. Set use: fifo for the upstream
// throughput-oriented behavior that batches same-model requests.
if config.Routing.Scheduler.Use == "" {
config.Routing.Scheduler.Use = "serial"
}
switch config.Routing.Scheduler.Use {
case "fifo", "serial":
default:
return Config{}, fmt.Errorf("routing.scheduler.use: unknown scheduler %q (valid: fifo, serial)", config.Routing.Scheduler.Use)
}
for modelID := range config.Routing.Scheduler.Settings.Fifo.Priority {
if _, found := config.RealModelName(modelID); !found {
return Config{}, fmt.Errorf("routing.scheduler.settings.fifo.priority references unknown model %q", modelID)
}
}
serialCfg := config.Routing.Scheduler.Settings.Serial
if d := serialCfg.AgingDivisor; d != nil && *d < 0 {
return Config{}, fmt.Errorf("routing.scheduler.settings.serial.agingDivisor: must be >= 0, got %d (0 disables aging)", *d)
}
// Capped below the 100-point band spacing so the bonus can only break
// near-ties. A bonus that could reach the next band would let a batch job
// on the loaded model outrank an interactive request that needs a swap.
if b := serialCfg.SwapAffinityBonus; b != nil && (*b < 0 || *b >= 100) {
return Config{}, fmt.Errorf("routing.scheduler.settings.serial.swapAffinityBonus: must be between 0 and 99, got %d (priority bands are 100 apart and the bonus must never cross one)", *b)
}
// Clean up hooks preload
if len(config.Hooks.OnStartup.Preload) > 0 {
var toPreload []string
for _, modelID := range config.Hooks.OnStartup.Preload {
modelID = strings.TrimSpace(modelID)
if modelID == "" {
continue
}
if real, found := config.RealModelName(modelID); found {
toPreload = append(toPreload, real)
}
}
config.Hooks.OnStartup.Preload = toPreload
}
// Validate API keys (env macros already substituted at string level)
for i, apikey := range config.RequiredAPIKeys {
if apikey == "" {
return Config{}, fmt.Errorf("empty api key found in apiKeys")
}
if strings.Contains(apikey, " ") {
return Config{}, fmt.Errorf("apiKeys[%d]: api key cannot contain spaces", i)
}
config.RequiredAPIKeys[i] = apikey
}
// Process peers with global macro substitution
for peerName, peerConfig := range config.Peers {
// Substitute global macros (LIFO order)
for i := len(config.Macros) - 1; i >= 0; i-- {
entry := config.Macros[i]
macroSlug := fmt.Sprintf("${%s}", entry.Name)
macroStr := fmt.Sprintf("%v", entry.Value)
peerConfig.ApiKey = strings.ReplaceAll(peerConfig.ApiKey, macroSlug, macroStr)
peerConfig.Filters.StripParams = strings.ReplaceAll(peerConfig.Filters.StripParams, macroSlug, macroStr)
// Substitute in setParams (type-preserving)
if len(peerConfig.Filters.SetParams) > 0 {
result, err := substituteMacroInValue(peerConfig.Filters.SetParams, entry.Name, entry.Value)
if err != nil {
return Config{}, fmt.Errorf("peers.%s.filters.setParams: %w", peerName, err)
}
peerConfig.Filters.SetParams = result.(map[string]any)
}
}
// Validate no unknown macros remain
if matches := macroPatternRegex.FindAllStringSubmatch(peerConfig.ApiKey, -1); len(matches) > 0 {
return Config{}, fmt.Errorf("peers.%s.apiKey: unknown macro '${%s}'", peerName, matches[0][1])
}
if matches := macroPatternRegex.FindAllStringSubmatch(peerConfig.Filters.StripParams, -1); len(matches) > 0 {
return Config{}, fmt.Errorf("peers.%s.filters.stripParams: unknown macro '${%s}'", peerName, matches[0][1])
}
if len(peerConfig.Filters.SetParams) > 0 {
if err := validateNestedForUnknownMacros(peerConfig.Filters.SetParams, fmt.Sprintf("peers.%s.filters.setParams", peerName)); err != nil {
return Config{}, err
}
}
config.Peers[peerName] = peerConfig
}
return config, nil
}
+198
View File
@@ -0,0 +1,198 @@
package config
import (
"fmt"
"os"
"regexp"
"strings"
"gopkg.in/yaml.v3"
)
var (
macroNameRegex = regexp.MustCompile(`^[a-zA-Z0-9_-]+$`)
macroPatternRegex = regexp.MustCompile(`\$\{([a-zA-Z0-9_-]+)\}`)
envMacroRegex = regexp.MustCompile(`\$\{env\.([a-zA-Z_][a-zA-Z0-9_]*)\}`)
)
// validateMacro validates macro name and value constraints
func validateMacro(name string, value any) error {
if len(name) >= 64 {
return fmt.Errorf("macro name '%s' exceeds maximum length of 63 characters", name)
}
if !macroNameRegex.MatchString(name) {
return fmt.Errorf("macro name '%s' contains invalid characters, must match pattern ^[a-zA-Z0-9_-]+$", name)
}
// Validate that value is a scalar type
switch v := value.(type) {
case string:
// Check for self-reference
macroSlug := fmt.Sprintf("${%s}", name)
if strings.Contains(v, macroSlug) {
return fmt.Errorf("macro '%s' contains self-reference", name)
}
case int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64, float32, float64, bool:
// These types are allowed
default:
return fmt.Errorf("macro '%s' has invalid type %T, must be a scalar type (string, int, float, or bool)", name, value)
}
switch name {
case "PORT", "MODEL_ID":
return fmt.Errorf("macro name '%s' is reserved", name)
}
return nil
}
// validateNestedForUnknownMacros recursively checks for any remaining macro references in nested structures
func validateNestedForUnknownMacros(value any, context string) error {
switch v := value.(type) {
case string:
matches := macroPatternRegex.FindAllStringSubmatch(v, -1)
for _, match := range matches {
macroName := match[1]
return fmt.Errorf("%s: unknown macro '${%s}'", context, macroName)
}
// Check for unsubstituted env macros
envMatches := envMacroRegex.FindAllStringSubmatch(v, -1)
for _, match := range envMatches {
varName := match[1]
return fmt.Errorf("%s: environment variable '%s' not set", context, varName)
}
return nil
case map[string]any:
for _, val := range v {
if err := validateNestedForUnknownMacros(val, context); err != nil {
return err
}
}
return nil
case []any:
for _, val := range v {
if err := validateNestedForUnknownMacros(val, context); err != nil {
return err
}
}
return nil
default:
// Scalar types don't contain macros
return nil
}
}
// substituteMacroInValue recursively substitutes a single macro in a value structure
// This is called once per macro, allowing LIFO substitution order
func substituteMacroInValue(value any, macroName string, macroValue any) (any, error) {
macroSlug := fmt.Sprintf("${%s}", macroName)
macroStr := fmt.Sprintf("%v", macroValue)
switch v := value.(type) {
case string:
// Check if this is a direct macro substitution
if v == macroSlug {
return macroValue, nil
}
// Handle string interpolation
if strings.Contains(v, macroSlug) {
return strings.ReplaceAll(v, macroSlug, macroStr), nil
}
return v, nil
case map[string]any:
// Recursively process map values
newMap := make(map[string]any)
for key, val := range v {
newVal, err := substituteMacroInValue(val, macroName, macroValue)
if err != nil {
return nil, err
}
newMap[key] = newVal
}
return newMap, nil
case []any:
// Recursively process slice elements
newSlice := make([]any, len(v))
for i, val := range v {
newVal, err := substituteMacroInValue(val, macroName, macroValue)
if err != nil {
return nil, err
}
newSlice[i] = newVal
}
return newSlice, nil
default:
// Return scalar types as-is
return value, nil
}
}
// substituteEnvMacros replaces ${env.VAR_NAME} with environment variable values.
// Returns error if any referenced env var is not set or contains invalid characters.
// Env macros inside YAML comments are ignored by unmarshalling the YAML first
// (which strips comments) and only checking the comment-free version for macros.
func substituteEnvMacros(s string) (string, error) {
// Unmarshal and remarshal to strip YAML comments
var raw any
if err := yaml.Unmarshal([]byte(s), &raw); err != nil {
// If YAML is invalid, fall back to scanning the original string
// so the user gets the env var error rather than a confusing YAML parse error
return substituteEnvMacrosInString(s, s)
}
clean, err := yaml.Marshal(raw)
if err != nil {
return substituteEnvMacrosInString(s, s)
}
return substituteEnvMacrosInString(s, string(clean))
}
// substituteEnvMacrosInString finds ${env.VAR} macros in scanStr and substitutes
// them in target. This separation allows scanning comment-free YAML while
// substituting in the original string.
func substituteEnvMacrosInString(target, scanStr string) (string, error) {
result := target
matches := envMacroRegex.FindAllStringSubmatch(scanStr, -1)
for _, match := range matches {
fullMatch := match[0] // ${env.VAR_NAME}
varName := match[1] // VAR_NAME
value, exists := os.LookupEnv(varName)
if !exists {
return "", fmt.Errorf("environment variable '%s' is not set", varName)
}
// Sanitize the value for safe YAML substitution
value, err := sanitizeEnvValueForYAML(value, varName)
if err != nil {
return "", err
}
result = strings.ReplaceAll(result, fullMatch, value)
}
return result, nil
}
// sanitizeEnvValueForYAML ensures an environment variable value is safe for YAML substitution.
// It rejects values with characters that break YAML structure and escapes quotes/backslashes
// for compatibility with double-quoted YAML strings.
func sanitizeEnvValueForYAML(value, varName string) (string, error) {
// Reject values that would break YAML structure regardless of quoting context
if strings.ContainsAny(value, "\n\r\x00") {
return "", fmt.Errorf("environment variable '%s' contains newlines or null bytes which are not allowed in YAML substitution", varName)
}
// Escape backslashes and double quotes for safe use in double-quoted YAML strings.
// In unquoted contexts, these escapes appear literally (harmless for most use cases).
// In double-quoted contexts, they are interpreted correctly.
value = strings.ReplaceAll(value, `\`, `\\`)
value = strings.ReplaceAll(value, `"`, `\"`)
return value, nil
}
+300
View File
@@ -0,0 +1,300 @@
package config
import (
"fmt"
"os"
"path/filepath"
"sort"
"strings"
"gopkg.in/yaml.v3"
)
// identityMapPaths is the set of dotted paths whose direct children are
// identity-keyed maps. A child key present in two sources is a hard error;
// such keys name discrete entities (a model, a group, a peer, etc.) and a
// duplicate means the user has split one entity across files by mistake.
var identityMapPaths = map[string]bool{
"models": true,
"groups": true,
"profiles": true,
"peers": true,
"matrix": true,
"routing.router.settings.groups": true,
"routing.router.settings.matrix": true,
}
// LoadConfigSources loads and merges configuration from -config (optional)
// and -config-dir (optional). At least one must be provided. The -config file
// is loaded first; *.yml/*.yaml files directly under -config-dir are then
// merged in sorted filename order. The merged document is passed through the
// existing LoadConfigFromReader pipeline unchanged.
func LoadConfigSources(configPath, configDir string) (Config, error) {
if configPath == "" && configDir == "" {
return Config{}, fmt.Errorf("at least one of -config or -config-dir must be provided")
}
var sourcePaths []string
if configPath != "" {
sourcePaths = append(sourcePaths, configPath)
}
if configDir != "" {
dirFiles, err := listYAMLFiles(configDir)
if err != nil {
return Config{}, fmt.Errorf("-config-dir %s: %w", configDir, err)
}
if configPath != "" {
absConfig, err := filepath.Abs(configPath)
if err != nil {
return Config{}, fmt.Errorf("failed to resolve -config path: %w", err)
}
for _, f := range dirFiles {
absF, err := filepath.Abs(f)
if err != nil {
return Config{}, fmt.Errorf("failed to resolve config dir file %s: %w", f, err)
}
if absConfig == absF {
return Config{}, fmt.Errorf("-config path %s is also present in -config-dir %s; remove it from one", configPath, configDir)
}
}
}
sourcePaths = append(sourcePaths, dirFiles...)
}
if len(sourcePaths) == 0 {
return Config{}, fmt.Errorf("no configuration sources found")
}
var merged *yaml.Node
for _, p := range sourcePaths {
node, err := parseSource(p)
if err != nil {
return Config{}, err
}
if node == nil {
continue // empty file
}
if merged == nil {
merged = node
continue
}
if err := mergeNodes(merged, node, "", p); err != nil {
return Config{}, err
}
}
if merged == nil {
// All sources were empty; run the pipeline on empty input so defaults
// and validation still apply (e.g. startPort, performance defaults).
return LoadConfigFromReader(strings.NewReader(""))
}
out, err := yaml.Marshal(merged)
if err != nil {
return Config{}, fmt.Errorf("failed to marshal merged config: %w", err)
}
return LoadConfigFromReader(strings.NewReader(string(out)))
}
// listYAMLFiles returns the top-level *.yml and *.yaml files in dir, sorted by
// filename for deterministic merge order. Subdirectories are not traversed.
func listYAMLFiles(dir string) ([]string, error) {
entries, err := os.ReadDir(dir)
if err != nil {
return nil, err
}
var files []string
for _, e := range entries {
if e.IsDir() {
continue
}
name := e.Name()
if !strings.HasSuffix(name, ".yml") && !strings.HasSuffix(name, ".yaml") {
continue
}
files = append(files, filepath.Join(dir, name))
}
sort.Strings(files)
return files, nil
}
// parseSource reads and parses one YAML config file into a root mapping node.
// Returns a nil node (no error) when the file is empty or contains only
// comments.
//
// Env macros (${env.VAR}) are substituted at the string level before YAML
// parsing so that flow-style constructs like [${env.API_KEY}] parse
// correctly — the brace would otherwise be interpreted as a flow mapping.
func parseSource(path string) (*yaml.Node, error) {
data, err := os.ReadFile(path)
if err != nil {
return nil, fmt.Errorf("failed to read config %s: %w", path, err)
}
yamlStr, err := substituteEnvMacros(string(data))
if err != nil {
return nil, fmt.Errorf("config %s: %w", path, err)
}
var doc yaml.Node
if err := yaml.Unmarshal([]byte(yamlStr), &doc); err != nil {
return nil, fmt.Errorf("failed to parse config %s: %w", path, err)
}
// yaml.Unmarshal into a yaml.Node yields a DocumentNode whose Content[0]
// is the actual root. Unwrap it so callers see the real top-level node.
root := &doc
if root.Kind == yaml.DocumentNode && len(root.Content) > 0 {
root = root.Content[0]
}
if root.Kind == 0 || root.Content == nil {
return nil, nil
}
if root.Kind != yaml.MappingNode {
return nil, fmt.Errorf("config %s: top-level YAML must be a mapping", path)
}
return root, nil
}
// mergeNodes merges src into dst (both MappingNodes) in place. Keys present in
// only one side are kept; shared keys are merged recursively under the rules
// in mergeValue. srcPath is included in error messages to identify the file
// that introduced the conflict.
func mergeNodes(dst, src *yaml.Node, path, srcPath string) error {
srcIdx := indexMapping(src)
// First pass: merge shared keys in place.
for i := 0; i+1 < len(dst.Content); i += 2 {
keyNode := dst.Content[i]
dstVal := dst.Content[i+1]
key := keyNode.Value
srcVal, ok := srcIdx[key]
if !ok {
continue // dst-only key, keep as-is
}
childPath := joinPath(path, key)
if identityMapPaths[childPath] {
// Identity-keyed map: each child key names a discrete entity
// (a model, group, peer, ...). A shared child key is a hard
// error; src-only children are appended in the second pass.
if err := mergeIdentityMap(dstVal, srcVal, childPath, key, srcPath); err != nil {
return err
}
continue
}
if err := mergeValue(dstVal, srcVal, childPath, srcPath); err != nil {
return err
}
}
// Second pass: append src-only keys.
dstIdx := indexMapping(dst)
for i := 0; i+1 < len(src.Content); i += 2 {
keyNode := src.Content[i]
srcVal := src.Content[i+1]
key := keyNode.Value
if _, ok := dstIdx[key]; ok {
continue // already merged above
}
keyCopy := *keyNode
valCopy := *srcVal
dst.Content = append(dst.Content, &keyCopy, &valCopy)
}
return nil
}
// mergeIdentityMap merges two identity-keyed mapping nodes (e.g. `models`,
// `groups`, `peers`). Any child key present in both sides is a duplicate
// entity and produces an error naming the conflicting key and source file.
// src-only keys are appended to dst.
func mergeIdentityMap(dst, src *yaml.Node, path, mapName, srcPath string) error {
if dst.Kind != yaml.MappingNode || src.Kind != yaml.MappingNode {
return fmt.Errorf("conflict at %q: expected a mapping, introduced by %s", path, srcPath)
}
dstIdx := indexMapping(dst)
for i := 0; i+1 < len(src.Content); i += 2 {
keyNode := src.Content[i]
srcVal := src.Content[i+1]
key := keyNode.Value
if _, dup := dstIdx[key]; dup {
return fmt.Errorf("duplicate %s %q found in %s (already defined in another config source)", mapName, key, srcPath)
}
keyCopy := *keyNode
valCopy := *srcVal
dst.Content = append(dst.Content, &keyCopy, &valCopy)
}
return nil
}
// mergeValue merges srcVal into dstVal (both pointing into the parent's
// Content slice). Mapping+Mapping recurses; Sequence+Sequence concatenates;
// Scalar+Scalar errors on value mismatch; null on either side yields to the
// non-null side.
func mergeValue(dstVal, srcVal *yaml.Node, path, srcPath string) error {
switch {
case dstVal.Kind == yaml.MappingNode && srcVal.Kind == yaml.MappingNode:
return mergeNodes(dstVal, srcVal, path, srcPath)
case dstVal.Kind == yaml.SequenceNode && srcVal.Kind == yaml.SequenceNode:
dstVal.Content = append(dstVal.Content, srcVal.Content...)
return nil
case dstVal.Kind == yaml.ScalarNode && srcVal.Kind == yaml.ScalarNode:
if isNullScalar(dstVal) {
*dstVal = *srcVal
return nil
}
if isNullScalar(srcVal) {
return nil
}
if dstVal.Value == srcVal.Value && dstVal.Tag == srcVal.Tag {
return nil
}
return fmt.Errorf("conflict at %q: %s sets a different value than a previous source", path, srcPath)
case isNull(dstVal):
*dstVal = *srcVal
return nil
case isNull(srcVal):
return nil
default:
return fmt.Errorf("conflict at %q: incompatible YAML node kinds (kind %d vs %d) introduced by %s", path, dstVal.Kind, srcVal.Kind, srcPath)
}
}
// isNull reports whether n represents a YAML null (empty or !!null).
func isNull(n *yaml.Node) bool {
if n == nil || n.Kind == 0 {
return true
}
return isNullScalar(n)
}
func isNullScalar(n *yaml.Node) bool {
return n.Kind == yaml.ScalarNode && (n.Tag == "!!null" || n.Tag == "") && n.Value == ""
}
// indexMapping builds a key -> value-node index for a mapping node.
func indexMapping(n *yaml.Node) map[string]*yaml.Node {
idx := make(map[string]*yaml.Node, len(n.Content)/2)
for i := 0; i+1 < len(n.Content); i += 2 {
idx[n.Content[i].Value] = n.Content[i+1]
}
return idx
}
func joinPath(parent, key string) string {
if parent == "" {
return key
}
return parent + "." + key
}
+304
View File
@@ -0,0 +1,304 @@
package config
import (
"os"
"path/filepath"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// writeYAML writes content to a file named name inside dir. Returns the full
// path of the written file.
func writeYAML(t *testing.T, dir, name, content string) string {
t.Helper()
p := filepath.Join(dir, name)
require.NoError(t, os.MkdirAll(filepath.Dir(p), 0o755))
require.NoError(t, os.WriteFile(p, []byte(content), 0o644))
return p
}
// modelCfg builds a single-model YAML snippet indented for nesting under a
// `models:` key. The proxy uses a fixed port so tests don't depend on
// ${PORT} allocation.
func modelCfg(id, cmd string) string {
return " " + id + ":\n cmd: " + cmd + "\n proxy: \"http://localhost:9999\"\n"
}
func TestLoadConfigSources_NeitherProvided(t *testing.T) {
_, err := LoadConfigSources("", "")
require.Error(t, err)
assert.Contains(t, err.Error(), "at least one of -config or -config-dir")
}
func TestLoadConfigSources_ConfigOnly(t *testing.T) {
dir := t.TempDir()
cfgPath := writeYAML(t, dir, "config.yaml", `
models:
`+modelCfg("model1", "echo hi")+`
groups:
group1:
members: ["model1"]
`)
cfg, err := LoadConfigSources(cfgPath, "")
require.NoError(t, err)
_, id, ok := cfg.FindConfig("model1")
require.True(t, ok)
assert.Equal(t, "model1", id)
}
func TestLoadConfigSources_DirOnly(t *testing.T) {
dir := t.TempDir()
writeYAML(t, dir, "a.yaml", "models:\n"+modelCfg("alpha", "echo a"))
writeYAML(t, dir, "b.yaml", "models:\n"+modelCfg("beta", "echo b"))
cfg, err := LoadConfigSources("", dir)
require.NoError(t, err)
for _, want := range []string{"alpha", "beta"} {
_, _, ok := cfg.FindConfig(want)
assert.True(t, ok, "model %s should be present", want)
}
}
func TestLoadConfigSources_ConfigPlusDirAdditive(t *testing.T) {
// -config lives outside -config-dir; both contribute models additively.
dir := t.TempDir()
cfgPath := writeYAML(t, dir, "config.yaml", "models:\n"+modelCfg("base", "echo base"))
cfgDir := t.TempDir()
writeYAML(t, cfgDir, "extra.yaml", "models:\n"+modelCfg("ext", "echo ext"))
cfg, err := LoadConfigSources(cfgPath, cfgDir)
require.NoError(t, err)
for _, want := range []string{"base", "ext"} {
_, _, ok := cfg.FindConfig(want)
assert.True(t, ok, "model %s should be present after merge", want)
}
}
// TestLoadConfigSources_ConfigInDirOverlap verifies that a -config file that
// is also a member of -config-dir is rejected.
func TestLoadConfigSources_ConfigInDirOverlap(t *testing.T) {
dir := t.TempDir()
cfgPath := writeYAML(t, dir, "main.yaml", "models:\n"+modelCfg("base", "echo base"))
_, err := LoadConfigSources(cfgPath, dir)
require.Error(t, err)
assert.Contains(t, err.Error(), "is also present in -config-dir")
}
func TestLoadConfigSources_DuplicateModelID(t *testing.T) {
dir := t.TempDir()
writeYAML(t, dir, "a.yaml", "models:\n"+modelCfg("dup", "echo a"))
writeYAML(t, dir, "b.yaml", "models:\n"+modelCfg("dup", "echo b"))
_, err := LoadConfigSources("", dir)
require.Error(t, err)
assert.Contains(t, err.Error(), `duplicate models "dup"`)
}
func TestLoadConfigSources_DuplicateGroupID(t *testing.T) {
dir := t.TempDir()
writeYAML(t, dir, "a.yaml", `
models:
`+modelCfg("m1", "echo m1")+"groups:\n g1:\n members: [m1]\n")
writeYAML(t, dir, "b.yaml", `
models:
`+modelCfg("m2", "echo m2")+"groups:\n g1:\n members: [m2]\n")
_, err := LoadConfigSources("", dir)
require.Error(t, err)
assert.Contains(t, err.Error(), `duplicate groups "g1"`)
}
func TestLoadConfigSources_DuplicatePeer(t *testing.T) {
dir := t.TempDir()
peerA := "peers:\n remote:\n proxy: http://x:1\n models: [m1]\n"
peerB := "peers:\n remote:\n proxy: http://x:2\n models: [m2]\n"
writeYAML(t, dir, "a.yaml", "models:\n"+modelCfg("m1", "echo m1")+"\n"+peerA)
writeYAML(t, dir, "b.yaml", "models:\n"+modelCfg("m2", "echo m2")+"\n"+peerB)
_, err := LoadConfigSources("", dir)
require.Error(t, err)
assert.Contains(t, err.Error(), `duplicate peers "remote"`)
}
func TestLoadConfigSources_ScalarConflict(t *testing.T) {
dir := t.TempDir()
writeYAML(t, dir, "a.yaml", "models:\n"+modelCfg("m1", "echo m1")+"\nglobalTTL: 100\n")
writeYAML(t, dir, "b.yaml", "models:\n"+modelCfg("m2", "echo m2")+"\nglobalTTL: 200\n")
_, err := LoadConfigSources("", dir)
require.Error(t, err)
assert.Contains(t, err.Error(), `conflict at "globalTTL"`)
}
func TestLoadConfigSources_ScalarSameValueNoConflict(t *testing.T) {
dir := t.TempDir()
writeYAML(t, dir, "a.yaml", "models:\n"+modelCfg("m1", "echo m1")+"\nglobalTTL: 100\n")
writeYAML(t, dir, "b.yaml", "models:\n"+modelCfg("m2", "echo m2")+"\nglobalTTL: 100\n")
cfg, err := LoadConfigSources("", dir)
require.NoError(t, err)
assert.Equal(t, 100, cfg.GlobalTTL)
}
func TestLoadConfigSources_MacrosConcatenate(t *testing.T) {
dir := t.TempDir()
writeYAML(t, dir, "a.yaml", "macros:\n LOW: 1\nmodels:\n"+modelCfg("m1", "echo ${LOW}"))
writeYAML(t, dir, "b.yaml", "macros:\n HIGH: 2\nmodels:\n"+modelCfg("m2", "echo ${HIGH}"))
cfg, err := LoadConfigSources("", dir)
require.NoError(t, err)
// Both macros are available globally after merge.
low, ok := cfg.Macros.Get("LOW")
require.True(t, ok)
assert.Equal(t, 1, low)
high, ok := cfg.Macros.Get("HIGH")
require.True(t, ok)
assert.Equal(t, 2, high)
}
func TestLoadConfigSources_APIKeysConcatenate(t *testing.T) {
dir := t.TempDir()
writeYAML(t, dir, "a.yaml", "models:\n"+modelCfg("m1", "echo m1")+"\napiKeys: [key-a]\n")
writeYAML(t, dir, "b.yaml", "models:\n"+modelCfg("m2", "echo m2")+"\napiKeys: [key-b]\n")
cfg, err := LoadConfigSources("", dir)
require.NoError(t, err)
assert.ElementsMatch(t, []string{"key-a", "key-b"}, cfg.RequiredAPIKeys)
}
func TestLoadConfigSources_RoutingGroupsMerge(t *testing.T) {
dir := t.TempDir()
writeYAML(t, dir, "a.yaml", `
models:
`+modelCfg("m1", "echo m1")+`
routing:
router:
settings:
groups:
groupA:
members: [m1]
`)
writeYAML(t, dir, "b.yaml", `
models:
`+modelCfg("m2", "echo m2")+`
routing:
router:
settings:
groups:
groupB:
members: [m2]
`)
cfg, err := LoadConfigSources("", dir)
require.NoError(t, err)
groups := cfg.Routing.Router.Settings.Groups
assert.Contains(t, groups, "groupA")
assert.Contains(t, groups, "groupB")
// default group added by pipeline for orphaned/leftover routing groups...
// here both groups reference distinct models
}
func TestLoadConfigSources_EnvMacrosSubstituted(t *testing.T) {
dir := t.TempDir()
// Use ${PORT} in cmd so the pipeline allocates a port and substitutes it;
// verifies env/macro substitution runs on the merged document.
writeYAML(t, dir, "a.yaml", "models:\n m1:\n cmd: serve --port ${PORT}\n proxy: \"http://localhost:${PORT}\"\n")
cfg, err := LoadConfigSources("", dir)
require.NoError(t, err)
m := cfg.Models["m1"]
assert.NotContains(t, m.Cmd, "${PORT}", "PORT macro should have been substituted")
assert.NotContains(t, m.Proxy, "${PORT}", "PORT macro should have been substituted in proxy")
}
func TestLoadConfigSources_EnvMacroInFlowStyleList(t *testing.T) {
// Regression: flow-style lists with ${env.*} must parse. Previously
// parseSource unmarshalled before env substitution, so the brace in
// [${env.API_KEY}] was misread as a flow mapping and parsing failed.
dir := t.TempDir()
writeYAML(t, dir, "a.yaml", "models:\n m1:\n cmd: echo hi\n proxy: \"http://localhost:9999\"\n")
writeYAML(t, dir, "keys.yaml", "apiKeys: [${env.TEST_API_KEY}]\nmodels:\n m2:\n cmd: echo hi\n proxy: \"http://localhost:9998\"\n")
t.Setenv("TEST_API_KEY", "secret123")
cfg, err := LoadConfigSources("", dir)
require.NoError(t, err)
assert.Contains(t, cfg.RequiredAPIKeys, "secret123")
}
func TestLoadConfigSources_SortedOrderDeterministic(t *testing.T) {
// Two files defining distinct models, scanned in z..a order by filename.
// Determine merged result is the same regardless of how the FS returns them.
dir := t.TempDir()
writeYAML(t, dir, "z.yaml", "models:\n"+modelCfg("zmodel", "echo z"))
writeYAML(t, dir, "a.yaml", "models:\n"+modelCfg("amodel", "echo a"))
const runs = 3
for i := 0; i < runs; i++ {
cfg, err := LoadConfigSources("", dir)
require.NoError(t, err)
// startPort-based allocation: first allocated model gets 5800.
// Sorted order means amodel gets 5800, zmodel gets 5801.
_, _, ok := cfg.FindConfig("amodel")
assert.True(t, ok)
_, _, ok = cfg.FindConfig("zmodel")
assert.True(t, ok)
}
}
func TestLoadConfigSources_EmptyDirWithConfig(t *testing.T) {
dir := t.TempDir()
cfgDir := t.TempDir()
cfgPath := writeYAML(t, dir, "main.yaml", "models:\n"+modelCfg("m1", "echo m1"))
cfg, err := LoadConfigSources(cfgPath, cfgDir)
require.NoError(t, err)
assert.Contains(t, cfg.Models, "m1")
}
func TestLoadConfigSources_EmptyDirOnly(t *testing.T) {
// An empty -config-dir with no -config is an error: there is nothing to
// load and silently producing an empty config would mask the misconfig.
cfgDir := t.TempDir()
_, err := LoadConfigSources("", cfgDir)
require.Error(t, err)
assert.Contains(t, err.Error(), "no configuration sources found")
}
func TestLoadConfigSources_AssertNoUnknownMacrosAfterMerge(t *testing.T) {
// Macros defined in one file should not satisfy unknown-macro validation in
// another — they do, because merge concats global macros before validation
// runs. This test documents that a macro from file A is usable in file B.
dir := t.TempDir()
writeYAML(t, dir, "macros.yaml", "macros:\n SHARED: hello\nmodels:\n"+modelCfg("dummy", "echo dummy"))
writeYAML(t, dir, "use.yaml", "models:\n"+modelCfg("user", "echo ${SHARED}"))
cfg, err := LoadConfigSources("", dir)
require.NoError(t, err)
m := cfg.Models["user"]
assert.Contains(t, m.Cmd, "hello")
assert.NotContains(t, m.Cmd, "${SHARED}")
}
func TestLoadConfigSources_KindMismatchErrors(t *testing.T) {
dir := t.TempDir()
writeYAML(t, dir, "a.yaml", "startPort: 5800\nmodels:\n"+modelCfg("m1", "echo m1"))
writeYAML(t, dir, "b.yaml", "startPort: [5800, 5801]\nmodels:\n"+modelCfg("m2", "echo m2"))
_, err := LoadConfigSources("", dir)
require.Error(t, err)
assert.Contains(t, err.Error(), "incompatible YAML node kinds")
}
func TestLoadConfigSources_NullYieldsToValue(t *testing.T) {
// File A: routing.router block absent (null on root for routing);
// file B: defines routing.router.settings.groups. Merge should keep B's.
dir := t.TempDir()
writeYAML(t, dir, "a.yaml", "models:\n"+modelCfg("m1", "echo m1"))
writeYAML(t, dir, "b.yaml", "routing:\n router:\n settings:\n groups:\n g1:\n members: [m1]\nmodels:\n"+modelCfg("m2", "echo m2"))
cfg, err := LoadConfigSources("", dir)
require.NoError(t, err)
assert.Contains(t, cfg.Routing.Router.Settings.Groups, "g1")
}
+3 -2
View File
@@ -14,6 +14,7 @@ var validModalities = map[string]struct{}{
"text": {},
"audio": {},
"image": {},
"video": {},
}
// ModelCapConfig defines what modalities and features a model supports.
@@ -37,12 +38,12 @@ func (c ModelCapConfig) Empty() bool {
func (c ModelCapConfig) Validate() error {
for _, m := range c.In {
if _, ok := validModalities[m]; !ok {
return fmt.Errorf("capabilities.in: invalid modality %q, must be one of: text, audio, image", m)
return fmt.Errorf("capabilities.in: invalid modality %q, must be one of: text, audio, image, video", m)
}
}
for _, m := range c.Out {
if _, ok := validModalities[m]; !ok {
return fmt.Errorf("capabilities.out: invalid modality %q, must be one of: text, audio, image", m)
return fmt.Errorf("capabilities.out: invalid modality %q, must be one of: text, audio, image, video", m)
}
}
if c.Context < 0 {
+11 -6
View File
@@ -296,20 +296,25 @@ func TestConfig_ModelCapabilities_Validate(t *testing.T) {
assert.NoError(t, caps.Validate())
})
t.Run("valid_video_modality", func(t *testing.T) {
caps := ModelCapConfig{In: []string{"text", "image"}, Out: []string{"video"}}
assert.NoError(t, caps.Validate())
})
t.Run("invalid_in_modality", func(t *testing.T) {
caps := ModelCapConfig{In: []string{"video"}}
caps := ModelCapConfig{In: []string{"smell"}}
err := caps.Validate()
assert.Error(t, err)
assert.Contains(t, err.Error(), "capabilities.in")
assert.Contains(t, err.Error(), "video")
assert.Contains(t, err.Error(), "smell")
})
t.Run("invalid_out_modality", func(t *testing.T) {
caps := ModelCapConfig{Out: []string{"video"}}
caps := ModelCapConfig{Out: []string{"smell"}}
err := caps.Validate()
assert.Error(t, err)
assert.Contains(t, err.Error(), "capabilities.out")
assert.Contains(t, err.Error(), "video")
assert.Contains(t, err.Error(), "smell")
})
t.Run("negative_context", func(t *testing.T) {
@@ -327,10 +332,10 @@ models:
capabilities:
in:
- text
- video
- smell
`
_, err := LoadConfigFromReader(strings.NewReader(content))
assert.Error(t, err)
assert.Contains(t, err.Error(), "video")
assert.Contains(t, err.Error(), "smell")
})
}
+55
View File
@@ -0,0 +1,55 @@
package config
import (
"fmt"
"regexp"
"gopkg.in/yaml.v3"
)
// DefaultUpstreamIgnorePathsPattern is the default regular expression applied
// to upstream.ignorePaths when the section is empty or absent from the config.
// It matches common static-asset suffixes so requests for .js/.css/.png/etc.
// files do not trigger a model swap.
const DefaultUpstreamIgnorePathsPattern = `.*\.(js|json|css|png|gif|jpg|jpeg|ico|txt)$`
// DefaultUpstreamIgnorePaths returns the default compiled ignore paths used
// when upstream.ignorePaths is not specified in the config. The returned slice
// is fresh so callers may mutate it without affecting other configs.
func DefaultUpstreamIgnorePaths() []*regexp.Regexp {
return []*regexp.Regexp{regexp.MustCompile(DefaultUpstreamIgnorePathsPattern)}
}
// UpstreamConfig controls behaviour of the /upstream passthrough endpoint.
type UpstreamConfig struct {
// IgnorePaths is a slice of compiled regular expressions. Any request to
// /upstream/<model>/<path> whose remaining path matches any of these
// expressions will be ignored and not trigger a swap. When the config
// does not specify any patterns, DefaultUpstreamIgnorePaths is applied.
IgnorePaths []*regexp.Regexp `yaml:"-"`
}
// rawUpstreamConfig is the intermediate form used to unmarshal the YAML into
// plain strings, which are then compiled into *regexp.Regexp.
type rawUpstreamConfig struct {
IgnorePaths []string `yaml:"ignorePaths"`
}
// UnmarshalYAML compiles each ignorePaths entry into a *regexp.Regexp. If any
// entry fails to compile, an error is returned.
func (u *UpstreamConfig) UnmarshalYAML(value *yaml.Node) error {
var raw rawUpstreamConfig
if err := value.Decode(&raw); err != nil {
return err
}
patterns := make([]*regexp.Regexp, 0, len(raw.IgnorePaths))
for _, p := range raw.IgnorePaths {
re, err := regexp.Compile(p)
if err != nil {
return fmt.Errorf("upstream.ignorePaths: invalid regular expression %q: %w", p, err)
}
patterns = append(patterns, re)
}
u.IgnorePaths = patterns
return nil
}
+88
View File
@@ -0,0 +1,88 @@
package config
import (
"regexp"
"strings"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
const upstreamConfigHeader = `
models:
model1:
cmd: path/to/cmd --arg1 one
proxy: "http://localhost:8080"
`
func TestConfig_UpstreamIgnorePaths_DefaultWhenAbsent(t *testing.T) {
// When upstream is not specified at all, the default pattern is applied.
content := upstreamConfigHeader
cfg, err := LoadConfigFromReader(strings.NewReader(content))
require.NoError(t, err)
require.Len(t, cfg.Upstream.IgnorePaths, 1)
def := cfg.Upstream.IgnorePaths[0]
assert.IsType(t, &regexp.Regexp{}, def)
assert.Equal(t, DefaultUpstreamIgnorePathsPattern, def.String())
// The default matches common static-asset suffixes.
assert.True(t, def.MatchString("/foo.js"))
assert.True(t, def.MatchString("/bar/baz.json"))
assert.True(t, def.MatchString("/static/img.png"))
assert.True(t, def.MatchString("/notes.txt"))
assert.True(t, def.MatchString("/favicon.ico"))
// And does not match inference API paths.
assert.False(t, def.MatchString("/v1/chat/completions"))
assert.False(t, def.MatchString("/v1/models"))
assert.False(t, def.MatchString("/health"))
}
func TestConfig_UpstreamIgnorePaths_DefaultWhenSectionEmpty(t *testing.T) {
// When upstream is present but ignorePaths is omitted, the default is still
// applied.
content := `upstream: {}` + "\n" + upstreamConfigHeader
cfg, err := LoadConfigFromReader(strings.NewReader(content))
require.NoError(t, err)
require.Len(t, cfg.Upstream.IgnorePaths, 1)
assert.Equal(t, DefaultUpstreamIgnorePathsPattern, cfg.Upstream.IgnorePaths[0].String())
}
func TestConfig_UpstreamIgnorePaths_Compiles(t *testing.T) {
content := `
upstream:
ignorePaths:
- ".*\\.(js|json|css|png|gif|jpg|jpeg|txt)$"
- "^/static/.*"
` + upstreamConfigHeader
cfg, err := LoadConfigFromReader(strings.NewReader(content))
require.NoError(t, err)
require.Len(t, cfg.Upstream.IgnorePaths, 2)
// Verify the patterns are compiled into *regexp.Regexp and match as expected.
assert.True(t, cfg.Upstream.IgnorePaths[0].MatchString("/foo.js"))
assert.True(t, cfg.Upstream.IgnorePaths[0].MatchString("/bar/baz.json"))
assert.False(t, cfg.Upstream.IgnorePaths[0].MatchString("/v1/chat/completions"))
assert.True(t, cfg.Upstream.IgnorePaths[1].MatchString("/static/foo.png"))
assert.False(t, cfg.Upstream.IgnorePaths[1].MatchString("/v1/chat/completions"))
// Confirm the type is *regexp.Regexp to satisfy the API contract.
for _, re := range cfg.Upstream.IgnorePaths {
assert.IsType(t, &regexp.Regexp{}, re)
}
}
func TestConfig_UpstreamIgnorePaths_InvalidRegexReturnsError(t *testing.T) {
content := `
upstream:
ignorePaths:
- "[invalid("
` + upstreamConfigHeader
_, err := LoadConfigFromReader(strings.NewReader(content))
require.Error(t, err)
assert.Contains(t, err.Error(), "upstream.ignorePaths")
assert.Contains(t, err.Error(), "invalid regular expression")
}
+14 -2
View File
@@ -338,6 +338,17 @@ func (b *baseRouter) Handles(model string) bool {
return ok
}
// SchedulerStats returns a snapshot of the scheduler's queue. ok is false when
// the configured scheduler does not report stats (only "serial" does today), in
// which case callers should omit the metrics rather than emit zeroes.
func (b *baseRouter) SchedulerStats() (scheduler.QueueStats, bool) {
reporter, ok := b.schedule.(scheduler.StatsReporter)
if !ok {
return scheduler.QueueStats{}, false
}
return reporter.QueueStats(), true
}
func (b *baseRouter) ProcessLogger(modelID string) (*logmon.Monitor, bool) {
if p, ok := b.processes[modelID]; ok {
return p.Logger(), true
@@ -420,8 +431,9 @@ func (b *baseRouter) ServeHTTP(w http.ResponseWriter, req *http.Request) {
}
hr := scheduler.HandlerReq{
Model: data.ModelID,
Ctx: req.Context(),
Model: data.ModelID,
Priority: data.Priority,
Ctx: req.Context(),
// Unbuffered: a successful send on Respond proves the waiter is
// alive and consuming. grant() relies on this to avoid handing a
// handleFunc to a cancelled waiter and leaking the inFlight count.
+72
View File
@@ -5,6 +5,7 @@ import (
"io"
"net/http"
"net/http/httptest"
"sync"
"testing"
"time"
@@ -12,6 +13,7 @@ import (
"github.com/mostlygeek/llama-swap/internal/logmon"
"github.com/mostlygeek/llama-swap/internal/process"
"github.com/mostlygeek/llama-swap/internal/router/scheduler"
"github.com/mostlygeek/llama-swap/internal/shared"
)
// These tests cover baseRouter's own machinery — the run loop, process
@@ -218,6 +220,76 @@ func TestBaseRouter_ContextCancel(t *testing.T) {
}
}
// TestBaseRouter_SerialPriorityHeader is the end-to-end path for issue #9: the
// X-LlamaSwap-Priority header on an HTTP request reaches the serial scheduler
// and changes which queued request runs next.
//
// "hold" occupies the single slot. batch and urgent queue behind it in that
// order; when the slot frees, urgent must go first despite arriving later.
// urgent blocks inside its handler, so observing urgent's handler entry proves
// the scheduler chose it — if batch had won, batch's (unblocked) handler would
// already have run by then.
func TestBaseRouter_SerialPriorityHeader(t *testing.T) {
hold := newFakeProcess("hold")
hold.autoReady = true
hold.serveBlock = make(chan struct{})
batch := newFakeProcess("batch")
batch.autoReady = true
urgent := newFakeProcess("urgent")
urgent.autoReady = true
urgent.serveBlock = make(chan struct{})
conf := config.Config{HealthCheckTimeout: 5}
conf.Routing.Scheduler.Use = "serial"
b, err := newBaseRouter("test", conf, map[string]process.Process{
"hold": hold, "batch": batch, "urgent": urgent,
}, logmon.NewWriter(io.Discard), &stubPlanner{})
if err != nil {
t.Fatalf("newBaseRouter: %v", err)
}
b.testProcessed = make(chan struct{}, 64)
go b.run()
releaseHold := sync.OnceFunc(func() { close(hold.serveBlock) })
releaseUrgent := sync.OnceFunc(func() { close(urgent.serveBlock) })
t.Cleanup(func() {
releaseHold()
releaseUrgent()
if !b.shuttingDown.Load() {
_ = b.Shutdown(time.Second)
}
})
serve := func(model, priority string) {
r := newRequest(model)
if priority != "" {
r.Header.Set(shared.PriorityHeader, priority)
}
go b.ServeHTTP(httptest.NewRecorder(), r)
}
serve("hold", "")
waitProcessed(t, b.testProcessed, 2) // OnRequest, then the swap completing
<-hold.serveStarted
// Both queue behind hold; batch arrives first.
serve("batch", "batch")
waitProcessed(t, b.testProcessed, 1)
serve("urgent", "interactive")
waitProcessed(t, b.testProcessed, 1)
releaseHold()
select {
case <-urgent.serveStarted:
case <-time.After(2 * time.Second):
t.Fatal("interactive request never started")
}
if got := batch.serveCalls.Load(); got != 0 {
t.Errorf("batch serveCalls=%d want 0 — the batch job ran before the interactive request", got)
}
}
func TestBaseRouter_ModelNotFound(t *testing.T) {
a := newFakeProcess("a")
b := newTestBase(t, map[string]process.Process{"a": a}, &stubPlanner{})
+5
View File
@@ -6,6 +6,7 @@ import (
"github.com/mostlygeek/llama-swap/internal/logmon"
"github.com/mostlygeek/llama-swap/internal/process"
"github.com/mostlygeek/llama-swap/internal/router/scheduler"
"github.com/mostlygeek/llama-swap/internal/shared"
)
@@ -49,4 +50,8 @@ type LocalRouter interface {
// modelID must be a real (non-alias) config key. Returns false when the
// model is not known to this router.
ProcessLogger(modelID string) (*logmon.Monitor, bool)
// SchedulerStats returns a snapshot of the scheduler's queue for metrics.
// ok is false when the configured scheduler does not report stats.
SchedulerStats() (scheduler.QueueStats, bool)
}
+13 -6
View File
@@ -280,7 +280,7 @@ func (s *FIFO) grantHandler(req HandlerReq, modelID string) {
return
}
if err := shared.SetReqData(req.Ctx, "fifo_priority", strconv.Itoa(s.cfg.Priority[req.Model])); err != nil {
if err := shared.SetReqData(req.Ctx, "fifo_priority", strconv.Itoa(s.priorityOf(req))); err != nil {
s.logger.Debugf("failed to set fifo_priority metadata: %v", err)
}
@@ -311,15 +311,22 @@ func (s *FIFO) startSwap(initial HandlerReq, evict, running []string) {
s.effects.StartSwap(initial.Model, evict)
}
// priorityOf is a request's effective queue priority: the caller's
// X-LlamaSwap-Priority plus the model's configured priority. Both default to 0,
// so a deployment that uses neither keeps plain FIFO order.
func (s *FIFO) priorityOf(req HandlerReq) int {
return req.Priority + s.cfg.Priority[req.Model]
}
// enqueue inserts req into the queue in priority order: it goes just before the
// first queued item whose priority is strictly lower, so higher-priority models
// are serviced first while equal-priority requests keep their arrival (FIFO)
// order. Priorities come from the FifoConfig; unlisted models default to 0.
// first queued item whose priority is strictly lower, so higher-priority
// requests are serviced first while equal-priority requests keep their arrival
// (FIFO) order.
func (s *FIFO) enqueue(req HandlerReq) {
p := s.cfg.Priority[req.Model]
p := s.priorityOf(req)
i := len(s.queued)
for j, q := range s.queued {
if s.cfg.Priority[q.Model] < p {
if s.priorityOf(q) < p {
i = j
break
}
+28
View File
@@ -571,6 +571,34 @@ func TestFIFO_PriorityQueueOrder(t *testing.T) {
}
}
// TestFIFO_RequestPriorityOrder verifies the caller's X-LlamaSwap-Priority is
// added to the model's configured priority, so a batch caller for a
// high-priority model still queues behind an interactive caller for a
// low-priority one.
func TestFIFO_RequestPriorityOrder(t *testing.T) {
eff := newFakeEffects()
for _, m := range []string{"z", "hot", "cold"} {
eff.states[m] = process.StateStopped
}
planner := &stubPlanner{evict: map[string][]string{"z": {"hot", "cold"}}}
cfg := config.FifoConfig{Priority: map[string]int{"hot": 10}}
s := NewFIFO("test", logmon.NewWriter(io.Discard), planner, cfg, nil, eff)
s.OnRequest(req("z")) // StartSwap(z, [hot, cold]) — everything else queues
s.OnRequest(HandlerReq{Model: "hot", Priority: shared.PriorityBatch}) // 10 - 100 = -90
s.OnRequest(HandlerReq{Model: "cold", Priority: shared.PriorityInteractive}) // 0 + 100 = 100
got := make([]string, len(s.queued))
for i, q := range s.queued {
got[i] = q.Model
}
want := []string{"cold", "hot"}
if len(got) != len(want) || got[0] != want[0] || got[1] != want[1] {
t.Fatalf("queue=%v want %v", got, want)
}
}
// TestFIFO_OnCancel_QueuedRequest verifies that cancelling a queued request
// 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.
+50 -4
View File
@@ -92,9 +92,14 @@ type Effects interface {
StopProcesses(timeout time.Duration, ids []string)
}
// New returns a Scheduler selected by conf.Routing.Scheduler.Use, configured
// from conf and bound to the given planner and effects. Currently only "fifo"
// (the default) is supported.
// New returns a Scheduler selected by conf.Routing.Scheduler.Use, configured from
// conf and bound to the given planner and effects. Supported values are "fifo"
// (throughput-oriented, batches same-model requests) and "serial" (strict
// one-model-at-a-time, highest-priority-first).
//
// The deployment default is applied by config loading (LoadConfig sets Use to
// "serial" when unset). The "" fallback here is the library default and remains
// "fifo" so callers that build a Config directly keep the original behavior.
func New(conf config.Config, name string, logger *logmon.Monitor, planner Swapper, eff Effects) (Scheduler, error) {
use := conf.Routing.Scheduler.Use
if use == "" {
@@ -103,6 +108,9 @@ func New(conf config.Config, name string, logger *logmon.Monitor, planner Swappe
switch use {
case "fifo":
return NewFIFO(name, logger, planner, conf.Routing.Scheduler.Settings.Fifo, conf.Models, eff), nil
case "serial":
// Serial ignores the group planner: it always evicts every other model.
return NewSerial(name, logger, conf.Routing.Scheduler.Settings.Serial, eff), nil
default:
return nil, fmt.Errorf("unsupported scheduler type: %q", use)
}
@@ -110,12 +118,50 @@ func New(conf config.Config, name string, logger *logmon.Monitor, planner Swappe
// HandlerReq is one in-flight ServeHTTP request waiting for a routing decision.
type HandlerReq struct {
Model string
Model string
// Priority is the caller's scheduling priority from the
// X-LlamaSwap-Priority header. Higher is more urgent, 0 is normal, and
// negative is batch work. Schedulers use it to order their queues.
Priority int
Ctx context.Context
Respond chan HandlerResp
PositionCh chan int
}
// BandStats is the queue state for one priority band.
type BandStats struct {
// Depth is how many requests are waiting in this band.
Depth int
// OldestWait is how long the longest-waiting request in this band has
// been queued, or 0 when the band is empty.
OldestWait time.Duration
// Dispatched counts requests in this band handed off since start.
Dispatched uint64
}
// QueueStats is a point-in-time snapshot of a scheduler's queue, for metrics.
// The reorder counters answer "did the scoring function actually change
// anything?" — without them, tuning aging and swap affinity is guesswork.
type QueueStats struct {
// Bands is keyed by shared.PriorityBand and always carries an entry for
// every band in shared.PriorityBands, so metric series stay stable.
Bands map[string]BandStats
// AgingReorders counts dispatches where the aging term changed which
// request was picked.
AgingReorders uint64
// AffinityReorders counts dispatches where the swap-affinity term changed
// which request was picked.
AffinityReorders uint64
}
// StatsReporter is implemented by schedulers that can report queue statistics.
// It is optional: callers type-assert and skip schedulers that do not provide
// it. Unlike the Scheduler methods, QueueStats is called from arbitrary
// goroutines (the /metrics handler) and must be safe for concurrent use.
type StatsReporter interface {
QueueStats() QueueStats
}
// HandlerResp is the routing decision returned to a HandlerReq's caller: either
// a handler to serve with, or an error.
type HandlerResp struct {
+536
View File
@@ -0,0 +1,536 @@
package scheduler
import (
"fmt"
"sort"
"strconv"
"sync"
"time"
"github.com/mostlygeek/llama-swap/internal/config"
"github.com/mostlygeek/llama-swap/internal/logmon"
"github.com/mostlygeek/llama-swap/internal/process"
"github.com/mostlygeek/llama-swap/internal/shared"
)
// Serial is a strict one-model-at-a-time scheduler for a size-1 resource: at
// most one request runs at any instant, and when the next request targets a
// model other than the one loaded, every other running model is evicted and the
// target is loaded before it runs. A single model occupies memory at a time, at
// the cost of throughput.
//
// Serial ignores group/eviction policy entirely: it always evicts every other
// running model, regardless of how groups are configured. That is what makes
// the single-model guarantee a property of the scheduler rather than of the
// config.
//
// # Dispatch order
//
// Scheduling is non-preemptive — a running job is never interrupted, because a
// generation cannot be stopped mid-sample without discarding the work. Priority
// applies at dispatch time: whenever the slot frees, the highest scoring
// waiting request goes next, where
//
// score = request priority + swap affinity + aging
//
// - request priority is the caller's X-LlamaSwap-Priority (0 by default).
// Bands sit 100 apart: +100 interactive, 0 normal, -100 batch.
// - swap affinity is a bounded bonus for a request that can run without a
// model swap. Bounded well below the band spacing, so it only breaks
// near-ties and can never promote batch work above interactive.
// - aging is waited_seconds / agingDivisor, and is unbounded on purpose: it
// is the only term allowed to cross a band, which is what stops a batch job
// from starving behind an endless stream of normal traffic.
//
// Equal scores keep arrival order, so with default priorities and swap affinity
// disabled this degrades to exact FIFO.
//
// The consequence of being non-preemptive is worth stating plainly: an
// interactive request can still wait up to one full job duration. Priority
// alone does not fix that — the complementary lever is client-side, where a
// batch producer submits one unit at a time and re-queues so gaps are frequent.
//
// Like FIFO, every Scheduler method runs on the router's single run-loop
// goroutine, so no internal locking is needed for queue state. QueueStats is
// the exception: it is read from the /metrics handler, so the stats mirror it
// reads is guarded by a mutex.
type Serial struct {
name string
logger *logmon.Monitor
effects Effects
// agingDivisor is how long a request must wait to gain one priority
// point. Zero disables aging.
agingDivisor time.Duration
// swapAffinityBonus is added to a request that needs no model swap. Zero
// disables the term.
swapAffinityBonus int
// now is the clock, injectable so tests can drive aging deterministically.
now func() time.Time
// queued holds waiting requests in arrival order. Dispatch picks by score
// rather than popping the head, but the slice itself is never reordered so
// index order remains arrival order — which is what makes equal scores
// break in favour of the earlier request.
queued []queuedReq
// active is the one request currently being processed (loading or serving),
// or nil when idle. phase is meaningful only while active != nil.
active *HandlerReq
phase serialPhase
mu sync.Mutex
stats serialStats
}
// queuedReq is one waiting request plus the arrival time aging is measured from.
type queuedReq struct {
req HandlerReq
enqueued time.Time
}
// serialStats mirrors queue state and dispatch counters for QueueStats. It is
// written on the run loop and read by the metrics handler, both under mu.
type serialStats struct {
depth map[string]int
oldest map[string]time.Time
dispatched map[string]uint64
agingReorders uint64
affinityReorders uint64
}
// serialPhase is the lifecycle stage of the active request.
type serialPhase int
const (
phaseIdle serialPhase = iota
phaseSwapping // waiting for OnSwapDone for active.Model
phaseServing // waiting for OnServeDone for active.Model
)
// NewSerial builds a Serial scheduler. It takes no Swapper: eviction is always
// "stop every other running model", so the group planner is not consulted.
func NewSerial(name string, logger *logmon.Monitor, cfg config.SerialConfig, eff Effects) *Serial {
return &Serial{
name: name,
logger: logger,
effects: eff,
agingDivisor: time.Duration(cfg.GetAgingDivisor()) * time.Second,
swapAffinityBonus: cfg.GetSwapAffinityBonus(),
now: time.Now,
stats: serialStats{
depth: make(map[string]int),
oldest: make(map[string]time.Time),
dispatched: make(map[string]uint64),
},
}
}
// OnRequest validates the model and appends the request to the queue, stamping
// its arrival time so aging can be measured, then tries to start the next job.
// Unknown models fail immediately.
func (s *Serial) OnRequest(req HandlerReq) {
if _, ok := s.effects.ModelState(req.Model); !ok {
s.logger.Debugf("%s: model %s not handled by this router", s.name, req.Model)
s.effects.GrantError(req, ErrModelNotFound)
return
}
s.queued = append(s.queued, queuedReq{req: req, enqueued: s.now()})
s.queueChanged()
s.startNext()
}
// startNext begins processing the highest scoring waiting request when nothing
// is active. It fast-paths a request whose model is already the sole
// loaded-and-ready process; otherwise it launches a swap that evicts every other
// running model first. The loop skips over requests for models that vanished
// (e.g. a config reload) and requests whose caller disconnected before they
// could be served.
func (s *Serial) startNext() {
if s.active != nil {
return // a job is already loading or serving
}
for len(s.queued) > 0 {
idx, score := s.pick()
q := s.queued[idx]
s.queued = append(s.queued[:idx], s.queued[idx+1:]...)
s.queueChanged()
req := q.req
state, ok := s.effects.ModelState(req.Model)
if !ok {
s.effects.GrantError(req, ErrModelNotFound)
continue
}
waited := s.now().Sub(q.enqueued)
s.annotate(req, score, waited)
r := req
s.active = &r
evict := s.otherRunning(req.Model)
if state == process.StateReady && len(evict) == 0 {
// Already loaded and the only model running — serve immediately.
s.logger.Debugf("%s: serving model %s (already loaded, priority %d, score %d, waited %s)",
s.name, req.Model, req.Priority, score, waited.Round(time.Millisecond))
if s.serve() {
s.recordDispatch(req.Priority)
return
}
continue // caller gone; pick the next request
}
s.logger.Debugf("%s: swapping to model %s (priority %d, score %d, waited %s), evicting %v",
s.name, req.Model, req.Priority, score, waited.Round(time.Millisecond), evict)
s.phase = phaseSwapping
s.recordDispatch(req.Priority)
s.effects.StartSwap(req.Model, evict)
return
}
}
// pick returns the index of the highest scoring queued request and its score.
// It also accounts for whether the aging and swap-affinity terms changed the
// outcome: each is recomputed with that term removed, and a different winner
// means the term reordered this dispatch. The queue is never empty here.
func (s *Serial) pick() (int, int) {
now := s.now()
resident, _ := s.noSwapModel()
idx, score := s.best(now, resident, true, true)
if len(s.queued) > 1 {
if other, _ := s.best(now, resident, false, true); other != idx {
s.countReorder(true)
}
if other, _ := s.best(now, resident, true, false); other != idx {
s.countReorder(false)
}
}
return idx, score
}
// best returns the index and score of the highest scoring queued request, with
// the aging and swap-affinity terms individually switchable so pick can measure
// their effect. Ties keep the earliest arrival because the queue is in arrival
// order and the comparison is strictly greater-than.
func (s *Serial) best(now time.Time, resident string, withAging, withAffinity bool) (int, int) {
bestIdx, bestScore := 0, 0
for i, q := range s.queued {
score := q.req.Priority
if withAffinity && resident != "" && q.req.Model == resident {
score += s.swapAffinityBonus
}
if withAging {
score += s.aging(q, now)
}
if i == 0 || score > bestScore {
bestIdx, bestScore = i, score
}
}
return bestIdx, bestScore
}
// aging converts how long a request has waited into priority points. It is
// unbounded: a request that waits long enough eventually outranks anything.
func (s *Serial) aging(q queuedReq, now time.Time) int {
if s.agingDivisor <= 0 {
return 0
}
waited := now.Sub(q.enqueued)
if waited <= 0 {
return 0
}
return int(waited / s.agingDivisor)
}
// noSwapModel returns the model that can be served without any swap: the sole
// running process, when it is ready. Anything else — nothing loaded, a process
// still loading, or more than one running — means every request costs a swap,
// so no request earns the affinity bonus.
func (s *Serial) noSwapModel() (string, bool) {
if s.swapAffinityBonus == 0 {
return "", false
}
running := s.effects.RunningModels()
if len(running) != 1 {
return "", false
}
for id, state := range running {
if state == process.StateReady {
return id, true
}
}
return "", false
}
// annotate records the dispatch decision on the request's context so it reaches
// the activity log. This is how "why did my request take 20 minutes" gets
// answered after the fact.
func (s *Serial) annotate(req HandlerReq, score int, waited time.Duration) {
if req.Ctx == nil {
return
}
fields := [...]struct{ key, value string }{
{"serial_priority", strconv.Itoa(req.Priority)},
{"serial_band", shared.PriorityBand(req.Priority)},
{"serial_score", strconv.Itoa(score)},
{"serial_queue_wait_ms", strconv.FormatInt(waited.Milliseconds(), 10)},
}
for _, f := range fields {
if err := shared.SetReqData(req.Ctx, f.key, f.value); err != nil {
// Every key shares one context, so the first failure means none
// of them will land.
s.logger.Debugf("%s: failed to set %s metadata: %v", s.name, f.key, err)
return
}
}
}
// serve hands the active request its tracked handler. It returns true when the
// request is now serving (await OnServeDone); false when the caller had already
// disconnected, in which case active is cleared so the next job can start.
func (s *Serial) serve() bool {
if s.effects.GrantServe(*s.active, s.active.Model) {
s.phase = phaseServing
return true
}
s.logger.Debugf("%s: caller for model %s gone before serve", s.name, s.active.Model)
s.active = nil
s.phase = phaseIdle
return false
}
// OnSwapDone fires when the load for the active request completes. On success the
// request is served; on failure its caller receives the error and the queue
// advances. A SwapDone that does not match the active load (e.g. its request was
// unloaded or cancelled mid-load) is ignored.
func (s *Serial) OnSwapDone(ev SwapDone) {
if s.active == nil || s.phase != phaseSwapping || s.active.Model != ev.ModelID {
return
}
if ev.Err != nil {
s.logger.Debugf("%s: swap for model %s failed: %v", s.name, ev.ModelID, ev.Err)
s.effects.GrantError(*s.active, ev.Err)
s.active = nil
s.phase = phaseIdle
s.startNext()
return
}
if !s.serve() {
s.startNext() // caller vanished while the model loaded; move on
}
}
// OnServeDone fires when the active request's handler returns. The slot is freed
// and the next queued request begins.
func (s *Serial) OnServeDone(ev ServeDoneEvent) {
if s.active == nil || s.phase != phaseServing {
return
}
s.active = nil
s.phase = phaseIdle
s.startNext()
}
// OnCancel removes a disconnected client's request from the queue. A request that
// is already active is left to finish: if it was loading, OnSwapDone's serve()
// will find the caller gone (GrantServe false) and advance; if it was serving,
// its handler returns normally and reaches OnServeDone.
func (s *Serial) OnCancel(req HandlerReq) {
if len(s.queued) == 0 {
return
}
kept := s.queued[:0]
removed := false
for _, q := range s.queued {
if q.req.Respond == req.Respond {
removed = true
continue
}
kept = append(kept, q)
}
s.queued = kept
if removed {
s.logger.Debugf("%s: cancelled request for model %s pruned from queue", s.name, req.Model)
s.queueChanged()
}
}
// OnUnload reconciles state for an unload, stops the targeted processes, and
// advances the queue. It mirrors the FIFO contract: queued requests for unloaded
// models are failed; an active *loading* request for an unloaded model is failed
// (its swap goroutine is left to finish and its SwapDone is then ignored); an
// active *serving* request is left for its handler to end when StopProcesses
// kills the upstream. The Stop is synchronous so callers of Unload can rely on
// the processes being stopped on return.
func (s *Serial) OnUnload(targets []string, timeout time.Duration) {
unloadErr := fmt.Errorf("%s: model unloaded", s.name)
targetSet := make(map[string]bool, len(targets))
for _, id := range targets {
targetSet[id] = true
}
if s.active != nil && s.phase == phaseSwapping && targetSet[s.active.Model] {
s.effects.GrantError(*s.active, unloadErr)
s.active = nil
s.phase = phaseIdle
}
if len(s.queued) > 0 {
kept := s.queued[:0]
for _, q := range s.queued {
if targetSet[q.req.Model] {
s.effects.GrantError(q.req, unloadErr)
continue
}
kept = append(kept, q)
}
s.queued = kept
s.queueChanged()
}
s.effects.StopProcesses(timeout, targets)
// A still-serving active request advances via OnServeDone when its killed
// handler returns; only start the next job when nothing is active now.
if s.active == nil {
s.startNext()
}
}
// OnShutdown grants err to every request the scheduler still holds: an active
// loading request and all queued requests. A serving request is torn down with
// its process by the baseRouter.
func (s *Serial) OnShutdown(err error) {
if s.active != nil && s.phase == phaseSwapping {
s.effects.GrantError(*s.active, err)
s.active = nil
s.phase = phaseIdle
}
for _, q := range s.queued {
s.effects.GrantError(q.req, err)
}
s.queued = nil
s.queueChanged()
}
// otherRunning returns every running model except target, sorted for
// deterministic eviction.
func (s *Serial) otherRunning(target string) []string {
var out []string
for id := range s.effects.RunningModels() {
if id != target {
out = append(out, id)
}
}
sort.Strings(out)
return out
}
// queueChanged is called after every mutation of s.queued. It refreshes the
// stats mirror and tells each waiter its new position.
func (s *Serial) queueChanged() {
s.syncStats()
s.broadcastPositions()
}
// broadcastPositions sends each waiter its 1-indexed place in dispatch order
// rather than arrival order, so a queued caller sees the position priority
// actually earned it. The ranking is a snapshot: aging keeps moving, so a
// position only holds until the next queue change.
func (s *Serial) broadcastPositions() {
if len(s.queued) == 0 {
return
}
now := s.now()
resident, _ := s.noSwapModel()
scores := make([]int, len(s.queued))
order := make([]int, len(s.queued))
for i, q := range s.queued {
score := q.req.Priority
if resident != "" && q.req.Model == resident {
score += s.swapAffinityBonus
}
scores[i] = score + s.aging(q, now)
order[i] = i
}
// Stable so equal scores keep arrival order, matching best().
sort.SliceStable(order, func(a, b int) bool { return scores[order[a]] > scores[order[b]] })
ranked := make([]HandlerReq, len(order))
for rank, i := range order {
ranked[rank] = s.queued[i].req
}
broadcastQueuePositions(ranked)
}
// syncStats refreshes the per-band depth and oldest-arrival mirror that
// QueueStats reads. Enqueue times are stored rather than durations so the wait
// is computed fresh at read time instead of going stale between queue changes.
func (s *Serial) syncStats() {
depth := make(map[string]int, len(shared.PriorityBands))
oldest := make(map[string]time.Time, len(shared.PriorityBands))
for _, q := range s.queued {
band := shared.PriorityBand(q.req.Priority)
depth[band]++
if t, ok := oldest[band]; !ok || q.enqueued.Before(t) {
oldest[band] = q.enqueued
}
}
s.mu.Lock()
defer s.mu.Unlock()
s.stats.depth = depth
s.stats.oldest = oldest
}
// recordDispatch counts one request handed off, bucketed by band.
func (s *Serial) recordDispatch(priority int) {
band := shared.PriorityBand(priority)
s.mu.Lock()
defer s.mu.Unlock()
s.stats.dispatched[band]++
}
// countReorder records that aging (or swap affinity) changed which request a
// dispatch picked.
func (s *Serial) countReorder(aging bool) {
s.mu.Lock()
defer s.mu.Unlock()
if aging {
s.stats.agingReorders++
} else {
s.stats.affinityReorders++
}
}
// QueueStats implements StatsReporter. It is safe to call from any goroutine.
func (s *Serial) QueueStats() QueueStats {
now := s.now()
s.mu.Lock()
defer s.mu.Unlock()
bands := make(map[string]BandStats, len(shared.PriorityBands))
for _, band := range shared.PriorityBands {
bs := BandStats{
Depth: s.stats.depth[band],
Dispatched: s.stats.dispatched[band],
}
if t, ok := s.stats.oldest[band]; ok {
if wait := now.Sub(t); wait > 0 {
bs.OldestWait = wait
}
}
bands[band] = bs
}
return QueueStats{
Bands: bands,
AgingReorders: s.stats.agingReorders,
AffinityReorders: s.stats.affinityReorders,
}
}
+731
View File
@@ -0,0 +1,731 @@
package scheduler
import (
"context"
"errors"
"io"
"testing"
"time"
"github.com/mostlygeek/llama-swap/internal/config"
"github.com/mostlygeek/llama-swap/internal/logmon"
"github.com/mostlygeek/llama-swap/internal/process"
"github.com/mostlygeek/llama-swap/internal/shared"
)
// Serial methods all run on the router's single run-loop goroutine, so these
// tests drive them directly and synchronously, reusing fakeEffects and the
// req/reqCh helpers from fifo_test.go. A load completes via OnSwapDone and a
// served request finishes via OnServeDone — the events the run loop delivers.
// newSerial builds a Serial with the production defaults: aging at one point
// per minute, swap affinity at +10.
func newSerial(eff Effects) *Serial {
return NewSerial("test", logmon.NewWriter(io.Discard), 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.
func lastStart(t *testing.T, eff *fakeEffects) startRec {
t.Helper()
if len(eff.starts) == 0 {
t.Fatal("no StartSwap recorded")
}
return eff.starts[len(eff.starts)-1]
}
func sameSet(a, b []string) bool {
if len(a) != len(b) {
return false
}
m := map[string]int{}
for _, x := range a {
m[x]++
}
for _, x := range b {
m[x]--
}
for _, v := range m {
if v != 0 {
return false
}
}
return true
}
// servedOrder returns the model IDs of every successful serve grant in order.
func servedOrder(eff *fakeEffects) []string {
var out []string
for _, g := range eff.grants {
if g.err == nil && g.serve {
out = append(out, g.model)
}
}
return out
}
func TestSerial_FastPath_AlreadyLoaded(t *testing.T) {
eff := newFakeEffects()
eff.states["a"] = process.StateReady
s := newSerial(eff)
s.OnRequest(req("a"))
if got := len(eff.starts); got != 0 {
t.Errorf("StartSwap calls=%d want 0 (already loaded, no swap)", got)
}
if got := eff.served("a"); got != 1 {
t.Errorf("served(a)=%d want 1", got)
}
}
func TestSerial_ColdStart_LoadsThenServes(t *testing.T) {
eff := newFakeEffects()
eff.states["a"] = process.StateStopped
s := newSerial(eff)
s.OnRequest(req("a"))
if got := eff.startsFor("a"); got != 1 {
t.Fatalf("StartSwap(a)=%d want 1", got)
}
if got := eff.served("a"); got != 0 {
t.Errorf("served(a)=%d want 0 before load completes", got)
}
eff.states["a"] = process.StateReady
s.OnSwapDone(SwapDone{ModelID: "a"})
if got := eff.served("a"); got != 1 {
t.Errorf("served(a)=%d want 1 after load", got)
}
}
func TestSerial_UnknownModel(t *testing.T) {
eff := newFakeEffects() // no states => unknown
s := newSerial(eff)
s.OnRequest(req("ghost"))
if len(eff.starts) != 0 {
t.Errorf("StartSwap calls=%d want 0", len(eff.starts))
}
if eff.errored("ghost") != 1 {
t.Fatalf("errored(ghost)=%d want 1", eff.errored("ghost"))
}
if !errors.Is(eff.grants[0].err, ErrModelNotFound) {
t.Errorf("err=%v want ErrModelNotFound", eff.grants[0].err)
}
}
func TestSerial_EvictsEveryOtherModel(t *testing.T) {
eff := newFakeEffects()
eff.states["x"] = process.StateReady // already running
eff.states["y"] = process.StateReady // also running (e.g. left over)
eff.states["a"] = process.StateStopped
s := newSerial(eff)
s.OnRequest(req("a"))
st := lastStart(t, eff)
if st.model != "a" {
t.Fatalf("loading %s want a", st.model)
}
if !sameSet(st.evict, []string{"x", "y"}) {
t.Errorf("evict=%v want [x y] (serial evicts ALL other models)", st.evict)
}
}
// TestSerial_OneJobAtATime verifies a second request waits while the first is
// serving, and only starts after the first finishes.
func TestSerial_OneJobAtATime(t *testing.T) {
eff := newFakeEffects()
eff.states["a"] = process.StateReady
eff.states["b"] = process.StateStopped
s := newSerial(eff)
s.OnRequest(req("a")) // served immediately
s.OnRequest(req("b")) // must wait — a is serving
if got := eff.startsFor("b"); got != 0 {
t.Fatalf("StartSwap(b)=%d want 0 while a is serving", got)
}
if got := eff.served("a"); got != 1 {
t.Fatalf("served(a)=%d want 1", got)
}
// a finishes -> b may now load (evicting a).
s.OnServeDone(ServeDoneEvent{ModelID: "a"})
if got := eff.startsFor("b"); got != 1 {
t.Fatalf("StartSwap(b)=%d want 1 after a finished", got)
}
if st := lastStart(t, eff); !sameSet(st.evict, []string{"a"}) {
t.Errorf("b evict=%v want [a]", st.evict)
}
}
// TestSerial_SameModelConsecutive_NoReload verifies back-to-back requests for the
// already-loaded model run without a reload, one after another.
func TestSerial_SameModelConsecutive_NoReload(t *testing.T) {
eff := newFakeEffects()
eff.states["a"] = process.StateStopped
s := newSerial(eff)
s.OnRequest(req("a")) // cold load
s.OnRequest(req("a")) // queued behind the first
eff.states["a"] = process.StateReady
s.OnSwapDone(SwapDone{ModelID: "a"}) // first serves
if got := eff.served("a"); got != 1 {
t.Fatalf("served(a)=%d want 1 (one at a time)", got)
}
s.OnServeDone(ServeDoneEvent{ModelID: "a"}) // first done -> second serves
if got := eff.served("a"); got != 2 {
t.Fatalf("served(a)=%d want 2", got)
}
if got := eff.startsFor("a"); got != 1 {
t.Errorf("StartSwap(a)=%d want 1 (second request must not reload)", got)
}
}
// TestSerial_StrictArrivalOrder covers the scoring function's degenerate case:
// with equal priorities and swap affinity disabled, qwen36, qwen35, sdxl,
// qwen36 execute in EXACTLY that order with evictions between each model
// switch, including reloading qwen36 at the end even though it ran first.
// Aging cannot reorder them either — they all arrive at the same instant, so
// they age at the same rate.
func TestSerial_StrictArrivalOrder(t *testing.T) {
eff := newFakeEffects()
for _, m := range []string{"qwen36", "qwen35", "sdxl"} {
eff.states[m] = process.StateStopped
}
s := newSerialCfg(eff, config.DefaultAgingDivisor, 0)
for _, m := range []string{"qwen36", "qwen35", "sdxl", "qwen36"} {
s.OnRequest(req(m))
}
// Only the first job starts loading; the rest wait their turn.
if len(eff.starts) != 1 || eff.starts[0].model != "qwen36" {
t.Fatalf("starts=%+v want only [qwen36] loading first", eff.starts)
}
// step completes the current model's load+serve and returns control to the
// scheduler, which must start the next queued model.
step := func(model string, wantEvict []string) {
t.Helper()
st := lastStart(t, eff)
if st.model != model {
t.Fatalf("loading %q want %q", st.model, model)
}
if !sameSet(st.evict, wantEvict) {
t.Fatalf("loading %q evict=%v want %v", model, st.evict, wantEvict)
}
// Simulate the eviction + load actually happening.
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})
}
step("qwen36", nil) // cold load, nothing else running
step("qwen35", []string{"qwen36"}) // evict qwen36
step("sdxl", []string{"qwen35"}) // evict qwen35
step("qwen36", []string{"sdxl"}) // RELOAD qwen36, evict sdxl
want := []string{"qwen36", "qwen35", "sdxl", "qwen36"}
if got := servedOrder(eff); !sameOrder(got, want) {
t.Fatalf("serve order=%v want %v", got, want)
}
}
func sameOrder(a, b []string) bool {
if len(a) != len(b) {
return false
}
for i := range a {
if a[i] != b[i] {
return false
}
}
return true
}
// stepModel completes the current load+serve for model and returns control to
// the scheduler so it dispatches the next queued request.
func stepModel(t *testing.T, s *Serial, eff *fakeEffects, model string) {
t.Helper()
if len(eff.starts) > 0 {
if st := eff.starts[len(eff.starts)-1]; st.model == model {
for _, e := range st.evict {
eff.states[e] = process.StateStopped
}
eff.states[model] = process.StateReady
s.OnSwapDone(SwapDone{ModelID: model})
}
}
s.OnServeDone(ServeDoneEvent{ModelID: model})
}
// TestSerial_HigherPriorityDispatchesFirst is the point of the whole exercise:
// a batch job already queued must yield to an interactive request that arrives
// later, because dispatch order is by score, not arrival.
func TestSerial_HigherPriorityDispatchesFirst(t *testing.T) {
eff := newFakeEffects()
for _, m := range []string{"running", "batch", "interactive"} {
eff.states[m] = process.StateStopped
}
s := newSerial(eff)
s.OnRequest(reqP("running", shared.PriorityNormal)) // dispatches immediately
s.OnRequest(reqP("batch", shared.PriorityBatch)) // queued first...
s.OnRequest(reqP("interactive", shared.PriorityInteractive)) // ...but this jumps it
stepModel(t, s, eff, "running")
if got := eff.startsFor("interactive"); got != 1 {
t.Fatalf("StartSwap(interactive)=%d want 1 (must overtake the queued batch job)", got)
}
if got := eff.startsFor("batch"); got != 0 {
t.Fatalf("StartSwap(batch)=%d want 0 (still waiting behind interactive)", got)
}
stepModel(t, s, eff, "interactive")
if got := eff.startsFor("batch"); got != 1 {
t.Fatalf("StartSwap(batch)=%d want 1 once nothing outranks it", got)
}
}
// TestSerial_TierOffsetBreaksTieWithinBand verifies the composition rule the
// design depends on: a small per-caller offset orders requests inside a band
// and never promotes one across a band. A "max member" batch job at -98 beats
// other batch work but still loses to every normal request at 0.
func TestSerial_TierOffsetBreaksTieWithinBand(t *testing.T) {
eff := newFakeEffects()
for _, m := range []string{"running", "free", "max", "normal"} {
eff.states[m] = process.StateStopped
}
s := newSerialCfg(eff, 0, 0) // isolate the priority term
s.OnRequest(reqP("running", 0))
s.OnRequest(reqP("free", shared.PriorityBatch)) // -100
s.OnRequest(reqP("max", shared.PriorityBatch+2)) // -98, max tier
s.OnRequest(reqP("normal", shared.PriorityNormal)) // 0
stepModel(t, s, eff, "running")
if got := eff.startsFor("normal"); got != 1 {
t.Fatalf("StartSwap(normal)=%d want 1 (a tier bonus must not cross a band)", got)
}
stepModel(t, s, eff, "normal")
if got := eff.startsFor("max"); got != 1 {
t.Fatalf("StartSwap(max)=%d want 1 (max tier outranks free within the batch band)", got)
}
stepModel(t, s, eff, "max")
if got := eff.startsFor("free"); got != 1 {
t.Fatalf("StartSwap(free)=%d want 1 (last)", got)
}
}
// TestSerial_AgingPreventsStarvation verifies the one term allowed to cross
// bands: a batch job that has waited long enough eventually beats an
// interactive request that arrived just now.
func TestSerial_AgingPreventsStarvation(t *testing.T) {
eff := newFakeEffects()
for _, m := range []string{"running", "batch", "interactive"} {
eff.states[m] = process.StateStopped
}
s := newSerialCfg(eff, 60, 0) // one point per minute, no affinity
clock := newFakeClock(s)
s.OnRequest(reqP("running", shared.PriorityNormal))
s.OnRequest(reqP("batch", shared.PriorityBatch)) // -100
// The batch job waits long enough to gain 201 points: -100 + 201 = 101,
// just past a fresh interactive request at +100.
clock.advance(201 * time.Minute)
s.OnRequest(reqP("interactive", shared.PriorityInteractive))
stepModel(t, s, eff, "running")
if got := eff.startsFor("batch"); got != 1 {
t.Fatalf("StartSwap(batch)=%d want 1 (aging must eventually beat a fresh interactive request)", got)
}
if got := eff.startsFor("interactive"); got != 0 {
t.Fatalf("StartSwap(interactive)=%d want 0 (outranked by the aged batch job)", got)
}
}
// TestSerial_AgingCannotCrossBandTooEarly is the other half of aging: a batch
// job that has waited only a little still loses to interactive traffic.
func TestSerial_AgingCannotCrossBandTooEarly(t *testing.T) {
eff := newFakeEffects()
for _, m := range []string{"running", "batch", "interactive"} {
eff.states[m] = process.StateStopped
}
s := newSerialCfg(eff, 60, 0)
clock := newFakeClock(s)
s.OnRequest(reqP("running", shared.PriorityNormal))
s.OnRequest(reqP("batch", shared.PriorityBatch))
clock.advance(30 * time.Minute) // -100 + 30 = -70
s.OnRequest(reqP("interactive", shared.PriorityInteractive))
stepModel(t, s, eff, "running")
if got := eff.startsFor("interactive"); got != 1 {
t.Fatalf("StartSwap(interactive)=%d want 1 (30 minutes is not enough aging)", got)
}
}
// TestSerial_SwapAffinity_PrefersResidentModel verifies the bounded bonus keeps
// the model cache from thrashing: with equal priorities, a queued request for
// the already-loaded model runs before earlier requests that would each force a
// cold load.
func TestSerial_SwapAffinity_PrefersResidentModel(t *testing.T) {
eff := newFakeEffects()
for _, m := range []string{"qwen36", "qwen35", "sdxl"} {
eff.states[m] = process.StateStopped
}
s := newSerial(eff) // affinity +10 by default
for _, m := range []string{"qwen36", "qwen35", "sdxl", "qwen36"} {
s.OnRequest(req(m))
}
stepModel(t, s, eff, "qwen36") // qwen36 now resident and ready
if got := eff.served("qwen36"); got != 2 {
t.Fatalf("served(qwen36)=%d want 2 (the trailing qwen36 should skip the swap queue)", got)
}
if len(eff.starts) != 1 {
t.Fatalf("starts=%+v want only the initial qwen36 load (no reload)", eff.starts)
}
stepModel(t, s, eff, "qwen36")
if got := eff.startsFor("qwen35"); got != 1 {
t.Fatalf("StartSwap(qwen35)=%d want 1 once no resident request remains", got)
}
}
// TestSerial_SwapAffinity_CannotCrossBand verifies the bonus is bounded: an
// interactive request that needs a cold load still beats a batch job for the
// already-resident model.
func TestSerial_SwapAffinity_CannotCrossBand(t *testing.T) {
eff := newFakeEffects()
eff.states["resident"] = process.StateReady
eff.states["cold"] = process.StateStopped
s := newSerialCfg(eff, 0, 99) // the largest bonus config permits
s.OnRequest(reqP("resident", shared.PriorityNormal)) // dispatches immediately
s.OnRequest(reqP("resident", shared.PriorityBatch)) // queued, would get +99
s.OnRequest(reqP("cold", shared.PriorityInteractive))
s.OnServeDone(ServeDoneEvent{ModelID: "resident"})
if got := eff.startsFor("cold"); got != 1 {
t.Fatalf("StartSwap(cold)=%d want 1 (-100+99 must still lose to +100)", got)
}
}
// TestSerial_DispatchSetsMetadata verifies the dispatch decision is recorded on
// the request context, which is what puts it in the activity log.
func TestSerial_DispatchSetsMetadata(t *testing.T) {
eff := newFakeEffects()
eff.states["hold"] = process.StateReady
eff.states["a"] = process.StateStopped
s := newSerialCfg(eff, 60, 0)
clock := newFakeClock(s)
s.OnRequest(req("hold")) // serves immediately, occupying the slot
ctx := shared.SetContext(context.Background(), shared.ReqContextData{ModelID: "a", Metadata: make(map[string]string)})
s.OnRequest(HandlerReq{Model: "a", Priority: shared.PriorityBatch, Ctx: ctx})
clock.advance(2 * time.Minute) // -100 + 2 aging points
s.OnServeDone(ServeDoneEvent{ModelID: "hold"})
// a is dispatched (and annotated) here; finish its load so it is granted.
eff.states["hold"] = process.StateStopped
eff.states["a"] = process.StateReady
s.OnSwapDone(SwapDone{ModelID: "a"})
data, ok := shared.ReadContext(eff.lastServeReq.Ctx)
if !ok {
t.Fatal("context data missing from granted request")
}
want := map[string]string{
"serial_priority": "-100",
"serial_band": shared.BandBatch,
"serial_score": "-98",
"serial_queue_wait_ms": "120000",
}
for k, v := range want {
if got := data.Metadata[k]; got != v {
t.Errorf("%s = %q, want %q", k, got, v)
}
}
}
// TestSerial_QueueStats checks the numbers the /metrics endpoint reports:
// per-band depth and wait while queued, dispatch counts after the fact, and the
// reorder counters that say whether the scoring terms changed anything.
func TestSerial_QueueStats(t *testing.T) {
eff := newFakeEffects()
for _, m := range []string{"running", "batch", "interactive"} {
eff.states[m] = process.StateStopped
}
s := newSerialCfg(eff, 60, 0)
clock := newFakeClock(s)
s.OnRequest(reqP("running", shared.PriorityNormal)) // dispatched, not queued
s.OnRequest(reqP("batch", shared.PriorityBatch))
clock.advance(5 * time.Minute)
s.OnRequest(reqP("interactive", shared.PriorityInteractive))
stats := s.QueueStats()
if got := stats.Bands[shared.BandBatch].Depth; got != 1 {
t.Errorf("batch depth=%d want 1", got)
}
if got := stats.Bands[shared.BandBatch].OldestWait; got != 5*time.Minute {
t.Errorf("batch oldest wait=%s want 5m", got)
}
if got := stats.Bands[shared.BandInteractive].Depth; got != 1 {
t.Errorf("interactive depth=%d want 1", got)
}
if got := stats.Bands[shared.BandNormal].Depth; got != 0 {
t.Errorf("normal depth=%d want 0 (it dispatched immediately)", got)
}
if got := stats.Bands[shared.BandNormal].Dispatched; got != 1 {
t.Errorf("normal dispatched=%d want 1", got)
}
// Dispatching interactive over the older batch job is priority doing its
// job, not aging: with only 5 minutes of aging the winner is unchanged.
stepModel(t, s, eff, "running")
stats = s.QueueStats()
if got := stats.Bands[shared.BandInteractive].Dispatched; got != 1 {
t.Errorf("interactive dispatched=%d want 1", got)
}
if stats.AgingReorders != 0 {
t.Errorf("aging reorders=%d want 0", stats.AgingReorders)
}
if stats.AffinityReorders != 0 {
t.Errorf("affinity reorders=%d want 0 (affinity disabled)", stats.AffinityReorders)
}
}
// TestSerial_QueueStats_CountsReorders verifies each scoring term is credited
// when it actually changes the dispatch decision.
func TestSerial_QueueStats_CountsReorders(t *testing.T) {
t.Run("aging", func(t *testing.T) {
eff := newFakeEffects()
for _, m := range []string{"running", "old", "new"} {
eff.states[m] = process.StateStopped
}
s := newSerialCfg(eff, 60, 0)
clock := newFakeClock(s)
s.OnRequest(reqP("running", 0))
s.OnRequest(reqP("old", -10))
clock.advance(30 * time.Minute) // old: -10 + 30 = 20
s.OnRequest(reqP("new", 0)) // new: 0
stepModel(t, s, eff, "running")
if got := eff.startsFor("old"); got != 1 {
t.Fatalf("StartSwap(old)=%d want 1 (aging promoted it)", got)
}
if got := s.QueueStats().AgingReorders; got != 1 {
t.Errorf("aging reorders=%d want 1", got)
}
})
t.Run("swap affinity", func(t *testing.T) {
eff := newFakeEffects()
eff.states["resident"] = process.StateReady
eff.states["cold"] = process.StateStopped
s := newSerialCfg(eff, 0, 10)
s.OnRequest(req("resident")) // dispatches immediately
s.OnRequest(reqP("cold", 5)) // higher priority, but needs a swap
s.OnRequest(reqP("resident", 0)) // +10 affinity beats it
s.OnServeDone(ServeDoneEvent{ModelID: "resident"})
if got := eff.served("resident"); got != 2 {
t.Fatalf("served(resident)=%d want 2 (affinity outranked the cold request)", got)
}
if got := s.QueueStats().AffinityReorders; got != 1 {
t.Errorf("affinity reorders=%d want 1", got)
}
})
}
func TestSerial_SwapError_FailsCallerAndAdvances(t *testing.T) {
eff := newFakeEffects()
eff.states["a"] = process.StateStopped
eff.states["b"] = process.StateStopped
s := newSerial(eff)
s.OnRequest(req("a"))
s.OnRequest(req("b")) // queued behind a
// a's load fails: its caller is errored and b proceeds.
s.OnSwapDone(SwapDone{ModelID: "a", Err: errors.New("boom")})
if eff.errored("a") != 1 {
t.Fatalf("errored(a)=%d want 1", eff.errored("a"))
}
if got := eff.startsFor("b"); got != 1 {
t.Fatalf("StartSwap(b)=%d want 1 after a's load failed", got)
}
}
// TestSerial_GrantServeFalse_Advances verifies that when the active request's
// caller has disconnected by serve time, the queue advances to the next request.
func TestSerial_GrantServeFalse_Advances(t *testing.T) {
eff := newFakeEffects()
eff.states["a"] = process.StateStopped
eff.states["b"] = process.StateStopped
eff.serveResult["a"] = false // a's caller is gone by grant time
s := newSerial(eff)
s.OnRequest(req("a"))
s.OnRequest(req("b")) // queued
eff.states["a"] = process.StateReady
s.OnSwapDone(SwapDone{ModelID: "a"}) // grant fails -> advance to b
if got := eff.served("a"); got != 0 {
t.Errorf("served(a)=%d want 0 (caller gone)", got)
}
if got := eff.startsFor("b"); got != 1 {
t.Fatalf("StartSwap(b)=%d want 1 (advanced after gone caller)", got)
}
}
func TestSerial_OnCancel_QueuedRequest(t *testing.T) {
eff := newFakeEffects()
eff.states["a"] = process.StateStopped
eff.states["b"] = process.StateStopped
s := newSerial(eff)
s.OnRequest(reqCh("a")) // starts loading a
cancelled := reqCh("b")
s.OnRequest(cancelled) // queued behind a
if len(s.queued) != 1 {
t.Fatalf("queued=%d want 1", len(s.queued))
}
s.OnCancel(cancelled)
if len(s.queued) != 0 {
t.Fatalf("queued=%d want 0 after cancel", len(s.queued))
}
// a completes; b is gone, so nothing starts for it.
eff.states["a"] = process.StateReady
s.OnSwapDone(SwapDone{ModelID: "a"})
s.OnServeDone(ServeDoneEvent{ModelID: "a"})
if got := eff.startsFor("b"); got != 0 {
t.Errorf("StartSwap(b)=%d want 0 (cancelled before its turn)", got)
}
}
func TestSerial_OnShutdown_FailsQueuedAndActiveLoad(t *testing.T) {
eff := newFakeEffects()
eff.states["a"] = process.StateStopped
eff.states["b"] = process.StateStopped
eff.states["c"] = process.StateStopped
s := newSerial(eff)
s.OnRequest(req("a")) // active (loading)
s.OnRequest(req("b")) // queued
s.OnRequest(req("c")) // queued
s.OnShutdown(errors.New("shutting down"))
if got := eff.errored(""); got != 3 {
t.Errorf("error grants=%d want 3 (active load + 2 queued)", got)
}
if len(s.queued) != 0 {
t.Errorf("queued=%d want 0 after shutdown", len(s.queued))
}
}
// TestSerial_OnUnload_WhileServing verifies that unloading the model that is
// actively serving does not strand the queue: OnUnload stops the process but
// leaves the active request to end via OnServeDone, which then advances.
func TestSerial_OnUnload_WhileServing(t *testing.T) {
eff := newFakeEffects()
eff.states["a"] = process.StateReady
eff.states["b"] = process.StateStopped
s := newSerial(eff)
s.OnRequest(req("a")) // served immediately (a ready)
s.OnRequest(req("b")) // queued behind a
if got := eff.served("a"); got != 1 {
t.Fatalf("served(a)=%d want 1", got)
}
// Unload a while it is serving: the process is stopped, but the queue must
// not advance yet — the active serve is still outstanding.
s.OnUnload([]string{"a"}, time.Second)
if len(eff.stops) != 1 || !sameSet(eff.stops[0].ids, []string{"a"}) {
t.Errorf("StopProcesses=%+v want one call stopping [a]", eff.stops)
}
if got := eff.startsFor("b"); got != 0 {
t.Fatalf("StartSwap(b)=%d want 0 before the serving request ends", got)
}
// The killed handler returns -> OnServeDone advances to b.
eff.states["a"] = process.StateStopped
s.OnServeDone(ServeDoneEvent{ModelID: "a"})
if got := eff.startsFor("b"); got != 1 {
t.Fatalf("StartSwap(b)=%d want 1 after the serving request ended", got)
}
}
func TestSerial_OnUnload_DropsQueuedAndStops(t *testing.T) {
eff := newFakeEffects()
eff.states["a"] = process.StateStopped
eff.states["b"] = process.StateStopped
s := newSerial(eff)
s.OnRequest(req("a")) // active (loading a)
s.OnRequest(req("b")) // queued
// Unload a: its active load is failed and a is stopped.
s.OnUnload([]string{"a"}, time.Second)
if eff.errored("a") != 1 {
t.Errorf("errored(a)=%d want 1 (active load failed)", eff.errored("a"))
}
if len(eff.stops) != 1 || !sameSet(eff.stops[0].ids, []string{"a"}) {
t.Errorf("StopProcesses=%+v want one call stopping [a]", eff.stops)
}
// b was queued and not unloaded; with a's load cancelled it now starts.
if got := eff.startsFor("b"); got != 1 {
t.Errorf("StartSwap(b)=%d want 1 after unload advanced the queue", got)
}
}
+45 -31
View File
@@ -2,6 +2,7 @@ package server
import (
"encoding/json"
"fmt"
"net/http"
"sort"
"strings"
@@ -9,6 +10,7 @@ import (
"github.com/mostlygeek/llama-swap/internal/config"
"github.com/mostlygeek/llama-swap/internal/event"
"github.com/mostlygeek/llama-swap/internal/process"
"github.com/mostlygeek/llama-swap/internal/shared"
)
@@ -86,6 +88,12 @@ func renderCapabilities(caps config.ModelCapConfig) (arch map[string]any, capsMa
if contains(caps.In, "image") && contains(caps.Out, "image") {
capsMap["image_to_image"] = true
}
if contains(caps.In, "text") && contains(caps.Out, "video") {
capsMap["video_generation"] = true
}
if contains(caps.In, "image") && contains(caps.Out, "video") {
capsMap["image_to_video"] = true
}
}
if caps.Tools {
@@ -285,15 +293,25 @@ func (s *Server) startPreload() {
}()
}
// handleMetrics serves Prometheus-format performance metrics. Returns 503 when
// performance monitoring is disabled.
// handleMetrics serves Prometheus-format metrics: system/GPU performance plus
// the scheduler's request queue. Returns 503 only when neither source is
// available — performance monitoring disabled and a scheduler that does not
// report queue stats.
func (s *Server) handleMetrics(w http.ResponseWriter, r *http.Request) {
if s.perf == nil {
queueStats, hasQueueStats := s.local.SchedulerStats()
if s.perf == nil && !hasQueueStats {
w.WriteHeader(http.StatusServiceUnavailable)
w.Write([]byte("# performance monitor not available\n"))
return
}
s.perf.MetricsHandler().ServeHTTP(w, r)
w.Header().Set("Content-Type", "text/plain; version=0.0.4; charset=utf-8")
if s.perf != nil {
s.perf.MetricsHandler().ServeHTTP(w, r)
}
if hasQueueStats {
writeSchedulerMetrics(w, queueStats)
}
}
func handleHealth(w http.ResponseWriter, r *http.Request) {
@@ -314,7 +332,7 @@ func handleUpstreamRedirect(w http.ResponseWriter, r *http.Request) {
func (s *Server) handleUpstream(w http.ResponseWriter, r *http.Request) {
upstreamPath := r.PathValue("upstreamPath")
searchName, modelID, remainingPath, found := findModelInPath(s.cfg, "/"+upstreamPath)
searchName, modelID, remainingPath, found := shared.FindModelInPath(s.cfg, "/"+upstreamPath)
if !found {
shared.SendResponse(w, r, http.StatusNotFound, "model not found")
return
@@ -340,6 +358,28 @@ func (s *Server) handleUpstream(w http.ResponseWriter, r *http.Request) {
// Pin the resolved model so the router skips body/query extraction.
*r = *r.WithContext(shared.SetContext(r.Context(), shared.ReqContextData{Model: searchName, ModelID: modelID, Metadata: make(map[string]string)}))
// If the path matches an upstream.ignorePaths entry and the model is
// not already loaded, refuse the request without triggering a swap. The
// server was not able to process the response because the model was not
// already loaded.
for _, re := range s.cfg.Upstream.IgnorePaths {
if !re.MatchString(remainingPath) {
continue
}
if s.local.Handles(modelID) {
state, ok := s.local.RunningModels()[modelID]
if !ok || state != process.StateReady {
shared.SendResponse(w, r, http.StatusConflict,
fmt.Sprintf("model %s is not loaded; path matches upstream.ignorePaths", modelID))
return
}
}
// Either the model is already loaded (no swap would be triggered)
// or this is a peer model (peer proxying never swaps). Fall through
// to normal dispatch.
break
}
switch {
case s.local.Handles(modelID):
s.local.ServeHTTP(w, r)
@@ -349,29 +389,3 @@ func (s *Server) handleUpstream(w http.ResponseWriter, r *http.Request) {
shared.SendResponse(w, r, http.StatusNotFound, "no router for model "+modelID)
}
}
// findModelInPath walks a slash-separated path, building up segments until one
// matches a configured model. This resolves model names that contain slashes
// (e.g. "author/model"). Returns the matched name, its real model ID, the
// remaining path, and whether a match was found.
func findModelInPath(cfg config.Config, path string) (searchName, realName, remainingPath string, found bool) {
parts := strings.Split(strings.TrimSpace(path), "/")
name := ""
for i, part := range parts {
if part == "" {
continue
}
if name == "" {
name = part
} else {
name = name + "/" + part
}
if modelID, ok := cfg.RealModelName(name); ok {
return name, modelID, "/" + strings.Join(parts[i+1:], "/"), true
}
}
return "", "", "", false
}
+211 -2
View File
@@ -2,11 +2,19 @@ package server
import (
"encoding/json"
"io"
"net/http"
"net/http/httptest"
"regexp"
"strings"
"testing"
"time"
"github.com/mostlygeek/llama-swap/internal/config"
"github.com/mostlygeek/llama-swap/internal/logmon"
"github.com/mostlygeek/llama-swap/internal/process"
"github.com/mostlygeek/llama-swap/internal/router/scheduler"
"github.com/mostlygeek/llama-swap/internal/shared"
)
func TestServer_HandleListModels(t *testing.T) {
@@ -78,6 +86,7 @@ func TestServer_HandleListModels_Aliases(t *testing.T) {
func TestServer_FindModelInPath(t *testing.T) {
cfg := config.Config{Models: map[string]config.ModelConfig{
"author": {},
"author/model": {},
"simple": {},
}}
@@ -91,13 +100,14 @@ func TestServer_FindModelInPath(t *testing.T) {
{"/simple/v1/chat", "simple", "/v1/chat", true},
{"/author/model/v1/chat", "author/model", "/v1/chat", true},
{"/author/model", "author/model", "/", true},
{"/author/v1/chat", "author", "/v1/chat", true},
{"/missing/v1", "", "", false},
{"/", "", "", false},
}
for _, c := range cases {
name, _, rem, found := findModelInPath(cfg, c.path)
name, _, rem, found := shared.FindModelInPath(cfg, c.path)
if found != c.wantFound || name != c.wantName || (found && rem != c.wantRem) {
t.Errorf("findModelInPath(%q) = (%q,%q,%v), want (%q,%q,%v)",
t.Errorf("FindModelInPath(%q) = (%q,%q,%v), want (%q,%q,%v)",
c.path, name, rem, found, c.wantName, c.wantRem, c.wantFound)
}
}
@@ -133,7 +143,168 @@ func TestServer_HandleUpstream(t *testing.T) {
})
}
func upstreamMetricsServer(response string) *Server {
cfg := config.Config{Models: map[string]config.ModelConfig{"m1": {}}}
proxylog := logmon.NewWriter(io.Discard)
s := &Server{
cfg: cfg,
muxlog: logmon.NewWriter(io.Discard),
proxylog: proxylog,
upstreamlog: logmon.NewWriter(io.Discard),
inflight: &inflightCounter{},
metrics: newMetricsMonitor(proxylog, 10, 0),
local: newStubRouter([]string{"m1"}, response),
peer: newStubRouter(nil, ""),
}
s.routes()
return s
}
func TestServer_HandleUpstream_IgnorePaths(t *testing.T) {
// Compile a pattern that matches static asset suffixes.
pattern := regexp.MustCompile(`.*\.(js|json|css|png|gif|jpg|jpeg|txt)$`)
t.Run("matched path, model not loaded, returns 409", func(t *testing.T) {
local := newStubRouter([]string{"m1"}, "upstream-body")
// running is nil/empty: model is not in RunningModels() => not loaded.
s := newTestServer(local, newStubRouter(nil, ""))
s.cfg = config.Config{
Models: map[string]config.ModelConfig{"m1": {}},
Upstream: config.UpstreamConfig{
IgnorePaths: []*regexp.Regexp{pattern},
},
}
w := httptest.NewRecorder()
s.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/upstream/m1/foo.js", nil))
if w.Code != http.StatusConflict {
t.Fatalf("status = %d, want %d (body=%q)", w.Code, http.StatusConflict, w.Body.String())
}
if !strings.Contains(w.Body.String(), "not loaded") {
t.Errorf("body = %q, want it to contain 'not loaded'", w.Body.String())
}
})
t.Run("matched path, model already loaded, serves normally", func(t *testing.T) {
local := newStubRouter([]string{"m1"}, "upstream-body")
local.running = map[string]process.ProcessState{"m1": process.StateReady}
s := newTestServer(local, newStubRouter(nil, ""))
s.cfg = config.Config{
Models: map[string]config.ModelConfig{"m1": {}},
Upstream: config.UpstreamConfig{
IgnorePaths: []*regexp.Regexp{pattern},
},
}
w := httptest.NewRecorder()
s.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/upstream/m1/foo.js", nil))
if w.Code != http.StatusOK || w.Body.String() != "upstream-body" {
t.Fatalf("status=%d body=%q, want 200 'upstream-body'", w.Code, w.Body.String())
}
})
t.Run("non-matched path, model not loaded, serves normally", func(t *testing.T) {
local := newStubRouter([]string{"m1"}, "upstream-body")
s := newTestServer(local, newStubRouter(nil, ""))
s.cfg = config.Config{
Models: map[string]config.ModelConfig{"m1": {}},
Upstream: config.UpstreamConfig{
IgnorePaths: []*regexp.Regexp{pattern},
},
}
w := httptest.NewRecorder()
s.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/upstream/m1/v1/chat/completions", nil))
if w.Code != http.StatusOK || w.Body.String() != "upstream-body" {
t.Fatalf("status=%d body=%q, want 200 'upstream-body'", w.Code, w.Body.String())
}
})
t.Run("matched path, peer model, serves normally", func(t *testing.T) {
// Peer routers do not appear via RunningModels on the local router;
// they should fall through to normal dispatch without 409.
local := newStubRouter(nil, "")
peer := newStubRouter([]string{"m1"}, "peer-body")
s := newTestServer(local, peer)
s.cfg = config.Config{
Models: map[string]config.ModelConfig{"m1": {}},
Upstream: config.UpstreamConfig{
IgnorePaths: []*regexp.Regexp{pattern},
},
}
w := httptest.NewRecorder()
s.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/upstream/m1/foo.js", nil))
if w.Code != http.StatusOK || w.Body.String() != "peer-body" {
t.Fatalf("status=%d body=%q, want 200 'peer-body'", w.Code, w.Body.String())
}
})
}
func TestServer_HandleUpstream_MetricsRecordsSupportedPath(t *testing.T) {
resp := `{"usage":{"prompt_tokens":3,"completion_tokens":5}}`
s := upstreamMetricsServer(resp)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodPost, "/upstream/m1/v1/chat/completions", strings.NewReader(`{}`))
req.Header.Set("Content-Type", "application/json")
s.ServeHTTP(w, req)
if w.Code != http.StatusOK || w.Body.String() != resp {
t.Fatalf("status=%d body=%q", w.Code, w.Body.String())
}
entries := s.metrics.getMetrics()
if len(entries) != 1 {
t.Fatalf("want 1 metrics entry, got %d", len(entries))
}
if entries[0].Model != "m1" {
t.Errorf("model = %q, want m1", entries[0].Model)
}
if entries[0].ReqPath != "/v1/chat/completions" {
t.Errorf("req_path = %q, want /v1/chat/completions", entries[0].ReqPath)
}
if entries[0].Tokens.InputTokens != 3 || entries[0].Tokens.OutputTokens != 5 {
t.Errorf("tokens = %+v, want input=3 output=5", entries[0].Tokens)
}
}
func TestServer_HandleUpstream_MetricsSkipsUnsupportedPath(t *testing.T) {
s := upstreamMetricsServer("ok")
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodPost, "/upstream/m1/probe", strings.NewReader(`{}`))
req.Header.Set("Content-Type", "application/json")
s.ServeHTTP(w, req)
if w.Code != http.StatusOK || w.Body.String() != "ok" {
t.Fatalf("status=%d body=%q", w.Code, w.Body.String())
}
if len(s.metrics.getMetrics()) != 0 {
t.Errorf("want no metrics entries for unsupported path, got %d", len(s.metrics.getMetrics()))
}
}
func TestServer_HandleUpstream_MetricsSkipsGET(t *testing.T) {
s := upstreamMetricsServer(`{"usage":{}}`)
w := httptest.NewRecorder()
s.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/upstream/m1/v1/chat/completions", nil))
if w.Code != http.StatusOK {
t.Fatalf("status=%d", w.Code)
}
if len(s.metrics.getMetrics()) != 0 {
t.Errorf("want no metrics entries for GET upstream, got %d", len(s.metrics.getMetrics()))
}
}
func TestServer_HandleMetrics_Unavailable(t *testing.T) {
// No perf monitor and a scheduler that reports no queue stats: nothing to
// serve.
s := newTestServer(newStubRouter(nil, ""), newStubRouter(nil, ""))
w := httptest.NewRecorder()
@@ -143,6 +314,44 @@ func TestServer_HandleMetrics_Unavailable(t *testing.T) {
}
}
// TestServer_HandleMetrics_SchedulerQueue verifies the scheduler queue is
// exported even when performance monitoring is off, and that every band gets a
// series so the label set stays stable while the queue drains.
func TestServer_HandleMetrics_SchedulerQueue(t *testing.T) {
local := newStubRouter(nil, "")
local.queueStats = &scheduler.QueueStats{
Bands: map[string]scheduler.BandStats{
shared.BandInteractive: {Depth: 2, OldestWait: 3 * time.Second, Dispatched: 41},
shared.BandBatch: {Depth: 1, Dispatched: 7},
},
AgingReorders: 5,
AffinityReorders: 9,
}
s := newTestServer(local, newStubRouter(nil, ""))
w := httptest.NewRecorder()
s.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/metrics", nil))
if w.Code != http.StatusOK {
t.Fatalf("status = %d, want 200", w.Code)
}
body := w.Body.String()
for _, want := range []string{
`llamaswap_scheduler_queue_depth{band="interactive"} 2`,
`llamaswap_scheduler_queue_depth{band="normal"} 0`,
`llamaswap_scheduler_queue_depth{band="batch"} 1`,
`llamaswap_scheduler_queue_oldest_wait_seconds{band="interactive"} 3`,
`llamaswap_scheduler_dispatched_total{band="interactive"} 41`,
`llamaswap_scheduler_dispatched_total{band="batch"} 7`,
`llamaswap_scheduler_reorders_total{cause="aging"} 5`,
`llamaswap_scheduler_reorders_total{cause="swap_affinity"} 9`,
} {
if !strings.Contains(body, want) {
t.Errorf("missing %q in:\n%s", want, body)
}
}
}
func TestServer_Redirects(t *testing.T) {
s := newTestServer(newStubRouter(nil, ""), newStubRouter(nil, ""))
+4 -1
View File
@@ -105,7 +105,9 @@ func (s *Server) handleAPIMetrics(w http.ResponseWriter, r *http.Request) {
// filtered to samples after the ?after=<RFC3339> timestamp.
func (s *Server) handleAPIPerformance(w http.ResponseWriter, r *http.Request) {
if s.perf == nil {
shared.SendResponse(w, r, http.StatusServiceUnavailable, "performance monitor not available")
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusServiceUnavailable)
json.NewEncoder(w).Encode(map[string]bool{"enabled": false})
return
}
@@ -136,6 +138,7 @@ func (s *Server) handleAPIPerformance(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(map[string]any{
"enabled": true,
"sys_stats": sysStats,
"gpu_stats": gpuStats,
})
+5 -1
View File
@@ -37,7 +37,8 @@ const (
)
// captureFieldsByPath overrides the default capture mask for routes carrying
// large binary payloads (audio/image) where storing the full body is wasteful.
// large binary payloads (audio/image/video) where storing the full body is
// wasteful.
var captureFieldsByPath = map[string]captureFields{
"/v1/audio/speech": captureReqAll | captureRespHeaders,
"/v1/audio/voices": captureReqHeaders | captureRespAll,
@@ -46,6 +47,9 @@ var captureFieldsByPath = map[string]captureFields{
"/v1/images/edits": captureReqHeaders | captureRespHeaders,
"/sdapi/v1/txt2img": captureReqAll | captureRespHeaders,
"/sdapi/v1/img2img": captureReqHeaders | captureRespHeaders,
// video: the request multipart can carry a conditioning frame and the
// response body IS the clip — headers only, both directions.
"/v1/videos/sync": captureReqHeaders | captureRespHeaders,
}
// captureFieldsFor returns the capture mask for a request path. Unlisted routes
+1 -1
View File
@@ -76,7 +76,7 @@ func (s *Server) getLogger(logMonitorID string) (*logmon.Monitor, error) {
case "upstream":
return s.upstreamlog, nil
default:
if _, modelID, _, found := findModelInPath(s.cfg, "/"+logMonitorID); found {
if _, modelID, _, found := shared.FindModelInPath(s.cfg, "/"+logMonitorID); found {
if log, ok := s.local.ProcessLogger(modelID); ok {
return log, nil
}
+115 -24
View File
@@ -25,6 +25,8 @@ import (
// TokenMetrics holds token usage and performance metrics.
type TokenMetrics struct {
CachedTokens int `json:"cache_tokens"`
DraftTokens int `json:"draft_tokens"`
DraftAccTokens int `json:"draft_acc_tokens"`
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
PromptPerSecond float64 `json:"prompt_per_second"`
@@ -42,6 +44,7 @@ type ActivityLogEntry struct {
Tokens TokenMetrics `json:"tokens"`
DurationMs int `json:"duration_ms"`
HasCapture bool `json:"has_capture"`
ErrorMsg string `json:"error_msg,omitempty"`
Metadata map[string]string `json:"metadata,omitempty"`
}
@@ -123,9 +126,11 @@ func (mp *metricsMonitor) getMetricsJSON() ([]byte, error) {
}
// record parses a completed response body and stores/emits an activity entry.
// When captures are enabled, a zstd+CBOR capture is stored for successful
// requests, with cf controlling which request/response parts are retained.
// reqBody and reqHeaders are the request data buffered before dispatch.
// Successful requests store a zstd+CBOR capture (when enabled) with cf
// controlling which parts are retained. Failed (non-200) requests capture the
// request only and set ErrorMsg to a description of the failure, so the error
// can be inspected without storing unreadable raw response bytes. reqBody and
// reqHeaders are the request data buffered before dispatch.
func (mp *metricsMonitor) record(modelID string, r *http.Request, recorder *responseBodyCopier, cf captureFields, reqBody []byte, reqHeaders map[string]string) {
tm := ActivityLogEntry{
Timestamp: time.Now(),
@@ -150,7 +155,13 @@ func (mp *metricsMonitor) record(modelID string, r *http.Request, recorder *resp
if recorder.Status() != http.StatusOK {
mp.logger.Warnf("non-200 response, recording partial metrics: status=%d, path=%s", recorder.Status(), r.URL.Path)
queueAndEmit()
decoded, decErr := mp.decodeResponseBody(recorder, r.URL.Path)
tm.ErrorMsg = failedErrorMessage(recorder.Status(), decoded, decErr)
tm.ID = mp.queueMetrics(tm)
// Capture the request only; the failure is surfaced via ErrorMsg
// rather than storing the (possibly undisplayable) response body.
tm.HasCapture = mp.storeCapture(tm.ID, r, recorder, cf&^captureRespBody, reqBody, reqHeaders, nil)
mp.emitMetric(tm)
return
}
@@ -165,6 +176,7 @@ func (mp *metricsMonitor) record(modelID string, r *http.Request, recorder *resp
decoded, err := decompressBody(body, encoding)
if err != nil {
mp.logger.Warnf("metrics: decompression failed: %v, path=%s, recording minimal metrics", err, r.URL.Path)
tm.ErrorMsg = fmt.Sprintf("response decompression failed: %v", err)
queueAndEmit()
return
}
@@ -203,28 +215,99 @@ func (mp *metricsMonitor) record(modelID string, r *http.Request, recorder *resp
}
tm.ID = mp.queueMetrics(tm)
if mp.enableCaptures {
capture := ReqRespCapture{
ID: tm.ID,
ReqPath: r.URL.Path,
ReqHeaders: reqHeaders,
}
if cf&captureReqBody != 0 {
capture.ReqBody = reqBody
}
if cf&captureRespHeaders != 0 {
capture.RespHeaders = headerMap(recorder.Header())
redactHeaders(capture.RespHeaders)
delete(capture.RespHeaders, "Content-Encoding")
}
if cf&captureRespBody != 0 {
capture.RespBody = body
}
if mp.addCapture(capture) {
tm.HasCapture = true
tm.HasCapture = mp.storeCapture(tm.ID, r, recorder, cf, reqBody, reqHeaders, body)
mp.emitMetric(tm)
}
// storeCapture assembles a ReqRespCapture for id, honoring the captureFields
// mask, and stores it when captures are enabled. body is the response body to
// capture (already decompressed by the caller); pass nil to omit it. Returns
// true if a capture was stored.
func (mp *metricsMonitor) storeCapture(id int, r *http.Request, recorder *responseBodyCopier, cf captureFields, reqBody []byte, reqHeaders map[string]string, body []byte) bool {
if !mp.enableCaptures {
return false
}
capture := ReqRespCapture{
ID: id,
ReqPath: r.URL.Path,
ReqHeaders: reqHeaders,
}
if cf&captureReqBody != 0 {
capture.ReqBody = reqBody
}
if cf&captureRespHeaders != 0 {
capture.RespHeaders = headerMap(recorder.Header())
redactHeaders(capture.RespHeaders)
delete(capture.RespHeaders, "Content-Encoding")
}
if cf&captureRespBody != 0 {
capture.RespBody = body
}
return mp.addCapture(capture)
}
// decodeResponseBody returns the buffered response body, decompressing it when
// the upstream set a Content-Encoding we recognize. On decompression failure it
// logs a warning and returns an error so the caller can record a description
// (via ErrorMsg) instead of storing unreadable raw bytes.
func (mp *metricsMonitor) decodeResponseBody(recorder *responseBodyCopier, path string) ([]byte, error) {
body := recorder.body.Bytes()
if len(body) == 0 {
return nil, nil
}
encoding := recorder.Header().Get("Content-Encoding")
if encoding == "" {
return body, nil
}
decoded, err := decompressBody(body, encoding)
if err != nil {
mp.logger.Warnf("metrics: response decompression failed: %v, path=%s", err, path)
return nil, err
}
return decoded, nil
}
// errorMessagePaths lists JSON paths where a human-readable error message can
// live across OpenAI- and llama.cpp-style error responses.
var errorMessagePaths = []string{"error.message", "error", "message", "detail"}
// extractErrorMessage pulls a human-readable error string from a JSON error
// response. Returns "" if no message is found or the body is not valid JSON.
func extractErrorMessage(body []byte) string {
if !gjson.ValidBytes(body) {
return ""
}
parsed := gjson.ParseBytes(body)
for _, path := range errorMessagePaths {
v := parsed.Get(path)
if v.Exists() && v.Type == gjson.String {
if s := strings.TrimSpace(v.String()); s != "" {
return s
}
}
}
mp.emitMetric(tm)
return ""
}
// failedErrorMessage builds a human-readable description for a non-200 response.
// It prefers an error message parsed from the (decompressed) body and falls back
// to the HTTP status text. A non-nil decErr indicates the body could not be
// decoded, in which case the decode error is described instead.
func failedErrorMessage(status int, body []byte, decErr error) string {
const maxLen = 500
if decErr != nil {
return fmt.Sprintf("response decode failed: %v", decErr)
}
if msg := extractErrorMessage(body); msg != "" {
if len(msg) > maxLen {
msg = msg[:maxLen] + "..."
}
return msg
}
if text := http.StatusText(status); text != "" {
return fmt.Sprintf("%d %s", status, text)
}
return fmt.Sprintf("HTTP %d", status)
}
// usagePaths lists the JSON paths where a per-event usage object can live.
@@ -345,6 +428,8 @@ func buildMetrics(modelID string, start time.Time, inputTokens, outputTokens, ca
durationMs := wallDurationMs
tokensPerSecond := -1.0
promptPerSecond := -1.0
draftTokens := -1
draftAccTokens := -1
if timings.Exists() {
inputTokens = timings.Get("prompt_n").Int()
@@ -358,6 +443,10 @@ func buildMetrics(modelID string, start time.Time, inputTokens, outputTokens, ca
if cachedValue := timings.Get("cache_n"); cachedValue.Exists() {
cachedTokens = cachedValue.Int()
}
if timings.Get("draft_n").Exists() && timings.Get("draft_n_accepted").Exists() {
draftTokens = int(timings.Get("draft_n").Int())
draftAccTokens = int(timings.Get("draft_n_accepted").Int())
}
}
return ActivityLogEntry{
@@ -365,6 +454,8 @@ func buildMetrics(modelID string, start time.Time, inputTokens, outputTokens, ca
Model: modelID,
Tokens: TokenMetrics{
CachedTokens: int(cachedTokens),
DraftTokens: draftTokens,
DraftAccTokens: draftAccTokens,
InputTokens: int(inputTokens),
OutputTokens: int(outputTokens),
PromptPerSecond: promptPerSecond,
+22 -2
View File
@@ -4,6 +4,7 @@ import (
"bytes"
"io"
"net/http"
"strings"
"github.com/mostlygeek/llama-swap/internal/chain"
"github.com/mostlygeek/llama-swap/internal/config"
@@ -21,8 +22,27 @@ func CreateMetricsMiddleware(mm *metricsMonitor, cfg config.Config) chain.Middle
return
}
// Determine the model-routed endpoint path. Regular routes are
// already meterable; /upstream/<model>/<path> is metered only when
// the remaining path matches a model-dispatched endpoint.
checkPath := r.URL.Path
if strings.HasPrefix(r.URL.Path, "/upstream/") {
var found bool
_, _, checkPath, found = shared.FindModelInPath(cfg, strings.TrimPrefix(r.URL.Path, "/upstream"))
if !found {
next.ServeHTTP(w, r)
return
}
}
if !isMetricsRecordPath(checkPath) {
next.ServeHTTP(w, r)
return
}
// Resolve the model now so downstream dispatch hits the context
// fast path; FetchContext restores the request body.
// fast path; FetchContext restores the request body for regular
// routes and extracts the model from the URL for /upstream routes.
data, err := shared.FetchContext(r, cfg)
if err != nil {
shared.SendError(w, r, shared.ErrNoModelInContext)
@@ -31,7 +51,7 @@ func CreateMetricsMiddleware(mm *metricsMonitor, cfg config.Config) chain.Middle
// Buffer the request body/headers for capture before dispatch
// consumes them.
cf := captureFieldsFor(r.URL.Path)
cf := captureFieldsFor(checkPath)
var reqBody []byte
var reqHeaders map[string]string
if mm.enableCaptures {
+39
View File
@@ -0,0 +1,39 @@
package server
import (
"fmt"
"io"
"github.com/mostlygeek/llama-swap/internal/router/scheduler"
"github.com/mostlygeek/llama-swap/internal/shared"
)
// writeSchedulerMetrics emits the request queue in Prometheus text format.
//
// The reorder counters are the ones that matter for tuning: they say how often
// aging or swap affinity actually changed which request went next. Without
// them, choosing agingDivisor and swapAffinityBonus is guesswork.
func writeSchedulerMetrics(w io.Writer, stats scheduler.QueueStats) {
fmt.Fprintf(w, "# HELP llamaswap_scheduler_queue_depth Requests waiting in the scheduler queue, by priority band\n")
fmt.Fprintf(w, "# TYPE llamaswap_scheduler_queue_depth gauge\n")
for _, band := range shared.PriorityBands {
fmt.Fprintf(w, "llamaswap_scheduler_queue_depth{band=%q} %d\n", band, stats.Bands[band].Depth)
}
fmt.Fprintf(w, "# HELP llamaswap_scheduler_queue_oldest_wait_seconds How long the longest-waiting queued request has waited, by priority band\n")
fmt.Fprintf(w, "# TYPE llamaswap_scheduler_queue_oldest_wait_seconds gauge\n")
for _, band := range shared.PriorityBands {
fmt.Fprintf(w, "llamaswap_scheduler_queue_oldest_wait_seconds{band=%q} %g\n", band, stats.Bands[band].OldestWait.Seconds())
}
fmt.Fprintf(w, "# HELP llamaswap_scheduler_dispatched_total Requests dispatched by the scheduler, by priority band\n")
fmt.Fprintf(w, "# TYPE llamaswap_scheduler_dispatched_total counter\n")
for _, band := range shared.PriorityBands {
fmt.Fprintf(w, "llamaswap_scheduler_dispatched_total{band=%q} %d\n", band, stats.Bands[band].Dispatched)
}
fmt.Fprintf(w, "# HELP llamaswap_scheduler_reorders_total Dispatches where a scoring term changed which request was picked\n")
fmt.Fprintf(w, "# TYPE llamaswap_scheduler_reorders_total counter\n")
fmt.Fprintf(w, "llamaswap_scheduler_reorders_total{cause=\"aging\"} %d\n", stats.AgingReorders)
fmt.Fprintf(w, "llamaswap_scheduler_reorders_total{cause=\"swap_affinity\"} %d\n", stats.AffinityReorders)
}
+206
View File
@@ -1,12 +1,15 @@
package server
import (
"io"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/mostlygeek/llama-swap/internal/config"
"github.com/mostlygeek/llama-swap/internal/logmon"
"github.com/mostlygeek/llama-swap/internal/shared"
"github.com/tidwall/gjson"
)
@@ -87,6 +90,172 @@ func TestMetricsMonitor_RecordMetadata(t *testing.T) {
}
}
func TestMetricsMonitor_RecordFailedRequestCapture(t *testing.T) {
mm := newMetricsMonitor(logmon.NewWriter(io.Discard), 10, 5)
r := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
reqHeaders := map[string]string{"content-type": "application/json"}
w := httptest.NewRecorder()
copier := newBodyCopier(w)
copier.Header().Set("Content-Type", "application/json")
copier.WriteHeader(http.StatusBadGateway)
copier.Write([]byte(`{"error":{"message":"model unavailable"}}`))
reqBody := []byte(`{"model":"m","messages":[]}`)
mm.record("m", r, copier, captureAll, reqBody, reqHeaders)
entries := mm.getMetrics()
if len(entries) != 1 {
t.Fatalf("want 1 entry, got %d", len(entries))
}
entry := entries[0]
if entry.RespStatusCode != http.StatusBadGateway {
t.Errorf("status = %d, want %d", entry.RespStatusCode, http.StatusBadGateway)
}
if entry.ErrorMsg != "model unavailable" {
t.Errorf("error_msg = %q, want extracted message", entry.ErrorMsg)
}
if !entry.HasCapture {
t.Fatal("failed request should capture the request so it can be inspected")
}
got := mm.getCaptureByID(entry.ID)
if got == nil {
t.Fatal("capture not found")
}
if string(got.ReqBody) != `{"model":"m","messages":[]}` {
t.Errorf("req body = %q", got.ReqBody)
}
if len(got.RespBody) != 0 {
t.Errorf("resp body stored for failed request (len=%d); want none", len(got.RespBody))
}
if got.RespHeaders["Content-Type"] != "application/json" {
t.Errorf("resp Content-Type = %q", got.RespHeaders["Content-Type"])
}
}
func TestMetricsMonitor_RecordFailedRequestStatusFallback(t *testing.T) {
// Non-JSON error body: ErrorMsg falls back to the HTTP status text.
mm := newMetricsMonitor(logmon.NewWriter(io.Discard), 10, 5)
r := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
w := httptest.NewRecorder()
copier := newBodyCopier(w)
copier.WriteHeader(http.StatusBadGateway)
copier.Write([]byte("<html>upstream down</html>"))
mm.record("m", r, copier, captureAll, nil, nil)
entries := mm.getMetrics()
if len(entries) != 1 {
t.Fatalf("want 1 entry, got %d", len(entries))
}
if entries[0].ErrorMsg != "502 Bad Gateway" {
t.Errorf("error_msg = %q, want status text", entries[0].ErrorMsg)
}
}
func TestMetricsMonitor_RecordFailedRequestCaptureDisabled(t *testing.T) {
mm := newMetricsMonitor(logmon.NewWriter(io.Discard), 10, 0) // captures disabled
r := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
w := httptest.NewRecorder()
copier := newBodyCopier(w)
copier.WriteHeader(http.StatusInternalServerError)
copier.Write([]byte(`{"error":"boom"}`))
mm.record("m", r, copier, captureAll, []byte("req"), nil)
entries := mm.getMetrics()
if len(entries) != 1 {
t.Fatalf("want 1 entry, got %d", len(entries))
}
if entries[0].HasCapture {
t.Fatal("captures disabled, HasCapture should be false")
}
// ErrorMsg is independent of whether captures are enabled.
if entries[0].ErrorMsg != "boom" {
t.Errorf("error_msg = %q, want boom", entries[0].ErrorMsg)
}
if mm.getCaptureByID(entries[0].ID) != nil {
t.Fatal("no capture should be stored when disabled")
}
}
func TestMetricsMonitor_RecordDecompressionFailureSetsErrorMsg(t *testing.T) {
mm := newMetricsMonitor(logmon.NewWriter(io.Discard), 10, 5)
r := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
w := httptest.NewRecorder()
copier := newBodyCopier(w)
copier.Header().Set("Content-Encoding", "gzip")
copier.WriteHeader(http.StatusOK)
copier.Write([]byte("not-really-gzip"))
mm.record("m", r, copier, captureAll, []byte("req"), nil)
entries := mm.getMetrics()
if len(entries) != 1 {
t.Fatalf("want 1 entry, got %d", len(entries))
}
if entries[0].ErrorMsg == "" {
t.Fatal("expected ErrorMsg for decompression failure")
}
// Raw bytes must not be stored when the body could not be decoded.
if entries[0].HasCapture {
t.Fatal("decompression failure should not store a capture")
}
}
func TestMetricsMonitor_DecodeResponseBody(t *testing.T) {
mm := newMetricsMonitor(logmon.NewWriter(io.Discard), 10, 5)
// No Content-Encoding: body returned unchanged.
w := httptest.NewRecorder()
copier := newBodyCopier(w)
copier.Write([]byte("plain"))
got, err := mm.decodeResponseBody(copier, "/p")
if err != nil || string(got) != "plain" {
t.Fatalf("plain body = %q, err = %v", got, err)
}
// Bogus gzip payload: returns an error and no body (no raw bytes kept).
w2 := httptest.NewRecorder()
copier2 := newBodyCopier(w2)
copier2.Header().Set("Content-Encoding", "gzip")
copier2.Write([]byte("not-really-gzip"))
got, err = mm.decodeResponseBody(copier2, "/p")
if err == nil {
t.Fatal("expected decompression error")
}
if got != nil {
t.Errorf("expected nil body on failure, got %q", got)
}
}
func TestServer_ExtractErrorMessage(t *testing.T) {
cases := []struct {
name string
body string
want string
}{
{"openai object", `{"error":{"message":"rate limited"}}`, "rate limited"},
{"string error", `{"error":"bad request"}`, "bad request"},
{"message field", `{"message":"nope"}`, "nope"},
{"detail field", `{"detail":"oops"}`, "oops"},
{"object error ignored", `{"error":{"code":42}}`, ""},
{"no error", `{"usage":{}}`, ""},
{"invalid json", `not-json`, ""},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
if got := extractErrorMessage([]byte(tc.body)); got != tc.want {
t.Errorf("extractErrorMessage = %q, want %q", got, tc.want)
}
})
}
}
func TestServer_ParseMetrics_Infill(t *testing.T) {
// /infill responses are arrays; timings live in the last element.
body := `[{"content":"a"},{"content":"b","timings":{"prompt_n":5,"predicted_n":9,"prompt_ms":10,"predicted_ms":20}}]`
@@ -103,3 +272,40 @@ func TestServer_ParseMetrics_Infill(t *testing.T) {
t.Fatalf("tokens = %+v", entry.Tokens)
}
}
// TestServer_MetricsMiddleware_UpstreamAudioCaptureSkipsRespBody verifies that
// an /upstream/<model>/v1/audio/speech request uses the path-specific capture
// mask (headers only) rather than falling back to captureAll.
func TestServer_MetricsMiddleware_UpstreamAudioCaptureSkipsRespBody(t *testing.T) {
mm := newMetricsMonitor(logmon.NewWriter(io.Discard), 100, 5)
cfg := config.Config{Models: map[string]config.ModelConfig{"m1": {}}}
inner := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "audio/mpeg")
w.WriteHeader(http.StatusOK)
w.Write([]byte("BINARY-AUDIO-DATA"))
})
handler := CreateMetricsMiddleware(mm, cfg)(inner)
req := httptest.NewRequest(http.MethodPost, "/upstream/m1/v1/audio/speech", strings.NewReader(`{"model":"m1"}`))
handler.ServeHTTP(httptest.NewRecorder(), req)
entries := mm.getMetrics()
if len(entries) == 0 {
t.Fatal("no metrics recorded")
}
last := entries[len(entries)-1]
if !last.HasCapture {
t.Fatal("expected capture to be stored")
}
cap := mm.getCaptureByID(last.ID)
if cap == nil {
t.Fatal("capture not found")
}
if len(cap.RespBody) != 0 {
t.Errorf("RespBody stored for /upstream audio route (len=%d); want path-specific mask to skip body", len(cap.RespBody))
}
if len(cap.RespHeaders) == 0 {
t.Error("RespHeaders not stored; want captureRespHeaders mask")
}
}
+34 -2
View File
@@ -80,6 +80,15 @@ var modelPostJSONRoutes = []string{
var modelPostFormRoutes = []string{
"/v1/audio/transcriptions",
"/v1/images/edits",
// video generation: ONLY the blocking sync shape (vLLM-Omni
// /v1/videos/sync — the response body is the clip). The async
// OpenAI-style POST /v1/videos is deliberately NOT dispatched: its
// poll/download companions (GET /v1/videos/{id}[/content]) carry no
// model field to route by, so registering creation alone would start
// unretrievable upstream jobs. Extraction is content-type driven, so
// JSON bodies dispatch here too.
"/v1/videos/sync",
}
// modelGetRoutes are model-dispatched GET endpoints (the model arrives as a
@@ -89,6 +98,27 @@ var modelGetRoutes = []string{
"/sdapi/v1/loras",
}
// isMetricsRecordPath reports whether path is one of the model-dispatched
// endpoints that the metrics middleware records in the activity log.
func isMetricsRecordPath(path string) bool {
for _, p := range modelPostJSONRoutes {
if p == path {
return true
}
}
for _, p := range modelPostFormRoutes {
if p == path {
return true
}
}
for _, p := range modelGetRoutes {
if p == path {
return true
}
}
return false
}
// BuildInfo carries version metadata surfaced by GET /api/version.
type BuildInfo struct {
Version string
@@ -219,9 +249,11 @@ func (s *Server) routes() {
mux.Handle("GET /unload", apiChain.ThenFunc(s.handleUnload))
mux.Handle("GET /running", apiChain.ThenFunc(s.handleRunning))
// Upstream passthrough.
// Upstream passthrough. Meter only the model-dispatched endpoints that can
// produce token usage/timings.
upstreamChain := apiChain.Append(CreateMetricsMiddleware(s.metrics, s.cfg))
mux.HandleFunc("GET /upstream", handleUpstreamRedirect)
mux.Handle("/upstream/{upstreamPath...}", apiChain.ThenFunc(s.handleUpstream))
mux.Handle("/upstream/{upstreamPath...}", upstreamChain.ThenFunc(s.handleUpstream))
// API group (API-key protected) consumed by the UI.
mux.Handle("POST /api/models/unload", apiChain.ThenFunc(s.handleAPIUnloadAll))
+64
View File
@@ -1,9 +1,11 @@
package server
import (
"bytes"
"context"
"encoding/json"
"io"
"mime/multipart"
"net/http"
"net/http/httptest"
"strings"
@@ -16,6 +18,7 @@ import (
"github.com/mostlygeek/llama-swap/internal/logmon"
"github.com/mostlygeek/llama-swap/internal/process"
"github.com/mostlygeek/llama-swap/internal/router"
"github.com/mostlygeek/llama-swap/internal/router/scheduler"
"github.com/mostlygeek/llama-swap/internal/shared"
)
@@ -27,6 +30,7 @@ type stubRouter struct {
running map[string]process.ProcessState
unloadCalls atomic.Int32
loggers map[string]*logmon.Monitor
queueStats *scheduler.QueueStats
}
func newStubRouter(models []string, response string) *stubRouter {
@@ -55,6 +59,13 @@ func (s *stubRouter) ProcessLogger(modelID string) (*logmon.Monitor, bool) {
return nil, false
}
func (s *stubRouter) SchedulerStats() (scheduler.QueueStats, bool) {
if s.queueStats == nil {
return scheduler.QueueStats{}, false
}
return *s.queueStats, true
}
// newTestServer wires a Server with stub routers and a built mux.
func newTestServer(local router.LocalRouter, peer router.Router) *Server {
ctx, cancel := context.WithCancel(context.Background())
@@ -339,3 +350,56 @@ func TestServer_LogStream_UnknownID_Returns400(t *testing.T) {
t.Errorf("status=%d want 400", w.Code)
}
}
func TestServer_VideoRoutesDispatch(t *testing.T) {
s := newTestServer(
newStubRouter([]string{"videogen-model"}, "video response"),
newStubRouter(nil, ""),
)
// Multipart form on /v1/videos/sync.
var buf bytes.Buffer
mw := multipart.NewWriter(&buf)
mw.WriteField("model", "videogen-model")
mw.WriteField("prompt", "a cat")
mw.Close()
req := httptest.NewRequest(http.MethodPost, "/v1/videos/sync", &buf)
req.Header.Set("Content-Type", mw.FormDataContentType())
w := httptest.NewRecorder()
s.ServeHTTP(w, req)
if w.Code != http.StatusOK || w.Body.String() != "video response" {
t.Fatalf("/v1/videos/sync: status=%d body=%q", w.Code, w.Body.String())
}
// JSON bodies dispatch too (extraction is content-type driven).
jsonReq := httptest.NewRequest(http.MethodPost, "/v1/videos/sync",
strings.NewReader(`{"model":"videogen-model","prompt":"a cat"}`))
jsonReq.Header.Set("Content-Type", "application/json")
w = httptest.NewRecorder()
s.ServeHTTP(w, jsonReq)
if w.Code != http.StatusOK || w.Body.String() != "video response" {
t.Fatalf("/v1/videos/sync json: status=%d body=%q", w.Code, w.Body.String())
}
// Unknown model on the video route still 404s.
nopeReq := httptest.NewRequest(http.MethodPost, "/v1/videos/sync",
strings.NewReader(`{"model":"nope","prompt":"a cat"}`))
nopeReq.Header.Set("Content-Type", "application/json")
w = httptest.NewRecorder()
s.ServeHTTP(w, nopeReq)
if w.Code != http.StatusNotFound {
t.Fatalf("/v1/videos/sync unknown model: status=%d want 404", w.Code)
}
// The async job-creation route is NOT dispatched (no poll routes to
// complete the flow) — a stray OpenAI-videos client gets a clean 404
// instead of an unretrievable upstream job.
asyncReq := httptest.NewRequest(http.MethodPost, "/v1/videos",
strings.NewReader(`{"model":"videogen-model","prompt":"a cat"}`))
asyncReq.Header.Set("Content-Type", "application/json")
w = httptest.NewRecorder()
s.ServeHTTP(w, asyncReq)
if w.Code != http.StatusNotFound {
t.Fatalf("POST /v1/videos: status=%d want 404 (async family not routed)", w.Code)
}
}
+70 -4
View File
@@ -26,6 +26,9 @@ type ReqContextData struct {
ModelID string
Streaming bool
SendLoadingState bool
// Priority is the caller's scheduling priority, read from PriorityHeader.
// Higher is more urgent; 0 (PriorityNormal) is the default. See priority.go.
Priority int
// Metadata is a request-scoped key/value bag that handlers may mutate
// while processing. The metrics middleware copies it into ActivityLogEntry.
Metadata map[string]string
@@ -91,16 +94,24 @@ func SendResponse(w http.ResponseWriter, r *http.Request, status int, message st
w.Write(resp)
}
// FetchContext will attempt to get the model id from the context then
// from the model body. If it extracts the model from the body it will
// store the model in the context for downstream handlers. An error
// will be returned when model can not be fetch from either location.
// FetchContext will attempt to get the model id from the context, then
// from an /upstream/<model> path prefix, then from the request body/query.
// If it extracts the model it will store it in the context for downstream
// handlers. An error will be returned when a model cannot be identified.
func FetchContext(r *http.Request, cfg config.Config) (ReqContextData, error) {
data, ok := ReadContext(r.Context())
if ok {
return data, nil
}
if strings.HasPrefix(r.URL.Path, "/upstream/") {
if data, ok := extractUpstreamContext(r, cfg); ok {
*r = *r.WithContext(SetContext(r.Context(), data))
return data, nil
}
return ReqContextData{}, ErrNoModelInContext
}
if data, err := extractContext(r); err == nil && data.Model != "" {
realName, _ := cfg.RealModelName(data.Model)
if realName == "" {
@@ -110,6 +121,7 @@ func FetchContext(r *http.Request, cfg config.Config) (ReqContextData, error) {
if mc, ok := cfg.Models[realName]; ok {
data.SendLoadingState = mc.SendLoadingState != nil && *mc.SendLoadingState
}
data.Priority = RequestPriority(r)
*r = *r.WithContext(SetContext(r.Context(), data))
return data, nil
}
@@ -117,6 +129,60 @@ func FetchContext(r *http.Request, cfg config.Config) (ReqContextData, error) {
return ReqContextData{}, ErrNoModelInContext
}
// extractUpstreamContext resolves the model from an /upstream/<model>/... path.
func extractUpstreamContext(r *http.Request, cfg config.Config) (ReqContextData, bool) {
searchName, realName, _, found := FindModelInPath(cfg, strings.TrimPrefix(r.URL.Path, "/upstream"))
if !found {
return ReqContextData{}, false
}
return ReqContextData{
Model: searchName,
ModelID: realName,
ApiKey: ExtractAPIKey(r),
Streaming: r.URL.Query().Get("stream") == "true",
SendLoadingState: sendLoadingState(cfg, realName),
Priority: RequestPriority(r),
Metadata: make(map[string]string),
}, true
}
// sendLoadingState reports whether the configured model wants loading-state SSEs.
func sendLoadingState(cfg config.Config, modelID string) bool {
if mc, ok := cfg.Models[modelID]; ok {
return mc.SendLoadingState != nil && *mc.SendLoadingState
}
return false
}
// FindModelInPath walks a slash-separated path, building up segments until one
// matches a configured model. This resolves model names that contain slashes
// (e.g. "author/model"). Returns the matched name, its real model ID, the
// remaining path, and whether a match was found.
func FindModelInPath(cfg config.Config, path string) (searchName, realName, remainingPath string, found bool) {
parts := strings.Split(strings.TrimSpace(path), "/")
name := ""
for i, part := range parts {
if part == "" {
continue
}
if name == "" {
name = part
} else {
name = name + "/" + part
}
if modelID, ok := cfg.RealModelName(name); ok {
searchName = name
realName = modelID
remainingPath = "/" + strings.Join(parts[i+1:], "/")
found = true
}
}
return
}
func SetContext(ctx context.Context, data ReqContextData) context.Context {
return context.WithValue(ctx, ReqContextKey, data)
}
+67
View File
@@ -11,6 +11,8 @@ import (
"net/url"
"strings"
"testing"
"github.com/mostlygeek/llama-swap/internal/config"
)
func TestExtractContext_GET(t *testing.T) {
@@ -456,3 +458,68 @@ func TestServer_ExtractAPIKey(t *testing.T) {
})
}
}
func TestFetchContext_UpstreamPath(t *testing.T) {
cfg := config.Config{
Models: map[string]config.ModelConfig{
"m1": {},
"author/model": {},
"real": {Aliases: []string{"nick"}},
},
}
cases := []struct {
name string
path string
wantModel string
wantModelID string
wantErr bool
}{
{"known model", "/upstream/m1/v1/chat/completions", "m1", "m1", false},
{"model with slash", "/upstream/author/model/v1/chat", "author/model", "author/model", false},
{"unknown model", "/upstream/nope/v1/chat/completions", "", "", true},
{"bare model path", "/upstream/m1/", "m1", "m1", false},
}
for _, c := range cases {
t.Run(c.name, func(t *testing.T) {
r := httptest.NewRequest(http.MethodPost, c.path, strings.NewReader(`{}`))
data, err := FetchContext(r, cfg)
if (err != nil) != c.wantErr {
t.Fatalf("wantErr=%v got err=%v", c.wantErr, err)
}
if c.wantErr {
return
}
if data.Model != c.wantModel {
t.Errorf("model = %q, want %q", data.Model, c.wantModel)
}
if data.ModelID != c.wantModelID {
t.Errorf("modelID = %q, want %q", data.ModelID, c.wantModelID)
}
if data.Metadata == nil {
t.Error("metadata map not initialized")
}
})
}
}
func TestFetchContext_UpstreamPath_DoesNotReadBody(t *testing.T) {
cfg := config.Config{Models: map[string]config.ModelConfig{"m1": {}}}
body := `{"model":"should-not-matter"}`
r := httptest.NewRequest(http.MethodPost, "/upstream/m1/v1/chat/completions", strings.NewReader(body))
_, err := FetchContext(r, cfg)
if err != nil {
t.Fatalf("FetchContext: %v", err)
}
// The body should be untouched so the upstream handler can still read it.
got, err := io.ReadAll(r.Body)
if err != nil {
t.Fatalf("read body: %v", err)
}
if string(got) != body {
t.Errorf("body was consumed: %q", string(got))
}
}
+89
View File
@@ -0,0 +1,89 @@
package shared
import (
"net/http"
"strconv"
"strings"
)
// PriorityHeader carries the caller's scheduling priority. Priority is a
// property of the caller's intent, not of the model — the same model serves
// both an interactive request and a batch job — so it travels as a header
// rather than as per-model configuration, and leaves the OpenAI-compatible
// request body untouched.
const PriorityHeader = "X-LlamaSwap-Priority"
// Priority band anchors. The number is the interface: callers may send any
// signed integer, and these names are conveniences that resolve to a value.
//
// Bands are spaced 100 apart so a consumer can add a small per-caller offset
// (a subscription tier, say) on top of a band without ever promoting a request
// across one: a "max member" batch job at -98 still loses to every normal
// request at 0.
const (
PriorityInteractive = 100 // a human is waiting
PriorityNormal = 0 // the default
PriorityBatch = -100
)
// bandHalfWidth is how far a priority may sit from a band anchor and still be
// reported as that band. It is half the 100-point band spacing, so every
// integer belongs to exactly one band.
const bandHalfWidth = 50
// Band names used for metrics labels, ordered most to least urgent.
const (
BandInteractive = "interactive"
BandNormal = "normal"
BandBatch = "batch"
)
// PriorityBands is every band name, in descending order of urgency. Metrics
// emit a series per band regardless of whether any request currently occupies
// it, so the label set stays stable.
var PriorityBands = []string{BandInteractive, BandNormal, BandBatch}
var priorityAliases = map[string]int{
"interactive": PriorityInteractive,
"normal": PriorityNormal,
"batch": PriorityBatch,
}
// ParsePriority resolves a PriorityHeader value to a signed integer. A numeric
// value is used as-is; a recognised alias resolves to its anchor. Anything
// absent or unparseable is PriorityNormal, so a malformed header degrades to
// the default rather than failing the request.
//
// Values are deliberately not clamped: the caller composes band, tier offset
// and anything else it wants into one number, and llama-swap honours it.
func ParsePriority(value string) int {
value = strings.TrimSpace(value)
if value == "" {
return PriorityNormal
}
if n, err := strconv.Atoi(value); err == nil {
return n
}
if n, ok := priorityAliases[strings.ToLower(value)]; ok {
return n
}
return PriorityNormal
}
// RequestPriority reads PriorityHeader off r and resolves it. See ParsePriority.
func RequestPriority(r *http.Request) int {
return ParsePriority(r.Header.Get(PriorityHeader))
}
// PriorityBand buckets a priority into the band it belongs to, for metrics
// labels. Values beyond the anchors saturate: +1000 is still "interactive".
func PriorityBand(priority int) string {
switch {
case priority >= PriorityInteractive-bandHalfWidth:
return BandInteractive
case priority <= PriorityBatch+bandHalfWidth:
return BandBatch
default:
return BandNormal
}
}
+118
View File
@@ -0,0 +1,118 @@
package shared
import (
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/mostlygeek/llama-swap/internal/config"
)
func TestShared_ParsePriority(t *testing.T) {
tests := []struct {
name string
value string
want int
}{
{"absent", "", PriorityNormal},
{"zero", "0", 0},
{"positive", "100", 100},
{"negative", "-100", -100},
{"explicit plus", "+42", 42},
{"surrounding space", " 100 ", 100},
{"large value is not clamped", "100000", 100000},
{"alias interactive", "interactive", PriorityInteractive},
{"alias normal", "normal", PriorityNormal},
{"alias batch", "batch", PriorityBatch},
{"alias is case insensitive", "Batch", PriorityBatch},
{"tier offset on a band", "-98", -98},
{"unknown alias falls back to normal", "urgent", PriorityNormal},
{"garbage falls back to normal", "!!", PriorityNormal},
{"float falls back to normal", "1.5", PriorityNormal},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
if got := ParsePriority(tc.value); got != tc.want {
t.Errorf("ParsePriority(%q)=%d want %d", tc.value, got, tc.want)
}
})
}
}
func TestShared_RequestPriority(t *testing.T) {
r := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
if got := RequestPriority(r); got != PriorityNormal {
t.Errorf("no header: got %d want %d", got, PriorityNormal)
}
r.Header.Set(PriorityHeader, "batch")
if got := RequestPriority(r); got != PriorityBatch {
t.Errorf("batch header: got %d want %d", got, PriorityBatch)
}
}
// TestShared_PriorityBand pins the bucketing the metrics labels rely on. The
// important cases are the tier offsets: a max member's batch job at -98 must
// still report as batch, not normal.
func TestShared_PriorityBand(t *testing.T) {
tests := []struct {
priority int
want string
}{
{PriorityInteractive, BandInteractive},
{PriorityInteractive + 2, BandInteractive}, // max tier, interactive
{1000, BandInteractive}, // saturates
{50, BandInteractive}, // lower edge
{49, BandNormal},
{PriorityNormal, BandNormal},
{2, BandNormal}, // max tier, normal
{-49, BandNormal},
{-50, BandBatch}, // upper edge
{PriorityBatch + 2, BandBatch},
{PriorityBatch, BandBatch},
{-1000, BandBatch}, // saturates
}
for _, tc := range tests {
if got := PriorityBand(tc.priority); got != tc.want {
t.Errorf("PriorityBand(%d)=%q want %q", tc.priority, got, tc.want)
}
}
}
// TestShared_FetchContext_Priority verifies the header reaches the request
// context, which is what carries it to the scheduler. Both entry points into
// FetchContext are covered: the normal body-parsed path and /upstream/<model>.
func TestShared_FetchContext_Priority(t *testing.T) {
cfg := config.Config{Models: map[string]config.ModelConfig{"m1": {}}}
cases := []struct {
name string
path string
header string
want int
}{
{"body path, no header", "/v1/chat/completions", "", PriorityNormal},
{"body path, alias", "/v1/chat/completions", "interactive", PriorityInteractive},
{"body path, number", "/v1/chat/completions", "-98", -98},
{"upstream path", "/upstream/m1/v1/chat/completions", "batch", PriorityBatch},
}
for _, c := range cases {
t.Run(c.name, func(t *testing.T) {
r := httptest.NewRequest(http.MethodPost, c.path, strings.NewReader(`{"model":"m1"}`))
r.Header.Set("Content-Type", "application/json")
if c.header != "" {
r.Header.Set(PriorityHeader, c.header)
}
data, err := FetchContext(r, cfg)
if err != nil {
t.Fatalf("FetchContext: %v", err)
}
if data.Priority != c.want {
t.Errorf("Priority=%d want %d", data.Priority, c.want)
}
})
}
}
+137
View File
@@ -0,0 +1,137 @@
package configwatcher
import (
"context"
"os"
"path/filepath"
"sort"
"strings"
"time"
)
// DirWatcher polls a directory for changes to its set of *.yml / *.yaml files.
// It fires OnChange when a file is added, removed, or has its mod time/size
// change. Like Watcher it is poll-based so it works in Docker bind-mounts and
// k8s ConfigMap projections where inotify is unreliable.
//
// The baseline poll establishes initial state and does not fire OnChange.
type DirWatcher struct {
Path string
Interval time.Duration
OnChange func()
}
// dirSnapshot is an ordered map of file name -> file state. The ordering is
// derived from sorted filenames so two snapshots compare deterministically
// regardless of readdir order. exists reflects whether the directory was
// readable at scan time; a missing directory yields exists=false.
type dirSnapshot struct {
exists bool
names []string
states map[string]snapshot
}
func newDirSnapshot() dirSnapshot {
return dirSnapshot{states: make(map[string]snapshot)}
}
// equal reports whether two snapshots describe the same file set and per-file
// state. A missing directory (exists=false) is treated as equal to any other
// missing directory regardless of cached names.
func (s dirSnapshot) equal(other dirSnapshot) bool {
if !s.exists && !other.exists {
return true
}
if s.exists != other.exists {
return false
}
if len(s.names) != len(other.names) {
return false
}
for i, n := range s.names {
if other.names[i] != n {
return false
}
}
for _, n := range s.names {
a, b := s.states[n], other.states[n]
if a.exists != b.exists || a.size != b.size || !a.modTime.Equal(b.modTime) {
return false
}
}
return true
}
// Run blocks until ctx is canceled. It polls Path on Interval and invokes
// OnChange whenever the directory's YAML file set changes.
//
// Policy mirrors the single-file Watcher: disappearance (directory missing or
// empty) is treated as a transient rename-style write and stays quiet; the
// transition back to present-with-content fires OnChange.
func (w *DirWatcher) Run(ctx context.Context) {
interval := w.Interval
if interval <= 0 {
interval = DefaultInterval
}
prev := scanDir(w.Path)
ticker := time.NewTicker(interval)
defer ticker.Stop()
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
cur := scanDir(w.Path)
// Suppress transitions involving an empty or missing directory —
// these are treated as transient rename-style writes, mirroring
// the single-file Watcher. Only present-with-content →
// present-with-content (changed) or no-content →
// present-with-content fires OnChange.
prevHasContent := prev.exists && len(prev.names) > 0
curHasContent := cur.exists && len(cur.names) > 0
if curHasContent && (!prevHasContent || !prev.equal(cur)) && w.OnChange != nil {
w.OnChange()
}
prev = cur
}
}
}
// scanDir returns a snapshot of the *.yml/*.yaml files in dir. If the
// directory cannot be read (missing, permission denied) the snapshot reports
// exists=false; the next successful scan will detect the recovery and fire
// OnChange.
func scanDir(dir string) dirSnapshot {
snap := newDirSnapshot()
entries, err := os.ReadDir(dir)
if err != nil {
return snap // exists=false
}
snap.exists = true
for _, e := range entries {
if e.IsDir() {
continue
}
name := e.Name()
if !strings.HasSuffix(name, ".yml") && !strings.HasSuffix(name, ".yaml") {
continue
}
fi, err := os.Stat(filepath.Join(dir, name))
if err != nil {
// File disappeared between ReadDir and Stat; skip it — the
// next poll will observe the removal cleanly.
continue
}
snap.names = append(snap.names, name)
snap.states[name] = snapshot{
exists: true,
modTime: fi.ModTime(),
size: fi.Size(),
}
}
sort.Strings(snap.names)
return snap
}
+199
View File
@@ -0,0 +1,199 @@
package configwatcher
import (
"context"
"os"
"path/filepath"
"sync/atomic"
"testing"
"time"
"github.com/stretchr/testify/require"
)
// startDirWatcher launches w.Run in a goroutine and returns a function that
// cancels the context and waits for Run to return.
func startDirWatcher(t *testing.T, w *DirWatcher) func() {
t.Helper()
ctx, cancel := context.WithCancel(context.Background())
done := make(chan struct{})
go func() {
w.Run(ctx)
close(done)
}()
return func() {
cancel()
select {
case <-done:
case <-time.After(2 * time.Second):
t.Fatal("DirWatcher did not stop within 2s of cancel")
}
}
}
func writeYAMLInDir(t *testing.T, dir, name, content string) {
t.Helper()
require.NoError(t, os.WriteFile(filepath.Join(dir, name), []byte(content), 0o644))
}
func TestDirWatcher_NoFireOnBaseline(t *testing.T) {
dir := t.TempDir()
writeYAMLInDir(t, dir, "a.yaml", "a")
var n int64
stop := startDirWatcher(t, &DirWatcher{
Path: dir,
Interval: testInterval,
OnChange: func() { atomic.AddInt64(&n, 1) },
})
defer stop()
time.Sleep(testInterval * 5)
require.Equal(t, int64(0), atomic.LoadInt64(&n), "baseline poll must not fire")
}
func TestDirWatcher_DetectsFileAdd(t *testing.T) {
dir := t.TempDir()
writeYAMLInDir(t, dir, "a.yaml", "a")
var n int64
stop := startDirWatcher(t, &DirWatcher{
Path: dir,
Interval: testInterval,
OnChange: func() { atomic.AddInt64(&n, 1) },
})
defer stop()
time.Sleep(testInterval * 2)
writeYAMLInDir(t, dir, "b.yaml", "b")
require.True(t, waitForCount(t, &n, 1, time.Second), "callback should fire when a file is added")
}
func TestDirWatcher_DetectsFileRemoval(t *testing.T) {
dir := t.TempDir()
writeYAMLInDir(t, dir, "a.yaml", "a")
writeYAMLInDir(t, dir, "b.yaml", "b")
var n int64
stop := startDirWatcher(t, &DirWatcher{
Path: dir,
Interval: testInterval,
OnChange: func() { atomic.AddInt64(&n, 1) },
})
defer stop()
time.Sleep(testInterval * 2)
require.NoError(t, os.Remove(filepath.Join(dir, "b.yaml")))
require.True(t, waitForCount(t, &n, 1, time.Second), "callback should fire when a file is removed")
}
func TestDirWatcher_DetectsModTimeChange(t *testing.T) {
dir := t.TempDir()
writeYAMLInDir(t, dir, "a.yaml", "a")
base := time.Now().Add(-1 * time.Hour).Truncate(time.Second)
require.NoError(t, os.Chtimes(filepath.Join(dir, "a.yaml"), base, base))
var n int64
stop := startDirWatcher(t, &DirWatcher{
Path: dir,
Interval: testInterval,
OnChange: func() { atomic.AddInt64(&n, 1) },
})
defer stop()
time.Sleep(testInterval * 2)
require.NoError(t, os.Chtimes(filepath.Join(dir, "a.yaml"), base.Add(10*time.Second), base.Add(10*time.Second)))
require.True(t, waitForCount(t, &n, 1, time.Second), "callback should fire after mtime change")
}
func TestDirWatcher_IgnoresNonYAMLFiles(t *testing.T) {
dir := t.TempDir()
writeYAMLInDir(t, dir, "a.yaml", "a")
var n int64
stop := startDirWatcher(t, &DirWatcher{
Path: dir,
Interval: testInterval,
OnChange: func() { atomic.AddInt64(&n, 1) },
})
defer stop()
time.Sleep(testInterval * 2)
// Adding a .txt file must not fire.
require.NoError(t, os.WriteFile(filepath.Join(dir, "notes.txt"), []byte("hi"), 0o644))
time.Sleep(testInterval * 4)
require.Equal(t, int64(0), atomic.LoadInt64(&n), "non-YAML files must be ignored")
// Adding a .yml file must fire.
writeYAMLInDir(t, dir, "b.yml", "b")
require.True(t, waitForCount(t, &n, 1, time.Second), "callback should fire for *.yml files")
}
func TestDirWatcher_MissingDirRecovers(t *testing.T) {
dir := t.TempDir()
writeYAMLInDir(t, dir, "a.yaml", "a")
var n int64
stop := startDirWatcher(t, &DirWatcher{
Path: dir,
Interval: testInterval,
OnChange: func() { atomic.AddInt64(&n, 1) },
})
defer stop()
time.Sleep(testInterval * 2)
// Remove the directory. No fire expected on disappearance alone.
require.NoError(t, os.RemoveAll(dir))
time.Sleep(testInterval * 3)
require.Equal(t, int64(0), atomic.LoadInt64(&n), "directory removal alone must not fire")
// Recreate the directory and a YAML file; the recovery should fire.
require.NoError(t, os.MkdirAll(dir, 0o755))
writeYAMLInDir(t, dir, "recovered.yaml", "r")
require.True(t, waitForCount(t, &n, 1, time.Second), "callback should fire when dir returns with content")
}
func TestDirWatcher_EmptyDirSuppressedThenRecovers(t *testing.T) {
// Present-with-content → empty (all YAML removed, dir still exists)
// must stay quiet — treated as transient per the documented policy.
// The transition back to content fires.
dir := t.TempDir()
writeYAMLInDir(t, dir, "a.yaml", "a")
var n int64
stop := startDirWatcher(t, &DirWatcher{
Path: dir,
Interval: testInterval,
OnChange: func() { atomic.AddInt64(&n, 1) },
})
defer stop()
time.Sleep(testInterval * 2)
// Remove the only YAML file. Dir still exists but is empty of YAML.
require.NoError(t, os.Remove(filepath.Join(dir, "a.yaml")))
time.Sleep(testInterval * 4)
require.Equal(t, int64(0), atomic.LoadInt64(&n), "emptying the directory must not fire")
// Add a YAML file back; transition to present-with-content fires.
writeYAMLInDir(t, dir, "c.yaml", "c")
require.True(t, waitForCount(t, &n, 1, time.Second), "callback should fire when content returns")
}
func TestDirWatcher_ContextCancelStopsRun(t *testing.T) {
dir := t.TempDir()
writeYAMLInDir(t, dir, "a.yaml", "a")
w := &DirWatcher{Path: dir, Interval: testInterval}
ctx, cancel := context.WithCancel(context.Background())
done := make(chan struct{})
go func() { w.Run(ctx); close(done) }()
time.Sleep(testInterval * 2)
cancel()
select {
case <-done:
case <-time.After(2 * time.Second):
t.Fatal("Run did not return within 2s of cancel")
}
}
+37 -19
View File
@@ -55,7 +55,8 @@ var logTimeFormats = map[string]string{
}
func main() {
flagConfig := flag.String("config", "", "path to config file (required)")
flagConfig := flag.String("config", "", "path to config file")
flagConfigDir := flag.String("config-dir", "", "directory of *.yml/*.yaml config files (additive to -config)")
flagListen := flag.String("listen", "", "listen address (default :8080 or :8443 for TLS)")
flagCertFile := flag.String("tls-cert-file", "", "TLS certificate file")
flagKeyFile := flag.String("tls-key-file", "", "TLS key file")
@@ -68,8 +69,8 @@ func main() {
os.Exit(0)
}
if *flagConfig == "" {
slog.Error("-config is required")
if *flagConfig == "" && *flagConfigDir == "" {
slog.Error("at least one of -config or -config-dir must be provided")
os.Exit(1)
}
@@ -88,10 +89,9 @@ func main() {
}
}
configPath := *flagConfig
cfg, err := config.LoadConfig(configPath)
cfg, err := config.LoadConfigSources(*flagConfig, *flagConfigDir)
if err != nil {
slog.Error("failed to load config", "path", configPath, "error", err)
slog.Error("failed to load config", "config", *flagConfig, "config-dir", *flagConfigDir, "error", err)
os.Exit(1)
}
@@ -187,7 +187,7 @@ func main() {
proxyLog.Info("reloading configuration")
newCfg, err := config.LoadConfig(configPath)
newCfg, err := config.LoadConfigSources(*flagConfig, *flagConfigDir)
if err != nil {
proxyLog.Warnf("failed to reload config: %v", err)
return
@@ -230,19 +230,37 @@ func main() {
defer watcherCancel()
if *flagWatchConfig {
absConfigPath, err := filepath.Abs(configPath)
if err != nil {
slog.Error("watch-config: failed to resolve config path", "error", err)
os.Exit(1)
}
proxyLog.Info("watching configuration for changes (poll-based, 2s interval)")
go func() {
(&configwatcher.Watcher{
Path: absConfigPath,
Interval: configwatcher.DefaultInterval,
OnChange: reload,
}).Run(watcherCtx)
}()
if *flagConfig != "" {
absConfigPath, err := filepath.Abs(*flagConfig)
if err != nil {
slog.Error("watch-config: failed to resolve config path", "error", err)
os.Exit(1)
}
go func() {
(&configwatcher.Watcher{
Path: absConfigPath,
Interval: configwatcher.DefaultInterval,
OnChange: reload,
}).Run(watcherCtx)
}()
}
if *flagConfigDir != "" {
absConfigDir, err := filepath.Abs(*flagConfigDir)
if err != nil {
slog.Error("watch-config: failed to resolve config-dir path", "error", err)
os.Exit(1)
}
go func() {
(&configwatcher.DirWatcher{
Path: absConfigDir,
Interval: configwatcher.DefaultInterval,
OnChange: reload,
}).Run(watcherCtx)
}()
}
}
sigChan := make(chan os.Signal, 1)
+2 -1
View File
@@ -8,7 +8,7 @@
import Performance from "./routes/Performance.svelte";
import Playground from "./routes/Playground.svelte";
import PlaygroundStub from "./routes/PlaygroundStub.svelte";
import { enableAPIEvents } from "./stores/api";
import { enableAPIEvents, checkPerformanceEnabled } from "./stores/api";
import { initScreenWidth, initSystemThemeListener, isDarkMode, appTitle, connectionState } from "./stores/theme";
import { currentRoute } from "./stores/route";
@@ -39,6 +39,7 @@
const cleanupScreenWidth = initScreenWidth();
const cleanupSystemTheme = initSystemThemeListener();
enableAPIEvents(true);
checkPerformanceEnabled();
return () => {
cleanupScreenWidth();
+13 -10
View File
@@ -3,6 +3,7 @@
import { screenWidth, toggleTheme, themeMode, appTitle, isNarrow } from "../stores/theme";
import { currentRoute } from "../stores/route";
import { playgroundActivity } from "../stores/playgroundActivity";
import { performanceEnabled } from "../stores/api";
import ConnectionStatus from "./ConnectionStatus.svelte";
function handleTitleChange(newTitle: string): void {
@@ -84,16 +85,18 @@
>
Logs
</a>
<a
href="/performance"
use:link
class="text-gray-600 hover:text-black dark:text-gray-300 dark:hover:text-gray-100 p-1 whitespace-nowrap"
class:font-semibold={isActive("/performance", $currentRoute)}
class:underline={isActive("/performance", $currentRoute)}
class:underline-offset-4={isActive("/performance", $currentRoute)}
>
Performance
</a>
{#if $performanceEnabled}
<a
href="/performance"
use:link
class="text-gray-600 hover:text-black dark:text-gray-300 dark:hover:text-gray-100 p-1 whitespace-nowrap"
class:font-semibold={isActive("/performance", $currentRoute)}
class:underline={isActive("/performance", $currentRoute)}
class:underline-offset-4={isActive("/performance", $currentRoute)}
>
Performance
</a>
{/if}
<button onclick={toggleTheme} title="Toggle theme (current: {$themeMode})">
{#if $themeMode === "system"}
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="currentColor" class="w-5 h-5">
@@ -0,0 +1,359 @@
<script lang="ts">
import { models } from "../../stores/api";
import { persistentStore } from "../../stores/persistent";
import { generateVideoSync } from "../../lib/videoApi";
import { playgroundStores } from "../../stores/playgroundActivity";
import ModelSelector from "./ModelSelector.svelte";
import ExpandableTextarea from "./ExpandableTextarea.svelte";
const selectedModelStore = persistentStore<string>("playground-video-model", "");
const selectedSizeStore = persistentStore<string>("playground-video-size", "");
const negativePromptStore = persistentStore<string>("playground-video-negative-prompt", "");
const numFramesStore = persistentStore<number>("playground-video-num-frames", 0);
const fpsStore = persistentStore<number>("playground-video-fps", 0);
const stepsStore = persistentStore<number>("playground-video-steps", 0);
const guidanceStore = persistentStore<number>("playground-video-guidance", 0);
const seedStore = persistentStore<number>("playground-video-seed", -1);
let prompt = $state("");
let isGenerating = $state(false);
let videoUrl = $state<string | null>(null);
let videoMime = $state("video/mp4");
let videoBytes = $state(0);
let elapsedSeconds = $state(0);
let error = $state<string | null>(null);
let abortController = $state<AbortController | null>(null);
let showSettings = $state(false);
let initImage = $state<File | null>(null);
let initImagePreview = $state<string | null>(null);
let fileInput = $state<HTMLInputElement | null>(null);
let timer: ReturnType<typeof setInterval> | null = null;
let hasModels = $derived($models.some((m) => !m.unlisted));
$effect(() => {
playgroundStores.videoGenerating.set(isGenerating);
});
function onInitImageChange(event: Event) {
const input = event.target as HTMLInputElement;
const file = input.files?.[0] ?? null;
setInitImage(file);
}
function setInitImage(file: File | null) {
if (initImagePreview) URL.revokeObjectURL(initImagePreview);
initImage = file;
initImagePreview = file ? URL.createObjectURL(file) : null;
}
function clearInitImage() {
setInitImage(null);
if (fileInput) fileInput.value = "";
}
async function generate() {
const trimmedPrompt = prompt.trim();
if (!trimmedPrompt || !$selectedModelStore || isGenerating) return;
isGenerating = true;
error = null;
elapsedSeconds = 0;
abortController = new AbortController();
timer = setInterval(() => (elapsedSeconds += 1), 1000);
try {
const result = await generateVideoSync(
$selectedModelStore,
trimmedPrompt,
{
negativePrompt: $negativePromptStore || undefined,
size: $selectedSizeStore || undefined,
numFrames: $numFramesStore,
fps: $fpsStore,
steps: $stepsStore,
guidanceScale: $guidanceStore,
seed: $seedStore,
initImage,
},
abortController.signal
);
if (videoUrl) URL.revokeObjectURL(videoUrl);
videoUrl = result.url;
videoMime = result.mime;
videoBytes = result.sizeBytes;
} catch (err) {
if (err instanceof Error && err.name === "AbortError") {
// User cancelled
} else {
error = err instanceof Error ? err.message : "An error occurred";
}
} finally {
isGenerating = false;
abortController = null;
if (timer) clearInterval(timer);
timer = null;
}
}
function cancelGeneration() {
abortController?.abort();
}
function clearVideo() {
if (videoUrl) URL.revokeObjectURL(videoUrl);
videoUrl = null;
error = null;
prompt = "";
}
function downloadVideo() {
if (!videoUrl) return;
const ext = videoMime.includes("webm") ? "webm" : "mp4";
const link = document.createElement("a");
link.href = videoUrl;
link.download = `generated-video-${Date.now()}.${ext}`;
document.body.appendChild(link);
link.click();
document.body.removeChild(link);
}
function formatBytes(n: number): string {
if (n > 1024 * 1024) return `${(n / 1024 / 1024).toFixed(1)} MB`;
if (n > 1024) return `${(n / 1024).toFixed(0)} KB`;
return `${n} B`;
}
function formatElapsed(s: number): string {
const m = Math.floor(s / 60);
return m > 0 ? `${m}m ${s % 60}s` : `${s}s`;
}
function handleKeyDown(event: KeyboardEvent) {
if (event.key === "Enter" && !event.shiftKey) {
event.preventDefault();
generate();
}
}
</script>
<div class="flex flex-col h-full">
<!-- Model selector and options -->
<div class="shrink-0 flex flex-wrap gap-2 mb-4">
<ModelSelector
bind:value={$selectedModelStore}
placeholder="Select a video model..."
disabled={isGenerating}
capabilities={["video_generation", "image_to_video"]}
matchAny={true}
/>
<select
class="px-3 py-2 rounded border border-gray-200 dark:border-white/10 bg-surface focus:outline-none focus:ring-2 focus:ring-primary"
bind:value={$selectedSizeStore}
disabled={isGenerating}
>
<option value="">Model default size</option>
<optgroup label="Landscape">
<option value="1280x704">1280x704</option>
<option value="832x480">832x480</option>
</optgroup>
<optgroup label="Portrait">
<option value="704x1280">704x1280</option>
<option value="480x832">480x832</option>
</optgroup>
<optgroup label="Square">
<option value="704x704">704x704</option>
</optgroup>
</select>
<button
class="px-3 py-2 rounded border border-gray-200 dark:border-white/10 bg-surface hover:bg-secondary-hover transition-colors"
onclick={() => (showSettings = !showSettings)}
>
{showSettings ? "Hide Settings" : "Settings"}
</button>
</div>
<!-- Settings panel -->
{#if showSettings}
<div class="shrink-0 mb-4 p-4 rounded border border-gray-200 dark:border-white/10 bg-surface">
<div class="grid grid-cols-2 md:grid-cols-5 gap-3 mb-3">
<label class="flex flex-col gap-1">
<span class="text-xs text-txtsecondary">Frames (0 = default)</span>
<input
type="number"
class="px-2 py-1 rounded border border-gray-200 dark:border-white/10 bg-surface focus:outline-none focus:ring-2 focus:ring-primary"
bind:value={$numFramesStore}
min="0"
max="1000"
/>
</label>
<label class="flex flex-col gap-1">
<span class="text-xs text-txtsecondary">FPS (0 = default)</span>
<input
type="number"
class="px-2 py-1 rounded border border-gray-200 dark:border-white/10 bg-surface focus:outline-none focus:ring-2 focus:ring-primary"
bind:value={$fpsStore}
min="0"
max="60"
/>
</label>
<label class="flex flex-col gap-1">
<span class="text-xs text-txtsecondary">Steps (0 = default)</span>
<input
type="number"
class="px-2 py-1 rounded border border-gray-200 dark:border-white/10 bg-surface focus:outline-none focus:ring-2 focus:ring-primary"
bind:value={$stepsStore}
min="0"
max="150"
/>
</label>
<label class="flex flex-col gap-1">
<span class="text-xs text-txtsecondary">Guidance (0 = default)</span>
<input
type="number"
class="px-2 py-1 rounded border border-gray-200 dark:border-white/10 bg-surface focus:outline-none focus:ring-2 focus:ring-primary"
bind:value={$guidanceStore}
min="0"
max="30"
step="0.5"
/>
</label>
<label class="flex flex-col gap-1">
<span class="text-xs text-txtsecondary">Seed (-1 = random)</span>
<input
type="number"
class="px-2 py-1 rounded border border-gray-200 dark:border-white/10 bg-surface focus:outline-none focus:ring-2 focus:ring-primary"
bind:value={$seedStore}
min="-1"
/>
</label>
</div>
<label class="flex flex-col gap-1">
<span class="text-xs text-txtsecondary">Negative Prompt</span>
<textarea
class="px-2 py-1 rounded border border-gray-200 dark:border-white/10 bg-surface focus:outline-none focus:ring-2 focus:ring-primary resize-y text-sm"
bind:value={$negativePromptStore}
rows="2"
placeholder="Elements to avoid..."
></textarea>
</label>
</div>
{/if}
<!-- Empty state for no models configured -->
{#if !hasModels}
<div class="flex-1 flex items-center justify-center text-txtsecondary">
<p>No models configured. Add models to your configuration to generate videos.</p>
</div>
{:else}
<!-- Video display area -->
<div
class="flex-1 overflow-auto mb-4 flex items-center justify-center bg-surface border border-gray-200 dark:border-white/10 rounded"
>
{#if isGenerating}
<div class="text-center text-txtsecondary">
<div
class="inline-block w-8 h-8 border-4 border-primary border-t-transparent rounded-full animate-spin mb-2"
></div>
<p>Generating video… {formatElapsed(elapsedSeconds)}</p>
<p class="text-xs mt-1">
Video generation takes minutes — a cold model load takes longer still.
</p>
</div>
{:else if error}
<div class="text-center text-red-500 p-4">
<p class="font-medium">Error</p>
<p class="text-sm mt-1 break-all">{error}</p>
</div>
{:else if videoUrl}
<div class="relative max-w-full max-h-full flex items-center justify-center">
<!-- svelte-ignore a11y_media_has_caption -->
<video src={videoUrl} class="max-w-full max-h-full object-contain" controls autoplay loop></video>
<button
class="absolute bottom-2 right-2 p-2 bg-black/60 hover:bg-black/80 text-white rounded-full transition-colors"
onclick={downloadVideo}
aria-label="Download video"
title="Download ({formatBytes(videoBytes)})"
>
<svg class="w-5 h-5" fill="none" stroke="currentColor" viewBox="0 0 24 24">
<path
stroke-linecap="round"
stroke-linejoin="round"
stroke-width="2"
d="M4 16v1a3 3 0 003 3h10a3 3 0 003-3v-1m-4-4l-4 4m0 0l-4-4m4 4V4"
></path>
</svg>
</button>
</div>
{:else}
<div class="text-center text-txtsecondary">
<p>Enter a prompt below to generate a short video clip</p>
<p class="text-xs mt-1">Attach an image to animate it (image-to-video)</p>
</div>
{/if}
</div>
<!-- Init image chip -->
{#if initImagePreview}
<div class="shrink-0 flex items-center gap-2 mb-2">
<img src={initImagePreview} alt="Conditioning frame" class="h-12 rounded border border-gray-200 dark:border-white/10" />
<span class="text-xs text-txtsecondary">Animating this image (image-to-video)</span>
<button
class="px-1.5 py-0.5 text-xs rounded border border-gray-200 dark:border-white/10 hover:bg-red-500 hover:text-white hover:border-red-500 transition-colors"
onclick={clearInitImage}
disabled={isGenerating}
>
Remove
</button>
</div>
{/if}
<!-- Prompt input area -->
<div class="shrink-0 flex flex-col md:flex-row gap-2">
<ExpandableTextarea
bind:value={prompt}
placeholder="Describe the video you want to generate..."
rows={3}
onkeydown={handleKeyDown}
disabled={isGenerating || !$selectedModelStore}
/>
<div class="flex flex-row md:flex-col gap-2">
{#if isGenerating}
<button class="btn bg-red-500 hover:bg-red-600 text-white flex-1 md:flex-none" onclick={cancelGeneration}>
Cancel
</button>
{:else}
<button
class="btn bg-primary text-btn-primary-text hover:opacity-90 flex-1 md:flex-none"
onclick={generate}
disabled={!prompt.trim() || !$selectedModelStore}
>
Generate
</button>
<label
class="btn flex-1 md:flex-none text-center cursor-pointer {isGenerating ? 'opacity-50 pointer-events-none' : ''}"
>
Image…
<input
type="file"
accept="image/*"
class="hidden"
bind:this={fileInput}
onchange={onInitImageChange}
/>
</label>
<button
class="btn flex-1 md:flex-none"
onclick={clearVideo}
disabled={!videoUrl && !error && !prompt.trim()}
>
Clear
</button>
{/if}
</div>
</div>
{/if}
</div>
+5
View File
@@ -8,6 +8,8 @@ export interface ModelCapabilities {
audio_speech?: boolean;
image_generation?: boolean;
image_to_image?: boolean;
video_generation?: boolean;
image_to_video?: boolean;
function_calling?: boolean;
reranker?: boolean;
}
@@ -25,6 +27,8 @@ export interface Model {
export interface TokenMetrics {
cache_tokens: number;
draft_tokens: number;
draft_acc_tokens: number;
input_tokens: number;
output_tokens: number;
prompt_per_second: number;
@@ -41,6 +45,7 @@ export interface ActivityLogEntry {
tokens: TokenMetrics;
duration_ms: number;
has_capture: boolean;
error_msg?: string;
metadata?: Record<string, string>;
}
+57
View File
@@ -0,0 +1,57 @@
export interface VideoGenerationOptions {
negativePrompt?: string;
size?: string; // "WxH"
numFrames?: number;
fps?: number;
steps?: number;
guidanceScale?: number;
seed?: number; // -1 = random (omitted)
initImage?: File | null; // image-to-video conditioning frame
}
// POST /v1/videos/sync (multipart; llama-swap routes by the `model` field).
// The response body IS the encoded clip — returns an object URL for <video>.
// Callers must URL.revokeObjectURL() the result when done with it.
export async function generateVideoSync(
model: string,
prompt: string,
options: VideoGenerationOptions = {},
signal?: AbortSignal
): Promise<{ url: string; mime: string; sizeBytes: number }> {
const form = new FormData();
form.set("model", model);
form.set("prompt", prompt);
if (options.negativePrompt) form.set("negative_prompt", options.negativePrompt);
if (options.size) {
const [w, h] = options.size.split("x").map(Number);
if (w > 0 && h > 0) {
// Both conventions: vLLM-Omni reads width/height, OpenAI-shaped
// upstreams read size; they can never disagree.
form.set("width", String(w));
form.set("height", String(h));
form.set("size", options.size);
}
}
if (options.numFrames && options.numFrames > 0) form.set("num_frames", String(options.numFrames));
if (options.fps && options.fps > 0) form.set("fps", String(options.fps));
if (options.steps && options.steps > 0) form.set("num_inference_steps", String(options.steps));
if (options.guidanceScale && options.guidanceScale > 0)
form.set("guidance_scale", String(options.guidanceScale));
if (options.seed !== undefined && options.seed >= 0) form.set("seed", String(options.seed));
if (options.initImage) form.set("input_reference", options.initImage, options.initImage.name);
const response = await fetch("/v1/videos/sync", {
method: "POST",
body: form,
signal,
});
if (!response.ok) {
const errorText = await response.text();
throw new Error(`Video API error: ${response.status} - ${errorText}`);
}
const blob = await response.blob();
const mime = response.headers.get("Content-Type") || "video/mp4";
return { url: URL.createObjectURL(blob), mime, sizeBytes: blob.size };
}
+17 -2
View File
@@ -21,11 +21,12 @@
{ key: "time", label: "Time", defaultVisible: true },
{ key: "model", label: "Model", defaultVisible: true },
{ key: "req_path", label: "Path", defaultVisible: false },
{ key: "resp_status_code", label: "Status", defaultVisible: false },
{ key: "resp_status_code", label: "Status", defaultVisible: true },
{ key: "resp_content_type", label: "Content-Type", defaultVisible: false },
{ key: "cached", label: "Cached", defaultVisible: true },
{ key: "prompt", label: "Prompt", defaultVisible: true },
{ key: "generated", label: "Generated", defaultVisible: true },
{ key: "drafted", label: "Drafted", defaultVisible: false },
{ key: "prompt_speed", label: "Prompt Speed", defaultVisible: true },
{ key: "gen_speed", label: "Gen Speed", defaultVisible: true },
{ key: "duration", label: "Duration", defaultVisible: true },
@@ -158,6 +159,10 @@
return speed < 0 ? "unknown" : speed.toFixed(2) + " t/s";
}
function formatDrafted(drafted: number, accepted: number): string {
return drafted > 0 ? (accepted * 100 / drafted).toFixed(1) + "% (" + accepted + "/" + drafted + ")" : "-";
}
function formatDuration(ms: number): string {
return (ms / 1000).toFixed(2) + "s";
}
@@ -273,6 +278,8 @@
Cached <Tooltip content="prompt tokens from cache" />
{:else if key === "prompt"}
Prompt <Tooltip content="new prompt tokens processed" />
{:else if key === "drafted"}
Drafted <Tooltip content="acceptance rate (accepted/drafted)" />
{:else}
{columnLabelMap[key] ?? key}
{/if}
@@ -301,7 +308,13 @@
{:else if key === "req_path"}
{metric.req_path || "-"}
{:else if key === "resp_status_code"}
{metric.resp_status_code || "-"}
{#if metric.error_msg}
<span class="text-red-500 dark:text-red-400 cursor-help" title={metric.error_msg}>
{metric.resp_status_code || "-"}
</span>
{:else}
{metric.resp_status_code || "-"}
{/if}
{:else if key === "resp_content_type"}
{metric.resp_content_type || "-"}
{:else if key === "cached"}
@@ -310,6 +323,8 @@
{metric.tokens.input_tokens.toLocaleString()}
{:else if key === "generated"}
{metric.tokens.output_tokens.toLocaleString()}
{:else if key === "drafted"}
{formatDrafted(metric.tokens.draft_tokens, metric.tokens.draft_acc_tokens)}
{:else if key === "prompt_speed"}
{formatSpeed(metric.tokens.prompt_per_second)}
{:else if key === "gen_speed"}
+6 -1
View File
@@ -2,12 +2,13 @@
import { persistentStore } from "../stores/persistent";
import ChatInterface from "../components/playground/ChatInterface.svelte";
import ImageInterface from "../components/playground/ImageInterface.svelte";
import VideoInterface from "../components/playground/VideoInterface.svelte";
import AudioInterface from "../components/playground/AudioInterface.svelte";
import SpeechInterface from "../components/playground/SpeechInterface.svelte";
import RerankInterface from "../components/playground/RerankInterface.svelte";
import ConcurrencyInterface from "../components/playground/ConcurrencyInterface.svelte";
type Tab = "chat" | "images" | "speech" | "audio" | "rerank" | "concurrency";
type Tab = "chat" | "images" | "video" | "speech" | "audio" | "rerank" | "concurrency";
const selectedTabStore = persistentStore<Tab>("playground-selected-tab", "chat");
let mobileMenuOpen = $state(false);
@@ -15,6 +16,7 @@
const tabs: { id: Tab; label: string }[] = [
{ id: "chat", label: "Chat" },
{ id: "images", label: "Images" },
{ id: "video", label: "Video" },
{ id: "speech", label: "Speech" },
{ id: "audio", label: "Transcription" },
{ id: "rerank", label: "Rerank" },
@@ -92,6 +94,9 @@
<div class="h-full" class:tab-hidden={$selectedTabStore !== "images"}>
<ImageInterface />
</div>
<div class="h-full" class:tab-hidden={$selectedTabStore !== "video"}>
<VideoInterface />
</div>
<div class="h-full" class:tab-hidden={$selectedTabStore !== "speech"}>
<SpeechInterface />
</div>
+15
View File
@@ -19,6 +19,7 @@ export const proxyLogs = writable<string>("");
export const upstreamLogs = writable<string>("");
export const metrics = writable<ActivityLogEntry[]>([]);
export const inFlightRequests = writable<number>(0);
export const performanceEnabled = writable<boolean>(false);
export const versionInfo = writable<VersionInfo>({
build_date: "unknown",
commit: "unknown",
@@ -210,6 +211,20 @@ export async function getCapture(id: number): Promise<ReqRespCapture | null> {
}
}
export async function checkPerformanceEnabled(): Promise<void> {
try {
const response = await fetch("/api/performance");
if (!response.ok) {
performanceEnabled.set(false);
return;
}
const data = await response.json();
performanceEnabled.set(data.enabled);
} catch {
performanceEnabled.set(false);
}
}
export async function fetchPerformance(after?: string): Promise<PerformanceResponse | null> {
try {
const url = after ? `/api/performance?after=${encodeURIComponent(after)}` : "/api/performance";
+4 -2
View File
@@ -2,18 +2,20 @@ import { writable, derived } from "svelte/store";
const chatStreaming = writable(false);
const imageGenerating = writable(false);
const videoGenerating = writable(false);
const speechGenerating = writable(false);
const audioTranscribing = writable(false);
const rerankLoading = writable(false);
export const playgroundActivity = derived(
[chatStreaming, imageGenerating, speechGenerating, audioTranscribing, rerankLoading],
([$chat, $image, $speech, $audio, $rerank]) => $chat || $image || $speech || $audio || $rerank
[chatStreaming, imageGenerating, videoGenerating, speechGenerating, audioTranscribing, rerankLoading],
([$chat, $image, $video, $speech, $audio, $rerank]) => $chat || $image || $video || $speech || $audio || $rerank
);
export const playgroundStores = {
chatStreaming,
imageGenerating,
videoGenerating,
speechGenerating,
audioTranscribing,
rerankLoading,