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>
104 lines
2.8 KiB
Go
104 lines
2.8 KiB
Go
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,
|
|
}}}
|
|
}
|