Ilker ManapandClaude Opus 5.5
117f417ef5
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>
87 lines
2.5 KiB
Go
87 lines
2.5 KiB
Go
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:
|
|
}
|
|
}
|
|
}
|