Compare commits
18
Commits
v228
..
c87893ba6f
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
c87893ba6f | ||
|
|
0358fe321e | ||
|
|
2f0bb2556e | ||
|
|
fdd89d6580 | ||
|
|
1934efb0b5 | ||
|
|
b83edeb9ef | ||
|
|
7886758053 | ||
|
|
d297852b9e | ||
|
|
cecb89880e | ||
|
|
0292c90ca1 | ||
|
|
617c7dc6b9 | ||
|
|
542b79dacf | ||
|
|
0a25b3bd31 | ||
|
|
32bc781326 | ||
|
|
316ad63f76 | ||
|
|
e37077a963 | ||
|
|
eff9b60434 | ||
|
|
9bcddad91b |
@@ -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"
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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
@@ -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
@@ -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
|
||||
|
||||
@@ -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" ]
|
||||
@@ -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.
|
||||
@@ -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")
|
||||
}
|
||||
+57
-661
@@ -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.
|
||||
@@ -178,12 +168,66 @@ type SchedulerConfig struct {
|
||||
|
||||
type SchedulerSettings struct {
|
||||
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
|
||||
}
|
||||
|
||||
@@ -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",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
@@ -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")
|
||||
})
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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, ®exp.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, ®exp.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")
|
||||
}
|
||||
@@ -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
|
||||
@@ -421,6 +432,7 @@ func (b *baseRouter) ServeHTTP(w http.ResponseWriter, req *http.Request) {
|
||||
|
||||
hr := scheduler.HandlerReq{
|
||||
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
|
||||
|
||||
@@ -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{})
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -111,11 +119,49 @@ func New(conf config.Config, name string, logger *logmon.Monitor, planner Swappe
|
||||
// HandlerReq is one in-flight ServeHTTP request waiting for a routing decision.
|
||||
type HandlerReq struct {
|
||||
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 {
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
+43
-3
@@ -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,16 +293,26 @@ 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
|
||||
}
|
||||
|
||||
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) {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
@@ -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)
|
||||
|
||||
@@ -5,11 +5,15 @@ import (
|
||||
"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"
|
||||
)
|
||||
|
||||
@@ -156,6 +160,91 @@ func upstreamMetricsServer(response string) *Server {
|
||||
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)
|
||||
@@ -214,6 +303,8 @@ func TestServer_HandleUpstream_MetricsSkipsGET(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestServer_HandleMetrics_Unavailable(t *testing.T) {
|
||||
// No perf monitor and a scheduler that reports no queue stats: nothing to
|
||||
// serve.
|
||||
s := newTestServer(newStubRouter(nil, ""), newStubRouter(nil, ""))
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
@@ -223,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, ""))
|
||||
|
||||
|
||||
@@ -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,
|
||||
})
|
||||
|
||||
@@ -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
|
||||
|
||||
+100
-9
@@ -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,9 +215,20 @@ func (mp *metricsMonitor) record(modelID string, r *http.Request, recorder *resp
|
||||
}
|
||||
|
||||
tm.ID = mp.queueMetrics(tm)
|
||||
if mp.enableCaptures {
|
||||
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: tm.ID,
|
||||
ID: id,
|
||||
ReqPath: r.URL.Path,
|
||||
ReqHeaders: reqHeaders,
|
||||
}
|
||||
@@ -220,11 +243,71 @@ func (mp *metricsMonitor) record(modelID string, r *http.Request, recorder *resp
|
||||
if cf&captureRespBody != 0 {
|
||||
capture.RespBody = body
|
||||
}
|
||||
if mp.addCapture(capture) {
|
||||
tm.HasCapture = true
|
||||
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,
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -90,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}}]`
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -26,6 +26,9 @@ type ReqContextData struct {
|
||||
ModelID string
|
||||
Streaming bool
|
||||
SendLoadingState bool
|
||||
// Priority is the caller's scheduling priority, read from PriorityHeader.
|
||||
// Higher is more urgent; 0 (PriorityNormal) is the default. See priority.go.
|
||||
Priority int
|
||||
// Metadata is a request-scoped key/value bag that handlers may mutate
|
||||
// while processing. The metrics middleware copies it into ActivityLogEntry.
|
||||
Metadata map[string]string
|
||||
@@ -118,6 +121,7 @@ func FetchContext(r *http.Request, cfg config.Config) (ReqContextData, error) {
|
||||
if mc, ok := cfg.Models[realName]; ok {
|
||||
data.SendLoadingState = mc.SendLoadingState != nil && *mc.SendLoadingState
|
||||
}
|
||||
data.Priority = RequestPriority(r)
|
||||
*r = *r.WithContext(SetContext(r.Context(), data))
|
||||
return data, nil
|
||||
}
|
||||
@@ -137,6 +141,7 @@ func extractUpstreamContext(r *http.Request, cfg config.Config) (ReqContextData,
|
||||
ApiKey: ExtractAPIKey(r),
|
||||
Streaming: r.URL.Query().Get("stream") == "true",
|
||||
SendLoadingState: sendLoadingState(cfg, realName),
|
||||
Priority: RequestPriority(r),
|
||||
Metadata: make(map[string]string),
|
||||
}, true
|
||||
}
|
||||
|
||||
@@ -0,0 +1,89 @@
|
||||
package shared
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// PriorityHeader carries the caller's scheduling priority. Priority is a
|
||||
// property of the caller's intent, not of the model — the same model serves
|
||||
// both an interactive request and a batch job — so it travels as a header
|
||||
// rather than as per-model configuration, and leaves the OpenAI-compatible
|
||||
// request body untouched.
|
||||
const PriorityHeader = "X-LlamaSwap-Priority"
|
||||
|
||||
// Priority band anchors. The number is the interface: callers may send any
|
||||
// signed integer, and these names are conveniences that resolve to a value.
|
||||
//
|
||||
// Bands are spaced 100 apart so a consumer can add a small per-caller offset
|
||||
// (a subscription tier, say) on top of a band without ever promoting a request
|
||||
// across one: a "max member" batch job at -98 still loses to every normal
|
||||
// request at 0.
|
||||
const (
|
||||
PriorityInteractive = 100 // a human is waiting
|
||||
PriorityNormal = 0 // the default
|
||||
PriorityBatch = -100
|
||||
)
|
||||
|
||||
// bandHalfWidth is how far a priority may sit from a band anchor and still be
|
||||
// reported as that band. It is half the 100-point band spacing, so every
|
||||
// integer belongs to exactly one band.
|
||||
const bandHalfWidth = 50
|
||||
|
||||
// Band names used for metrics labels, ordered most to least urgent.
|
||||
const (
|
||||
BandInteractive = "interactive"
|
||||
BandNormal = "normal"
|
||||
BandBatch = "batch"
|
||||
)
|
||||
|
||||
// PriorityBands is every band name, in descending order of urgency. Metrics
|
||||
// emit a series per band regardless of whether any request currently occupies
|
||||
// it, so the label set stays stable.
|
||||
var PriorityBands = []string{BandInteractive, BandNormal, BandBatch}
|
||||
|
||||
var priorityAliases = map[string]int{
|
||||
"interactive": PriorityInteractive,
|
||||
"normal": PriorityNormal,
|
||||
"batch": PriorityBatch,
|
||||
}
|
||||
|
||||
// ParsePriority resolves a PriorityHeader value to a signed integer. A numeric
|
||||
// value is used as-is; a recognised alias resolves to its anchor. Anything
|
||||
// absent or unparseable is PriorityNormal, so a malformed header degrades to
|
||||
// the default rather than failing the request.
|
||||
//
|
||||
// Values are deliberately not clamped: the caller composes band, tier offset
|
||||
// and anything else it wants into one number, and llama-swap honours it.
|
||||
func ParsePriority(value string) int {
|
||||
value = strings.TrimSpace(value)
|
||||
if value == "" {
|
||||
return PriorityNormal
|
||||
}
|
||||
if n, err := strconv.Atoi(value); err == nil {
|
||||
return n
|
||||
}
|
||||
if n, ok := priorityAliases[strings.ToLower(value)]; ok {
|
||||
return n
|
||||
}
|
||||
return PriorityNormal
|
||||
}
|
||||
|
||||
// RequestPriority reads PriorityHeader off r and resolves it. See ParsePriority.
|
||||
func RequestPriority(r *http.Request) int {
|
||||
return ParsePriority(r.Header.Get(PriorityHeader))
|
||||
}
|
||||
|
||||
// PriorityBand buckets a priority into the band it belongs to, for metrics
|
||||
// labels. Values beyond the anchors saturate: +1000 is still "interactive".
|
||||
func PriorityBand(priority int) string {
|
||||
switch {
|
||||
case priority >= PriorityInteractive-bandHalfWidth:
|
||||
return BandInteractive
|
||||
case priority <= PriorityBatch+bandHalfWidth:
|
||||
return BandBatch
|
||||
default:
|
||||
return BandNormal
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,118 @@
|
||||
package shared
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/mostlygeek/llama-swap/internal/config"
|
||||
)
|
||||
|
||||
func TestShared_ParsePriority(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
value string
|
||||
want int
|
||||
}{
|
||||
{"absent", "", PriorityNormal},
|
||||
{"zero", "0", 0},
|
||||
{"positive", "100", 100},
|
||||
{"negative", "-100", -100},
|
||||
{"explicit plus", "+42", 42},
|
||||
{"surrounding space", " 100 ", 100},
|
||||
{"large value is not clamped", "100000", 100000},
|
||||
{"alias interactive", "interactive", PriorityInteractive},
|
||||
{"alias normal", "normal", PriorityNormal},
|
||||
{"alias batch", "batch", PriorityBatch},
|
||||
{"alias is case insensitive", "Batch", PriorityBatch},
|
||||
{"tier offset on a band", "-98", -98},
|
||||
{"unknown alias falls back to normal", "urgent", PriorityNormal},
|
||||
{"garbage falls back to normal", "!!", PriorityNormal},
|
||||
{"float falls back to normal", "1.5", PriorityNormal},
|
||||
}
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
if got := ParsePriority(tc.value); got != tc.want {
|
||||
t.Errorf("ParsePriority(%q)=%d want %d", tc.value, got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestShared_RequestPriority(t *testing.T) {
|
||||
r := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
||||
if got := RequestPriority(r); got != PriorityNormal {
|
||||
t.Errorf("no header: got %d want %d", got, PriorityNormal)
|
||||
}
|
||||
|
||||
r.Header.Set(PriorityHeader, "batch")
|
||||
if got := RequestPriority(r); got != PriorityBatch {
|
||||
t.Errorf("batch header: got %d want %d", got, PriorityBatch)
|
||||
}
|
||||
}
|
||||
|
||||
// TestShared_PriorityBand pins the bucketing the metrics labels rely on. The
|
||||
// important cases are the tier offsets: a max member's batch job at -98 must
|
||||
// still report as batch, not normal.
|
||||
func TestShared_PriorityBand(t *testing.T) {
|
||||
tests := []struct {
|
||||
priority int
|
||||
want string
|
||||
}{
|
||||
{PriorityInteractive, BandInteractive},
|
||||
{PriorityInteractive + 2, BandInteractive}, // max tier, interactive
|
||||
{1000, BandInteractive}, // saturates
|
||||
{50, BandInteractive}, // lower edge
|
||||
{49, BandNormal},
|
||||
{PriorityNormal, BandNormal},
|
||||
{2, BandNormal}, // max tier, normal
|
||||
{-49, BandNormal},
|
||||
{-50, BandBatch}, // upper edge
|
||||
{PriorityBatch + 2, BandBatch},
|
||||
{PriorityBatch, BandBatch},
|
||||
{-1000, BandBatch}, // saturates
|
||||
}
|
||||
for _, tc := range tests {
|
||||
if got := PriorityBand(tc.priority); got != tc.want {
|
||||
t.Errorf("PriorityBand(%d)=%q want %q", tc.priority, got, tc.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestShared_FetchContext_Priority verifies the header reaches the request
|
||||
// context, which is what carries it to the scheduler. Both entry points into
|
||||
// FetchContext are covered: the normal body-parsed path and /upstream/<model>.
|
||||
func TestShared_FetchContext_Priority(t *testing.T) {
|
||||
cfg := config.Config{Models: map[string]config.ModelConfig{"m1": {}}}
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
path string
|
||||
header string
|
||||
want int
|
||||
}{
|
||||
{"body path, no header", "/v1/chat/completions", "", PriorityNormal},
|
||||
{"body path, alias", "/v1/chat/completions", "interactive", PriorityInteractive},
|
||||
{"body path, number", "/v1/chat/completions", "-98", -98},
|
||||
{"upstream path", "/upstream/m1/v1/chat/completions", "batch", PriorityBatch},
|
||||
}
|
||||
|
||||
for _, c := range cases {
|
||||
t.Run(c.name, func(t *testing.T) {
|
||||
r := httptest.NewRequest(http.MethodPost, c.path, strings.NewReader(`{"model":"m1"}`))
|
||||
r.Header.Set("Content-Type", "application/json")
|
||||
if c.header != "" {
|
||||
r.Header.Set(PriorityHeader, c.header)
|
||||
}
|
||||
|
||||
data, err := FetchContext(r, cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("FetchContext: %v", err)
|
||||
}
|
||||
if data.Priority != c.want {
|
||||
t.Errorf("Priority=%d want %d", data.Priority, c.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
+27
-9
@@ -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,12 +230,14 @@ func main() {
|
||||
defer watcherCancel()
|
||||
|
||||
if *flagWatchConfig {
|
||||
absConfigPath, err := filepath.Abs(configPath)
|
||||
proxyLog.Info("watching configuration for changes (poll-based, 2s interval)")
|
||||
|
||||
if *flagConfig != "" {
|
||||
absConfigPath, err := filepath.Abs(*flagConfig)
|
||||
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,
|
||||
@@ -245,6 +247,22 @@ func main() {
|
||||
}()
|
||||
}
|
||||
|
||||
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)
|
||||
signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP)
|
||||
|
||||
|
||||
@@ -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();
|
||||
|
||||
@@ -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,6 +85,7 @@
|
||||
>
|
||||
Logs
|
||||
</a>
|
||||
{#if $performanceEnabled}
|
||||
<a
|
||||
href="/performance"
|
||||
use:link
|
||||
@@ -94,6 +96,7 @@
|
||||
>
|
||||
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>
|
||||
@@ -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>;
|
||||
}
|
||||
|
||||
|
||||
@@ -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 };
|
||||
}
|
||||
@@ -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"}
|
||||
{#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"}
|
||||
|
||||
@@ -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>
|
||||
|
||||
@@ -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";
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user