Survive a control plane restart without orphaning live engines
Not persisting replica state left one hole: restarting the control plane made it forget every replica while the engines kept running. They held GPUs the inventory believed were free, so the next placement failed to allocate, with an error naming the new process rather than the old one. The fix belongs on the agent, not in a database. A row can say "running" about a process that died; the fleet cannot. - the supervisor writes the full shard spec, not just a pid and command line. A pid identifies a process; the spec is what lets it be described back to the control plane and matched to a replica. - Adopt takes over an engine a previous agent left running. It is not this process's child, so it cannot be waited on: it is watched by polling and stopped by signalling its group directly. Its earlier output is gone and the log says so rather than presenting emptiness as silence. - every Register reports what the agent is running, not only the first, because a reconnect may be a control plane restart and has to recover the same way. - the reconciler rebuilds replicas from those reports and reclaims their GPUs. Anything whose model or profile left the catalog is stopped; nothing will route to it and it is holding cards. - an engine the agent cannot identify is killed. A process holding GPUs that nothing can describe is worse than no process, because placement hands those GPUs out again. Stop was also wrong for adopted shards: it checked for an exec.Cmd, which they do not have, and reported them stopped without killing anything. TestAgentReconnects had been passing for no reason. httptest's CloseClientConnections does nothing to an h2c stream — prior-knowledge h2c hijacks the connection and httptest stops tracking hijacked ones — so the agent never reconnected, and the assertion on the connection map passed on a stale entry. Tests now sever connections through their own listener and count opened streams. Plan §11 also records what multi-node still needs, since placement landing early makes it look further along than it is: workers never become Ready because the agent health-checks every shard while only rank 0 serves HTTP; RayCommand is written and tested but never called; and startup is fire-and-forget. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
1 parent
40427b2705
commit
75a8e76706
13 files changed
+1312
-145
No files matched your search
@@ -134,6 +134,42 @@ 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.
|
||||
|
||||
### Surviving a control plane restart
|
||||
|
||||
Not persisting replica state left one hole that had to be closed before it
|
||||
mattered: restarting the control plane orphaned live engines. They kept holding
|
||||
GPUs the inventory believed were free, so the next placement failed to allocate
|
||||
with an error naming the new process rather than the old one.
|
||||
|
||||
The fix is on the agent, not in a database, and it cannot disagree with reality
|
||||
the way a stale row can:
|
||||
|
||||
- The supervisor now writes the full shard spec, not just a pid and command
|
||||
line. A pid identifies a process; the spec is what lets it be described back
|
||||
to the control plane and matched to a replica.
|
||||
- `Adopt` takes over a still-running engine from a previous agent. It cannot be
|
||||
waited on — it is not this process's child — so it is watched by polling, and
|
||||
stopped by signalling its group directly. Its earlier output is gone, and the
|
||||
log says so rather than presenting emptiness as silence.
|
||||
- Every `Register`, not only the first, reports what the agent is running. A
|
||||
reconnect may be a control plane restart and has to recover the same way.
|
||||
- The reconciler rebuilds replicas from those reports and reclaims their GPUs.
|
||||
Anything whose model or profile left the catalog is stopped: nothing will
|
||||
route to it and it is holding cards.
|
||||
- An engine the agent cannot identify — a corrupt state file — is killed. A
|
||||
process holding GPUs that nothing can describe is worse than no process,
|
||||
because placement will hand those GPUs out again.
|
||||
|
||||
One mechanism covers agent restarts and control plane restarts.
|
||||
|
||||
### A test that was passing for no reason
|
||||
|
||||
`TestAgentReconnects` used `httptest.CloseClientConnections`, which does nothing
|
||||
to an h2c stream: prior-knowledge h2c hijacks the connection, and httptest stops
|
||||
tracking hijacked connections. The agent never reconnected and the assertion —
|
||||
on the connection map, which still held a stale entry — passed anyway. Tests now
|
||||
sever connections through their own listener and count opened streams.
|
||||
|
||||
### Still missing for phase 2
|
||||
|
||||
- mTLS. The stream is plaintext HTTP/2 today, so it must stay on a trusted
|
||||
|
||||
+287
-71
@@ -90,7 +90,7 @@ func (x ShardStatus_State) Number() protoreflect.EnumNumber {
|
||||
|
||||
// Deprecated: Use ShardStatus_State.Descriptor instead.
|
||||
func (ShardStatus_State) EnumDescriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{6, 0}
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{7, 0}
|
||||
}
|
||||
|
||||
// AgentMessage is anything a node tells the control plane.
|
||||
@@ -327,8 +327,17 @@ type Register struct {
|
||||
// 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
|
||||
// running_shards are the engines this agent is already supervising, including
|
||||
// ones adopted from a previous agent.
|
||||
//
|
||||
// This is what lets a control plane restart recover instead of orphaning live
|
||||
// engines. Without it the control plane forgets every replica, places new
|
||||
// ones, and the old engines keep holding GPUs the inventory believes are free
|
||||
// — so the new engines fail to allocate, with an error naming the wrong
|
||||
// process (plan §6.5).
|
||||
RunningShards []*RunningShard `protobuf:"bytes,9,rep,name=running_shards,json=runningShards,proto3" json:"running_shards,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
}
|
||||
|
||||
func (x *Register) Reset() {
|
||||
@@ -417,6 +426,169 @@ func (x *Register) GetEngineVersions() map[string]string {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (x *Register) GetRunningShards() []*RunningShard {
|
||||
if x != nil {
|
||||
return x.RunningShards
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// RunningShard describes an engine already running on a node.
|
||||
//
|
||||
// It carries enough for the control plane to rebuild the replica it belongs to,
|
||||
// which is more than a shard id: the replica's identity, its pool and fabric,
|
||||
// and the leader endpoint clients reach it through.
|
||||
type RunningShard 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"`
|
||||
GpuIndices []int32 `protobuf:"varint,5,rep,packed,name=gpu_indices,json=gpuIndices,proto3" json:"gpu_indices,omitempty"`
|
||||
Hcas []string `protobuf:"bytes,6,rep,name=hcas,proto3" json:"hcas,omitempty"`
|
||||
HttpPort int32 `protobuf:"varint,7,opt,name=http_port,json=httpPort,proto3" json:"http_port,omitempty"`
|
||||
State ShardStatus_State `protobuf:"varint,8,opt,name=state,proto3,enum=llmgw.v1.ShardStatus_State" json:"state,omitempty"`
|
||||
// model and profile name the catalog entry this shard was started from.
|
||||
Model string `protobuf:"bytes,9,opt,name=model,proto3" json:"model,omitempty"`
|
||||
Profile string `protobuf:"bytes,10,opt,name=profile,proto3" json:"profile,omitempty"`
|
||||
// pool and fabric are echoed back from the StartShard that created it; the
|
||||
// agent has no way to derive them.
|
||||
Pool string `protobuf:"bytes,11,opt,name=pool,proto3" json:"pool,omitempty"`
|
||||
Fabric string `protobuf:"bytes,12,opt,name=fabric,proto3" json:"fabric,omitempty"`
|
||||
LeaderAddr string `protobuf:"bytes,13,opt,name=leader_addr,json=leaderAddr,proto3" json:"leader_addr,omitempty"`
|
||||
LeaderPort int32 `protobuf:"varint,14,opt,name=leader_port,json=leaderPort,proto3" json:"leader_port,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
}
|
||||
|
||||
func (x *RunningShard) Reset() {
|
||||
*x = RunningShard{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[3]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
|
||||
func (x *RunningShard) String() string {
|
||||
return protoimpl.X.MessageStringOf(x)
|
||||
}
|
||||
|
||||
func (*RunningShard) ProtoMessage() {}
|
||||
|
||||
func (x *RunningShard) 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 RunningShard.ProtoReflect.Descriptor instead.
|
||||
func (*RunningShard) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{3}
|
||||
}
|
||||
|
||||
func (x *RunningShard) GetReplicaId() string {
|
||||
if x != nil {
|
||||
return x.ReplicaId
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *RunningShard) GetShardId() string {
|
||||
if x != nil {
|
||||
return x.ShardId
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *RunningShard) GetRank() int32 {
|
||||
if x != nil {
|
||||
return x.Rank
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func (x *RunningShard) GetNnodes() int32 {
|
||||
if x != nil {
|
||||
return x.Nnodes
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func (x *RunningShard) GetGpuIndices() []int32 {
|
||||
if x != nil {
|
||||
return x.GpuIndices
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (x *RunningShard) GetHcas() []string {
|
||||
if x != nil {
|
||||
return x.Hcas
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (x *RunningShard) GetHttpPort() int32 {
|
||||
if x != nil {
|
||||
return x.HttpPort
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func (x *RunningShard) GetState() ShardStatus_State {
|
||||
if x != nil {
|
||||
return x.State
|
||||
}
|
||||
return ShardStatus_STATE_UNSPECIFIED
|
||||
}
|
||||
|
||||
func (x *RunningShard) GetModel() string {
|
||||
if x != nil {
|
||||
return x.Model
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *RunningShard) GetProfile() string {
|
||||
if x != nil {
|
||||
return x.Profile
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *RunningShard) GetPool() string {
|
||||
if x != nil {
|
||||
return x.Pool
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *RunningShard) GetFabric() string {
|
||||
if x != nil {
|
||||
return x.Fabric
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *RunningShard) GetLeaderAddr() string {
|
||||
if x != nil {
|
||||
return x.LeaderAddr
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *RunningShard) GetLeaderPort() int32 {
|
||||
if x != nil {
|
||||
return x.LeaderPort
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
type GPU struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
Index int32 `protobuf:"varint,1,opt,name=index,proto3" json:"index,omitempty"`
|
||||
@@ -442,7 +614,7 @@ type GPU struct {
|
||||
|
||||
func (x *GPU) Reset() {
|
||||
*x = GPU{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[3]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[4]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
@@ -454,7 +626,7 @@ func (x *GPU) String() string {
|
||||
func (*GPU) ProtoMessage() {}
|
||||
|
||||
func (x *GPU) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[3]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[4]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
@@ -467,7 +639,7 @@ func (x *GPU) ProtoReflect() protoreflect.Message {
|
||||
|
||||
// Deprecated: Use GPU.ProtoReflect.Descriptor instead.
|
||||
func (*GPU) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{3}
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{4}
|
||||
}
|
||||
|
||||
func (x *GPU) GetIndex() int32 {
|
||||
@@ -582,7 +754,7 @@ type FabricAttachment struct {
|
||||
|
||||
func (x *FabricAttachment) Reset() {
|
||||
*x = FabricAttachment{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[4]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[5]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
@@ -594,7 +766,7 @@ func (x *FabricAttachment) String() string {
|
||||
func (*FabricAttachment) ProtoMessage() {}
|
||||
|
||||
func (x *FabricAttachment) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[4]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[5]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
@@ -607,7 +779,7 @@ func (x *FabricAttachment) ProtoReflect() protoreflect.Message {
|
||||
|
||||
// Deprecated: Use FabricAttachment.ProtoReflect.Descriptor instead.
|
||||
func (*FabricAttachment) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{4}
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{5}
|
||||
}
|
||||
|
||||
func (x *FabricAttachment) GetFabric() string {
|
||||
@@ -645,7 +817,7 @@ type Heartbeat struct {
|
||||
|
||||
func (x *Heartbeat) Reset() {
|
||||
*x = Heartbeat{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[5]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[6]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
@@ -657,7 +829,7 @@ func (x *Heartbeat) String() string {
|
||||
func (*Heartbeat) ProtoMessage() {}
|
||||
|
||||
func (x *Heartbeat) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[5]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[6]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
@@ -670,7 +842,7 @@ func (x *Heartbeat) ProtoReflect() protoreflect.Message {
|
||||
|
||||
// Deprecated: Use Heartbeat.ProtoReflect.Descriptor instead.
|
||||
func (*Heartbeat) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{5}
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{6}
|
||||
}
|
||||
|
||||
func (x *Heartbeat) GetNodeId() string {
|
||||
@@ -712,7 +884,7 @@ type ShardStatus struct {
|
||||
|
||||
func (x *ShardStatus) Reset() {
|
||||
*x = ShardStatus{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[6]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[7]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
@@ -724,7 +896,7 @@ func (x *ShardStatus) String() string {
|
||||
func (*ShardStatus) ProtoMessage() {}
|
||||
|
||||
func (x *ShardStatus) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[6]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[7]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
@@ -737,7 +909,7 @@ func (x *ShardStatus) ProtoReflect() protoreflect.Message {
|
||||
|
||||
// Deprecated: Use ShardStatus.ProtoReflect.Descriptor instead.
|
||||
func (*ShardStatus) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{6}
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{7}
|
||||
}
|
||||
|
||||
func (x *ShardStatus) GetNodeId() string {
|
||||
@@ -808,13 +980,17 @@ type StartShard struct {
|
||||
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
|
||||
// pool and fabric are control plane concepts the agent cannot derive. They
|
||||
// are sent so the agent can echo them back in RunningShard after a restart.
|
||||
Pool string `protobuf:"bytes,13,opt,name=pool,proto3" json:"pool,omitempty"`
|
||||
Fabric string `protobuf:"bytes,14,opt,name=fabric,proto3" json:"fabric,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
}
|
||||
|
||||
func (x *StartShard) Reset() {
|
||||
*x = StartShard{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[7]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[8]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
@@ -826,7 +1002,7 @@ func (x *StartShard) String() string {
|
||||
func (*StartShard) ProtoMessage() {}
|
||||
|
||||
func (x *StartShard) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[7]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[8]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
@@ -839,7 +1015,7 @@ func (x *StartShard) ProtoReflect() protoreflect.Message {
|
||||
|
||||
// Deprecated: Use StartShard.ProtoReflect.Descriptor instead.
|
||||
func (*StartShard) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{7}
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{8}
|
||||
}
|
||||
|
||||
func (x *StartShard) GetReplicaId() string {
|
||||
@@ -926,6 +1102,20 @@ func (x *StartShard) GetStartupTimeoutMs() int64 {
|
||||
return 0
|
||||
}
|
||||
|
||||
func (x *StartShard) GetPool() string {
|
||||
if x != nil {
|
||||
return x.Pool
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *StartShard) GetFabric() string {
|
||||
if x != nil {
|
||||
return x.Fabric
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
type EngineProfile struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
Name string `protobuf:"bytes,1,opt,name=name,proto3" json:"name,omitempty"`
|
||||
@@ -945,7 +1135,7 @@ type EngineProfile struct {
|
||||
|
||||
func (x *EngineProfile) Reset() {
|
||||
*x = EngineProfile{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[8]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[9]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
@@ -957,7 +1147,7 @@ func (x *EngineProfile) String() string {
|
||||
func (*EngineProfile) ProtoMessage() {}
|
||||
|
||||
func (x *EngineProfile) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[8]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[9]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
@@ -970,7 +1160,7 @@ func (x *EngineProfile) ProtoReflect() protoreflect.Message {
|
||||
|
||||
// Deprecated: Use EngineProfile.ProtoReflect.Descriptor instead.
|
||||
func (*EngineProfile) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{8}
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{9}
|
||||
}
|
||||
|
||||
func (x *EngineProfile) GetName() string {
|
||||
@@ -1060,7 +1250,7 @@ type StopShard struct {
|
||||
|
||||
func (x *StopShard) Reset() {
|
||||
*x = StopShard{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[9]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[10]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
@@ -1072,7 +1262,7 @@ func (x *StopShard) String() string {
|
||||
func (*StopShard) ProtoMessage() {}
|
||||
|
||||
func (x *StopShard) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[9]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[10]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
@@ -1085,7 +1275,7 @@ func (x *StopShard) ProtoReflect() protoreflect.Message {
|
||||
|
||||
// Deprecated: Use StopShard.ProtoReflect.Descriptor instead.
|
||||
func (*StopShard) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{9}
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{10}
|
||||
}
|
||||
|
||||
func (x *StopShard) GetShardId() string {
|
||||
@@ -1105,7 +1295,7 @@ type DrainShard struct {
|
||||
|
||||
func (x *DrainShard) Reset() {
|
||||
*x = DrainShard{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[10]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[11]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
@@ -1117,7 +1307,7 @@ func (x *DrainShard) String() string {
|
||||
func (*DrainShard) ProtoMessage() {}
|
||||
|
||||
func (x *DrainShard) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[10]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[11]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
@@ -1130,7 +1320,7 @@ func (x *DrainShard) ProtoReflect() protoreflect.Message {
|
||||
|
||||
// Deprecated: Use DrainShard.ProtoReflect.Descriptor instead.
|
||||
func (*DrainShard) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{10}
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{11}
|
||||
}
|
||||
|
||||
func (x *DrainShard) GetShardId() string {
|
||||
@@ -1152,7 +1342,7 @@ type Ack struct {
|
||||
|
||||
func (x *Ack) Reset() {
|
||||
*x = Ack{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[11]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[12]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
@@ -1164,7 +1354,7 @@ func (x *Ack) String() string {
|
||||
func (*Ack) ProtoMessage() {}
|
||||
|
||||
func (x *Ack) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[11]
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[12]
|
||||
if x != nil {
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
if ms.LoadMessageInfo() == nil {
|
||||
@@ -1177,7 +1367,7 @@ func (x *Ack) ProtoReflect() protoreflect.Message {
|
||||
|
||||
// Deprecated: Use Ack.ProtoReflect.Descriptor instead.
|
||||
func (*Ack) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{11}
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{12}
|
||||
}
|
||||
|
||||
func (x *Ack) GetMessageId() string {
|
||||
@@ -1216,7 +1406,7 @@ const file_llmgw_v1_agent_proto_rawDesc = "" +
|
||||
"\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" +
|
||||
"\x03msg\"\xa0\x04\n" +
|
||||
"\bRegister\x12\x17\n" +
|
||||
"\anode_id\x18\x01 \x01(\tR\x06nodeId\x12\x1a\n" +
|
||||
"\bhostname\x18\x02 \x01(\tR\bhostname\x12\x1b\n" +
|
||||
@@ -1225,13 +1415,34 @@ const file_llmgw_v1_agent_proto_rawDesc = "" +
|
||||
"\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" +
|
||||
"\x0fengine_versions\x18\b \x03(\v2&.llmgw.v1.Register.EngineVersionsEntryR\x0eengineVersions\x12=\n" +
|
||||
"\x0erunning_shards\x18\t \x03(\v2\x16.llmgw.v1.RunningShardR\rrunningShards\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" +
|
||||
"\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"\x97\x03\n" +
|
||||
"\fRunningShard\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" +
|
||||
"\vgpu_indices\x18\x05 \x03(\x05R\n" +
|
||||
"gpuIndices\x12\x12\n" +
|
||||
"\x04hcas\x18\x06 \x03(\tR\x04hcas\x12\x1b\n" +
|
||||
"\thttp_port\x18\a \x01(\x05R\bhttpPort\x121\n" +
|
||||
"\x05state\x18\b \x01(\x0e2\x1b.llmgw.v1.ShardStatus.StateR\x05state\x12\x14\n" +
|
||||
"\x05model\x18\t \x01(\tR\x05model\x12\x18\n" +
|
||||
"\aprofile\x18\n" +
|
||||
" \x01(\tR\aprofile\x12\x12\n" +
|
||||
"\x04pool\x18\v \x01(\tR\x04pool\x12\x16\n" +
|
||||
"\x06fabric\x18\f \x01(\tR\x06fabric\x12\x1f\n" +
|
||||
"\vleader_addr\x18\r \x01(\tR\n" +
|
||||
"leaderAddr\x12\x1f\n" +
|
||||
"\vleader_port\x18\x0e \x01(\x05R\n" +
|
||||
"leaderPort\"\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" +
|
||||
@@ -1276,7 +1487,7 @@ const file_llmgw_v1_agent_proto_rawDesc = "" +
|
||||
"\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" +
|
||||
"\fSTATE_FAILED\x10\a\"\xba\x03\n" +
|
||||
"\n" +
|
||||
"StartShard\x12\x1d\n" +
|
||||
"\n" +
|
||||
@@ -1295,7 +1506,9 @@ const file_llmgw_v1_agent_proto_rawDesc = "" +
|
||||
"\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" +
|
||||
"\x12startup_timeout_ms\x18\f \x01(\x03R\x10startupTimeoutMs\x12\x12\n" +
|
||||
"\x04pool\x18\r \x01(\tR\x04pool\x12\x16\n" +
|
||||
"\x06fabric\x18\x0e \x01(\tR\x06fabric\"\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" +
|
||||
@@ -1341,47 +1554,50 @@ func file_llmgw_v1_agent_proto_rawDescGZIP() []byte {
|
||||
}
|
||||
|
||||
var file_llmgw_v1_agent_proto_enumTypes = make([]protoimpl.EnumInfo, 1)
|
||||
var file_llmgw_v1_agent_proto_msgTypes = make([]protoimpl.MessageInfo, 15)
|
||||
var file_llmgw_v1_agent_proto_msgTypes = make([]protoimpl.MessageInfo, 16)
|
||||
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
|
||||
(*RunningShard)(nil), // 4: llmgw.v1.RunningShard
|
||||
(*GPU)(nil), // 5: llmgw.v1.GPU
|
||||
(*FabricAttachment)(nil), // 6: llmgw.v1.FabricAttachment
|
||||
(*Heartbeat)(nil), // 7: llmgw.v1.Heartbeat
|
||||
(*ShardStatus)(nil), // 8: llmgw.v1.ShardStatus
|
||||
(*StartShard)(nil), // 9: llmgw.v1.StartShard
|
||||
(*EngineProfile)(nil), // 10: llmgw.v1.EngineProfile
|
||||
(*StopShard)(nil), // 11: llmgw.v1.StopShard
|
||||
(*DrainShard)(nil), // 12: llmgw.v1.DrainShard
|
||||
(*Ack)(nil), // 13: llmgw.v1.Ack
|
||||
nil, // 14: llmgw.v1.Register.LabelsEntry
|
||||
nil, // 15: llmgw.v1.Register.EngineVersionsEntry
|
||||
nil, // 16: 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
|
||||
7, // 1: llmgw.v1.AgentMessage.heartbeat:type_name -> llmgw.v1.Heartbeat
|
||||
8, // 2: llmgw.v1.AgentMessage.shard:type_name -> llmgw.v1.ShardStatus
|
||||
9, // 3: llmgw.v1.ControlMessage.start:type_name -> llmgw.v1.StartShard
|
||||
11, // 4: llmgw.v1.ControlMessage.stop:type_name -> llmgw.v1.StopShard
|
||||
12, // 5: llmgw.v1.ControlMessage.drain:type_name -> llmgw.v1.DrainShard
|
||||
13, // 6: llmgw.v1.ControlMessage.ack:type_name -> llmgw.v1.Ack
|
||||
14, // 7: llmgw.v1.Register.labels:type_name -> llmgw.v1.Register.LabelsEntry
|
||||
5, // 8: llmgw.v1.Register.gpus:type_name -> llmgw.v1.GPU
|
||||
6, // 9: llmgw.v1.Register.fabrics:type_name -> llmgw.v1.FabricAttachment
|
||||
15, // 10: llmgw.v1.Register.engine_versions:type_name -> llmgw.v1.Register.EngineVersionsEntry
|
||||
4, // 11: llmgw.v1.Register.running_shards:type_name -> llmgw.v1.RunningShard
|
||||
0, // 12: llmgw.v1.RunningShard.state:type_name -> llmgw.v1.ShardStatus.State
|
||||
0, // 13: llmgw.v1.ShardStatus.state:type_name -> llmgw.v1.ShardStatus.State
|
||||
10, // 14: llmgw.v1.StartShard.profile:type_name -> llmgw.v1.EngineProfile
|
||||
16, // 15: llmgw.v1.EngineProfile.env:type_name -> llmgw.v1.EngineProfile.EnvEntry
|
||||
1, // 16: llmgw.v1.AgentService.Connect:input_type -> llmgw.v1.AgentMessage
|
||||
2, // 17: llmgw.v1.AgentService.Connect:output_type -> llmgw.v1.ControlMessage
|
||||
17, // [17:18] is the sub-list for method output_type
|
||||
16, // [16:17] is the sub-list for method input_type
|
||||
16, // [16:16] is the sub-list for extension type_name
|
||||
16, // [16:16] is the sub-list for extension extendee
|
||||
0, // [0:16] is the sub-list for field type_name
|
||||
}
|
||||
|
||||
func init() { file_llmgw_v1_agent_proto_init() }
|
||||
@@ -1406,7 +1622,7 @@ func file_llmgw_v1_agent_proto_init() {
|
||||
GoPackagePath: reflect.TypeOf(x{}).PkgPath(),
|
||||
RawDescriptor: unsafe.Slice(unsafe.StringData(file_llmgw_v1_agent_proto_rawDesc), len(file_llmgw_v1_agent_proto_rawDesc)),
|
||||
NumEnums: 1,
|
||||
NumMessages: 15,
|
||||
NumMessages: 16,
|
||||
NumExtensions: 0,
|
||||
NumServices: 1,
|
||||
},
|
||||
|
||||
@@ -57,6 +57,43 @@ message Register {
|
||||
// engine_versions decides which CLI flags and metric names apply; both drift
|
||||
// between engine releases (plan §14).
|
||||
map<string, string> engine_versions = 8;
|
||||
|
||||
// running_shards are the engines this agent is already supervising, including
|
||||
// ones adopted from a previous agent.
|
||||
//
|
||||
// This is what lets a control plane restart recover instead of orphaning live
|
||||
// engines. Without it the control plane forgets every replica, places new
|
||||
// ones, and the old engines keep holding GPUs the inventory believes are free
|
||||
// — so the new engines fail to allocate, with an error naming the wrong
|
||||
// process (plan §6.5).
|
||||
repeated RunningShard running_shards = 9;
|
||||
}
|
||||
|
||||
// RunningShard describes an engine already running on a node.
|
||||
//
|
||||
// It carries enough for the control plane to rebuild the replica it belongs to,
|
||||
// which is more than a shard id: the replica's identity, its pool and fabric,
|
||||
// and the leader endpoint clients reach it through.
|
||||
message RunningShard {
|
||||
string replica_id = 1;
|
||||
string shard_id = 2;
|
||||
int32 rank = 3;
|
||||
int32 nnodes = 4;
|
||||
repeated int32 gpu_indices = 5;
|
||||
repeated string hcas = 6;
|
||||
int32 http_port = 7;
|
||||
ShardStatus.State state = 8;
|
||||
|
||||
// model and profile name the catalog entry this shard was started from.
|
||||
string model = 9;
|
||||
string profile = 10;
|
||||
// pool and fabric are echoed back from the StartShard that created it; the
|
||||
// agent has no way to derive them.
|
||||
string pool = 11;
|
||||
string fabric = 12;
|
||||
|
||||
string leader_addr = 13;
|
||||
int32 leader_port = 14;
|
||||
}
|
||||
|
||||
message GPU {
|
||||
@@ -142,6 +179,11 @@ message StartShard {
|
||||
|
||||
EngineProfile profile = 11;
|
||||
int64 startup_timeout_ms = 12;
|
||||
|
||||
// pool and fabric are control plane concepts the agent cannot derive. They
|
||||
// are sent so the agent can echo them back in RunningShard after a restart.
|
||||
string pool = 13;
|
||||
string fabric = 14;
|
||||
}
|
||||
|
||||
message EngineProfile {
|
||||
|
||||
+97
-5
@@ -149,6 +149,7 @@ func (a *Agent) Run(ctx context.Context) error {
|
||||
// an operator can see it and why it is useless.
|
||||
a.log.Error("hardware discovery failed", "err", err)
|
||||
}
|
||||
a.recoverOrphans()
|
||||
|
||||
attempt := 0
|
||||
for ctx.Err() == nil {
|
||||
@@ -185,6 +186,93 @@ func (a *Agent) reconnectDelay(attempt int) time.Duration {
|
||||
return d/2 + time.Duration(rand.Int64N(int64(d/2))) //nolint:gosec // jitter
|
||||
}
|
||||
|
||||
// recoverOrphans takes back engines a previous agent left running.
|
||||
//
|
||||
// Engines deliberately outlive their agent — the router talks to them directly,
|
||||
// so an agent restart must not interrupt served traffic. The cost is that they
|
||||
// have to be found again, and anything that cannot be identified has to be
|
||||
// killed: a process holding GPUs that nothing can describe is worse than no
|
||||
// process at all, because placement will hand those GPUs out again.
|
||||
func (a *Agent) recoverOrphans() {
|
||||
orphans, err := a.sup.FindOrphans()
|
||||
if err != nil {
|
||||
a.log.Warn("scanning for leftover engines failed", "err", err)
|
||||
return
|
||||
}
|
||||
|
||||
for _, o := range orphans {
|
||||
switch {
|
||||
case !o.Alive:
|
||||
// The process is gone; only its state file remains.
|
||||
if err := a.sup.KillOrphan(o); err != nil {
|
||||
a.log.Warn("cleaning up a dead shard's state failed", "shard", o.ShardID, "err", err)
|
||||
}
|
||||
case o.ReplicaID == "":
|
||||
// Alive, but its state file could not be read, so it cannot be
|
||||
// described to the control plane or matched to a replica.
|
||||
a.log.Warn("killing an unidentifiable engine", "shard", o.ShardID, "pid", o.PID)
|
||||
if err := a.sup.KillOrphan(o); err != nil {
|
||||
a.log.Warn("killing it failed", "shard", o.ShardID, "err", err)
|
||||
}
|
||||
default:
|
||||
sh, err := a.sup.Adopt(o)
|
||||
if err != nil {
|
||||
a.log.Warn("adopting a shard failed", "shard", o.ShardID, "err", err)
|
||||
continue
|
||||
}
|
||||
a.mu.Lock()
|
||||
a.shardReplica[o.ShardID] = o.ReplicaID
|
||||
a.mu.Unlock()
|
||||
a.log.Info("adopted a running engine",
|
||||
"shard", o.ShardID, "replica", o.ReplicaID, "pid", o.PID)
|
||||
_ = sh
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// runningShards describes what this agent is supervising, for registration.
|
||||
func (a *Agent) runningShards() []*llmgwv1.RunningShard {
|
||||
var out []*llmgwv1.RunningShard
|
||||
for _, st := range a.sup.Statuses() {
|
||||
if st.State.Terminal() {
|
||||
continue
|
||||
}
|
||||
sh, err := a.sup.Get(st.ID)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
spec := sh.Spec
|
||||
out = append(out, &llmgwv1.RunningShard{
|
||||
ReplicaId: spec.ReplicaID,
|
||||
ShardId: st.ID,
|
||||
Rank: int32(spec.Rank),
|
||||
Nnodes: int32(spec.NNodes),
|
||||
GpuIndices: int32Slice(spec.GPUIndices),
|
||||
Hcas: spec.HCAs,
|
||||
HttpPort: int32(spec.HTTPPort),
|
||||
State: agentsvc.StateToWire(st.State),
|
||||
Model: spec.Profile.ServedModelName,
|
||||
Profile: spec.Profile.Name,
|
||||
Pool: spec.Pool,
|
||||
Fabric: spec.Fabric,
|
||||
LeaderAddr: a.cfg.MgmtAddr,
|
||||
LeaderPort: int32(spec.LeaderPort),
|
||||
})
|
||||
}
|
||||
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
|
||||
}
|
||||
|
||||
// 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)
|
||||
@@ -208,12 +296,16 @@ func (a *Agent) session(ctx context.Context) error {
|
||||
inv := a.inv
|
||||
a.mu.Unlock()
|
||||
|
||||
reg := agentsvc.RegisterFromInventory(
|
||||
a.cfg.NodeID, a.cfg.Hostname, a.cfg.MgmtAddr, a.cfg.Version,
|
||||
a.cfg.Labels, inv, a.cfg.Fabrics)
|
||||
// Every registration reports what is already running, not just the first
|
||||
// one: a reconnect may be a control plane restart, and it has to recover
|
||||
// the same way (plan §6.5).
|
||||
reg.RunningShards = a.runningShards()
|
||||
|
||||
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),
|
||||
},
|
||||
Msg: &llmgwv1.AgentMessage_Register{Register: reg},
|
||||
}); err != nil {
|
||||
return fmt.Errorf("registering: %w", err)
|
||||
}
|
||||
|
||||
@@ -12,6 +12,9 @@ import (
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"syscall"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@@ -25,8 +28,15 @@ import (
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
|
||||
)
|
||||
|
||||
// statusRecorder captures what the control plane hears about shards.
|
||||
type statusRecorder struct{ ch chan status }
|
||||
// statusRecorder captures what the control plane hears about shards, and what
|
||||
// it is asked to adopt.
|
||||
type statusRecorder struct {
|
||||
ch chan status
|
||||
|
||||
mu sync.Mutex
|
||||
adopted []reconciler.AdoptedShard
|
||||
unclaim []string // shard ids to report back as unclaimed
|
||||
}
|
||||
|
||||
type status struct {
|
||||
replicaID, shardID string
|
||||
@@ -43,6 +53,32 @@ func (r *statusRecorder) OnShardStatus(replicaID, shardID string, st supervisor.
|
||||
}
|
||||
}
|
||||
|
||||
func (r *statusRecorder) Adopt(shards []reconciler.AdoptedShard) []reconciler.AdoptedShard {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.adopted = append(r.adopted, shards...)
|
||||
|
||||
var unclaimed []reconciler.AdoptedShard
|
||||
for _, s := range shards {
|
||||
if slices.Contains(r.unclaim, s.ShardID) {
|
||||
unclaimed = append(unclaimed, s)
|
||||
}
|
||||
}
|
||||
return unclaimed
|
||||
}
|
||||
|
||||
func (r *statusRecorder) adoptedShards() []reconciler.AdoptedShard {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
return slices.Clone(r.adopted)
|
||||
}
|
||||
|
||||
func (r *statusRecorder) markUnclaimed(ids ...string) {
|
||||
r.mu.Lock()
|
||||
r.unclaim = append(r.unclaim, ids...)
|
||||
r.mu.Unlock()
|
||||
}
|
||||
|
||||
func (r *statusRecorder) await(t *testing.T, limit time.Duration, pred func(status) bool) status {
|
||||
t.Helper()
|
||||
deadline := time.After(limit)
|
||||
@@ -115,22 +151,84 @@ type harness struct {
|
||||
rec *statusRecorder
|
||||
srv *httptest.Server
|
||||
ag *Agent
|
||||
ln *dropListener
|
||||
// streams counts how many times an agent has opened the Connect stream,
|
||||
// which is the only honest way to tell a reconnect from a stale entry.
|
||||
streams *atomic.Int64
|
||||
}
|
||||
|
||||
// dropListener tracks accepted connections so a test can sever them.
|
||||
//
|
||||
// httptest's CloseClientConnections does not work here: h2c with prior
|
||||
// knowledge hijacks the connection, and httptest stops tracking hijacked
|
||||
// connections. A test relying on it would pass without ever reconnecting.
|
||||
type dropListener struct {
|
||||
net.Listener
|
||||
mu sync.Mutex
|
||||
conns []net.Conn
|
||||
}
|
||||
|
||||
func (l *dropListener) Accept() (net.Conn, error) {
|
||||
c, err := l.Listener.Accept()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
l.mu.Lock()
|
||||
l.conns = append(l.conns, c)
|
||||
l.mu.Unlock()
|
||||
return c, nil
|
||||
}
|
||||
|
||||
// drop severs every connection, as a control plane restart would.
|
||||
func (l *dropListener) drop() {
|
||||
l.mu.Lock()
|
||||
conns := l.conns
|
||||
l.conns = nil
|
||||
l.mu.Unlock()
|
||||
for _, c := range conns {
|
||||
_ = c.Close()
|
||||
}
|
||||
}
|
||||
|
||||
// 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()
|
||||
return startWith(t, slog.New(slog.DiscardHandler), cfgFns...)
|
||||
}
|
||||
|
||||
// startLogged is start with logging on, for diagnosing protocol behaviour.
|
||||
func startLogged(t *testing.T, cfgFns ...func(*Config)) *harness {
|
||||
t.Helper()
|
||||
return startWith(t, slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelDebug})), cfgFns...)
|
||||
}
|
||||
|
||||
func startWith(t *testing.T, log *slog.Logger, cfgFns ...func(*Config)) *harness {
|
||||
t.Helper()
|
||||
|
||||
inv := inventory.NewStore()
|
||||
rec := newRecorder()
|
||||
svc := agentsvc.New(inv, rec, slog.New(slog.DiscardHandler))
|
||||
svc := agentsvc.New(inv, rec, log)
|
||||
|
||||
path, handler := svc.Handler()
|
||||
mux := http.NewServeMux()
|
||||
mux.Handle(path, handler)
|
||||
|
||||
var streams atomic.Int64
|
||||
counting := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
streams.Add(1)
|
||||
mux.ServeHTTP(w, r)
|
||||
})
|
||||
|
||||
raw, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ln := &dropListener{Listener: raw}
|
||||
// connect's bidirectional streams need HTTP/2; h2c provides it without TLS.
|
||||
srv := httptest.NewServer(h2c.NewHandler(mux, &http2.Server{}))
|
||||
srv := httptest.NewUnstartedServer(h2c.NewHandler(counting, &http2.Server{}))
|
||||
srv.Listener = ln
|
||||
srv.Start()
|
||||
|
||||
cfg := DefaultConfig()
|
||||
cfg.NodeID = "n1"
|
||||
@@ -156,7 +254,7 @@ func start(t *testing.T, cfgFns ...func(*Config)) *harness {
|
||||
},
|
||||
}}
|
||||
|
||||
ag, err := New(cfg, client, slog.New(slog.DiscardHandler))
|
||||
ag, err := New(cfg, client, log)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
@@ -164,7 +262,7 @@ func start(t *testing.T, cfgFns ...func(*Config)) *harness {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
go func() { _ = ag.Run(ctx) }()
|
||||
|
||||
h := &harness{svc: svc, inv: inv, rec: rec, srv: srv, ag: ag}
|
||||
h := &harness{svc: svc, inv: inv, rec: rec, srv: srv, ag: ag, ln: ln, streams: &streams}
|
||||
t.Cleanup(func() {
|
||||
cancel()
|
||||
ag.Supervisor().StopAll(context.Background())
|
||||
@@ -351,11 +449,16 @@ func TestDispatchToUnconnectedNode(t *testing.T) {
|
||||
func TestAgentReconnects(t *testing.T) {
|
||||
// A control plane restart must not require touching every GPU node.
|
||||
h := start(t)
|
||||
h.srv.CloseClientConnections()
|
||||
|
||||
before := h.streams.Load()
|
||||
h.ln.drop()
|
||||
|
||||
// Counting streams, not connected nodes: a stale entry in the connection
|
||||
// map would make a node look connected without any reconnect having
|
||||
// happened, and an assertion on it would pass vacuously.
|
||||
deadline := time.Now().Add(15 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if len(h.svc.ConnectedNodes()) == 1 {
|
||||
if h.streams.Load() > before && 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)
|
||||
}
|
||||
@@ -363,7 +466,7 @@ func TestAgentReconnects(t *testing.T) {
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
t.Fatal("the agent never reconnected")
|
||||
t.Fatalf("the agent never reconnected (streams %d -> %d)", before, h.streams.Load())
|
||||
}
|
||||
|
||||
func TestIdempotentStart(t *testing.T) {
|
||||
@@ -401,3 +504,97 @@ func TestIdempotentStart(t *testing.T) {
|
||||
t.Errorf("a second process was started: pid %d then %d", pid, sh2.Status().PID)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRunningShardsReportedOnReconnect covers the control plane restart hole
|
||||
// (plan §6.5): a reconnecting agent has to describe what it is already running,
|
||||
// or the control plane concludes the cluster is empty and places duplicates
|
||||
// while the originals keep holding GPUs.
|
||||
func TestRunningShardsReportedOnReconnect(t *testing.T) {
|
||||
h := start(t, func(c *Config) {
|
||||
c.Engines.VLLM.Bin = fakeEngineScript(t, `sleep 60`)
|
||||
})
|
||||
|
||||
cmd := shardCommand(t, "shard-live")
|
||||
cmd.Pool = "ib-a/h100-80g"
|
||||
cmd.Fabric = "ib-a"
|
||||
cmd.Spec.HCAs = []string{"mlx5_0"}
|
||||
|
||||
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-live" && s.state == supervisor.Starting
|
||||
})
|
||||
|
||||
// Drop the stream, as a control plane restart would.
|
||||
h.ln.drop()
|
||||
|
||||
deadline := time.Now().Add(15 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if got := h.rec.adoptedShards(); len(got) > 0 {
|
||||
a := got[0]
|
||||
if a.ShardID != "shard-live" || a.ReplicaID != "rep-1" {
|
||||
t.Fatalf("wrong shard reported: %+v", a)
|
||||
}
|
||||
// Everything the control plane needs to rebuild the replica.
|
||||
if a.Model != "test-model" || a.Profile != "p" {
|
||||
t.Errorf("catalog identity lost: model=%q profile=%q", a.Model, a.Profile)
|
||||
}
|
||||
if a.Pool != "ib-a/h100-80g" || a.Fabric != "ib-a" {
|
||||
t.Errorf("pool/fabric lost: %q %q", a.Pool, a.Fabric)
|
||||
}
|
||||
if !slices.Equal(a.HCAs, []string{"mlx5_0"}) {
|
||||
t.Errorf("HCAs lost: %v", a.HCAs)
|
||||
}
|
||||
if a.Endpoint == "" {
|
||||
t.Error("rank 0 reported no endpoint; the router would have nowhere to send")
|
||||
}
|
||||
if len(a.GPUIndices) == 0 {
|
||||
t.Error("GPU indices lost; the inventory could not reclaim them")
|
||||
}
|
||||
return
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
t.Fatal("the reconnecting agent never reported its running shard")
|
||||
}
|
||||
|
||||
// TestUnclaimedShardIsStopped covers the other half: a shard whose model left
|
||||
// the catalog is holding GPUs nothing will ever route to.
|
||||
func TestUnclaimedShardIsStopped(t *testing.T) {
|
||||
h := start(t, func(c *Config) {
|
||||
c.Engines.VLLM.Bin = fakeEngineScript(t, `sleep 60`)
|
||||
})
|
||||
|
||||
cmd := shardCommand(t, "shard-orphaned")
|
||||
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-orphaned" && s.state == supervisor.Starting
|
||||
})
|
||||
|
||||
sh, err := h.ag.Supervisor().Get("shard-orphaned")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
pid := sh.Status().PID
|
||||
|
||||
// The control plane will refuse to claim it on reconnect.
|
||||
h.rec.markUnclaimed("shard-orphaned")
|
||||
h.ln.drop()
|
||||
|
||||
deadline := time.Now().Add(20 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if !processAliveForTest(pid) {
|
||||
return
|
||||
}
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
}
|
||||
t.Error("an unclaimed shard kept running and holding its GPUs")
|
||||
}
|
||||
|
||||
// processAliveForTest mirrors the supervisor's liveness check.
|
||||
func processAliveForTest(pid int) bool {
|
||||
return syscall.Kill(pid, 0) == nil
|
||||
}
|
||||
@@ -66,6 +66,12 @@ type ShardSpec struct {
|
||||
|
||||
HTTPPort int
|
||||
Profile Profile
|
||||
|
||||
// Pool and Fabric are control plane concepts the agent cannot derive. They
|
||||
// are carried so the agent can report them back after a restart, and are
|
||||
// not used when building a command.
|
||||
Pool string
|
||||
Fabric string
|
||||
}
|
||||
|
||||
// Leader reports whether this shard serves HTTP for the replica.
|
||||
|
||||
@@ -3,6 +3,7 @@ package supervisor
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
@@ -10,7 +11,6 @@ import (
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"syscall"
|
||||
@@ -82,6 +82,9 @@ type Shard struct {
|
||||
cmd *exec.Cmd
|
||||
cancel context.CancelFunc
|
||||
done chan struct{}
|
||||
// adopted marks a shard inherited from a previous agent: it has no
|
||||
// exec.Cmd, so it is watched by polling and killed by signal alone.
|
||||
adopted bool
|
||||
}
|
||||
|
||||
// Status is a snapshot of a shard, safe to send to the control plane.
|
||||
@@ -217,7 +220,7 @@ func (s *Supervisor) launch(ctx context.Context, sh *Shard, cmd engines.Command)
|
||||
sh.state = Starting
|
||||
sh.mu.Unlock()
|
||||
|
||||
if err := s.writePIDFile(sh); err != nil {
|
||||
if err := s.writeState(sh); err != nil {
|
||||
s.log.Warn("could not write pid file; adoption after a restart will not work",
|
||||
"shard", sh.ID, "err", err)
|
||||
}
|
||||
@@ -268,7 +271,7 @@ func (s *Supervisor) wait(sh *Shard) {
|
||||
final, code := sh.state, sh.exitCode
|
||||
sh.mu.Unlock()
|
||||
|
||||
s.removePIDFile(sh)
|
||||
s.removeState(sh)
|
||||
if final == Failed {
|
||||
s.log.Error("shard failed", "shard", sh.ID, "exit_code", code, "err", err,
|
||||
"last_output", strings.Join(sh.Logs(5), " | "))
|
||||
@@ -330,13 +333,15 @@ func (s *Supervisor) Stop(ctx context.Context, id string) error {
|
||||
}
|
||||
|
||||
sh.mu.Lock()
|
||||
c, pid := sh.cmd, sh.pid
|
||||
pid := sh.pid
|
||||
if sh.state.Live() {
|
||||
sh.state = Draining // marks the exit as expected for wait()
|
||||
sh.state = Draining // marks the exit as expected for wait() and poll()
|
||||
}
|
||||
sh.mu.Unlock()
|
||||
|
||||
if c == nil || pid == 0 {
|
||||
// A pid of zero means the process never started; an adopted shard has a pid
|
||||
// but no exec.Cmd, and must still be killed by signal.
|
||||
if pid == 0 {
|
||||
sh.setState(Stopped, nil)
|
||||
return nil
|
||||
}
|
||||
@@ -347,18 +352,26 @@ func (s *Supervisor) Stop(ctx context.Context, id string) error {
|
||||
s.log.Warn("SIGTERM failed", "shard", id, "err", err)
|
||||
}
|
||||
|
||||
// force kills the group outright. exec.Cmd's cancel does this for shards we
|
||||
// started; an adopted shard has none, so signal its group directly.
|
||||
force := func() {
|
||||
if sh.cancel != nil {
|
||||
sh.cancel()
|
||||
return
|
||||
}
|
||||
if err := killGroup(pid, syscall.SIGKILL); err != nil && !errors.Is(err, os.ErrProcessDone) {
|
||||
s.log.Warn("SIGKILL failed", "shard", id, "err", err)
|
||||
}
|
||||
}
|
||||
|
||||
select {
|
||||
case <-sh.done:
|
||||
return nil
|
||||
case <-time.After(s.cfg.StopTimeout):
|
||||
s.log.Warn("shard did not exit in time, killing", "shard", id, "timeout", s.cfg.StopTimeout)
|
||||
if sh.cancel != nil {
|
||||
sh.cancel()
|
||||
}
|
||||
force()
|
||||
case <-ctx.Done():
|
||||
if sh.cancel != nil {
|
||||
sh.cancel()
|
||||
}
|
||||
force()
|
||||
return ctx.Err()
|
||||
}
|
||||
|
||||
@@ -424,39 +437,68 @@ func (s *Supervisor) Statuses() []Status {
|
||||
return out
|
||||
}
|
||||
|
||||
// --- pid files, for adoption after an agent restart ---
|
||||
// --- state files, for adoption after an agent restart ---
|
||||
|
||||
func (s *Supervisor) pidPath(id string) string {
|
||||
return filepath.Join(s.cfg.WorkDir, "shard-"+id+".pid")
|
||||
func (s *Supervisor) statePath(id string) string {
|
||||
return filepath.Join(s.cfg.WorkDir, "shard-"+id+".json")
|
||||
}
|
||||
|
||||
// writePIDFile records enough to recognise the process later.
|
||||
// shardState is what survives an agent restart.
|
||||
//
|
||||
// The pid alone is not enough: pids are reused, and adopting an unrelated
|
||||
// process would be worse than leaking the original. The command line is stored
|
||||
// alongside it as a fingerprint.
|
||||
func (s *Supervisor) writePIDFile(sh *Shard) error {
|
||||
content := fmt.Sprintf("%d\n%s\n", sh.pid, strings.Join(sh.cmd.Args, "\x00"))
|
||||
return os.WriteFile(s.pidPath(sh.ID), []byte(content), 0o600)
|
||||
// The pid alone is not enough for two separate reasons. Pids are reused, so the
|
||||
// command line is kept as a fingerprint — adopting an unrelated process would
|
||||
// be worse than leaking the original. And the full spec is kept because a
|
||||
// recovered shard has to be described back to the control plane: without it the
|
||||
// control plane cannot tell which replica the process belongs to, and would
|
||||
// place a duplicate while the original keeps holding the GPUs (plan §6.5).
|
||||
type shardState struct {
|
||||
ShardID string `json:"shard_id"`
|
||||
ReplicaID string `json:"replica_id"`
|
||||
PID int `json:"pid"`
|
||||
Cmdline []string `json:"cmdline"`
|
||||
StartedAt time.Time `json:"started_at"`
|
||||
Spec engines.ShardSpec `json:"spec"`
|
||||
}
|
||||
|
||||
func (s *Supervisor) removePIDFile(sh *Shard) {
|
||||
_ = os.Remove(s.pidPath(sh.ID))
|
||||
// writeState records everything needed to recognise and describe the process later.
|
||||
func (s *Supervisor) writeState(sh *Shard) error {
|
||||
b, err := json.Marshal(shardState{
|
||||
ShardID: sh.ID,
|
||||
ReplicaID: sh.Spec.ReplicaID,
|
||||
PID: sh.pid,
|
||||
Cmdline: sh.cmd.Args,
|
||||
StartedAt: sh.started,
|
||||
Spec: sh.Spec,
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// Written atomically: a half-written state file after a crash would be
|
||||
// indistinguishable from a corrupt one, and the shard would be killed.
|
||||
tmp := s.statePath(sh.ID) + ".tmp"
|
||||
if err := os.WriteFile(tmp, b, 0o600); err != nil {
|
||||
return err
|
||||
}
|
||||
return os.Rename(tmp, s.statePath(sh.ID))
|
||||
}
|
||||
|
||||
func (s *Supervisor) removeState(sh *Shard) {
|
||||
_ = os.Remove(s.statePath(sh.ID))
|
||||
}
|
||||
|
||||
// Orphan is an engine process left behind by a previous agent.
|
||||
type Orphan struct {
|
||||
ShardID string
|
||||
PID int
|
||||
Cmdline []string
|
||||
// Alive reports whether the process still exists and matches its fingerprint.
|
||||
ShardID string
|
||||
ReplicaID string
|
||||
PID int
|
||||
Cmdline []string
|
||||
Spec engines.ShardSpec
|
||||
StartedAt time.Time
|
||||
// Alive reports whether the process still exists.
|
||||
Alive bool
|
||||
}
|
||||
|
||||
// FindOrphans scans the work directory for processes a previous agent started.
|
||||
//
|
||||
// The control plane decides what happens next (plan §5.4): a shard of a
|
||||
// multi-node replica cannot be adopted in isolation, because its peers are gone.
|
||||
func (s *Supervisor) FindOrphans() ([]Orphan, error) {
|
||||
entries, err := os.ReadDir(s.cfg.WorkDir)
|
||||
if err != nil {
|
||||
@@ -466,35 +508,101 @@ func (s *Supervisor) FindOrphans() ([]Orphan, error) {
|
||||
var out []Orphan
|
||||
for _, e := range entries {
|
||||
name := e.Name()
|
||||
if !strings.HasPrefix(name, "shard-") || !strings.HasSuffix(name, ".pid") {
|
||||
if !strings.HasPrefix(name, "shard-") || !strings.HasSuffix(name, ".json") {
|
||||
continue
|
||||
}
|
||||
id := strings.TrimSuffix(strings.TrimPrefix(name, "shard-"), ".pid")
|
||||
|
||||
b, err := os.ReadFile(filepath.Join(s.cfg.WorkDir, name))
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
lines := strings.SplitN(strings.TrimSpace(string(b)), "\n", 2)
|
||||
pid, err := strconv.Atoi(strings.TrimSpace(lines[0]))
|
||||
if err != nil {
|
||||
var st shardState
|
||||
if err := json.Unmarshal(b, &st); err != nil || st.ShardID == "" || st.PID <= 0 {
|
||||
// A corrupt state file describes a process that cannot be
|
||||
// identified, so it cannot be adopted safely. Leave it for
|
||||
// KillOrphan to clean up by shard id from the filename.
|
||||
id := strings.TrimSuffix(strings.TrimPrefix(name, "shard-"), ".json")
|
||||
out = append(out, Orphan{ShardID: id})
|
||||
continue
|
||||
}
|
||||
var cmdline []string
|
||||
if len(lines) > 1 {
|
||||
cmdline = strings.Split(lines[1], "\x00")
|
||||
}
|
||||
|
||||
out = append(out, Orphan{
|
||||
ShardID: id,
|
||||
PID: pid,
|
||||
Cmdline: cmdline,
|
||||
Alive: processAlive(pid),
|
||||
ShardID: st.ShardID,
|
||||
ReplicaID: st.ReplicaID,
|
||||
PID: st.PID,
|
||||
Cmdline: st.Cmdline,
|
||||
Spec: st.Spec,
|
||||
StartedAt: st.StartedAt,
|
||||
Alive: processAlive(st.PID),
|
||||
})
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// Adopt takes over a still-running engine from a previous agent.
|
||||
//
|
||||
// The process cannot be waited on — it is not this process's child — so it is
|
||||
// watched by polling instead. That is the whole difference between an adopted
|
||||
// shard and a started one; everything downstream treats them the same.
|
||||
func (s *Supervisor) Adopt(o Orphan) (*Shard, error) {
|
||||
if !o.Alive {
|
||||
return nil, fmt.Errorf("shard %s: process %d is gone", o.ShardID, o.PID)
|
||||
}
|
||||
|
||||
s.mu.Lock()
|
||||
if _, exists := s.shards[o.ShardID]; exists {
|
||||
s.mu.Unlock()
|
||||
return nil, ErrAlreadyRunning
|
||||
}
|
||||
sh := &Shard{
|
||||
ID: o.ShardID,
|
||||
Spec: o.Spec,
|
||||
state: Starting,
|
||||
pid: o.PID,
|
||||
logs: newRingBuffer(s.cfg.LogLines),
|
||||
started: o.StartedAt,
|
||||
done: make(chan struct{}),
|
||||
adopted: true,
|
||||
}
|
||||
s.shards[o.ShardID] = sh
|
||||
s.mu.Unlock()
|
||||
|
||||
// Output is gone with the previous agent; say so rather than presenting an
|
||||
// empty log as if the engine had been silent.
|
||||
sh.logs.add("[adopted from a previous agent; earlier output is not available]")
|
||||
|
||||
go s.poll(sh)
|
||||
return sh, nil
|
||||
}
|
||||
|
||||
// poll watches an adopted process, which cannot be reaped with Wait.
|
||||
func (s *Supervisor) poll(sh *Shard) {
|
||||
defer close(sh.done)
|
||||
|
||||
// Polled rather than waited on, so the interval bounds how quickly a stop
|
||||
// is observed. Short enough that Stop does not sit out its whole timeout.
|
||||
t := time.NewTicker(200 * time.Millisecond)
|
||||
defer t.Stop()
|
||||
for range t.C {
|
||||
if processAlive(sh.pid) {
|
||||
continue
|
||||
}
|
||||
sh.mu.Lock()
|
||||
expected := sh.state == Draining || sh.state == Stopped
|
||||
if expected {
|
||||
sh.state = Stopped
|
||||
} else {
|
||||
sh.state = Failed
|
||||
sh.err = errors.New("adopted engine exited")
|
||||
}
|
||||
sh.mu.Unlock()
|
||||
|
||||
s.removeState(sh)
|
||||
s.log.Info("adopted shard exited", "shard", sh.ID, "expected", expected)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// KillOrphan terminates a leftover process group and removes its pid file.
|
||||
func (s *Supervisor) KillOrphan(o Orphan) error {
|
||||
if o.Alive {
|
||||
@@ -513,9 +621,9 @@ func (s *Supervisor) KillOrphan(o Orphan) error {
|
||||
_ = killGroup(o.PID, syscall.SIGKILL)
|
||||
}
|
||||
}
|
||||
// An already-missing pid file is success: the previous agent may have
|
||||
// An already-missing state file is success: the previous agent may have
|
||||
// reaped the process itself between the scan and this call.
|
||||
if err := os.Remove(s.pidPath(o.ShardID)); err != nil && !errors.Is(err, os.ErrNotExist) {
|
||||
if err := os.Remove(s.statePath(o.ShardID)); err != nil && !errors.Is(err, os.ErrNotExist) {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
|
||||
@@ -9,6 +9,7 @@ import (
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"syscall"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@@ -307,7 +308,8 @@ func TestLogsAreCaptured(t *testing.T) {
|
||||
func TestFindOrphans(t *testing.T) {
|
||||
s, _ := newSupervisor(t)
|
||||
|
||||
sh, err := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand("sleep 30"))
|
||||
spec := engines.ShardSpec{ReplicaID: "rep-1", ShardID: "a", HTTPPort: 20100}
|
||||
sh, err := s.Start(t.Context(), "a", spec, shellCommand("sleep 30"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
@@ -336,6 +338,14 @@ func TestFindOrphans(t *testing.T) {
|
||||
if len(o.Cmdline) == 0 || !strings.Contains(strings.Join(o.Cmdline, " "), "sleep 30") {
|
||||
t.Errorf("cmdline fingerprint missing: %v", o.Cmdline)
|
||||
}
|
||||
// And the spec, without which the control plane could not tell which
|
||||
// replica the process belongs to and would place a duplicate.
|
||||
if o.ReplicaID != "rep-1" {
|
||||
t.Errorf("replica id not recovered: %q", o.ReplicaID)
|
||||
}
|
||||
if o.Spec.HTTPPort != 20100 {
|
||||
t.Errorf("spec not recovered: %+v", o.Spec)
|
||||
}
|
||||
|
||||
if err := s2.KillOrphan(o); err != nil {
|
||||
t.Fatal(err)
|
||||
@@ -350,20 +360,20 @@ func TestFindOrphans(t *testing.T) {
|
||||
if processAlive(pid) {
|
||||
t.Error("the orphan survived KillOrphan")
|
||||
}
|
||||
if _, err := os.Stat(s2.pidPath("a")); !os.IsNotExist(err) {
|
||||
t.Error("the pid file was left behind")
|
||||
if _, err := os.Stat(s2.statePath("a")); !os.IsNotExist(err) {
|
||||
t.Error("the state file was left behind")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPIDFileRemovedOnExit(t *testing.T) {
|
||||
func TestStateFileRemovedOnExit(t *testing.T) {
|
||||
s, _ := newSupervisor(t)
|
||||
|
||||
sh, _ := s.Start(t.Context(), "a", engines.ShardSpec{}, shellCommand("exit 0"))
|
||||
waitState(t, sh, Failed, 3*time.Second)
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
|
||||
if _, err := os.Stat(s.pidPath("a")); !os.IsNotExist(err) {
|
||||
t.Error("the pid file outlived the process; a later scan would find a dead orphan")
|
||||
if _, err := os.Stat(s.statePath("a")); !os.IsNotExist(err) {
|
||||
t.Error("the state file outlived the process; a later scan would find a dead orphan")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -407,3 +417,118 @@ func TestRingBuffer(t *testing.T) {
|
||||
t.Errorf("tail(2)=%q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdoptTakesOverRunningShard(t *testing.T) {
|
||||
// The control plane restart fix depends on this: a shard that survived the
|
||||
// agent must come back under supervision, not be restarted or leaked.
|
||||
s, _ := newSupervisor(t)
|
||||
|
||||
spec := engines.ShardSpec{ReplicaID: "rep-1", ShardID: "a", HTTPPort: 20200}
|
||||
sh, err := s.Start(t.Context(), "a", spec, shellCommand("sleep 30"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
pid := sh.Status().PID
|
||||
|
||||
// A new supervisor over the same work dir, as after an agent restart.
|
||||
s2, err := New(Config{WorkDir: s.cfg.WorkDir, StopTimeout: 3 * time.Second},
|
||||
&scriptedHealth{}, slog.New(slog.DiscardHandler))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
orphans, err := s2.FindOrphans()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(orphans) != 1 || !orphans[0].Alive {
|
||||
t.Fatalf("orphan not found: %+v", orphans)
|
||||
}
|
||||
|
||||
adopted, err := s2.Adopt(orphans[0])
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := adopted.Status().PID; got != pid {
|
||||
t.Errorf("adopted pid %d, expected %d", got, pid)
|
||||
}
|
||||
if adopted.Spec.ReplicaID != "rep-1" {
|
||||
t.Errorf("spec lost in adoption: %+v", adopted.Spec)
|
||||
}
|
||||
// The earlier output is gone; saying so beats presenting an empty log as if
|
||||
// the engine had been silent.
|
||||
if logs := strings.Join(adopted.Logs(5), " "); !strings.Contains(logs, "adopted") {
|
||||
t.Errorf("adoption was not noted in the log: %q", logs)
|
||||
}
|
||||
|
||||
// Stopping an adopted shard must actually kill it: it has no exec.Cmd, so
|
||||
// the signal path is the only one available.
|
||||
if err := s2.Stop(t.Context(), "a"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
deadline := time.Now().Add(5 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
if !processAlive(pid) {
|
||||
return
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
t.Error("the adopted process survived Stop")
|
||||
}
|
||||
|
||||
func TestAdoptRefusesDeadProcess(t *testing.T) {
|
||||
s, _ := newSupervisor(t)
|
||||
if _, err := s.Adopt(Orphan{ShardID: "gone", PID: 999999, Alive: false}); err == nil {
|
||||
t.Error("a dead process was adopted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdoptedExitIsReported(t *testing.T) {
|
||||
// An adopted engine dying has to be noticed; nothing waits on it.
|
||||
s, _ := newSupervisor(t)
|
||||
|
||||
spec := engines.ShardSpec{ReplicaID: "rep-1", ShardID: "a"}
|
||||
sh, err := s.Start(t.Context(), "a", spec, shellCommand("sleep 30"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
pid := sh.Status().PID
|
||||
|
||||
s2, err := New(Config{WorkDir: s.cfg.WorkDir, StopTimeout: 2 * time.Second},
|
||||
&scriptedHealth{}, slog.New(slog.DiscardHandler))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
orphans, _ := s2.FindOrphans()
|
||||
adopted, err := s2.Adopt(orphans[0])
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// Kill it behind the supervisor's back.
|
||||
_ = killGroup(pid, syscall.SIGKILL)
|
||||
waitState(t, adopted, Failed, 5*time.Second)
|
||||
}
|
||||
|
||||
func TestCorruptStateFileIsCleanedUp(t *testing.T) {
|
||||
// A half-written state file describes a process that cannot be identified,
|
||||
// so it cannot be adopted safely.
|
||||
s, _ := newSupervisor(t)
|
||||
if err := os.WriteFile(s.statePath("broken"), []byte("{not json"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
orphans, err := s.FindOrphans()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(orphans) != 1 || orphans[0].ShardID != "broken" || orphans[0].Alive {
|
||||
t.Fatalf("corrupt state file not reported for cleanup: %+v", orphans)
|
||||
}
|
||||
if err := s.KillOrphan(orphans[0]); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := os.Stat(s.statePath("broken")); !os.IsNotExist(err) {
|
||||
t.Error("the corrupt state file was left behind")
|
||||
}
|
||||
}
|
||||
@@ -2,6 +2,7 @@
|
||||
package agentsvc
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
||||
@@ -87,6 +88,33 @@ func RegisterFromInventory(nodeID, hostname, mgmtAddr, agentVersion string,
|
||||
return r
|
||||
}
|
||||
|
||||
// AdoptedFromRegister converts an agent's running-shard report into the
|
||||
// reconciler's view.
|
||||
func AdoptedFromRegister(nodeID string, shards []*llmgwv1.RunningShard) []reconciler.AdoptedShard {
|
||||
out := make([]reconciler.AdoptedShard, 0, len(shards))
|
||||
for _, s := range shards {
|
||||
a := reconciler.AdoptedShard{
|
||||
NodeID: nodeID,
|
||||
ReplicaID: s.GetReplicaId(),
|
||||
ShardID: s.GetShardId(),
|
||||
Rank: int(s.GetRank()),
|
||||
GPUIndices: intSlice(s.GetGpuIndices()),
|
||||
HCAs: s.GetHcas(),
|
||||
HTTPPort: int(s.GetHttpPort()),
|
||||
State: StateFromWire(s.GetState()),
|
||||
Model: s.GetModel(),
|
||||
Profile: s.GetProfile(),
|
||||
Pool: s.GetPool(),
|
||||
Fabric: s.GetFabric(),
|
||||
}
|
||||
if a.Rank == 0 && s.GetLeaderAddr() != "" {
|
||||
a.Endpoint = fmt.Sprintf("http://%s:%d", s.GetLeaderAddr(), s.GetHttpPort())
|
||||
}
|
||||
out = append(out, a)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// StartShardFromCommand converts a reconciler command into a wire message.
|
||||
func StartShardFromCommand(cmd reconciler.ShardCommand) *llmgwv1.StartShard {
|
||||
p := cmd.Spec.Profile
|
||||
@@ -109,6 +137,8 @@ func StartShardFromCommand(cmd reconciler.ShardCommand) *llmgwv1.StartShard {
|
||||
ModelPath: p.ModelPath, ServedModelName: p.ServedModelName,
|
||||
},
|
||||
StartupTimeoutMs: cmd.StartupTimeout.Milliseconds(),
|
||||
Pool: cmd.Pool,
|
||||
Fabric: cmd.Fabric,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -126,6 +156,8 @@ func ShardSpecFromStart(m *llmgwv1.StartShard) (engines.ShardSpec, time.Duration
|
||||
HCAs: m.GetHcas(),
|
||||
MgmtInterface: m.GetMgmtInterface(),
|
||||
HTTPPort: int(m.GetHttpPort()),
|
||||
Pool: m.GetPool(),
|
||||
Fabric: m.GetFabric(),
|
||||
Profile: engines.Profile{
|
||||
Name: p.GetName(), Engine: p.GetEngine(), GPUClass: p.GetGpuClass(),
|
||||
TP: int(p.GetTp()), PP: int(p.GetPp()),
|
||||
|
||||
@@ -19,9 +19,11 @@ import (
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
|
||||
)
|
||||
|
||||
// StatusSink receives shard transitions. The reconciler implements it.
|
||||
// StatusSink receives shard transitions and recovers replica state from what
|
||||
// agents report they are already running. The reconciler implements it.
|
||||
type StatusSink interface {
|
||||
OnShardStatus(replicaID, shardID string, state supervisor.State, errMsg string)
|
||||
Adopt(shards []reconciler.AdoptedShard) (unclaimed []reconciler.AdoptedShard)
|
||||
}
|
||||
|
||||
// Service is the control plane's AgentService implementation, and at the same
|
||||
@@ -87,6 +89,7 @@ func (s *Service) Connect(
|
||||
s.inv.Upsert(node)
|
||||
s.log.Info("agent registered",
|
||||
"node", nodeID, "gpus", len(node.GPUs), "fabrics", node.FabricNames(),
|
||||
"running_shards", len(reg.GetRunningShards()),
|
||||
"agent_version", reg.GetAgentVersion())
|
||||
|
||||
c := s.attach(ctx, nodeID)
|
||||
@@ -98,6 +101,11 @@ func (s *Service) Connect(
|
||||
return err
|
||||
}
|
||||
|
||||
// Recover replica state from what this node is already running, before any
|
||||
// reconciliation pass can conclude the cluster is empty and place
|
||||
// duplicates (plan §6.5).
|
||||
s.adopt(ctx, nodeID, reg.GetRunningShards(), c)
|
||||
|
||||
// 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
|
||||
@@ -123,6 +131,28 @@ func (s *Service) Connect(
|
||||
}
|
||||
}
|
||||
|
||||
// adopt hands the agent's running shards to the reconciler and stops whatever
|
||||
// nothing in the catalog claims.
|
||||
func (s *Service) adopt(ctx context.Context, nodeID string, shards []*llmgwv1.RunningShard, c *conn) {
|
||||
if len(shards) == 0 {
|
||||
return
|
||||
}
|
||||
unclaimed := s.status.Adopt(AdoptedFromRegister(nodeID, shards))
|
||||
for _, u := range unclaimed {
|
||||
// The model or profile left the catalog while this was running. Nobody
|
||||
// will ever route to it, and it is holding GPUs.
|
||||
s.log.Warn("stopping an unclaimed shard",
|
||||
"node", nodeID, "shard", u.ShardID, "model", u.Model, "profile", u.Profile)
|
||||
select {
|
||||
case c.send <- &llmgwv1.ControlMessage{
|
||||
Msg: &llmgwv1.ControlMessage_Stop{Stop: &llmgwv1.StopShard{ShardId: u.ShardID}},
|
||||
}:
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Service) writeLoop(ctx context.Context,
|
||||
stream *connect.BidiStream[llmgwv1.AgentMessage, llmgwv1.ControlMessage], c *conn,
|
||||
) {
|
||||
|
||||
@@ -24,6 +24,11 @@ type ShardCommand struct {
|
||||
Spec engines.ShardSpec
|
||||
// StartupTimeout comes from the profile; model loading dominates it.
|
||||
StartupTimeout time.Duration
|
||||
// Pool and Fabric are control plane concepts the agent cannot derive. They
|
||||
// travel with the command so the agent can report them back after a
|
||||
// control plane restart (plan §6.5).
|
||||
Pool string
|
||||
Fabric string
|
||||
}
|
||||
|
||||
// Dispatcher sends commands to agents.
|
||||
@@ -436,6 +441,8 @@ func (r *Reconciler) startReplica(ctx context.Context, m catalog.Model, pl *plac
|
||||
Profile: espec,
|
||||
},
|
||||
StartupTimeout: prof.StartupTimeout,
|
||||
Pool: pl.Pool,
|
||||
Fabric: pl.Fabric,
|
||||
}
|
||||
if err := r.disp.StartShard(ctx, cmd); err != nil {
|
||||
r.markShard(id, s.ID, supervisor.Failed, err.Error())
|
||||
@@ -518,6 +525,95 @@ func (r *Reconciler) stopReplica(ctx context.Context, id string) {
|
||||
r.mu.Unlock()
|
||||
}
|
||||
|
||||
// Adopt rebuilds replica state from shards an agent reports it is already
|
||||
// running (plan §6.5).
|
||||
//
|
||||
// This is what makes a control plane restart survivable. Without it the
|
||||
// reconciler sees an empty cluster, places duplicates, and the engines already
|
||||
// running keep holding GPUs the inventory believes are free — so the duplicates
|
||||
// fail to allocate, with an error naming the new process rather than the old.
|
||||
//
|
||||
// Shards whose model or profile is no longer in the catalog are returned for
|
||||
// the caller to stop: nothing will ever claim them.
|
||||
func (r *Reconciler) Adopt(shards []AdoptedShard) (unclaimed []AdoptedShard) {
|
||||
byReplica := map[string][]AdoptedShard{}
|
||||
for _, s := range shards {
|
||||
byReplica[s.ReplicaID] = append(byReplica[s.ReplicaID], s)
|
||||
}
|
||||
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
|
||||
for replicaID, group := range byReplica {
|
||||
if _, known := r.replicas[replicaID]; known {
|
||||
// Already tracked — an agent reconnecting rather than a control
|
||||
// plane restart. Its status reports keep the state current.
|
||||
continue
|
||||
}
|
||||
|
||||
first := group[0]
|
||||
m, ok := r.cat.Resolve(first.Model)
|
||||
if !ok {
|
||||
unclaimed = append(unclaimed, group...)
|
||||
continue
|
||||
}
|
||||
prof, ok := m.ProfileByName(first.Profile)
|
||||
if !ok {
|
||||
unclaimed = append(unclaimed, group...)
|
||||
continue
|
||||
}
|
||||
|
||||
rep := &Replica{
|
||||
ID: replicaID, Model: m.Name, Profile: prof.Name,
|
||||
Fabric: first.Fabric, Pool: first.Pool,
|
||||
Weight: prof.Weight, CreatedAt: r.now(),
|
||||
Endpoint: first.Endpoint,
|
||||
}
|
||||
for _, s := range group {
|
||||
rep.Shards = append(rep.Shards, Shard{
|
||||
ID: s.ShardID, NodeID: s.NodeID, Rank: s.Rank,
|
||||
GPUIndices: s.GPUIndices, HCAs: s.HCAs,
|
||||
HTTPPort: s.HTTPPort, State: s.State,
|
||||
})
|
||||
// Reclaim the GPUs, or the next placement will hand them out twice.
|
||||
r.inv.Allocate(s.NodeID, inventory.Allocation{
|
||||
ShardID: s.ShardID, ReplicaID: replicaID, GPUIndices: s.GPUIndices,
|
||||
})
|
||||
}
|
||||
// The endpoint comes from the leader; a group missing rank 0 is a
|
||||
// partial replica and will be failed by deriveState below.
|
||||
for _, s := range group {
|
||||
if s.Rank == 0 && s.Endpoint != "" {
|
||||
rep.Endpoint = s.Endpoint
|
||||
}
|
||||
}
|
||||
rep.deriveState()
|
||||
r.replicas[replicaID] = rep
|
||||
|
||||
r.log.Info("adopted replica from the fleet",
|
||||
"replica", replicaID, "model", m.Name, "shards", len(group), "state", rep.State)
|
||||
}
|
||||
return unclaimed
|
||||
}
|
||||
|
||||
// AdoptedShard is one running shard an agent reported at registration.
|
||||
type AdoptedShard struct {
|
||||
NodeID string
|
||||
ReplicaID string
|
||||
ShardID string
|
||||
Rank int
|
||||
GPUIndices []int
|
||||
HCAs []string
|
||||
HTTPPort int
|
||||
State supervisor.State
|
||||
Model string
|
||||
Profile string
|
||||
Pool string
|
||||
Fabric string
|
||||
// Endpoint is the leader's base URL, set on the rank 0 shard.
|
||||
Endpoint string
|
||||
}
|
||||
|
||||
// OnShardStatus records a status report from an agent.
|
||||
func (r *Reconciler) OnShardStatus(replicaID, shardID string, st supervisor.State, errMsg string) {
|
||||
r.markShard(replicaID, shardID, st, errMsg)
|
||||
|
||||
@@ -596,3 +596,171 @@ func TestReplicaStateStrings(t *testing.T) {
|
||||
t.Error("Counts is wrong")
|
||||
}
|
||||
}
|
||||
|
||||
// --- adoption after a control plane restart (plan §6.5) ---
|
||||
|
||||
func adopted(replicaID, shardID, nodeID string, rank int, gpus []int) AdoptedShard {
|
||||
a := AdoptedShard{
|
||||
ReplicaID: replicaID, ShardID: shardID, NodeID: nodeID, Rank: rank,
|
||||
GPUIndices: gpus, HTTPPort: 20000, State: supervisor.Ready,
|
||||
Model: "llama", Profile: "tp4", Pool: "none/h100-80g",
|
||||
}
|
||||
if rank == 0 {
|
||||
a.Endpoint = "http://10.0.0." + nodeID + ":20000"
|
||||
}
|
||||
return a
|
||||
}
|
||||
|
||||
func TestAdoptRebuildsReplica(t *testing.T) {
|
||||
// The control plane restarted. Nothing is known, but a node reports a live
|
||||
// replica. Without this the reconciler would place a duplicate while the
|
||||
// original keeps holding GPUs.
|
||||
h := newHarness(t, testModel(1), gpuNode("1", 8))
|
||||
|
||||
unclaimed := h.r.Adopt([]AdoptedShard{adopted("llama-99", "llama-99-r0", "1", 0, []int{0, 1, 2, 3})})
|
||||
if len(unclaimed) != 0 {
|
||||
t.Fatalf("a valid shard was reported unclaimed: %+v", unclaimed)
|
||||
}
|
||||
|
||||
reps := h.r.Replicas()
|
||||
if len(reps) != 1 {
|
||||
t.Fatalf("expected 1 adopted replica, got %d", len(reps))
|
||||
}
|
||||
rep := reps[0]
|
||||
if rep.ID != "llama-99" || rep.Model != "llama" || rep.Profile != "tp4" {
|
||||
t.Errorf("identity wrong: %+v", rep)
|
||||
}
|
||||
if rep.State != Ready {
|
||||
t.Errorf("state=%s; every shard reported Ready", rep.State)
|
||||
}
|
||||
if rep.Endpoint == "" {
|
||||
t.Error("no endpoint; the router would have nowhere to send")
|
||||
}
|
||||
if rep.Weight <= 0 {
|
||||
t.Error("weight not taken from the catalog")
|
||||
}
|
||||
|
||||
// The GPUs must be reclaimed, or the next placement hands them out twice.
|
||||
if free := len(h.inv.FreeGPUs("1")); free != 4 {
|
||||
t.Errorf("node has %d free GPUs, expected 4 after adoption", free)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdoptPreventsDuplicatePlacement(t *testing.T) {
|
||||
// The whole point: after adopting, a reconciliation pass must see the
|
||||
// desired count already met.
|
||||
h := newHarness(t, testModel(1), gpuNode("1", 8))
|
||||
|
||||
h.r.Adopt([]AdoptedShard{adopted("llama-99", "llama-99-r0", "1", 0, []int{0, 1, 2, 3})})
|
||||
h.disp.reset()
|
||||
h.r.Reconcile(t.Context())
|
||||
|
||||
if got := len(h.disp.startedShards()); got != 0 {
|
||||
t.Errorf("a duplicate replica was placed despite adoption (%d shards started)", got)
|
||||
}
|
||||
if got := len(h.r.Replicas()); got != 1 {
|
||||
t.Errorf("expected 1 replica, got %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdoptMultiNodeReplica(t *testing.T) {
|
||||
h := newHarness(t, testModel(1, catalog.Profile{
|
||||
Name: "2node", GPUClass: "h100-80g", TP: 16, Nodes: 2, GPUsPerNode: 8,
|
||||
}), gpuNode("1", 8, "ib-a"), gpuNode("2", 8, "ib-a"))
|
||||
|
||||
shards := []AdoptedShard{
|
||||
adopted("llama-7", "llama-7-r0", "1", 0, []int{0, 1, 2, 3, 4, 5, 6, 7}),
|
||||
adopted("llama-7", "llama-7-r1", "2", 1, []int{0, 1, 2, 3, 4, 5, 6, 7}),
|
||||
}
|
||||
for i := range shards {
|
||||
shards[i].Profile = "2node"
|
||||
shards[i].Fabric = "ib-a"
|
||||
shards[i].Pool = "ib-a/h100-80g"
|
||||
}
|
||||
|
||||
if unclaimed := h.r.Adopt(shards); len(unclaimed) != 0 {
|
||||
t.Fatalf("unclaimed: %+v", unclaimed)
|
||||
}
|
||||
reps := h.r.Replicas()
|
||||
if len(reps) != 1 || len(reps[0].Shards) != 2 {
|
||||
t.Fatalf("the two shards were not grouped into one replica: %+v", reps)
|
||||
}
|
||||
if reps[0].Fabric != "ib-a" {
|
||||
t.Errorf("fabric lost: %q", reps[0].Fabric)
|
||||
}
|
||||
// Both nodes' GPUs must be reclaimed.
|
||||
for _, id := range []string{"1", "2"} {
|
||||
if free := len(h.inv.FreeGPUs(id)); free != 0 {
|
||||
t.Errorf("node %s has %d free GPUs after adoption", id, free)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdoptPartialReplicaFails(t *testing.T) {
|
||||
// A multi-node replica missing a shard is useless: the survivors wait on
|
||||
// collectives that will never complete while holding GPUs.
|
||||
h := newHarness(t, testModel(1, catalog.Profile{
|
||||
Name: "2node", GPUClass: "h100-80g", TP: 16, Nodes: 2, GPUsPerNode: 8,
|
||||
}), gpuNode("1", 8, "ib-a"), gpuNode("2", 8, "ib-a"))
|
||||
|
||||
s := adopted("llama-7", "llama-7-r0", "1", 0, []int{0, 1, 2, 3, 4, 5, 6, 7})
|
||||
s.Profile = "2node"
|
||||
s.State = supervisor.Failed
|
||||
h.r.Adopt([]AdoptedShard{s})
|
||||
|
||||
reps := h.r.Replicas()
|
||||
if len(reps) != 1 || reps[0].State != Failed {
|
||||
t.Fatalf("a failed adopted shard did not fail its replica: %+v", reps)
|
||||
}
|
||||
|
||||
// And the next pass releases its GPUs.
|
||||
h.r.Reconcile(t.Context())
|
||||
if free := len(h.inv.FreeGPUs("1")); free != 8 {
|
||||
t.Errorf("node has %d free GPUs; the failed adopted replica still holds some", free)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdoptUnknownModelIsUnclaimed(t *testing.T) {
|
||||
// The model left the catalog while the engine kept running. Nothing will
|
||||
// ever route to it, and it is holding GPUs.
|
||||
h := newHarness(t, testModel(1), gpuNode("1", 8))
|
||||
|
||||
s := adopted("ghost-1", "ghost-1-r0", "1", 0, []int{0})
|
||||
s.Model = "a-model-that-was-removed"
|
||||
|
||||
unclaimed := h.r.Adopt([]AdoptedShard{s})
|
||||
if len(unclaimed) != 1 || unclaimed[0].ShardID != "ghost-1-r0" {
|
||||
t.Fatalf("an orphaned shard was not reported unclaimed: %+v", unclaimed)
|
||||
}
|
||||
if got := len(h.r.Replicas()); got != 0 {
|
||||
t.Errorf("a replica was created for an unknown model: %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdoptUnknownProfileIsUnclaimed(t *testing.T) {
|
||||
h := newHarness(t, testModel(1), gpuNode("1", 8))
|
||||
|
||||
s := adopted("llama-9", "llama-9-r0", "1", 0, []int{0})
|
||||
s.Profile = "a-profile-that-was-removed"
|
||||
|
||||
if unclaimed := h.r.Adopt([]AdoptedShard{s}); len(unclaimed) != 1 {
|
||||
t.Errorf("a shard with an unknown profile was claimed: %+v", unclaimed)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdoptIsIdempotent(t *testing.T) {
|
||||
// An agent reconnecting re-reports the same shards; that must not double
|
||||
// the allocation.
|
||||
h := newHarness(t, testModel(1), gpuNode("1", 8))
|
||||
s := []AdoptedShard{adopted("llama-99", "llama-99-r0", "1", 0, []int{0, 1, 2, 3})}
|
||||
|
||||
h.r.Adopt(s)
|
||||
h.r.Adopt(s)
|
||||
|
||||
if got := len(h.r.Replicas()); got != 1 {
|
||||
t.Errorf("re-reporting created %d replicas", got)
|
||||
}
|
||||
if free := len(h.inv.FreeGPUs("1")); free != 4 {
|
||||
t.Errorf("node has %d free GPUs; the allocation was applied twice", free)
|
||||
}
|
||||
}
|
||||
+26
-7
@@ -1259,10 +1259,10 @@ work_dir: /var/lib/llmgw-agent
|
||||
remaining item.
|
||||
|
||||
### Phase 3 — Control plane (2–3 weeks)
|
||||
- [ ] `Register` carries the agent's running shards, so a control plane restart
|
||||
- [x] `Register` carries the agent's running shards, so a control plane restart
|
||||
rebuilds replica state from the fleet instead of orphaning live engines (§6.5)
|
||||
- [ ] Wire the supervisor's orphan scan: a shard nothing claims after one
|
||||
reconciliation pass is killed
|
||||
- [x] Orphan adoption: an agent takes back engines a previous agent left running,
|
||||
and kills whatever it cannot identify or the catalog will not claim
|
||||
- [x] Inventory: nodes, hardware, fabric attachments, allocations
|
||||
- [x] Placement: profile priority, bin packing, NVLink-aware GPU selection,
|
||||
`explain-placement` rejection reasons, property-based invariant tests
|
||||
@@ -1280,10 +1280,29 @@ work_dir: /var/lib/llmgw-agent
|
||||
of the transport.
|
||||
|
||||
### Phase 4 — Multi-node and fabrics (2 weeks)
|
||||
- [ ] Fabric identity (config plus GID prefix), verification
|
||||
- [ ] Multi-node placement (the fabric constraint), port allocation, rank-ordered startup
|
||||
- [ ] NCCL environment generation, HCA selection
|
||||
- [ ] Replica-level failure handling (one dead shard kills the replica)
|
||||
- [x] Fabric identity from the agent's explicit labels
|
||||
- [x] Multi-node placement (the fabric constraint), port allocation, rank ordering
|
||||
- [x] NCCL environment generation, HCA selection
|
||||
- [x] Replica-level failure handling (one dead shard kills the replica)
|
||||
|
||||
Placement and environment generation landed early, with phase 3's work. **The
|
||||
launch path has not, and a multi-node replica cannot currently come up.** Three
|
||||
specific things are missing:
|
||||
|
||||
- [ ] **Workers never become Ready.** The agent health-checks every shard at
|
||||
`127.0.0.1:<port>`, but in a multi-node replica only rank 0 serves HTTP.
|
||||
Workers time out, report Failed, and take the replica with them. The fix
|
||||
needs the engine adapter to say which ranks serve — `ServesHTTP(spec)` —
|
||||
so a worker is judged by liveness and the replica's readiness is rank 0's.
|
||||
- [ ] **Ray is never started.** `VLLM.RayCommand` is written and tested but
|
||||
nothing calls it, so `launch_mode: ray` runs `vllm serve` against a
|
||||
cluster that does not exist. The supervisor assumes one process per
|
||||
shard; Ray needs two, sequenced.
|
||||
- [ ] **Startup is fire-and-forget.** Workers are *sent* before rank 0 but
|
||||
nothing waits for them. SGLang tolerates this because its ranks
|
||||
rendezvous; Ray does not.
|
||||
|
||||
- **Exit criterion:** a 2-node replica on each fabric; killing a node rebuilds the replica on the same fabric and never spreads it across two.
|
||||
- **Exit criterion:** a 2-node replica on each fabric; killing a node rebuilds the replica on the same fabric and never spreads it to the other one.
|
||||
|
||||
### Phase 5 — Production features (3 weeks)
|
||||
|
||||
Reference in new issue
Block a user