From 117f417ef503be479e0459be387e98cc56cc4220 Mon Sep 17 00:00:00 2001 From: Ilker Manap Date: Mon, 5 Oct 2026 01:21:11 +0200 Subject: [PATCH] Precise KV-index routing: vLLM KV cache events through the agent, routed by leading blocks held MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- api/gen/llmgw/v1/agent.pb.go | 594 ++++++++++++++-------- api/proto/llmgw/v1/agent.proto | 30 ++ docs/admin-api.openapi.json | 4 + docs/install.md | 59 +++ docs/reference.md | 59 +++ go.mod | 4 + go.sum | 8 + internal/agent/agent.go | 20 + internal/agent/engines/engine.go | 8 + internal/agent/engines/kvevents_test.go | 47 ++ internal/agent/engines/vllm.go | 11 + internal/agent/kv.go | 86 ++++ internal/agent/kv_test.go | 103 ++++ internal/cli/server.go | 6 + internal/config/config.go | 21 + internal/control/agentsvc/convert.go | 2 + internal/control/agentsvc/kv.go | 103 ++++ internal/control/agentsvc/kv_test.go | 29 ++ internal/control/agentsvc/service.go | 16 + internal/control/catalog/catalog.go | 31 ++ internal/control/catalog/kvevents_test.go | 40 ++ internal/docgen/reference_gen.go | 7 + internal/kvevents/decode.go | 224 ++++++++ internal/kvevents/hash.go | 133 +++++ internal/kvevents/interop_test.go | 49 ++ internal/kvevents/kvevents_test.go | 386 ++++++++++++++ internal/kvevents/mirror.go | 308 +++++++++++ internal/kvevents/subscriber.go | 172 +++++++ internal/kvevents/testdata/pypub.py | 77 +++ internal/router/admin.go | 3 + internal/router/kvindex/decision.go | 33 ++ internal/router/kvindex/index.go | 132 +++++ internal/router/kvindex/kvindex_test.go | 67 +++ internal/router/kvindex/tokenize.go | 146 ++++++ internal/router/kvroute.go | 151 ++++++ internal/router/kvroute_test.go | 285 +++++++++++ internal/router/metrics.go | 22 + internal/router/prefix/scorer.go | 8 + internal/router/router.go | 87 ++-- internal/router/scheduler/kv_test.go | 42 ++ internal/router/scheduler/prefix.go | 91 +++- internal/router/scheduler/scheduler.go | 13 + notes/c24-kvindex.md | 42 ++ 43 files changed, 3496 insertions(+), 263 deletions(-) create mode 100644 internal/agent/engines/kvevents_test.go create mode 100644 internal/agent/kv.go create mode 100644 internal/agent/kv_test.go create mode 100644 internal/control/agentsvc/kv.go create mode 100644 internal/control/agentsvc/kv_test.go create mode 100644 internal/control/catalog/kvevents_test.go create mode 100644 internal/kvevents/decode.go create mode 100644 internal/kvevents/hash.go create mode 100644 internal/kvevents/interop_test.go create mode 100644 internal/kvevents/kvevents_test.go create mode 100644 internal/kvevents/mirror.go create mode 100644 internal/kvevents/subscriber.go create mode 100644 internal/kvevents/testdata/pypub.py create mode 100644 internal/router/kvindex/decision.go create mode 100644 internal/router/kvindex/index.go create mode 100644 internal/router/kvindex/kvindex_test.go create mode 100644 internal/router/kvindex/tokenize.go create mode 100644 internal/router/kvroute.go create mode 100644 internal/router/kvroute_test.go create mode 100644 internal/router/scheduler/kv_test.go create mode 100644 notes/c24-kvindex.md diff --git a/api/gen/llmgw/v1/agent.pb.go b/api/gen/llmgw/v1/agent.pb.go index 42bd922..5365209 100644 --- a/api/gen/llmgw/v1/agent.pb.go +++ b/api/gen/llmgw/v1/agent.pb.go @@ -90,7 +90,7 @@ func (x ShardStatus_State) Number() protoreflect.EnumNumber { // Deprecated: Use ShardStatus_State.Descriptor instead. func (ShardStatus_State) EnumDescriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{22, 0} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{23, 0} } // AgentMessage is anything a node tells the control plane. @@ -105,6 +105,7 @@ type AgentMessage struct { // *AgentMessage_FabricScan // *AgentMessage_FabricTest // *AgentMessage_Upgrade + // *AgentMessage_KvEvents Msg isAgentMessage_Msg `protobuf_oneof:"msg"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache @@ -210,6 +211,15 @@ func (x *AgentMessage) GetUpgrade() *UpgradeReply { return nil } +func (x *AgentMessage) GetKvEvents() *KVEvents { + if x != nil { + if x, ok := x.Msg.(*AgentMessage_KvEvents); ok { + return x.KvEvents + } + } + return nil +} + type isAgentMessage_Msg interface { isAgentMessage_Msg() } @@ -242,6 +252,10 @@ type AgentMessage_Upgrade struct { Upgrade *UpgradeReply `protobuf:"bytes,7,opt,name=upgrade,proto3,oneof"` } +type AgentMessage_KvEvents struct { + KvEvents *KVEvents `protobuf:"bytes,8,opt,name=kv_events,json=kvEvents,proto3,oneof"` +} + func (*AgentMessage_Register) isAgentMessage_Msg() {} func (*AgentMessage_Heartbeat) isAgentMessage_Msg() {} @@ -256,6 +270,109 @@ func (*AgentMessage_FabricTest) isAgentMessage_Msg() {} func (*AgentMessage_Upgrade) isAgentMessage_Msg() {} +func (*AgentMessage_KvEvents) isAgentMessage_Msg() {} + +// 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. +type KVEvents struct { + state protoimpl.MessageState `protogen:"open.v1"` + ReplicaId string `protobuf:"bytes,1,opt,name=replica_id,json=replicaId,proto3" json:"replica_id,omitempty"` + // block_size is how many tokens make one of the replica's blocks. + BlockSize int32 `protobuf:"varint,2,opt,name=block_size,json=blockSize,proto3" json:"block_size,omitempty"` + // 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. + ClearFirst bool `protobuf:"varint,3,opt,name=clear_first,json=clearFirst,proto3" json:"clear_first,omitempty"` + // stored and removed are block hashes as vLLM publishes them: the low 64 + // bits of the sha256_cbor digest. + Stored []uint64 `protobuf:"fixed64,4,rep,packed,name=stored,proto3" json:"stored,omitempty"` + Removed []uint64 `protobuf:"fixed64,5,rep,packed,name=removed,proto3" json:"removed,omitempty"` + // 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. + Problem string `protobuf:"bytes,6,opt,name=problem,proto3" json:"problem,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *KVEvents) Reset() { + *x = KVEvents{} + mi := &file_llmgw_v1_agent_proto_msgTypes[1] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *KVEvents) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*KVEvents) ProtoMessage() {} + +func (x *KVEvents) ProtoReflect() protoreflect.Message { + mi := &file_llmgw_v1_agent_proto_msgTypes[1] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use KVEvents.ProtoReflect.Descriptor instead. +func (*KVEvents) Descriptor() ([]byte, []int) { + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{1} +} + +func (x *KVEvents) GetReplicaId() string { + if x != nil { + return x.ReplicaId + } + return "" +} + +func (x *KVEvents) GetBlockSize() int32 { + if x != nil { + return x.BlockSize + } + return 0 +} + +func (x *KVEvents) GetClearFirst() bool { + if x != nil { + return x.ClearFirst + } + return false +} + +func (x *KVEvents) GetStored() []uint64 { + if x != nil { + return x.Stored + } + return nil +} + +func (x *KVEvents) GetRemoved() []uint64 { + if x != nil { + return x.Removed + } + return nil +} + +func (x *KVEvents) GetProblem() string { + if x != nil { + return x.Problem + } + return "" +} + // ControlMessage is anything the control plane tells a node. type ControlMessage struct { state protoimpl.MessageState `protogen:"open.v1"` @@ -285,7 +402,7 @@ type ControlMessage struct { func (x *ControlMessage) Reset() { *x = ControlMessage{} - mi := &file_llmgw_v1_agent_proto_msgTypes[1] + mi := &file_llmgw_v1_agent_proto_msgTypes[2] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -297,7 +414,7 @@ func (x *ControlMessage) String() string { func (*ControlMessage) ProtoMessage() {} func (x *ControlMessage) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[1] + mi := &file_llmgw_v1_agent_proto_msgTypes[2] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -310,7 +427,7 @@ func (x *ControlMessage) ProtoReflect() protoreflect.Message { // Deprecated: Use ControlMessage.ProtoReflect.Descriptor instead. func (*ControlMessage) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{1} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{2} } func (x *ControlMessage) GetMsg() isControlMessage_Msg { @@ -509,7 +626,7 @@ type ControlPlanes struct { func (x *ControlPlanes) Reset() { *x = ControlPlanes{} - mi := &file_llmgw_v1_agent_proto_msgTypes[2] + mi := &file_llmgw_v1_agent_proto_msgTypes[3] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -521,7 +638,7 @@ func (x *ControlPlanes) String() string { func (*ControlPlanes) ProtoMessage() {} func (x *ControlPlanes) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[2] + mi := &file_llmgw_v1_agent_proto_msgTypes[3] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -534,7 +651,7 @@ func (x *ControlPlanes) ProtoReflect() protoreflect.Message { // Deprecated: Use ControlPlanes.ProtoReflect.Descriptor instead. func (*ControlPlanes) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{2} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{3} } func (x *ControlPlanes) GetServers() []string { @@ -634,7 +751,7 @@ type Register struct { func (x *Register) Reset() { *x = Register{} - mi := &file_llmgw_v1_agent_proto_msgTypes[3] + mi := &file_llmgw_v1_agent_proto_msgTypes[4] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -646,7 +763,7 @@ func (x *Register) String() string { func (*Register) ProtoMessage() {} func (x *Register) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[3] + mi := &file_llmgw_v1_agent_proto_msgTypes[4] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -659,7 +776,7 @@ func (x *Register) ProtoReflect() protoreflect.Message { // Deprecated: Use Register.ProtoReflect.Descriptor instead. func (*Register) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{3} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{4} } func (x *Register) GetNodeId() string { @@ -811,7 +928,7 @@ type Host struct { func (x *Host) Reset() { *x = Host{} - mi := &file_llmgw_v1_agent_proto_msgTypes[4] + mi := &file_llmgw_v1_agent_proto_msgTypes[5] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -823,7 +940,7 @@ func (x *Host) String() string { func (*Host) ProtoMessage() {} func (x *Host) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[4] + mi := &file_llmgw_v1_agent_proto_msgTypes[5] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -836,7 +953,7 @@ func (x *Host) ProtoReflect() protoreflect.Message { // Deprecated: Use Host.ProtoReflect.Descriptor instead. func (*Host) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{4} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{5} } func (x *Host) GetOsId() string { @@ -903,7 +1020,7 @@ type Storage struct { func (x *Storage) Reset() { *x = Storage{} - mi := &file_llmgw_v1_agent_proto_msgTypes[5] + mi := &file_llmgw_v1_agent_proto_msgTypes[6] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -915,7 +1032,7 @@ func (x *Storage) String() string { func (*Storage) ProtoMessage() {} func (x *Storage) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[5] + mi := &file_llmgw_v1_agent_proto_msgTypes[6] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -928,7 +1045,7 @@ func (x *Storage) ProtoReflect() protoreflect.Message { // Deprecated: Use Storage.ProtoReflect.Descriptor instead. func (*Storage) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{5} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{6} } func (x *Storage) GetPath() string { @@ -995,7 +1112,7 @@ type RunningShard struct { func (x *RunningShard) Reset() { *x = RunningShard{} - mi := &file_llmgw_v1_agent_proto_msgTypes[6] + mi := &file_llmgw_v1_agent_proto_msgTypes[7] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1007,7 +1124,7 @@ func (x *RunningShard) String() string { func (*RunningShard) ProtoMessage() {} func (x *RunningShard) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[6] + mi := &file_llmgw_v1_agent_proto_msgTypes[7] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1020,7 +1137,7 @@ func (x *RunningShard) ProtoReflect() protoreflect.Message { // Deprecated: Use RunningShard.ProtoReflect.Descriptor instead. func (*RunningShard) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{6} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{7} } func (x *RunningShard) GetReplicaId() string { @@ -1173,7 +1290,7 @@ type GPU struct { func (x *GPU) Reset() { *x = GPU{} - mi := &file_llmgw_v1_agent_proto_msgTypes[7] + mi := &file_llmgw_v1_agent_proto_msgTypes[8] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1185,7 +1302,7 @@ func (x *GPU) String() string { func (*GPU) ProtoMessage() {} func (x *GPU) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[7] + mi := &file_llmgw_v1_agent_proto_msgTypes[8] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1198,7 +1315,7 @@ func (x *GPU) ProtoReflect() protoreflect.Message { // Deprecated: Use GPU.ProtoReflect.Descriptor instead. func (*GPU) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{7} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{8} } func (x *GPU) GetIndex() int32 { @@ -1374,7 +1491,7 @@ type MIGDevice struct { func (x *MIGDevice) Reset() { *x = MIGDevice{} - mi := &file_llmgw_v1_agent_proto_msgTypes[8] + mi := &file_llmgw_v1_agent_proto_msgTypes[9] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1386,7 +1503,7 @@ func (x *MIGDevice) String() string { func (*MIGDevice) ProtoMessage() {} func (x *MIGDevice) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[8] + mi := &file_llmgw_v1_agent_proto_msgTypes[9] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1399,7 +1516,7 @@ func (x *MIGDevice) ProtoReflect() protoreflect.Message { // Deprecated: Use MIGDevice.ProtoReflect.Descriptor instead. func (*MIGDevice) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{8} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{9} } func (x *MIGDevice) GetParentIndex() int32 { @@ -1465,7 +1582,7 @@ type FabricAttachment struct { func (x *FabricAttachment) Reset() { *x = FabricAttachment{} - mi := &file_llmgw_v1_agent_proto_msgTypes[9] + mi := &file_llmgw_v1_agent_proto_msgTypes[10] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1477,7 +1594,7 @@ func (x *FabricAttachment) String() string { func (*FabricAttachment) ProtoMessage() {} func (x *FabricAttachment) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[9] + mi := &file_llmgw_v1_agent_proto_msgTypes[10] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1490,7 +1607,7 @@ func (x *FabricAttachment) ProtoReflect() protoreflect.Message { // Deprecated: Use FabricAttachment.ProtoReflect.Descriptor instead. func (*FabricAttachment) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{9} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{10} } func (x *FabricAttachment) GetFabric() string { @@ -1532,7 +1649,7 @@ type LocalWeights struct { func (x *LocalWeights) Reset() { *x = LocalWeights{} - mi := &file_llmgw_v1_agent_proto_msgTypes[10] + mi := &file_llmgw_v1_agent_proto_msgTypes[11] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1544,7 +1661,7 @@ func (x *LocalWeights) String() string { func (*LocalWeights) ProtoMessage() {} func (x *LocalWeights) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[10] + mi := &file_llmgw_v1_agent_proto_msgTypes[11] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1557,7 +1674,7 @@ func (x *LocalWeights) ProtoReflect() protoreflect.Message { // Deprecated: Use LocalWeights.ProtoReflect.Descriptor instead. func (*LocalWeights) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{10} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{11} } func (x *LocalWeights) GetPath() string { @@ -1613,7 +1730,7 @@ type NodeLoad struct { func (x *NodeLoad) Reset() { *x = NodeLoad{} - mi := &file_llmgw_v1_agent_proto_msgTypes[11] + mi := &file_llmgw_v1_agent_proto_msgTypes[12] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1625,7 +1742,7 @@ func (x *NodeLoad) String() string { func (*NodeLoad) ProtoMessage() {} func (x *NodeLoad) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[11] + mi := &file_llmgw_v1_agent_proto_msgTypes[12] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1638,7 +1755,7 @@ func (x *NodeLoad) ProtoReflect() protoreflect.Message { // Deprecated: Use NodeLoad.ProtoReflect.Descriptor instead. func (*NodeLoad) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{11} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{12} } func (x *NodeLoad) GetCpuPercent() float64 { @@ -1771,7 +1888,7 @@ type GPUSensor struct { func (x *GPUSensor) Reset() { *x = GPUSensor{} - mi := &file_llmgw_v1_agent_proto_msgTypes[12] + mi := &file_llmgw_v1_agent_proto_msgTypes[13] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -1783,7 +1900,7 @@ func (x *GPUSensor) String() string { func (*GPUSensor) ProtoMessage() {} func (x *GPUSensor) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[12] + mi := &file_llmgw_v1_agent_proto_msgTypes[13] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -1796,7 +1913,7 @@ func (x *GPUSensor) ProtoReflect() protoreflect.Message { // Deprecated: Use GPUSensor.ProtoReflect.Descriptor instead. func (*GPUSensor) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{12} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{13} } func (x *GPUSensor) GetIndex() int32 { @@ -2010,7 +2127,7 @@ type ShardLoad struct { func (x *ShardLoad) Reset() { *x = ShardLoad{} - mi := &file_llmgw_v1_agent_proto_msgTypes[13] + mi := &file_llmgw_v1_agent_proto_msgTypes[14] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -2022,7 +2139,7 @@ func (x *ShardLoad) String() string { func (*ShardLoad) ProtoMessage() {} func (x *ShardLoad) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[13] + mi := &file_llmgw_v1_agent_proto_msgTypes[14] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -2035,7 +2152,7 @@ func (x *ShardLoad) ProtoReflect() protoreflect.Message { // Deprecated: Use ShardLoad.ProtoReflect.Descriptor instead. func (*ShardLoad) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{13} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{14} } func (x *ShardLoad) GetShardId() string { @@ -2103,7 +2220,7 @@ type IBPort struct { func (x *IBPort) Reset() { *x = IBPort{} - mi := &file_llmgw_v1_agent_proto_msgTypes[14] + mi := &file_llmgw_v1_agent_proto_msgTypes[15] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -2115,7 +2232,7 @@ func (x *IBPort) String() string { func (*IBPort) ProtoMessage() {} func (x *IBPort) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[14] + mi := &file_llmgw_v1_agent_proto_msgTypes[15] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -2128,7 +2245,7 @@ func (x *IBPort) ProtoReflect() protoreflect.Message { // Deprecated: Use IBPort.ProtoReflect.Descriptor instead. func (*IBPort) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{14} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{15} } func (x *IBPort) GetHca() string { @@ -2271,7 +2388,7 @@ type SubnetManager struct { func (x *SubnetManager) Reset() { *x = SubnetManager{} - mi := &file_llmgw_v1_agent_proto_msgTypes[15] + mi := &file_llmgw_v1_agent_proto_msgTypes[16] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -2283,7 +2400,7 @@ func (x *SubnetManager) String() string { func (*SubnetManager) ProtoMessage() {} func (x *SubnetManager) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[15] + mi := &file_llmgw_v1_agent_proto_msgTypes[16] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -2296,7 +2413,7 @@ func (x *SubnetManager) ProtoReflect() protoreflect.Message { // Deprecated: Use SubnetManager.ProtoReflect.Descriptor instead. func (*SubnetManager) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{15} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{16} } func (x *SubnetManager) GetGuid() string { @@ -2339,7 +2456,7 @@ type IMEX struct { func (x *IMEX) Reset() { *x = IMEX{} - mi := &file_llmgw_v1_agent_proto_msgTypes[16] + mi := &file_llmgw_v1_agent_proto_msgTypes[17] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -2351,7 +2468,7 @@ func (x *IMEX) String() string { func (*IMEX) ProtoMessage() {} func (x *IMEX) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[16] + mi := &file_llmgw_v1_agent_proto_msgTypes[17] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -2364,7 +2481,7 @@ func (x *IMEX) ProtoReflect() protoreflect.Message { // Deprecated: Use IMEX.ProtoReflect.Descriptor instead. func (*IMEX) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{16} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{17} } func (x *IMEX) GetState() string { @@ -2396,7 +2513,7 @@ type GPUFabricState struct { func (x *GPUFabricState) Reset() { *x = GPUFabricState{} - mi := &file_llmgw_v1_agent_proto_msgTypes[17] + mi := &file_llmgw_v1_agent_proto_msgTypes[18] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -2408,7 +2525,7 @@ func (x *GPUFabricState) String() string { func (*GPUFabricState) ProtoMessage() {} func (x *GPUFabricState) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[17] + mi := &file_llmgw_v1_agent_proto_msgTypes[18] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -2421,7 +2538,7 @@ func (x *GPUFabricState) ProtoReflect() protoreflect.Message { // Deprecated: Use GPUFabricState.ProtoReflect.Descriptor instead. func (*GPUFabricState) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{17} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{18} } func (x *GPUFabricState) GetIndex() int32 { @@ -2536,7 +2653,7 @@ type Heartbeat struct { func (x *Heartbeat) Reset() { *x = Heartbeat{} - mi := &file_llmgw_v1_agent_proto_msgTypes[18] + mi := &file_llmgw_v1_agent_proto_msgTypes[19] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -2548,7 +2665,7 @@ func (x *Heartbeat) String() string { func (*Heartbeat) ProtoMessage() {} func (x *Heartbeat) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[18] + mi := &file_llmgw_v1_agent_proto_msgTypes[19] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -2561,7 +2678,7 @@ func (x *Heartbeat) ProtoReflect() protoreflect.Message { // Deprecated: Use Heartbeat.ProtoReflect.Descriptor instead. func (*Heartbeat) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{18} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{19} } func (x *Heartbeat) GetNodeId() string { @@ -2716,7 +2833,7 @@ type Provision struct { func (x *Provision) Reset() { *x = Provision{} - mi := &file_llmgw_v1_agent_proto_msgTypes[19] + mi := &file_llmgw_v1_agent_proto_msgTypes[20] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -2728,7 +2845,7 @@ func (x *Provision) String() string { func (*Provision) ProtoMessage() {} func (x *Provision) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[19] + mi := &file_llmgw_v1_agent_proto_msgTypes[20] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -2741,7 +2858,7 @@ func (x *Provision) ProtoReflect() protoreflect.Message { // Deprecated: Use Provision.ProtoReflect.Descriptor instead. func (*Provision) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{19} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{20} } func (x *Provision) GetEngines() []*EngineWant { @@ -2765,7 +2882,7 @@ type EngineWant struct { func (x *EngineWant) Reset() { *x = EngineWant{} - mi := &file_llmgw_v1_agent_proto_msgTypes[20] + mi := &file_llmgw_v1_agent_proto_msgTypes[21] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -2777,7 +2894,7 @@ func (x *EngineWant) String() string { func (*EngineWant) ProtoMessage() {} func (x *EngineWant) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[20] + mi := &file_llmgw_v1_agent_proto_msgTypes[21] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -2790,7 +2907,7 @@ func (x *EngineWant) ProtoReflect() protoreflect.Message { // Deprecated: Use EngineWant.ProtoReflect.Descriptor instead. func (*EngineWant) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{20} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{21} } func (x *EngineWant) GetEngine() string { @@ -2830,7 +2947,7 @@ type ProvisionStatus struct { func (x *ProvisionStatus) Reset() { *x = ProvisionStatus{} - mi := &file_llmgw_v1_agent_proto_msgTypes[21] + mi := &file_llmgw_v1_agent_proto_msgTypes[22] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -2842,7 +2959,7 @@ func (x *ProvisionStatus) String() string { func (*ProvisionStatus) ProtoMessage() {} func (x *ProvisionStatus) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[21] + mi := &file_llmgw_v1_agent_proto_msgTypes[22] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -2855,7 +2972,7 @@ func (x *ProvisionStatus) ProtoReflect() protoreflect.Message { // Deprecated: Use ProvisionStatus.ProtoReflect.Descriptor instead. func (*ProvisionStatus) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{21} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{22} } func (x *ProvisionStatus) GetEngine() string { @@ -2919,7 +3036,7 @@ type ShardStatus struct { func (x *ShardStatus) Reset() { *x = ShardStatus{} - mi := &file_llmgw_v1_agent_proto_msgTypes[22] + mi := &file_llmgw_v1_agent_proto_msgTypes[23] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -2931,7 +3048,7 @@ func (x *ShardStatus) String() string { func (*ShardStatus) ProtoMessage() {} func (x *ShardStatus) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[22] + mi := &file_llmgw_v1_agent_proto_msgTypes[23] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -2944,7 +3061,7 @@ func (x *ShardStatus) ProtoReflect() protoreflect.Message { // Deprecated: Use ShardStatus.ProtoReflect.Descriptor instead. func (*ShardStatus) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{22} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{23} } func (x *ShardStatus) GetNodeId() string { @@ -3018,7 +3135,7 @@ type FetchProgress struct { func (x *FetchProgress) Reset() { *x = FetchProgress{} - mi := &file_llmgw_v1_agent_proto_msgTypes[23] + mi := &file_llmgw_v1_agent_proto_msgTypes[24] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -3030,7 +3147,7 @@ func (x *FetchProgress) String() string { func (*FetchProgress) ProtoMessage() {} func (x *FetchProgress) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[23] + mi := &file_llmgw_v1_agent_proto_msgTypes[24] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -3043,7 +3160,7 @@ func (x *FetchProgress) ProtoReflect() protoreflect.Message { // Deprecated: Use FetchProgress.ProtoReflect.Descriptor instead. func (*FetchProgress) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{23} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{24} } func (x *FetchProgress) GetSource() string { @@ -3106,7 +3223,7 @@ type StartShard struct { func (x *StartShard) Reset() { *x = StartShard{} - mi := &file_llmgw_v1_agent_proto_msgTypes[24] + mi := &file_llmgw_v1_agent_proto_msgTypes[25] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -3118,7 +3235,7 @@ func (x *StartShard) String() string { func (*StartShard) ProtoMessage() {} func (x *StartShard) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[24] + mi := &file_llmgw_v1_agent_proto_msgTypes[25] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -3131,7 +3248,7 @@ func (x *StartShard) ProtoReflect() protoreflect.Message { // Deprecated: Use StartShard.ProtoReflect.Descriptor instead. func (*StartShard) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{24} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{25} } func (x *StartShard) GetReplicaId() string { @@ -3268,14 +3385,17 @@ type EngineProfile struct { ServedModelName string `protobuf:"bytes,11,opt,name=served_model_name,json=servedModelName,proto3" json:"served_model_name,omitempty"` // adapters are LoRA adapters loaded beside the model, in order: the index // is the id a request applies one by. - Adapters []*LoraAdapter `protobuf:"bytes,12,rep,name=adapters,proto3" json:"adapters,omitempty"` + Adapters []*LoraAdapter `protobuf:"bytes,12,rep,name=adapters,proto3" json:"adapters,omitempty"` + // kv_events has a vLLM engine publish its KV cache events to the agent, for + // the router's precise index. + KvEvents bool `protobuf:"varint,13,opt,name=kv_events,json=kvEvents,proto3" json:"kv_events,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } func (x *EngineProfile) Reset() { *x = EngineProfile{} - mi := &file_llmgw_v1_agent_proto_msgTypes[25] + mi := &file_llmgw_v1_agent_proto_msgTypes[26] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -3287,7 +3407,7 @@ func (x *EngineProfile) String() string { func (*EngineProfile) ProtoMessage() {} func (x *EngineProfile) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[25] + mi := &file_llmgw_v1_agent_proto_msgTypes[26] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -3300,7 +3420,7 @@ func (x *EngineProfile) ProtoReflect() protoreflect.Message { // Deprecated: Use EngineProfile.ProtoReflect.Descriptor instead. func (*EngineProfile) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{25} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{26} } func (x *EngineProfile) GetName() string { @@ -3387,6 +3507,13 @@ func (x *EngineProfile) GetAdapters() []*LoraAdapter { return nil } +func (x *EngineProfile) GetKvEvents() bool { + if x != nil { + return x.KvEvents + } + return false +} + // RemoveWeights deletes cached model files a node no longer needs. // // The paths are computed by the control plane from the catalog, but the agent @@ -3403,7 +3530,7 @@ type RemoveWeights struct { func (x *RemoveWeights) Reset() { *x = RemoveWeights{} - mi := &file_llmgw_v1_agent_proto_msgTypes[26] + mi := &file_llmgw_v1_agent_proto_msgTypes[27] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -3415,7 +3542,7 @@ func (x *RemoveWeights) String() string { func (*RemoveWeights) ProtoMessage() {} func (x *RemoveWeights) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[26] + mi := &file_llmgw_v1_agent_proto_msgTypes[27] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -3428,7 +3555,7 @@ func (x *RemoveWeights) ProtoReflect() protoreflect.Message { // Deprecated: Use RemoveWeights.ProtoReflect.Descriptor instead. func (*RemoveWeights) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{26} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{27} } func (x *RemoveWeights) GetPaths() []string { @@ -3448,7 +3575,7 @@ type StopShard struct { func (x *StopShard) Reset() { *x = StopShard{} - mi := &file_llmgw_v1_agent_proto_msgTypes[27] + mi := &file_llmgw_v1_agent_proto_msgTypes[28] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -3460,7 +3587,7 @@ func (x *StopShard) String() string { func (*StopShard) ProtoMessage() {} func (x *StopShard) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[27] + mi := &file_llmgw_v1_agent_proto_msgTypes[28] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -3473,7 +3600,7 @@ func (x *StopShard) ProtoReflect() protoreflect.Message { // Deprecated: Use StopShard.ProtoReflect.Descriptor instead. func (*StopShard) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{27} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{28} } func (x *StopShard) GetShardId() string { @@ -3493,7 +3620,7 @@ type DrainShard struct { func (x *DrainShard) Reset() { *x = DrainShard{} - mi := &file_llmgw_v1_agent_proto_msgTypes[28] + mi := &file_llmgw_v1_agent_proto_msgTypes[29] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -3505,7 +3632,7 @@ func (x *DrainShard) String() string { func (*DrainShard) ProtoMessage() {} func (x *DrainShard) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[28] + mi := &file_llmgw_v1_agent_proto_msgTypes[29] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -3518,7 +3645,7 @@ func (x *DrainShard) ProtoReflect() protoreflect.Message { // Deprecated: Use DrainShard.ProtoReflect.Descriptor instead. func (*DrainShard) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{28} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{29} } func (x *DrainShard) GetShardId() string { @@ -3540,7 +3667,7 @@ type Ack struct { func (x *Ack) Reset() { *x = Ack{} - mi := &file_llmgw_v1_agent_proto_msgTypes[29] + mi := &file_llmgw_v1_agent_proto_msgTypes[30] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -3552,7 +3679,7 @@ func (x *Ack) String() string { func (*Ack) ProtoMessage() {} func (x *Ack) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[29] + mi := &file_llmgw_v1_agent_proto_msgTypes[30] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -3565,7 +3692,7 @@ func (x *Ack) ProtoReflect() protoreflect.Message { // Deprecated: Use Ack.ProtoReflect.Descriptor instead. func (*Ack) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{29} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{30} } func (x *Ack) GetMessageId() string { @@ -3609,7 +3736,7 @@ type LogsRequest struct { func (x *LogsRequest) Reset() { *x = LogsRequest{} - mi := &file_llmgw_v1_agent_proto_msgTypes[30] + mi := &file_llmgw_v1_agent_proto_msgTypes[31] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -3621,7 +3748,7 @@ func (x *LogsRequest) String() string { func (*LogsRequest) ProtoMessage() {} func (x *LogsRequest) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[30] + mi := &file_llmgw_v1_agent_proto_msgTypes[31] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -3634,7 +3761,7 @@ func (x *LogsRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use LogsRequest.ProtoReflect.Descriptor instead. func (*LogsRequest) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{30} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{31} } func (x *LogsRequest) GetRequestId() string { @@ -3674,7 +3801,7 @@ type LogsReply struct { func (x *LogsReply) Reset() { *x = LogsReply{} - mi := &file_llmgw_v1_agent_proto_msgTypes[31] + mi := &file_llmgw_v1_agent_proto_msgTypes[32] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -3686,7 +3813,7 @@ func (x *LogsReply) String() string { func (*LogsReply) ProtoMessage() {} func (x *LogsReply) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[31] + mi := &file_llmgw_v1_agent_proto_msgTypes[32] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -3699,7 +3826,7 @@ func (x *LogsReply) ProtoReflect() protoreflect.Message { // Deprecated: Use LogsReply.ProtoReflect.Descriptor instead. func (*LogsReply) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{31} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{32} } func (x *LogsReply) GetRequestId() string { @@ -3745,7 +3872,7 @@ type FabricScanRequest struct { func (x *FabricScanRequest) Reset() { *x = FabricScanRequest{} - mi := &file_llmgw_v1_agent_proto_msgTypes[32] + mi := &file_llmgw_v1_agent_proto_msgTypes[33] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -3757,7 +3884,7 @@ func (x *FabricScanRequest) String() string { func (*FabricScanRequest) ProtoMessage() {} func (x *FabricScanRequest) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[32] + mi := &file_llmgw_v1_agent_proto_msgTypes[33] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -3770,7 +3897,7 @@ func (x *FabricScanRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use FabricScanRequest.ProtoReflect.Descriptor instead. func (*FabricScanRequest) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{32} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{33} } func (x *FabricScanRequest) GetRequestId() string { @@ -3809,7 +3936,7 @@ type FabricScanReply struct { func (x *FabricScanReply) Reset() { *x = FabricScanReply{} - mi := &file_llmgw_v1_agent_proto_msgTypes[33] + mi := &file_llmgw_v1_agent_proto_msgTypes[34] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -3821,7 +3948,7 @@ func (x *FabricScanReply) String() string { func (*FabricScanReply) ProtoMessage() {} func (x *FabricScanReply) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[33] + mi := &file_llmgw_v1_agent_proto_msgTypes[34] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -3834,7 +3961,7 @@ func (x *FabricScanReply) ProtoReflect() protoreflect.Message { // Deprecated: Use FabricScanReply.ProtoReflect.Descriptor instead. func (*FabricScanReply) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{33} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{34} } func (x *FabricScanReply) GetRequestId() string { @@ -3890,7 +4017,7 @@ type FabricTestRequest struct { func (x *FabricTestRequest) Reset() { *x = FabricTestRequest{} - mi := &file_llmgw_v1_agent_proto_msgTypes[34] + mi := &file_llmgw_v1_agent_proto_msgTypes[35] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -3902,7 +4029,7 @@ func (x *FabricTestRequest) String() string { func (*FabricTestRequest) ProtoMessage() {} func (x *FabricTestRequest) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[34] + mi := &file_llmgw_v1_agent_proto_msgTypes[35] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -3915,7 +4042,7 @@ func (x *FabricTestRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use FabricTestRequest.ProtoReflect.Descriptor instead. func (*FabricTestRequest) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{34} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{35} } func (x *FabricTestRequest) GetRequestId() string { @@ -3984,7 +4111,7 @@ type FabricTestReply struct { func (x *FabricTestReply) Reset() { *x = FabricTestReply{} - mi := &file_llmgw_v1_agent_proto_msgTypes[35] + mi := &file_llmgw_v1_agent_proto_msgTypes[36] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -3996,7 +4123,7 @@ func (x *FabricTestReply) String() string { func (*FabricTestReply) ProtoMessage() {} func (x *FabricTestReply) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[35] + mi := &file_llmgw_v1_agent_proto_msgTypes[36] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -4009,7 +4136,7 @@ func (x *FabricTestReply) ProtoReflect() protoreflect.Message { // Deprecated: Use FabricTestReply.ProtoReflect.Descriptor instead. func (*FabricTestReply) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{35} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{36} } func (x *FabricTestReply) GetRequestId() string { @@ -4069,7 +4196,7 @@ type UpgradeAgent struct { func (x *UpgradeAgent) Reset() { *x = UpgradeAgent{} - mi := &file_llmgw_v1_agent_proto_msgTypes[36] + mi := &file_llmgw_v1_agent_proto_msgTypes[37] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -4081,7 +4208,7 @@ func (x *UpgradeAgent) String() string { func (*UpgradeAgent) ProtoMessage() {} func (x *UpgradeAgent) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[36] + mi := &file_llmgw_v1_agent_proto_msgTypes[37] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -4094,7 +4221,7 @@ func (x *UpgradeAgent) ProtoReflect() protoreflect.Message { // Deprecated: Use UpgradeAgent.ProtoReflect.Descriptor instead. func (*UpgradeAgent) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{36} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{37} } func (x *UpgradeAgent) GetRequestId() string { @@ -4131,7 +4258,7 @@ type UpgradeReply struct { func (x *UpgradeReply) Reset() { *x = UpgradeReply{} - mi := &file_llmgw_v1_agent_proto_msgTypes[37] + mi := &file_llmgw_v1_agent_proto_msgTypes[38] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -4143,7 +4270,7 @@ func (x *UpgradeReply) String() string { func (*UpgradeReply) ProtoMessage() {} func (x *UpgradeReply) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[37] + mi := &file_llmgw_v1_agent_proto_msgTypes[38] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -4156,7 +4283,7 @@ func (x *UpgradeReply) ProtoReflect() protoreflect.Message { // Deprecated: Use UpgradeReply.ProtoReflect.Descriptor instead. func (*UpgradeReply) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{37} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{38} } func (x *UpgradeReply) GetRequestId() string { @@ -4192,7 +4319,7 @@ type LoraAdapter struct { func (x *LoraAdapter) Reset() { *x = LoraAdapter{} - mi := &file_llmgw_v1_agent_proto_msgTypes[38] + mi := &file_llmgw_v1_agent_proto_msgTypes[39] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -4204,7 +4331,7 @@ func (x *LoraAdapter) String() string { func (*LoraAdapter) ProtoMessage() {} func (x *LoraAdapter) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[38] + mi := &file_llmgw_v1_agent_proto_msgTypes[39] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -4217,7 +4344,7 @@ func (x *LoraAdapter) ProtoReflect() protoreflect.Message { // Deprecated: Use LoraAdapter.ProtoReflect.Descriptor instead. func (*LoraAdapter) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{38} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{39} } func (x *LoraAdapter) GetName() string { @@ -4238,7 +4365,7 @@ var File_llmgw_v1_agent_proto protoreflect.FileDescriptor const file_llmgw_v1_agent_proto_rawDesc = "" + "\n" + - "\x14llmgw/v1/agent.proto\x12\bllmgw.v1\"\x86\x03\n" + + "\x14llmgw/v1/agent.proto\x12\bllmgw.v1\"\xb9\x03\n" + "\fAgentMessage\x120\n" + "\bregister\x18\x01 \x01(\v2\x12.llmgw.v1.RegisterH\x00R\bregister\x123\n" + "\theartbeat\x18\x02 \x01(\v2\x13.llmgw.v1.HeartbeatH\x00R\theartbeat\x12-\n" + @@ -4248,8 +4375,19 @@ const file_llmgw_v1_agent_proto_rawDesc = "" + "fabricScan\x12<\n" + "\vfabric_test\x18\x06 \x01(\v2\x19.llmgw.v1.FabricTestReplyH\x00R\n" + "fabricTest\x122\n" + - "\aupgrade\x18\a \x01(\v2\x16.llmgw.v1.UpgradeReplyH\x00R\aupgradeB\x05\n" + - "\x03msg\"\xd0\x05\n" + + "\aupgrade\x18\a \x01(\v2\x16.llmgw.v1.UpgradeReplyH\x00R\aupgrade\x121\n" + + "\tkv_events\x18\b \x01(\v2\x12.llmgw.v1.KVEventsH\x00R\bkvEventsB\x05\n" + + "\x03msg\"\xb5\x01\n" + + "\bKVEvents\x12\x1d\n" + + "\n" + + "replica_id\x18\x01 \x01(\tR\treplicaId\x12\x1d\n" + + "\n" + + "block_size\x18\x02 \x01(\x05R\tblockSize\x12\x1f\n" + + "\vclear_first\x18\x03 \x01(\bR\n" + + "clearFirst\x12\x16\n" + + "\x06stored\x18\x04 \x03(\x06R\x06stored\x12\x18\n" + + "\aremoved\x18\x05 \x03(\x06R\aremoved\x12\x18\n" + + "\aproblem\x18\x06 \x01(\tR\aproblem\"\xd0\x05\n" + "\x0eControlMessage\x12,\n" + "\x05start\x18\x01 \x01(\v2\x14.llmgw.v1.StartShardH\x00R\x05start\x12)\n" + "\x04stop\x18\x02 \x01(\v2\x13.llmgw.v1.StopShardH\x00R\x04stop\x12,\n" + @@ -4568,7 +4706,7 @@ const file_llmgw_v1_agent_proto_rawDesc = "" + "profile_fp\x18\x0f \x01(\tR\tprofileFp\x12\x17\n" + "\aapi_key\x18\x10 \x01(\tR\x06apiKey\x12\x1f\n" + "\vcpu_threads\x18\x11 \x01(\x05R\n" + - "cpuThreads\"\xb0\x03\n" + + "cpuThreads\"\xcd\x03\n" + "\rEngineProfile\x12\x12\n" + "\x04name\x18\x01 \x01(\tR\x04name\x12\x16\n" + "\x06engine\x18\x02 \x01(\tR\x06engine\x12\x1b\n" + @@ -4583,7 +4721,8 @@ const file_llmgw_v1_agent_proto_rawDesc = "" + "model_path\x18\n" + " \x01(\tR\tmodelPath\x12*\n" + "\x11served_model_name\x18\v \x01(\tR\x0fservedModelName\x121\n" + - "\badapters\x18\f \x03(\v2\x15.llmgw.v1.LoraAdapterR\badapters\x1a6\n" + + "\badapters\x18\f \x03(\v2\x15.llmgw.v1.LoraAdapterR\badapters\x12\x1b\n" + + "\tkv_events\x18\r \x01(\bR\bkvEvents\x1a6\n" + "\bEnvEntry\x12\x10\n" + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"%\n" + @@ -4668,114 +4807,116 @@ func file_llmgw_v1_agent_proto_rawDescGZIP() []byte { } var file_llmgw_v1_agent_proto_enumTypes = make([]protoimpl.EnumInfo, 1) -var file_llmgw_v1_agent_proto_msgTypes = make([]protoimpl.MessageInfo, 46) +var file_llmgw_v1_agent_proto_msgTypes = make([]protoimpl.MessageInfo, 47) var file_llmgw_v1_agent_proto_goTypes = []any{ (ShardStatus_State)(0), // 0: llmgw.v1.ShardStatus.State (*AgentMessage)(nil), // 1: llmgw.v1.AgentMessage - (*ControlMessage)(nil), // 2: llmgw.v1.ControlMessage - (*ControlPlanes)(nil), // 3: llmgw.v1.ControlPlanes - (*Register)(nil), // 4: llmgw.v1.Register - (*Host)(nil), // 5: llmgw.v1.Host - (*Storage)(nil), // 6: llmgw.v1.Storage - (*RunningShard)(nil), // 7: llmgw.v1.RunningShard - (*GPU)(nil), // 8: llmgw.v1.GPU - (*MIGDevice)(nil), // 9: llmgw.v1.MIGDevice - (*FabricAttachment)(nil), // 10: llmgw.v1.FabricAttachment - (*LocalWeights)(nil), // 11: llmgw.v1.LocalWeights - (*NodeLoad)(nil), // 12: llmgw.v1.NodeLoad - (*GPUSensor)(nil), // 13: llmgw.v1.GPUSensor - (*ShardLoad)(nil), // 14: llmgw.v1.ShardLoad - (*IBPort)(nil), // 15: llmgw.v1.IBPort - (*SubnetManager)(nil), // 16: llmgw.v1.SubnetManager - (*IMEX)(nil), // 17: llmgw.v1.IMEX - (*GPUFabricState)(nil), // 18: llmgw.v1.GPUFabricState - (*Heartbeat)(nil), // 19: llmgw.v1.Heartbeat - (*Provision)(nil), // 20: llmgw.v1.Provision - (*EngineWant)(nil), // 21: llmgw.v1.EngineWant - (*ProvisionStatus)(nil), // 22: llmgw.v1.ProvisionStatus - (*ShardStatus)(nil), // 23: llmgw.v1.ShardStatus - (*FetchProgress)(nil), // 24: llmgw.v1.FetchProgress - (*StartShard)(nil), // 25: llmgw.v1.StartShard - (*EngineProfile)(nil), // 26: llmgw.v1.EngineProfile - (*RemoveWeights)(nil), // 27: llmgw.v1.RemoveWeights - (*StopShard)(nil), // 28: llmgw.v1.StopShard - (*DrainShard)(nil), // 29: llmgw.v1.DrainShard - (*Ack)(nil), // 30: llmgw.v1.Ack - (*LogsRequest)(nil), // 31: llmgw.v1.LogsRequest - (*LogsReply)(nil), // 32: llmgw.v1.LogsReply - (*FabricScanRequest)(nil), // 33: llmgw.v1.FabricScanRequest - (*FabricScanReply)(nil), // 34: llmgw.v1.FabricScanReply - (*FabricTestRequest)(nil), // 35: llmgw.v1.FabricTestRequest - (*FabricTestReply)(nil), // 36: llmgw.v1.FabricTestReply - (*UpgradeAgent)(nil), // 37: llmgw.v1.UpgradeAgent - (*UpgradeReply)(nil), // 38: llmgw.v1.UpgradeReply - (*LoraAdapter)(nil), // 39: llmgw.v1.LoraAdapter - nil, // 40: llmgw.v1.ControlMessage.TraceEntry - nil, // 41: llmgw.v1.Register.LabelsEntry - nil, // 42: llmgw.v1.Register.EngineVersionsEntry - nil, // 43: llmgw.v1.GPU.NicPathsEntry - nil, // 44: llmgw.v1.IBPort.CountersEntry - nil, // 45: llmgw.v1.Heartbeat.EngineVersionsEntry - nil, // 46: llmgw.v1.EngineProfile.EnvEntry + (*KVEvents)(nil), // 2: llmgw.v1.KVEvents + (*ControlMessage)(nil), // 3: llmgw.v1.ControlMessage + (*ControlPlanes)(nil), // 4: llmgw.v1.ControlPlanes + (*Register)(nil), // 5: llmgw.v1.Register + (*Host)(nil), // 6: llmgw.v1.Host + (*Storage)(nil), // 7: llmgw.v1.Storage + (*RunningShard)(nil), // 8: llmgw.v1.RunningShard + (*GPU)(nil), // 9: llmgw.v1.GPU + (*MIGDevice)(nil), // 10: llmgw.v1.MIGDevice + (*FabricAttachment)(nil), // 11: llmgw.v1.FabricAttachment + (*LocalWeights)(nil), // 12: llmgw.v1.LocalWeights + (*NodeLoad)(nil), // 13: llmgw.v1.NodeLoad + (*GPUSensor)(nil), // 14: llmgw.v1.GPUSensor + (*ShardLoad)(nil), // 15: llmgw.v1.ShardLoad + (*IBPort)(nil), // 16: llmgw.v1.IBPort + (*SubnetManager)(nil), // 17: llmgw.v1.SubnetManager + (*IMEX)(nil), // 18: llmgw.v1.IMEX + (*GPUFabricState)(nil), // 19: llmgw.v1.GPUFabricState + (*Heartbeat)(nil), // 20: llmgw.v1.Heartbeat + (*Provision)(nil), // 21: llmgw.v1.Provision + (*EngineWant)(nil), // 22: llmgw.v1.EngineWant + (*ProvisionStatus)(nil), // 23: llmgw.v1.ProvisionStatus + (*ShardStatus)(nil), // 24: llmgw.v1.ShardStatus + (*FetchProgress)(nil), // 25: llmgw.v1.FetchProgress + (*StartShard)(nil), // 26: llmgw.v1.StartShard + (*EngineProfile)(nil), // 27: llmgw.v1.EngineProfile + (*RemoveWeights)(nil), // 28: llmgw.v1.RemoveWeights + (*StopShard)(nil), // 29: llmgw.v1.StopShard + (*DrainShard)(nil), // 30: llmgw.v1.DrainShard + (*Ack)(nil), // 31: llmgw.v1.Ack + (*LogsRequest)(nil), // 32: llmgw.v1.LogsRequest + (*LogsReply)(nil), // 33: llmgw.v1.LogsReply + (*FabricScanRequest)(nil), // 34: llmgw.v1.FabricScanRequest + (*FabricScanReply)(nil), // 35: llmgw.v1.FabricScanReply + (*FabricTestRequest)(nil), // 36: llmgw.v1.FabricTestRequest + (*FabricTestReply)(nil), // 37: llmgw.v1.FabricTestReply + (*UpgradeAgent)(nil), // 38: llmgw.v1.UpgradeAgent + (*UpgradeReply)(nil), // 39: llmgw.v1.UpgradeReply + (*LoraAdapter)(nil), // 40: llmgw.v1.LoraAdapter + nil, // 41: llmgw.v1.ControlMessage.TraceEntry + nil, // 42: llmgw.v1.Register.LabelsEntry + nil, // 43: llmgw.v1.Register.EngineVersionsEntry + nil, // 44: llmgw.v1.GPU.NicPathsEntry + nil, // 45: llmgw.v1.IBPort.CountersEntry + nil, // 46: llmgw.v1.Heartbeat.EngineVersionsEntry + nil, // 47: llmgw.v1.EngineProfile.EnvEntry } var file_llmgw_v1_agent_proto_depIdxs = []int32{ - 4, // 0: llmgw.v1.AgentMessage.register:type_name -> llmgw.v1.Register - 19, // 1: llmgw.v1.AgentMessage.heartbeat:type_name -> llmgw.v1.Heartbeat - 23, // 2: llmgw.v1.AgentMessage.shard:type_name -> llmgw.v1.ShardStatus - 32, // 3: llmgw.v1.AgentMessage.logs:type_name -> llmgw.v1.LogsReply - 34, // 4: llmgw.v1.AgentMessage.fabric_scan:type_name -> llmgw.v1.FabricScanReply - 36, // 5: llmgw.v1.AgentMessage.fabric_test:type_name -> llmgw.v1.FabricTestReply - 38, // 6: llmgw.v1.AgentMessage.upgrade:type_name -> llmgw.v1.UpgradeReply - 25, // 7: llmgw.v1.ControlMessage.start:type_name -> llmgw.v1.StartShard - 28, // 8: llmgw.v1.ControlMessage.stop:type_name -> llmgw.v1.StopShard - 29, // 9: llmgw.v1.ControlMessage.drain:type_name -> llmgw.v1.DrainShard - 30, // 10: llmgw.v1.ControlMessage.ack:type_name -> llmgw.v1.Ack - 27, // 11: llmgw.v1.ControlMessage.remove_weights:type_name -> llmgw.v1.RemoveWeights - 20, // 12: llmgw.v1.ControlMessage.provision:type_name -> llmgw.v1.Provision - 31, // 13: llmgw.v1.ControlMessage.logs:type_name -> llmgw.v1.LogsRequest - 33, // 14: llmgw.v1.ControlMessage.fabric_scan:type_name -> llmgw.v1.FabricScanRequest - 35, // 15: llmgw.v1.ControlMessage.fabric_test:type_name -> llmgw.v1.FabricTestRequest - 37, // 16: llmgw.v1.ControlMessage.upgrade:type_name -> llmgw.v1.UpgradeAgent - 3, // 17: llmgw.v1.ControlMessage.control_planes:type_name -> llmgw.v1.ControlPlanes - 40, // 18: llmgw.v1.ControlMessage.trace:type_name -> llmgw.v1.ControlMessage.TraceEntry - 41, // 19: llmgw.v1.Register.labels:type_name -> llmgw.v1.Register.LabelsEntry - 8, // 20: llmgw.v1.Register.gpus:type_name -> llmgw.v1.GPU - 10, // 21: llmgw.v1.Register.fabrics:type_name -> llmgw.v1.FabricAttachment - 42, // 22: llmgw.v1.Register.engine_versions:type_name -> llmgw.v1.Register.EngineVersionsEntry - 7, // 23: llmgw.v1.Register.running_shards:type_name -> llmgw.v1.RunningShard - 11, // 24: llmgw.v1.Register.weights:type_name -> llmgw.v1.LocalWeights - 6, // 25: llmgw.v1.Register.storage:type_name -> llmgw.v1.Storage - 5, // 26: llmgw.v1.Register.host:type_name -> llmgw.v1.Host - 17, // 27: llmgw.v1.Register.imex:type_name -> llmgw.v1.IMEX - 0, // 28: llmgw.v1.RunningShard.state:type_name -> llmgw.v1.ShardStatus.State - 43, // 29: llmgw.v1.GPU.nic_paths:type_name -> llmgw.v1.GPU.NicPathsEntry - 9, // 30: llmgw.v1.GPU.mig:type_name -> llmgw.v1.MIGDevice - 14, // 31: llmgw.v1.NodeLoad.shards:type_name -> llmgw.v1.ShardLoad - 13, // 32: llmgw.v1.NodeLoad.gpus:type_name -> llmgw.v1.GPUSensor - 44, // 33: llmgw.v1.IBPort.counters:type_name -> llmgw.v1.IBPort.CountersEntry - 16, // 34: llmgw.v1.IBPort.subnet_managers:type_name -> llmgw.v1.SubnetManager - 6, // 35: llmgw.v1.Heartbeat.storage:type_name -> llmgw.v1.Storage - 11, // 36: llmgw.v1.Heartbeat.weights:type_name -> llmgw.v1.LocalWeights - 45, // 37: llmgw.v1.Heartbeat.engine_versions:type_name -> llmgw.v1.Heartbeat.EngineVersionsEntry - 22, // 38: llmgw.v1.Heartbeat.provision:type_name -> llmgw.v1.ProvisionStatus - 12, // 39: llmgw.v1.Heartbeat.load:type_name -> llmgw.v1.NodeLoad - 15, // 40: llmgw.v1.Heartbeat.ib_ports:type_name -> llmgw.v1.IBPort - 10, // 41: llmgw.v1.Heartbeat.fabrics:type_name -> llmgw.v1.FabricAttachment - 18, // 42: llmgw.v1.Heartbeat.gpu_fabric:type_name -> llmgw.v1.GPUFabricState - 17, // 43: llmgw.v1.Heartbeat.imex:type_name -> llmgw.v1.IMEX - 21, // 44: llmgw.v1.Provision.engines:type_name -> llmgw.v1.EngineWant - 0, // 45: llmgw.v1.ShardStatus.state:type_name -> llmgw.v1.ShardStatus.State - 24, // 46: llmgw.v1.ShardStatus.fetch:type_name -> llmgw.v1.FetchProgress - 26, // 47: llmgw.v1.StartShard.profile:type_name -> llmgw.v1.EngineProfile - 46, // 48: llmgw.v1.EngineProfile.env:type_name -> llmgw.v1.EngineProfile.EnvEntry - 39, // 49: llmgw.v1.EngineProfile.adapters:type_name -> llmgw.v1.LoraAdapter - 1, // 50: llmgw.v1.AgentService.Connect:input_type -> llmgw.v1.AgentMessage - 2, // 51: llmgw.v1.AgentService.Connect:output_type -> llmgw.v1.ControlMessage - 51, // [51:52] is the sub-list for method output_type - 50, // [50:51] is the sub-list for method input_type - 50, // [50:50] is the sub-list for extension type_name - 50, // [50:50] is the sub-list for extension extendee - 0, // [0:50] is the sub-list for field type_name + 5, // 0: llmgw.v1.AgentMessage.register:type_name -> llmgw.v1.Register + 20, // 1: llmgw.v1.AgentMessage.heartbeat:type_name -> llmgw.v1.Heartbeat + 24, // 2: llmgw.v1.AgentMessage.shard:type_name -> llmgw.v1.ShardStatus + 33, // 3: llmgw.v1.AgentMessage.logs:type_name -> llmgw.v1.LogsReply + 35, // 4: llmgw.v1.AgentMessage.fabric_scan:type_name -> llmgw.v1.FabricScanReply + 37, // 5: llmgw.v1.AgentMessage.fabric_test:type_name -> llmgw.v1.FabricTestReply + 39, // 6: llmgw.v1.AgentMessage.upgrade:type_name -> llmgw.v1.UpgradeReply + 2, // 7: llmgw.v1.AgentMessage.kv_events:type_name -> llmgw.v1.KVEvents + 26, // 8: llmgw.v1.ControlMessage.start:type_name -> llmgw.v1.StartShard + 29, // 9: llmgw.v1.ControlMessage.stop:type_name -> llmgw.v1.StopShard + 30, // 10: llmgw.v1.ControlMessage.drain:type_name -> llmgw.v1.DrainShard + 31, // 11: llmgw.v1.ControlMessage.ack:type_name -> llmgw.v1.Ack + 28, // 12: llmgw.v1.ControlMessage.remove_weights:type_name -> llmgw.v1.RemoveWeights + 21, // 13: llmgw.v1.ControlMessage.provision:type_name -> llmgw.v1.Provision + 32, // 14: llmgw.v1.ControlMessage.logs:type_name -> llmgw.v1.LogsRequest + 34, // 15: llmgw.v1.ControlMessage.fabric_scan:type_name -> llmgw.v1.FabricScanRequest + 36, // 16: llmgw.v1.ControlMessage.fabric_test:type_name -> llmgw.v1.FabricTestRequest + 38, // 17: llmgw.v1.ControlMessage.upgrade:type_name -> llmgw.v1.UpgradeAgent + 4, // 18: llmgw.v1.ControlMessage.control_planes:type_name -> llmgw.v1.ControlPlanes + 41, // 19: llmgw.v1.ControlMessage.trace:type_name -> llmgw.v1.ControlMessage.TraceEntry + 42, // 20: llmgw.v1.Register.labels:type_name -> llmgw.v1.Register.LabelsEntry + 9, // 21: llmgw.v1.Register.gpus:type_name -> llmgw.v1.GPU + 11, // 22: llmgw.v1.Register.fabrics:type_name -> llmgw.v1.FabricAttachment + 43, // 23: llmgw.v1.Register.engine_versions:type_name -> llmgw.v1.Register.EngineVersionsEntry + 8, // 24: llmgw.v1.Register.running_shards:type_name -> llmgw.v1.RunningShard + 12, // 25: llmgw.v1.Register.weights:type_name -> llmgw.v1.LocalWeights + 7, // 26: llmgw.v1.Register.storage:type_name -> llmgw.v1.Storage + 6, // 27: llmgw.v1.Register.host:type_name -> llmgw.v1.Host + 18, // 28: llmgw.v1.Register.imex:type_name -> llmgw.v1.IMEX + 0, // 29: llmgw.v1.RunningShard.state:type_name -> llmgw.v1.ShardStatus.State + 44, // 30: llmgw.v1.GPU.nic_paths:type_name -> llmgw.v1.GPU.NicPathsEntry + 10, // 31: llmgw.v1.GPU.mig:type_name -> llmgw.v1.MIGDevice + 15, // 32: llmgw.v1.NodeLoad.shards:type_name -> llmgw.v1.ShardLoad + 14, // 33: llmgw.v1.NodeLoad.gpus:type_name -> llmgw.v1.GPUSensor + 45, // 34: llmgw.v1.IBPort.counters:type_name -> llmgw.v1.IBPort.CountersEntry + 17, // 35: llmgw.v1.IBPort.subnet_managers:type_name -> llmgw.v1.SubnetManager + 7, // 36: llmgw.v1.Heartbeat.storage:type_name -> llmgw.v1.Storage + 12, // 37: llmgw.v1.Heartbeat.weights:type_name -> llmgw.v1.LocalWeights + 46, // 38: llmgw.v1.Heartbeat.engine_versions:type_name -> llmgw.v1.Heartbeat.EngineVersionsEntry + 23, // 39: llmgw.v1.Heartbeat.provision:type_name -> llmgw.v1.ProvisionStatus + 13, // 40: llmgw.v1.Heartbeat.load:type_name -> llmgw.v1.NodeLoad + 16, // 41: llmgw.v1.Heartbeat.ib_ports:type_name -> llmgw.v1.IBPort + 11, // 42: llmgw.v1.Heartbeat.fabrics:type_name -> llmgw.v1.FabricAttachment + 19, // 43: llmgw.v1.Heartbeat.gpu_fabric:type_name -> llmgw.v1.GPUFabricState + 18, // 44: llmgw.v1.Heartbeat.imex:type_name -> llmgw.v1.IMEX + 22, // 45: llmgw.v1.Provision.engines:type_name -> llmgw.v1.EngineWant + 0, // 46: llmgw.v1.ShardStatus.state:type_name -> llmgw.v1.ShardStatus.State + 25, // 47: llmgw.v1.ShardStatus.fetch:type_name -> llmgw.v1.FetchProgress + 27, // 48: llmgw.v1.StartShard.profile:type_name -> llmgw.v1.EngineProfile + 47, // 49: llmgw.v1.EngineProfile.env:type_name -> llmgw.v1.EngineProfile.EnvEntry + 40, // 50: llmgw.v1.EngineProfile.adapters:type_name -> llmgw.v1.LoraAdapter + 1, // 51: llmgw.v1.AgentService.Connect:input_type -> llmgw.v1.AgentMessage + 3, // 52: llmgw.v1.AgentService.Connect:output_type -> llmgw.v1.ControlMessage + 52, // [52:53] is the sub-list for method output_type + 51, // [51:52] is the sub-list for method input_type + 51, // [51:51] is the sub-list for extension type_name + 51, // [51:51] is the sub-list for extension extendee + 0, // [0:51] is the sub-list for field type_name } func init() { file_llmgw_v1_agent_proto_init() } @@ -4791,8 +4932,9 @@ func file_llmgw_v1_agent_proto_init() { (*AgentMessage_FabricScan)(nil), (*AgentMessage_FabricTest)(nil), (*AgentMessage_Upgrade)(nil), + (*AgentMessage_KvEvents)(nil), } - file_llmgw_v1_agent_proto_msgTypes[1].OneofWrappers = []any{ + file_llmgw_v1_agent_proto_msgTypes[2].OneofWrappers = []any{ (*ControlMessage_Start)(nil), (*ControlMessage_Stop)(nil), (*ControlMessage_Drain)(nil), @@ -4811,7 +4953,7 @@ func file_llmgw_v1_agent_proto_init() { GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: unsafe.Slice(unsafe.StringData(file_llmgw_v1_agent_proto_rawDesc), len(file_llmgw_v1_agent_proto_rawDesc)), NumEnums: 1, - NumMessages: 46, + NumMessages: 47, NumExtensions: 0, NumServices: 1, }, diff --git a/api/proto/llmgw/v1/agent.proto b/api/proto/llmgw/v1/agent.proto index 90447fb..b312e35 100644 --- a/api/proto/llmgw/v1/agent.proto +++ b/api/proto/llmgw/v1/agent.proto @@ -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. diff --git a/docs/admin-api.openapi.json b/docs/admin-api.openapi.json index 74c449b..6582335 100644 --- a/docs/admin-api.openapi.json +++ b/docs/admin-api.openapi.json @@ -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." } } }, diff --git a/docs/install.md b/docs/install.md index 114315a..21bfd0c 100644 --- a/docs/install.md +++ b/docs/install.md @@ -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 diff --git a/docs/reference.md b/docs/reference.md index d8a2484..ebe9d51 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -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` diff --git a/go.mod b/go.mod index 10cba6a..5734aec 100644 --- a/go.mod +++ b/go.mod @@ -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 diff --git a/go.sum b/go.sum index 6211a65..4690486 100644 --- a/go.sum +++ b/go.sum @@ -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= diff --git a/internal/agent/agent.go b/internal/agent/agent.go index b792358..287d129 100644 --- a/internal/agent/agent.go +++ b/internal/agent/agent.go @@ -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) diff --git a/internal/agent/engines/engine.go b/internal/agent/engines/engine.go index a370cc3..eb80178 100644 --- a/internal/agent/engines/engine.go +++ b/internal/agent/engines/engine.go @@ -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. diff --git a/internal/agent/engines/kvevents_test.go b/internal/agent/engines/kvevents_test.go new file mode 100644 index 0000000..0cb66a3 --- /dev/null +++ b/internal/agent/engines/kvevents_test.go @@ -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") + } +} diff --git a/internal/agent/engines/vllm.go b/internal/agent/engines/vllm.go index 1fff60b..c49f792 100644 --- a/internal/agent/engines/vllm.go +++ b/internal/agent/engines/vllm.go @@ -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...) diff --git a/internal/agent/kv.go b/internal/agent/kv.go new file mode 100644 index 0000000..07d6409 --- /dev/null +++ b/internal/agent/kv.go @@ -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: + } + } +} diff --git a/internal/agent/kv_test.go b/internal/agent/kv_test.go new file mode 100644 index 0000000..214473a --- /dev/null +++ b/internal/agent/kv_test.go @@ -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") + } +} diff --git a/internal/cli/server.go b/internal/cli/server.go index 79a572e..3ec15e9 100644 --- a/internal/cli/server.go +++ b/internal/cli/server.go @@ -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) diff --git a/internal/config/config.go b/internal/config/config.go index 20f5585..5c46ad9 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -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"` diff --git a/internal/control/agentsvc/convert.go b/internal/control/agentsvc/convert.go index 67a07e0..8e1cfae 100644 --- a/internal/control/agentsvc/convert.go +++ b/internal/control/agentsvc/convert.go @@ -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 diff --git a/internal/control/agentsvc/kv.go b/internal/control/agentsvc/kv.go new file mode 100644 index 0000000..1eaf597 --- /dev/null +++ b/internal/control/agentsvc/kv.go @@ -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, + }}} +} diff --git a/internal/control/agentsvc/kv_test.go b/internal/control/agentsvc/kv_test.go new file mode 100644 index 0000000..1151929 --- /dev/null +++ b/internal/control/agentsvc/kv_test.go @@ -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) + } +} diff --git a/internal/control/agentsvc/service.go b/internal/control/agentsvc/service.go index 5e79d7b..ee03240 100644 --- a/internal/control/agentsvc/service.go +++ b/internal/control/agentsvc/service.go @@ -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) } } diff --git a/internal/control/catalog/catalog.go b/internal/control/catalog/catalog.go index ded29dc..97e9982 100644 --- a/internal/control/catalog/catalog.go +++ b/internal/control/catalog/catalog.go @@ -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 { diff --git a/internal/control/catalog/kvevents_test.go b/internal/control/catalog/kvevents_test.go new file mode 100644 index 0000000..532a992 --- /dev/null +++ b/internal/control/catalog/kvevents_test.go @@ -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") + } +} diff --git a/internal/docgen/reference_gen.go b/internal/docgen/reference_gen.go index abd4942..8d2686e 100644 --- a/internal/docgen/reference_gen.go +++ b/internal/docgen/reference_gen.go @@ -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: ""}, diff --git a/internal/kvevents/decode.go b/internal/kvevents/decode.go new file mode 100644 index 0000000..48afad6 --- /dev/null +++ b/internal/kvevents/decode.go @@ -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 +} diff --git a/internal/kvevents/hash.go b/internal/kvevents/hash.go new file mode 100644 index 0000000..f483698 --- /dev/null +++ b/internal/kvevents/hash.go @@ -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) + } +} diff --git a/internal/kvevents/interop_test.go b/internal/kvevents/interop_test.go new file mode 100644 index 0000000..a42e287 --- /dev/null +++ b/internal/kvevents/interop_test.go @@ -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) + } +} diff --git a/internal/kvevents/kvevents_test.go b/internal/kvevents/kvevents_test.go new file mode 100644 index 0000000..364ec55 --- /dev/null +++ b/internal/kvevents/kvevents_test.go @@ -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") + } +} diff --git a/internal/kvevents/mirror.go b/internal/kvevents/mirror.go new file mode 100644 index 0000000..3c28d46 --- /dev/null +++ b/internal/kvevents/mirror.go @@ -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 +} diff --git a/internal/kvevents/subscriber.go b/internal/kvevents/subscriber.go new file mode 100644 index 0000000..dc71e55 --- /dev/null +++ b/internal/kvevents/subscriber.go @@ -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", + } +} diff --git a/internal/kvevents/testdata/pypub.py b/internal/kvevents/testdata/pypub.py new file mode 100644 index 0000000..14ba058 --- /dev/null +++ b/internal/kvevents/testdata/pypub.py @@ -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) diff --git a/internal/router/admin.go b/internal/router/admin.go index 4e74cd2..012df89 100644 --- a/internal/router/admin.go +++ b/internal/router/admin.go @@ -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 diff --git a/internal/router/kvindex/decision.go b/internal/router/kvindex/decision.go new file mode 100644 index 0000000..f2dbf72 --- /dev/null +++ b/internal/router/kvindex/decision.go @@ -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" +) diff --git a/internal/router/kvindex/index.go b/internal/router/kvindex/index.go new file mode 100644 index 0000000..eeadeb5 --- /dev/null +++ b/internal/router/kvindex/index.go @@ -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 +} diff --git a/internal/router/kvindex/kvindex_test.go b/internal/router/kvindex/kvindex_test.go new file mode 100644 index 0000000..0c45f9a --- /dev/null +++ b/internal/router/kvindex/kvindex_test.go @@ -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) + } + } +} diff --git a/internal/router/kvindex/tokenize.go b/internal/router/kvindex/tokenize.go new file mode 100644 index 0000000..5928c70 --- /dev/null +++ b/internal/router/kvindex/tokenize.go @@ -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 +} diff --git a/internal/router/kvroute.go b/internal/router/kvroute.go new file mode 100644 index 0000000..908bab7 --- /dev/null +++ b/internal/router/kvroute.go @@ -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) + } +} diff --git a/internal/router/kvroute_test.go b/internal/router/kvroute_test.go new file mode 100644 index 0000000..e9b9164 --- /dev/null +++ b/internal/router/kvroute_test.go @@ -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) + } +} diff --git a/internal/router/metrics.go b/internal/router/metrics.go index cd31eb7..5c67391 100644 --- a/internal/router/metrics.go +++ b/internal/router/metrics.go @@ -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" diff --git a/internal/router/prefix/scorer.go b/internal/router/prefix/scorer.go index f60e739..5463a36 100644 --- a/internal/router/prefix/scorer.go +++ b/internal/router/prefix/scorer.go @@ -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 { diff --git a/internal/router/router.go b/internal/router/router.go index d129cf6..1ec4efd 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -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 diff --git a/internal/router/scheduler/kv_test.go b/internal/router/scheduler/kv_test.go new file mode 100644 index 0000000..829fe15 --- /dev/null +++ b/internal/router/scheduler/kv_test.go @@ -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) + } +} diff --git a/internal/router/scheduler/prefix.go b/internal/router/scheduler/prefix.go index 3c64ef1..0241142 100644 --- a/internal/router/scheduler/prefix.go +++ b/internal/router/scheduler/prefix.go @@ -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) diff --git a/internal/router/scheduler/scheduler.go b/internal/router/scheduler/scheduler.go index 4ab93fe..52d6b62 100644 --- a/internal/router/scheduler/scheduler.go +++ b/internal/router/scheduler/scheduler.go @@ -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 diff --git a/notes/c24-kvindex.md b/notes/c24-kvindex.md new file mode 100644 index 0000000..48c1318 --- /dev/null +++ b/notes/c24-kvindex.md @@ -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.