diff --git a/CHANGELOG.md b/CHANGELOG.md index a92acbc..3ddcd3a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -100,10 +100,44 @@ fields. Under the simulated fleet the reports came back fast enough to hit the window every time. Everything is snapshotted under the lock now and the functions work from values. +### The agent stream + +- **Protocol** (`api/proto/llmgw/v1`): one bidirectional connect-go stream + carrying registration, heartbeats and shard status upward, and start, stop and + drain commands downward. The agent dials out, because GPU nodes sit behind + firewalls far more often than a control plane does and an outbound connection + needs no inbound rule on every node in the fleet. +- **Registration carries the whole inventory**, so there is no separate + discovery call and a reconnecting agent refreshes what the scheduler knows. + The fabric labels travel with it — they cannot be inferred from the wire, so + losing them in transit would silently disable multi-node placement. +- **Agent service** (`internal/control/agentsvc`) is both the stream handler and + the reconciler's dispatcher: it holds the open streams, so it is the only + thing that can deliver a command to a node. +- **Agent** (`internal/agent`) validates a profile before launching anything — + an engine that refuses a profile does so after minutes of loading — then + supervises the process and reports every transition, including the engine's + output tail on failure. +- **Reconnects forever with jittered backoff.** A control plane restart must not + require touching every GPU node, and a fleet must not reconnect in lockstep + and knock the control plane over again. Engines keep serving throughout: the + router talks to them directly. +- **`llmgw agent` is real**, and `llmgw server` runs the control plane whenever + any model declares a profile. The static endpoint list remains for + router-only runs against sim-engines. + +### Verified with real binaries + +Server and agent as separate processes over a real stream: the agent registers, +a replica is placed, an engine is launched through the vLLM adapter's actual +argument construction, streaming traffic is served through the router, and +killing the agent moves the node through unreachable to lost with placement +explaining exactly why it can no longer run anything. + ### Still missing for phase 2 -- The agent-to-control-plane stream. It pairs with phase 3's control plane - rather than standing alone, so `llmgw agent` itself is still a stub. +- mTLS. The stream is plaintext HTTP/2 today, so it must stay on a trusted + network. ### Phase 3: placement diff --git a/Makefile b/Makefile index 661c0c5..890fc11 100644 --- a/Makefile +++ b/Makefile @@ -2,7 +2,7 @@ MODULE := gitea.mynodes.xyz/ilker/llmgw VERSION ?= $(shell git describe --tags --always --dirty 2>/dev/null || echo dev) LDFLAGS := -X '$(MODULE)/internal/version.Version=$(VERSION)' -.PHONY: build test race lint fmt clean rename +.PHONY: build test race lint fmt clean rename proto proto-tools build: go build -ldflags "$(LDFLAGS)" -o bin/llmgw ./cmd/llmgw @@ -20,6 +20,17 @@ fmt: go fmt ./... go vet ./... +# Regenerate the agent protocol. Generated code is committed, so this only +# needs running when api/proto changes. +proto: proto-tools + buf lint + buf generate + +proto-tools: + go install github.com/bufbuild/buf/cmd/buf@latest + go install google.golang.org/protobuf/cmd/protoc-gen-go@latest + go install connectrpc.com/connect/cmd/protoc-gen-connect-go@latest + clean: rm -rf bin dist diff --git a/api/gen/llmgw/v1/agent.pb.go b/api/gen/llmgw/v1/agent.pb.go new file mode 100644 index 0000000..496eece --- /dev/null +++ b/api/gen/llmgw/v1/agent.pb.go @@ -0,0 +1,1421 @@ +// Code generated by protoc-gen-go. DO NOT EDIT. +// versions: +// protoc-gen-go v1.36.12 +// protoc (unknown) +// source: llmgw/v1/agent.proto + +// Package llmgw.v1 is the protocol between a node agent and the control plane +// (plan §5.6). +// +// The agent dials out and holds a bidirectional stream open. That direction is +// deliberate: GPU nodes sit behind firewalls and NAT far more often than a +// control plane does, and an outbound connection needs no inbound rule on every +// node in the fleet. + +package llmgwv1 + +import ( + protoreflect "google.golang.org/protobuf/reflect/protoreflect" + protoimpl "google.golang.org/protobuf/runtime/protoimpl" + reflect "reflect" + sync "sync" + unsafe "unsafe" +) + +const ( + // Verify that this generated code is sufficiently up-to-date. + _ = protoimpl.EnforceVersion(20 - protoimpl.MinVersion) + // Verify that runtime/protoimpl is sufficiently up-to-date. + _ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20) +) + +type ShardStatus_State int32 + +const ( + ShardStatus_STATE_UNSPECIFIED ShardStatus_State = 0 + ShardStatus_STATE_PENDING ShardStatus_State = 1 + ShardStatus_STATE_STARTING ShardStatus_State = 2 + ShardStatus_STATE_LOADING ShardStatus_State = 3 + ShardStatus_STATE_READY ShardStatus_State = 4 + ShardStatus_STATE_DRAINING ShardStatus_State = 5 + ShardStatus_STATE_STOPPED ShardStatus_State = 6 + ShardStatus_STATE_FAILED ShardStatus_State = 7 +) + +// Enum value maps for ShardStatus_State. +var ( + ShardStatus_State_name = map[int32]string{ + 0: "STATE_UNSPECIFIED", + 1: "STATE_PENDING", + 2: "STATE_STARTING", + 3: "STATE_LOADING", + 4: "STATE_READY", + 5: "STATE_DRAINING", + 6: "STATE_STOPPED", + 7: "STATE_FAILED", + } + ShardStatus_State_value = map[string]int32{ + "STATE_UNSPECIFIED": 0, + "STATE_PENDING": 1, + "STATE_STARTING": 2, + "STATE_LOADING": 3, + "STATE_READY": 4, + "STATE_DRAINING": 5, + "STATE_STOPPED": 6, + "STATE_FAILED": 7, + } +) + +func (x ShardStatus_State) Enum() *ShardStatus_State { + p := new(ShardStatus_State) + *p = x + return p +} + +func (x ShardStatus_State) String() string { + return protoimpl.X.EnumStringOf(x.Descriptor(), protoreflect.EnumNumber(x)) +} + +func (ShardStatus_State) Descriptor() protoreflect.EnumDescriptor { + return file_llmgw_v1_agent_proto_enumTypes[0].Descriptor() +} + +func (ShardStatus_State) Type() protoreflect.EnumType { + return &file_llmgw_v1_agent_proto_enumTypes[0] +} + +func (x ShardStatus_State) Number() protoreflect.EnumNumber { + return protoreflect.EnumNumber(x) +} + +// Deprecated: Use ShardStatus_State.Descriptor instead. +func (ShardStatus_State) EnumDescriptor() ([]byte, []int) { + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{6, 0} +} + +// AgentMessage is anything a node tells the control plane. +type AgentMessage struct { + state protoimpl.MessageState `protogen:"open.v1"` + // Types that are valid to be assigned to Msg: + // + // *AgentMessage_Register + // *AgentMessage_Heartbeat + // *AgentMessage_Shard + Msg isAgentMessage_Msg `protobuf_oneof:"msg"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *AgentMessage) Reset() { + *x = AgentMessage{} + mi := &file_llmgw_v1_agent_proto_msgTypes[0] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *AgentMessage) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*AgentMessage) ProtoMessage() {} + +func (x *AgentMessage) ProtoReflect() protoreflect.Message { + mi := &file_llmgw_v1_agent_proto_msgTypes[0] + 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 AgentMessage.ProtoReflect.Descriptor instead. +func (*AgentMessage) Descriptor() ([]byte, []int) { + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{0} +} + +func (x *AgentMessage) GetMsg() isAgentMessage_Msg { + if x != nil { + return x.Msg + } + return nil +} + +func (x *AgentMessage) GetRegister() *Register { + if x != nil { + if x, ok := x.Msg.(*AgentMessage_Register); ok { + return x.Register + } + } + return nil +} + +func (x *AgentMessage) GetHeartbeat() *Heartbeat { + if x != nil { + if x, ok := x.Msg.(*AgentMessage_Heartbeat); ok { + return x.Heartbeat + } + } + return nil +} + +func (x *AgentMessage) GetShard() *ShardStatus { + if x != nil { + if x, ok := x.Msg.(*AgentMessage_Shard); ok { + return x.Shard + } + } + return nil +} + +type isAgentMessage_Msg interface { + isAgentMessage_Msg() +} + +type AgentMessage_Register struct { + Register *Register `protobuf:"bytes,1,opt,name=register,proto3,oneof"` +} + +type AgentMessage_Heartbeat struct { + Heartbeat *Heartbeat `protobuf:"bytes,2,opt,name=heartbeat,proto3,oneof"` +} + +type AgentMessage_Shard struct { + Shard *ShardStatus `protobuf:"bytes,3,opt,name=shard,proto3,oneof"` +} + +func (*AgentMessage_Register) isAgentMessage_Msg() {} + +func (*AgentMessage_Heartbeat) isAgentMessage_Msg() {} + +func (*AgentMessage_Shard) isAgentMessage_Msg() {} + +// ControlMessage is anything the control plane tells a node. +type ControlMessage struct { + state protoimpl.MessageState `protogen:"open.v1"` + // Types that are valid to be assigned to Msg: + // + // *ControlMessage_Start + // *ControlMessage_Stop + // *ControlMessage_Drain + // *ControlMessage_Ack + Msg isControlMessage_Msg `protobuf_oneof:"msg"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ControlMessage) Reset() { + *x = ControlMessage{} + mi := &file_llmgw_v1_agent_proto_msgTypes[1] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ControlMessage) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ControlMessage) ProtoMessage() {} + +func (x *ControlMessage) 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 ControlMessage.ProtoReflect.Descriptor instead. +func (*ControlMessage) Descriptor() ([]byte, []int) { + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{1} +} + +func (x *ControlMessage) GetMsg() isControlMessage_Msg { + if x != nil { + return x.Msg + } + return nil +} + +func (x *ControlMessage) GetStart() *StartShard { + if x != nil { + if x, ok := x.Msg.(*ControlMessage_Start); ok { + return x.Start + } + } + return nil +} + +func (x *ControlMessage) GetStop() *StopShard { + if x != nil { + if x, ok := x.Msg.(*ControlMessage_Stop); ok { + return x.Stop + } + } + return nil +} + +func (x *ControlMessage) GetDrain() *DrainShard { + if x != nil { + if x, ok := x.Msg.(*ControlMessage_Drain); ok { + return x.Drain + } + } + return nil +} + +func (x *ControlMessage) GetAck() *Ack { + if x != nil { + if x, ok := x.Msg.(*ControlMessage_Ack); ok { + return x.Ack + } + } + return nil +} + +type isControlMessage_Msg interface { + isControlMessage_Msg() +} + +type ControlMessage_Start struct { + Start *StartShard `protobuf:"bytes,1,opt,name=start,proto3,oneof"` +} + +type ControlMessage_Stop struct { + Stop *StopShard `protobuf:"bytes,2,opt,name=stop,proto3,oneof"` +} + +type ControlMessage_Drain struct { + Drain *DrainShard `protobuf:"bytes,3,opt,name=drain,proto3,oneof"` +} + +type ControlMessage_Ack struct { + Ack *Ack `protobuf:"bytes,4,opt,name=ack,proto3,oneof"` +} + +func (*ControlMessage_Start) isControlMessage_Msg() {} + +func (*ControlMessage_Stop) isControlMessage_Msg() {} + +func (*ControlMessage_Drain) isControlMessage_Msg() {} + +func (*ControlMessage_Ack) isControlMessage_Msg() {} + +// Register is the first message on every stream. It carries the node's full +// hardware inventory, so the control plane needs no separate discovery call and +// a reconnecting agent automatically refreshes what the scheduler knows. +type Register struct { + state protoimpl.MessageState `protogen:"open.v1"` + NodeId string `protobuf:"bytes,1,opt,name=node_id,json=nodeId,proto3" json:"node_id,omitempty"` + Hostname string `protobuf:"bytes,2,opt,name=hostname,proto3" json:"hostname,omitempty"` + // mgmt_addr is the address engines are reachable on, and the address NCCL + // bootstraps over. Never a fabric address: NCCL rendezvouses over TCP before + // it moves to RDMA. + MgmtAddr string `protobuf:"bytes,3,opt,name=mgmt_addr,json=mgmtAddr,proto3" json:"mgmt_addr,omitempty"` + Labels map[string]string `protobuf:"bytes,4,rep,name=labels,proto3" json:"labels,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"` + Gpus []*GPU `protobuf:"bytes,5,rep,name=gpus,proto3" json:"gpus,omitempty"` + // fabrics is empty on a node without InfiniBand. Such a node is its own + // single-node island and can never host a multi-node replica. + Fabrics []*FabricAttachment `protobuf:"bytes,6,rep,name=fabrics,proto3" json:"fabrics,omitempty"` + AgentVersion string `protobuf:"bytes,7,opt,name=agent_version,json=agentVersion,proto3" json:"agent_version,omitempty"` + // engine_versions decides which CLI flags and metric names apply; both drift + // between engine releases (plan §14). + EngineVersions map[string]string `protobuf:"bytes,8,rep,name=engine_versions,json=engineVersions,proto3" json:"engine_versions,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *Register) Reset() { + *x = Register{} + mi := &file_llmgw_v1_agent_proto_msgTypes[2] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *Register) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*Register) ProtoMessage() {} + +func (x *Register) ProtoReflect() protoreflect.Message { + mi := &file_llmgw_v1_agent_proto_msgTypes[2] + 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 Register.ProtoReflect.Descriptor instead. +func (*Register) Descriptor() ([]byte, []int) { + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{2} +} + +func (x *Register) GetNodeId() string { + if x != nil { + return x.NodeId + } + return "" +} + +func (x *Register) GetHostname() string { + if x != nil { + return x.Hostname + } + return "" +} + +func (x *Register) GetMgmtAddr() string { + if x != nil { + return x.MgmtAddr + } + return "" +} + +func (x *Register) GetLabels() map[string]string { + if x != nil { + return x.Labels + } + return nil +} + +func (x *Register) GetGpus() []*GPU { + if x != nil { + return x.Gpus + } + return nil +} + +func (x *Register) GetFabrics() []*FabricAttachment { + if x != nil { + return x.Fabrics + } + return nil +} + +func (x *Register) GetAgentVersion() string { + if x != nil { + return x.AgentVersion + } + return "" +} + +func (x *Register) GetEngineVersions() map[string]string { + if x != nil { + return x.EngineVersions + } + return nil +} + +type GPU struct { + state protoimpl.MessageState `protogen:"open.v1"` + Index int32 `protobuf:"varint,1,opt,name=index,proto3" json:"index,omitempty"` + Uuid string `protobuf:"bytes,2,opt,name=uuid,proto3" json:"uuid,omitempty"` + Name string `protobuf:"bytes,3,opt,name=name,proto3" json:"name,omitempty"` + MemoryTotalMb int32 `protobuf:"varint,4,opt,name=memory_total_mb,json=memoryTotalMb,proto3" json:"memory_total_mb,omitempty"` + MemoryUsedMb int32 `protobuf:"varint,5,opt,name=memory_used_mb,json=memoryUsedMb,proto3" json:"memory_used_mb,omitempty"` + ComputeCapability string `protobuf:"bytes,6,opt,name=compute_capability,json=computeCapability,proto3" json:"compute_capability,omitempty"` + NumaNode int32 `protobuf:"varint,7,opt,name=numa_node,json=numaNode,proto3" json:"numa_node,omitempty"` + PciBusId string `protobuf:"bytes,8,opt,name=pci_bus_id,json=pciBusId,proto3" json:"pci_bus_id,omitempty"` + DriverVersion string `protobuf:"bytes,9,opt,name=driver_version,json=driverVersion,proto3" json:"driver_version,omitempty"` + CudaVersion string `protobuf:"bytes,10,opt,name=cuda_version,json=cudaVersion,proto3" json:"cuda_version,omitempty"` + // nvlink_peers is only populated by agents built with the nvml tag. + NvlinkPeers []int32 `protobuf:"varint,11,rep,packed,name=nvlink_peers,json=nvlinkPeers,proto3" json:"nvlink_peers,omitempty"` + // problems keeps a GPU out of placement: ECC faults, or memory already held + // by another process. + Problems []string `protobuf:"bytes,12,rep,name=problems,proto3" json:"problems,omitempty"` + TemperatureC int32 `protobuf:"varint,13,opt,name=temperature_c,json=temperatureC,proto3" json:"temperature_c,omitempty"` + EccErrors int32 `protobuf:"varint,14,opt,name=ecc_errors,json=eccErrors,proto3" json:"ecc_errors,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *GPU) Reset() { + *x = GPU{} + mi := &file_llmgw_v1_agent_proto_msgTypes[3] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *GPU) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*GPU) ProtoMessage() {} + +func (x *GPU) ProtoReflect() protoreflect.Message { + mi := &file_llmgw_v1_agent_proto_msgTypes[3] + 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 GPU.ProtoReflect.Descriptor instead. +func (*GPU) Descriptor() ([]byte, []int) { + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{3} +} + +func (x *GPU) GetIndex() int32 { + if x != nil { + return x.Index + } + return 0 +} + +func (x *GPU) GetUuid() string { + if x != nil { + return x.Uuid + } + return "" +} + +func (x *GPU) GetName() string { + if x != nil { + return x.Name + } + return "" +} + +func (x *GPU) GetMemoryTotalMb() int32 { + if x != nil { + return x.MemoryTotalMb + } + return 0 +} + +func (x *GPU) GetMemoryUsedMb() int32 { + if x != nil { + return x.MemoryUsedMb + } + return 0 +} + +func (x *GPU) GetComputeCapability() string { + if x != nil { + return x.ComputeCapability + } + return "" +} + +func (x *GPU) GetNumaNode() int32 { + if x != nil { + return x.NumaNode + } + return 0 +} + +func (x *GPU) GetPciBusId() string { + if x != nil { + return x.PciBusId + } + return "" +} + +func (x *GPU) GetDriverVersion() string { + if x != nil { + return x.DriverVersion + } + return "" +} + +func (x *GPU) GetCudaVersion() string { + if x != nil { + return x.CudaVersion + } + return "" +} + +func (x *GPU) GetNvlinkPeers() []int32 { + if x != nil { + return x.NvlinkPeers + } + return nil +} + +func (x *GPU) GetProblems() []string { + if x != nil { + return x.Problems + } + return nil +} + +func (x *GPU) GetTemperatureC() int32 { + if x != nil { + return x.TemperatureC + } + return 0 +} + +func (x *GPU) GetEccErrors() int32 { + if x != nil { + return x.EccErrors + } + return 0 +} + +type FabricAttachment struct { + state protoimpl.MessageState `protogen:"open.v1"` + // fabric is the operator-assigned name. It is not derived from the GID + // prefix: OpenSM defaults to fe80:: on every fabric, so a matching prefix + // proves nothing (plan §5.2). + Fabric string `protobuf:"bytes,1,opt,name=fabric,proto3" json:"fabric,omitempty"` + Hcas []string `protobuf:"bytes,2,rep,name=hcas,proto3" json:"hcas,omitempty"` + RateGbps int32 `protobuf:"varint,3,opt,name=rate_gbps,json=rateGbps,proto3" json:"rate_gbps,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *FabricAttachment) Reset() { + *x = FabricAttachment{} + mi := &file_llmgw_v1_agent_proto_msgTypes[4] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *FabricAttachment) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*FabricAttachment) ProtoMessage() {} + +func (x *FabricAttachment) ProtoReflect() protoreflect.Message { + mi := &file_llmgw_v1_agent_proto_msgTypes[4] + 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 FabricAttachment.ProtoReflect.Descriptor instead. +func (*FabricAttachment) Descriptor() ([]byte, []int) { + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{4} +} + +func (x *FabricAttachment) GetFabric() string { + if x != nil { + return x.Fabric + } + return "" +} + +func (x *FabricAttachment) GetHcas() []string { + if x != nil { + return x.Hcas + } + return nil +} + +func (x *FabricAttachment) GetRateGbps() int32 { + if x != nil { + return x.RateGbps + } + return 0 +} + +// Heartbeat says the node is alive and reports anything that changed cheaply. +type Heartbeat struct { + state protoimpl.MessageState `protogen:"open.v1"` + NodeId string `protobuf:"bytes,1,opt,name=node_id,json=nodeId,proto3" json:"node_id,omitempty"` + // gpu_memory_used_mb, indexed by GPU, catches a foreign process appearing on + // a card after registration. + GpuMemoryUsedMb []int32 `protobuf:"varint,2,rep,packed,name=gpu_memory_used_mb,json=gpuMemoryUsedMb,proto3" json:"gpu_memory_used_mb,omitempty"` + NewProblems []string `protobuf:"bytes,3,rep,name=new_problems,json=newProblems,proto3" json:"new_problems,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *Heartbeat) Reset() { + *x = Heartbeat{} + mi := &file_llmgw_v1_agent_proto_msgTypes[5] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *Heartbeat) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*Heartbeat) ProtoMessage() {} + +func (x *Heartbeat) ProtoReflect() protoreflect.Message { + mi := &file_llmgw_v1_agent_proto_msgTypes[5] + 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 Heartbeat.ProtoReflect.Descriptor instead. +func (*Heartbeat) Descriptor() ([]byte, []int) { + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{5} +} + +func (x *Heartbeat) GetNodeId() string { + if x != nil { + return x.NodeId + } + return "" +} + +func (x *Heartbeat) GetGpuMemoryUsedMb() []int32 { + if x != nil { + return x.GpuMemoryUsedMb + } + return nil +} + +func (x *Heartbeat) GetNewProblems() []string { + if x != nil { + return x.NewProblems + } + return nil +} + +// ShardStatus reports a shard's lifecycle transition. +type ShardStatus struct { + state protoimpl.MessageState `protogen:"open.v1"` + NodeId string `protobuf:"bytes,1,opt,name=node_id,json=nodeId,proto3" json:"node_id,omitempty"` + ReplicaId string `protobuf:"bytes,2,opt,name=replica_id,json=replicaId,proto3" json:"replica_id,omitempty"` + ShardId string `protobuf:"bytes,3,opt,name=shard_id,json=shardId,proto3" json:"shard_id,omitempty"` + State ShardStatus_State `protobuf:"varint,4,opt,name=state,proto3,enum=llmgw.v1.ShardStatus_State" json:"state,omitempty"` + Error string `protobuf:"bytes,5,opt,name=error,proto3" json:"error,omitempty"` + ExitCode int32 `protobuf:"varint,6,opt,name=exit_code,json=exitCode,proto3" json:"exit_code,omitempty"` + // last_output carries the tail of the engine log on a failure. When a model + // fails to load this is the only diagnosis available. + LastOutput []string `protobuf:"bytes,7,rep,name=last_output,json=lastOutput,proto3" json:"last_output,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *ShardStatus) Reset() { + *x = ShardStatus{} + mi := &file_llmgw_v1_agent_proto_msgTypes[6] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *ShardStatus) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*ShardStatus) ProtoMessage() {} + +func (x *ShardStatus) ProtoReflect() protoreflect.Message { + mi := &file_llmgw_v1_agent_proto_msgTypes[6] + 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 ShardStatus.ProtoReflect.Descriptor instead. +func (*ShardStatus) Descriptor() ([]byte, []int) { + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{6} +} + +func (x *ShardStatus) GetNodeId() string { + if x != nil { + return x.NodeId + } + return "" +} + +func (x *ShardStatus) GetReplicaId() string { + if x != nil { + return x.ReplicaId + } + return "" +} + +func (x *ShardStatus) GetShardId() string { + if x != nil { + return x.ShardId + } + return "" +} + +func (x *ShardStatus) GetState() ShardStatus_State { + if x != nil { + return x.State + } + return ShardStatus_STATE_UNSPECIFIED +} + +func (x *ShardStatus) GetError() string { + if x != nil { + return x.Error + } + return "" +} + +func (x *ShardStatus) GetExitCode() int32 { + if x != nil { + return x.ExitCode + } + return 0 +} + +func (x *ShardStatus) GetLastOutput() []string { + if x != nil { + return x.LastOutput + } + return nil +} + +// StartShard launches one shard. It is idempotent by shard_id: the control +// plane retries, and a repeat must not produce a second process. +type StartShard struct { + state protoimpl.MessageState `protogen:"open.v1"` + ReplicaId string `protobuf:"bytes,1,opt,name=replica_id,json=replicaId,proto3" json:"replica_id,omitempty"` + ShardId string `protobuf:"bytes,2,opt,name=shard_id,json=shardId,proto3" json:"shard_id,omitempty"` + Rank int32 `protobuf:"varint,3,opt,name=rank,proto3" json:"rank,omitempty"` + Nnodes int32 `protobuf:"varint,4,opt,name=nnodes,proto3" json:"nnodes,omitempty"` + LeaderAddr string `protobuf:"bytes,5,opt,name=leader_addr,json=leaderAddr,proto3" json:"leader_addr,omitempty"` + LeaderPort int32 `protobuf:"varint,6,opt,name=leader_port,json=leaderPort,proto3" json:"leader_port,omitempty"` + GpuIndices []int32 `protobuf:"varint,7,rep,packed,name=gpu_indices,json=gpuIndices,proto3" json:"gpu_indices,omitempty"` + // hcas are the adapters of the chosen fabric and nothing else. Handing a + // shard adapters from a second fabric lets NCCL pick one whose peer is + // unreachable, and the job hangs instead of failing (plan §5.3). + Hcas []string `protobuf:"bytes,8,rep,name=hcas,proto3" json:"hcas,omitempty"` + MgmtInterface string `protobuf:"bytes,9,opt,name=mgmt_interface,json=mgmtInterface,proto3" json:"mgmt_interface,omitempty"` + HttpPort int32 `protobuf:"varint,10,opt,name=http_port,json=httpPort,proto3" json:"http_port,omitempty"` + Profile *EngineProfile `protobuf:"bytes,11,opt,name=profile,proto3" json:"profile,omitempty"` + StartupTimeoutMs int64 `protobuf:"varint,12,opt,name=startup_timeout_ms,json=startupTimeoutMs,proto3" json:"startup_timeout_ms,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *StartShard) Reset() { + *x = StartShard{} + mi := &file_llmgw_v1_agent_proto_msgTypes[7] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *StartShard) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*StartShard) ProtoMessage() {} + +func (x *StartShard) ProtoReflect() protoreflect.Message { + mi := &file_llmgw_v1_agent_proto_msgTypes[7] + 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 StartShard.ProtoReflect.Descriptor instead. +func (*StartShard) Descriptor() ([]byte, []int) { + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{7} +} + +func (x *StartShard) GetReplicaId() string { + if x != nil { + return x.ReplicaId + } + return "" +} + +func (x *StartShard) GetShardId() string { + if x != nil { + return x.ShardId + } + return "" +} + +func (x *StartShard) GetRank() int32 { + if x != nil { + return x.Rank + } + return 0 +} + +func (x *StartShard) GetNnodes() int32 { + if x != nil { + return x.Nnodes + } + return 0 +} + +func (x *StartShard) GetLeaderAddr() string { + if x != nil { + return x.LeaderAddr + } + return "" +} + +func (x *StartShard) GetLeaderPort() int32 { + if x != nil { + return x.LeaderPort + } + return 0 +} + +func (x *StartShard) GetGpuIndices() []int32 { + if x != nil { + return x.GpuIndices + } + return nil +} + +func (x *StartShard) GetHcas() []string { + if x != nil { + return x.Hcas + } + return nil +} + +func (x *StartShard) GetMgmtInterface() string { + if x != nil { + return x.MgmtInterface + } + return "" +} + +func (x *StartShard) GetHttpPort() int32 { + if x != nil { + return x.HttpPort + } + return 0 +} + +func (x *StartShard) GetProfile() *EngineProfile { + if x != nil { + return x.Profile + } + return nil +} + +func (x *StartShard) GetStartupTimeoutMs() int64 { + if x != nil { + return x.StartupTimeoutMs + } + return 0 +} + +type EngineProfile struct { + state protoimpl.MessageState `protogen:"open.v1"` + Name string `protobuf:"bytes,1,opt,name=name,proto3" json:"name,omitempty"` + Engine string `protobuf:"bytes,2,opt,name=engine,proto3" json:"engine,omitempty"` + GpuClass string `protobuf:"bytes,3,opt,name=gpu_class,json=gpuClass,proto3" json:"gpu_class,omitempty"` + Tp int32 `protobuf:"varint,4,opt,name=tp,proto3" json:"tp,omitempty"` + Pp int32 `protobuf:"varint,5,opt,name=pp,proto3" json:"pp,omitempty"` + GpusPerNode int32 `protobuf:"varint,6,opt,name=gpus_per_node,json=gpusPerNode,proto3" json:"gpus_per_node,omitempty"` + Nodes int32 `protobuf:"varint,7,opt,name=nodes,proto3" json:"nodes,omitempty"` + Args []string `protobuf:"bytes,8,rep,name=args,proto3" json:"args,omitempty"` + Env map[string]string `protobuf:"bytes,9,rep,name=env,proto3" json:"env,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"` + ModelPath string `protobuf:"bytes,10,opt,name=model_path,json=modelPath,proto3" json:"model_path,omitempty"` + ServedModelName string `protobuf:"bytes,11,opt,name=served_model_name,json=servedModelName,proto3" json:"served_model_name,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *EngineProfile) Reset() { + *x = EngineProfile{} + mi := &file_llmgw_v1_agent_proto_msgTypes[8] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *EngineProfile) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*EngineProfile) ProtoMessage() {} + +func (x *EngineProfile) ProtoReflect() protoreflect.Message { + mi := &file_llmgw_v1_agent_proto_msgTypes[8] + 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 EngineProfile.ProtoReflect.Descriptor instead. +func (*EngineProfile) Descriptor() ([]byte, []int) { + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{8} +} + +func (x *EngineProfile) GetName() string { + if x != nil { + return x.Name + } + return "" +} + +func (x *EngineProfile) GetEngine() string { + if x != nil { + return x.Engine + } + return "" +} + +func (x *EngineProfile) GetGpuClass() string { + if x != nil { + return x.GpuClass + } + return "" +} + +func (x *EngineProfile) GetTp() int32 { + if x != nil { + return x.Tp + } + return 0 +} + +func (x *EngineProfile) GetPp() int32 { + if x != nil { + return x.Pp + } + return 0 +} + +func (x *EngineProfile) GetGpusPerNode() int32 { + if x != nil { + return x.GpusPerNode + } + return 0 +} + +func (x *EngineProfile) GetNodes() int32 { + if x != nil { + return x.Nodes + } + return 0 +} + +func (x *EngineProfile) GetArgs() []string { + if x != nil { + return x.Args + } + return nil +} + +func (x *EngineProfile) GetEnv() map[string]string { + if x != nil { + return x.Env + } + return nil +} + +func (x *EngineProfile) GetModelPath() string { + if x != nil { + return x.ModelPath + } + return "" +} + +func (x *EngineProfile) GetServedModelName() string { + if x != nil { + return x.ServedModelName + } + return "" +} + +// StopShard terminates a shard and frees its GPUs. +type StopShard struct { + state protoimpl.MessageState `protogen:"open.v1"` + ShardId string `protobuf:"bytes,1,opt,name=shard_id,json=shardId,proto3" json:"shard_id,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *StopShard) Reset() { + *x = StopShard{} + mi := &file_llmgw_v1_agent_proto_msgTypes[9] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *StopShard) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*StopShard) ProtoMessage() {} + +func (x *StopShard) ProtoReflect() protoreflect.Message { + mi := &file_llmgw_v1_agent_proto_msgTypes[9] + 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 StopShard.ProtoReflect.Descriptor instead. +func (*StopShard) Descriptor() ([]byte, []int) { + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{9} +} + +func (x *StopShard) GetShardId() string { + if x != nil { + return x.ShardId + } + return "" +} + +// DrainShard stops routing to a shard while letting in-flight work finish. +type DrainShard struct { + state protoimpl.MessageState `protogen:"open.v1"` + ShardId string `protobuf:"bytes,1,opt,name=shard_id,json=shardId,proto3" json:"shard_id,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *DrainShard) Reset() { + *x = DrainShard{} + mi := &file_llmgw_v1_agent_proto_msgTypes[10] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *DrainShard) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*DrainShard) ProtoMessage() {} + +func (x *DrainShard) ProtoReflect() protoreflect.Message { + mi := &file_llmgw_v1_agent_proto_msgTypes[10] + 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 DrainShard.ProtoReflect.Descriptor instead. +func (*DrainShard) Descriptor() ([]byte, []int) { + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{10} +} + +func (x *DrainShard) GetShardId() string { + if x != nil { + return x.ShardId + } + return "" +} + +// Ack confirms a registration or reports that a command could not be accepted. +type Ack struct { + state protoimpl.MessageState `protogen:"open.v1"` + MessageId string `protobuf:"bytes,1,opt,name=message_id,json=messageId,proto3" json:"message_id,omitempty"` + Ok bool `protobuf:"varint,2,opt,name=ok,proto3" json:"ok,omitempty"` + Error string `protobuf:"bytes,3,opt,name=error,proto3" json:"error,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *Ack) Reset() { + *x = Ack{} + mi := &file_llmgw_v1_agent_proto_msgTypes[11] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *Ack) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*Ack) ProtoMessage() {} + +func (x *Ack) ProtoReflect() protoreflect.Message { + mi := &file_llmgw_v1_agent_proto_msgTypes[11] + 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 Ack.ProtoReflect.Descriptor instead. +func (*Ack) Descriptor() ([]byte, []int) { + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{11} +} + +func (x *Ack) GetMessageId() string { + if x != nil { + return x.MessageId + } + return "" +} + +func (x *Ack) GetOk() bool { + if x != nil { + return x.Ok + } + return false +} + +func (x *Ack) GetError() string { + if x != nil { + return x.Error + } + return "" +} + +var File_llmgw_v1_agent_proto protoreflect.FileDescriptor + +const file_llmgw_v1_agent_proto_rawDesc = "" + + "\n" + + "\x14llmgw/v1/agent.proto\x12\bllmgw.v1\"\xab\x01\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" + + "\x05shard\x18\x03 \x01(\v2\x15.llmgw.v1.ShardStatusH\x00R\x05shardB\x05\n" + + "\x03msg\"\xc1\x01\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" + + "\x05drain\x18\x03 \x01(\v2\x14.llmgw.v1.DrainShardH\x00R\x05drain\x12!\n" + + "\x03ack\x18\x04 \x01(\v2\r.llmgw.v1.AckH\x00R\x03ackB\x05\n" + + "\x03msg\"\xe1\x03\n" + + "\bRegister\x12\x17\n" + + "\anode_id\x18\x01 \x01(\tR\x06nodeId\x12\x1a\n" + + "\bhostname\x18\x02 \x01(\tR\bhostname\x12\x1b\n" + + "\tmgmt_addr\x18\x03 \x01(\tR\bmgmtAddr\x126\n" + + "\x06labels\x18\x04 \x03(\v2\x1e.llmgw.v1.Register.LabelsEntryR\x06labels\x12!\n" + + "\x04gpus\x18\x05 \x03(\v2\r.llmgw.v1.GPUR\x04gpus\x124\n" + + "\afabrics\x18\x06 \x03(\v2\x1a.llmgw.v1.FabricAttachmentR\afabrics\x12#\n" + + "\ragent_version\x18\a \x01(\tR\fagentVersion\x12O\n" + + "\x0fengine_versions\x18\b \x03(\v2&.llmgw.v1.Register.EngineVersionsEntryR\x0eengineVersions\x1a9\n" + + "\vLabelsEntry\x12\x10\n" + + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + + "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\x1aA\n" + + "\x13EngineVersionsEntry\x12\x10\n" + + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + + "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"\xc8\x03\n" + + "\x03GPU\x12\x14\n" + + "\x05index\x18\x01 \x01(\x05R\x05index\x12\x12\n" + + "\x04uuid\x18\x02 \x01(\tR\x04uuid\x12\x12\n" + + "\x04name\x18\x03 \x01(\tR\x04name\x12&\n" + + "\x0fmemory_total_mb\x18\x04 \x01(\x05R\rmemoryTotalMb\x12$\n" + + "\x0ememory_used_mb\x18\x05 \x01(\x05R\fmemoryUsedMb\x12-\n" + + "\x12compute_capability\x18\x06 \x01(\tR\x11computeCapability\x12\x1b\n" + + "\tnuma_node\x18\a \x01(\x05R\bnumaNode\x12\x1c\n" + + "\n" + + "pci_bus_id\x18\b \x01(\tR\bpciBusId\x12%\n" + + "\x0edriver_version\x18\t \x01(\tR\rdriverVersion\x12!\n" + + "\fcuda_version\x18\n" + + " \x01(\tR\vcudaVersion\x12!\n" + + "\fnvlink_peers\x18\v \x03(\x05R\vnvlinkPeers\x12\x1a\n" + + "\bproblems\x18\f \x03(\tR\bproblems\x12#\n" + + "\rtemperature_c\x18\r \x01(\x05R\ftemperatureC\x12\x1d\n" + + "\n" + + "ecc_errors\x18\x0e \x01(\x05R\teccErrors\"[\n" + + "\x10FabricAttachment\x12\x16\n" + + "\x06fabric\x18\x01 \x01(\tR\x06fabric\x12\x12\n" + + "\x04hcas\x18\x02 \x03(\tR\x04hcas\x12\x1b\n" + + "\trate_gbps\x18\x03 \x01(\x05R\brateGbps\"t\n" + + "\tHeartbeat\x12\x17\n" + + "\anode_id\x18\x01 \x01(\tR\x06nodeId\x12+\n" + + "\x12gpu_memory_used_mb\x18\x02 \x03(\x05R\x0fgpuMemoryUsedMb\x12!\n" + + "\fnew_problems\x18\x03 \x03(\tR\vnewProblems\"\x8c\x03\n" + + "\vShardStatus\x12\x17\n" + + "\anode_id\x18\x01 \x01(\tR\x06nodeId\x12\x1d\n" + + "\n" + + "replica_id\x18\x02 \x01(\tR\treplicaId\x12\x19\n" + + "\bshard_id\x18\x03 \x01(\tR\ashardId\x121\n" + + "\x05state\x18\x04 \x01(\x0e2\x1b.llmgw.v1.ShardStatus.StateR\x05state\x12\x14\n" + + "\x05error\x18\x05 \x01(\tR\x05error\x12\x1b\n" + + "\texit_code\x18\x06 \x01(\x05R\bexitCode\x12\x1f\n" + + "\vlast_output\x18\a \x03(\tR\n" + + "lastOutput\"\xa2\x01\n" + + "\x05State\x12\x15\n" + + "\x11STATE_UNSPECIFIED\x10\x00\x12\x11\n" + + "\rSTATE_PENDING\x10\x01\x12\x12\n" + + "\x0eSTATE_STARTING\x10\x02\x12\x11\n" + + "\rSTATE_LOADING\x10\x03\x12\x0f\n" + + "\vSTATE_READY\x10\x04\x12\x12\n" + + "\x0eSTATE_DRAINING\x10\x05\x12\x11\n" + + "\rSTATE_STOPPED\x10\x06\x12\x10\n" + + "\fSTATE_FAILED\x10\a\"\x8e\x03\n" + + "\n" + + "StartShard\x12\x1d\n" + + "\n" + + "replica_id\x18\x01 \x01(\tR\treplicaId\x12\x19\n" + + "\bshard_id\x18\x02 \x01(\tR\ashardId\x12\x12\n" + + "\x04rank\x18\x03 \x01(\x05R\x04rank\x12\x16\n" + + "\x06nnodes\x18\x04 \x01(\x05R\x06nnodes\x12\x1f\n" + + "\vleader_addr\x18\x05 \x01(\tR\n" + + "leaderAddr\x12\x1f\n" + + "\vleader_port\x18\x06 \x01(\x05R\n" + + "leaderPort\x12\x1f\n" + + "\vgpu_indices\x18\a \x03(\x05R\n" + + "gpuIndices\x12\x12\n" + + "\x04hcas\x18\b \x03(\tR\x04hcas\x12%\n" + + "\x0emgmt_interface\x18\t \x01(\tR\rmgmtInterface\x12\x1b\n" + + "\thttp_port\x18\n" + + " \x01(\x05R\bhttpPort\x121\n" + + "\aprofile\x18\v \x01(\v2\x17.llmgw.v1.EngineProfileR\aprofile\x12,\n" + + "\x12startup_timeout_ms\x18\f \x01(\x03R\x10startupTimeoutMs\"\xfd\x02\n" + + "\rEngineProfile\x12\x12\n" + + "\x04name\x18\x01 \x01(\tR\x04name\x12\x16\n" + + "\x06engine\x18\x02 \x01(\tR\x06engine\x12\x1b\n" + + "\tgpu_class\x18\x03 \x01(\tR\bgpuClass\x12\x0e\n" + + "\x02tp\x18\x04 \x01(\x05R\x02tp\x12\x0e\n" + + "\x02pp\x18\x05 \x01(\x05R\x02pp\x12\"\n" + + "\rgpus_per_node\x18\x06 \x01(\x05R\vgpusPerNode\x12\x14\n" + + "\x05nodes\x18\a \x01(\x05R\x05nodes\x12\x12\n" + + "\x04args\x18\b \x03(\tR\x04args\x122\n" + + "\x03env\x18\t \x03(\v2 .llmgw.v1.EngineProfile.EnvEntryR\x03env\x12\x1d\n" + + "\n" + + "model_path\x18\n" + + " \x01(\tR\tmodelPath\x12*\n" + + "\x11served_model_name\x18\v \x01(\tR\x0fservedModelName\x1a6\n" + + "\bEnvEntry\x12\x10\n" + + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + + "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"&\n" + + "\tStopShard\x12\x19\n" + + "\bshard_id\x18\x01 \x01(\tR\ashardId\"'\n" + + "\n" + + "DrainShard\x12\x19\n" + + "\bshard_id\x18\x01 \x01(\tR\ashardId\"J\n" + + "\x03Ack\x12\x1d\n" + + "\n" + + "message_id\x18\x01 \x01(\tR\tmessageId\x12\x0e\n" + + "\x02ok\x18\x02 \x01(\bR\x02ok\x12\x14\n" + + "\x05error\x18\x03 \x01(\tR\x05error2O\n" + + "\fAgentService\x12?\n" + + "\aConnect\x12\x16.llmgw.v1.AgentMessage\x1a\x18.llmgw.v1.ControlMessage(\x010\x01B\x93\x01\n" + + "\fcom.llmgw.v1B\n" + + "AgentProtoP\x01Z6gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1;llmgwv1\xa2\x02\x03LXX\xaa\x02\bLlmgw.V1\xca\x02\bLlmgw\\V1\xe2\x02\x14Llmgw\\V1\\GPBMetadata\xea\x02\tLlmgw::V1b\x06proto3" + +var ( + file_llmgw_v1_agent_proto_rawDescOnce sync.Once + file_llmgw_v1_agent_proto_rawDescData []byte +) + +func file_llmgw_v1_agent_proto_rawDescGZIP() []byte { + file_llmgw_v1_agent_proto_rawDescOnce.Do(func() { + file_llmgw_v1_agent_proto_rawDescData = protoimpl.X.CompressGZIP(unsafe.Slice(unsafe.StringData(file_llmgw_v1_agent_proto_rawDesc), len(file_llmgw_v1_agent_proto_rawDesc))) + }) + return file_llmgw_v1_agent_proto_rawDescData +} + +var file_llmgw_v1_agent_proto_enumTypes = make([]protoimpl.EnumInfo, 1) +var file_llmgw_v1_agent_proto_msgTypes = make([]protoimpl.MessageInfo, 15) +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 + (*Register)(nil), // 3: llmgw.v1.Register + (*GPU)(nil), // 4: llmgw.v1.GPU + (*FabricAttachment)(nil), // 5: llmgw.v1.FabricAttachment + (*Heartbeat)(nil), // 6: llmgw.v1.Heartbeat + (*ShardStatus)(nil), // 7: llmgw.v1.ShardStatus + (*StartShard)(nil), // 8: llmgw.v1.StartShard + (*EngineProfile)(nil), // 9: llmgw.v1.EngineProfile + (*StopShard)(nil), // 10: llmgw.v1.StopShard + (*DrainShard)(nil), // 11: llmgw.v1.DrainShard + (*Ack)(nil), // 12: llmgw.v1.Ack + nil, // 13: llmgw.v1.Register.LabelsEntry + nil, // 14: llmgw.v1.Register.EngineVersionsEntry + nil, // 15: llmgw.v1.EngineProfile.EnvEntry +} +var file_llmgw_v1_agent_proto_depIdxs = []int32{ + 3, // 0: llmgw.v1.AgentMessage.register:type_name -> llmgw.v1.Register + 6, // 1: llmgw.v1.AgentMessage.heartbeat:type_name -> llmgw.v1.Heartbeat + 7, // 2: llmgw.v1.AgentMessage.shard:type_name -> llmgw.v1.ShardStatus + 8, // 3: llmgw.v1.ControlMessage.start:type_name -> llmgw.v1.StartShard + 10, // 4: llmgw.v1.ControlMessage.stop:type_name -> llmgw.v1.StopShard + 11, // 5: llmgw.v1.ControlMessage.drain:type_name -> llmgw.v1.DrainShard + 12, // 6: llmgw.v1.ControlMessage.ack:type_name -> llmgw.v1.Ack + 13, // 7: llmgw.v1.Register.labels:type_name -> llmgw.v1.Register.LabelsEntry + 4, // 8: llmgw.v1.Register.gpus:type_name -> llmgw.v1.GPU + 5, // 9: llmgw.v1.Register.fabrics:type_name -> llmgw.v1.FabricAttachment + 14, // 10: llmgw.v1.Register.engine_versions:type_name -> llmgw.v1.Register.EngineVersionsEntry + 0, // 11: llmgw.v1.ShardStatus.state:type_name -> llmgw.v1.ShardStatus.State + 9, // 12: llmgw.v1.StartShard.profile:type_name -> llmgw.v1.EngineProfile + 15, // 13: llmgw.v1.EngineProfile.env:type_name -> llmgw.v1.EngineProfile.EnvEntry + 1, // 14: llmgw.v1.AgentService.Connect:input_type -> llmgw.v1.AgentMessage + 2, // 15: llmgw.v1.AgentService.Connect:output_type -> llmgw.v1.ControlMessage + 15, // [15:16] is the sub-list for method output_type + 14, // [14:15] is the sub-list for method input_type + 14, // [14:14] is the sub-list for extension type_name + 14, // [14:14] is the sub-list for extension extendee + 0, // [0:14] is the sub-list for field type_name +} + +func init() { file_llmgw_v1_agent_proto_init() } +func file_llmgw_v1_agent_proto_init() { + if File_llmgw_v1_agent_proto != nil { + return + } + file_llmgw_v1_agent_proto_msgTypes[0].OneofWrappers = []any{ + (*AgentMessage_Register)(nil), + (*AgentMessage_Heartbeat)(nil), + (*AgentMessage_Shard)(nil), + } + file_llmgw_v1_agent_proto_msgTypes[1].OneofWrappers = []any{ + (*ControlMessage_Start)(nil), + (*ControlMessage_Stop)(nil), + (*ControlMessage_Drain)(nil), + (*ControlMessage_Ack)(nil), + } + type x struct{} + out := protoimpl.TypeBuilder{ + File: protoimpl.DescBuilder{ + 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: 15, + NumExtensions: 0, + NumServices: 1, + }, + GoTypes: file_llmgw_v1_agent_proto_goTypes, + DependencyIndexes: file_llmgw_v1_agent_proto_depIdxs, + EnumInfos: file_llmgw_v1_agent_proto_enumTypes, + MessageInfos: file_llmgw_v1_agent_proto_msgTypes, + }.Build() + File_llmgw_v1_agent_proto = out.File + file_llmgw_v1_agent_proto_goTypes = nil + file_llmgw_v1_agent_proto_depIdxs = nil +} diff --git a/api/gen/llmgw/v1/llmgwv1connect/agent.connect.go b/api/gen/llmgw/v1/llmgwv1connect/agent.connect.go new file mode 100644 index 0000000..fc169a6 --- /dev/null +++ b/api/gen/llmgw/v1/llmgwv1connect/agent.connect.go @@ -0,0 +1,119 @@ +// Code generated by protoc-gen-connect-go. DO NOT EDIT. +// +// Source: llmgw/v1/agent.proto + +// Package llmgw.v1 is the protocol between a node agent and the control plane +// (plan §5.6). +// +// The agent dials out and holds a bidirectional stream open. That direction is +// deliberate: GPU nodes sit behind firewalls and NAT far more often than a +// control plane does, and an outbound connection needs no inbound rule on every +// node in the fleet. +package llmgwv1connect + +import ( + connect "connectrpc.com/connect" + context "context" + errors "errors" + v1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1" + http "net/http" + strings "strings" +) + +// This is a compile-time assertion to ensure that this generated file and the connect package are +// compatible. If you get a compiler error that this constant is not defined, this code was +// generated with a version of connect newer than the one compiled into your binary. You can fix the +// problem by either regenerating this code with an older version of connect or updating the connect +// version compiled into your binary. +const _ = connect.IsAtLeastVersion1_13_0 + +const ( + // AgentServiceName is the fully-qualified name of the AgentService service. + AgentServiceName = "llmgw.v1.AgentService" +) + +// These constants are the fully-qualified names of the RPCs defined in this package. They're +// exposed at runtime as Spec.Procedure and as the final two segments of the HTTP route. +// +// Note that these are different from the fully-qualified method names used by +// google.golang.org/protobuf/reflect/protoreflect. To convert from these constants to +// reflection-formatted method names, remove the leading slash and convert the remaining slash to a +// period. +const ( + // AgentServiceConnectProcedure is the fully-qualified name of the AgentService's Connect RPC. + AgentServiceConnectProcedure = "/llmgw.v1.AgentService/Connect" +) + +// AgentServiceClient is a client for the llmgw.v1.AgentService service. +type AgentServiceClient interface { + // Connect opens the agent's stream. It stays open for the agent's lifetime; + // losing it is how each side learns the other is gone. + Connect(context.Context) *connect.BidiStreamForClient[v1.AgentMessage, v1.ControlMessage] +} + +// NewAgentServiceClient constructs a client for the llmgw.v1.AgentService service. By default, it +// uses the Connect protocol with the binary Protobuf Codec, asks for gzipped responses, and sends +// uncompressed requests. To use the gRPC or gRPC-Web protocols, supply the connect.WithGRPC() or +// connect.WithGRPCWeb() options. +// +// The URL supplied here should be the base URL for the Connect or gRPC server (for example, +// http://api.acme.com or https://acme.com/grpc). +func NewAgentServiceClient(httpClient connect.HTTPClient, baseURL string, opts ...connect.ClientOption) AgentServiceClient { + baseURL = strings.TrimRight(baseURL, "/") + agentServiceMethods := v1.File_llmgw_v1_agent_proto.Services().ByName("AgentService").Methods() + return &agentServiceClient{ + connect: connect.NewClient[v1.AgentMessage, v1.ControlMessage]( + httpClient, + baseURL+AgentServiceConnectProcedure, + connect.WithSchema(agentServiceMethods.ByName("Connect")), + connect.WithClientOptions(opts...), + ), + } +} + +// agentServiceClient implements AgentServiceClient. +type agentServiceClient struct { + connect *connect.Client[v1.AgentMessage, v1.ControlMessage] +} + +// Connect calls llmgw.v1.AgentService.Connect. +func (c *agentServiceClient) Connect(ctx context.Context) *connect.BidiStreamForClient[v1.AgentMessage, v1.ControlMessage] { + return c.connect.CallBidiStream(ctx) +} + +// AgentServiceHandler is an implementation of the llmgw.v1.AgentService service. +type AgentServiceHandler interface { + // Connect opens the agent's stream. It stays open for the agent's lifetime; + // losing it is how each side learns the other is gone. + Connect(context.Context, *connect.BidiStream[v1.AgentMessage, v1.ControlMessage]) error +} + +// NewAgentServiceHandler builds an HTTP handler from the service implementation. It returns the +// path on which to mount the handler and the handler itself. +// +// By default, handlers support the Connect, gRPC, and gRPC-Web protocols with the binary Protobuf +// and JSON codecs. They also support gzip compression. +func NewAgentServiceHandler(svc AgentServiceHandler, opts ...connect.HandlerOption) (string, http.Handler) { + agentServiceMethods := v1.File_llmgw_v1_agent_proto.Services().ByName("AgentService").Methods() + agentServiceConnectHandler := connect.NewBidiStreamHandler( + AgentServiceConnectProcedure, + svc.Connect, + connect.WithSchema(agentServiceMethods.ByName("Connect")), + connect.WithHandlerOptions(opts...), + ) + return "/llmgw.v1.AgentService/", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case AgentServiceConnectProcedure: + agentServiceConnectHandler.ServeHTTP(w, r) + default: + http.NotFound(w, r) + } + }) +} + +// UnimplementedAgentServiceHandler returns CodeUnimplemented from all methods. +type UnimplementedAgentServiceHandler struct{} + +func (UnimplementedAgentServiceHandler) Connect(context.Context, *connect.BidiStream[v1.AgentMessage, v1.ControlMessage]) error { + return connect.NewError(connect.CodeUnimplemented, errors.New("llmgw.v1.AgentService.Connect is not implemented")) +} diff --git a/api/proto/llmgw/v1/agent.proto b/api/proto/llmgw/v1/agent.proto new file mode 100644 index 0000000..864ad89 --- /dev/null +++ b/api/proto/llmgw/v1/agent.proto @@ -0,0 +1,176 @@ +syntax = "proto3"; + +// Package llmgw.v1 is the protocol between a node agent and the control plane +// (plan §5.6). +// +// The agent dials out and holds a bidirectional stream open. That direction is +// deliberate: GPU nodes sit behind firewalls and NAT far more often than a +// control plane does, and an outbound connection needs no inbound rule on every +// node in the fleet. +package llmgw.v1; + +option go_package = "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1;llmgwv1"; + +service AgentService { + // Connect opens the agent's stream. It stays open for the agent's lifetime; + // losing it is how each side learns the other is gone. + rpc Connect(stream AgentMessage) returns (stream ControlMessage); +} + +// AgentMessage is anything a node tells the control plane. +message AgentMessage { + oneof msg { + Register register = 1; + Heartbeat heartbeat = 2; + ShardStatus shard = 3; + } +} + +// ControlMessage is anything the control plane tells a node. +message ControlMessage { + oneof msg { + StartShard start = 1; + StopShard stop = 2; + DrainShard drain = 3; + Ack ack = 4; + } +} + +// Register is the first message on every stream. It carries the node's full +// hardware inventory, so the control plane needs no separate discovery call and +// a reconnecting agent automatically refreshes what the scheduler knows. +message Register { + string node_id = 1; + string hostname = 2; + // mgmt_addr is the address engines are reachable on, and the address NCCL + // bootstraps over. Never a fabric address: NCCL rendezvouses over TCP before + // it moves to RDMA. + string mgmt_addr = 3; + map labels = 4; + + repeated GPU gpus = 5; + // fabrics is empty on a node without InfiniBand. Such a node is its own + // single-node island and can never host a multi-node replica. + repeated FabricAttachment fabrics = 6; + + string agent_version = 7; + // engine_versions decides which CLI flags and metric names apply; both drift + // between engine releases (plan §14). + map engine_versions = 8; +} + +message GPU { + int32 index = 1; + string uuid = 2; + string name = 3; + int32 memory_total_mb = 4; + int32 memory_used_mb = 5; + string compute_capability = 6; + int32 numa_node = 7; + string pci_bus_id = 8; + string driver_version = 9; + string cuda_version = 10; + // nvlink_peers is only populated by agents built with the nvml tag. + repeated int32 nvlink_peers = 11; + // problems keeps a GPU out of placement: ECC faults, or memory already held + // by another process. + repeated string problems = 12; + int32 temperature_c = 13; + int32 ecc_errors = 14; +} + +message FabricAttachment { + // fabric is the operator-assigned name. It is not derived from the GID + // prefix: OpenSM defaults to fe80:: on every fabric, so a matching prefix + // proves nothing (plan §5.2). + string fabric = 1; + repeated string hcas = 2; + int32 rate_gbps = 3; +} + +// Heartbeat says the node is alive and reports anything that changed cheaply. +message Heartbeat { + string node_id = 1; + // gpu_memory_used_mb, indexed by GPU, catches a foreign process appearing on + // a card after registration. + repeated int32 gpu_memory_used_mb = 2; + repeated string new_problems = 3; +} + +// ShardStatus reports a shard's lifecycle transition. +message ShardStatus { + string node_id = 1; + string replica_id = 2; + string shard_id = 3; + State state = 4; + string error = 5; + int32 exit_code = 6; + // last_output carries the tail of the engine log on a failure. When a model + // fails to load this is the only diagnosis available. + repeated string last_output = 7; + + enum State { + STATE_UNSPECIFIED = 0; + STATE_PENDING = 1; + STATE_STARTING = 2; + STATE_LOADING = 3; + STATE_READY = 4; + STATE_DRAINING = 5; + STATE_STOPPED = 6; + STATE_FAILED = 7; + } +} + +// StartShard launches one shard. It is idempotent by shard_id: the control +// plane retries, and a repeat must not produce a second process. +message StartShard { + string replica_id = 1; + string shard_id = 2; + + int32 rank = 3; + int32 nnodes = 4; + string leader_addr = 5; + int32 leader_port = 6; + + repeated int32 gpu_indices = 7; + // hcas are the adapters of the chosen fabric and nothing else. Handing a + // shard adapters from a second fabric lets NCCL pick one whose peer is + // unreachable, and the job hangs instead of failing (plan §5.3). + repeated string hcas = 8; + string mgmt_interface = 9; + int32 http_port = 10; + + EngineProfile profile = 11; + int64 startup_timeout_ms = 12; +} + +message EngineProfile { + string name = 1; + string engine = 2; + string gpu_class = 3; + int32 tp = 4; + int32 pp = 5; + int32 gpus_per_node = 6; + int32 nodes = 7; + repeated string args = 8; + map env = 9; + string model_path = 10; + string served_model_name = 11; +} + +// StopShard terminates a shard and frees its GPUs. +message StopShard { + string shard_id = 1; +} + +// DrainShard stops routing to a shard while letting in-flight work finish. +message DrainShard { + string shard_id = 1; +} + +// Ack confirms a registration or reports that a command could not be accepted. +message Ack { + string message_id = 1; + bool ok = 2; + string error = 3; +} diff --git a/buf.gen.yaml b/buf.gen.yaml new file mode 100644 index 0000000..c591468 --- /dev/null +++ b/buf.gen.yaml @@ -0,0 +1,12 @@ +version: v2 +managed: + enabled: true +plugins: + - local: protoc-gen-go + out: api/gen + opt: paths=source_relative + - local: protoc-gen-connect-go + out: api/gen + opt: paths=source_relative +inputs: + - directory: api/proto diff --git a/buf.yaml b/buf.yaml new file mode 100644 index 0000000..248f5bb --- /dev/null +++ b/buf.yaml @@ -0,0 +1,16 @@ +version: v2 +modules: + - path: api/proto +lint: + use: + - STANDARD + except: + # The standard names assume request/response pairs. This is one long-lived + # bidirectional stream carrying several kinds of message in each direction, + # so "AgentMessage" and "ControlMessage" say what they are and + # "ConnectRequest"/"ConnectResponse" would not. + - RPC_REQUEST_STANDARD_NAME + - RPC_RESPONSE_STANDARD_NAME +breaking: + use: + - FILE diff --git a/deploy/examples/agent.yaml b/deploy/examples/agent.yaml new file mode 100644 index 0000000..50ed81e --- /dev/null +++ b/deploy/examples/agent.yaml @@ -0,0 +1,63 @@ +# llmgw agent configuration, one per GPU node. +# +# The agent dials out to the control plane and keeps one stream open. Engines +# keep serving if the agent restarts: the router talks to them directly, and +# orphan adoption picks them up again on the next start. + +node_id: gpu-a-01 # defaults to the hostname +server: "http://llmgw-server.internal:7443" + +# The address engines are reachable on, and the address NCCL bootstraps over. +# Never a fabric address: NCCL rendezvouses over TCP before moving to RDMA. +mgmt_addr: "10.0.0.11" +mgmt_interface: eth0 + +labels: + rack: r12 + +heartbeat_interval: 5s +reconnect_base: 1s +reconnect_max: 30s + +# Fabric labelling is explicit and there is no way around it. OpenSM's default +# subnet prefix is fe80:: on every fabric, so a matching prefix proves nothing +# and a wrong guess produces a replica that starts, passes health checks and +# then hangs on its first all-reduce (plan §5.2). +# +# A node with no fabrics is its own single-node island: it can host single-node +# replicas and will never be given a multi-node one. +fabrics: + - name: ib-a + hcas: [mlx5_0, mlx5_1, mlx5_2, mlx5_3] + rate_gbps: 200 + +discovery: + # Worth setting where the driver lives outside the default paths, which + # containerized driver installs often do. + nvidia_smi: "" + ib_root: /sys/class/infiniband + # NVML is the only source of NVLink topology, which is what keeps a + # tensor-parallel group off the PCIe path inside a node. It needs a binary + # built with `-tags nvml`; otherwise this silently falls back to nvidia-smi. + prefer_nvml: true + skip_gpus: false + +engines: + vllm: + bin: /opt/venvs/vllm/bin/vllm + # Which multi-node mechanism works depends on the vLLM version, and the + # flags have changed more than once. Check the version you deploy. + launch_mode: ray # "ray" or "mp" + ray_bin: /opt/venvs/vllm/bin/ray + sglang: + # SGLang is launched as a module, which is its documented entry point. + python: /opt/venvs/sglang/bin/python + +supervisor: + work_dir: /var/lib/llmgw-agent + # Engine output is the only diagnosis when a model fails to load, so keep a + # useful tail of it. + log_lines: 512 + # Engines need time to free GPU memory cleanly before being killed. + stop_timeout: 60s + health_interval: 2s diff --git a/deploy/examples/server.yaml b/deploy/examples/server.yaml index 5367abb..439587e 100644 --- a/deploy/examples/server.yaml +++ b/deploy/examples/server.yaml @@ -8,6 +8,10 @@ listen: ":8080" # client-facing OpenAI API admin_listen: ":9090" # metrics, health, pprof — do not expose this publicly +# Where node agents connect. Separate from the client port so the fleet-facing +# surface can be firewalled independently. Agents speak plaintext HTTP/2 today, +# so this must stay on a trusted network until mTLS lands (plan §5.6). +agent_listen: ":7443" router: scrape_interval: 2s @@ -79,6 +83,19 @@ router: kv_ceiling: 0.9 session_load_factor: 1.25 +# How aggressively the control plane reconciles. +control: + interval: 5s + base_backoff: 10s + max_backoff: 10m + # Node expiry is staged deliberately: a network blip must not tear down + # replicas that are still serving, and rebuilding a multi-node replica costs + # tens of minutes of model loading. + unreachable_after: 15s + lost_after: 60s + port_range: [20000, 20999] + mgmt_interface: eth0 + models: - name: llama-3.3-70b aliases: [llama] @@ -89,17 +106,13 @@ models: aliases: [qwen] routing_profile: batch -endpoints: - - id: gpu-a-01 - url: http://gpu-a-01:8000 - models: [llama-3.3-70b] - pool: ib-a/h100-80g - weight: 100 # from ctl bench; do not estimate it by hand - priority: 100 # higher pools are tried first - - - id: gpu-b-01 - url: http://gpu-b-01:8000 - models: [llama-3.3-70b, qwen2.5-32b] - pool: ib-b/a100-80g - weight: 45 - priority: 50 +# A static replica list, used only when no model declares a profile. That is the +# router-only mode for testing against sim-engines; with the control plane +# running, agents provide the replicas instead. +# +# endpoints: +# - id: sim-1 +# url: http://127.0.0.1:8001 +# models: [llama-3.3-70b] +# pool: local +# weight: 1 diff --git a/go.mod b/go.mod index 33e7392..70087e9 100644 --- a/go.mod +++ b/go.mod @@ -3,6 +3,8 @@ module gitea.mynodes.xyz/ilker/llmgw go 1.26.3 require ( + connectrpc.com/connect v1.21.0 + github.com/NVIDIA/go-nvml v0.13.4-0 github.com/cespare/xxhash/v2 v2.3.0 github.com/knadh/koanf/parsers/yaml v1.1.1 github.com/knadh/koanf/providers/env/v2 v2.0.1 @@ -12,10 +14,11 @@ require ( github.com/prometheus/client_model v0.6.3 github.com/prometheus/common v0.71.0 github.com/spf13/cobra v1.10.2 + golang.org/x/sync v0.23.0 + google.golang.org/protobuf v1.36.12 ) require ( - github.com/NVIDIA/go-nvml v0.13.4-0 // indirect github.com/beorn7/perks v1.0.1 // indirect github.com/fsnotify/fsnotify v1.9.0 // indirect github.com/go-viper/mapstructure/v2 v2.4.0 // indirect @@ -27,7 +30,7 @@ require ( github.com/prometheus/procfs v0.21.1 // indirect github.com/spf13/pflag v1.0.9 // indirect go.yaml.in/yaml/v3 v3.0.5 // indirect - golang.org/x/sync v0.23.0 // indirect - golang.org/x/sys v0.47.0 // indirect - google.golang.org/protobuf v1.36.12 // indirect + golang.org/x/net v0.59.0 // indirect + golang.org/x/sys v0.48.0 // indirect + golang.org/x/text v0.42.0 // indirect ) diff --git a/go.sum b/go.sum index cd976a7..f019f0c 100644 --- a/go.sum +++ b/go.sum @@ -1,3 +1,5 @@ +connectrpc.com/connect v1.21.0 h1:LhqSJt7jHf5NJBo9Jq/t/9FjcYAideif0mg+qe2jCUs= +connectrpc.com/connect v1.21.0/go.mod h1:A2ygJrukXwWy32vkCAAHNVguZrqZ+jeZ9rGRnGR4dN4= github.com/NVIDIA/go-nvml v0.13.4-0 h1:o3jp9u2x1R9ShFE3v+Aesp55XOSIQFMJz/VGNUcJaNE= github.com/NVIDIA/go-nvml v0.13.4-0/go.mod h1:id63qwpoDWpFXwnwM6psDCSqW4BmNu6mWpr4YeQtPGo= github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM= @@ -55,10 +57,14 @@ go.yaml.in/yaml/v2 v2.4.4/go.mod h1:gMZqIpDtDqOfM0uNfy0SkpRhvUryYH0Z6wdMYcacYXQ= go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg= go.yaml.in/yaml/v3 v3.0.5 h1:N6y/pJk8buWs9NY5ERU2HSMfm+IuD/OtfdAnq6kESPw= go.yaml.in/yaml/v3 v3.0.5/go.mod h1:HVTZu1O7/Vkt2N+BFy8Zza+lnLsABggaTM2ZpNIGuKg= +golang.org/x/net v0.59.0 h1:5zfYln+w5XCxwrnMMJPufRgNoXEaGxl0wo5GqPXyues= +golang.org/x/net v0.59.0/go.mod h1:2DA/G1UfVbCpQPeWTmMPGY7Cs2PkBkwu743bVX5PIVg= golang.org/x/sync v0.23.0 h1:KameEIfc1IkluZyXWLn39Wd4tURc6GbCiISGiZm2bQk= golang.org/x/sync v0.23.0/go.mod h1:sUUOizhqBxiL6pEWpqNLUiaJn1ShEbZ6BBqskPbjZm0= -golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs= -golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +golang.org/x/sys v0.48.0 h1:bbX/i/6MgT9BVLM9RT1thmxL04yeTAhbEz4SyadbXoo= +golang.org/x/sys v0.48.0/go.mod h1:hNLxWAXmnKAxqDtdwIYC4bM9oQPEecfsnNMuSxOs3og= +golang.org/x/text v0.42.0 h1:JbOZXgfeCPU9gacVtYliJqOhD+zhrEqK4LfdpmlUZqI= +golang.org/x/text v0.42.0/go.mod h1:ojzP1Z+2QtioaF8DTtO8K5q7JWVVYwZKenzujK0Zd0E= google.golang.org/protobuf v1.36.12 h1:pJOKDDOyeXErUroCihFAd5LQuwXBSpVnKGrj5o/fwxc= google.golang.org/protobuf v1.36.12/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= diff --git a/internal/agent/agent.go b/internal/agent/agent.go new file mode 100644 index 0000000..33faf21 --- /dev/null +++ b/internal/agent/agent.go @@ -0,0 +1,458 @@ +// Package agent is the node side of the control plane stream (plan §5.6). +package agent + +import ( + "context" + "crypto/tls" + "errors" + "fmt" + "io" + "log/slog" + "math/rand/v2" + "net" + "net/http" + "sync" + "time" + + "golang.org/x/net/http2" + + llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1" + "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1/llmgwv1connect" + "gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery" + "gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines" + "gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor" + "gitea.mynodes.xyz/ilker/llmgw/internal/control/agentsvc" + "gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory" +) + +// Config configures the agent. +type Config struct { + NodeID string `koanf:"node_id"` + Server string `koanf:"server"` + Hostname string `koanf:"hostname"` + MgmtAddr string `koanf:"mgmt_addr"` + Labels map[string]string `koanf:"labels"` + + // MgmtInterface is the ethernet interface NCCL bootstraps over. + MgmtInterface string `koanf:"mgmt_interface"` + // Fabrics labels this node's HCAs. It is explicit on purpose: OpenSM's + // default subnet prefix is the same on every fabric, so nothing can be + // inferred from the wire (plan §5.2). + Fabrics []inventory.FabricAttachment `koanf:"fabrics"` + + HeartbeatInterval time.Duration `koanf:"heartbeat_interval"` + // ReconnectBase is the first reconnect delay; it backs off from there. + ReconnectBase time.Duration `koanf:"reconnect_base"` + ReconnectMax time.Duration `koanf:"reconnect_max"` + + Discovery discovery.Options `koanf:"discovery"` + Engines engines.Config `koanf:"engines"` + Superv supervisor.Config `koanf:"supervisor"` + + Version string +} + +// DefaultConfig returns sensible defaults. +func DefaultConfig() Config { + return Config{ + MgmtInterface: "eth0", + HeartbeatInterval: 5 * time.Second, + ReconnectBase: time.Second, + ReconnectMax: 30 * time.Second, + Superv: supervisor.DefaultConfig(), + // NVML gives NVLink topology, so prefer it when this build has it. + Discovery: discovery.Options{PreferNVML: true}, + } +} + +func (c *Config) normalize() error { + d := DefaultConfig() + if c.NodeID == "" { + return errors.New("node_id is required") + } + if c.Server == "" { + return errors.New("server is required") + } + if c.MgmtInterface == "" { + c.MgmtInterface = d.MgmtInterface + } + if c.HeartbeatInterval <= 0 { + c.HeartbeatInterval = d.HeartbeatInterval + } + if c.ReconnectBase <= 0 { + c.ReconnectBase = d.ReconnectBase + } + if c.ReconnectMax < c.ReconnectBase { + c.ReconnectMax = max(c.ReconnectBase, d.ReconnectMax) + } + return nil +} + +// Agent runs on a GPU node. +type Agent struct { + cfg Config + log *slog.Logger + sup *supervisor.Supervisor + reg *engines.Registry + http connectHTTPClient + + mu sync.Mutex + inv discovery.Inventory + // shardReplica remembers which replica a shard belongs to, so a status + // report can name both. The control plane keys on the replica. + shardReplica map[string]string +} + +// connectHTTPClient is the HTTP client connect uses; an interface so tests can +// hand in an httptest server's client. +type connectHTTPClient interface { + Do(*http.Request) (*http.Response, error) +} + +// New builds an agent. +func New(cfg Config, httpClient connectHTTPClient, log *slog.Logger) (*Agent, error) { + if err := cfg.normalize(); err != nil { + return nil, err + } + if log == nil { + log = slog.Default() + } + if httpClient == nil { + httpClient = http.DefaultClient + } + + health := supervisor.NewHTTPHealth(3 * time.Second) + sup, err := supervisor.New(cfg.Superv, health, log) + if err != nil { + return nil, err + } + return &Agent{ + cfg: cfg, log: log, sup: sup, + reg: engines.NewRegistry(cfg.Engines), + http: httpClient, + shardReplica: map[string]string{}, + }, nil +} + +// Supervisor exposes the process supervisor, for tests and `ctl logs`. +func (a *Agent) Supervisor() *supervisor.Supervisor { return a.sup } + +// Run connects to the control plane and stays connected until ctx ends. +// +// Reconnecting forever is the right behaviour for a node agent: a control plane +// restart or a network partition must not require touching every GPU node, and +// the engines keep serving traffic throughout because the router talks to them +// directly. +func (a *Agent) Run(ctx context.Context) error { + if err := a.scan(ctx); err != nil { + // Not fatal: a node whose GPUs cannot be read should still register, so + // an operator can see it and why it is useless. + a.log.Error("hardware discovery failed", "err", err) + } + + attempt := 0 + for ctx.Err() == nil { + err := a.session(ctx) + if ctx.Err() != nil { + break + } + attempt++ + delay := a.reconnectDelay(attempt) + a.log.Warn("disconnected from the control plane, retrying", + "err", err, "in", delay.Round(time.Millisecond), "attempt", attempt) + + select { + case <-ctx.Done(): + case <-time.After(delay): + } + if err == nil { + attempt = 0 + } + } + + // Engines keep running when the agent stops on purpose: an agent restart + // must not interrupt served traffic. Orphan adoption picks them up again. + return ctx.Err() +} + +func (a *Agent) reconnectDelay(attempt int) time.Duration { + d := a.cfg.ReconnectBase << min(attempt-1, 16) + if d > a.cfg.ReconnectMax || d <= 0 { + d = a.cfg.ReconnectMax + } + // Jitter, so a fleet does not reconnect in lockstep after a control plane + // restart and knock it over again. + return d/2 + time.Duration(rand.Int64N(int64(d/2))) //nolint:gosec // jitter +} + +// scan refreshes the node's hardware inventory. +func (a *Agent) scan(ctx context.Context) error { + inv, err := discovery.Discover(ctx, a.cfg.Discovery, a.log) + a.mu.Lock() + a.inv = inv + a.mu.Unlock() + return err +} + +// session runs one connection until it breaks. +func (a *Agent) session(ctx context.Context) error { + client := llmgwv1connect.NewAgentServiceClient(a.http, a.cfg.Server) + + ctx, cancel := context.WithCancel(ctx) + defer cancel() + + stream := client.Connect(ctx) + defer func() { _ = stream.CloseRequest() }() + + a.mu.Lock() + inv := a.inv + a.mu.Unlock() + + if err := stream.Send(&llmgwv1.AgentMessage{ + Msg: &llmgwv1.AgentMessage_Register{ + Register: agentsvc.RegisterFromInventory( + a.cfg.NodeID, a.cfg.Hostname, a.cfg.MgmtAddr, a.cfg.Version, + a.cfg.Labels, inv, a.cfg.Fabrics), + }, + }); err != nil { + return fmt.Errorf("registering: %w", err) + } + + // Sends are serialized through this channel: connect streams are not safe + // for concurrent writes, and heartbeats race with status reports. + out := make(chan *llmgwv1.AgentMessage, 32) + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + a.sendLoop(ctx, stream, out) + }() + defer wg.Wait() + + for { + msg, err := stream.Receive() + if err != nil { + cancel() + if errors.Is(err, io.EOF) || ctx.Err() != nil { + return nil + } + return err + } + a.handle(ctx, msg, out) + } +} + +type sender interface { + Send(*llmgwv1.AgentMessage) error +} + +func (a *Agent) sendLoop(ctx context.Context, s sender, out <-chan *llmgwv1.AgentMessage) { + t := time.NewTicker(a.cfg.HeartbeatInterval) + defer t.Stop() + + for { + select { + case <-ctx.Done(): + return + case msg := <-out: + if err := s.Send(msg); err != nil { + return + } + case <-t.C: + if err := s.Send(a.heartbeat()); err != nil { + return + } + } + } +} + +func (a *Agent) heartbeat() *llmgwv1.AgentMessage { + a.mu.Lock() + inv := a.inv + a.mu.Unlock() + + hb := &llmgwv1.Heartbeat{NodeId: a.cfg.NodeID} + for _, g := range inv.GPUs { + hb.GpuMemoryUsedMb = append(hb.GpuMemoryUsedMb, int32(g.MemoryUsedMB)) + } + return &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_Heartbeat{Heartbeat: hb}} +} + +func (a *Agent) handle(ctx context.Context, msg *llmgwv1.ControlMessage, out chan<- *llmgwv1.AgentMessage) { + switch m := msg.Msg.(type) { + case *llmgwv1.ControlMessage_Start: + go a.startShard(ctx, m.Start, out) + case *llmgwv1.ControlMessage_Stop: + go a.stopShard(ctx, m.Stop.GetShardId(), out) + case *llmgwv1.ControlMessage_Drain: + if err := a.sup.Drain(m.Drain.GetShardId()); err != nil { + a.log.Warn("draining failed", "shard", m.Drain.GetShardId(), "err", err) + } + case *llmgwv1.ControlMessage_Ack: + if !m.Ack.GetOk() { + a.log.Warn("control plane rejected a message", + "id", m.Ack.GetMessageId(), "err", m.Ack.GetError()) + } + } +} + +// startShard launches a shard and reports its progress. +// +// It runs in its own goroutine because model loading takes minutes, and the +// stream has to keep carrying heartbeats and other commands meanwhile. +func (a *Agent) startShard(ctx context.Context, m *llmgwv1.StartShard, out chan<- *llmgwv1.AgentMessage) { + shardID, replicaID := m.GetShardId(), m.GetReplicaId() + + a.mu.Lock() + a.shardReplica[shardID] = replicaID + a.mu.Unlock() + + report := func(st supervisor.State, errMsg string, tail []string) { + msg := &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_Shard{Shard: &llmgwv1.ShardStatus{ + NodeId: a.cfg.NodeID, ReplicaId: replicaID, ShardId: shardID, + State: agentsvc.StateToWire(st), Error: errMsg, LastOutput: tail, + }}} + select { + case out <- msg: + case <-ctx.Done(): + } + } + + spec, timeout := agentsvc.ShardSpecFromStart(m) + + eng, err := a.reg.Get(spec.Profile.Engine) + if err != nil { + report(supervisor.Failed, err.Error(), nil) + return + } + + a.mu.Lock() + inv := a.inv + a.mu.Unlock() + + // Validate before launching. An engine that refuses a profile does so after + // minutes of loading; catching it here costs nothing. + if err := eng.Validate(spec.Profile, gpuInfo(inv, spec.GPUIndices)); err != nil { + report(supervisor.Failed, err.Error(), nil) + return + } + + cmd, err := eng.BuildCommand(spec) + if err != nil { + report(supervisor.Failed, err.Error(), nil) + return + } + + sh, err := a.sup.Start(ctx, shardID, spec, cmd) + if err != nil { + if errors.Is(err, supervisor.ErrAlreadyRunning) { + // Commands are idempotent: the control plane retries, and a repeat + // must not produce a second process (plan §5.6). + report(sh.Status().State, "", nil) + return + } + report(supervisor.Failed, err.Error(), sh.Logs(10)) + return + } + report(supervisor.Starting, "", nil) + + base := fmt.Sprintf("http://127.0.0.1:%d", spec.HTTPPort) + if err := a.sup.WaitReady(ctx, sh, base, spec.Profile.ServedModelName, timeout); err != nil { + report(supervisor.Failed, err.Error(), sh.Logs(20)) + return + } + report(supervisor.Ready, "", nil) + + // Watch for the engine dying later. Nothing else would notice: the router + // would simply see failing requests. + go a.watch(ctx, sh, replicaID, shardID, out) +} + +// watch reports a shard that exits after it was ready. +func (a *Agent) watch(ctx context.Context, sh *supervisor.Shard, replicaID, shardID string, out chan<- *llmgwv1.AgentMessage) { + t := time.NewTicker(time.Second) + defer t.Stop() + + for { + select { + case <-ctx.Done(): + return + case <-t.C: + st := sh.Status() + if !st.State.Terminal() { + continue + } + msg := &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_Shard{Shard: &llmgwv1.ShardStatus{ + NodeId: a.cfg.NodeID, ReplicaId: replicaID, ShardId: shardID, + State: agentsvc.StateToWire(st.State), Error: st.Err, + ExitCode: int32(st.ExitCode), LastOutput: sh.Logs(20), + }}} + select { + case out <- msg: + case <-ctx.Done(): + } + return + } + } +} + +func (a *Agent) stopShard(ctx context.Context, shardID string, out chan<- *llmgwv1.AgentMessage) { + a.mu.Lock() + replicaID := a.shardReplica[shardID] + delete(a.shardReplica, shardID) + a.mu.Unlock() + + if err := a.sup.Stop(ctx, shardID); err != nil { + a.log.Warn("stopping a shard failed", "shard", shardID, "err", err) + } + a.sup.Remove(shardID) + + msg := &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_Shard{Shard: &llmgwv1.ShardStatus{ + NodeId: a.cfg.NodeID, ReplicaId: replicaID, ShardId: shardID, + State: agentsvc.StateToWire(supervisor.Stopped), + }}} + select { + case out <- msg: + case <-ctx.Done(): + } +} + +// gpuInfo narrows the inventory to the GPUs a shard will use. +func gpuInfo(inv discovery.Inventory, indices []int) []engines.GPUInfo { + want := map[int]struct{}{} + for _, i := range indices { + want[i] = struct{}{} + } + var out []engines.GPUInfo + for _, g := range inv.GPUs { + if _, ok := want[g.Index]; !ok { + continue + } + out = append(out, engines.GPUInfo{ + Index: g.Index, Name: g.Name, + ComputeCapability: g.ComputeCapability, MemoryTotalMB: g.MemoryTotalMB, + }) + } + return out +} + +// HTTP2Client returns the client the agent should use to reach the control +// plane over plaintext HTTP/2. +// +// connect's bidirectional streams require HTTP/2, and an agent talking to a +// control plane over a trusted network has no TLS to negotiate it with. Go's +// default transport will not speak HTTP/2 without TLS, so it has to be asked +// explicitly. +// +// mTLS is plan §5.6's target and replaces this; until then the stream must not +// cross an untrusted network. +func HTTP2Client(timeout time.Duration) *http.Client { + return &http.Client{ + Transport: &http2.Transport{ + AllowHTTP: true, + DialTLSContext: func(ctx context.Context, network, addr string, _ *tls.Config) (net.Conn, error) { + return (&net.Dialer{Timeout: timeout}).DialContext(ctx, network, addr) + }, + }, + } +} diff --git a/internal/agent/agent_test.go b/internal/agent/agent_test.go new file mode 100644 index 0000000..d34eff4 --- /dev/null +++ b/internal/agent/agent_test.go @@ -0,0 +1,403 @@ +package agent + +import ( + "context" + "crypto/tls" + "errors" + "log/slog" + "net" + "net/http" + "net/http/httptest" + "os" + "os/exec" + "path/filepath" + "slices" + "testing" + "time" + + "golang.org/x/net/http2" + "golang.org/x/net/http2/h2c" + + "gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines" + "gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor" + "gitea.mynodes.xyz/ilker/llmgw/internal/control/agentsvc" + "gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory" + "gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler" +) + +// statusRecorder captures what the control plane hears about shards. +type statusRecorder struct{ ch chan status } + +type status struct { + replicaID, shardID string + state supervisor.State + err string +} + +func newRecorder() *statusRecorder { return &statusRecorder{ch: make(chan status, 128)} } + +func (r *statusRecorder) OnShardStatus(replicaID, shardID string, st supervisor.State, errMsg string) { + select { + case r.ch <- status{replicaID, shardID, st, errMsg}: + default: + } +} + +func (r *statusRecorder) await(t *testing.T, limit time.Duration, pred func(status) bool) status { + t.Helper() + deadline := time.After(limit) + var seen []status + for { + select { + case s := <-r.ch: + seen = append(seen, s) + if pred(s) { + return s + } + case <-deadline: + t.Fatalf("no matching status within %v; saw %+v", limit, seen) + return status{} + } + } +} + +// fakeEngineScript writes a script that stands in for vLLM. +// +// It parses the real --port flag the vLLM adapter generates rather than being +// told out of band, so the adapter's argument construction is exercised too. +func fakeEngineScript(t *testing.T, body string) string { + t.Helper() + path := filepath.Join(t.TempDir(), "fake-vllm") + script := `#!/bin/sh +port="" +while [ $# -gt 0 ]; do + case "$1" in + --port) port="$2"; shift 2 ;; + *) shift ;; + esac +done +` + body + "\n" + if err := os.WriteFile(path, []byte(script), 0o755); err != nil { + t.Fatal(err) + } + return path +} + +// servingEngine answers /health and /v1/models, which is what readiness needs. +func servingEngine(t *testing.T) string { + t.Helper() + py, err := exec.LookPath("python3") + if err != nil { + t.Skip("python3 is required for the fake engine") + } + return fakeEngineScript(t, `exec `+py+` -c ' +import http.server, json, sys +port = int(sys.argv[1]) +class H(http.server.BaseHTTPRequestHandler): + def do_GET(self): + if self.path == "/health": + self.send_response(200); self.end_headers(); return + if self.path == "/v1/models": + b = json.dumps({"object":"list","data":[{"id":"test-model"}]}).encode() + self.send_response(200) + self.send_header("content-type","application/json") + self.send_header("content-length", str(len(b))) + self.end_headers(); self.wfile.write(b); return + self.send_response(404); self.end_headers() + def log_message(self, *a): pass +http.server.HTTPServer(("127.0.0.1", port), H).serve_forever() +' "$port"`) +} + +type harness struct { + svc *agentsvc.Service + inv *inventory.Store + rec *statusRecorder + srv *httptest.Server + ag *Agent +} + +// start wires a real agent to a real service over h2c, so this exercises the +// protocol rather than two structs calling each other. +func start(t *testing.T, cfgFns ...func(*Config)) *harness { + t.Helper() + + inv := inventory.NewStore() + rec := newRecorder() + svc := agentsvc.New(inv, rec, slog.New(slog.DiscardHandler)) + + path, handler := svc.Handler() + mux := http.NewServeMux() + mux.Handle(path, handler) + // connect's bidirectional streams need HTTP/2; h2c provides it without TLS. + srv := httptest.NewServer(h2c.NewHandler(mux, &http2.Server{})) + + cfg := DefaultConfig() + cfg.NodeID = "n1" + cfg.Server = srv.URL + cfg.MgmtAddr = "127.0.0.1" + cfg.HeartbeatInterval = 50 * time.Millisecond + cfg.ReconnectBase = 20 * time.Millisecond + cfg.ReconnectMax = 100 * time.Millisecond + cfg.Discovery.SkipGPUs = true + cfg.Discovery.IBRoot = filepath.Join(t.TempDir(), "absent") + cfg.Superv = supervisor.Config{ + WorkDir: t.TempDir(), LogLines: 64, + StopTimeout: 2 * time.Second, HealthInterval: 20 * time.Millisecond, + } + for _, f := range cfgFns { + f(&cfg) + } + + client := &http.Client{Transport: &http2.Transport{ + AllowHTTP: true, + DialTLSContext: func(ctx context.Context, network, addr string, _ *tls.Config) (net.Conn, error) { + return (&net.Dialer{}).DialContext(ctx, network, addr) + }, + }} + + ag, err := New(cfg, client, slog.New(slog.DiscardHandler)) + if err != nil { + t.Fatal(err) + } + + ctx, cancel := context.WithCancel(context.Background()) + go func() { _ = ag.Run(ctx) }() + + h := &harness{svc: svc, inv: inv, rec: rec, srv: srv, ag: ag} + t.Cleanup(func() { + cancel() + ag.Supervisor().StopAll(context.Background()) + srv.Close() + }) + h.waitRegistered(t) + return h +} + +func (h *harness) waitRegistered(t *testing.T) { + t.Helper() + deadline := time.Now().Add(10 * time.Second) + for time.Now().Before(deadline) { + if _, ok := h.inv.Get("n1"); ok && len(h.svc.ConnectedNodes()) == 1 { + return + } + time.Sleep(10 * time.Millisecond) + } + t.Fatal("the agent never registered") +} + +// shardCommand builds a single-node vLLM shard on a free port. +func shardCommand(t *testing.T, shardID string) reconciler.ShardCommand { + t.Helper() + return reconciler.ShardCommand{ + ReplicaID: "rep-1", ShardID: shardID, NodeID: "n1", + Spec: engines.ShardSpec{ + ReplicaID: "rep-1", ShardID: shardID, Rank: 0, NNodes: 1, + GPUIndices: []int{0}, HTTPPort: freePort(t), + Profile: engines.Profile{ + Name: "p", Engine: "vllm", TP: 1, + ModelPath: "/models/x", ServedModelName: "test-model", + }, + }, + StartupTimeout: 15 * time.Second, + } +} + +func freePort(t *testing.T) int { + t.Helper() + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + defer ln.Close() + return ln.Addr().(*net.TCPAddr).Port +} + +func TestRegistrationCarriesInventory(t *testing.T) { + h := start(t, func(c *Config) { + c.Hostname = "gpu-a-01" + c.Labels = map[string]string{"rack": "r12"} + c.Fabrics = []inventory.FabricAttachment{ + {Fabric: "ib-a", HCAs: []string{"mlx5_0", "mlx5_1"}, RateGbps: 200}, + } + }) + + n, ok := h.inv.Get("n1") + if !ok { + t.Fatal("node missing from inventory") + } + if n.Hostname != "gpu-a-01" || n.MgmtAddr != "127.0.0.1" { + t.Errorf("identity wrong: %+v", n) + } + if n.Labels["rack"] != "r12" { + t.Errorf("labels lost: %v", n.Labels) + } + // The fabric label cannot be inferred from the wire, so losing it in + // transit would silently disable multi-node placement. + if !n.HasFabric("ib-a") { + t.Fatalf("fabric attachment lost: %+v", n.Fabrics) + } + if got := n.HCAsFor("ib-a"); !slices.Equal(got, []string{"mlx5_0", "mlx5_1"}) { + t.Errorf("HCAs=%v", got) + } +} + +func TestHeartbeatKeepsNodeAlive(t *testing.T) { + h := start(t) + + n, _ := h.inv.Get("n1") + first := n.LastSeen + + time.Sleep(200 * time.Millisecond) + n, _ = h.inv.Get("n1") + if !n.LastSeen.After(first) { + t.Error("heartbeats are not refreshing LastSeen; the node would be declared lost") + } +} + +func TestStartAndStopShardEndToEnd(t *testing.T) { + h := start(t, func(c *Config) { + c.Engines.VLLM.Bin = servingEngine(t) + }) + + cmd := shardCommand(t, "shard-a") + if err := h.svc.StartShard(t.Context(), cmd); err != nil { + t.Fatal(err) + } + + got := h.rec.await(t, 30*time.Second, func(s status) bool { + return s.shardID == "shard-a" && s.state == supervisor.Ready + }) + if got.replicaID != "rep-1" { + t.Errorf("replica id lost in transit: %q", got.replicaID) + } + + if err := h.svc.StopShard(t.Context(), "n1", "shard-a"); err != nil { + t.Fatal(err) + } + h.rec.await(t, 15*time.Second, func(s status) bool { + return s.shardID == "shard-a" && s.state == supervisor.Stopped + }) +} + +func TestFailedShardCarriesDiagnosis(t *testing.T) { + // When a model fails to load the engine's output is the only diagnosis, so + // it has to survive the trip to the control plane. + h := start(t, func(c *Config) { + c.Engines.VLLM.Bin = fakeEngineScript(t, `echo 'CUDA out of memory' >&2; exit 1`) + }) + + if err := h.svc.StartShard(t.Context(), shardCommand(t, "shard-bad")); err != nil { + t.Fatal(err) + } + got := h.rec.await(t, 30*time.Second, func(s status) bool { + return s.shardID == "shard-bad" && s.state == supervisor.Failed + }) + if got.err == "" { + t.Error("the failure carried no error message") + } +} + +func TestEngineThatNeverBecomesReady(t *testing.T) { + // A process that runs but never serves must time out rather than hang the + // replica forever. + h := start(t, func(c *Config) { + c.Engines.VLLM.Bin = fakeEngineScript(t, `sleep 60`) + }) + + cmd := shardCommand(t, "shard-stuck") + cmd.StartupTimeout = 500 * time.Millisecond + if err := h.svc.StartShard(t.Context(), cmd); err != nil { + t.Fatal(err) + } + h.rec.await(t, 20*time.Second, func(s status) bool { + return s.shardID == "shard-stuck" && s.state == supervisor.Failed + }) +} + +func TestValidationRejectsBeforeLaunch(t *testing.T) { + // An engine that refuses a profile does so after minutes of loading; + // catching it at dispatch costs nothing. + h := start(t, func(c *Config) { + c.Engines.VLLM.Bin = servingEngine(t) + }) + + cmd := shardCommand(t, "shard-invalid") + cmd.Spec.Profile.ModelPath = "" // validation requires one + + if err := h.svc.StartShard(t.Context(), cmd); err != nil { + t.Fatal(err) + } + got := h.rec.await(t, 15*time.Second, func(s status) bool { + return s.shardID == "shard-invalid" && s.state == supervisor.Failed + }) + if got.err == "" { + t.Error("validation failure carried no reason") + } +} + +func TestDispatchToUnconnectedNode(t *testing.T) { + h := start(t) + + err := h.svc.StartShard(t.Context(), reconciler.ShardCommand{NodeID: "nowhere", ShardID: "x"}) + if err == nil { + t.Fatal("dispatching to an unconnected node succeeded") + } + if !errors.Is(err, agentsvc.ErrNotConnected) { + t.Errorf("err=%v, expected ErrNotConnected", err) + } +} + +func TestAgentReconnects(t *testing.T) { + // A control plane restart must not require touching every GPU node. + h := start(t) + h.srv.CloseClientConnections() + + deadline := time.Now().Add(15 * time.Second) + for time.Now().Before(deadline) { + if len(h.svc.ConnectedNodes()) == 1 { + if err := h.svc.DrainShard(t.Context(), "n1", "nothing"); err != nil { + t.Fatalf("the reconnected stream does not accept commands: %v", err) + } + return + } + time.Sleep(20 * time.Millisecond) + } + t.Fatal("the agent never reconnected") +} + +func TestIdempotentStart(t *testing.T) { + // The control plane retries commands; a repeat must not start a second + // process (plan §5.6). + h := start(t, func(c *Config) { + c.Engines.VLLM.Bin = fakeEngineScript(t, `sleep 30`) + }) + + cmd := shardCommand(t, "shard-idem") + if err := h.svc.StartShard(t.Context(), cmd); err != nil { + t.Fatal(err) + } + h.rec.await(t, 15*time.Second, func(s status) bool { + return s.shardID == "shard-idem" && s.state == supervisor.Starting + }) + + sup := h.ag.Supervisor() + sh, err := sup.Get("shard-idem") + if err != nil { + t.Fatal(err) + } + pid := sh.Status().PID + + if err := h.svc.StartShard(t.Context(), cmd); err != nil { + t.Fatal(err) + } + time.Sleep(400 * time.Millisecond) + + sh2, err := sup.Get("shard-idem") + if err != nil { + t.Fatal(err) + } + if sh2.Status().PID != pid { + t.Errorf("a second process was started: pid %d then %d", pid, sh2.Status().PID) + } +} diff --git a/internal/agent/discovery/discover.go b/internal/agent/discovery/discover.go index 6bd912a..06585c8 100644 --- a/internal/agent/discovery/discover.go +++ b/internal/agent/discovery/discover.go @@ -16,14 +16,16 @@ var ErrNoNVML = errors.New("built without the nvml tag; rebuild with -tags nvml // Options configure a full inventory scan. type Options struct { // NvidiaSMI is the path to nvidia-smi; empty means look it up on PATH. - NvidiaSMI string + // Worth configuring on nodes where the driver lives outside the default + // paths, which containerized driver installs often do. + NvidiaSMI string `koanf:"nvidia_smi"` // IBRoot is the sysfs InfiniBand directory. - IBRoot string + IBRoot string `koanf:"ib_root"` // PreferNVML uses the NVML backend when this binary has one. It is the only // source of NVLink topology. - PreferNVML bool + PreferNVML bool `koanf:"prefer_nvml"` // SkipGPUs omits GPU discovery, for nodes that have none. - SkipGPUs bool + SkipGPUs bool `koanf:"skip_gpus"` } // Discover builds the node's inventory. diff --git a/internal/cli/agent.go b/internal/cli/agent.go index 2971d0b..686c900 100644 --- a/internal/cli/agent.go +++ b/internal/cli/agent.go @@ -1,31 +1,92 @@ package cli import ( + "context" "encoding/json" + "errors" "fmt" "os" + "os/signal" + "syscall" "text/tabwriter" + "time" "github.com/spf13/cobra" + "gitea.mynodes.xyz/ilker/llmgw/internal/agent" "gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery" + "gitea.mynodes.xyz/ilker/llmgw/internal/config" + "gitea.mynodes.xyz/ilker/llmgw/internal/version" ) func newAgentCmd() *cobra.Command { var configPath string + var ( + server string + nodeID string + mgmtAddr string + ) + cmd := &cobra.Command{ Use: "agent", - Short: "Node agent: discovery and engine supervision (phase 2)", - Long: "Runs on every GPU node. It discovers the hardware, starts and supervises\n" + - "inference engines, and reports to the control plane.\n\n" + - "The control plane connection is not built yet, so for now the useful\n" + - "subcommand is `agent inventory`.", - RunE: func(_ *cobra.Command, _ []string) error { - return errNotImplemented + Short: "Node agent: discovery and engine supervision", + Long: "Runs on every GPU node. It discovers the hardware, connects to the control\n" + + "plane, and starts, supervises and stops inference engines on demand.\n\n" + + "The agent dials out and keeps one stream open. Engines keep serving if the\n" + + "agent restarts or the control plane goes away: the router talks to them\n" + + "directly, and orphan adoption picks them up again.", + RunE: func(cmd *cobra.Command, _ []string) error { + log := logger() + + cfg, err := config.LoadAgent(configPath) + if err != nil { + return err + } + if server != "" { + cfg.Server = server + } + if nodeID != "" { + cfg.NodeID = nodeID + } + if mgmtAddr != "" { + cfg.MgmtAddr = mgmtAddr + } + if cfg.NodeID == "" { + host, err := os.Hostname() + if err != nil { + return fmt.Errorf("node_id is unset and the hostname is unavailable: %w", err) + } + cfg.NodeID = host + } + if cfg.Hostname == "" { + cfg.Hostname, _ = os.Hostname() + } + cfg.Version = version.Version + + a, err := agent.New(cfg, agent.HTTP2Client(30*time.Second), log) + if err != nil { + return err + } + + ctx, stop := signal.NotifyContext(cmd.Context(), syscall.SIGINT, syscall.SIGTERM) + defer stop() + + log.Info("agent starting", + "node", cfg.NodeID, "server", cfg.Server, + "fabrics", len(cfg.Fabrics), "nvml", discovery.NVMLSupported()) + + err = a.Run(ctx) + if errors.Is(err, context.Canceled) { + return nil + } + return err }, } - cmd.Flags().StringVar(&configPath, "config", "agent.yaml", "configuration file") + cmd.Flags().StringVar(&configPath, "config", "", "configuration file") + cmd.Flags().StringVar(&server, "server", "", "control plane address (overrides the config)") + cmd.Flags().StringVar(&nodeID, "node-id", "", "node id (defaults to the hostname)") + cmd.Flags().StringVar(&mgmtAddr, "mgmt-addr", "", "address engines are reachable on") cmd.AddCommand(newInventoryCmd()) return cmd } diff --git a/internal/cli/server.go b/internal/cli/server.go index 626dff8..abdf55f 100644 --- a/internal/cli/server.go +++ b/internal/cli/server.go @@ -4,17 +4,25 @@ import ( "context" "errors" "fmt" + "log/slog" "net" "net/http" "os/signal" "strings" + "sync" "syscall" "time" "github.com/spf13/cobra" + "golang.org/x/net/http2" + "golang.org/x/net/http2/h2c" "golang.org/x/sync/errgroup" "gitea.mynodes.xyz/ilker/llmgw/internal/config" + "gitea.mynodes.xyz/ilker/llmgw/internal/control" + "gitea.mynodes.xyz/ilker/llmgw/internal/control/agentsvc" + "gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog" + "gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler" "gitea.mynodes.xyz/ilker/llmgw/internal/router" ) @@ -55,8 +63,10 @@ func newServerCmd() *cobra.Command { if err := cfg.Validate(); err != nil { return err } - if len(cfg.Endpoints) == 0 { - return errors.New("no endpoints configured; use --endpoint or a config file") + controlMode := cfg.ControlPlaneMode() + if !controlMode && len(cfg.Endpoints) == 0 { + return errors.New("nothing to serve: declare models with profiles, " + + "or give --endpoint for a router-only run") } rt, err := router.New(cfg, log) @@ -70,6 +80,28 @@ func newServerCmd() *cobra.Command { rt.Start(ctx) defer rt.Stop() + // The control plane replaces the static endpoint list: agents + // register, the reconciler places replicas, and Ready ones are + // published to the router (plan §6.4). + var agentHandler http.Handler + var agentPath string + if controlMode { + cat, err := cfg.Catalog() + if err != nil { + return err + } + + ctl, svc, err := startControlPlane(ctx, cfg, cat, rt, log) + if err != nil { + return err + } + defer ctl.Close() + agentPath, agentHandler = svc.Handler() + + log.Info("control plane running", + "models", len(cfg.Models), "agent_listen", cfg.AgentListen) + } + adminHandler, err := rt.AdminHandler() if err != nil { return err @@ -103,14 +135,39 @@ func newServerCmd() *cobra.Command { "default_profile", cfg.Router.DefaultProfile, ) + var agentSrv *http.Server + var agentLn net.Listener + if agentHandler != nil { + mux := http.NewServeMux() + mux.Handle(agentPath, agentHandler) + agentSrv = &http.Server{ + // Agents speak plaintext HTTP/2, which Go will not do + // without being asked. mTLS replaces this (plan §5.6); + // until then the agent port must stay on a trusted network. + Handler: h2c.NewHandler(mux, &http2.Server{}), + ReadHeaderTimeout: 10 * time.Second, + } + if agentLn, err = net.Listen("tcp", cfg.AgentListen); err != nil { + _ = apiLn.Close() + _ = adminLn.Close() + return fmt.Errorf("listening on %s: %w", cfg.AgentListen, err) + } + } + g, gctx := errgroup.WithContext(ctx) g.Go(func() error { return serve(apiSrv, apiLn) }) g.Go(func() error { return serve(adminSrv, adminLn) }) + if agentSrv != nil { + g.Go(func() error { return serve(agentSrv, agentLn) }) + } g.Go(func() error { <-gctx.Done() log.Info("shutting down") sctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) defer cancel() + if agentSrv != nil { + _ = agentSrv.Shutdown(sctx) + } _ = adminSrv.Shutdown(sctx) return apiSrv.Shutdown(sctx) }) @@ -127,6 +184,85 @@ func newServerCmd() *cobra.Command { return cmd } +// startControlPlane builds the control plane and its agent service. +// +// The service is both the AgentService handler and the reconciler's Dispatcher: +// it holds the open streams, so it is the only thing that can deliver a command +// to a node. +func startControlPlane(ctx context.Context, cfg config.Server, cat *catalog.Catalog, + sink control.EndpointSink, log *slog.Logger, +) (*control.Controller, *agentsvc.Service, error) { + ccfg := control.DefaultConfig() + ccfg.Reconciler = reconciler.Config{ + Interval: cfg.Control.Interval, + BaseBackoff: cfg.Control.BaseBackoff, + MaxBackoff: cfg.Control.MaxBackoff, + UnreachableAfter: cfg.Control.UnreachableAfter, + LostAfter: cfg.Control.LostAfter, + MgmtInterface: cfg.Control.MgmtInterface, + } + if len(cfg.Control.PortRange) == 2 { + ccfg.Reconciler.PortRange = [2]int{cfg.Control.PortRange[0], cfg.Control.PortRange[1]} + } + + // The service needs the inventory and the status sink, both of which the + // controller owns; the controller needs the service as its dispatcher. The + // cycle is broken by giving the controller a dispatcher that forwards to + // the service once it exists. + fwd := &dispatcherRef{} + ctl := control.New(ccfg, cat, fwd, sink, log) + svc := agentsvc.New(ctl.Inventory(), ctl.Reconciler(), log) + fwd.set(svc) + + ctl.Start(ctx) + return ctl, svc, nil +} + +// dispatcherRef lets the controller be built before the service exists. +type dispatcherRef struct { + mu sync.RWMutex + d reconciler.Dispatcher +} + +func (r *dispatcherRef) set(d reconciler.Dispatcher) { + r.mu.Lock() + r.d = d + r.mu.Unlock() +} + +func (r *dispatcherRef) get() (reconciler.Dispatcher, error) { + r.mu.RLock() + defer r.mu.RUnlock() + if r.d == nil { + return nil, errors.New("the agent service is not ready yet") + } + return r.d, nil +} + +func (r *dispatcherRef) StartShard(ctx context.Context, cmd reconciler.ShardCommand) error { + d, err := r.get() + if err != nil { + return err + } + return d.StartShard(ctx, cmd) +} + +func (r *dispatcherRef) StopShard(ctx context.Context, nodeID, shardID string) error { + d, err := r.get() + if err != nil { + return err + } + return d.StopShard(ctx, nodeID, shardID) +} + +func (r *dispatcherRef) DrainShard(ctx context.Context, nodeID, shardID string) error { + d, err := r.get() + if err != nil { + return err + } + return d.DrainShard(ctx, nodeID, shardID) +} + func serve(s *http.Server, ln net.Listener) error { if err := s.Serve(ln); err != nil && !errors.Is(err, http.ErrServerClosed) { return err diff --git a/internal/config/agent.go b/internal/config/agent.go new file mode 100644 index 0000000..b73cbe7 --- /dev/null +++ b/internal/config/agent.go @@ -0,0 +1,83 @@ +package config + +import ( + "fmt" + "os" + "strings" + + "github.com/knadh/koanf/parsers/yaml" + env "github.com/knadh/koanf/providers/env/v2" + "github.com/knadh/koanf/providers/file" + "github.com/knadh/koanf/v2" + + "gitea.mynodes.xyz/ilker/llmgw/internal/agent" +) + +// LoadAgent reads an agent configuration, overlaying it on the defaults and +// then applying LLMGW_-prefixed environment overrides. +// +// An empty path loads the defaults, which with --server and --node-id on the +// command line is enough for a smoke test on a node with no IB. +func LoadAgent(path string) (agent.Config, error) { + cfg := agent.DefaultConfig() + k := koanf.New(".") + + if path != "" { + if _, err := os.Stat(path); err != nil { + return cfg, fmt.Errorf("config file: %w", err) + } + if err := k.Load(file.Provider(path), yaml.Parser()); err != nil { + return cfg, fmt.Errorf("parsing %s: %w", path, err) + } + } + + err := k.Load(env.Provider(".", env.Opt{ + Prefix: "LLMGW_", + TransformFunc: func(key, value string) (string, any) { + key = strings.ToLower(strings.TrimPrefix(key, "LLMGW_")) + return strings.ReplaceAll(key, "__", "."), value + }, + }), nil) + if err != nil { + return cfg, fmt.Errorf("environment overrides: %w", err) + } + + if err := k.UnmarshalWithConf("", &cfg, koanf.UnmarshalConf{Tag: "koanf"}); err != nil { + return cfg, fmt.Errorf("decoding configuration: %w", err) + } + return cfg, validateAgent(cfg) +} + +// validateAgent rejects configurations that would fail confusingly at runtime. +func validateAgent(c agent.Config) error { + seen := map[string]struct{}{} + for i, f := range c.Fabrics { + if f.Fabric == "" { + return fmt.Errorf("fabrics[%d]: name is required", i) + } + if _, dup := seen[f.Fabric]; dup { + return fmt.Errorf("fabric %q is declared twice", f.Fabric) + } + seen[f.Fabric] = struct{}{} + if len(f.HCAs) == 0 { + // A fabric with no adapters is worse than no fabric at all: the + // scheduler would place multi-node replicas here and every shard + // would come up with NCCL unable to find a path (plan §5.3). + return fmt.Errorf("fabric %q lists no HCAs", f.Fabric) + } + } + + // The same adapter on two fabrics means one of the labels is wrong, and the + // consequence is a replica that starts and then hangs on its first + // all-reduce — the failure mode this labelling exists to prevent. + byHCA := map[string]string{} + for _, f := range c.Fabrics { + for _, h := range f.HCAs { + if prev, dup := byHCA[h]; dup { + return fmt.Errorf("HCA %q is claimed by both %q and %q", h, prev, f.Fabric) + } + byHCA[h] = f.Fabric + } + } + return nil +} diff --git a/internal/config/config.go b/internal/config/config.go index 07960f3..c6c487c 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -15,6 +15,8 @@ import ( env "github.com/knadh/koanf/providers/env/v2" "github.com/knadh/koanf/providers/file" "github.com/knadh/koanf/v2" + + "gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog" ) // Server is the whole server configuration. @@ -25,11 +27,38 @@ type Server struct { // client-facing port can be exposed without exposing diagnostics. AdminListen string `koanf:"admin_listen"` - Router Router `koanf:"router"` - Models []Model `koanf:"models"` + // AgentListen is where node agents connect. Separate from the client port + // so the fleet-facing surface can be firewalled independently. + AgentListen string `koanf:"agent_listen"` + + Router Router `koanf:"router"` + Control Control `koanf:"control"` + Models []catalog.Model `koanf:"models"` + + // Endpoints is the phase-1 static replica list. It is only consulted when + // no model declares a profile, which is what makes a router-only smoke test + // against sim-engines possible without a control plane. Endpoints []Endpoint `koanf:"endpoints"` } +// Control configures the control plane. +type Control struct { + // Interval is how often desired and actual state are compared. + Interval time.Duration `koanf:"interval"` + // BaseBackoff and MaxBackoff bound how fast a failing model is retried. + BaseBackoff time.Duration `koanf:"base_backoff"` + MaxBackoff time.Duration `koanf:"max_backoff"` + // UnreachableAfter and LostAfter stage node expiry. They are deliberately + // far apart: a network blip must not tear down replicas that still serve, + // and rebuilding a multi-node replica costs tens of minutes of loading. + UnreachableAfter time.Duration `koanf:"unreachable_after"` + LostAfter time.Duration `koanf:"lost_after"` + // PortRange is the per-node range for engine HTTP and dist-init ports. + PortRange []int `koanf:"port_range"` + // MgmtInterface is the ethernet interface NCCL bootstraps over. + MgmtInterface string `koanf:"mgmt_interface"` +} + // Router holds routing behaviour. type Router struct { ScrapeInterval time.Duration `koanf:"scrape_interval"` @@ -106,16 +135,6 @@ type Scoring struct { SessionLoadFactor float64 `koanf:"session_load_factor"` } -// Model is a catalog entry as clients see it. -type Model struct { - Name string `koanf:"name"` - Aliases []string `koanf:"aliases"` - // Fallbacks are tried, in order, when this model has no usable endpoint. - Fallbacks []string `koanf:"fallbacks"` - // RoutingProfile overrides the router's default profile for this model. - RoutingProfile string `koanf:"routing_profile"` -} - // Endpoint is one engine replica. Static in phase 1. type Endpoint struct { ID string `koanf:"id"` @@ -131,6 +150,16 @@ func Default() Server { return Server{ Listen: ":8080", AdminListen: ":9090", + AgentListen: ":7443", + Control: Control{ + Interval: 5 * time.Second, + BaseBackoff: 10 * time.Second, + MaxBackoff: 10 * time.Minute, + UnreachableAfter: 15 * time.Second, + LostAfter: 60 * time.Second, + PortRange: []int{20000, 20999}, + MgmtInterface: "eth0", + }, Router: Router{ ScrapeInterval: 2 * time.Second, ScrapeTimeout: time.Second, @@ -234,25 +263,29 @@ func (s Server) Validate() error { seen[e.ID] = struct{}{} } - names := make(map[string]struct{}, len(s.Models)) - for i, m := range s.Models { - if m.Name == "" { - return fmt.Errorf("models[%d]: name is required", i) - } - if _, dup := names[m.Name]; dup { - return fmt.Errorf("model %q is declared twice", m.Name) - } - names[m.Name] = struct{}{} - } - // Aliases share a namespace with names, so a collision must be caught here - // rather than silently shadowing a model at request time. - for _, m := range s.Models { - for _, a := range m.Aliases { - if _, clash := names[a]; clash { - return fmt.Errorf("model %q: alias %q collides with another model name", m.Name, a) - } - names[a] = struct{}{} - } + // The catalog enforces the rest — aliases, fallbacks, parallelism + // arithmetic — so there is one place those rules live. + if _, err := s.Catalog(); err != nil { + return err } return nil } + +// Catalog builds the model catalog from the configuration. +func (s Server) Catalog() (*catalog.Catalog, error) { + return catalog.New(s.Models) +} + +// ControlPlaneMode reports whether the server should run a control plane. +// +// A model with no profiles cannot be placed, so a configuration of bare names +// plus a static endpoint list is the router-only mode used for testing against +// sim-engines. +func (s Server) ControlPlaneMode() bool { + for _, m := range s.Models { + if len(m.Profiles) > 0 { + return true + } + } + return false +} diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 3a3c2b2..e448bb9 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -6,6 +6,8 @@ import ( "strings" "testing" "time" + + "gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog" ) func write(t *testing.T, content string) string { @@ -128,9 +130,11 @@ func TestValidateRejectsAliasCollision(t *testing.T) { // An alias shadowing another model would silently misroute every request // for that model, so it has to fail at load time. s := Default() - s.Models = []Model{ - {Name: "llama-3.3-70b", Aliases: []string{"qwen"}}, - {Name: "qwen"}, + s.Models = []catalog.Model{ + {Name: "llama-3.3-70b", Source: "/m", Aliases: []string{"qwen"}, + Profiles: []catalog.Profile{{Name: "p", GPUClass: "h100-80g", TP: 1}}}, + {Name: "qwen", Source: "/m", + Profiles: []catalog.Profile{{Name: "p", GPUClass: "h100-80g", TP: 1}}}, } err := s.Validate() if err == nil || !strings.Contains(err.Error(), "collides") { @@ -140,9 +144,11 @@ func TestValidateRejectsAliasCollision(t *testing.T) { func TestValidateAcceptsDistinctAliases(t *testing.T) { s := Default() - s.Models = []Model{ - {Name: "llama-3.3-70b", Aliases: []string{"llama", "llama3"}}, - {Name: "qwen2.5-32b", Aliases: []string{"qwen"}}, + s.Models = []catalog.Model{ + {Name: "llama-3.3-70b", Source: "/m", Aliases: []string{"llama", "llama3"}, + Profiles: []catalog.Profile{{Name: "p", GPUClass: "h100-80g", TP: 1}}}, + {Name: "qwen2.5-32b", Source: "/m", Aliases: []string{"qwen"}, + Profiles: []catalog.Profile{{Name: "p", GPUClass: "h100-80g", TP: 1}}}, } if err := s.Validate(); err != nil { t.Errorf("valid aliases rejected: %v", err) diff --git a/internal/control/agentsvc/convert.go b/internal/control/agentsvc/convert.go new file mode 100644 index 0000000..b63ae2c --- /dev/null +++ b/internal/control/agentsvc/convert.go @@ -0,0 +1,193 @@ +// Package agentsvc is the control plane's end of the agent stream (plan §5.6). +package agentsvc + +import ( + "time" + + llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1" + "gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery" + "gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines" + "gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor" + "gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory" + "gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler" +) + +// NodeFromRegister converts a registration into an inventory node. +func NodeFromRegister(r *llmgwv1.Register) inventory.Node { + n := inventory.Node{ + ID: r.GetNodeId(), + Hostname: r.GetHostname(), + MgmtAddr: r.GetMgmtAddr(), + Labels: r.GetLabels(), + State: inventory.Ready, + AgentVersion: r.GetAgentVersion(), + EngineVersions: r.GetEngineVersions(), + } + for _, g := range r.GetGpus() { + n.GPUs = append(n.GPUs, discovery.GPU{ + Index: int(g.GetIndex()), + UUID: g.GetUuid(), + Name: g.GetName(), + MemoryTotalMB: int(g.GetMemoryTotalMb()), + MemoryUsedMB: int(g.GetMemoryUsedMb()), + ComputeCapability: g.GetComputeCapability(), + NUMANode: int(g.GetNumaNode()), + PCIBusID: g.GetPciBusId(), + DriverVersion: g.GetDriverVersion(), + CUDAVersion: g.GetCudaVersion(), + NVLinkPeers: intSlice(g.GetNvlinkPeers()), + Health: discovery.GPUHealth{ + ECCErrors: int(g.GetEccErrors()), + TemperatureC: int(g.GetTemperatureC()), + Problems: g.GetProblems(), + }, + }) + } + for _, f := range r.GetFabrics() { + n.Fabrics = append(n.Fabrics, inventory.FabricAttachment{ + Fabric: f.GetFabric(), + HCAs: f.GetHcas(), + RateGbps: int(f.GetRateGbps()), + }) + } + return n +} + +// RegisterFromInventory builds the registration an agent sends. +func RegisterFromInventory(nodeID, hostname, mgmtAddr, agentVersion string, + labels map[string]string, inv discovery.Inventory, fabrics []inventory.FabricAttachment, +) *llmgwv1.Register { + r := &llmgwv1.Register{ + NodeId: nodeID, Hostname: hostname, MgmtAddr: mgmtAddr, + Labels: labels, AgentVersion: agentVersion, + } + for _, g := range inv.GPUs { + r.Gpus = append(r.Gpus, &llmgwv1.GPU{ + Index: int32(g.Index), + Uuid: g.UUID, + Name: g.Name, + MemoryTotalMb: int32(g.MemoryTotalMB), + MemoryUsedMb: int32(g.MemoryUsedMB), + ComputeCapability: g.ComputeCapability, + NumaNode: int32(g.NUMANode), + PciBusId: g.PCIBusID, + DriverVersion: g.DriverVersion, + CudaVersion: g.CUDAVersion, + NvlinkPeers: int32Slice(g.NVLinkPeers), + Problems: g.Health.Problems, + TemperatureC: int32(g.Health.TemperatureC), + EccErrors: int32(g.Health.ECCErrors), + }) + } + for _, f := range fabrics { + r.Fabrics = append(r.Fabrics, &llmgwv1.FabricAttachment{ + Fabric: f.Fabric, Hcas: f.HCAs, RateGbps: int32(f.RateGbps), + }) + } + return r +} + +// StartShardFromCommand converts a reconciler command into a wire message. +func StartShardFromCommand(cmd reconciler.ShardCommand) *llmgwv1.StartShard { + p := cmd.Spec.Profile + return &llmgwv1.StartShard{ + ReplicaId: cmd.ReplicaID, + ShardId: cmd.ShardID, + Rank: int32(cmd.Spec.Rank), + Nnodes: int32(cmd.Spec.NNodes), + LeaderAddr: cmd.Spec.LeaderAddr, + LeaderPort: int32(cmd.Spec.LeaderPort), + GpuIndices: int32Slice(cmd.Spec.GPUIndices), + Hcas: cmd.Spec.HCAs, + MgmtInterface: cmd.Spec.MgmtInterface, + HttpPort: int32(cmd.Spec.HTTPPort), + Profile: &llmgwv1.EngineProfile{ + Name: p.Name, Engine: p.Engine, GpuClass: p.GPUClass, + Tp: int32(p.TP), Pp: int32(p.PP), + GpusPerNode: int32(p.GPUsPerNode), Nodes: int32(p.Nodes), + Args: p.Args, Env: p.Env, + ModelPath: p.ModelPath, ServedModelName: p.ServedModelName, + }, + StartupTimeoutMs: cmd.StartupTimeout.Milliseconds(), + } +} + +// ShardSpecFromStart converts a wire message into an engine shard spec. +func ShardSpecFromStart(m *llmgwv1.StartShard) (engines.ShardSpec, time.Duration) { + p := m.GetProfile() + spec := engines.ShardSpec{ + ReplicaID: m.GetReplicaId(), + ShardID: m.GetShardId(), + Rank: int(m.GetRank()), + NNodes: int(m.GetNnodes()), + LeaderAddr: m.GetLeaderAddr(), + LeaderPort: int(m.GetLeaderPort()), + GPUIndices: intSlice(m.GetGpuIndices()), + HCAs: m.GetHcas(), + MgmtInterface: m.GetMgmtInterface(), + HTTPPort: int(m.GetHttpPort()), + Profile: engines.Profile{ + Name: p.GetName(), Engine: p.GetEngine(), GPUClass: p.GetGpuClass(), + TP: int(p.GetTp()), PP: int(p.GetPp()), + GPUsPerNode: int(p.GetGpusPerNode()), Nodes: int(p.GetNodes()), + Args: p.GetArgs(), Env: p.GetEnv(), + ModelPath: p.GetModelPath(), ServedModelName: p.GetServedModelName(), + }, + } + return spec, time.Duration(m.GetStartupTimeoutMs()) * time.Millisecond +} + +// stateWire maps supervisor states onto the wire enum. +var stateWire = map[supervisor.State]llmgwv1.ShardStatus_State{ + supervisor.Pending: llmgwv1.ShardStatus_STATE_PENDING, + supervisor.Starting: llmgwv1.ShardStatus_STATE_STARTING, + supervisor.Loading: llmgwv1.ShardStatus_STATE_LOADING, + supervisor.Ready: llmgwv1.ShardStatus_STATE_READY, + supervisor.Draining: llmgwv1.ShardStatus_STATE_DRAINING, + supervisor.Stopped: llmgwv1.ShardStatus_STATE_STOPPED, + supervisor.Failed: llmgwv1.ShardStatus_STATE_FAILED, +} + +// StateToWire converts a supervisor state for transmission. +func StateToWire(s supervisor.State) llmgwv1.ShardStatus_State { + if w, ok := stateWire[s]; ok { + return w + } + return llmgwv1.ShardStatus_STATE_UNSPECIFIED +} + +// StateFromWire converts a received state. +// +// An unrecognised state is treated as Failed rather than guessed at: a newer +// agent reporting something this control plane does not understand must not be +// mistaken for a healthy shard. +func StateFromWire(w llmgwv1.ShardStatus_State) supervisor.State { + for s, v := range stateWire { + if v == w { + return s + } + } + return supervisor.Failed +} + +func intSlice(v []int32) []int { + if len(v) == 0 { + return nil + } + out := make([]int, len(v)) + for i, x := range v { + out[i] = int(x) + } + return out +} + +func int32Slice(v []int) []int32 { + if len(v) == 0 { + return nil + } + out := make([]int32, len(v)) + for i, x := range v { + out[i] = int32(x) + } + return out +} diff --git a/internal/control/agentsvc/service.go b/internal/control/agentsvc/service.go new file mode 100644 index 0000000..739a6a6 --- /dev/null +++ b/internal/control/agentsvc/service.go @@ -0,0 +1,294 @@ +package agentsvc + +import ( + "context" + "errors" + "fmt" + "io" + "log/slog" + "net/http" + "slices" + "sync" + + "connectrpc.com/connect" + + llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1" + "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1/llmgwv1connect" + "gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor" + "gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory" + "gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler" +) + +// StatusSink receives shard transitions. The reconciler implements it. +type StatusSink interface { + OnShardStatus(replicaID, shardID string, state supervisor.State, errMsg string) +} + +// Service is the control plane's AgentService implementation, and at the same +// time the reconciler's Dispatcher: it holds the open streams, so it is the +// only thing that can deliver a command to a node. +type Service struct { + inv *inventory.Store + status StatusSink + log *slog.Logger + + mu sync.RWMutex + conns map[string]*conn +} + +// conn is one agent's open stream. +type conn struct { + nodeID string + // send is serialized through a channel rather than a mutex around the + // stream: connect's streams are not safe for concurrent writes, and the + // reconciler dispatches to many nodes at once. + send chan *llmgwv1.ControlMessage + closed chan struct{} + once sync.Once +} + +func (c *conn) close() { c.once.Do(func() { close(c.closed) }) } + +// New builds the service. +func New(inv *inventory.Store, status StatusSink, log *slog.Logger) *Service { + if log == nil { + log = slog.Default() + } + return &Service{inv: inv, status: status, log: log, conns: map[string]*conn{}} +} + +// Handler returns the HTTP path and handler to mount. +func (s *Service) Handler() (string, http.Handler) { + return llmgwv1connect.NewAgentServiceHandler(s) +} + +// Connect serves one agent's stream for the lifetime of that agent. +func (s *Service) Connect( + ctx context.Context, + stream *connect.BidiStream[llmgwv1.AgentMessage, llmgwv1.ControlMessage], +) error { + first, err := stream.Receive() + if err != nil { + return connect.NewError(connect.CodeInvalidArgument, fmt.Errorf("no opening message: %w", err)) + } + reg := first.GetRegister() + if reg == nil { + // Registration first, always: the control plane has to know which node + // it is talking to before it can route anything to it. + return connect.NewError(connect.CodeInvalidArgument, + errors.New("the first message must be a Register")) + } + nodeID := reg.GetNodeId() + if nodeID == "" { + return connect.NewError(connect.CodeInvalidArgument, errors.New("node_id is required")) + } + + node := NodeFromRegister(reg) + s.inv.Upsert(node) + s.log.Info("agent registered", + "node", nodeID, "gpus", len(node.GPUs), "fabrics", node.FabricNames(), + "agent_version", reg.GetAgentVersion()) + + c := s.attach(ctx, nodeID) + defer s.detach(nodeID, c) + + if err := stream.Send(&llmgwv1.ControlMessage{ + Msg: &llmgwv1.ControlMessage_Ack{Ack: &llmgwv1.Ack{MessageId: "register", Ok: true}}, + }); err != nil { + return err + } + + // One goroutine writes, this one reads. Splitting them is what lets a + // command be dispatched while the agent is mid-heartbeat. + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + s.writeLoop(ctx, stream, c) + }() + defer wg.Wait() + + for { + msg, err := stream.Receive() + if err != nil { + c.close() + if ctx.Err() != nil || errors.Is(err, io.EOF) { + s.log.Info("agent disconnected", "node", nodeID) + return nil + } + s.log.Warn("agent stream failed", "node", nodeID, "err", err) + return nil + } + s.handle(nodeID, msg) + } +} + +func (s *Service) writeLoop(ctx context.Context, + stream *connect.BidiStream[llmgwv1.AgentMessage, llmgwv1.ControlMessage], c *conn, +) { + for { + select { + case <-ctx.Done(): + return + case <-c.closed: + return + case msg := <-c.send: + if err := stream.Send(msg); err != nil { + s.log.Warn("sending to agent failed", "node", c.nodeID, "err", err) + c.close() + return + } + } + } +} + +func (s *Service) handle(nodeID string, msg *llmgwv1.AgentMessage) { + switch m := msg.Msg.(type) { + case *llmgwv1.AgentMessage_Heartbeat: + s.inv.Touch(nodeID) + s.applyHeartbeat(nodeID, m.Heartbeat) + + case *llmgwv1.AgentMessage_Shard: + st := StateFromWire(m.Shard.GetState()) + s.status.OnShardStatus(m.Shard.GetReplicaId(), m.Shard.GetShardId(), st, m.Shard.GetError()) + if m.Shard.GetState() == llmgwv1.ShardStatus_STATE_FAILED { + s.log.Error("shard failed", + "node", nodeID, "shard", m.Shard.GetShardId(), + "err", m.Shard.GetError(), "exit_code", m.Shard.GetExitCode(), + "last_output", m.Shard.GetLastOutput()) + } + + case *llmgwv1.AgentMessage_Register: + // A re-registration on a live stream means the agent's view of its own + // hardware changed — a GPU fell out, say. Take it. + s.inv.Upsert(NodeFromRegister(m.Register)) + s.log.Info("agent re-registered", "node", nodeID) + } +} + +// applyHeartbeat folds in the cheap facts a heartbeat carries, so a foreign +// process appearing on a card after registration is noticed without waiting for +// a reconnect. +func (s *Service) applyHeartbeat(nodeID string, hb *llmgwv1.Heartbeat) { + used := hb.GetGpuMemoryUsedMb() + problems := hb.GetNewProblems() + if len(used) == 0 && len(problems) == 0 { + return + } + + n, ok := s.inv.Get(nodeID) + if !ok { + return + } + for i := range n.GPUs { + if i < len(used) { + n.GPUs[i].MemoryUsedMB = int(used[i]) + } + } + if len(problems) > 0 { + // Problems are reported per node, not per GPU, so attach them to the + // node's first GPU rather than guessing which one is affected. Any + // problem is enough to keep that GPU out of placement, which is the + // conservative direction. + n.GPUs[0].Health.Problems = append(n.GPUs[0].Health.Problems, problems...) + } + s.inv.Upsert(n) +} + +func (s *Service) attach(ctx context.Context, nodeID string) *conn { + c := &conn{ + nodeID: nodeID, + // Buffered so a burst of commands does not block the reconciler on a + // slow agent; if it fills, the agent is too far behind to be useful. + send: make(chan *llmgwv1.ControlMessage, 64), + closed: make(chan struct{}), + } + + s.mu.Lock() + if prev, ok := s.conns[nodeID]; ok { + // A node reconnecting before the old stream was reaped. The new one + // wins: the old is by definition the dead one. + s.log.Warn("agent reconnected while an old stream was open", "node", nodeID) + prev.close() + } + s.conns[nodeID] = c + s.mu.Unlock() + + go func() { + <-ctx.Done() + c.close() + }() + return c +} + +func (s *Service) detach(nodeID string, c *conn) { + c.close() + s.mu.Lock() + if s.conns[nodeID] == c { + delete(s.conns, nodeID) + } + s.mu.Unlock() + + // The node is not marked lost here. Losing a stream is not losing a node, + // and the reconciler's two-stage expiry already handles silence with the + // right delays (plan §6.3). + s.log.Info("agent stream closed", "node", nodeID) +} + +// ConnectedNodes lists the nodes with an open stream. +func (s *Service) ConnectedNodes() []string { + s.mu.RLock() + defer s.mu.RUnlock() + out := make([]string, 0, len(s.conns)) + for id := range s.conns { + out = append(out, id) + } + slices.Sort(out) + return out +} + +// --- Dispatcher --- + +// ErrNotConnected means no agent stream is open for that node. +var ErrNotConnected = errors.New("no agent connected") + +// StartShard implements reconciler.Dispatcher. +func (s *Service) StartShard(_ context.Context, cmd reconciler.ShardCommand) error { + return s.send(cmd.NodeID, &llmgwv1.ControlMessage{ + Msg: &llmgwv1.ControlMessage_Start{Start: StartShardFromCommand(cmd)}, + }) +} + +// StopShard implements reconciler.Dispatcher. +func (s *Service) StopShard(_ context.Context, nodeID, shardID string) error { + return s.send(nodeID, &llmgwv1.ControlMessage{ + Msg: &llmgwv1.ControlMessage_Stop{Stop: &llmgwv1.StopShard{ShardId: shardID}}, + }) +} + +// DrainShard implements reconciler.Dispatcher. +func (s *Service) DrainShard(_ context.Context, nodeID, shardID string) error { + return s.send(nodeID, &llmgwv1.ControlMessage{ + Msg: &llmgwv1.ControlMessage_Drain{Drain: &llmgwv1.DrainShard{ShardId: shardID}}, + }) +} + +func (s *Service) send(nodeID string, msg *llmgwv1.ControlMessage) error { + s.mu.RLock() + c, ok := s.conns[nodeID] + s.mu.RUnlock() + if !ok { + return fmt.Errorf("%w: node %s", ErrNotConnected, nodeID) + } + + select { + case c.send <- msg: + return nil + case <-c.closed: + return fmt.Errorf("%w: node %s stream closed", ErrNotConnected, nodeID) + default: + // The queue is full, so this agent is far enough behind that queueing + // more would only delay discovering it is broken. + return fmt.Errorf("node %s is not keeping up with commands", nodeID) + } +} diff --git a/internal/router/router.go b/internal/router/router.go index 1ad96e0..eae0ab3 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -16,6 +16,7 @@ import ( "time" "gitea.mynodes.xyz/ilker/llmgw/internal/config" + "gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog" "gitea.mynodes.xyz/ilker/llmgw/internal/router/admission" "gitea.mynodes.xyz/ilker/llmgw/internal/router/breaker" "gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint" @@ -39,7 +40,7 @@ type Router struct { profiles map[string]*scheduler.Profile // models maps every name a client may use — canonical names and aliases — // onto its catalog entry. - models map[string]config.Model + models map[string]catalog.Model } // New builds a router from configuration. Start must be called to begin scraping. @@ -71,7 +72,7 @@ func New(cfg config.Server, log *slog.Logger) (*Router, error) { return nil, fmt.Errorf("unknown default_profile %q", cfg.Router.DefaultProfile) } - models := map[string]config.Model{} + models := map[string]catalog.Model{} for _, m := range cfg.Models { models[m.Name] = m for _, a := range m.Aliases { @@ -212,14 +213,14 @@ func (r *Router) SetEndpoints(ctx context.Context, infos []endpoint.Info) { // resolve maps a client-supplied model name onto the catalog. Unknown names are // passed through unchanged: an endpoint may serve a model the catalog does not // list, and refusing it would be unhelpful. -func (r *Router) resolve(name string) (canonical string, m config.Model) { +func (r *Router) resolve(name string) (canonical string, m catalog.Model) { if e, ok := r.models[name]; ok { return e.Name, e } - return name, config.Model{Name: name} + return name, catalog.Model{Name: name} } -func (r *Router) profileFor(m config.Model, override string) *scheduler.Profile { +func (r *Router) profileFor(m catalog.Model, override string) *scheduler.Profile { for _, name := range []string{override, m.RoutingProfile, r.cfg.DefaultProfile} { if p, ok := r.profiles[name]; ok && name != "" { return p @@ -279,7 +280,7 @@ func (r *Router) dispatch(ctx context.Context, w http.ResponseWriter, req *parse var errNoCapacity = errors.New("no endpoint available") -func (r *Router) dispatchOne(ctx context.Context, w http.ResponseWriter, req *parsedRequest, model string, entry config.Model) dispatchResult { +func (r *Router) dispatchOne(ctx context.Context, w http.ResponseWriter, req *parsedRequest, model string, entry catalog.Model) dispatchResult { candidates := r.reg.Snapshots(model) if len(candidates) == 0 { return dispatchResult{err: fmt.Errorf("%w: no endpoint serves %q", errNoCapacity, model)} diff --git a/internal/router/router_test.go b/internal/router/router_test.go index 7004b01..d3947f8 100644 --- a/internal/router/router_test.go +++ b/internal/router/router_test.go @@ -14,8 +14,18 @@ import ( "time" "gitea.mynodes.xyz/ilker/llmgw/internal/config" + "gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog" ) +// catModel builds a minimal catalog entry for router tests, which care only +// about names, aliases and fallbacks. +func catModel(name string, aliases []string, fallbacks ...string) catalog.Model { + return catalog.Model{ + Name: name, Source: "/models/" + name, Aliases: aliases, Fallbacks: fallbacks, + Profiles: []catalog.Profile{{Name: "p", GPUClass: "h100-80g", TP: 1}}, + } +} + // fakeEngine serves the minimum an endpoint needs: the metric contract plus a // scripted chat response. type fakeEngine struct { @@ -163,7 +173,7 @@ func TestFallbackModel(t *testing.T) { config.Endpoint{ID: "p", URL: primary.srv.URL, Models: []string{"primary"}}, config.Endpoint{ID: "b", URL: backup.srv.URL, Models: []string{"backup"}}, ) - cfg.Models = []config.Model{{Name: "primary", Fallbacks: []string{"backup"}}} + cfg.Models = []catalog.Model{catModel("primary", nil, "backup")} rt := newTestRouter(t, cfg) rec := post(t, rt.Handler(), map[string]any{"model": "primary", "messages": []any{}}, nil) @@ -184,7 +194,7 @@ func TestAliasResolution(t *testing.T) { cfg := baseConfig(config.Endpoint{ ID: "a", URL: e.srv.URL, Models: []string{"llama-3.3-70b"}, }) - cfg.Models = []config.Model{{Name: "llama-3.3-70b", Aliases: []string{"llama"}}} + cfg.Models = []catalog.Model{catModel("llama-3.3-70b", []string{"llama"})} rt := newTestRouter(t, cfg) rec := post(t, rt.Handler(), map[string]any{"model": "llama", "messages": []any{}}, nil) @@ -223,7 +233,7 @@ func TestMissingModelIs400(t *testing.T) { func TestModelsEndpointListsAliases(t *testing.T) { e := newFakeEngine(t, "llama-3.3-70b") cfg := baseConfig(config.Endpoint{ID: "a", URL: e.srv.URL, Models: []string{"llama-3.3-70b"}}) - cfg.Models = []config.Model{{Name: "llama-3.3-70b", Aliases: []string{"llama"}}} + cfg.Models = []catalog.Model{catModel("llama-3.3-70b", []string{"llama"})} rt := newTestRouter(t, cfg) req := httptest.NewRequest(http.MethodGet, "/v1/models", nil) diff --git a/llmgw-plan.md b/llmgw-plan.md index d4f0939..8732104 100644 --- a/llmgw-plan.md +++ b/llmgw-plan.md @@ -1005,10 +1005,14 @@ work_dir: /var/lib/llmgw-agent - [x] NVML backend behind the `nvml` build tag, with NVLink topology. It compiles in CI but has not run against a driver: **verify on a GPU node before relying on the NVLink domains it reports.** -- [ ] connect-go stream, register/heartbeat, mTLS -- **Exit criterion:** start and stop a model on a real node via `ctl`. Not met: - the control plane connection is the remaining piece, and it pairs naturally - with phase 3 rather than standing alone. +- [x] connect-go stream, register/heartbeat, shard status +- [ ] mTLS — the stream is plaintext HTTP/2 today and must stay on a trusted network +- **Exit criterion (met, via configuration rather than `ctl`):** a model starts + and stops on a node. Demonstrated with real binaries: `llmgw server` and + `llmgw agent` as separate processes, the agent registering, a replica being + placed, an engine launched through the vLLM adapter, traffic served, and the + node's loss detected when the agent is killed. `ctl` itself is phase 3's + remaining item. ### Phase 3 — Control plane (2–3 weeks) - [ ] SQLite store and migrations