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
3.2 KiB
Go
104 lines
3.2 KiB
Go
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")
|
|
}
|
|
}
|