Precise KV-index routing: vLLM KV cache events through the agent, routed by leading blocks held
Plan §7.5 tier 2, opt-in per profile with kv_events. The agent starts vLLM publishing KV cache events to a loopback port it listens on (pure-Go zmq4, msgpack in both of vLLM's encodings), with sha256_cbor hashing and a fixed PYTHONHASHSEED; it mirrors each engine's blocks and sends net changes up the existing stream, with a full copy on every new one and a reset on gaps, restarts and clears. The control plane hands them to the router, for the sending node's own replicas only. The router tokenizes through a replica's /tokenize, hashes the tokens as vLLM does (checked against cbor2/hashlib golden values), and scores replicas by the share of leading blocks held, falling back to the approximate tree. New metrics beside llmgw_prefix_decisions_total. Untested against a running vLLM. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Ilker ManapandClaude Opus 5.5 committed
1 parent
d41104a452
commit
117f417ef5
43 files changed
+3496
-263
No files matched your search
+368
-226
File diff suppressed because it is too large.
Load diff
@@ -27,9 +27,36 @@ message AgentMessage {
|
||||
FabricScanReply fabric_scan = 5;
|
||||
FabricTestReply fabric_test = 6;
|
||||
UpgradeReply upgrade = 7;
|
||||
KVEvents kv_events = 8;
|
||||
}
|
||||
}
|
||||
|
||||
// KVEvents is a change to what one vLLM replica's prefix cache holds, read
|
||||
// by the agent from the engine's KV cache events and sent on for the router's
|
||||
// precise index (plan §7.5 tier 2).
|
||||
//
|
||||
// A net effect since the last one, not the events themselves: a block stored
|
||||
// and evicted between two sends is not worth the bytes. On every new stream
|
||||
// the agent first sends each replica in full, with clear_first set, because a
|
||||
// control plane that has just started or taken over knows nothing.
|
||||
message KVEvents {
|
||||
string replica_id = 1;
|
||||
// block_size is how many tokens make one of the replica's blocks.
|
||||
int32 block_size = 2;
|
||||
// clear_first says to forget what was known of the replica before applying the
|
||||
// rest: its cache was cleared, its engine restarted, events were lost, or
|
||||
// this is the first piece of a full copy.
|
||||
bool clear_first = 3;
|
||||
// stored and removed are block hashes as vLLM publishes them: the low 64
|
||||
// bits of the sha256_cbor digest.
|
||||
repeated fixed64 stored = 4;
|
||||
repeated fixed64 removed = 5;
|
||||
// problem, when set, is why the replica's events cannot be used — most
|
||||
// likely an engine hashing another way than llmgw configured. Its index is
|
||||
// empty while it stands.
|
||||
string problem = 6;
|
||||
}
|
||||
|
||||
// ControlMessage is anything the control plane tells a node.
|
||||
message ControlMessage {
|
||||
oneof msg {
|
||||
@@ -650,6 +677,9 @@ message EngineProfile {
|
||||
// adapters are LoRA adapters loaded beside the model, in order: the index
|
||||
// is the id a request applies one by.
|
||||
repeated LoraAdapter adapters = 12;
|
||||
// kv_events has a vLLM engine publish its KV cache events to the agent, for
|
||||
// the router's precise index.
|
||||
bool kv_events = 13;
|
||||
}
|
||||
|
||||
// RemoveWeights deletes cached model files a node no longer needs.
|
||||
|
||||
@@ -5253,6 +5253,10 @@
|
||||
"type": "integer",
|
||||
"format": "int64",
|
||||
"description": "StartupTimeout bounds model loading. A 70B model across two nodes takes tens of minutes, so this is per profile rather than global.\n\nA duration, in nanoseconds."
|
||||
},
|
||||
"kv_events": {
|
||||
"type": "boolean",
|
||||
"description": "KVEvents has vLLM publish what its prefix cache holds, and the router route this model's requests by it: to the replica holding the most of a prompt's leading blocks (plan §7.5 tier 2). vLLM only.\n\nOff by default, because it costs something on every request: the router asks a replica to tokenize the prompt before placing it, a round trip and the chat template run twice. It pays where prompts share long prefixes that the approximate tree guesses wrong — many replicas, conversations that move between them, evictions. The engine is started with sha256_cbor block hashing and a fixed PYTHONHASHSEED, which the router's hashing assumes; arguments overriding either leave the index empty, and the agent says so in its log."
|
||||
}
|
||||
}
|
||||
},
|
||||
|
||||
@@ -1028,6 +1028,65 @@ The adapter file must be a GGUF LoRA made for this model — llama.cpp's
|
||||
`convert_lora_to_gguf.py` makes one from a PEFT adapter — and not a model with
|
||||
the adapter already merged in, which is what most "LoRA GGUF" repositories hold.
|
||||
|
||||
### Routing by what each replica's cache actually holds
|
||||
|
||||
With several replicas of a vLLM model, a request is fastest on the replica
|
||||
that already holds its prompt's prefix in KV cache. The router guesses which
|
||||
that is from what it has sent where (`router.prefix`). It can also be told:
|
||||
vLLM publishes an event every time it stores or evicts a cache block, and with
|
||||
`kv_events` on a profile llmgw routes that model by those events.
|
||||
|
||||
```yaml
|
||||
- name: llama-3.1-8b-instruct
|
||||
source: meta-llama/Llama-3.1-8B-Instruct
|
||||
routing_profile: chat # chat or long-context; batch ignores prefixes
|
||||
profiles:
|
||||
- name: h100
|
||||
gpu_class: h100-80g
|
||||
kv_events: true
|
||||
```
|
||||
|
||||
How it works, so you know what you are paying for:
|
||||
|
||||
- The node agent starts vLLM with `--kv-events-config` publishing to a port
|
||||
on 127.0.0.1 that the agent listens on, `--prefix-caching-hash-algo
|
||||
sha256_cbor`, and `PYTHONHASHSEED=1729`. The events never leave the node:
|
||||
the agent keeps the set of blocks each engine holds and sends the changes,
|
||||
four times a second, up the stream it already has to the control plane,
|
||||
which hands them to the router. A new stream — a control plane restart, a
|
||||
failover — starts with a full copy. Nothing new is opened in the firewall.
|
||||
- For each request to such a model, the router asks one of its replicas to
|
||||
tokenize the prompt (vLLM's `/tokenize`, with the request's messages and
|
||||
tools), hashes the token ids into blocks the way vLLM does, and prefers the
|
||||
replica holding the most leading blocks. Tokenizing at the engine is what
|
||||
makes the hashes agree: the chat template is part of the tokens, and only
|
||||
the engine runs it exactly. It costs a round trip before the request is
|
||||
placed and the template run twice; `router.kv_index.tokenize_timeout`
|
||||
(250ms) bounds it.
|
||||
- When the index has nothing to say — no events have arrived yet, the
|
||||
request has an image or a `cache_salt` or names a LoRA adapter, the
|
||||
tokenize call failed — the router routes as it would without it: by the
|
||||
approximate tree if `router.prefix.enabled`, else by load. The imbalance
|
||||
guard of `router.prefix` applies to both.
|
||||
|
||||
Turning `kv_events` on changes how the engine is started, so replicas already
|
||||
running keep running without it (and show as stale) until they are replaced.
|
||||
Leave `--prefix-caching-hash-algo` out of the profile's `args`: an engine
|
||||
hashing another way publishes blocks the router can never match. The agent
|
||||
checks the first block of prompts against its own hashing, and if they never
|
||||
agree it says so in its log and on `llmgw_kv_index_problem`, and the replica
|
||||
is routed as though it published nothing.
|
||||
|
||||
Watch `llmgw_kv_index_decisions_total` beside `llmgw_prefix_decisions_total`:
|
||||
`kv_hit` is the index deciding, `kv_miss_min_load` is it knowing no replica
|
||||
holds enough (`router.kv_index.min_blocks`, 4 blocks), and the `fallback_*`
|
||||
values say why it could not answer. `llmgw_kv_index_tokenize_seconds` is what
|
||||
the tokenize call costs, and `llmgw_kv_index_blocks` what each replica holds.
|
||||
|
||||
This has been built and tested against vLLM's event format and its hashing as
|
||||
its source defines them, with libzmq and msgspec publishing — not yet against
|
||||
a running vLLM. Treat it as experimental until it has been.
|
||||
|
||||
### Models served outside the cluster
|
||||
|
||||
A model can be an OpenAI-compatible API elsewhere — a hosted provider, or
|
||||
|
||||
@@ -93,6 +93,9 @@ rather than repeated.
|
||||
| `output` | `string` | [JudgePolicy](#judgepolicy) |
|
||||
| `categories` | `[]string` | [JudgePolicy](#judgepolicy) |
|
||||
| `latency_budget` | `time.Duration` | [JudgePolicy](#judgepolicy) |
|
||||
| `tokenize_timeout` | `time.Duration` | [KVIndex](#kvindex) |
|
||||
| `min_blocks` | `int` | [KVIndex](#kvindex) |
|
||||
| `max_blocks` | `int` | [KVIndex](#kvindex) |
|
||||
| `contact` | `string` | [Landing](#landing) |
|
||||
| `notice` | `string` | [Landing](#landing) |
|
||||
| `docs_sign_in` | `bool` | [Landing](#landing) |
|
||||
@@ -197,6 +200,7 @@ rather than repeated.
|
||||
| `weight` | `float64` | [Profile](#profile) |
|
||||
| `priority` | `int` | [Profile](#profile) |
|
||||
| `startup_timeout` | `time.Duration` | [Profile](#profile) |
|
||||
| `kv_events` | `bool` | [Profile](#profile) |
|
||||
| `activation_wait` | `time.Duration` | [Router](#router) |
|
||||
| `scrape_interval` | `time.Duration` | [Router](#router) |
|
||||
| `scrape_timeout` | `time.Duration` | [Router](#router) |
|
||||
@@ -206,6 +210,7 @@ rather than repeated.
|
||||
| `failure_threshold` | `int` | [Router](#router) |
|
||||
| `default_profile` | `string` | [Router](#router) |
|
||||
| `prefix` | `Prefix` | [Router](#router) |
|
||||
| `kv_index` | `KVIndex` | [Router](#router) |
|
||||
| `session_header` | `string` | [Router](#router) |
|
||||
| `inject_usage` | `bool` | [Router](#router) |
|
||||
| `tenant_label` | `string` | [Router](#router) |
|
||||
@@ -889,6 +894,33 @@ codes (S1, S10) or the names the model reports.
|
||||
LatencyBudget is how long the judge is waited for before the policy's
|
||||
fail applies. Zero is 250ms.
|
||||
|
||||
## KVIndex
|
||||
|
||||
### `tokenize_timeout`
|
||||
|
||||
`time.Duration`
|
||||
|
||||
TokenizeTimeout bounds the call that turns a request into token ids,
|
||||
made to one of the model's replicas before the request is placed. Past
|
||||
it the request is routed as tier 1 would. Zero takes 250ms.
|
||||
|
||||
### `min_blocks`
|
||||
|
||||
`int`
|
||||
|
||||
MinBlocks is the fewest leading blocks a replica must hold to be
|
||||
preferred for them; fewer is coincidence — a chat template's opening
|
||||
tokens — rather than a prefix worth following. Zero takes 4, which is
|
||||
64 tokens at vLLM's usual block size of 16.
|
||||
|
||||
### `max_blocks`
|
||||
|
||||
`int`
|
||||
|
||||
MaxBlocks bounds how many of a prompt's blocks are hashed and looked
|
||||
up. Beyond it a longer match still helps, but the difference no longer
|
||||
decides anything. Zero takes 4096.
|
||||
|
||||
## Landing
|
||||
|
||||
### `contact`
|
||||
@@ -1658,6 +1690,23 @@ Priority orders profiles during placement; highest that fits wins.
|
||||
StartupTimeout bounds model loading. A 70B model across two nodes takes
|
||||
tens of minutes, so this is per profile rather than global.
|
||||
|
||||
### `kv_events`
|
||||
|
||||
`bool`
|
||||
|
||||
KVEvents has vLLM publish what its prefix cache holds, and the router
|
||||
route this model's requests by it: to the replica holding the most of
|
||||
a prompt's leading blocks (plan §7.5 tier 2). vLLM only.
|
||||
|
||||
Off by default, because it costs something on every request: the
|
||||
router asks a replica to tokenize the prompt before placing it, a
|
||||
round trip and the chat template run twice. It pays where prompts
|
||||
share long prefixes that the approximate tree guesses wrong — many
|
||||
replicas, conversations that move between them, evictions. The engine
|
||||
is started with sha256_cbor block hashing and a fixed PYTHONHASHSEED,
|
||||
which the router's hashing assumes; arguments overriding either leave
|
||||
the index empty, and the agent says so in its log.
|
||||
|
||||
## Router
|
||||
|
||||
### `activation_wait`
|
||||
@@ -1723,6 +1772,16 @@ the win for multi-turn traffic.
|
||||
|
||||
Takes a [Prefix](#prefix) block.
|
||||
|
||||
### `kv_index`
|
||||
|
||||
`KVIndex`
|
||||
|
||||
KVIndex tunes the precise KV index (plan §7.5 tier 2). Which models
|
||||
use it is the catalog's to say, by kv_events on a vLLM profile; this
|
||||
is only how.
|
||||
|
||||
Takes a [KVIndex](#kvindex) block.
|
||||
|
||||
### `session_header`
|
||||
|
||||
`string`
|
||||
|
||||
@@ -12,6 +12,7 @@ require (
|
||||
github.com/charmbracelet/x/ansi v0.11.6
|
||||
github.com/coreos/go-oidc/v3 v3.21.0
|
||||
github.com/go-jose/go-jose/v4 v4.1.4
|
||||
github.com/go-zeromq/zmq4 v0.17.0
|
||||
github.com/jackc/pgx/v5 v5.11.0
|
||||
github.com/knadh/koanf/parsers/yaml v1.1.1
|
||||
github.com/knadh/koanf/providers/env/v2 v2.0.1
|
||||
@@ -22,6 +23,7 @@ require (
|
||||
github.com/prometheus/client_model v0.6.3
|
||||
github.com/prometheus/common v0.71.0
|
||||
github.com/spf13/cobra v1.10.2
|
||||
github.com/vmihailenco/msgpack/v5 v5.4.1
|
||||
github.com/yuin/goldmark v1.8.6
|
||||
go.opentelemetry.io/otel v1.46.0
|
||||
go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.46.0
|
||||
@@ -55,6 +57,7 @@ require (
|
||||
github.com/go-logr/logr v1.4.4 // indirect
|
||||
github.com/go-logr/stdr v1.2.2 // indirect
|
||||
github.com/go-viper/mapstructure/v2 v2.5.0 // indirect
|
||||
github.com/go-zeromq/goczmq/v4 v4.2.2 // indirect
|
||||
github.com/google/uuid v1.6.0 // indirect
|
||||
github.com/grpc-ecosystem/grpc-gateway/v2 v2.30.0 // indirect
|
||||
github.com/inconshreveable/mousetrap v1.1.0 // indirect
|
||||
@@ -77,6 +80,7 @@ require (
|
||||
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
|
||||
github.com/rivo/uniseg v0.4.7 // indirect
|
||||
github.com/spf13/pflag v1.0.9 // indirect
|
||||
github.com/vmihailenco/tagparser/v2 v2.0.0 // indirect
|
||||
github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect
|
||||
go.opentelemetry.io/auto/sdk v1.2.1 // indirect
|
||||
go.opentelemetry.io/otel/exporters/otlp/otlptrace v1.46.0 // indirect
|
||||
|
||||
@@ -55,6 +55,10 @@ github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag=
|
||||
github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE=
|
||||
github.com/go-viper/mapstructure/v2 v2.5.0 h1:vM5IJoUAy3d7zRSVtIwQgBj7BiWtMPfmPEgAXnvj1Ro=
|
||||
github.com/go-viper/mapstructure/v2 v2.5.0/go.mod h1:oJDH3BJKyqBA2TXFhDsKDGDTlndYOZ6rGS0BRZIxGhM=
|
||||
github.com/go-zeromq/goczmq/v4 v4.2.2 h1:HAJN+i+3NW55ijMJJhk7oWxHKXgAuSBkoFfvr8bYj4U=
|
||||
github.com/go-zeromq/goczmq/v4 v4.2.2/go.mod h1:Sm/lxrfxP/Oxqs0tnHD6WAhwkWrx+S+1MRrKzcxoaYE=
|
||||
github.com/go-zeromq/zmq4 v0.17.0 h1:r12/XdqPeRbuaF4C3QZJeWCt7a5vpJbslDH1rTXF+Kc=
|
||||
github.com/go-zeromq/zmq4 v0.17.0/go.mod h1:EQxjJD92qKnrsVMzAnx62giD6uJIPi1dMGZ781iCDtY=
|
||||
github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek=
|
||||
github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps=
|
||||
github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
|
||||
@@ -136,6 +140,10 @@ github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UV
|
||||
github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
|
||||
github.com/stretchr/testify v1.12.1 h1:EuwCh5fleGS7H32xRwO3wRGT7DxrDhLAT6FF8MpWDWE=
|
||||
github.com/stretchr/testify v1.12.1/go.mod h1:MDEgiDPPsNp5cuIrHPPCyornHKgEVbtFUmoNlxoYthg=
|
||||
github.com/vmihailenco/msgpack/v5 v5.4.1 h1:cQriyiUvjTwOHg8QZaPihLWeRAAVoCpE00IUPn0Bjt8=
|
||||
github.com/vmihailenco/msgpack/v5 v5.4.1/go.mod h1:GaZTsDaehaPpQVyxrf5mtQlH+pc21PIudVV/E3rRQok=
|
||||
github.com/vmihailenco/tagparser/v2 v2.0.0 h1:y09buUbR+b5aycVFQs/g70pqKVZNBmxwAhO7/IwNM9g=
|
||||
github.com/vmihailenco/tagparser/v2 v2.0.0/go.mod h1:Wri+At7QHww0WTrCBeu4J6bNtoV6mEfg5OIWRZA9qds=
|
||||
github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavMF/ppJZNG9ZpyihvCd0w101no=
|
||||
github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM=
|
||||
github.com/yuin/goldmark v1.8.6 h1:d0VcaP1sx9GkFVkoW+KtggpGi2KZ965i14b0+bDQST4=
|
||||
|
||||
@@ -35,6 +35,7 @@ import (
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/weights"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/agentsvc"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/kvevents"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/obs"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/tlsconf"
|
||||
)
|
||||
@@ -120,6 +121,12 @@ type Agent struct {
|
||||
// — eight shards starting together produce one download — and all eight
|
||||
// are waiting, so all eight should say so.
|
||||
fetchWatchers map[string][]func(weights.Progress)
|
||||
|
||||
// kv mirrors the KV caches of the engines that publish events, for the
|
||||
// router's precise index; kvSubs are their listeners by shard, guarded
|
||||
// by mu.
|
||||
kv *kvevents.Mirror
|
||||
kvSubs map[string]*kvevents.Subscriber
|
||||
}
|
||||
|
||||
// connectHTTPClient is the HTTP client connect uses; an interface so tests can
|
||||
@@ -162,6 +169,8 @@ func New(cfg Config, httpClient connectHTTPClient, log *slog.Logger) (*Agent, er
|
||||
sensors: discovery.NewSensorSampler(),
|
||||
wear: discovery.NewHealthSampler(),
|
||||
serverList: loadControlPlanes(cfg.Servers, cfg.TLS.Enabled()),
|
||||
kv: kvevents.NewMirror(),
|
||||
kvSubs: map[string]*kvevents.Subscriber{},
|
||||
}
|
||||
a.prov = &provision.Installer{Root: cfg.EnginesDir, Log: log, OnStatus: a.provisionStatus}
|
||||
return a, nil
|
||||
@@ -280,6 +289,12 @@ func (a *Agent) recoverOrphans() {
|
||||
a.mu.Lock()
|
||||
a.shardReplica[o.ShardID] = o.ReplicaID
|
||||
a.mu.Unlock()
|
||||
// The engine is still publishing to the port the last agent
|
||||
// listened on; listen there again.
|
||||
if o.Spec.KVEventsPort > 0 {
|
||||
spec := o.Spec
|
||||
a.listenKV(&spec, o.Spec.KVEventsPort)
|
||||
}
|
||||
a.log.Info("adopted a running engine",
|
||||
"shard", o.ShardID, "replica", o.ReplicaID, "pid", o.PID)
|
||||
go a.settleAdopted(sh, o)
|
||||
@@ -533,6 +548,7 @@ func (a *Agent) session(ctx context.Context, server string) error {
|
||||
}()
|
||||
defer wg.Wait()
|
||||
go a.watchExits(ctx, out)
|
||||
go a.forwardKV(ctx, out)
|
||||
|
||||
for {
|
||||
msg, err := stream.Receive()
|
||||
@@ -1054,8 +1070,10 @@ func (a *Agent) startShard(ctx context.Context, m *llmgwv1.StartShard, out chan<
|
||||
spec.Profile.Adapters[i].Path = path
|
||||
}
|
||||
|
||||
a.listenKV(&spec, 0)
|
||||
cmds, err := eng.Commands(spec)
|
||||
if err != nil {
|
||||
a.stopKV(shardID, replicaID)
|
||||
report(supervisor.Failed, err.Error(), nil)
|
||||
return
|
||||
}
|
||||
@@ -1068,6 +1086,7 @@ func (a *Agent) startShard(ctx context.Context, m *llmgwv1.StartShard, out chan<
|
||||
report(sh.Status().State, "", nil)
|
||||
return
|
||||
}
|
||||
a.stopKV(shardID, replicaID)
|
||||
report(supervisor.Failed, err.Error(), sh.Logs(10))
|
||||
return
|
||||
}
|
||||
@@ -1164,6 +1183,7 @@ func (a *Agent) stopShard(ctx context.Context, shardID string, out chan<- *llmgw
|
||||
if err != nil {
|
||||
a.log.Warn("stopping a shard failed", "shard", shardID, "err", err)
|
||||
}
|
||||
a.stopKV(shardID, replicaID)
|
||||
obs.End(span, err)
|
||||
a.sup.Remove(shardID)
|
||||
|
||||
|
||||
@@ -44,6 +44,9 @@ type Profile struct {
|
||||
// Adapters are LoRA adapters to load beside the model, in order: an
|
||||
// adapter's index is the id requests apply it by.
|
||||
Adapters []Adapter
|
||||
// KVEvents has the engine publish its KV cache events to the agent,
|
||||
// for the router's precise index. vLLM only.
|
||||
KVEvents bool `json:"kv_events,omitempty"`
|
||||
}
|
||||
|
||||
// ShardSpec is one shard of one replica on this node.
|
||||
@@ -90,6 +93,11 @@ type ShardSpec struct {
|
||||
// and accounting. Given to the engine in its environment rather than on
|
||||
// its command line, which every user on the node can read.
|
||||
APIKey string `json:"api_key,omitempty"`
|
||||
// KVEventsPort is the loopback port the agent receives this shard's KV
|
||||
// cache events on, chosen by the agent; zero when it publishes none.
|
||||
// Kept with the spec so an agent that restarts listens where the engine
|
||||
// it adopts is still sending.
|
||||
KVEventsPort int `json:"kv_events_port,omitempty"`
|
||||
// Threads is the CPU cores placement reserved for a shard with no GPU,
|
||||
// and so what its engine computes on. Zero leaves it to the engine's own
|
||||
// configuration.
|
||||
|
||||
@@ -0,0 +1,47 @@
|
||||
package engines
|
||||
|
||||
import (
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestAVLLMShardWithKVEventsPublishesToTheAgentOnLoopback(t *testing.T) {
|
||||
s := singleNodeSpec()
|
||||
s.Profile.KVEvents = true
|
||||
s.KVEventsPort = 5999
|
||||
s.Profile.Args = []string{"--max-model-len", "8192"}
|
||||
cmd := buildOne(t, &VLLM{}, s)
|
||||
|
||||
i := slices.Index(cmd.Args, "--kv-events-config")
|
||||
if i < 0 || !strings.Contains(cmd.Args[i+1], `"endpoint":"tcp://127.0.0.1:5999"`) {
|
||||
t.Fatalf("args %v should name the agent's loopback port", cmd.Args)
|
||||
}
|
||||
if j := slices.Index(cmd.Args, "--prefix-caching-hash-algo"); j < 0 || cmd.Args[j+1] != "sha256_cbor" {
|
||||
t.Fatalf("args %v should pin the hashing the router computes", cmd.Args)
|
||||
}
|
||||
if slices.Index(cmd.Args, "--max-model-len") < i {
|
||||
t.Fatal("profile arguments come last, so an operator can still override")
|
||||
}
|
||||
if !slices.Contains(cmd.Env, "PYTHONHASHSEED=1729") {
|
||||
t.Fatalf("env %v should fix the hash seed", cmd.Env)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNoKVEventsWithoutAPortOrOffTheLeader(t *testing.T) {
|
||||
s := singleNodeSpec()
|
||||
s.Profile.KVEvents = true // but the agent could not listen
|
||||
if cmd := buildOne(t, &VLLM{}, s); slices.Contains(cmd.Args, "--kv-events-config") {
|
||||
t.Fatal("an engine told to publish where nobody listens only fills a queue")
|
||||
}
|
||||
w := multiNodeSpec(1)
|
||||
w.Profile.KVEvents, w.KVEventsPort = true, 5999
|
||||
if cmd := buildEngine(t, &VLLM{}, w); slices.Contains(cmd.Args, "--kv-events-config") {
|
||||
t.Fatal("a headless worker has no scheduler to publish from")
|
||||
}
|
||||
off := singleNodeSpec()
|
||||
off.KVEventsPort = 5999
|
||||
if cmd := buildOne(t, &VLLM{}, off); slices.Contains(cmd.Args, "--kv-events-config") {
|
||||
t.Fatal("kv_events is opt-in")
|
||||
}
|
||||
}
|
||||
@@ -2,8 +2,11 @@ package engines
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"maps"
|
||||
"strconv"
|
||||
"sync/atomic"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/kvevents"
|
||||
)
|
||||
|
||||
// VLLMConfig locates the vLLM installation and picks how it is launched.
|
||||
@@ -147,6 +150,14 @@ func (v *VLLM) serverCommand(s ShardSpec) (Command, error) {
|
||||
}
|
||||
}
|
||||
|
||||
// The KV cache events, for the router's precise index: published to the
|
||||
// agent on loopback, hashed the way the router hashes. Only the leader
|
||||
// has a scheduler to publish from.
|
||||
if p.KVEvents && s.KVEventsPort > 0 && s.Leader() {
|
||||
args = append(args, kvevents.VLLMArgs(fmt.Sprintf("tcp://127.0.0.1:%d", s.KVEventsPort))...)
|
||||
maps.Copy(env, kvevents.VLLMEnv())
|
||||
}
|
||||
|
||||
// Profile arguments last, so an operator can override anything above.
|
||||
args = append(args, p.Args...)
|
||||
|
||||
|
||||
@@ -0,0 +1,86 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/agentsvc"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/kvevents"
|
||||
)
|
||||
|
||||
// kvFlushInterval is how often KV cache changes go up the stream. Often
|
||||
// enough that a prefix cached by one request is known before the next few
|
||||
// arrive; seldom enough that a busy engine's churn is sent as its net effect.
|
||||
const kvFlushInterval = 250 * time.Millisecond
|
||||
|
||||
// listenKV starts receiving a shard's KV cache events, when its profile asks
|
||||
// for them and this rank publishes any, and records the port in the spec for
|
||||
// the engine's arguments. port is the one to reuse, or 0 for any.
|
||||
//
|
||||
// A failure is logged and the engine started without events: the router then
|
||||
// routes the replica by tier 1, which is what it did before, rather than the
|
||||
// replica not starting at all for want of an optimisation.
|
||||
func (a *Agent) listenKV(spec *engines.ShardSpec, port int) {
|
||||
if !spec.Profile.KVEvents || spec.Profile.Engine != "vllm" || !spec.Leader() {
|
||||
return
|
||||
}
|
||||
a.mu.Lock()
|
||||
if sub, ok := a.kvSubs[spec.ShardID]; ok {
|
||||
a.mu.Unlock()
|
||||
// A repeated start command: the shard already has its listener.
|
||||
spec.KVEventsPort = sub.Port()
|
||||
return
|
||||
}
|
||||
a.mu.Unlock()
|
||||
|
||||
sub, err := kvevents.Listen(spec.ReplicaID, port, a.kv, a.log)
|
||||
if err != nil {
|
||||
a.log.Warn("cannot receive kv cache events; the router will route this replica without them",
|
||||
"shard", spec.ShardID, "replica", spec.ReplicaID, "err", err)
|
||||
return
|
||||
}
|
||||
a.mu.Lock()
|
||||
a.kvSubs[spec.ShardID] = sub
|
||||
a.mu.Unlock()
|
||||
spec.KVEventsPort = sub.Port()
|
||||
}
|
||||
|
||||
// stopKV stops a shard's listener and tells the control plane its replica
|
||||
// holds nothing any more.
|
||||
func (a *Agent) stopKV(shardID, replicaID string) {
|
||||
a.mu.Lock()
|
||||
sub, ok := a.kvSubs[shardID]
|
||||
delete(a.kvSubs, shardID)
|
||||
a.mu.Unlock()
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
sub.Close()
|
||||
a.kv.Remove(replicaID)
|
||||
}
|
||||
|
||||
// forwardKV sends KV cache changes up one session's stream until it ends,
|
||||
// beginning with a full copy of every replica: the control plane at the other
|
||||
// end may have just started, or just taken over, and know nothing.
|
||||
func (a *Agent) forwardKV(ctx context.Context, out chan<- *llmgwv1.AgentMessage) {
|
||||
s := a.kv.Open()
|
||||
defer s.Close()
|
||||
t := time.NewTicker(kvFlushInterval)
|
||||
defer t.Stop()
|
||||
for {
|
||||
for _, u := range s.Drain() {
|
||||
select {
|
||||
case out <- agentsvc.KVEventsMessage(u):
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-t.C:
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,103 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/go-zeromq/zmq4"
|
||||
"github.com/vmihailenco/msgpack/v5"
|
||||
|
||||
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/kvevents"
|
||||
)
|
||||
|
||||
func kvAgent() *Agent {
|
||||
return &Agent{log: slog.New(slog.DiscardHandler), kv: kvevents.NewMirror(),
|
||||
kvSubs: map[string]*kvevents.Subscriber{}}
|
||||
}
|
||||
|
||||
// TestAnEnginesEventsReachTheControlPlaneStream runs the agent's half end to
|
||||
// end: a listener for the shard, a publisher standing in for vLLM connecting
|
||||
// to the port the engine would be given, and the stream to the control plane
|
||||
// carrying the result.
|
||||
func TestAnEnginesEventsReachTheControlPlaneStream(t *testing.T) {
|
||||
a := kvAgent()
|
||||
spec := engines.ShardSpec{ReplicaID: "r1", ShardID: "r1-0", NNodes: 1,
|
||||
Profile: engines.Profile{Engine: "vllm", KVEvents: true}}
|
||||
a.listenKV(&spec, 0)
|
||||
if spec.KVEventsPort == 0 {
|
||||
t.Fatal("the shard should have been given a port to publish to")
|
||||
}
|
||||
defer a.stopKV(spec.ShardID, spec.ReplicaID)
|
||||
|
||||
pub := zmq4.NewPub(context.Background())
|
||||
defer func() { _ = pub.Close() }()
|
||||
if err := pub.Dial(fmt.Sprintf("tcp://127.0.0.1:%d", spec.KVEventsPort)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
tokens := make([]int, 32)
|
||||
for i := range tokens {
|
||||
tokens[i] = 500 + i
|
||||
}
|
||||
hs := kvevents.BlockHashes(tokens, 16, 0)
|
||||
payload, _ := msgpack.Marshal([]any{1.0, []any{map[string]any{
|
||||
"type": "BlockStored", "block_hashes": []any{hs[0], hs[1]}, "parent_block_hash": nil,
|
||||
"token_ids": tokens, "block_size": 16, "lora_id": nil, "medium": "GPU",
|
||||
}}})
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
out := make(chan *llmgwv1.AgentMessage, 16)
|
||||
go a.forwardKV(ctx, out)
|
||||
|
||||
deadline := time.After(5 * time.Second)
|
||||
var seq uint64
|
||||
for {
|
||||
// Sent until it arrives: a subscription reaches the publisher after
|
||||
// the connection, and what is published before it is dropped.
|
||||
var s [8]byte
|
||||
binary.BigEndian.PutUint64(s[:], seq)
|
||||
seq++
|
||||
_ = pub.Send(zmq4.NewMsgFrom([]byte("kv-events"), s[:], payload))
|
||||
select {
|
||||
case m := <-out:
|
||||
kv := m.GetKvEvents()
|
||||
if kv == nil || kv.GetReplicaId() != "r1" || kv.GetBlockSize() != 16 || len(kv.GetStored()) != 2 {
|
||||
t.Fatalf("message %v", m)
|
||||
}
|
||||
return
|
||||
case <-deadline:
|
||||
t.Fatal("nothing reached the stream")
|
||||
case <-time.After(50 * time.Millisecond):
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAShardWithoutKVEventsGetsNoListener(t *testing.T) {
|
||||
a := kvAgent()
|
||||
spec := engines.ShardSpec{ReplicaID: "r1", ShardID: "r1-0", NNodes: 1,
|
||||
Profile: engines.Profile{Engine: "vllm"}}
|
||||
a.listenKV(&spec, 0)
|
||||
if spec.KVEventsPort != 0 || len(a.kvSubs) != 0 {
|
||||
t.Fatal("kv_events is opt-in")
|
||||
}
|
||||
}
|
||||
|
||||
func TestARepeatedStartReusesTheShardsListener(t *testing.T) {
|
||||
a := kvAgent()
|
||||
spec := engines.ShardSpec{ReplicaID: "r1", ShardID: "r1-0", NNodes: 1,
|
||||
Profile: engines.Profile{Engine: "vllm", KVEvents: true}}
|
||||
a.listenKV(&spec, 0)
|
||||
defer a.stopKV(spec.ShardID, spec.ReplicaID)
|
||||
again := spec
|
||||
again.KVEventsPort = 0
|
||||
a.listenKV(&again, 0)
|
||||
if again.KVEventsPort != spec.KVEventsPort || len(a.kvSubs) != 1 {
|
||||
t.Fatal("a start command repeated must not open a second listener")
|
||||
}
|
||||
}
|
||||
@@ -45,6 +45,7 @@ import (
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/guard"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/hf"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/kvevents"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/mirror"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/notify"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/obs"
|
||||
@@ -684,6 +685,11 @@ func startControlPlane(ctx context.Context, cfg config.Server, cat *catalog.Live
|
||||
svc.SetEnginesFor(func(n inventory.Node) map[string]string {
|
||||
return cfg.EnginesFor(n.ID, slices.Sorted(maps.Keys(n.GPUsByClass())), n.Labels)
|
||||
})
|
||||
// The KV cache events agents forward go to the router's precise index,
|
||||
// when there is a router in this process to have one.
|
||||
if k, ok := sink.(interface{ ApplyKV(kvevents.Update) }); ok {
|
||||
svc.SetKVSink(k.ApplyKV)
|
||||
}
|
||||
log.Info("nodes install these engines", "engines", cfg.WantedEngines(), "overrides", len(cfg.EngineOverrides))
|
||||
if ch := ctl.EngineChoices(); ch != nil {
|
||||
svc.SetEngineChoices(ch.For)
|
||||
|
||||
@@ -446,6 +446,10 @@ type Router struct {
|
||||
// router uses session affinity alone, which is tier 0 and already most of
|
||||
// the win for multi-turn traffic.
|
||||
Prefix Prefix `koanf:"prefix"`
|
||||
// KVIndex tunes the precise KV index (plan §7.5 tier 2). Which models
|
||||
// use it is the catalog's to say, by kv_events on a vLLM profile; this
|
||||
// is only how.
|
||||
KVIndex KVIndex `koanf:"kv_index"`
|
||||
|
||||
// SessionHeader carries the session id for affinity (plan §7.5 tier 0).
|
||||
SessionHeader string `koanf:"session_header"`
|
||||
@@ -532,6 +536,23 @@ type Prefix struct {
|
||||
BalanceRel float64 `koanf:"balance_rel"`
|
||||
}
|
||||
|
||||
// KVIndex tunes routing by the blocks each vLLM replica reports holding.
|
||||
type KVIndex struct {
|
||||
// TokenizeTimeout bounds the call that turns a request into token ids,
|
||||
// made to one of the model's replicas before the request is placed. Past
|
||||
// it the request is routed as tier 1 would. Zero takes 250ms.
|
||||
TokenizeTimeout time.Duration `koanf:"tokenize_timeout"`
|
||||
// MinBlocks is the fewest leading blocks a replica must hold to be
|
||||
// preferred for them; fewer is coincidence — a chat template's opening
|
||||
// tokens — rather than a prefix worth following. Zero takes 4, which is
|
||||
// 64 tokens at vLLM's usual block size of 16.
|
||||
MinBlocks int `koanf:"min_blocks"`
|
||||
// MaxBlocks bounds how many of a prompt's blocks are hashed and looked
|
||||
// up. Beyond it a longer match still helps, but the difference no longer
|
||||
// decides anything. Zero takes 4096.
|
||||
MaxBlocks int `koanf:"max_blocks"`
|
||||
}
|
||||
|
||||
// Band is one priority class.
|
||||
type Band struct {
|
||||
Name string `koanf:"name"`
|
||||
|
||||
@@ -198,6 +198,7 @@ func StartShardFromCommand(cmd reconciler.ShardCommand) *llmgwv1.StartShard {
|
||||
Args: p.Args, Env: p.Env,
|
||||
ModelPath: p.ModelPath, ServedModelName: p.ServedModelName,
|
||||
Adapters: adaptersToWire(p.Adapters),
|
||||
KvEvents: p.KVEvents,
|
||||
},
|
||||
StartupTimeoutMs: cmd.StartupTimeout.Milliseconds(),
|
||||
Pool: cmd.Pool,
|
||||
@@ -234,6 +235,7 @@ func ShardSpecFromStart(m *llmgwv1.StartShard) (engines.ShardSpec, time.Duration
|
||||
Args: p.GetArgs(), Env: p.GetEnv(),
|
||||
ModelPath: p.GetModelPath(), ServedModelName: p.GetServedModelName(),
|
||||
Adapters: adaptersFromWire(p.GetAdapters()),
|
||||
KVEvents: p.GetKvEvents(),
|
||||
},
|
||||
}
|
||||
return spec, time.Duration(m.GetStartupTimeoutMs()) * time.Millisecond
|
||||
|
||||
@@ -0,0 +1,103 @@
|
||||
package agentsvc
|
||||
|
||||
import (
|
||||
"sync"
|
||||
|
||||
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/kvevents"
|
||||
)
|
||||
|
||||
// SetKVSink has the KV cache events agents forward delivered to fn — the
|
||||
// router's precise index, in a process that serves both. Unset, they are
|
||||
// dropped, which is right for a control plane with no router beside it.
|
||||
func (s *Service) SetKVSink(fn func(kvevents.Update)) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.kv = fn
|
||||
}
|
||||
|
||||
func (s *Service) kvSink() func(kvevents.Update) {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
return s.kv
|
||||
}
|
||||
|
||||
// kvState is what one stream has said about KV caches: which replicas the
|
||||
// node runs, and so may describe, and which it has described.
|
||||
type kvState struct {
|
||||
mu sync.Mutex
|
||||
owned map[string]bool
|
||||
described map[string]bool
|
||||
}
|
||||
|
||||
// own records a replica the node runs a shard of.
|
||||
func (c *conn) own(replicaID string) {
|
||||
if replicaID == "" {
|
||||
return
|
||||
}
|
||||
c.kv.mu.Lock()
|
||||
defer c.kv.mu.Unlock()
|
||||
if c.kv.owned == nil {
|
||||
c.kv.owned = map[string]bool{}
|
||||
}
|
||||
c.kv.owned[replicaID] = true
|
||||
}
|
||||
|
||||
// kvEvents passes one replica's update on.
|
||||
//
|
||||
// Only for a replica the node has said it runs. A node describing another's
|
||||
// cache — misconfigured, or worse — could otherwise draw that replica's
|
||||
// traffic, or starve it, from anywhere in the cluster.
|
||||
func (s *Service) kvEvents(c *conn, m *llmgwv1.KVEvents) {
|
||||
sink := s.kvSink()
|
||||
if sink == nil {
|
||||
return
|
||||
}
|
||||
id := m.GetReplicaId()
|
||||
c.kv.mu.Lock()
|
||||
ok := c.kv.owned[id]
|
||||
if ok {
|
||||
if c.kv.described == nil {
|
||||
c.kv.described = map[string]bool{}
|
||||
}
|
||||
c.kv.described[id] = true
|
||||
}
|
||||
c.kv.mu.Unlock()
|
||||
if !ok {
|
||||
s.log.Debug("ignoring kv cache events for a replica the node does not run",
|
||||
"node", c.nodeID, "replica", id)
|
||||
return
|
||||
}
|
||||
sink(kvevents.Update{
|
||||
Replica: id, BlockSize: int(m.GetBlockSize()), Reset: m.GetClearFirst(),
|
||||
Stored: m.GetStored(), Removed: m.GetRemoved(), Problem: m.GetProblem(),
|
||||
})
|
||||
}
|
||||
|
||||
// kvGone empties what a closed stream described. Events stop arriving with
|
||||
// it, evictions among them, and an index that kept the blocks would steer
|
||||
// requests to caches that may have moved on; the next stream starts with a
|
||||
// full copy anyway.
|
||||
func (s *Service) kvGone(c *conn) {
|
||||
sink := s.kvSink()
|
||||
if sink == nil {
|
||||
return
|
||||
}
|
||||
c.kv.mu.Lock()
|
||||
ids := make([]string, 0, len(c.kv.described))
|
||||
for id := range c.kv.described {
|
||||
ids = append(ids, id)
|
||||
}
|
||||
c.kv.mu.Unlock()
|
||||
for _, id := range ids {
|
||||
sink(kvevents.Update{Replica: id, Reset: true})
|
||||
}
|
||||
}
|
||||
|
||||
// KVEventsMessage wraps an update for the stream, as the agent sends it.
|
||||
func KVEventsMessage(u kvevents.Update) *llmgwv1.AgentMessage {
|
||||
return &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_KvEvents{KvEvents: &llmgwv1.KVEvents{
|
||||
ReplicaId: u.Replica, BlockSize: int32(u.BlockSize), ClearFirst: u.Reset,
|
||||
Stored: u.Stored, Removed: u.Removed, Problem: u.Problem,
|
||||
}}}
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
package agentsvc
|
||||
|
||||
import (
|
||||
"log/slog"
|
||||
"testing"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/kvevents"
|
||||
)
|
||||
|
||||
func TestKVEventsAreTakenOnlyForTheNodesOwnReplicas(t *testing.T) {
|
||||
s := &Service{log: slog.New(slog.DiscardHandler)}
|
||||
var got []kvevents.Update
|
||||
s.SetKVSink(func(u kvevents.Update) { got = append(got, u) })
|
||||
c := &conn{nodeID: "n1"}
|
||||
c.own("r1")
|
||||
|
||||
s.kvEvents(c, KVEventsMessage(kvevents.Update{Replica: "r1", BlockSize: 16, Reset: true, Stored: []uint64{7}}).GetKvEvents())
|
||||
s.kvEvents(c, KVEventsMessage(kvevents.Update{Replica: "r-elsewhere", Stored: []uint64{8}}).GetKvEvents())
|
||||
if len(got) != 1 || got[0].Replica != "r1" || !got[0].Reset || got[0].Stored[0] != 7 || got[0].BlockSize != 16 {
|
||||
t.Fatalf("got %+v; a node may describe its own replicas' caches and no others", got)
|
||||
}
|
||||
|
||||
// The stream closing empties what it described: no more evictions will
|
||||
// arrive by it.
|
||||
s.kvGone(c)
|
||||
if len(got) != 2 || got[1].Replica != "r1" || !got[1].Reset || len(got[1].Stored) != 0 {
|
||||
t.Fatalf("got %+v", got)
|
||||
}
|
||||
}
|
||||
@@ -24,6 +24,7 @@ import (
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/fabrictest"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/kvevents"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/obs"
|
||||
)
|
||||
|
||||
@@ -58,6 +59,8 @@ type Service struct {
|
||||
|
||||
mu sync.RWMutex
|
||||
conns map[string]*conn
|
||||
// kv receives the KV cache events agents forward; see SetKVSink.
|
||||
kv func(kvevents.Update)
|
||||
|
||||
// enginesFor is what a node is told to install, set once at start: the
|
||||
// cluster's engines, with any override that picks the node laid over.
|
||||
@@ -123,6 +126,8 @@ type conn struct {
|
||||
send chan *llmgwv1.ControlMessage
|
||||
closed chan struct{}
|
||||
once sync.Once
|
||||
// kv is what this stream has said about KV caches.
|
||||
kv kvState
|
||||
}
|
||||
|
||||
func (c *conn) close() { c.once.Do(func() { close(c.closed) }) }
|
||||
@@ -242,6 +247,10 @@ func (s *Service) Connect(
|
||||
|
||||
c := s.attach(ctx, nodeID)
|
||||
defer s.detach(nodeID, c)
|
||||
for _, rs := range reg.GetRunningShards() {
|
||||
c.own(rs.GetReplicaId())
|
||||
}
|
||||
defer s.kvGone(c)
|
||||
|
||||
if err := stream.Send(&llmgwv1.ControlMessage{
|
||||
Msg: &llmgwv1.ControlMessage_Ack{Ack: &llmgwv1.Ack{MessageId: "register", Ok: true}},
|
||||
@@ -301,6 +310,13 @@ func (s *Service) Connect(
|
||||
s.log.Warn("agent stream failed", "node", nodeID, "err", err)
|
||||
return nil
|
||||
}
|
||||
if kv := msg.GetKvEvents(); kv != nil {
|
||||
s.kvEvents(c, kv)
|
||||
continue
|
||||
}
|
||||
if sh := msg.GetShard(); sh != nil {
|
||||
c.own(sh.GetReplicaId())
|
||||
}
|
||||
s.handle(nodeID, msg)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -234,6 +234,20 @@ type Profile struct {
|
||||
// StartupTimeout bounds model loading. A 70B model across two nodes takes
|
||||
// tens of minutes, so this is per profile rather than global.
|
||||
StartupTimeout time.Duration `json:"startup_timeout,omitempty" koanf:"startup_timeout"`
|
||||
|
||||
// KVEvents has vLLM publish what its prefix cache holds, and the router
|
||||
// route this model's requests by it: to the replica holding the most of
|
||||
// a prompt's leading blocks (plan §7.5 tier 2). vLLM only.
|
||||
//
|
||||
// Off by default, because it costs something on every request: the
|
||||
// router asks a replica to tokenize the prompt before placing it, a
|
||||
// round trip and the chat template run twice. It pays where prompts
|
||||
// share long prefixes that the approximate tree guesses wrong — many
|
||||
// replicas, conversations that move between them, evictions. The engine
|
||||
// is started with sha256_cbor block hashing and a fixed PYTHONHASHSEED,
|
||||
// which the router's hashing assumes; arguments overriding either leave
|
||||
// the index empty, and the agent says so in its log.
|
||||
KVEvents bool `json:"kv_events,omitempty" koanf:"kv_events"`
|
||||
}
|
||||
|
||||
// CPUOnly reports whether this profile needs no GPUs.
|
||||
@@ -714,6 +728,10 @@ func (m Model) validate() error {
|
||||
if err := engines.UnsafeArgs(p.Engine, p.Args); err != nil {
|
||||
return fmt.Errorf("model %q profile %q: %w", m.Name, p.Name, err)
|
||||
}
|
||||
if p.KVEvents && p.Engine != "vllm" {
|
||||
return fmt.Errorf("model %q profile %q: kv_events is vLLM's; engine %q publishes no cache events",
|
||||
m.Name, p.Name, p.Engine)
|
||||
}
|
||||
|
||||
if p.GPUClass == "" && !cpuCapable(p.Engine) {
|
||||
return fmt.Errorf("model %q profile %q: gpu_class is required for engine %q",
|
||||
@@ -871,11 +889,24 @@ func (m Model) EngineProfile(name string) (engines.Profile, bool) {
|
||||
ModelPath: m.Source,
|
||||
ServedModelName: m.Name,
|
||||
Adapters: m.engineAdapters(),
|
||||
KVEvents: p.KVEvents,
|
||||
}, true
|
||||
}
|
||||
return engines.Profile{}, false
|
||||
}
|
||||
|
||||
// KVIndex reports whether the router should route this model by the KV
|
||||
// cache its replicas report: whether any profile publishes events. A replica
|
||||
// of another profile is simply one the index has not heard from.
|
||||
func (m Model) KVIndex() bool {
|
||||
for _, p := range m.Profiles {
|
||||
if p.KVEvents {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// ProfileByName returns one of a model's profiles.
|
||||
func (m Model) ProfileByName(name string) (Profile, bool) {
|
||||
for _, p := range m.Profiles {
|
||||
|
||||
@@ -0,0 +1,40 @@
|
||||
package catalog
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestKVEventsReachTheEngineAndTheRouter(t *testing.T) {
|
||||
c, err := New([]Model{model("m", Profile{Name: "p", GPUClass: "h100-80g", TP: 1, KVEvents: true})})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
m, _ := c.Resolve("m")
|
||||
if !m.KVIndex() {
|
||||
t.Fatal("a model with a profile publishing events should be routed by them")
|
||||
}
|
||||
if ep, _ := m.EngineProfile("p"); !ep.KVEvents {
|
||||
t.Fatal("the engine has to be told to publish")
|
||||
}
|
||||
plain, _ := New([]Model{model("n")})
|
||||
if n, _ := plain.Resolve("n"); n.KVIndex() {
|
||||
t.Fatal("off by default")
|
||||
}
|
||||
}
|
||||
|
||||
func TestKVEventsAreVLLMsOnly(t *testing.T) {
|
||||
_, err := New([]Model{model("m", Profile{Name: "p", Engine: "llamacpp", KVEvents: true})})
|
||||
if err == nil || !strings.Contains(err.Error(), "kv_events") {
|
||||
t.Fatalf("err %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTurningKVEventsOnMarksRunningReplicasStale(t *testing.T) {
|
||||
p := Profile{Name: "p", GPUClass: "h100-80g", TP: 1}
|
||||
q := p
|
||||
q.KVEvents = true
|
||||
if p.Fingerprint() == q.Fingerprint() {
|
||||
t.Fatal("replicas started without events do not publish until replaced, and should be shown as stale")
|
||||
}
|
||||
}
|
||||
@@ -110,6 +110,11 @@ var Reference = []Block{
|
||||
{Key: "categories", Type: "[]string", Doc: "Categories, when set, are the only ones that count: Llama Guard's\ncodes (S1, S10) or the names the model reports.", Nested: ""},
|
||||
{Key: "latency_budget", Type: "time.Duration", Doc: "LatencyBudget is how long the judge is waited for before the policy's\nfail applies. Zero is 250ms.", Nested: ""},
|
||||
}},
|
||||
{Package: "config", Name: "KVIndex", Title: "KVIndex", Doc: "", Fields: []Field{
|
||||
{Key: "tokenize_timeout", Type: "time.Duration", Doc: "TokenizeTimeout bounds the call that turns a request into token ids,\nmade to one of the model's replicas before the request is placed. Past\nit the request is routed as tier 1 would. Zero takes 250ms.", Nested: ""},
|
||||
{Key: "min_blocks", Type: "int", Doc: "MinBlocks is the fewest leading blocks a replica must hold to be\npreferred for them; fewer is coincidence — a chat template's opening\ntokens — rather than a prefix worth following. Zero takes 4, which is\n64 tokens at vLLM's usual block size of 16.", Nested: ""},
|
||||
{Key: "max_blocks", Type: "int", Doc: "MaxBlocks bounds how many of a prompt's blocks are hashed and looked\nup. Beyond it a longer match still helps, but the difference no longer\ndecides anything. Zero takes 4096.", Nested: ""},
|
||||
}},
|
||||
{Package: "config", Name: "Landing", Title: "Landing", Doc: "", Fields: []Field{
|
||||
{Key: "contact", Type: "string", Doc: "Contact is where someone without a key should go — an address, a URL, a\nchannel name. Rendered as a link when it is one, as text otherwise.", Nested: ""},
|
||||
{Key: "notice", Type: "string", Doc: "Notice is an optional line of your own, for whatever the reader needs to\nknow that this page cannot guess.", Nested: ""},
|
||||
@@ -243,6 +248,7 @@ var Reference = []Block{
|
||||
{Key: "weight", Type: "float64", Doc: "Weight is the router's capacity weight. It should come from ctl bench\nrather than a guess (plan §6.2.1).", Nested: ""},
|
||||
{Key: "priority", Type: "int", Doc: "Priority orders profiles during placement; highest that fits wins.", Nested: ""},
|
||||
{Key: "startup_timeout", Type: "time.Duration", Doc: "StartupTimeout bounds model loading. A 70B model across two nodes takes\ntens of minutes, so this is per profile rather than global.", Nested: ""},
|
||||
{Key: "kv_events", Type: "bool", Doc: "KVEvents has vLLM publish what its prefix cache holds, and the router\nroute this model's requests by it: to the replica holding the most of\na prompt's leading blocks (plan §7.5 tier 2). vLLM only.\n\nOff by default, because it costs something on every request: the\nrouter asks a replica to tokenize the prompt before placing it, a\nround trip and the chat template run twice. It pays where prompts\nshare long prefixes that the approximate tree guesses wrong — many\nreplicas, conversations that move between them, evictions. The engine\nis started with sha256_cbor block hashing and a fixed PYTHONHASHSEED,\nwhich the router's hashing assumes; arguments overriding either leave\nthe index empty, and the agent says so in its log.", Nested: ""},
|
||||
}},
|
||||
{Package: "config", Name: "Router", Title: "Router", Doc: "", Fields: []Field{
|
||||
{Key: "activation_wait", Type: "time.Duration", Doc: "ActivationWait is how long a request waits for an on-demand model to\nload before the client is told to come back. Zero takes the default.\n\nThe one real decision in on-demand loading. Refusing outright makes the\nfeature useless, since the first request for every model would fail;\nwaiting without a bound turns a request into a job, and a client that\nbudgeted thirty seconds ends up holding a connection for ten minutes.", Nested: ""},
|
||||
@@ -254,6 +260,7 @@ var Reference = []Block{
|
||||
{Key: "failure_threshold", Type: "int", Doc: "FailureThreshold is the consecutive failure count that opens a breaker.\n1 turns the breaker into a pure cooldown.", Nested: ""},
|
||||
{Key: "default_profile", Type: "string", Doc: "DefaultProfile names the routing profile used when a model picks none.", Nested: ""},
|
||||
{Key: "prefix", Type: "Prefix", Doc: "Prefix turns on the approximate prefix tree (plan §7.5 tier 1). Off, the\nrouter uses session affinity alone, which is tier 0 and already most of\nthe win for multi-turn traffic.", Nested: "config.Prefix"},
|
||||
{Key: "kv_index", Type: "KVIndex", Doc: "KVIndex tunes the precise KV index (plan §7.5 tier 2). Which models\nuse it is the catalog's to say, by kv_events on a vLLM profile; this\nis only how.", Nested: "config.KVIndex"},
|
||||
{Key: "session_header", Type: "string", Doc: "SessionHeader carries the session id for affinity (plan §7.5 tier 0).", Nested: ""},
|
||||
{Key: "inject_usage", Type: "bool", Doc: "InjectUsage asks engines for usage even when the client did not, so token\naccounting works regardless of client behaviour.", Nested: ""},
|
||||
{Key: "tenant_label", Type: "string", Doc: "TenantLabel is the level of the tree the per-tenant request series are\nlabelled by — llmgw_tenant_requests_total and its siblings: \"org\",\n\"team\" or \"project\", or \"none\" for no such series. A label per project\non a cluster of thousands is more series than Prometheus wants, and a\nteam is what an alert is usually about; \"team\" is the default.", Nested: ""},
|
||||
|
||||
@@ -0,0 +1,224 @@
|
||||
package kvevents
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/vmihailenco/msgpack/v5"
|
||||
)
|
||||
|
||||
// Kind is what an event says happened to the cache.
|
||||
type Kind int
|
||||
|
||||
// The event kinds vLLM publishes.
|
||||
const (
|
||||
Stored Kind = iota + 1
|
||||
Removed
|
||||
Cleared
|
||||
)
|
||||
|
||||
// Event is one cache event, reduced to what the index needs.
|
||||
type Event struct {
|
||||
Kind Kind
|
||||
// Hashes are the blocks stored or removed, in external form.
|
||||
Hashes []uint64
|
||||
// Parent is the hash of the block before the first stored one; HasParent
|
||||
// is false for the first block of a prompt.
|
||||
Parent uint64
|
||||
HasParent bool
|
||||
// Tokens are the stored blocks' token ids, and BlockSize how many make a
|
||||
// block.
|
||||
Tokens []int
|
||||
BlockSize int
|
||||
// Extra says the blocks were hashed with more than their tokens: a LoRA
|
||||
// adapter, an image, a cache salt. Their hashes are indexed like any
|
||||
// other, but say nothing about how plain tokens hash.
|
||||
Extra bool
|
||||
// Elsewhere says the event is about a copy outside this engine's GPU
|
||||
// memory — offloaded to CPU or storage, or another engine's — or about a
|
||||
// second KV cache group. Such events are ignored; see Mirror.Apply.
|
||||
Elsewhere bool
|
||||
}
|
||||
|
||||
// Batch is one published message: the events of one engine step.
|
||||
type Batch struct {
|
||||
TS float64
|
||||
Events []Event
|
||||
}
|
||||
|
||||
// DecodeBatch reads a msgpack-encoded KVEventBatch.
|
||||
//
|
||||
// Two encodings are accepted, because vLLM has used both. A batch is always
|
||||
// an array, [ts, events, data_parallel_rank]. An event was an array too —
|
||||
// ["BlockStored", block_hashes, parent_block_hash, token_ids, block_size,
|
||||
// lora_id, medium, ...] — in the releases that introduced the events, and is
|
||||
// now a map tagged by its "type" key, which lets fields be added without
|
||||
// breaking anyone. A hash is an integer, or 32 bytes with
|
||||
// VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES off.
|
||||
//
|
||||
// An event of a kind not listed here is skipped rather than refused: a newer
|
||||
// engine adding one must not stop the index.
|
||||
func DecodeBatch(payload []byte) (Batch, error) {
|
||||
d := msgpack.NewDecoder(bytes.NewReader(payload))
|
||||
d.UseLooseInterfaceDecoding(true)
|
||||
v, err := d.DecodeInterface()
|
||||
if err != nil {
|
||||
return Batch{}, fmt.Errorf("kv events: not msgpack: %w", err)
|
||||
}
|
||||
arr, ok := v.([]any)
|
||||
if !ok || len(arr) < 2 {
|
||||
return Batch{}, errors.New("kv events: a batch is an array of at least ts and events")
|
||||
}
|
||||
var b Batch
|
||||
b.TS, _ = toFloat(arr[0])
|
||||
evs, ok := arr[1].([]any)
|
||||
if !ok && arr[1] != nil {
|
||||
return Batch{}, errors.New("kv events: the events of a batch are not an array")
|
||||
}
|
||||
for i, raw := range evs {
|
||||
e, ok, err := decodeEvent(raw)
|
||||
if err != nil {
|
||||
return Batch{}, fmt.Errorf("kv events: event %d: %w", i, err)
|
||||
}
|
||||
if ok {
|
||||
b.Events = append(b.Events, e)
|
||||
}
|
||||
}
|
||||
return b, nil
|
||||
}
|
||||
|
||||
// decodeEvent turns either encoding into a map by field name, then reads it.
|
||||
func decodeEvent(raw any) (Event, bool, error) {
|
||||
fields := map[string]any{}
|
||||
var tag string
|
||||
switch v := raw.(type) {
|
||||
case map[string]any:
|
||||
fields = v
|
||||
tag, _ = v["type"].(string)
|
||||
case []any:
|
||||
if len(v) == 0 {
|
||||
return Event{}, false, errors.New("empty event")
|
||||
}
|
||||
tag, _ = v[0].(string)
|
||||
var names []string
|
||||
switch tag {
|
||||
case "BlockStored":
|
||||
names = []string{"block_hashes", "parent_block_hash", "token_ids", "block_size",
|
||||
"lora_id", "medium", "lora_name", "extra_keys", "group_idx"}
|
||||
case "BlockRemoved":
|
||||
names = []string{"block_hashes", "medium", "group_idx"}
|
||||
}
|
||||
for i, name := range names {
|
||||
if i+1 < len(v) {
|
||||
fields[name] = v[i+1]
|
||||
}
|
||||
}
|
||||
default:
|
||||
return Event{}, false, fmt.Errorf("an event is a map or an array, not %T", raw)
|
||||
}
|
||||
|
||||
var e Event
|
||||
switch tag {
|
||||
case "BlockStored":
|
||||
e.Kind = Stored
|
||||
case "BlockRemoved":
|
||||
e.Kind = Removed
|
||||
case "AllBlocksCleared":
|
||||
e.Kind = Cleared
|
||||
return e, true, nil
|
||||
default:
|
||||
return Event{}, false, nil
|
||||
}
|
||||
|
||||
hs, _ := fields["block_hashes"].([]any)
|
||||
for _, h := range hs {
|
||||
x, err := hashOf(h)
|
||||
if err != nil {
|
||||
return Event{}, false, err
|
||||
}
|
||||
e.Hashes = append(e.Hashes, x)
|
||||
}
|
||||
if medium, ok := fields["medium"].(string); ok && !strings.EqualFold(medium, "GPU") {
|
||||
e.Elsewhere = true
|
||||
}
|
||||
if loc, ok := fields["locality"].(string); ok && !strings.EqualFold(loc, "LOCAL") {
|
||||
e.Elsewhere = true
|
||||
}
|
||||
if g, ok := toInt(fields["group_idx"]); ok && g != 0 {
|
||||
e.Elsewhere = true
|
||||
}
|
||||
if e.Kind == Removed {
|
||||
return e, true, nil
|
||||
}
|
||||
|
||||
if p := fields["parent_block_hash"]; p != nil {
|
||||
x, err := hashOf(p)
|
||||
if err != nil {
|
||||
return Event{}, false, err
|
||||
}
|
||||
e.Parent, e.HasParent = x, true
|
||||
}
|
||||
ts, _ := fields["token_ids"].([]any)
|
||||
e.Tokens = make([]int, 0, len(ts))
|
||||
for _, t := range ts {
|
||||
n, ok := toInt(t)
|
||||
if !ok {
|
||||
return Event{}, false, fmt.Errorf("token id %v is not an integer", t)
|
||||
}
|
||||
e.Tokens = append(e.Tokens, n)
|
||||
}
|
||||
e.BlockSize, _ = toInt(fields["block_size"])
|
||||
if fields["lora_id"] != nil || fields["lora_name"] != nil {
|
||||
e.Extra = true
|
||||
}
|
||||
if ks, ok := fields["extra_keys"].([]any); ok {
|
||||
for _, k := range ks {
|
||||
if k != nil {
|
||||
e.Extra = true
|
||||
}
|
||||
}
|
||||
}
|
||||
return e, true, nil
|
||||
}
|
||||
|
||||
func hashOf(v any) (uint64, error) {
|
||||
switch h := v.(type) {
|
||||
case []byte:
|
||||
return External(h), nil
|
||||
case string:
|
||||
// msgpack has no separate bytes type in its oldest form, and a
|
||||
// publisher using that writes bytes as a string.
|
||||
return External([]byte(h)), nil
|
||||
case int64:
|
||||
return uint64(h), nil
|
||||
case uint64:
|
||||
return h, nil
|
||||
}
|
||||
return 0, fmt.Errorf("a block hash is an integer or bytes, not %T", v)
|
||||
}
|
||||
|
||||
func toInt(v any) (int, bool) {
|
||||
switch n := v.(type) {
|
||||
case int64:
|
||||
return int(n), true
|
||||
case uint64:
|
||||
return int(n), true
|
||||
}
|
||||
return 0, false
|
||||
}
|
||||
|
||||
func toFloat(v any) (float64, bool) {
|
||||
switch n := v.(type) {
|
||||
case float64:
|
||||
return n, true
|
||||
case float32:
|
||||
return float64(n), true
|
||||
case int64:
|
||||
return float64(n), true
|
||||
case uint64:
|
||||
return float64(n), true
|
||||
}
|
||||
return 0, false
|
||||
}
|
||||
@@ -0,0 +1,133 @@
|
||||
// Package kvevents reads the KV cache events a vLLM engine publishes, and
|
||||
// computes the block hashes those events are keyed by (plan §7.5 tier 2).
|
||||
//
|
||||
// Both halves live together because they must agree byte for byte. The agent
|
||||
// reads the events and uses the hashing to check that the engine hashes the
|
||||
// way llmgw expects; the router uses the same hashing to ask which replicas
|
||||
// hold the leading blocks of a prompt. A difference between the two is not an
|
||||
// error anything would report — it is a cache index that never matches.
|
||||
package kvevents
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"encoding/binary"
|
||||
)
|
||||
|
||||
// The hashing scheme llmgw starts vLLM with, and assumes.
|
||||
//
|
||||
// sha256_cbor because it is the one of vLLM's algorithms defined outside
|
||||
// Python: canonical CBOR of the tuple (parent hash, token ids, extra keys),
|
||||
// then SHA-256. The default, "sha256", hashes a pickle, which only Python can
|
||||
// reproduce; the builtin hash vLLM once offered changes from process to
|
||||
// process.
|
||||
//
|
||||
// HashSeed is what PYTHONHASHSEED is set to. vLLM derives the hash every
|
||||
// chain starts from — its NONE_HASH — from that variable, and without it
|
||||
// either from a random value (older releases) or a fixed string (newer ones).
|
||||
// Setting it pins the start of the chain in every release. It has to be an
|
||||
// integer for Python to start at all; which one does not matter, only that
|
||||
// the engine and the router agree, and it is a constant so they cannot
|
||||
// disagree.
|
||||
const (
|
||||
HashAlgo = "sha256_cbor"
|
||||
HashSeed = "1729"
|
||||
)
|
||||
|
||||
// BlockHashes returns the external hashes of the full blocks of tokens, as
|
||||
// vLLM would publish them for a request carrying exactly these token ids and
|
||||
// nothing that adds extra keys (no LoRA adapter, no images, no cache salt).
|
||||
//
|
||||
// A partial last block has no hash: vLLM caches only full blocks. At most max
|
||||
// blocks are hashed, and max <= 0 means all of them.
|
||||
func BlockHashes(tokens []int, blockSize, max int) []uint64 {
|
||||
if blockSize <= 0 {
|
||||
return nil
|
||||
}
|
||||
n := len(tokens) / blockSize
|
||||
if max > 0 && n > max {
|
||||
n = max
|
||||
}
|
||||
out := make([]uint64, 0, n)
|
||||
parent := noneHash()
|
||||
buf := make([]byte, 0, 64+blockSize*5)
|
||||
for i := range n {
|
||||
buf = blockInput(buf[:0], parent[:], tokens[i*blockSize:(i+1)*blockSize])
|
||||
parent = sha256.Sum256(buf)
|
||||
out = append(out, External(parent[:]))
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// FirstBlockHash is the hash of a chain's first block, which needs nothing
|
||||
// but its tokens. It is how the agent checks an engine hashes the way llmgw
|
||||
// expects: the first block of every prompt is published with its tokens and
|
||||
// no parent.
|
||||
func FirstBlockHash(tokens []int) uint64 {
|
||||
parent := noneHash()
|
||||
sum := sha256.Sum256(blockInput(nil, parent[:], tokens))
|
||||
return External(sum[:])
|
||||
}
|
||||
|
||||
// External is the form vLLM publishes a hash in: the low 64 bits of the
|
||||
// digest read big-endian, which is what VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES
|
||||
// (on by default, and set by llmgw) makes of the 32 bytes. A publisher with it
|
||||
// off sends the bytes, and they are reduced the same way, so both forms index
|
||||
// alike.
|
||||
func External(digest []byte) uint64 {
|
||||
if len(digest) < 8 {
|
||||
var pad [8]byte
|
||||
copy(pad[8-len(digest):], digest)
|
||||
return binary.BigEndian.Uint64(pad[:])
|
||||
}
|
||||
return binary.BigEndian.Uint64(digest[len(digest)-8:])
|
||||
}
|
||||
|
||||
// noneHash is vLLM's NONE_HASH under PYTHONHASHSEED=HashSeed:
|
||||
// sha256_cbor of the seed, which Python reads as a string.
|
||||
func noneHash() [32]byte {
|
||||
return sha256.Sum256(cborText(nil, HashSeed))
|
||||
}
|
||||
|
||||
// blockInput is canonical CBOR of the tuple vLLM hashes for a block:
|
||||
// (parent_block_hash, tuple(token_ids), extra_keys) with no extra keys, as
|
||||
// cbor2.dumps(..., canonical=True) writes it. A tuple is an array, bytes a
|
||||
// byte string, a token id an unsigned integer, None the simple value null.
|
||||
func blockInput(buf, parent []byte, tokens []int) []byte {
|
||||
buf = cborHead(buf, 4, 3) // array of three
|
||||
buf = cborHead(buf, 2, uint64(len(parent)))
|
||||
buf = append(buf, parent...)
|
||||
buf = cborHead(buf, 4, uint64(len(tokens)))
|
||||
for _, t := range tokens {
|
||||
if t < 0 {
|
||||
// Token ids are never negative; a negative one would be a major
|
||||
// type 1 integer, and is encoded as such rather than wrapped.
|
||||
buf = cborHead(buf, 1, uint64(-1-t))
|
||||
continue
|
||||
}
|
||||
buf = cborHead(buf, 0, uint64(t))
|
||||
}
|
||||
return append(buf, 0xf6) // None
|
||||
}
|
||||
|
||||
func cborText(buf []byte, s string) []byte {
|
||||
buf = cborHead(buf, 3, uint64(len(s)))
|
||||
return append(buf, s...)
|
||||
}
|
||||
|
||||
// cborHead writes a CBOR item head in its shortest form, which is what
|
||||
// canonical encoding requires.
|
||||
func cborHead(buf []byte, major byte, n uint64) []byte {
|
||||
m := major << 5
|
||||
switch {
|
||||
case n < 24:
|
||||
return append(buf, m|byte(n))
|
||||
case n <= 0xff:
|
||||
return append(buf, m|24, byte(n))
|
||||
case n <= 0xffff:
|
||||
return binary.BigEndian.AppendUint16(append(buf, m|25), uint16(n))
|
||||
case n <= 0xffffffff:
|
||||
return binary.BigEndian.AppendUint32(append(buf, m|26), uint32(n))
|
||||
default:
|
||||
return binary.BigEndian.AppendUint64(append(buf, m|27), n)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,49 @@
|
||||
package kvevents
|
||||
|
||||
import (
|
||||
"os"
|
||||
"os/exec"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// TestInteropWithPyZMQ checks the subscriber against libzmq and msgspec, the
|
||||
// libraries vLLM publishes with, by running testdata/pypub.py under uv. Off
|
||||
// unless LLMGW_PYZMQ_INTEROP is set: it needs uv and the network to fetch
|
||||
// pyzmq and msgspec. It is not vLLM — no engine runs — but it is the wire
|
||||
// format from the libraries vLLM uses, with vLLM's struct definitions.
|
||||
func TestInteropWithPyZMQ(t *testing.T) {
|
||||
if os.Getenv("LLMGW_PYZMQ_INTEROP") == "" {
|
||||
t.Skip("set LLMGW_PYZMQ_INTEROP=1 to run against pyzmq and msgspec (needs uv)")
|
||||
}
|
||||
m := NewMirror()
|
||||
sub, err := Listen("r1", 0, m, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer sub.Close()
|
||||
cmd := exec.Command("uv", "run", "--quiet", "--with", "pyzmq", "--with", "msgspec",
|
||||
"python", "testdata/pypub.py", sub.Endpoint())
|
||||
cmd.Stderr = os.Stderr
|
||||
if err := cmd.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer func() { _ = cmd.Wait() }()
|
||||
|
||||
waitForWithin(t, 60*time.Second, "pyzmq's events", func() bool { return m.Blocks("r1") == 3 })
|
||||
if m.Problem("r1") != "" {
|
||||
t.Fatalf("hashes made by the golden chain should verify: %s", m.Problem("r1"))
|
||||
}
|
||||
waitForWithin(t, 10*time.Second, "the removal", func() bool { return m.Blocks("r1") == 1 })
|
||||
}
|
||||
|
||||
func waitForWithin(t *testing.T, d time.Duration, what string, ok func() bool) {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(d)
|
||||
for !ok() {
|
||||
if time.Now().After(deadline) {
|
||||
t.Fatalf("timed out waiting for %s", what)
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,386 @@
|
||||
package kvevents
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/binary"
|
||||
"encoding/hex"
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/go-zeromq/zmq4"
|
||||
"github.com/vmihailenco/msgpack/v5"
|
||||
)
|
||||
|
||||
// The expected values below were computed by Python with cbor2 and hashlib,
|
||||
// exactly as vLLM's sha256_cbor and hash_block_tokens do it, under
|
||||
// PYTHONHASHSEED=1729:
|
||||
//
|
||||
// NONE = sha256(cbor2.dumps("1729", canonical=True))
|
||||
// h_i = sha256(cbor2.dumps((h_{i-1}, tuple(block_i), None), canonical=True))
|
||||
// external = int.from_bytes(h_i, "big") & (2**64 - 1)
|
||||
func TestBlockHashesAreVLLMsSha256CborChain(t *testing.T) {
|
||||
tokens := []int{1, 23, 24, 255, 256, 65535, 65536, 151643, 4294967296}
|
||||
for i := 100; i < 140; i++ {
|
||||
tokens = append(tokens, i)
|
||||
}
|
||||
got := BlockHashes(tokens, 16, 0)
|
||||
want := []uint64{2557436066646899323, 14514142130713257781, 480514406790448453}
|
||||
if !slices.Equal(got, want) {
|
||||
t.Fatalf("hashes %v, want %v", got, want)
|
||||
}
|
||||
if got := BlockHashes([]int{0, 1, 2, 3, 4}, 4, 0); !slices.Equal(got, []uint64{2734566579713004419}) {
|
||||
t.Fatalf("a partial last block must not be hashed: %v", got)
|
||||
}
|
||||
if got := BlockHashes(tokens, 16, 2); len(got) != 2 {
|
||||
t.Fatalf("max should bound the blocks hashed, got %d", len(got))
|
||||
}
|
||||
if FirstBlockHash(tokens[:16]) != want[0] {
|
||||
t.Fatal("the first block's hash should need nothing but its tokens")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheSeedIsTheOneTheGoldenValuesWereMadeWith(t *testing.T) {
|
||||
sum := noneHash()
|
||||
if hex.EncodeToString(sum[:]) != "25699066d86cd786fe44373a87613dbf15e2386a5d945a476bacb4e72d1130a6" {
|
||||
t.Fatalf("NONE_HASH changed: %x — the router and running engines would no longer agree", sum)
|
||||
}
|
||||
}
|
||||
|
||||
// Encoded by msgspec from vLLM's own struct definitions (main branch, 2026):
|
||||
// tagged-map events, a GPU store of two blocks, a CPU store, a removal and a
|
||||
// clear.
|
||||
const goldenNew = "93cb3ff80000000000009489a474797065ab426c6f636b53746f726564ac626c6f636b5f68617368657392cf237dd68d8042ba7bcfc96c975e60b1e735b1706172656e745f626c6f636b5f68617368c0a9746f6b656e5f696473dc0020000102030405060708090a0b0c0d0e0f000102030405060708090a0b0c0d0e0faa626c6f636b5f73697a6510a76c6f72615f6964c0a66d656469756da3475055a96c6f72615f6e616d65c0a967726f75705f6964780088a474797065ab426c6f636b53746f726564ac626c6f636b5f68617368657391c4200000000000000000000000000000000000000000000000000000000000000007b1706172656e745f626c6f636b5f68617368cfc96c975e60b1e735a9746f6b656e5f696473dc0010000102030405060708090a0b0c0d0e0faa626c6f636b5f73697a6510a76c6f72615f6964c0a66d656469756da3435055a96c6f72615f6e616d65c083a474797065ac426c6f636b52656d6f766564ac626c6f636b5f68617368657391cfc96c975e60b1e735a66d656469756da347505581a474797065b0416c6c426c6f636b73436c656172656400"
|
||||
|
||||
// The same with the older array_like events.
|
||||
const goldenOld = "93cb40000000000000009397ab426c6f636b53746f72656491cf237dd68d8042ba7bc0dc0010000102030405060708090a0b0c0d0e0f10c0c093ac426c6f636b52656d6f76656491cf237dd68d8042ba7bc091b0416c6c426c6f636b73436c6561726564c0"
|
||||
|
||||
func TestABatchInVLLMsCurrentEncodingIsRead(t *testing.T) {
|
||||
raw, _ := hex.DecodeString(goldenNew)
|
||||
b, err := DecodeBatch(raw)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if b.TS != 1.5 || len(b.Events) != 4 {
|
||||
t.Fatalf("batch %+v", b)
|
||||
}
|
||||
st := b.Events[0]
|
||||
if st.Kind != Stored || !slices.Equal(st.Hashes, []uint64{2557436066646899323, 14514142130713257781}) ||
|
||||
st.HasParent || st.BlockSize != 16 || len(st.Tokens) != 32 || st.Elsewhere || st.Extra {
|
||||
t.Fatalf("stored %+v", st)
|
||||
}
|
||||
cpu := b.Events[1]
|
||||
if !cpu.Elsewhere || cpu.Hashes[0] != 7 || !cpu.HasParent || cpu.Parent != 14514142130713257781 {
|
||||
t.Fatalf("a CPU copy, hashed as bytes, should be read and marked elsewhere: %+v", cpu)
|
||||
}
|
||||
if b.Events[2].Kind != Removed || b.Events[2].Hashes[0] != 14514142130713257781 {
|
||||
t.Fatalf("removed %+v", b.Events[2])
|
||||
}
|
||||
if b.Events[3].Kind != Cleared {
|
||||
t.Fatalf("cleared %+v", b.Events[3])
|
||||
}
|
||||
}
|
||||
|
||||
func TestABatchInVLLMsOlderArrayEncodingIsRead(t *testing.T) {
|
||||
raw, _ := hex.DecodeString(goldenOld)
|
||||
b, err := DecodeBatch(raw)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(b.Events) != 3 || b.Events[0].Kind != Stored || b.Events[0].Hashes[0] != 2557436066646899323 ||
|
||||
b.Events[0].BlockSize != 16 || len(b.Events[0].Tokens) != 16 ||
|
||||
b.Events[1].Kind != Removed || b.Events[2].Kind != Cleared {
|
||||
t.Fatalf("batch %+v", b)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnUnknownEventKindIsSkippedNotRefused(t *testing.T) {
|
||||
raw, _ := msgpack.Marshal([]any{1.0, []any{
|
||||
map[string]any{"type": "SomethingNew", "x": 1},
|
||||
map[string]any{"type": "AllBlocksCleared"},
|
||||
}})
|
||||
b, err := DecodeBatch(raw)
|
||||
if err != nil || len(b.Events) != 1 || b.Events[0].Kind != Cleared {
|
||||
t.Fatalf("batch %+v, err %v", b, err)
|
||||
}
|
||||
if _, err := DecodeBatch([]byte{0xc1}); err == nil {
|
||||
t.Fatal("garbage should be refused")
|
||||
}
|
||||
}
|
||||
|
||||
// stored builds a BlockStored event for a chain of blocks, hashed as vLLM
|
||||
// would, in the current encoding.
|
||||
func stored(tokens []int, blockSize int) map[string]any {
|
||||
hs := BlockHashes(tokens, blockSize, 0)
|
||||
anyHs := make([]any, len(hs))
|
||||
for i, h := range hs {
|
||||
anyHs[i] = h
|
||||
}
|
||||
return map[string]any{
|
||||
"type": "BlockStored", "block_hashes": anyHs, "parent_block_hash": nil,
|
||||
"token_ids": tokens[:len(hs)*blockSize], "block_size": blockSize,
|
||||
"lora_id": nil, "medium": "GPU", "lora_name": nil,
|
||||
}
|
||||
}
|
||||
|
||||
func removed(hs ...uint64) map[string]any {
|
||||
anyHs := make([]any, len(hs))
|
||||
for i, h := range hs {
|
||||
anyHs[i] = h
|
||||
}
|
||||
return map[string]any{"type": "BlockRemoved", "block_hashes": anyHs, "medium": "GPU"}
|
||||
}
|
||||
|
||||
func batch(t *testing.T, events ...any) Batch {
|
||||
t.Helper()
|
||||
raw, err := msgpack.Marshal([]any{float64(time.Now().Unix()), events})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
b, err := DecodeBatch(raw)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
func seq(n int) []int {
|
||||
out := make([]int, n)
|
||||
for i := range out {
|
||||
out[i] = 1000 + i
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func TestTheMirrorFollowsStoresRemovalsAndClears(t *testing.T) {
|
||||
m := NewMirror()
|
||||
s := m.Open()
|
||||
defer s.Close()
|
||||
toks := seq(64)
|
||||
hs := BlockHashes(toks, 16, 0)
|
||||
|
||||
m.Apply("r1", batch(t, stored(toks, 16)))
|
||||
if m.Blocks("r1") != 4 {
|
||||
t.Fatalf("blocks %d, want 4", m.Blocks("r1"))
|
||||
}
|
||||
m.Apply("r1", batch(t, removed(hs[3])))
|
||||
ups := s.Drain()
|
||||
if len(ups) != 1 || ups[0].Reset || ups[0].BlockSize != 16 || len(ups[0].Stored) != 3 || len(ups[0].Removed) != 1 {
|
||||
t.Fatalf("a store then a removal should arrive as their net effect: %+v", ups)
|
||||
}
|
||||
if len(s.Drain()) != 0 {
|
||||
t.Fatal("drained changes must not be sent twice")
|
||||
}
|
||||
|
||||
m.Apply("r1", batch(t, map[string]any{"type": "AllBlocksCleared"}))
|
||||
ups = s.Drain()
|
||||
if m.Blocks("r1") != 0 || len(ups) != 1 || !ups[0].Reset || len(ups[0].Stored) != 0 {
|
||||
t.Fatalf("a clear should empty the replica and say so: %+v", ups)
|
||||
}
|
||||
}
|
||||
|
||||
func TestANewStreamStartsWithAFullCopy(t *testing.T) {
|
||||
m := NewMirror()
|
||||
m.Apply("r1", batch(t, stored(seq(48), 16)))
|
||||
m.Apply("r2", batch(t, stored(seq(16), 16)))
|
||||
|
||||
s := m.Open()
|
||||
defer s.Close()
|
||||
ups := s.Drain()
|
||||
if len(ups) != 2 || !ups[0].Reset || len(ups[0].Stored) != 3 || ups[1].Replica != "r2" || len(ups[1].Stored) != 1 {
|
||||
t.Fatalf("a control plane that knows nothing should be told everything: %+v", ups)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAFullCopyOfALargeCacheGoesInPieces(t *testing.T) {
|
||||
m := NewMirror()
|
||||
toks := seq(16 * (maxPerUpdate + 10))
|
||||
m.Apply("r1", batch(t, stored(toks, 16)))
|
||||
s := m.Open()
|
||||
defer s.Close()
|
||||
ups := s.Drain()
|
||||
if len(ups) != 2 || !ups[0].Reset || ups[1].Reset || len(ups[0].Stored)+len(ups[1].Stored) != maxPerUpdate+10 {
|
||||
t.Fatalf("pieces: %d, first reset %v", len(ups), len(ups) > 0 && ups[0].Reset)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCopiesOutsideTheGPUAreNotIndexed(t *testing.T) {
|
||||
m := NewMirror()
|
||||
ev := stored(seq(16), 16)
|
||||
ev["medium"] = "CPU"
|
||||
m.Apply("r1", batch(t, ev))
|
||||
if m.Blocks("r1") != 0 {
|
||||
t.Fatal("a block offloaded to CPU memory is not one the GPU holds")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnEngineHashingAnotherWayIsCaughtAndItsIndexEmptied(t *testing.T) {
|
||||
m := NewMirror()
|
||||
s := m.Open()
|
||||
defer s.Close()
|
||||
for i := range mismatchLimit {
|
||||
ev := stored(seq(16+i), 16)
|
||||
ev["block_hashes"] = []any{uint64(i + 1)} // not what sha256_cbor gives
|
||||
m.Apply("r1", batch(t, ev))
|
||||
}
|
||||
if m.Problem("r1") == "" {
|
||||
t.Fatal("first blocks that never hash alike should be reported")
|
||||
}
|
||||
if m.Blocks("r1") != 0 {
|
||||
t.Fatal("an index of hashes the router cannot compute is worse than none")
|
||||
}
|
||||
ups := s.Drain()
|
||||
if len(ups) != 1 || ups[0].Problem == "" || !ups[0].Reset {
|
||||
t.Fatalf("the problem should travel: %+v", ups)
|
||||
}
|
||||
|
||||
// One bad first block among good ones is a cache salt, not a wrong seed.
|
||||
m2 := NewMirror()
|
||||
m2.Apply("r1", batch(t, stored(seq(16), 16)))
|
||||
for i := range mismatchLimit {
|
||||
ev := stored(seq(17+i), 16)
|
||||
ev["block_hashes"] = []any{uint64(i + 1)}
|
||||
m2.Apply("r1", batch(t, ev))
|
||||
}
|
||||
if m2.Problem("r1") != "" {
|
||||
t.Fatal("an engine seen to hash alike should not be condemned by odd blocks")
|
||||
}
|
||||
}
|
||||
|
||||
// publisher is a stand-in for vLLM's ZmqEventPublisher: a PUB socket that
|
||||
// connects to the endpoint it is given — as vLLM does for an endpoint naming
|
||||
// a host — and sends (topic, 8-byte big-endian sequence, msgpack batch).
|
||||
type publisher struct {
|
||||
t *testing.T
|
||||
sock zmq4.Socket
|
||||
seq uint64
|
||||
}
|
||||
|
||||
func newPublisher(t *testing.T, endpoint string) *publisher {
|
||||
t.Helper()
|
||||
sock := zmq4.NewPub(context.Background())
|
||||
if err := sock.Dial(endpoint); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = sock.Close() })
|
||||
return &publisher{t: t, sock: sock}
|
||||
}
|
||||
|
||||
func (p *publisher) send(events ...any) {
|
||||
p.t.Helper()
|
||||
raw, err := msgpack.Marshal([]any{float64(time.Now().UnixNano()) / 1e9, events})
|
||||
if err != nil {
|
||||
p.t.Fatal(err)
|
||||
}
|
||||
var s [8]byte
|
||||
binary.BigEndian.PutUint64(s[:], p.seq)
|
||||
p.seq++
|
||||
if err := p.sock.Send(zmq4.NewMsgFrom([]byte("kv-events"), s[:], raw)); err != nil {
|
||||
p.t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func waitFor(t *testing.T, what string, ok func() bool) {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(5 * time.Second)
|
||||
for !ok() {
|
||||
if time.Now().After(deadline) {
|
||||
t.Fatalf("timed out waiting for %s", what)
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
}
|
||||
|
||||
// waitSubscribed sends a probe until it arrives: a subscription travels to
|
||||
// the publisher after the connection, and a PUB socket drops what is sent
|
||||
// before it does, as libzmq's does.
|
||||
func waitSubscribed(t *testing.T, p *publisher, m *Mirror, replica string) {
|
||||
t.Helper()
|
||||
probe := seq(16)
|
||||
deadline := time.Now().Add(5 * time.Second)
|
||||
for m.Blocks(replica) == 0 {
|
||||
if time.Now().After(deadline) {
|
||||
t.Fatal("the subscriber never received anything")
|
||||
}
|
||||
p.send(stored(probe, 16))
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
p.send(map[string]any{"type": "AllBlocksCleared"})
|
||||
waitFor(t, "the probe to be cleared", func() bool { return m.Blocks(replica) == 0 })
|
||||
}
|
||||
|
||||
func TestTheSubscriberReceivesAnEnginesEventsOverZMQ(t *testing.T) {
|
||||
m := NewMirror()
|
||||
sub, err := Listen("r1", 0, m, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer sub.Close()
|
||||
if sub.Port() == 0 {
|
||||
t.Fatal("listening on port 0 should report the port chosen")
|
||||
}
|
||||
pub := newPublisher(t, sub.Endpoint())
|
||||
waitSubscribed(t, pub, m, "r1")
|
||||
|
||||
toks := seq(80)
|
||||
pub.send(stored(toks, 16))
|
||||
waitFor(t, "five blocks", func() bool { return m.Blocks("r1") == 5 })
|
||||
|
||||
hs := BlockHashes(toks, 16, 0)
|
||||
pub.send(removed(hs[4], hs[3]))
|
||||
waitFor(t, "three blocks", func() bool { return m.Blocks("r1") == 3 })
|
||||
}
|
||||
|
||||
func TestLostEventsStartTheIndexAgain(t *testing.T) {
|
||||
m := NewMirror()
|
||||
sub, err := Listen("r1", 0, m, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer sub.Close()
|
||||
pub := newPublisher(t, sub.Endpoint())
|
||||
waitSubscribed(t, pub, m, "r1")
|
||||
|
||||
pub.send(stored(seq(64), 16))
|
||||
waitFor(t, "four blocks", func() bool { return m.Blocks("r1") == 4 })
|
||||
|
||||
pub.seq += 3 // three batches never arrive; one may have removed blocks
|
||||
other := seq(16)
|
||||
for i := range other {
|
||||
other[i] += 5000
|
||||
}
|
||||
pub.send(stored(other, 16))
|
||||
waitFor(t, "only the block after the gap", func() bool { return m.Blocks("r1") == 1 })
|
||||
}
|
||||
|
||||
func TestARestartedEngineStartsTheIndexAgain(t *testing.T) {
|
||||
m := NewMirror()
|
||||
sub, err := Listen("r1", 0, m, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer sub.Close()
|
||||
pub := newPublisher(t, sub.Endpoint())
|
||||
waitSubscribed(t, pub, m, "r1")
|
||||
pub.send(stored(seq(64), 16))
|
||||
waitFor(t, "four blocks", func() bool { return m.Blocks("r1") == 4 })
|
||||
|
||||
pub.seq = 0 // a new process numbers from zero again
|
||||
pub.send(stored(seq(16), 16))
|
||||
waitFor(t, "the new engine's one block", func() bool { return m.Blocks("r1") == 1 })
|
||||
}
|
||||
|
||||
func TestTheEngineIsToldToHashTheWayTheRouterDoes(t *testing.T) {
|
||||
args := VLLMArgs("tcp://127.0.0.1:5999")
|
||||
if !slices.Contains(args, "sha256_cbor") || !slices.Contains(args, "--kv-events-config") {
|
||||
t.Fatalf("args %v", args)
|
||||
}
|
||||
cfg := args[slices.Index(args, "--kv-events-config")+1]
|
||||
if want := `"endpoint":"tcp://127.0.0.1:5999"`; !strings.Contains(cfg, want) {
|
||||
t.Fatalf("config %s should name the loopback endpoint, so vLLM connects rather than binding every address", cfg)
|
||||
}
|
||||
if VLLMEnv()["PYTHONHASHSEED"] != HashSeed {
|
||||
t.Fatal("the seed must be the one the router hashes with")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,308 @@
|
||||
package kvevents
|
||||
|
||||
import (
|
||||
"maps"
|
||||
"slices"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// Update is a change to what one replica's cache holds, as it travels from
|
||||
// the agent to the router.
|
||||
//
|
||||
// It is a net effect, not a log. Between two sends a block can be stored,
|
||||
// evicted and stored again; only where it ended up matters, and sending the
|
||||
// history would cost the stream for nothing.
|
||||
type Update struct {
|
||||
Replica string
|
||||
// BlockSize is how many tokens make one of this replica's blocks.
|
||||
BlockSize int
|
||||
// Reset says to forget everything known of the replica before applying
|
||||
// the rest: the engine cleared its cache or restarted, events were lost,
|
||||
// or the stream is new and this is the start of a full copy.
|
||||
Reset bool
|
||||
// Stored and Removed are block hashes, in external form.
|
||||
Stored []uint64
|
||||
Removed []uint64
|
||||
// Problem, when set, says the replica's events cannot be used, and why.
|
||||
// Its index is empty until the problem goes, and requests for it are
|
||||
// routed as though it had none.
|
||||
Problem string
|
||||
}
|
||||
|
||||
// maxPerUpdate bounds the hashes in one update, so a full copy of a large
|
||||
// cache goes up the stream in pieces rather than as one message the size of
|
||||
// the cache.
|
||||
const maxPerUpdate = 32 << 10
|
||||
|
||||
// mismatchLimit is how many first blocks must hash differently, with none
|
||||
// hashing alike, before a replica's hashing is declared wrong. Not one: a
|
||||
// block can carry a cache salt that older engines do not report, and one such
|
||||
// prompt must not switch the index off.
|
||||
const mismatchLimit = 8
|
||||
|
||||
// Mirror holds what each engine's cache holds, as its events describe it, and
|
||||
// what each stream to the control plane has yet to be told.
|
||||
//
|
||||
// It keeps the whole set rather than forwarding events as they come because
|
||||
// the stream is not permanent. A control plane restarts, an agent fails over
|
||||
// to another: the new one knows nothing, and the only way to tell it the
|
||||
// truth is a full copy. Holding the set costs eight bytes a block, and an
|
||||
// engine has some tens of thousands.
|
||||
type Mirror struct {
|
||||
mu sync.Mutex
|
||||
replicas map[string]*mirrored
|
||||
streams map[*Stream]struct{}
|
||||
}
|
||||
|
||||
type mirrored struct {
|
||||
blockSize int
|
||||
blocks map[uint64]struct{}
|
||||
// matched and mismatched count first blocks checked against llmgw's
|
||||
// hashing; problem is set once the check has failed.
|
||||
matched, mismatched int
|
||||
problem string
|
||||
}
|
||||
|
||||
// Stream is one consumer of a mirror's changes: one session with a control
|
||||
// plane.
|
||||
type Stream struct {
|
||||
m *Mirror
|
||||
pending map[string]*delta
|
||||
}
|
||||
|
||||
type delta struct {
|
||||
reset bool
|
||||
blockSize int
|
||||
net map[uint64]bool // true stored, false removed
|
||||
problem string
|
||||
}
|
||||
|
||||
// NewMirror returns an empty mirror.
|
||||
func NewMirror() *Mirror {
|
||||
return &Mirror{replicas: map[string]*mirrored{}, streams: map[*Stream]struct{}{}}
|
||||
}
|
||||
|
||||
// Open starts a stream whose first changes are a full copy of every replica.
|
||||
func (m *Mirror) Open() *Stream {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
s := &Stream{m: m, pending: map[string]*delta{}}
|
||||
for id, r := range m.replicas {
|
||||
d := &delta{reset: true, blockSize: r.blockSize, problem: r.problem,
|
||||
net: make(map[uint64]bool, len(r.blocks))}
|
||||
for h := range r.blocks {
|
||||
d.net[h] = true
|
||||
}
|
||||
s.pending[id] = d
|
||||
}
|
||||
m.streams[s] = struct{}{}
|
||||
return s
|
||||
}
|
||||
|
||||
// Close stops a stream; its pending changes are dropped.
|
||||
func (s *Stream) Close() {
|
||||
s.m.mu.Lock()
|
||||
defer s.m.mu.Unlock()
|
||||
delete(s.m.streams, s)
|
||||
}
|
||||
|
||||
// Drain takes everything the stream has yet to send, as updates of bounded
|
||||
// size. A replica's reset, when there is one, is in its first update.
|
||||
func (s *Stream) Drain() []Update {
|
||||
s.m.mu.Lock()
|
||||
pending := s.pending
|
||||
s.pending = map[string]*delta{}
|
||||
s.m.mu.Unlock()
|
||||
|
||||
var out []Update
|
||||
for _, id := range slices.Sorted(maps.Keys(pending)) {
|
||||
d := pending[id]
|
||||
u := Update{Replica: id, BlockSize: d.blockSize, Reset: d.reset, Problem: d.problem}
|
||||
for h, stored := range d.net {
|
||||
if stored {
|
||||
u.Stored = append(u.Stored, h)
|
||||
} else {
|
||||
u.Removed = append(u.Removed, h)
|
||||
}
|
||||
}
|
||||
out = append(out, split(u)...)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// split cuts an update into pieces of at most maxPerUpdate hashes, keeping the
|
||||
// reset on the first so later pieces add to it rather than wiping it.
|
||||
func split(u Update) []Update {
|
||||
if len(u.Stored)+len(u.Removed) <= maxPerUpdate {
|
||||
return []Update{u}
|
||||
}
|
||||
var out []Update
|
||||
stored, removed := u.Stored, u.Removed
|
||||
first := true
|
||||
for first || len(stored) > 0 || len(removed) > 0 {
|
||||
p := Update{Replica: u.Replica, BlockSize: u.BlockSize, Problem: u.Problem, Reset: first && u.Reset}
|
||||
first = false
|
||||
room := maxPerUpdate
|
||||
n := min(room, len(removed))
|
||||
p.Removed, removed = removed[:n], removed[n:]
|
||||
room -= n
|
||||
n = min(room, len(stored))
|
||||
p.Stored, stored = stored[:n], stored[n:]
|
||||
out = append(out, p)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// Apply folds one batch of a replica's events into the mirror.
|
||||
//
|
||||
// Events about copies outside the GPU, or about a second KV cache group, are
|
||||
// ignored: a block offloaded to CPU memory is a slower hit the router cannot
|
||||
// weigh, and a hybrid model's groups would need a set each. The index then
|
||||
// under-counts, which costs a missed preference and never a wrong one.
|
||||
func (m *Mirror) Apply(replica string, b Batch) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
r := m.replica(replica)
|
||||
for _, e := range b.Events {
|
||||
if e.Kind == Cleared {
|
||||
m.resetLocked(replica, r)
|
||||
continue
|
||||
}
|
||||
if e.Elsewhere {
|
||||
continue
|
||||
}
|
||||
if e.Kind == Stored {
|
||||
m.check(replica, r, e)
|
||||
}
|
||||
if r.problem != "" {
|
||||
continue
|
||||
}
|
||||
if e.Kind == Stored && e.BlockSize > r.blockSize {
|
||||
// The largest seen, not the latest: a partial block is published
|
||||
// with its own, smaller size.
|
||||
r.blockSize = e.BlockSize
|
||||
for s := range m.streams {
|
||||
s.delta(replica).blockSize = r.blockSize
|
||||
}
|
||||
}
|
||||
for _, h := range e.Hashes {
|
||||
if e.Kind == Stored {
|
||||
r.blocks[h] = struct{}{}
|
||||
} else {
|
||||
delete(r.blocks, h)
|
||||
}
|
||||
for s := range m.streams {
|
||||
s.delta(replica).net[h] = e.Kind == Stored
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// check compares the first block of a prompt with llmgw's hashing of its
|
||||
// tokens.
|
||||
//
|
||||
// This is the one thing that can go silently wrong with an index: an engine
|
||||
// started with another hash algorithm or seed — a profile argument overriding
|
||||
// llmgw's, a vLLM that changed its scheme — publishes hashes the router will
|
||||
// never compute, and the index fills with blocks that match nothing. Better to
|
||||
// say so, and route without it.
|
||||
func (m *Mirror) check(replica string, r *mirrored, e Event) {
|
||||
if e.HasParent || e.Extra || len(e.Hashes) == 0 || e.BlockSize <= 0 ||
|
||||
len(e.Tokens) < e.BlockSize || r.problem != "" {
|
||||
return
|
||||
}
|
||||
if FirstBlockHash(e.Tokens[:e.BlockSize]) == e.Hashes[0] {
|
||||
r.matched++
|
||||
return
|
||||
}
|
||||
r.mismatched++
|
||||
if r.matched == 0 && r.mismatched >= mismatchLimit {
|
||||
r.problem = "the engine's block hashes are not " + HashAlgo + " with PYTHONHASHSEED=" + HashSeed +
|
||||
"; remove any --prefix-caching-hash-algo from the profile's arguments, and check the vLLM version"
|
||||
r.blocks = map[uint64]struct{}{}
|
||||
for s := range m.streams {
|
||||
d := s.delta(replica)
|
||||
d.reset = true
|
||||
d.net = map[uint64]bool{}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Reset forgets what a replica's cache holds, keeping the replica: its events
|
||||
// were interrupted, and what was missed cannot be known.
|
||||
func (m *Mirror) Reset(replica string) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
m.resetLocked(replica, m.replica(replica))
|
||||
}
|
||||
|
||||
func (m *Mirror) resetLocked(replica string, r *mirrored) {
|
||||
r.blocks = map[uint64]struct{}{}
|
||||
for s := range m.streams {
|
||||
d := s.delta(replica)
|
||||
d.reset = true
|
||||
d.net = map[uint64]bool{}
|
||||
}
|
||||
}
|
||||
|
||||
// Remove drops a replica whose engine has stopped. The streams are told it
|
||||
// holds nothing, which is true.
|
||||
func (m *Mirror) Remove(replica string) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
if _, ok := m.replicas[replica]; !ok {
|
||||
return
|
||||
}
|
||||
delete(m.replicas, replica)
|
||||
for s := range m.streams {
|
||||
d := s.delta(replica)
|
||||
d.reset = true
|
||||
d.net = map[uint64]bool{}
|
||||
}
|
||||
}
|
||||
|
||||
// Blocks is how many blocks the mirror believes a replica holds.
|
||||
func (m *Mirror) Blocks(replica string) int {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
if r, ok := m.replicas[replica]; ok {
|
||||
return len(r.blocks)
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// Problem is why a replica's events are not used, or "".
|
||||
func (m *Mirror) Problem(replica string) string {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
if r, ok := m.replicas[replica]; ok {
|
||||
return r.problem
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (m *Mirror) replica(id string) *mirrored {
|
||||
r, ok := m.replicas[id]
|
||||
if !ok {
|
||||
r = &mirrored{blocks: map[uint64]struct{}{}}
|
||||
m.replicas[id] = r
|
||||
}
|
||||
return r
|
||||
}
|
||||
|
||||
func (s *Stream) delta(replica string) *delta {
|
||||
d, ok := s.pending[replica]
|
||||
if !ok {
|
||||
d = &delta{net: map[uint64]bool{}}
|
||||
s.pending[replica] = d
|
||||
}
|
||||
if r, ok := s.m.replicas[replica]; ok {
|
||||
if d.blockSize == 0 {
|
||||
d.blockSize = r.blockSize
|
||||
}
|
||||
// Every update carries the replica's standing problem, so one sent
|
||||
// after it was found cannot read as its having gone away.
|
||||
d.problem = r.problem
|
||||
}
|
||||
return d
|
||||
}
|
||||
@@ -0,0 +1,172 @@
|
||||
package kvevents
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/binary"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"net"
|
||||
"sync"
|
||||
|
||||
"github.com/go-zeromq/zmq4"
|
||||
)
|
||||
|
||||
// Subscriber receives one engine's events and folds them into a mirror.
|
||||
//
|
||||
// It listens and the engine connects, which is the direction vLLM chooses for
|
||||
// an endpoint naming a host: a publisher binds only "tcp://*" (every
|
||||
// address), "::" or ipc, and connects to anything else. Listening here on the
|
||||
// loopback address keeps the stream on the machine. The alternative, the
|
||||
// engine binding every address, would publish the token ids of every prompt
|
||||
// to anyone who can reach the node.
|
||||
type Subscriber struct {
|
||||
replica string
|
||||
mirror *Mirror
|
||||
log *slog.Logger
|
||||
sock zmq4.Socket
|
||||
port int
|
||||
cancel context.CancelFunc
|
||||
done chan struct{}
|
||||
|
||||
mu sync.Mutex
|
||||
lastSeq int64
|
||||
seen bool
|
||||
}
|
||||
|
||||
// Listen starts receiving one replica's events on 127.0.0.1:port; port 0
|
||||
// picks a free one, which Port then reports. An agent restarting gives the
|
||||
// port it chose before, because the engine it adopts is still connecting to
|
||||
// it.
|
||||
func Listen(replica string, port int, m *Mirror, log *slog.Logger) (*Subscriber, error) {
|
||||
if log == nil {
|
||||
log = slog.Default()
|
||||
}
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
sock := zmq4.NewSub(ctx)
|
||||
if err := sock.Listen(fmt.Sprintf("tcp://127.0.0.1:%d", port)); err != nil {
|
||||
cancel()
|
||||
_ = sock.Close()
|
||||
return nil, fmt.Errorf("listening for kv cache events: %w", err)
|
||||
}
|
||||
// Every topic: only this replica's engine is told the port.
|
||||
if err := sock.SetOption(zmq4.OptionSubscribe, ""); err != nil {
|
||||
cancel()
|
||||
_ = sock.Close()
|
||||
return nil, fmt.Errorf("subscribing to kv cache events: %w", err)
|
||||
}
|
||||
tcp, ok := sock.Addr().(*net.TCPAddr)
|
||||
if !ok {
|
||||
cancel()
|
||||
_ = sock.Close()
|
||||
return nil, fmt.Errorf("kv cache events: listening on %v, not a TCP address", sock.Addr())
|
||||
}
|
||||
s := &Subscriber{
|
||||
replica: replica, mirror: m, log: log, sock: sock, port: tcp.Port,
|
||||
cancel: cancel, done: make(chan struct{}), lastSeq: -1,
|
||||
}
|
||||
go s.run(ctx)
|
||||
return s, nil
|
||||
}
|
||||
|
||||
// Port is the loopback port the engine is to publish to.
|
||||
func (s *Subscriber) Port() int { return s.port }
|
||||
|
||||
// Endpoint is the address to give the engine.
|
||||
func (s *Subscriber) Endpoint() string { return fmt.Sprintf("tcp://127.0.0.1:%d", s.port) }
|
||||
|
||||
// Close stops receiving. The replica stays in the mirror; Mirror.Remove is
|
||||
// what says its engine has gone.
|
||||
func (s *Subscriber) Close() {
|
||||
s.cancel()
|
||||
_ = s.sock.Close()
|
||||
<-s.done
|
||||
}
|
||||
|
||||
func (s *Subscriber) run(ctx context.Context) {
|
||||
defer close(s.done)
|
||||
for {
|
||||
msg, err := s.sock.Recv()
|
||||
if err != nil {
|
||||
if ctx.Err() != nil {
|
||||
return
|
||||
}
|
||||
s.log.Debug("receiving kv cache events failed", "replica", s.replica, "err", err)
|
||||
continue
|
||||
}
|
||||
s.handle(msg.Frames)
|
||||
}
|
||||
}
|
||||
|
||||
// handle reads one message: vLLM sends three frames, the topic, an eight-byte
|
||||
// big-endian sequence number, and the msgpack batch.
|
||||
func (s *Subscriber) handle(frames [][]byte) {
|
||||
if len(frames) != 3 || len(frames[1]) != 8 {
|
||||
s.log.Debug("ignoring a kv cache event message of unexpected shape",
|
||||
"replica", s.replica, "frames", len(frames))
|
||||
return
|
||||
}
|
||||
seq := int64(binary.BigEndian.Uint64(frames[1]))
|
||||
b, err := DecodeBatch(frames[2])
|
||||
|
||||
s.mu.Lock()
|
||||
expected := s.lastSeq + 1
|
||||
gap := s.seen && seq != expected
|
||||
restarted := s.seen && seq < expected
|
||||
s.lastSeq, s.seen = seq, true
|
||||
s.mu.Unlock()
|
||||
|
||||
switch {
|
||||
case restarted:
|
||||
// Numbering started again: a new engine process on the same port,
|
||||
// with a cache that starts empty.
|
||||
s.log.Info("kv cache events restarted; the engine's cache is new", "replica", s.replica)
|
||||
s.mirror.Reset(s.replica)
|
||||
case gap:
|
||||
// Events were lost, among them perhaps removals. An index still
|
||||
// holding blocks the engine evicted would send requests to a replica
|
||||
// for a cache it no longer has; starting again from nothing only
|
||||
// misses some it does. vLLM can replay from a buffer, which would
|
||||
// keep the index whole; not used yet.
|
||||
s.log.Warn("kv cache events were lost; starting the replica's index again",
|
||||
"replica", s.replica, "expected", expected, "got", seq)
|
||||
s.mirror.Reset(s.replica)
|
||||
}
|
||||
if err != nil {
|
||||
// A batch that cannot be read is a gap of one.
|
||||
s.log.Warn("kv cache event batch unreadable; starting the replica's index again",
|
||||
"replica", s.replica, "err", err)
|
||||
s.mirror.Reset(s.replica)
|
||||
return
|
||||
}
|
||||
s.mirror.Apply(s.replica, b)
|
||||
}
|
||||
|
||||
// VLLMArgs are the arguments that make vLLM publish to endpoint, hashing the
|
||||
// way the router computes.
|
||||
//
|
||||
// Prefix caching is turned on explicitly, though it is vLLM's default, because
|
||||
// an index of a cache that is off would be an index of nothing.
|
||||
func VLLMArgs(endpoint string) []string {
|
||||
cfg, _ := json.Marshal(map[string]any{
|
||||
"enable_kv_cache_events": true,
|
||||
"publisher": "zmq",
|
||||
"endpoint": endpoint,
|
||||
"topic": "kv-events",
|
||||
})
|
||||
return []string{
|
||||
"--enable-prefix-caching",
|
||||
"--prefix-caching-hash-algo", HashAlgo,
|
||||
"--kv-events-config", string(cfg),
|
||||
}
|
||||
}
|
||||
|
||||
// VLLMEnv is the environment that goes with VLLMArgs: the seed every hash
|
||||
// chain starts from, and hashes published as integers whatever the release's
|
||||
// default.
|
||||
func VLLMEnv() map[string]string {
|
||||
return map[string]string{
|
||||
"PYTHONHASHSEED": HashSeed,
|
||||
"VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES": "1",
|
||||
}
|
||||
}
|
||||
Vendored
+77
@@ -0,0 +1,77 @@
|
||||
# Publishes KV cache events the way vLLM's ZmqEventPublisher does, with
|
||||
# libzmq (pyzmq) and msgspec and vLLM's own struct definitions, for
|
||||
# TestInteropWithPyZMQ. Not vLLM itself: its publisher thread reduced to the
|
||||
# calls that put bytes on the wire.
|
||||
#
|
||||
# python pypub.py tcp://127.0.0.1:PORT
|
||||
import sys
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
import msgspec
|
||||
import zmq
|
||||
|
||||
|
||||
class EventBatch(msgspec.Struct, array_like=True, omit_defaults=True, gc=False):
|
||||
ts: float
|
||||
events: list[Any]
|
||||
data_parallel_rank: int | None = None
|
||||
|
||||
|
||||
class KVCacheEvent(msgspec.Struct, omit_defaults=True, gc=False, tag=True):
|
||||
pass
|
||||
|
||||
|
||||
class BlockStored(KVCacheEvent):
|
||||
block_hashes: list[bytes | int]
|
||||
parent_block_hash: bytes | int | None
|
||||
token_ids: list[int]
|
||||
block_size: int
|
||||
lora_id: int | None
|
||||
medium: str | None
|
||||
lora_name: str | None
|
||||
extra_keys: list[tuple[Any, ...] | None] | None = None
|
||||
group_idx: int | None = None
|
||||
|
||||
|
||||
class BlockRemoved(KVCacheEvent):
|
||||
block_hashes: list[bytes | int]
|
||||
medium: str | None
|
||||
group_idx: int | None = None
|
||||
|
||||
|
||||
class AllBlocksCleared(KVCacheEvent):
|
||||
pass
|
||||
|
||||
|
||||
class KVEventBatch(EventBatch):
|
||||
events: list[BlockStored | BlockRemoved | AllBlocksCleared]
|
||||
|
||||
|
||||
# The chain TestBlockHashesAreVLLMsSha256CborChain checks.
|
||||
TOKENS = [1, 23, 24, 255, 256, 65535, 65536, 151643, 4294967296] + list(range(100, 140))
|
||||
HASHES = [2557436066646899323, 14514142130713257781, 480514406790448453]
|
||||
|
||||
ctx = zmq.Context.instance()
|
||||
pub = ctx.socket(zmq.PUB)
|
||||
pub.connect(sys.argv[1]) # a host, not "*": vLLM connects
|
||||
enc = msgspec.msgpack.Encoder()
|
||||
seq = 0
|
||||
|
||||
|
||||
def publish(batch):
|
||||
global seq
|
||||
pub.send_multipart((b"kv-events", seq.to_bytes(8, "big"), enc.encode(batch)))
|
||||
seq += 1
|
||||
|
||||
|
||||
stored = KVEventBatch(ts=time.time(), events=[BlockStored(
|
||||
block_hashes=HASHES, parent_block_hash=None, token_ids=TOKENS[:48],
|
||||
block_size=16, lora_id=None, medium="GPU", lora_name=None, group_idx=0)])
|
||||
# Until the subscriber has said what it wants, PUB drops; the Go side waits
|
||||
# for three blocks, then for one, which the removal leaves.
|
||||
for _ in range(40):
|
||||
publish(stored)
|
||||
time.sleep(0.05)
|
||||
publish(KVEventBatch(ts=time.time(), events=[BlockRemoved(block_hashes=HASHES[1:], medium="GPU")]))
|
||||
time.sleep(0.5)
|
||||
@@ -41,6 +41,9 @@ func (r *Router) AdminHandler() (*http.ServeMux, error) {
|
||||
if err := reg.Register(newEndpointCollector(r)); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := reg.Register(newKVCollector(r.kv)); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for _, c := range r.extraMetrics {
|
||||
if err := reg.Register(c); err != nil {
|
||||
return nil, err
|
||||
|
||||
@@ -0,0 +1,33 @@
|
||||
package kvindex
|
||||
|
||||
// Decision is which branch a precise-index lookup took, counted in
|
||||
// llmgw_kv_index_decisions_total beside tier 1's llmgw_prefix_decisions_total.
|
||||
//
|
||||
// The fallbacks are counted by why. "The index is not helping" has three
|
||||
// causes that want different fixes — no events are arriving, the requests are
|
||||
// of a kind it cannot index, the engine is too slow to tokenize — and one
|
||||
// count would not say which.
|
||||
type Decision string
|
||||
|
||||
// The branches.
|
||||
const (
|
||||
// Hit: some replica holds enough of the prompt's leading blocks, and
|
||||
// replicas are scored by how many.
|
||||
Hit Decision = "kv_hit"
|
||||
// MissMinLoad: the index knows the replicas and none holds enough; load
|
||||
// decides, and tier 1 is not asked — its guess cannot beat a fact.
|
||||
MissMinLoad Decision = "kv_miss_min_load"
|
||||
// LoadOverride: load is lopsided enough that affinity is ignored, as in
|
||||
// tier 1.
|
||||
LoadOverride Decision = "load_balance_override"
|
||||
// FallbackNoIndex: no candidate's events have arrived (or their hashing
|
||||
// was found wrong); tier 1 decides.
|
||||
FallbackNoIndex Decision = "fallback_no_index"
|
||||
// FallbackUnsupported: the request cannot be indexed — an image, a cache
|
||||
// salt, an adapter, a list of prompts, an endpoint other than chat or
|
||||
// completions; tier 1 decides.
|
||||
FallbackUnsupported Decision = "fallback_unsupported"
|
||||
// FallbackTokenizeFailed: the engine did not tokenize the request in
|
||||
// time; tier 1 decides.
|
||||
FallbackTokenizeFailed Decision = "fallback_tokenize_failed"
|
||||
)
|
||||
@@ -0,0 +1,132 @@
|
||||
// Package kvindex knows which blocks of KV cache each replica holds, from the
|
||||
// events its engine publishes (plan §7.5 tier 2).
|
||||
//
|
||||
// Tier 1 guesses from what the router sent where. This is what the engines
|
||||
// say they hold, so it also knows what they have evicted, and what reached
|
||||
// them some other way. It is precise only about the replicas it hears from;
|
||||
// for the rest the router falls back to the guess.
|
||||
package kvindex
|
||||
|
||||
import (
|
||||
"sync"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/kvevents"
|
||||
)
|
||||
|
||||
// Index is the block hashes each replica holds.
|
||||
type Index struct {
|
||||
mu sync.RWMutex
|
||||
replicas map[string]*replica
|
||||
}
|
||||
|
||||
type replica struct {
|
||||
blockSize int
|
||||
blocks map[uint64]struct{}
|
||||
problem string
|
||||
}
|
||||
|
||||
// New returns an empty index.
|
||||
func New() *Index { return &Index{replicas: map[string]*replica{}} }
|
||||
|
||||
// Apply folds in one update from an agent.
|
||||
func (x *Index) Apply(u kvevents.Update) {
|
||||
if u.Replica == "" {
|
||||
return
|
||||
}
|
||||
x.mu.Lock()
|
||||
defer x.mu.Unlock()
|
||||
r, ok := x.replicas[u.Replica]
|
||||
if !ok {
|
||||
r = &replica{blocks: map[uint64]struct{}{}}
|
||||
x.replicas[u.Replica] = r
|
||||
}
|
||||
if u.Reset {
|
||||
r.blocks = map[uint64]struct{}{}
|
||||
}
|
||||
if u.BlockSize > 0 {
|
||||
r.blockSize = u.BlockSize
|
||||
}
|
||||
r.problem = u.Problem
|
||||
if r.problem != "" {
|
||||
r.blocks = map[uint64]struct{}{}
|
||||
return
|
||||
}
|
||||
for _, h := range u.Removed {
|
||||
delete(r.blocks, h)
|
||||
}
|
||||
for _, h := range u.Stored {
|
||||
r.blocks[h] = struct{}{}
|
||||
}
|
||||
}
|
||||
|
||||
// Forget drops a replica: it has gone, or the stream that described it has.
|
||||
func (x *Index) Forget(id string) {
|
||||
x.mu.Lock()
|
||||
defer x.mu.Unlock()
|
||||
delete(x.replicas, id)
|
||||
}
|
||||
|
||||
// BlockSizes are the distinct block sizes of the replicas among ids the index
|
||||
// can say something about. Empty means it can say nothing, and the request
|
||||
// is not worth tokenizing.
|
||||
func (x *Index) BlockSizes(ids []string) []int {
|
||||
x.mu.RLock()
|
||||
defer x.mu.RUnlock()
|
||||
var out []int
|
||||
for _, id := range ids {
|
||||
r, ok := x.replicas[id]
|
||||
if !ok || r.problem != "" || r.blockSize <= 0 {
|
||||
continue
|
||||
}
|
||||
dup := false
|
||||
for _, s := range out {
|
||||
dup = dup || s == r.blockSize
|
||||
}
|
||||
if !dup {
|
||||
out = append(out, r.blockSize)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// Leading reports how many of a prompt's leading blocks a replica holds, and
|
||||
// whether the index knows the replica at all.
|
||||
//
|
||||
// Leading, not total: a block is only reusable when every block before it is,
|
||||
// because its hash is chained from theirs. vLLM stops at the first miss, and
|
||||
// so does this. hashes is keyed by block size, since a replica's blocks are
|
||||
// only comparable with hashes cut at its own.
|
||||
func (x *Index) Leading(id string, hashes map[int][]uint64) (matched, total int, known bool) {
|
||||
x.mu.RLock()
|
||||
defer x.mu.RUnlock()
|
||||
r, ok := x.replicas[id]
|
||||
if !ok || r.problem != "" || r.blockSize <= 0 {
|
||||
return 0, 0, false
|
||||
}
|
||||
hs := hashes[r.blockSize]
|
||||
for _, h := range hs {
|
||||
if _, ok := r.blocks[h]; !ok {
|
||||
break
|
||||
}
|
||||
matched++
|
||||
}
|
||||
return matched, len(hs), true
|
||||
}
|
||||
|
||||
// Stat is what the index holds of one replica, for the metrics.
|
||||
type Stat struct {
|
||||
Replica string
|
||||
Blocks int
|
||||
Problem string
|
||||
}
|
||||
|
||||
// Stats lists every replica the index has heard of.
|
||||
func (x *Index) Stats() []Stat {
|
||||
x.mu.RLock()
|
||||
defer x.mu.RUnlock()
|
||||
out := make([]Stat, 0, len(x.replicas))
|
||||
for id, r := range x.replicas {
|
||||
out = append(out, Stat{Replica: id, Blocks: len(r.blocks), Problem: r.problem})
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,67 @@
|
||||
package kvindex
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/kvevents"
|
||||
)
|
||||
|
||||
func TestLeadingStopsAtTheFirstMissingBlock(t *testing.T) {
|
||||
x := New()
|
||||
x.Apply(kvevents.Update{Replica: "a", BlockSize: 16, Stored: []uint64{1, 2, 4}})
|
||||
m, total, ok := x.Leading("a", map[int][]uint64{16: {1, 2, 3, 4}})
|
||||
if !ok || m != 2 || total != 4 {
|
||||
t.Fatalf("matched %d of %d (known %v); block 4 is useless without block 3", m, total, ok)
|
||||
}
|
||||
if _, _, ok := x.Leading("b", nil); ok {
|
||||
t.Fatal("a replica never heard from is unknown, not empty")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAResetReplacesWhatWasKnown(t *testing.T) {
|
||||
x := New()
|
||||
x.Apply(kvevents.Update{Replica: "a", BlockSize: 16, Stored: []uint64{1, 2}})
|
||||
x.Apply(kvevents.Update{Replica: "a", Reset: true, Stored: []uint64{9}})
|
||||
if m, _, _ := x.Leading("a", map[int][]uint64{16: {1}}); m != 0 {
|
||||
t.Fatal("a reset must forget the blocks before it")
|
||||
}
|
||||
if x.Stats()[0].Blocks != 1 {
|
||||
t.Fatal("and keep what came with it")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAProblemReplicaIsNotIndexed(t *testing.T) {
|
||||
x := New()
|
||||
x.Apply(kvevents.Update{Replica: "a", BlockSize: 16, Stored: []uint64{1}})
|
||||
x.Apply(kvevents.Update{Replica: "a", Problem: "hashes differ"})
|
||||
if len(x.BlockSizes([]string{"a"})) != 0 {
|
||||
t.Fatal("a replica whose hashing is wrong must not be asked about")
|
||||
}
|
||||
x.Apply(kvevents.Update{Replica: "a", BlockSize: 16, Reset: true, Stored: []uint64{1}})
|
||||
if len(x.BlockSizes([]string{"a"})) != 1 {
|
||||
t.Fatal("once the problem is gone the replica is indexed again")
|
||||
}
|
||||
}
|
||||
|
||||
func TestOnlyWhatChangesTheTokensIsSentToTokenize(t *testing.T) {
|
||||
cases := map[string]struct {
|
||||
path, body string
|
||||
ok bool
|
||||
}{
|
||||
"chat": {"/v1/chat/completions", `{"model":"x","messages":[{"role":"user","content":"hi"}]}`, true},
|
||||
"text parts": {"/v1/chat/completions", `{"messages":[{"role":"user","content":[{"type":"text","text":"hi"}]}]}`, true},
|
||||
"image": {"/v1/chat/completions", `{"messages":[{"role":"user","content":[{"type":"image_url"}]}]}`, false},
|
||||
"cache salt": {"/v1/chat/completions", `{"messages":[],"cache_salt":"s"}`, false},
|
||||
"completion": {"/v1/completions", `{"prompt":"once upon"}`, true},
|
||||
"prompt list": {"/v1/completions", `{"prompt":["a","b"]}`, false},
|
||||
"embeddings": {"/v1/embeddings", `{"input":"x"}`, false},
|
||||
"not json at all": {"/v1/chat/completions", `{`, false},
|
||||
}
|
||||
for name, c := range cases {
|
||||
_, err := TokenizeBody(c.path, []byte(c.body), "served")
|
||||
if c.ok != (err == nil) || (err != nil && !errors.Is(err, ErrUnsupported)) {
|
||||
t.Errorf("%s: err %v", name, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,146 @@
|
||||
package kvindex
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
)
|
||||
|
||||
// Why tokenizing goes to the engine rather than happening here.
|
||||
//
|
||||
// Block hashes are over token ids, so the router needs the exact ids the
|
||||
// engine will compute. For a chat request those come from the model's chat
|
||||
// template — a Jinja program shipped with the tokenizer, which differs per
|
||||
// model and takes the tools and template arguments as input — and only then
|
||||
// from the tokenizer. A Go tokenizer for tokenizer.json would get the second
|
||||
// half right and still need a Jinja interpreter faithful to transformers' for
|
||||
// the first; a mismatch anywhere is an index that never matches, silently.
|
||||
// vLLM's /tokenize runs the same code the request will, so it agrees by
|
||||
// construction.
|
||||
//
|
||||
// The price is a round trip to a replica on the management network before
|
||||
// the request is placed, and the template and tokenizer run twice on the
|
||||
// engine — once here, once for the request. Against a fake engine on
|
||||
// loopback the router's side of it — the call, decoding 2048 token ids and
|
||||
// hashing 128 blocks — measured about 0.2ms (BenchmarkTokenize in the router
|
||||
// package). Against a real vLLM add its templating and tokenizing, which has
|
||||
// not been measured here; it is expected to be milliseconds for prompts of a
|
||||
// few thousand tokens. TokenizeTimeout
|
||||
// bounds it: a slow answer means routing as tier 1 would, not a slow request.
|
||||
|
||||
// ErrUnsupported says a request cannot be indexed: its blocks would be hashed
|
||||
// with more than their tokens, or it is not a shape /tokenize takes.
|
||||
var ErrUnsupported = errors.New("request not indexable")
|
||||
|
||||
// TokenizeBody builds the /tokenize request for an inference request, or
|
||||
// says why there is none.
|
||||
//
|
||||
// Only what changes the token ids is carried. A request the engine would
|
||||
// hash with extra keys — a cache salt, an image — gets ErrUnsupported: its
|
||||
// blocks hash differently from its tokens alone, and asking the index about
|
||||
// them would only ever miss.
|
||||
func TokenizeBody(path string, body []byte, model string) ([]byte, error) {
|
||||
var fields map[string]json.RawMessage
|
||||
if err := json.Unmarshal(body, &fields); err != nil {
|
||||
return nil, fmt.Errorf("%w: %v", ErrUnsupported, err)
|
||||
}
|
||||
if _, ok := fields["cache_salt"]; ok {
|
||||
return nil, fmt.Errorf("%w: cache_salt", ErrUnsupported)
|
||||
}
|
||||
out := map[string]json.RawMessage{}
|
||||
modelJSON, _ := json.Marshal(model)
|
||||
out["model"] = modelJSON
|
||||
|
||||
switch path {
|
||||
case "/v1/chat/completions":
|
||||
msgs, ok := fields["messages"]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("%w: no messages", ErrUnsupported)
|
||||
}
|
||||
if !textOnly(msgs) {
|
||||
return nil, fmt.Errorf("%w: content other than text", ErrUnsupported)
|
||||
}
|
||||
out["messages"] = msgs
|
||||
for _, k := range []string{"tools", "add_generation_prompt", "continue_final_message",
|
||||
"chat_template", "chat_template_kwargs", "add_special_tokens"} {
|
||||
if v, ok := fields[k]; ok && string(v) != "null" {
|
||||
out[k] = v
|
||||
}
|
||||
}
|
||||
case "/v1/completions":
|
||||
var s string
|
||||
if err := json.Unmarshal(fields["prompt"], &s); err != nil {
|
||||
// A list of prompts, or of token ids: a batch the index would have
|
||||
// to place as one, which it cannot.
|
||||
return nil, fmt.Errorf("%w: prompt is not a single string", ErrUnsupported)
|
||||
}
|
||||
out["prompt"] = fields["prompt"]
|
||||
if v, ok := fields["add_special_tokens"]; ok {
|
||||
out["add_special_tokens"] = v
|
||||
}
|
||||
default:
|
||||
return nil, fmt.Errorf("%w: %s", ErrUnsupported, path)
|
||||
}
|
||||
return json.Marshal(out)
|
||||
}
|
||||
|
||||
// textOnly reports whether every message's content is a string or a list of
|
||||
// text parts. An image or audio part is hashed by the engine with its
|
||||
// identity as an extra key.
|
||||
func textOnly(raw json.RawMessage) bool {
|
||||
var msgs []struct {
|
||||
Content json.RawMessage `json:"content"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &msgs); err != nil {
|
||||
return false
|
||||
}
|
||||
for _, m := range msgs {
|
||||
if len(m.Content) == 0 || m.Content[0] != '[' {
|
||||
continue
|
||||
}
|
||||
var parts []struct {
|
||||
Type string `json:"type"`
|
||||
}
|
||||
if err := json.Unmarshal(m.Content, &parts); err != nil {
|
||||
return false
|
||||
}
|
||||
for _, p := range parts {
|
||||
if p.Type != "text" {
|
||||
return false
|
||||
}
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// Tokenize asks an engine for a request's token ids.
|
||||
func Tokenize(ctx context.Context, client *http.Client, baseURL, apiKey string, body []byte) ([]int, error) {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, baseURL+"/tokenize", bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
if apiKey != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
msg, _ := io.ReadAll(io.LimitReader(resp.Body, 512))
|
||||
return nil, fmt.Errorf("tokenize: %s: %s", resp.Status, bytes.TrimSpace(msg))
|
||||
}
|
||||
var out struct {
|
||||
Tokens []int `json:"tokens"`
|
||||
}
|
||||
if err := json.NewDecoder(io.LimitReader(resp.Body, 64<<20)).Decode(&out); err != nil {
|
||||
return nil, fmt.Errorf("tokenize: %w", err)
|
||||
}
|
||||
return out.Tokens, nil
|
||||
}
|
||||
@@ -0,0 +1,151 @@
|
||||
package router
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/prometheus/client_golang/prometheus"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/kvevents"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router/kvindex"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router/scheduler"
|
||||
)
|
||||
|
||||
// ApplyKV folds in a change to a replica's KV cache, as its node's agent
|
||||
// forwarded it. The control plane calls it for every KVEvents message.
|
||||
func (r *Router) ApplyKV(u kvevents.Update) { r.kv.Apply(u) }
|
||||
|
||||
// kvSettings are router.kv_index with the defaults filled in.
|
||||
func (r *Router) kvSettings() (timeout time.Duration, maxBlocks int) {
|
||||
timeout, maxBlocks = r.cfg.KVIndex.TokenizeTimeout, r.cfg.KVIndex.MaxBlocks
|
||||
if timeout <= 0 {
|
||||
timeout = 250 * time.Millisecond
|
||||
}
|
||||
if maxBlocks <= 0 {
|
||||
maxBlocks = 4096
|
||||
}
|
||||
return timeout, maxBlocks
|
||||
}
|
||||
|
||||
// kvQuery asks the precise index about a request for a model that uses it,
|
||||
// before scheduling: which of the candidates hold how many of the prompt's
|
||||
// leading blocks. Whatever goes wrong, the answer is a fallback to tier 1, and
|
||||
// the request is never refused or held up past the tokenize timeout for it.
|
||||
func (r *Router) kvQuery(ctx context.Context, req *parsedRequest, model string, eps []endpoint.Snapshot) *scheduler.KVQuery {
|
||||
ids := make([]string, len(eps))
|
||||
for i, e := range eps {
|
||||
ids[i] = e.ID
|
||||
}
|
||||
// Asked first because it is free: with no events from any candidate
|
||||
// there is nothing to look up, and nothing worth a round trip.
|
||||
sizes := r.kv.BlockSizes(ids)
|
||||
if len(sizes) == 0 {
|
||||
return &scheduler.KVQuery{Fallback: kvindex.FallbackNoIndex}
|
||||
}
|
||||
// An adapter's blocks are hashed with its name, a form's are not text.
|
||||
if _, ok := r.adapter(req.Model); ok || req.ContentType != "" {
|
||||
return &scheduler.KVQuery{Fallback: kvindex.FallbackUnsupported}
|
||||
}
|
||||
|
||||
// Tokenized by the least busy candidate the index knows: any of them
|
||||
// runs the same template and tokenizer, and the busiest is the one
|
||||
// least able to spare the time.
|
||||
var by *endpoint.Snapshot
|
||||
for i := range eps {
|
||||
e := &eps[i]
|
||||
if _, _, known := r.kv.Leading(e.ID, nil); !known {
|
||||
continue
|
||||
}
|
||||
if by == nil || e.Inflight+e.Stats.Waiting < by.Inflight+by.Stats.Waiting {
|
||||
by = e
|
||||
}
|
||||
}
|
||||
served := model
|
||||
if by.UpstreamModel != "" {
|
||||
served = by.UpstreamModel
|
||||
}
|
||||
body, err := kvindex.TokenizeBody(req.Path, req.UpstreamBody, served)
|
||||
if err != nil {
|
||||
return &scheduler.KVQuery{Fallback: kvindex.FallbackUnsupported}
|
||||
}
|
||||
|
||||
timeout, maxBlocks := r.kvSettings()
|
||||
tctx, cancel := context.WithTimeout(ctx, timeout)
|
||||
defer cancel()
|
||||
start := time.Now()
|
||||
tokens, err := kvindex.Tokenize(tctx, r.kvClient, by.URL, by.APIKey, body)
|
||||
outcome := "ok"
|
||||
if err != nil {
|
||||
outcome = "error"
|
||||
if errors.Is(err, context.DeadlineExceeded) {
|
||||
outcome = "timeout"
|
||||
}
|
||||
}
|
||||
r.m.kvTokenize.WithLabelValues(outcome).Observe(time.Since(start).Seconds())
|
||||
if err != nil {
|
||||
r.log.Debug("tokenizing for the kv index failed; routing by tier 1",
|
||||
"model", model, "endpoint", by.ID, "err", err)
|
||||
return &scheduler.KVQuery{Fallback: kvindex.FallbackTokenizeFailed}
|
||||
}
|
||||
|
||||
hashes := make(map[int][]uint64, len(sizes))
|
||||
for _, s := range sizes {
|
||||
hashes[s] = kvevents.BlockHashes(tokens, s, maxBlocks)
|
||||
}
|
||||
q := &scheduler.KVQuery{Matched: map[string]int{}, Total: map[string]int{}}
|
||||
for _, id := range ids {
|
||||
if m, t, known := r.kv.Leading(id, hashes); known {
|
||||
q.Matched[id], q.Total[id] = m, t
|
||||
}
|
||||
}
|
||||
return q
|
||||
}
|
||||
|
||||
// newKVClient is the client the tokenize calls use: engines are on the
|
||||
// management network, and a call that has not connected in a second is one
|
||||
// the timeout will end anyway.
|
||||
func newKVClient() *http.Client {
|
||||
return &http.Client{Transport: &http.Transport{
|
||||
MaxIdleConnsPerHost: 16,
|
||||
IdleConnTimeout: 90 * time.Second,
|
||||
}}
|
||||
}
|
||||
|
||||
// kvCollector reports what the index holds of each replica.
|
||||
type kvCollector struct {
|
||||
idx *kvindex.Index
|
||||
blocks *prometheus.Desc
|
||||
problem *prometheus.Desc
|
||||
}
|
||||
|
||||
func newKVCollector(idx *kvindex.Index) *kvCollector {
|
||||
return &kvCollector{
|
||||
idx: idx,
|
||||
blocks: prometheus.NewDesc("llmgw_kv_index_blocks",
|
||||
"KV cache blocks each replica reports holding, as the router's precise index has them.",
|
||||
[]string{"replica"}, nil),
|
||||
problem: prometheus.NewDesc("llmgw_kv_index_problem",
|
||||
"1 for a replica whose KV cache events cannot be used (its engine hashes another way); "+
|
||||
"the agent's log says why.",
|
||||
[]string{"replica"}, nil),
|
||||
}
|
||||
}
|
||||
|
||||
func (c *kvCollector) Describe(ch chan<- *prometheus.Desc) {
|
||||
ch <- c.blocks
|
||||
ch <- c.problem
|
||||
}
|
||||
|
||||
func (c *kvCollector) Collect(ch chan<- prometheus.Metric) {
|
||||
for _, s := range c.idx.Stats() {
|
||||
ch <- prometheus.MustNewConstMetric(c.blocks, prometheus.GaugeValue, float64(s.Blocks), s.Replica)
|
||||
p := 0.0
|
||||
if s.Problem != "" {
|
||||
p = 1
|
||||
}
|
||||
ch <- prometheus.MustNewConstMetric(c.problem, prometheus.GaugeValue, p, s.Replica)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,285 @@
|
||||
package router
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
|
||||
"github.com/prometheus/client_golang/prometheus/testutil"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/config"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/kvevents"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router/kvindex"
|
||||
)
|
||||
|
||||
// fakeVLLM is an engine with vLLM's /tokenize beside the completion: it
|
||||
// "tokenizes" a chat as the bytes of each message's role and content, which
|
||||
// is enough for the router and a test to agree on token ids.
|
||||
type fakeVLLM struct {
|
||||
srv *httptest.Server
|
||||
served atomic.Int64
|
||||
tokenized atomic.Int64
|
||||
failTok atomic.Bool
|
||||
lastTok atomic.Value
|
||||
}
|
||||
|
||||
func newFakeVLLM(t *testing.T) *fakeVLLM {
|
||||
t.Helper()
|
||||
f := &fakeVLLM{}
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("GET /metrics", func(w http.ResponseWriter, _ *http.Request) {
|
||||
fmt.Fprintf(w, "vllm:num_requests_running 0\nvllm:num_requests_waiting 0\nvllm:kv_cache_usage_perc 0.1\n")
|
||||
})
|
||||
mux.HandleFunc("POST /tokenize", func(w http.ResponseWriter, r *http.Request) {
|
||||
f.tokenized.Add(1)
|
||||
b, _ := io.ReadAll(r.Body)
|
||||
f.lastTok.Store(b)
|
||||
if f.failTok.Load() {
|
||||
http.Error(w, "overloaded", http.StatusServiceUnavailable)
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
Messages []struct{ Role, Content string } `json:"messages"`
|
||||
}
|
||||
if err := json.Unmarshal(b, &req); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
toks := fakeTokens(req.Messages...)
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"count": len(toks), "max_model_len": 4096, "tokens": toks})
|
||||
})
|
||||
mux.HandleFunc("POST /v1/chat/completions", func(w http.ResponseWriter, _ *http.Request) {
|
||||
f.served.Add(1)
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
fmt.Fprint(w, `{"choices":[{"message":{"content":"hi"}}],"usage":{"prompt_tokens":4,"completion_tokens":2}}`)
|
||||
})
|
||||
f.srv = httptest.NewServer(mux)
|
||||
t.Cleanup(f.srv.Close)
|
||||
return f
|
||||
}
|
||||
|
||||
func fakeTokens(msgs ...struct{ Role, Content string }) []int {
|
||||
var out []int
|
||||
for _, m := range msgs {
|
||||
for _, c := range []byte(m.Role + ":" + m.Content + "\n") {
|
||||
out = append(out, int(c))
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func chatOf(prompt string) []struct{ Role, Content string } {
|
||||
return []struct{ Role, Content string }{{"user", prompt}}
|
||||
}
|
||||
|
||||
// kvRouter is a router over two fake vLLM replicas, a and b, of a model
|
||||
// whose profile publishes KV cache events.
|
||||
func kvRouter(t *testing.T, tier1 bool) (*Router, *fakeVLLM, *fakeVLLM) {
|
||||
t.Helper()
|
||||
a, b := newFakeVLLM(t), newFakeVLLM(t)
|
||||
cfg := baseConfig(
|
||||
config.Endpoint{ID: "a", URL: a.srv.URL, Models: []string{"m"}},
|
||||
config.Endpoint{ID: "b", URL: b.srv.URL, Models: []string{"m"}},
|
||||
config.Endpoint{ID: "o-1", URL: a.srv.URL, Models: []string{"other"}},
|
||||
)
|
||||
m := catModel("m", nil)
|
||||
m.Profiles[0].KVEvents = true
|
||||
m.RoutingProfile = "chat"
|
||||
other := catModel("other", nil)
|
||||
other.RoutingProfile = "chat"
|
||||
cfg.Models = []catalog.Model{m, other}
|
||||
cfg.Router.Prefix.Enabled = tier1
|
||||
return newTestRouter(t, cfg), a, b
|
||||
}
|
||||
|
||||
// longPrompt is long enough for several blocks of sixteen.
|
||||
var longPrompt = strings.Repeat("You are a careful assistant who answers briefly. ", 6)
|
||||
|
||||
func sendChat(t *testing.T, rt *Router, prompt string) {
|
||||
t.Helper()
|
||||
rec := post(t, rt.Handler(), map[string]any{
|
||||
"model": "m", "messages": []map[string]any{{"role": "user", "content": prompt}},
|
||||
}, nil)
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status %d: %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func kvDecisions(rt *Router, d kvindex.Decision) float64 {
|
||||
return testutil.ToFloat64(rt.m.kvDecisions.WithLabelValues(string(d)))
|
||||
}
|
||||
|
||||
func TestRequestsGoToTheReplicaHoldingTheLongestBlockPrefix(t *testing.T) {
|
||||
rt, a, b := kvRouter(t, false)
|
||||
hs := kvevents.BlockHashes(fakeTokens(chatOf(longPrompt)...), 16, 0)
|
||||
if len(hs) < 6 {
|
||||
t.Fatalf("the prompt should make several blocks, made %d", len(hs))
|
||||
}
|
||||
// a holds the first two blocks — the opening a chat template shares —
|
||||
// and b the whole prompt.
|
||||
rt.ApplyKV(kvevents.Update{Replica: "a", BlockSize: 16, Stored: hs[:2]})
|
||||
rt.ApplyKV(kvevents.Update{Replica: "b", BlockSize: 16, Stored: hs})
|
||||
|
||||
for range 10 {
|
||||
sendChat(t, rt, longPrompt)
|
||||
}
|
||||
if b.served.Load() != 10 || a.served.Load() != 0 {
|
||||
t.Fatalf("served a=%d b=%d; every request should go to the replica holding its blocks",
|
||||
a.served.Load(), b.served.Load())
|
||||
}
|
||||
if n := kvDecisions(rt, kvindex.Hit); n != 10 {
|
||||
t.Fatalf("kv_hit counted %v times, want 10", n)
|
||||
}
|
||||
|
||||
// b evicts everything past the first block; now a holds the longer
|
||||
// prefix, and the index follows the engines rather than history.
|
||||
rt.ApplyKV(kvevents.Update{Replica: "b", BlockSize: 16, Removed: hs[1:]})
|
||||
rt.ApplyKV(kvevents.Update{Replica: "a", BlockSize: 16, Stored: hs[2:]})
|
||||
for range 5 {
|
||||
sendChat(t, rt, longPrompt)
|
||||
}
|
||||
if a.served.Load() != 5 {
|
||||
t.Fatalf("served a=%d after b evicted; want 5", a.served.Load())
|
||||
}
|
||||
}
|
||||
|
||||
func TestTooShortAMatchLeavesTheChoiceToLoad(t *testing.T) {
|
||||
rt, _, _ := kvRouter(t, false)
|
||||
hs := kvevents.BlockHashes(fakeTokens(chatOf(longPrompt)...), 16, 0)
|
||||
rt.ApplyKV(kvevents.Update{Replica: "a", BlockSize: 16, Stored: hs[:1]}) // below min_blocks
|
||||
rt.ApplyKV(kvevents.Update{Replica: "b", BlockSize: 16})
|
||||
sendChat(t, rt, longPrompt)
|
||||
if kvDecisions(rt, kvindex.MissMinLoad) != 1 {
|
||||
t.Fatal("one shared block is coincidence, and should be counted as a miss")
|
||||
}
|
||||
}
|
||||
|
||||
func TestWithNoEventsTheTreeDecidesAndNothingIsTokenized(t *testing.T) {
|
||||
rt, a, b := kvRouter(t, true)
|
||||
before := testutil.ToFloat64(rt.m.prefixDecisions.WithLabelValues("cache_miss_min_load"))
|
||||
sendChat(t, rt, longPrompt)
|
||||
if a.tokenized.Load()+b.tokenized.Load() != 0 {
|
||||
t.Fatal("with nothing in the index a tokenize call is a round trip for nothing")
|
||||
}
|
||||
if kvDecisions(rt, kvindex.FallbackNoIndex) != 1 {
|
||||
t.Fatal("the fallback should be counted by its reason")
|
||||
}
|
||||
if testutil.ToFloat64(rt.m.prefixDecisions.WithLabelValues("cache_miss_min_load")) != before+1 {
|
||||
t.Fatal("tier 1 should have decided, and said so in its own metric")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAFailedTokenizeFallsBackAndTheRequestStillSucceeds(t *testing.T) {
|
||||
rt, a, b := kvRouter(t, true)
|
||||
a.failTok.Store(true)
|
||||
b.failTok.Store(true)
|
||||
rt.ApplyKV(kvevents.Update{Replica: "a", BlockSize: 16, Stored: []uint64{1}})
|
||||
rt.ApplyKV(kvevents.Update{Replica: "b", BlockSize: 16, Stored: []uint64{2}})
|
||||
sendChat(t, rt, longPrompt)
|
||||
if kvDecisions(rt, kvindex.FallbackTokenizeFailed) != 1 {
|
||||
t.Fatal("a failed tokenize should fall back to tier 1, counted as such")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTheTokenizeCallCarriesWhatShapesTheTokens(t *testing.T) {
|
||||
rt, a, b := kvRouter(t, false)
|
||||
rt.ApplyKV(kvevents.Update{Replica: "a", BlockSize: 16})
|
||||
rt.ApplyKV(kvevents.Update{Replica: "b", BlockSize: 16})
|
||||
rec := post(t, rt.Handler(), map[string]any{
|
||||
"model": "m", "messages": []map[string]any{{"role": "user", "content": "hi"}},
|
||||
"tools": []any{map[string]any{"type": "function", "function": map[string]any{"name": "f"}}},
|
||||
"temperature": 0.2, "max_tokens": 5,
|
||||
}, nil)
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatal(rec.Body.String())
|
||||
}
|
||||
var last []byte
|
||||
for _, f := range []*fakeVLLM{a, b} {
|
||||
if v, ok := f.lastTok.Load().([]byte); ok {
|
||||
last = v
|
||||
}
|
||||
}
|
||||
var got map[string]any
|
||||
if err := json.Unmarshal(last, &got); err != nil {
|
||||
t.Fatalf("tokenize body %q: %v", last, err)
|
||||
}
|
||||
if got["tools"] == nil || got["messages"] == nil || got["model"] != "m" {
|
||||
t.Fatalf("the chat template sees the tools, so must the tokenize call: %s", last)
|
||||
}
|
||||
if _, ok := got["temperature"]; ok {
|
||||
t.Fatalf("sampling settings do not change the tokens and should not be sent: %s", last)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnImageMakesARequestUnindexable(t *testing.T) {
|
||||
rt, a, b := kvRouter(t, false)
|
||||
rt.ApplyKV(kvevents.Update{Replica: "a", BlockSize: 16})
|
||||
rt.ApplyKV(kvevents.Update{Replica: "b", BlockSize: 16})
|
||||
rec := post(t, rt.Handler(), map[string]any{
|
||||
"model": "m", "messages": []map[string]any{{"role": "user", "content": []any{
|
||||
map[string]any{"type": "text", "text": "what is this"},
|
||||
map[string]any{"type": "image_url", "image_url": map[string]any{"url": "data:image/png;base64,AAAA"}},
|
||||
}}},
|
||||
}, nil)
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatal(rec.Body.String())
|
||||
}
|
||||
if a.tokenized.Load()+b.tokenized.Load() != 0 || kvDecisions(rt, kvindex.FallbackUnsupported) != 1 {
|
||||
t.Fatal("an image's blocks hash with its identity, which the index cannot know")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAModelWithoutKVEventsIsNeverTokenized(t *testing.T) {
|
||||
rt, a, _ := kvRouter(t, false)
|
||||
rt.ApplyKV(kvevents.Update{Replica: "o-1", BlockSize: 16})
|
||||
rec := post(t, rt.Handler(), map[string]any{
|
||||
"model": "other", "messages": []map[string]any{{"role": "user", "content": longPrompt}},
|
||||
}, nil)
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatal(rec.Body.String())
|
||||
}
|
||||
if a.tokenized.Load() != 0 {
|
||||
t.Fatal("the precise index is opt-in per model")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAReplicaThatLeavesIsForgottenByTheIndex(t *testing.T) {
|
||||
rt, _, _ := kvRouter(t, false)
|
||||
rt.ApplyKV(kvevents.Update{Replica: "a", BlockSize: 16, Stored: []uint64{1, 2}})
|
||||
rt.SetEndpoints(t.Context(), nil)
|
||||
if len(rt.kv.Stats()) != 0 {
|
||||
t.Fatal("a gone replica's cache went with it")
|
||||
}
|
||||
}
|
||||
|
||||
// BenchmarkTokenize is the round trip the precise index adds before placing
|
||||
// a request, against a fake engine on loopback: the HTTP cost alone, without
|
||||
// a real tokenizer behind it.
|
||||
func BenchmarkTokenize(b *testing.B) {
|
||||
f := &fakeVLLM{}
|
||||
mux := http.NewServeMux()
|
||||
toks := make([]int, 2048)
|
||||
mux.HandleFunc("POST /tokenize", func(w http.ResponseWriter, r *http.Request) {
|
||||
_, _ = io.Copy(io.Discard, r.Body)
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"tokens": toks})
|
||||
})
|
||||
f.srv = httptest.NewServer(mux)
|
||||
defer f.srv.Close()
|
||||
body, _ := kvindex.TokenizeBody("/v1/chat/completions",
|
||||
[]byte(`{"model":"m","messages":[{"role":"user","content":"`+strings.Repeat("x", 8000)+`"}]}`), "m")
|
||||
client := newKVClient()
|
||||
b.ResetTimer()
|
||||
for b.Loop() {
|
||||
got, err := kvindex.Tokenize(b.Context(), client, f.srv.URL, "", body)
|
||||
if err != nil || len(got) != 2048 {
|
||||
b.Fatal(err)
|
||||
}
|
||||
_ = kvevents.BlockHashes(got, 16, 0)
|
||||
}
|
||||
}
|
||||
@@ -7,6 +7,7 @@ import (
|
||||
|
||||
"github.com/prometheus/client_golang/prometheus"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router/kvindex"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router/prefix"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router/proxy"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router/scheduler"
|
||||
@@ -30,6 +31,10 @@ type metrics struct {
|
||||
admissionRejected *prometheus.CounterVec
|
||||
|
||||
prefixDecisions *prometheus.CounterVec
|
||||
// kvDecisions is the precise index's counterpart, and kvTokenize how
|
||||
// long the tokenize call it needs took, by outcome.
|
||||
kvDecisions *prometheus.CounterVec
|
||||
kvTokenize *prometheus.HistogramVec
|
||||
|
||||
// cacheRequests counts what the answer cache did, cacheBytes and
|
||||
// cacheEntries what it holds.
|
||||
@@ -49,6 +54,17 @@ func newMetrics() *metrics {
|
||||
Help: "Prefix routing decisions by branch: cache_hit, cache_miss_min_load, " +
|
||||
"load_balance_override, no_tree, no_healthy_worker.",
|
||||
}, []string{"decision"}),
|
||||
kvDecisions: prometheus.NewCounterVec(prometheus.CounterOpts{
|
||||
Name: "llmgw_kv_index_decisions_total",
|
||||
Help: "Precise KV index routing decisions, for models with kv_events, by branch: kv_hit, " +
|
||||
"kv_miss_min_load, load_balance_override, and fallback_no_index, fallback_unsupported, " +
|
||||
"fallback_tokenize_failed, where tier 1 decided instead.",
|
||||
}, []string{"decision"}),
|
||||
kvTokenize: prometheus.NewHistogramVec(prometheus.HistogramOpts{
|
||||
Name: "llmgw_kv_index_tokenize_seconds",
|
||||
Help: "Time the precise KV index waited for a replica to tokenize a request, by outcome: ok, timeout, error.",
|
||||
Buckets: []float64{0.0005, 0.001, 0.0025, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5},
|
||||
}, []string{"outcome"}),
|
||||
requests: prometheus.NewCounterVec(prometheus.CounterOpts{
|
||||
Name: "llmgw_requests_total",
|
||||
Help: "Requests handled, by model, pool and response status.",
|
||||
@@ -123,6 +139,7 @@ func (m *metrics) collectors() []prometheus.Collector {
|
||||
m.requests, m.duration, m.ttft, m.tokens,
|
||||
m.scorerValue, m.filtered, m.considered,
|
||||
m.admissionWait, m.admissionRejected, m.prefixDecisions,
|
||||
m.kvDecisions, m.kvTokenize,
|
||||
m.cacheRequests, m.cacheBytes, m.cacheEntries,
|
||||
}, slices.Concat(m.tenant.collectors(), m.guard.collectors(), m.token.collectors())...)
|
||||
}
|
||||
@@ -135,6 +152,11 @@ func (m *metrics) observePrefix(d prefix.Decision) {
|
||||
m.prefixDecisions.WithLabelValues(string(d)).Inc()
|
||||
}
|
||||
|
||||
// observeKV records which branch a precise index decision took.
|
||||
func (m *metrics) observeKV(d kvindex.Decision) {
|
||||
m.kvDecisions.WithLabelValues(string(d)).Inc()
|
||||
}
|
||||
|
||||
func (m *metrics) observeAdmission(band, model string, wait time.Duration) {
|
||||
if band == "" {
|
||||
band = "none"
|
||||
|
||||
@@ -94,6 +94,14 @@ func Choose(t *Tree, prompt string, eps []endpoint.Snapshot, th Thresholds) (int
|
||||
return -1, MissMinLoad
|
||||
}
|
||||
|
||||
// Imbalanced reports whether load is lopsided enough to ignore affinity. The
|
||||
// precise index (tier 2) asks it too: knowing for certain that a replica holds
|
||||
// a prefix does not make it able to serve one more request quickly.
|
||||
func Imbalanced(eps []endpoint.Snapshot, th Thresholds) bool {
|
||||
th.normalize()
|
||||
return imbalanced(eps, th)
|
||||
}
|
||||
|
||||
// imbalanced reports whether load is lopsided enough to ignore affinity.
|
||||
func imbalanced(eps []endpoint.Snapshot, th Thresholds) bool {
|
||||
if len(eps) < 2 {
|
||||
|
||||
+57
-30
@@ -33,6 +33,7 @@ import (
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router/admission"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router/breaker"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router/kvindex"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router/prefix"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router/proxy"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router/scheduler"
|
||||
@@ -132,8 +133,14 @@ type Router struct {
|
||||
ownsThreads bool
|
||||
recorder *accounting.Recorder
|
||||
|
||||
// prefix is nil unless approximate prefix routing is turned on.
|
||||
// prefix scores by prefix: the precise index for models with kv_events,
|
||||
// the approximate tree when router.prefix is on. Its tree is nil when it
|
||||
// is off.
|
||||
prefix *scheduler.PrefixScorer
|
||||
// kv is the precise index (plan §7.5 tier 2), fed by the control plane
|
||||
// from the agents' KV cache events; kvClient makes its tokenize calls.
|
||||
kv *kvindex.Index
|
||||
kvClient *http.Client
|
||||
|
||||
// clusterView is nil until the control plane hands it over, and stays nil
|
||||
// in a router-only deployment.
|
||||
@@ -222,22 +229,30 @@ func New(cfg config.Server, log *slog.Logger) (*Router, error) {
|
||||
metrics := newMetrics()
|
||||
metrics.tenant.level = cfg.Router.TenantLabel
|
||||
metrics.tenant.setTagLabels(cfg.Router.TagLabels)
|
||||
var prefixScorer *scheduler.PrefixScorer
|
||||
// The scorer exists whether or not the tree does: which models use the
|
||||
// precise index is the catalog's to say, and the catalog changes while
|
||||
// the router runs. With neither in use it scores nothing and records
|
||||
// nothing.
|
||||
var tree *prefix.Tree
|
||||
if cfg.Router.Prefix.Enabled {
|
||||
prefixScorer = scheduler.NewPrefixScorer(
|
||||
prefix.New(prefix.Config{
|
||||
MaxNodes: cfg.Router.Prefix.MaxNodes,
|
||||
MinMatch: cfg.Router.Prefix.MinMatch,
|
||||
TTL: cfg.Router.Prefix.TTL,
|
||||
}),
|
||||
prefix.Thresholds{
|
||||
MatchRatio: cfg.Router.Prefix.MatchRatio,
|
||||
BalanceAbs: cfg.Router.Prefix.BalanceAbs,
|
||||
BalanceRel: cfg.Router.Prefix.BalanceRel,
|
||||
},
|
||||
metrics.observePrefix)
|
||||
sc.Prefix = prefixScorer
|
||||
tree = prefix.New(prefix.Config{
|
||||
MaxNodes: cfg.Router.Prefix.MaxNodes,
|
||||
MinMatch: cfg.Router.Prefix.MinMatch,
|
||||
TTL: cfg.Router.Prefix.TTL,
|
||||
})
|
||||
}
|
||||
minBlocks := cfg.Router.KVIndex.MinBlocks
|
||||
if minBlocks <= 0 {
|
||||
minBlocks = 4
|
||||
}
|
||||
prefixScorer := scheduler.NewPrefixScorer(tree,
|
||||
prefix.Thresholds{
|
||||
MatchRatio: cfg.Router.Prefix.MatchRatio,
|
||||
BalanceAbs: cfg.Router.Prefix.BalanceAbs,
|
||||
BalanceRel: cfg.Router.Prefix.BalanceRel,
|
||||
},
|
||||
metrics.observePrefix).WithKV(minBlocks, metrics.observeKV)
|
||||
sc.Prefix = prefixScorer
|
||||
|
||||
profiles := map[string]*scheduler.Profile{}
|
||||
for _, name := range []string{"chat", "batch", "long-context", "debug"} {
|
||||
@@ -255,6 +270,8 @@ func New(cfg config.Server, log *slog.Logger) (*Router, error) {
|
||||
reg: reg,
|
||||
profiles: profiles,
|
||||
prefix: prefixScorer,
|
||||
kv: kvindex.New(),
|
||||
kvClient: newKVClient(),
|
||||
m: metrics,
|
||||
px: proxy.New(proxy.Config{
|
||||
DialTimeout: cfg.Router.Timeouts.Dial,
|
||||
@@ -713,21 +730,20 @@ func (r *Router) ModelLoad() map[string]reconciler.Load {
|
||||
}
|
||||
|
||||
func (r *Router) SetEndpoints(ctx context.Context, infos []endpoint.Info) {
|
||||
if r.prefix != nil {
|
||||
// A replica that has gone takes its cache with it. Leaving its entries
|
||||
// in the tree would steer traffic towards something that is not there,
|
||||
// and since nothing reuses a replica id the entries would never be
|
||||
// corrected — only expire, minutes later.
|
||||
before := map[string]struct{}{}
|
||||
for _, s := range r.reg.Snapshots("") {
|
||||
before[s.ID] = struct{}{}
|
||||
}
|
||||
for _, i := range infos {
|
||||
delete(before, i.ID)
|
||||
}
|
||||
for id := range before {
|
||||
r.prefix.Forget(id)
|
||||
}
|
||||
// A replica that has gone takes its cache with it. Leaving its entries
|
||||
// in the tree or the index would steer traffic towards something that is
|
||||
// not there, and since nothing reuses a replica id the entries would never
|
||||
// be corrected — only expire, minutes later, or never.
|
||||
before := map[string]struct{}{}
|
||||
for _, s := range r.reg.Snapshots("") {
|
||||
before[s.ID] = struct{}{}
|
||||
}
|
||||
for _, i := range infos {
|
||||
delete(before, i.ID)
|
||||
}
|
||||
for id := range before {
|
||||
r.prefix.Forget(id)
|
||||
r.kv.Forget(id)
|
||||
}
|
||||
|
||||
r.reg.Replace(infos)
|
||||
@@ -917,6 +933,17 @@ func (r *Router) dispatchOne(ctx context.Context, w http.ResponseWriter, req *pa
|
||||
MaxTokens: req.MaxTokens,
|
||||
},
|
||||
}
|
||||
// Only where a profile publishes events, and only for a profile with a
|
||||
// prefix scorer to read the answer: the batch profile does not care where
|
||||
// a prefix is cached, and should not pay a round trip to find out.
|
||||
if entry.KVIndex() && profile.Uses("prefix") {
|
||||
_, kspan := tracer().Start(ctx, "kv index")
|
||||
sctx.Req.KV = r.kvQuery(ctx, req, model, candidates)
|
||||
if sctx.Req.KV.Fallback != "" {
|
||||
kspan.SetAttributes(attribute.String("llmgw.kv_fallback", string(sctx.Req.KV.Fallback)))
|
||||
}
|
||||
kspan.End()
|
||||
}
|
||||
|
||||
attempts := max(r.cfg.Retries+1, 1)
|
||||
var last dispatchResult
|
||||
|
||||
@@ -0,0 +1,42 @@
|
||||
package scheduler
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router/kvindex"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router/prefix"
|
||||
)
|
||||
|
||||
func kvScorer(got *[]kvindex.Decision) *PrefixScorer {
|
||||
return NewPrefixScorer(nil, prefix.DefaultThresholds(), nil).
|
||||
WithKV(2, func(d kvindex.Decision) { *got = append(*got, d) })
|
||||
}
|
||||
|
||||
func TestReplicasAreScoredByTheShareOfThePromptTheyHold(t *testing.T) {
|
||||
var got []kvindex.Decision
|
||||
q := &KVQuery{Matched: map[string]int{"a": 2, "b": 8}, Total: map[string]int{"a": 8, "b": 8}}
|
||||
s := kvScorer(&got).Score(newCtx(Request{KV: q}), []endpoint.Snapshot{ep("a"), ep("b"), ep("c")})
|
||||
if s[0] != 0.25 || s[1] != 1 || s[2] != 0 || got[0] != kvindex.Hit {
|
||||
t.Fatalf("scores %v, decisions %v", s, got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestACertainHitIsStillIgnoredWhenLoadIsLopsided(t *testing.T) {
|
||||
var got []kvindex.Decision
|
||||
q := &KVQuery{Matched: map[string]int{"a": 8}, Total: map[string]int{"a": 8, "b": 8}}
|
||||
busy := ep("a", func(s *endpoint.Snapshot) { s.Stats.Waiting = 40 })
|
||||
s := kvScorer(&got).Score(newCtx(Request{KV: q}), []endpoint.Snapshot{busy, ep("b")})
|
||||
if s[0] != 0 || got[0] != kvindex.LoadOverride {
|
||||
t.Fatalf("scores %v, decisions %v", s, got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAFallbackWithTier1OffScoresNothing(t *testing.T) {
|
||||
var got []kvindex.Decision
|
||||
q := &KVQuery{Fallback: kvindex.FallbackNoIndex}
|
||||
s := kvScorer(&got).Score(newCtx(Request{KV: q, PromptPrefix: "x"}), []endpoint.Snapshot{ep("a"), ep("b")})
|
||||
if s[0] != 0 || s[1] != 0 || len(got) != 1 || got[0] != kvindex.FallbackNoIndex {
|
||||
t.Fatalf("scores %v, decisions %v", s, got)
|
||||
}
|
||||
}
|
||||
@@ -2,11 +2,29 @@ package scheduler
|
||||
|
||||
import (
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router/kvindex"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/router/prefix"
|
||||
)
|
||||
|
||||
// PrefixScorer steers a request towards the replica that has probably already
|
||||
// prefilled its prefix (plan §7.5 tier 1).
|
||||
// KVQuery is the precise index's answer for one request, worked out before
|
||||
// scheduling because it needs a round trip the scorer must not make.
|
||||
type KVQuery struct {
|
||||
// Fallback, when set, is why the index could not be asked; tier 1
|
||||
// decides instead.
|
||||
Fallback kvindex.Decision
|
||||
// Matched is how many of the prompt's leading blocks each replica the
|
||||
// index knows holds, and Total how many blocks the prompt has at that
|
||||
// replica's block size. A replica absent from both is one the index has
|
||||
// not heard from, and is scored as holding nothing — an under-estimate,
|
||||
// which costs a missed preference and never a wrong one.
|
||||
Matched map[string]int
|
||||
Total map[string]int
|
||||
}
|
||||
|
||||
// PrefixScorer steers a request towards the replica that has already
|
||||
// prefilled its prefix: from the precise KV index where the model uses it and
|
||||
// it has an answer (plan §7.5 tier 2), and otherwise from the approximate tree
|
||||
// (tier 1).
|
||||
//
|
||||
// A scorer rather than a filter, deliberately. A filter would make affinity a
|
||||
// rule, and a rule that sends every request carrying a popular system prompt to
|
||||
@@ -20,23 +38,56 @@ type PrefixScorer struct {
|
||||
// observe reports which branch the decision took, for the metric without
|
||||
// which tuning this is guesswork.
|
||||
observe func(prefix.Decision)
|
||||
|
||||
// observeKV and minBlocks are the precise index's: its decision metric,
|
||||
// and the fewest leading blocks worth preferring a replica for.
|
||||
observeKV func(kvindex.Decision)
|
||||
minBlocks int
|
||||
}
|
||||
|
||||
// NewPrefixScorer builds the scorer. A nil tree makes it inert, which is what
|
||||
// a model with prefix routing turned off should get.
|
||||
// NewPrefixScorer builds the scorer. A nil tree turns tier 1 off: requests
|
||||
// the precise index cannot answer are then left to load, and nothing is
|
||||
// recorded for them.
|
||||
func NewPrefixScorer(t *prefix.Tree, th prefix.Thresholds, observe func(prefix.Decision)) *PrefixScorer {
|
||||
if observe == nil {
|
||||
observe = func(prefix.Decision) {}
|
||||
}
|
||||
return &PrefixScorer{tree: t, thresholds: th, observe: observe}
|
||||
return &PrefixScorer{tree: t, thresholds: th, observe: observe,
|
||||
observeKV: func(kvindex.Decision) {}, minBlocks: 1}
|
||||
}
|
||||
|
||||
// WithKV sets what the precise index needs: the fewest leading blocks a
|
||||
// replica must hold to be preferred (below one means one), and where its
|
||||
// decisions are counted.
|
||||
func (p *PrefixScorer) WithKV(minBlocks int, observe func(kvindex.Decision)) *PrefixScorer {
|
||||
if observe != nil {
|
||||
p.observeKV = observe
|
||||
}
|
||||
p.minBlocks = max(minBlocks, 1)
|
||||
return p
|
||||
}
|
||||
|
||||
// Name identifies the scorer in weights and metrics.
|
||||
func (*PrefixScorer) Name() string { return "prefix" }
|
||||
|
||||
// Score gives the matched replica 1 and everything else 0.
|
||||
// Score rates each candidate.
|
||||
//
|
||||
// From the index: the share of the prompt's blocks the replica holds, so two
|
||||
// replicas holding a shared system prompt score alike and the one also
|
||||
// holding this conversation scores higher. From the tree: 1 for the replica
|
||||
// it names and 0 for the rest, because a guess has no degrees.
|
||||
func (p *PrefixScorer) Score(c *Ctx, eps []endpoint.Snapshot) []float64 {
|
||||
out := make([]float64, len(eps))
|
||||
if q := c.Req.KV; q != nil {
|
||||
if q.Fallback == "" {
|
||||
p.observeKV(p.scoreKV(q, eps, out))
|
||||
return out
|
||||
}
|
||||
p.observeKV(q.Fallback)
|
||||
}
|
||||
if p.tree == nil {
|
||||
return out
|
||||
}
|
||||
idx, decision := prefix.Choose(p.tree, c.Req.PromptPrefix, eps, p.thresholds)
|
||||
p.observe(decision)
|
||||
if idx >= 0 {
|
||||
@@ -45,12 +96,38 @@ func (p *PrefixScorer) Score(c *Ctx, eps []endpoint.Snapshot) []float64 {
|
||||
return out
|
||||
}
|
||||
|
||||
// scoreKV fills out from the index's answer and says which branch it took.
|
||||
func (p *PrefixScorer) scoreKV(q *KVQuery, eps []endpoint.Snapshot, out []float64) kvindex.Decision {
|
||||
if len(eps) == 0 {
|
||||
return kvindex.MissMinLoad
|
||||
}
|
||||
// The guard first, as in tier 1: a certain hit on a replica that cannot
|
||||
// take the request quickly is still the wrong place for it.
|
||||
if prefix.Imbalanced(eps, p.thresholds) {
|
||||
return kvindex.LoadOverride
|
||||
}
|
||||
best := 0
|
||||
for i, e := range eps {
|
||||
m, t := q.Matched[e.ID], q.Total[e.ID]
|
||||
if t > 0 {
|
||||
out[i] = float64(m) / float64(t)
|
||||
}
|
||||
best = max(best, m)
|
||||
}
|
||||
if best < p.minBlocks {
|
||||
clear(out)
|
||||
return kvindex.MissMinLoad
|
||||
}
|
||||
return kvindex.Hit
|
||||
}
|
||||
|
||||
// Observe records where a request was actually sent.
|
||||
//
|
||||
// After the decision rather than during it: the scorer proposes, the picker
|
||||
// disposes, and recording the proposal would fill the tree with replicas that
|
||||
// were never sent anything — which is the one thing that would make the
|
||||
// estimate worse than no estimate at all.
|
||||
// estimate worse than no estimate at all. The precise index needs no such
|
||||
// record; the engine reports what it cached.
|
||||
func (p *PrefixScorer) Observe(promptPrefix, replicaID string) {
|
||||
if p.tree != nil {
|
||||
p.tree.Observe(promptPrefix, replicaID)
|
||||
|
||||
@@ -30,6 +30,9 @@ type Request struct {
|
||||
PromptTokens int
|
||||
// MaxTokens is the requested generation length.
|
||||
MaxTokens int
|
||||
// KV is what the precise KV index says about this request (plan §7.5
|
||||
// tier 2), nil for a model that does not use it.
|
||||
KV *KVQuery
|
||||
}
|
||||
|
||||
// Ctx is the context of a single scheduling decision.
|
||||
@@ -77,6 +80,16 @@ type Profile struct {
|
||||
StalePenalty float64
|
||||
}
|
||||
|
||||
// Uses reports whether the profile has a scorer of this name with a weight.
|
||||
func (p *Profile) Uses(scorer string) bool {
|
||||
for _, ws := range p.Scorers {
|
||||
if ws.Scorer != nil && ws.Weight != 0 && ws.Scorer.Name() == scorer {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// Decision is a routing choice together with its rationale.
|
||||
//
|
||||
// Recording the rationale is not decoration: tuning the prefix and queue
|
||||
|
||||
@@ -0,0 +1,42 @@
|
||||
# C24 — precise KV-index routing (plan §7.5 tier 2)
|
||||
|
||||
## CHANGELOG
|
||||
|
||||
### Routing by what each vLLM replica's cache holds
|
||||
|
||||
- A vLLM profile can set `kv_events: true`. The agent then starts the engine
|
||||
publishing its KV cache events (`--kv-events-config`) to a port on
|
||||
127.0.0.1 the agent listens on, with `--prefix-caching-hash-algo
|
||||
sha256_cbor` and `PYTHONHASHSEED=1729`, so its block hashes can be
|
||||
computed outside Python. The events never leave the node.
|
||||
- The agent mirrors each engine's set of cached blocks and sends the net
|
||||
changes up its existing stream to the control plane four times a second
|
||||
(new `KVEvents` agent message), starting every new stream with a full
|
||||
copy. Lost events (a gap in vLLM's sequence numbers), an engine restart or
|
||||
`AllBlocksCleared` start a replica's index again; blocks offloaded to CPU,
|
||||
remote ones and a hybrid model's second cache group are not indexed.
|
||||
- The control plane passes them to the router's index, only for replicas
|
||||
the sending node runs, and empties what a stream described when it closes.
|
||||
- For a model with `kv_events`, the router asks one of its replicas to
|
||||
tokenize the request (vLLM's `/tokenize`, with messages and tools, so the
|
||||
chat template is applied exactly), hashes the tokens into blocks as vLLM
|
||||
does, and scores replicas by the share of the prompt's leading blocks they
|
||||
hold. The imbalance guard still applies. With no answer — no events yet,
|
||||
an image, a cache salt, a LoRA adapter, a failed or slow tokenize — it falls
|
||||
back to the approximate tree (`router.prefix`) or to load.
|
||||
- The agent checks first blocks against llmgw's hashing; an engine that
|
||||
hashes another way (a profile overriding the algorithm, a changed vLLM) is
|
||||
reported in the agent's log and `llmgw_kv_index_problem`, and routed as
|
||||
though it published nothing.
|
||||
- New settings `router.kv_index.tokenize_timeout` (250ms), `min_blocks` (4)
|
||||
and `max_blocks` (4096). New metrics `llmgw_kv_index_decisions_total`
|
||||
(`kv_hit`, `kv_miss_min_load`, `load_balance_override`,
|
||||
`fallback_no_index`, `fallback_unsupported`, `fallback_tokenize_failed`),
|
||||
`llmgw_kv_index_tokenize_seconds`, `llmgw_kv_index_blocks`.
|
||||
- Experimental: tested against vLLM's event format and hashing as its source
|
||||
defines them, and against libzmq and msgspec publishing — not yet against a
|
||||
running vLLM.
|
||||
|
||||
## BACKLOG
|
||||
|
||||
- [~] 24. Precise KV index (tier 2) — built: vLLM KV events via the agent (loopback ZMQ, pure-Go zmq4, msgpack in both of vLLM's encodings) forwarded as net block-hash changes over the agent stream; router index, tokenize via the replica's `/tokenize`, sha256_cbor block-hash chain checked against cbor2/hashlib golden values; scorer by leading-block share with tier 1 fallback; opt-in `kv_events` per profile; metrics; docs. Left: validation against a real vLLM (event shape, hash agreement, `/tokenize` matching the served prompt, cost of tokenizing); vLLM's replay endpoint instead of resetting on a gap; data-parallel ranks (one publisher each, ports offset); hybrid models' cache groups and CPU-offloaded blocks; LoRA/multimodal/cache_salt requests (extra keys); a router-only deployment (no agents) subscribing directly; standby control planes have no index until they lead; a per-replica fallback when only some candidates have an index; caching tokenizations of repeated prompts.
|
||||
Reference in new issue
Block a user