Slurm beside llmgw, P2: the agent keeps its node's part
A Slurm message says whether the node runs slurmctld on the leader's files, slurmd configless, both or neither, and whether it is lent. The agent installs munge and Slurm from the machine's sources (22.05 or newer, for cgroup v2), writes the munge key, starts slurmd through a drop-in with --conf-server, resumes a lent node that llmgw drained and drains one that is not, never touching a drain a Slurm administrator made, and reports its version, slurmd -C, state and jobs in SlurmStatus. Told it has no part, it stops what it started. New nodes enter slurm.conf drained. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Ilker ManapandClaude Opus 5.5 committed
1 parent
d439d507d7
commit
35de47444c
8 files changed
+1347
-77
No files matched your search
+363
-73
@@ -108,6 +108,7 @@ type AgentMessage struct {
|
||||
// *AgentMessage_KvEvents
|
||||
// *AgentMessage_TrainStatus
|
||||
// *AgentMessage_TrainArtifact
|
||||
// *AgentMessage_SlurmStatus
|
||||
Msg isAgentMessage_Msg `protobuf_oneof:"msg"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
@@ -240,6 +241,15 @@ func (x *AgentMessage) GetTrainArtifact() *TrainArtifact {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (x *AgentMessage) GetSlurmStatus() *SlurmStatus {
|
||||
if x != nil {
|
||||
if x, ok := x.Msg.(*AgentMessage_SlurmStatus); ok {
|
||||
return x.SlurmStatus
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type isAgentMessage_Msg interface {
|
||||
isAgentMessage_Msg()
|
||||
}
|
||||
@@ -284,6 +294,10 @@ type AgentMessage_TrainArtifact struct {
|
||||
TrainArtifact *TrainArtifact `protobuf:"bytes,10,opt,name=train_artifact,json=trainArtifact,proto3,oneof"`
|
||||
}
|
||||
|
||||
type AgentMessage_SlurmStatus struct {
|
||||
SlurmStatus *SlurmStatus `protobuf:"bytes,11,opt,name=slurm_status,json=slurmStatus,proto3,oneof"`
|
||||
}
|
||||
|
||||
func (*AgentMessage_Register) isAgentMessage_Msg() {}
|
||||
|
||||
func (*AgentMessage_Heartbeat) isAgentMessage_Msg() {}
|
||||
@@ -304,6 +318,8 @@ func (*AgentMessage_TrainStatus) isAgentMessage_Msg() {}
|
||||
|
||||
func (*AgentMessage_TrainArtifact) isAgentMessage_Msg() {}
|
||||
|
||||
func (*AgentMessage_SlurmStatus) isAgentMessage_Msg() {}
|
||||
|
||||
// KVEvents is a change to what one vLLM replica's prefix cache holds, read
|
||||
// by the agent from the engine's KV cache events and sent on for the router's
|
||||
// precise index (plan §7.5 tier 2).
|
||||
@@ -426,6 +442,7 @@ type ControlMessage struct {
|
||||
// *ControlMessage_AgentSettings
|
||||
// *ControlMessage_TrainStart
|
||||
// *ControlMessage_TrainCancel
|
||||
// *ControlMessage_Slurm
|
||||
Msg isControlMessage_Msg `protobuf_oneof:"msg"`
|
||||
// trace is the W3C trace context (traceparent, and tracestate when there
|
||||
// is one) of the control plane's span that sent this command, so what the
|
||||
@@ -618,6 +635,15 @@ func (x *ControlMessage) GetTrainCancel() *TrainCancel {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (x *ControlMessage) GetSlurm() *Slurm {
|
||||
if x != nil {
|
||||
if x, ok := x.Msg.(*ControlMessage_Slurm); ok {
|
||||
return x.Slurm
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (x *ControlMessage) GetTrace() map[string]string {
|
||||
if x != nil {
|
||||
return x.Trace
|
||||
@@ -693,6 +719,10 @@ type ControlMessage_TrainCancel struct {
|
||||
TrainCancel *TrainCancel `protobuf:"bytes,17,opt,name=train_cancel,json=trainCancel,proto3,oneof"`
|
||||
}
|
||||
|
||||
type ControlMessage_Slurm struct {
|
||||
Slurm *Slurm `protobuf:"bytes,18,opt,name=slurm,proto3,oneof"`
|
||||
}
|
||||
|
||||
func (*ControlMessage_Start) isControlMessage_Msg() {}
|
||||
|
||||
func (*ControlMessage_Stop) isControlMessage_Msg() {}
|
||||
@@ -725,6 +755,8 @@ func (*ControlMessage_TrainStart) isControlMessage_Msg() {}
|
||||
|
||||
func (*ControlMessage_TrainCancel) isControlMessage_Msg() {}
|
||||
|
||||
func (*ControlMessage_Slurm) isControlMessage_Msg() {}
|
||||
|
||||
// RestartAgent has the agent start again, as llmgw restart asks of every
|
||||
// node: its engines keep serving and are taken over again.
|
||||
type RestartAgent struct {
|
||||
@@ -5251,11 +5283,234 @@ func (x *TrainArtifact) GetLast() bool {
|
||||
return false
|
||||
}
|
||||
|
||||
// Slurm is what this node is in the Slurm cluster run beside llmgw
|
||||
// (BACKLOG P): its controller, one of its nodes, both or neither. Sent when
|
||||
// an agent registers and whenever it changes; the agent installs what is
|
||||
// missing and keeps the daemons running as it says.
|
||||
type Slurm struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
// controller runs slurmctld here, on files.
|
||||
Controller bool `protobuf:"varint,1,opt,name=controller,proto3" json:"controller,omitempty"`
|
||||
// node runs slurmd here, configless, fetching its files from conf_server.
|
||||
Node bool `protobuf:"varint,2,opt,name=node,proto3" json:"node,omitempty"`
|
||||
// name is this node's NodeName in slurm.conf.
|
||||
Name string `protobuf:"bytes,3,opt,name=name,proto3" json:"name,omitempty"`
|
||||
// conf_server is slurmctld's host:port.
|
||||
ConfServer string `protobuf:"bytes,4,opt,name=conf_server,json=confServer,proto3" json:"conf_server,omitempty"`
|
||||
// munge_key is the key every machine of the cluster shares.
|
||||
MungeKey []byte `protobuf:"bytes,5,opt,name=munge_key,json=mungeKey,proto3" json:"munge_key,omitempty"`
|
||||
// files are the controller's, by name in /etc/slurm; empty for a node
|
||||
// that is not the controller.
|
||||
Files map[string]string `protobuf:"bytes,6,rep,name=files,proto3" json:"files,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"`
|
||||
// lent says Slurm may run jobs here now: resume the node if llmgw drained
|
||||
// it. Otherwise it is drained, its jobs left to finish.
|
||||
Lent bool `protobuf:"varint,7,opt,name=lent,proto3" json:"lent,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
}
|
||||
|
||||
func (x *Slurm) Reset() {
|
||||
*x = Slurm{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[48]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
|
||||
func (x *Slurm) String() string {
|
||||
return protoimpl.X.MessageStringOf(x)
|
||||
}
|
||||
|
||||
func (*Slurm) ProtoMessage() {}
|
||||
|
||||
func (x *Slurm) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[48]
|
||||
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 Slurm.ProtoReflect.Descriptor instead.
|
||||
func (*Slurm) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{48}
|
||||
}
|
||||
|
||||
func (x *Slurm) GetController() bool {
|
||||
if x != nil {
|
||||
return x.Controller
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (x *Slurm) GetNode() bool {
|
||||
if x != nil {
|
||||
return x.Node
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (x *Slurm) GetName() string {
|
||||
if x != nil {
|
||||
return x.Name
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *Slurm) GetConfServer() string {
|
||||
if x != nil {
|
||||
return x.ConfServer
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *Slurm) GetMungeKey() []byte {
|
||||
if x != nil {
|
||||
return x.MungeKey
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (x *Slurm) GetFiles() map[string]string {
|
||||
if x != nil {
|
||||
return x.Files
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (x *Slurm) GetLent() bool {
|
||||
if x != nil {
|
||||
return x.Lent
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// SlurmStatus is what this node's Slurm daemons are doing, sent when it
|
||||
// changes and every few minutes while the node has a part in Slurm.
|
||||
type SlurmStatus struct {
|
||||
state protoimpl.MessageState `protogen:"open.v1"`
|
||||
// version is Slurm's, as "23.11.4"; empty when it is not installed.
|
||||
Version string `protobuf:"bytes,1,opt,name=version,proto3" json:"version,omitempty"`
|
||||
// hardware is `slurmd -C` on a node, which slurm.conf must match.
|
||||
Hardware string `protobuf:"bytes,2,opt,name=hardware,proto3" json:"hardware,omitempty"`
|
||||
SlurmdRunning bool `protobuf:"varint,3,opt,name=slurmd_running,json=slurmdRunning,proto3" json:"slurmd_running,omitempty"`
|
||||
SlurmctldRunning bool `protobuf:"varint,4,opt,name=slurmctld_running,json=slurmctldRunning,proto3" json:"slurmctld_running,omitempty"`
|
||||
// state and reason are the controller's for this node, from scontrol.
|
||||
State string `protobuf:"bytes,5,opt,name=state,proto3" json:"state,omitempty"`
|
||||
Reason string `protobuf:"bytes,6,opt,name=reason,proto3" json:"reason,omitempty"`
|
||||
// jobs running on this node now.
|
||||
Jobs int32 `protobuf:"varint,7,opt,name=jobs,proto3" json:"jobs,omitempty"`
|
||||
// error is the last thing that went wrong keeping Slurm, empty when
|
||||
// nothing did.
|
||||
Error string `protobuf:"bytes,8,opt,name=error,proto3" json:"error,omitempty"`
|
||||
// files_sha256 is of the controller's files as written, so the leader
|
||||
// knows they arrived.
|
||||
FilesSha256 string `protobuf:"bytes,9,opt,name=files_sha256,json=filesSha256,proto3" json:"files_sha256,omitempty"`
|
||||
unknownFields protoimpl.UnknownFields
|
||||
sizeCache protoimpl.SizeCache
|
||||
}
|
||||
|
||||
func (x *SlurmStatus) Reset() {
|
||||
*x = SlurmStatus{}
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[49]
|
||||
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
|
||||
ms.StoreMessageInfo(mi)
|
||||
}
|
||||
|
||||
func (x *SlurmStatus) String() string {
|
||||
return protoimpl.X.MessageStringOf(x)
|
||||
}
|
||||
|
||||
func (*SlurmStatus) ProtoMessage() {}
|
||||
|
||||
func (x *SlurmStatus) ProtoReflect() protoreflect.Message {
|
||||
mi := &file_llmgw_v1_agent_proto_msgTypes[49]
|
||||
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 SlurmStatus.ProtoReflect.Descriptor instead.
|
||||
func (*SlurmStatus) Descriptor() ([]byte, []int) {
|
||||
return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{49}
|
||||
}
|
||||
|
||||
func (x *SlurmStatus) GetVersion() string {
|
||||
if x != nil {
|
||||
return x.Version
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *SlurmStatus) GetHardware() string {
|
||||
if x != nil {
|
||||
return x.Hardware
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *SlurmStatus) GetSlurmdRunning() bool {
|
||||
if x != nil {
|
||||
return x.SlurmdRunning
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (x *SlurmStatus) GetSlurmctldRunning() bool {
|
||||
if x != nil {
|
||||
return x.SlurmctldRunning
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (x *SlurmStatus) GetState() string {
|
||||
if x != nil {
|
||||
return x.State
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *SlurmStatus) GetReason() string {
|
||||
if x != nil {
|
||||
return x.Reason
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *SlurmStatus) GetJobs() int32 {
|
||||
if x != nil {
|
||||
return x.Jobs
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func (x *SlurmStatus) GetError() string {
|
||||
if x != nil {
|
||||
return x.Error
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (x *SlurmStatus) GetFilesSha256() string {
|
||||
if x != nil {
|
||||
return x.FilesSha256
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
var File_llmgw_v1_agent_proto protoreflect.FileDescriptor
|
||||
|
||||
const file_llmgw_v1_agent_proto_rawDesc = "" +
|
||||
"\n" +
|
||||
"\x14llmgw/v1/agent.proto\x12\bllmgw.v1\"\xb7\x04\n" +
|
||||
"\x14llmgw/v1/agent.proto\x12\bllmgw.v1\"\xf3\x04\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" +
|
||||
@@ -5269,7 +5524,8 @@ const file_llmgw_v1_agent_proto_rawDesc = "" +
|
||||
"\tkv_events\x18\b \x01(\v2\x12.llmgw.v1.KVEventsH\x00R\bkvEvents\x12:\n" +
|
||||
"\ftrain_status\x18\t \x01(\v2\x15.llmgw.v1.TrainStatusH\x00R\vtrainStatus\x12@\n" +
|
||||
"\x0etrain_artifact\x18\n" +
|
||||
" \x01(\v2\x17.llmgw.v1.TrainArtifactH\x00R\rtrainArtifactB\x05\n" +
|
||||
" \x01(\v2\x17.llmgw.v1.TrainArtifactH\x00R\rtrainArtifact\x12:\n" +
|
||||
"\fslurm_status\x18\v \x01(\v2\x15.llmgw.v1.SlurmStatusH\x00R\vslurmStatusB\x05\n" +
|
||||
"\x03msg\"\xb5\x01\n" +
|
||||
"\bKVEvents\x12\x1d\n" +
|
||||
"\n" +
|
||||
@@ -5280,7 +5536,7 @@ const file_llmgw_v1_agent_proto_rawDesc = "" +
|
||||
"clearFirst\x12\x16\n" +
|
||||
"\x06stored\x18\x04 \x03(\x06R\x06stored\x12\x18\n" +
|
||||
"\aremoved\x18\x05 \x03(\x06R\aremoved\x12\x18\n" +
|
||||
"\aproblem\x18\x06 \x01(\tR\aproblem\"\xf8\a\n" +
|
||||
"\aproblem\x18\x06 \x01(\tR\aproblem\"\xa1\b\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" +
|
||||
@@ -5301,7 +5557,8 @@ const file_llmgw_v1_agent_proto_rawDesc = "" +
|
||||
"\x0eagent_settings\x18\x0f \x01(\v2\x17.llmgw.v1.AgentSettingsH\x00R\ragentSettings\x127\n" +
|
||||
"\vtrain_start\x18\x10 \x01(\v2\x14.llmgw.v1.TrainStartH\x00R\n" +
|
||||
"trainStart\x12:\n" +
|
||||
"\ftrain_cancel\x18\x11 \x01(\v2\x15.llmgw.v1.TrainCancelH\x00R\vtrainCancel\x129\n" +
|
||||
"\ftrain_cancel\x18\x11 \x01(\v2\x15.llmgw.v1.TrainCancelH\x00R\vtrainCancel\x12'\n" +
|
||||
"\x05slurm\x18\x12 \x01(\v2\x0f.llmgw.v1.SlurmH\x00R\x05slurm\x129\n" +
|
||||
"\x05trace\x18\f \x03(\v2#.llmgw.v1.ControlMessage.TraceEntryR\x05trace\x1a8\n" +
|
||||
"\n" +
|
||||
"TraceEntry\x12\x10\n" +
|
||||
@@ -5754,7 +6011,32 @@ const file_llmgw_v1_agent_proto_rawDesc = "" +
|
||||
"\x04kind\x18\x02 \x01(\tR\x04kind\x12\x10\n" +
|
||||
"\x03seq\x18\x03 \x01(\x05R\x03seq\x12\x12\n" +
|
||||
"\x04data\x18\x04 \x01(\fR\x04data\x12\x12\n" +
|
||||
"\x04last\x18\x05 \x01(\bR\x04last2O\n" +
|
||||
"\x04last\x18\x05 \x01(\bR\x04last\"\x8d\x02\n" +
|
||||
"\x05Slurm\x12\x1e\n" +
|
||||
"\n" +
|
||||
"controller\x18\x01 \x01(\bR\n" +
|
||||
"controller\x12\x12\n" +
|
||||
"\x04node\x18\x02 \x01(\bR\x04node\x12\x12\n" +
|
||||
"\x04name\x18\x03 \x01(\tR\x04name\x12\x1f\n" +
|
||||
"\vconf_server\x18\x04 \x01(\tR\n" +
|
||||
"confServer\x12\x1b\n" +
|
||||
"\tmunge_key\x18\x05 \x01(\fR\bmungeKey\x120\n" +
|
||||
"\x05files\x18\x06 \x03(\v2\x1a.llmgw.v1.Slurm.FilesEntryR\x05files\x12\x12\n" +
|
||||
"\x04lent\x18\a \x01(\bR\x04lent\x1a8\n" +
|
||||
"\n" +
|
||||
"FilesEntry\x12\x10\n" +
|
||||
"\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" +
|
||||
"\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"\x92\x02\n" +
|
||||
"\vSlurmStatus\x12\x18\n" +
|
||||
"\aversion\x18\x01 \x01(\tR\aversion\x12\x1a\n" +
|
||||
"\bhardware\x18\x02 \x01(\tR\bhardware\x12%\n" +
|
||||
"\x0eslurmd_running\x18\x03 \x01(\bR\rslurmdRunning\x12+\n" +
|
||||
"\x11slurmctld_running\x18\x04 \x01(\bR\x10slurmctldRunning\x12\x14\n" +
|
||||
"\x05state\x18\x05 \x01(\tR\x05state\x12\x16\n" +
|
||||
"\x06reason\x18\x06 \x01(\tR\x06reason\x12\x12\n" +
|
||||
"\x04jobs\x18\a \x01(\x05R\x04jobs\x12\x14\n" +
|
||||
"\x05error\x18\b \x01(\tR\x05error\x12!\n" +
|
||||
"\ffiles_sha256\x18\t \x01(\tR\vfilesSha2562O\n" +
|
||||
"\fAgentService\x12?\n" +
|
||||
"\aConnect\x12\x16.llmgw.v1.AgentMessage\x1a\x18.llmgw.v1.ControlMessage(\x010\x01B\x93\x01\n" +
|
||||
"\fcom.llmgw.v1B\n" +
|
||||
@@ -5773,7 +6055,7 @@ 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, 56)
|
||||
var file_llmgw_v1_agent_proto_msgTypes = make([]protoimpl.MessageInfo, 59)
|
||||
var file_llmgw_v1_agent_proto_goTypes = []any{
|
||||
(ShardStatus_State)(0), // 0: llmgw.v1.ShardStatus.State
|
||||
(*AgentMessage)(nil), // 1: llmgw.v1.AgentMessage
|
||||
@@ -5824,14 +6106,17 @@ var file_llmgw_v1_agent_proto_goTypes = []any{
|
||||
(*TrainCancel)(nil), // 46: llmgw.v1.TrainCancel
|
||||
(*TrainStatus)(nil), // 47: llmgw.v1.TrainStatus
|
||||
(*TrainArtifact)(nil), // 48: llmgw.v1.TrainArtifact
|
||||
nil, // 49: llmgw.v1.ControlMessage.TraceEntry
|
||||
nil, // 50: llmgw.v1.Register.LabelsEntry
|
||||
nil, // 51: llmgw.v1.Register.EngineVersionsEntry
|
||||
nil, // 52: llmgw.v1.GPU.NicPathsEntry
|
||||
nil, // 53: llmgw.v1.IBPort.CountersEntry
|
||||
nil, // 54: llmgw.v1.Heartbeat.EngineVersionsEntry
|
||||
nil, // 55: llmgw.v1.EngineWant.OfflineSha256Entry
|
||||
nil, // 56: llmgw.v1.EngineProfile.EnvEntry
|
||||
(*Slurm)(nil), // 49: llmgw.v1.Slurm
|
||||
(*SlurmStatus)(nil), // 50: llmgw.v1.SlurmStatus
|
||||
nil, // 51: llmgw.v1.ControlMessage.TraceEntry
|
||||
nil, // 52: llmgw.v1.Register.LabelsEntry
|
||||
nil, // 53: llmgw.v1.Register.EngineVersionsEntry
|
||||
nil, // 54: llmgw.v1.GPU.NicPathsEntry
|
||||
nil, // 55: llmgw.v1.IBPort.CountersEntry
|
||||
nil, // 56: llmgw.v1.Heartbeat.EngineVersionsEntry
|
||||
nil, // 57: llmgw.v1.EngineWant.OfflineSha256Entry
|
||||
nil, // 58: llmgw.v1.EngineProfile.EnvEntry
|
||||
nil, // 59: llmgw.v1.Slurm.FilesEntry
|
||||
}
|
||||
var file_llmgw_v1_agent_proto_depIdxs = []int32{
|
||||
8, // 0: llmgw.v1.AgentMessage.register:type_name -> llmgw.v1.Register
|
||||
@@ -5844,64 +6129,67 @@ var file_llmgw_v1_agent_proto_depIdxs = []int32{
|
||||
2, // 7: llmgw.v1.AgentMessage.kv_events:type_name -> llmgw.v1.KVEvents
|
||||
47, // 8: llmgw.v1.AgentMessage.train_status:type_name -> llmgw.v1.TrainStatus
|
||||
48, // 9: llmgw.v1.AgentMessage.train_artifact:type_name -> llmgw.v1.TrainArtifact
|
||||
30, // 10: llmgw.v1.ControlMessage.start:type_name -> llmgw.v1.StartShard
|
||||
33, // 11: llmgw.v1.ControlMessage.stop:type_name -> llmgw.v1.StopShard
|
||||
34, // 12: llmgw.v1.ControlMessage.drain:type_name -> llmgw.v1.DrainShard
|
||||
35, // 13: llmgw.v1.ControlMessage.ack:type_name -> llmgw.v1.Ack
|
||||
32, // 14: llmgw.v1.ControlMessage.remove_weights:type_name -> llmgw.v1.RemoveWeights
|
||||
24, // 15: llmgw.v1.ControlMessage.provision:type_name -> llmgw.v1.Provision
|
||||
36, // 16: llmgw.v1.ControlMessage.logs:type_name -> llmgw.v1.LogsRequest
|
||||
38, // 17: llmgw.v1.ControlMessage.fabric_scan:type_name -> llmgw.v1.FabricScanRequest
|
||||
40, // 18: llmgw.v1.ControlMessage.fabric_test:type_name -> llmgw.v1.FabricTestRequest
|
||||
42, // 19: llmgw.v1.ControlMessage.upgrade:type_name -> llmgw.v1.UpgradeAgent
|
||||
7, // 20: llmgw.v1.ControlMessage.control_planes:type_name -> llmgw.v1.ControlPlanes
|
||||
5, // 21: llmgw.v1.ControlMessage.firewall:type_name -> llmgw.v1.Firewall
|
||||
4, // 22: llmgw.v1.ControlMessage.restart_agent:type_name -> llmgw.v1.RestartAgent
|
||||
6, // 23: llmgw.v1.ControlMessage.agent_settings:type_name -> llmgw.v1.AgentSettings
|
||||
45, // 24: llmgw.v1.ControlMessage.train_start:type_name -> llmgw.v1.TrainStart
|
||||
46, // 25: llmgw.v1.ControlMessage.train_cancel:type_name -> llmgw.v1.TrainCancel
|
||||
49, // 26: llmgw.v1.ControlMessage.trace:type_name -> llmgw.v1.ControlMessage.TraceEntry
|
||||
50, // 27: llmgw.v1.Register.labels:type_name -> llmgw.v1.Register.LabelsEntry
|
||||
12, // 28: llmgw.v1.Register.gpus:type_name -> llmgw.v1.GPU
|
||||
14, // 29: llmgw.v1.Register.fabrics:type_name -> llmgw.v1.FabricAttachment
|
||||
51, // 30: llmgw.v1.Register.engine_versions:type_name -> llmgw.v1.Register.EngineVersionsEntry
|
||||
11, // 31: llmgw.v1.Register.running_shards:type_name -> llmgw.v1.RunningShard
|
||||
15, // 32: llmgw.v1.Register.weights:type_name -> llmgw.v1.LocalWeights
|
||||
10, // 33: llmgw.v1.Register.storage:type_name -> llmgw.v1.Storage
|
||||
9, // 34: llmgw.v1.Register.host:type_name -> llmgw.v1.Host
|
||||
21, // 35: llmgw.v1.Register.imex:type_name -> llmgw.v1.IMEX
|
||||
0, // 36: llmgw.v1.RunningShard.state:type_name -> llmgw.v1.ShardStatus.State
|
||||
27, // 37: llmgw.v1.RunningShard.compute:type_name -> llmgw.v1.Compute
|
||||
52, // 38: llmgw.v1.GPU.nic_paths:type_name -> llmgw.v1.GPU.NicPathsEntry
|
||||
13, // 39: llmgw.v1.GPU.mig:type_name -> llmgw.v1.MIGDevice
|
||||
18, // 40: llmgw.v1.NodeLoad.shards:type_name -> llmgw.v1.ShardLoad
|
||||
17, // 41: llmgw.v1.NodeLoad.gpus:type_name -> llmgw.v1.GPUSensor
|
||||
53, // 42: llmgw.v1.IBPort.counters:type_name -> llmgw.v1.IBPort.CountersEntry
|
||||
20, // 43: llmgw.v1.IBPort.subnet_managers:type_name -> llmgw.v1.SubnetManager
|
||||
10, // 44: llmgw.v1.Heartbeat.storage:type_name -> llmgw.v1.Storage
|
||||
15, // 45: llmgw.v1.Heartbeat.weights:type_name -> llmgw.v1.LocalWeights
|
||||
54, // 46: llmgw.v1.Heartbeat.engine_versions:type_name -> llmgw.v1.Heartbeat.EngineVersionsEntry
|
||||
26, // 47: llmgw.v1.Heartbeat.provision:type_name -> llmgw.v1.ProvisionStatus
|
||||
16, // 48: llmgw.v1.Heartbeat.load:type_name -> llmgw.v1.NodeLoad
|
||||
19, // 49: llmgw.v1.Heartbeat.ib_ports:type_name -> llmgw.v1.IBPort
|
||||
14, // 50: llmgw.v1.Heartbeat.fabrics:type_name -> llmgw.v1.FabricAttachment
|
||||
22, // 51: llmgw.v1.Heartbeat.gpu_fabric:type_name -> llmgw.v1.GPUFabricState
|
||||
21, // 52: llmgw.v1.Heartbeat.imex:type_name -> llmgw.v1.IMEX
|
||||
25, // 53: llmgw.v1.Provision.engines:type_name -> llmgw.v1.EngineWant
|
||||
55, // 54: llmgw.v1.EngineWant.offline_sha256:type_name -> llmgw.v1.EngineWant.OfflineSha256Entry
|
||||
0, // 55: llmgw.v1.ShardStatus.state:type_name -> llmgw.v1.ShardStatus.State
|
||||
29, // 56: llmgw.v1.ShardStatus.fetch:type_name -> llmgw.v1.FetchProgress
|
||||
27, // 57: llmgw.v1.ShardStatus.compute:type_name -> llmgw.v1.Compute
|
||||
31, // 58: llmgw.v1.StartShard.profile:type_name -> llmgw.v1.EngineProfile
|
||||
56, // 59: llmgw.v1.EngineProfile.env:type_name -> llmgw.v1.EngineProfile.EnvEntry
|
||||
44, // 60: llmgw.v1.EngineProfile.adapters:type_name -> llmgw.v1.LoraAdapter
|
||||
1, // 61: llmgw.v1.AgentService.Connect:input_type -> llmgw.v1.AgentMessage
|
||||
3, // 62: llmgw.v1.AgentService.Connect:output_type -> llmgw.v1.ControlMessage
|
||||
62, // [62:63] is the sub-list for method output_type
|
||||
61, // [61:62] is the sub-list for method input_type
|
||||
61, // [61:61] is the sub-list for extension type_name
|
||||
61, // [61:61] is the sub-list for extension extendee
|
||||
0, // [0:61] is the sub-list for field type_name
|
||||
50, // 10: llmgw.v1.AgentMessage.slurm_status:type_name -> llmgw.v1.SlurmStatus
|
||||
30, // 11: llmgw.v1.ControlMessage.start:type_name -> llmgw.v1.StartShard
|
||||
33, // 12: llmgw.v1.ControlMessage.stop:type_name -> llmgw.v1.StopShard
|
||||
34, // 13: llmgw.v1.ControlMessage.drain:type_name -> llmgw.v1.DrainShard
|
||||
35, // 14: llmgw.v1.ControlMessage.ack:type_name -> llmgw.v1.Ack
|
||||
32, // 15: llmgw.v1.ControlMessage.remove_weights:type_name -> llmgw.v1.RemoveWeights
|
||||
24, // 16: llmgw.v1.ControlMessage.provision:type_name -> llmgw.v1.Provision
|
||||
36, // 17: llmgw.v1.ControlMessage.logs:type_name -> llmgw.v1.LogsRequest
|
||||
38, // 18: llmgw.v1.ControlMessage.fabric_scan:type_name -> llmgw.v1.FabricScanRequest
|
||||
40, // 19: llmgw.v1.ControlMessage.fabric_test:type_name -> llmgw.v1.FabricTestRequest
|
||||
42, // 20: llmgw.v1.ControlMessage.upgrade:type_name -> llmgw.v1.UpgradeAgent
|
||||
7, // 21: llmgw.v1.ControlMessage.control_planes:type_name -> llmgw.v1.ControlPlanes
|
||||
5, // 22: llmgw.v1.ControlMessage.firewall:type_name -> llmgw.v1.Firewall
|
||||
4, // 23: llmgw.v1.ControlMessage.restart_agent:type_name -> llmgw.v1.RestartAgent
|
||||
6, // 24: llmgw.v1.ControlMessage.agent_settings:type_name -> llmgw.v1.AgentSettings
|
||||
45, // 25: llmgw.v1.ControlMessage.train_start:type_name -> llmgw.v1.TrainStart
|
||||
46, // 26: llmgw.v1.ControlMessage.train_cancel:type_name -> llmgw.v1.TrainCancel
|
||||
49, // 27: llmgw.v1.ControlMessage.slurm:type_name -> llmgw.v1.Slurm
|
||||
51, // 28: llmgw.v1.ControlMessage.trace:type_name -> llmgw.v1.ControlMessage.TraceEntry
|
||||
52, // 29: llmgw.v1.Register.labels:type_name -> llmgw.v1.Register.LabelsEntry
|
||||
12, // 30: llmgw.v1.Register.gpus:type_name -> llmgw.v1.GPU
|
||||
14, // 31: llmgw.v1.Register.fabrics:type_name -> llmgw.v1.FabricAttachment
|
||||
53, // 32: llmgw.v1.Register.engine_versions:type_name -> llmgw.v1.Register.EngineVersionsEntry
|
||||
11, // 33: llmgw.v1.Register.running_shards:type_name -> llmgw.v1.RunningShard
|
||||
15, // 34: llmgw.v1.Register.weights:type_name -> llmgw.v1.LocalWeights
|
||||
10, // 35: llmgw.v1.Register.storage:type_name -> llmgw.v1.Storage
|
||||
9, // 36: llmgw.v1.Register.host:type_name -> llmgw.v1.Host
|
||||
21, // 37: llmgw.v1.Register.imex:type_name -> llmgw.v1.IMEX
|
||||
0, // 38: llmgw.v1.RunningShard.state:type_name -> llmgw.v1.ShardStatus.State
|
||||
27, // 39: llmgw.v1.RunningShard.compute:type_name -> llmgw.v1.Compute
|
||||
54, // 40: llmgw.v1.GPU.nic_paths:type_name -> llmgw.v1.GPU.NicPathsEntry
|
||||
13, // 41: llmgw.v1.GPU.mig:type_name -> llmgw.v1.MIGDevice
|
||||
18, // 42: llmgw.v1.NodeLoad.shards:type_name -> llmgw.v1.ShardLoad
|
||||
17, // 43: llmgw.v1.NodeLoad.gpus:type_name -> llmgw.v1.GPUSensor
|
||||
55, // 44: llmgw.v1.IBPort.counters:type_name -> llmgw.v1.IBPort.CountersEntry
|
||||
20, // 45: llmgw.v1.IBPort.subnet_managers:type_name -> llmgw.v1.SubnetManager
|
||||
10, // 46: llmgw.v1.Heartbeat.storage:type_name -> llmgw.v1.Storage
|
||||
15, // 47: llmgw.v1.Heartbeat.weights:type_name -> llmgw.v1.LocalWeights
|
||||
56, // 48: llmgw.v1.Heartbeat.engine_versions:type_name -> llmgw.v1.Heartbeat.EngineVersionsEntry
|
||||
26, // 49: llmgw.v1.Heartbeat.provision:type_name -> llmgw.v1.ProvisionStatus
|
||||
16, // 50: llmgw.v1.Heartbeat.load:type_name -> llmgw.v1.NodeLoad
|
||||
19, // 51: llmgw.v1.Heartbeat.ib_ports:type_name -> llmgw.v1.IBPort
|
||||
14, // 52: llmgw.v1.Heartbeat.fabrics:type_name -> llmgw.v1.FabricAttachment
|
||||
22, // 53: llmgw.v1.Heartbeat.gpu_fabric:type_name -> llmgw.v1.GPUFabricState
|
||||
21, // 54: llmgw.v1.Heartbeat.imex:type_name -> llmgw.v1.IMEX
|
||||
25, // 55: llmgw.v1.Provision.engines:type_name -> llmgw.v1.EngineWant
|
||||
57, // 56: llmgw.v1.EngineWant.offline_sha256:type_name -> llmgw.v1.EngineWant.OfflineSha256Entry
|
||||
0, // 57: llmgw.v1.ShardStatus.state:type_name -> llmgw.v1.ShardStatus.State
|
||||
29, // 58: llmgw.v1.ShardStatus.fetch:type_name -> llmgw.v1.FetchProgress
|
||||
27, // 59: llmgw.v1.ShardStatus.compute:type_name -> llmgw.v1.Compute
|
||||
31, // 60: llmgw.v1.StartShard.profile:type_name -> llmgw.v1.EngineProfile
|
||||
58, // 61: llmgw.v1.EngineProfile.env:type_name -> llmgw.v1.EngineProfile.EnvEntry
|
||||
44, // 62: llmgw.v1.EngineProfile.adapters:type_name -> llmgw.v1.LoraAdapter
|
||||
59, // 63: llmgw.v1.Slurm.files:type_name -> llmgw.v1.Slurm.FilesEntry
|
||||
1, // 64: llmgw.v1.AgentService.Connect:input_type -> llmgw.v1.AgentMessage
|
||||
3, // 65: llmgw.v1.AgentService.Connect:output_type -> llmgw.v1.ControlMessage
|
||||
65, // [65:66] is the sub-list for method output_type
|
||||
64, // [64:65] is the sub-list for method input_type
|
||||
64, // [64:64] is the sub-list for extension type_name
|
||||
64, // [64:64] is the sub-list for extension extendee
|
||||
0, // [0:64] is the sub-list for field type_name
|
||||
}
|
||||
|
||||
func init() { file_llmgw_v1_agent_proto_init() }
|
||||
@@ -5920,6 +6208,7 @@ func file_llmgw_v1_agent_proto_init() {
|
||||
(*AgentMessage_KvEvents)(nil),
|
||||
(*AgentMessage_TrainStatus)(nil),
|
||||
(*AgentMessage_TrainArtifact)(nil),
|
||||
(*AgentMessage_SlurmStatus)(nil),
|
||||
}
|
||||
file_llmgw_v1_agent_proto_msgTypes[2].OneofWrappers = []any{
|
||||
(*ControlMessage_Start)(nil),
|
||||
@@ -5938,6 +6227,7 @@ func file_llmgw_v1_agent_proto_init() {
|
||||
(*ControlMessage_AgentSettings)(nil),
|
||||
(*ControlMessage_TrainStart)(nil),
|
||||
(*ControlMessage_TrainCancel)(nil),
|
||||
(*ControlMessage_Slurm)(nil),
|
||||
}
|
||||
type x struct{}
|
||||
out := protoimpl.TypeBuilder{
|
||||
@@ -5945,7 +6235,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: 56,
|
||||
NumMessages: 59,
|
||||
NumExtensions: 0,
|
||||
NumServices: 1,
|
||||
},
|
||||
|
||||
@@ -30,6 +30,7 @@ message AgentMessage {
|
||||
KVEvents kv_events = 8;
|
||||
TrainStatus train_status = 9;
|
||||
TrainArtifact train_artifact = 10;
|
||||
SlurmStatus slurm_status = 11;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -78,6 +79,7 @@ message ControlMessage {
|
||||
AgentSettings agent_settings = 15;
|
||||
TrainStart train_start = 16;
|
||||
TrainCancel train_cancel = 17;
|
||||
Slurm slurm = 18;
|
||||
}
|
||||
// trace is the W3C trace context (traceparent, and tracestate when there
|
||||
// is one) of the control plane's span that sent this command, so what the
|
||||
@@ -986,3 +988,48 @@ message TrainArtifact {
|
||||
bytes data = 4;
|
||||
bool last = 5;
|
||||
}
|
||||
|
||||
// Slurm is what this node is in the Slurm cluster run beside llmgw
|
||||
// (BACKLOG P): its controller, one of its nodes, both or neither. Sent when
|
||||
// an agent registers and whenever it changes; the agent installs what is
|
||||
// missing and keeps the daemons running as it says.
|
||||
message Slurm {
|
||||
// controller runs slurmctld here, on files.
|
||||
bool controller = 1;
|
||||
// node runs slurmd here, configless, fetching its files from conf_server.
|
||||
bool node = 2;
|
||||
// name is this node's NodeName in slurm.conf.
|
||||
string name = 3;
|
||||
// conf_server is slurmctld's host:port.
|
||||
string conf_server = 4;
|
||||
// munge_key is the key every machine of the cluster shares.
|
||||
bytes munge_key = 5;
|
||||
// files are the controller's, by name in /etc/slurm; empty for a node
|
||||
// that is not the controller.
|
||||
map<string, string> files = 6;
|
||||
// lent says Slurm may run jobs here now: resume the node if llmgw drained
|
||||
// it. Otherwise it is drained, its jobs left to finish.
|
||||
bool lent = 7;
|
||||
}
|
||||
|
||||
// SlurmStatus is what this node's Slurm daemons are doing, sent when it
|
||||
// changes and every few minutes while the node has a part in Slurm.
|
||||
message SlurmStatus {
|
||||
// version is Slurm's, as "23.11.4"; empty when it is not installed.
|
||||
string version = 1;
|
||||
// hardware is `slurmd -C` on a node, which slurm.conf must match.
|
||||
string hardware = 2;
|
||||
bool slurmd_running = 3;
|
||||
bool slurmctld_running = 4;
|
||||
// state and reason are the controller's for this node, from scontrol.
|
||||
string state = 5;
|
||||
string reason = 6;
|
||||
// jobs running on this node now.
|
||||
int32 jobs = 7;
|
||||
// error is the last thing that went wrong keeping Slurm, empty when
|
||||
// nothing did.
|
||||
string error = 8;
|
||||
// files_sha256 is of the controller's files as written, so the leader
|
||||
// knows they arrived.
|
||||
string files_sha256 = 9;
|
||||
}
|
||||
+12
-1
@@ -35,6 +35,7 @@ import (
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/provision"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/slurmnode"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/train"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/weights"
|
||||
@@ -148,6 +149,10 @@ type Agent struct {
|
||||
|
||||
// trainer runs LoRA training jobs, one at a time (BACKLOG F4).
|
||||
trainer *train.Runner
|
||||
// slurm keeps the node's part in Slurm (BACKLOG P); slurmOut is where
|
||||
// its status goes, under mu.
|
||||
slurm *slurmnode.Keeper
|
||||
slurmOut chan<- *llmgwv1.AgentMessage
|
||||
}
|
||||
|
||||
// connectHTTPClient is the HTTP client connect uses; an interface so tests can
|
||||
@@ -206,6 +211,7 @@ func New(cfg Config, httpClient connectHTTPClient, log *slog.Logger) (*Agent, er
|
||||
}
|
||||
a.trainer = train.New(train.Env{EnginesDir: enginesDir, WorkDir: filepath.Join(cfg.Superv.WorkDir, "trainer"),
|
||||
RunAs: runAs, Log: log})
|
||||
a.slurm = slurmnode.New(slurmnode.Local{}, a.prov, slurmMarker, log)
|
||||
return a, nil
|
||||
}
|
||||
|
||||
@@ -232,6 +238,7 @@ func (a *Agent) Run(ctx context.Context) error {
|
||||
}
|
||||
go a.provisionLoop(ctx)
|
||||
go a.ensureDiags(ctx)
|
||||
go a.slurm.Run(ctx, a.sendSlurm)
|
||||
|
||||
attempt := 0
|
||||
next := 0
|
||||
@@ -578,7 +585,7 @@ func (a *Agent) session(ctx context.Context, server string) error {
|
||||
reg.EngineVersions = a.engines(parent)
|
||||
reg.EnginesProbed = true
|
||||
reg.InstallableEngines = slices.DeleteFunc(provision.Installable(), a.disabled)
|
||||
reg.Features = []string{FeatureFabricTest}
|
||||
reg.Features = []string{FeatureFabricTest, FeatureSlurm}
|
||||
reg.RunningShards = a.runningShards()
|
||||
a.mu.Lock()
|
||||
h := a.prov.Host
|
||||
@@ -628,6 +635,8 @@ func (a *Agent) session(ctx context.Context, server string) error {
|
||||
go a.watchExits(ctx, out, unready(reg.RunningShards))
|
||||
go a.forwardKV(ctx, out)
|
||||
go a.reportCompute(ctx, out)
|
||||
a.setSlurmOut(out)
|
||||
defer a.setSlurmOut(nil)
|
||||
|
||||
for {
|
||||
msg, err := stream.Receive()
|
||||
@@ -951,6 +960,8 @@ func (a *Agent) handle(ctx context.Context, msg *llmgwv1.ControlMessage, out cha
|
||||
go a.applySettings(m.AgentSettings)
|
||||
case *llmgwv1.ControlMessage_TrainStart:
|
||||
go a.trainJob(ctx, m.TrainStart, out)
|
||||
case *llmgwv1.ControlMessage_Slurm:
|
||||
a.slurm.Apply(m.Slurm)
|
||||
case *llmgwv1.ControlMessage_TrainCancel:
|
||||
a.trainer.Cancel(m.TrainCancel.GetJobId())
|
||||
case *llmgwv1.ControlMessage_RestartAgent:
|
||||
|
||||
@@ -0,0 +1,38 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
||||
)
|
||||
|
||||
// FeatureSlurm is the agent keeping its node's part in Slurm (BACKLOG P).
|
||||
const FeatureSlurm = "slurm"
|
||||
|
||||
// slurmMarker says llmgw started Slurm's daemons on this machine.
|
||||
var slurmMarker = "/var/lib/llmgw-agent/slurm"
|
||||
|
||||
// setSlurmOut is where Slurm's status goes now: this session's queue, or
|
||||
// nowhere between sessions.
|
||||
func (a *Agent) setSlurmOut(out chan<- *llmgwv1.AgentMessage) {
|
||||
a.mu.Lock()
|
||||
a.slurmOut = out
|
||||
a.mu.Unlock()
|
||||
if out != nil {
|
||||
a.slurm.Resend()
|
||||
}
|
||||
}
|
||||
|
||||
// sendSlurm queues a status without waiting: a full queue or no session
|
||||
// drops it, and the next session is sent it again.
|
||||
func (a *Agent) sendSlurm(st *llmgwv1.SlurmStatus) {
|
||||
a.mu.Lock()
|
||||
out := a.slurmOut
|
||||
a.mu.Unlock()
|
||||
if out == nil {
|
||||
return
|
||||
}
|
||||
select {
|
||||
case out <- &llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_SlurmStatus{SlurmStatus: st}}:
|
||||
default:
|
||||
a.log.Warn("Slurm's status not sent: the queue to the control plane is full")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,572 @@
|
||||
// Package slurmnode keeps this machine's part in the Slurm cluster run beside
|
||||
// llmgw (BACKLOG P), as the leader says: slurmctld on the controller, slurmd
|
||||
// on a node, munge on both, installed from the machine's own package sources.
|
||||
//
|
||||
// What the leader sends is wanted state, kept by Keeper's loop: sent again,
|
||||
// or checked again a minute later, it changes nothing that is already so.
|
||||
package slurmnode
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io/fs"
|
||||
"log/slog"
|
||||
"os"
|
||||
"os/exec"
|
||||
"os/user"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"google.golang.org/protobuf/proto"
|
||||
|
||||
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/slurm"
|
||||
)
|
||||
|
||||
// Host is what Keeper does to the machine; Local in the agent, a fake in
|
||||
// tests.
|
||||
type Host interface {
|
||||
Run(ctx context.Context, env []string, name string, args ...string) ([]byte, error)
|
||||
LookPath(name string) (string, error)
|
||||
ReadFile(path string) ([]byte, error)
|
||||
// WriteFile writes b at path with mode, owned by owner ("" for root),
|
||||
// making its directory.
|
||||
WriteFile(path string, b []byte, mode fs.FileMode, owner string) error
|
||||
// MkdirAll makes path, owned by owner.
|
||||
MkdirAll(path string, mode fs.FileMode, owner string) error
|
||||
Remove(path string) error
|
||||
}
|
||||
|
||||
// Installer installs packages by family, "deb" and "rpm".
|
||||
type Installer interface {
|
||||
InstallPackages(ctx context.Context, byFamily map[string][]string) error
|
||||
}
|
||||
|
||||
// Reasons llmgw drains a node with; one drained for another reason, by a
|
||||
// Slurm administrator, is never resumed by llmgw.
|
||||
const (
|
||||
ReasonServing = "llmgw: serving models"
|
||||
reasonPrefix = "llmgw"
|
||||
)
|
||||
|
||||
// DropIn is the unit drop-in that starts slurmd configless.
|
||||
const DropIn = "/etc/systemd/system/slurmd.service.d/llmgw.conf"
|
||||
|
||||
// Keeper keeps the machine as the last Slurm message says.
|
||||
type Keeper struct {
|
||||
host Host
|
||||
inst Installer
|
||||
log *slog.Logger
|
||||
marker string
|
||||
|
||||
mu sync.Mutex
|
||||
want *llmgwv1.Slurm
|
||||
kick chan struct{}
|
||||
// resend has the next status sent even if unchanged: a new stream's
|
||||
// control plane may never have had it.
|
||||
resend bool
|
||||
// installed is true once the packages were found or put there, so a
|
||||
// minute's check does not look again.
|
||||
installed map[string]bool
|
||||
files string
|
||||
}
|
||||
|
||||
// New is a Keeper; marker is a file it keeps while it runs Slurm here, so
|
||||
// an agent started again knows to stop what it started when told it has no
|
||||
// part any more.
|
||||
func New(host Host, inst Installer, marker string, log *slog.Logger) *Keeper {
|
||||
if log == nil {
|
||||
log = slog.Default()
|
||||
}
|
||||
return &Keeper{host: host, inst: inst, log: log, marker: marker,
|
||||
kick: make(chan struct{}, 1), installed: map[string]bool{}}
|
||||
}
|
||||
|
||||
// Apply is a new message from the leader: kept from now on.
|
||||
func (k *Keeper) Apply(m *llmgwv1.Slurm) {
|
||||
k.mu.Lock()
|
||||
k.want = proto.CloneOf(m)
|
||||
k.mu.Unlock()
|
||||
select {
|
||||
case k.kick <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
// Resend has the status sent again on the next pass, which it starts now.
|
||||
func (k *Keeper) Resend() {
|
||||
k.mu.Lock()
|
||||
k.resend = true
|
||||
k.mu.Unlock()
|
||||
select {
|
||||
case k.kick <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
// Every is how often the machine is checked again without a new message.
|
||||
var Every = time.Minute
|
||||
|
||||
// resendAfter is how often an unchanged status is sent anyway.
|
||||
const resendAfter = 5 * time.Minute
|
||||
|
||||
// Run keeps the machine until ctx ends, sending its status when it changes
|
||||
// and every few minutes while it has a part.
|
||||
func (k *Keeper) Run(ctx context.Context, send func(*llmgwv1.SlurmStatus)) {
|
||||
t := time.NewTicker(Every)
|
||||
defer t.Stop()
|
||||
var last *llmgwv1.SlurmStatus
|
||||
var sent time.Time
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-k.kick:
|
||||
case <-t.C:
|
||||
}
|
||||
k.mu.Lock()
|
||||
want, again := k.want, k.resend
|
||||
k.resend = false
|
||||
k.mu.Unlock()
|
||||
if want == nil {
|
||||
continue
|
||||
}
|
||||
st := k.Reconcile(ctx, want)
|
||||
if st == nil {
|
||||
last = nil
|
||||
continue
|
||||
}
|
||||
if again || !proto.Equal(st, last) || time.Since(sent) > resendAfter {
|
||||
send(st)
|
||||
last, sent = st, time.Now()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Reconcile makes the machine what want says once, and answers its status;
|
||||
// nil when it has no part in Slurm.
|
||||
func (k *Keeper) Reconcile(ctx context.Context, want *llmgwv1.Slurm) *llmgwv1.SlurmStatus {
|
||||
ctx, cancel := context.WithTimeout(ctx, 20*time.Minute)
|
||||
defer cancel()
|
||||
if !want.GetController() && !want.GetNode() {
|
||||
k.stop(ctx)
|
||||
return nil
|
||||
}
|
||||
st := &llmgwv1.SlurmStatus{}
|
||||
fail := func(what string, err error) *llmgwv1.SlurmStatus {
|
||||
st.Error = what + ": " + err.Error()
|
||||
k.log.Warn("keeping Slurm on this node: "+what, "err", err)
|
||||
return st
|
||||
}
|
||||
if err := k.ensurePackages(ctx, want); err != nil {
|
||||
return fail("installing Slurm", err)
|
||||
}
|
||||
out, err := k.host.Run(ctx, nil, "scontrol", "--version")
|
||||
if err != nil {
|
||||
return fail("reading Slurm's version", err)
|
||||
}
|
||||
v, err := slurm.ParseVersion(string(out))
|
||||
if err != nil {
|
||||
return fail("reading Slurm's version", err)
|
||||
}
|
||||
st.Version = strings.TrimPrefix(strings.TrimSpace(string(out)), "slurm ")
|
||||
if slurm.TooOld(v) {
|
||||
return fail("Slurm", fmt.Errorf("%d.%02d is older than %d.%02d, the first to run jobs in cgroup v2; "+
|
||||
"a newer release of the distribution, or Slurm's own packages, has one", v[0], v[1],
|
||||
slurm.MinVersion[0], slurm.MinVersion[1]))
|
||||
}
|
||||
if err := k.host.WriteFile(k.marker, []byte("llmgw runs Slurm here (BACKLOG P)\n"), 0o644, ""); err != nil {
|
||||
return fail("keeping its marker", err)
|
||||
}
|
||||
if err := k.munge(ctx, want.GetMungeKey()); err != nil {
|
||||
return fail("munge", err)
|
||||
}
|
||||
if want.GetController() {
|
||||
sum, err := k.controller(ctx, want.GetFiles())
|
||||
st.FilesSha256 = sum
|
||||
if err != nil {
|
||||
return fail("slurmctld", err)
|
||||
}
|
||||
st.SlurmctldRunning = k.active(ctx, "slurmctld")
|
||||
}
|
||||
if want.GetNode() {
|
||||
if err := k.node(ctx, want, st); err != nil {
|
||||
return fail("slurmd", err)
|
||||
}
|
||||
}
|
||||
return st
|
||||
}
|
||||
|
||||
func (k *Keeper) ensurePackages(ctx context.Context, want *llmgwv1.Slurm) error {
|
||||
need := map[string][2]string{
|
||||
"munged": {"munge", "munge"},
|
||||
"scontrol": {"slurm-client", "slurm"},
|
||||
}
|
||||
if want.GetController() {
|
||||
need["slurmctld"] = [2]string{"slurmctld", "slurm-slurmctld"}
|
||||
}
|
||||
if want.GetNode() {
|
||||
need["slurmd"] = [2]string{"slurmd", "slurm-slurmd"}
|
||||
}
|
||||
pkgs := map[string][]string{}
|
||||
for bin, p := range need {
|
||||
if k.installed[bin] {
|
||||
continue
|
||||
}
|
||||
if _, err := k.host.LookPath(bin); err == nil {
|
||||
k.installed[bin] = true
|
||||
continue
|
||||
}
|
||||
pkgs["deb"] = append(pkgs["deb"], p[0])
|
||||
pkgs["rpm"] = append(pkgs["rpm"], p[1])
|
||||
}
|
||||
if len(pkgs) == 0 {
|
||||
return nil
|
||||
}
|
||||
slices.Sort(pkgs["deb"])
|
||||
slices.Sort(pkgs["rpm"])
|
||||
if k.inst == nil {
|
||||
return fmt.Errorf("%s missing, and this agent installs nothing", strings.Join(pkgs["deb"], ", "))
|
||||
}
|
||||
k.log.Info("installing Slurm, as the cluster's settings say", "packages", pkgs["deb"])
|
||||
if err := k.inst.InstallPackages(ctx, pkgs); err != nil {
|
||||
return err
|
||||
}
|
||||
for bin := range need {
|
||||
if _, err := k.host.LookPath(bin); err != nil {
|
||||
return fmt.Errorf("%s is still missing after installing %s", bin, strings.Join(pkgs["deb"], ", "))
|
||||
}
|
||||
k.installed[bin] = true
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// munge writes the cluster's key and keeps munged running on it.
|
||||
func (k *Keeper) munge(ctx context.Context, key []byte) error {
|
||||
if len(key) < 32 {
|
||||
return errors.New("the leader sent no key")
|
||||
}
|
||||
changed, err := k.write(slurm.MungeKeyPath, key, 0o400, "munge")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if changed {
|
||||
k.log.Info("munge key written; munge starts again on it")
|
||||
return k.systemctl(ctx, "enable", "--now", "munge", "&&restart")
|
||||
}
|
||||
if !k.active(ctx, "munge") {
|
||||
return k.systemctl(ctx, "enable", "--now", "munge")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// controller writes the files and keeps slurmctld running on them,
|
||||
// answering their digest.
|
||||
func (k *Keeper) controller(ctx context.Context, files map[string]string) (string, error) {
|
||||
if files["slurm.conf"] == "" {
|
||||
return "", errors.New("the leader sent no slurm.conf")
|
||||
}
|
||||
for _, d := range []string{slurm.StateDir, slurm.LogDir} {
|
||||
if err := k.host.MkdirAll(d, 0o755, "slurm"); err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
names := make([]string, 0, len(files))
|
||||
for n := range files {
|
||||
names = append(names, n)
|
||||
}
|
||||
slices.Sort(names)
|
||||
h := sha256.New()
|
||||
changed := false
|
||||
for _, n := range names {
|
||||
if strings.ContainsAny(n, "/\\") || n == "" || n == "." || n == ".." {
|
||||
return "", fmt.Errorf("the leader sent a file named %q", n)
|
||||
}
|
||||
fmt.Fprintf(h, "%s\x00%s\x00", n, files[n])
|
||||
c, err := k.write(filepath.Join(slurm.ConfDir, n), []byte(files[n]), 0o644, "")
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
changed = changed || c
|
||||
}
|
||||
sum := hex.EncodeToString(h.Sum(nil))
|
||||
if changed || sum != k.files {
|
||||
if err := k.systemctl(ctx, "enable", "slurmctld"); err != nil {
|
||||
return sum, err
|
||||
}
|
||||
// Nodes added or taken away need a restart before 23.11 and a
|
||||
// reconfigure after it; a restart does both, and running jobs live
|
||||
// through it.
|
||||
k.log.Info("slurmctld starts again on new files")
|
||||
if err := k.systemctl(ctx, "restart", "slurmctld"); err != nil {
|
||||
return sum, err
|
||||
}
|
||||
if err := k.waitPing(ctx); err != nil {
|
||||
return sum, err
|
||||
}
|
||||
// Hands the configless slurmd the new files.
|
||||
if out, err := k.host.Run(ctx, nil, "scontrol", "reconfigure"); err != nil {
|
||||
return sum, fmt.Errorf("scontrol reconfigure: %v: %s", err, lastLine(out))
|
||||
}
|
||||
k.files = sum
|
||||
return sum, nil
|
||||
}
|
||||
if !k.active(ctx, "slurmctld") {
|
||||
return sum, k.systemctl(ctx, "enable", "--now", "slurmctld")
|
||||
}
|
||||
return sum, nil
|
||||
}
|
||||
|
||||
func (k *Keeper) waitPing(ctx context.Context) error {
|
||||
var last []byte
|
||||
for range 30 {
|
||||
out, err := k.host.Run(ctx, nil, "scontrol", "ping")
|
||||
if err == nil && bytes.Contains(out, []byte("UP")) {
|
||||
return nil
|
||||
}
|
||||
last = out
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case <-time.After(pingWait):
|
||||
}
|
||||
}
|
||||
return fmt.Errorf("slurmctld did not answer after starting: %s", lastLine(last))
|
||||
}
|
||||
|
||||
var pingWait = time.Second
|
||||
|
||||
// node keeps slurmd running configless and puts the node in or out of the
|
||||
// controller's hands as lent says.
|
||||
func (k *Keeper) node(ctx context.Context, want *llmgwv1.Slurm, st *llmgwv1.SlurmStatus) error {
|
||||
name, server := want.GetName(), want.GetConfServer()
|
||||
if name == "" || server == "" {
|
||||
return errors.New("the leader sent no name or controller")
|
||||
}
|
||||
if strings.ContainsAny(name+server, " \n\t\"'\\$`") {
|
||||
return fmt.Errorf("the leader sent a name or controller with spaces or quotes in it")
|
||||
}
|
||||
for _, d := range []string{slurm.SpoolDir, slurm.LogDir} {
|
||||
if err := k.host.MkdirAll(d, 0o755, "slurm"); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
bin, err := k.host.LookPath("slurmd")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
unit := "# Written by llmgw (BACKLOG P): slurmd fetches its files from slurmctld.\n" +
|
||||
"[Unit]\nConditionPathExists=\n" +
|
||||
"[Service]\nExecStart=\nExecStart=" + bin + " -D -s --conf-server " + server + " -N " + name + "\n"
|
||||
changed, err := k.write(DropIn, []byte(unit), 0o644, "")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if out, err := k.host.Run(ctx, nil, "slurmd", "-C"); err == nil {
|
||||
st.Hardware = strings.TrimSpace(string(out))
|
||||
}
|
||||
if changed {
|
||||
if err := k.systemctl(ctx, "daemon-reload"); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := k.systemctl(ctx, "enable", "slurmd"); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := k.systemctl(ctx, "restart", "slurmd"); err != nil {
|
||||
return err
|
||||
}
|
||||
} else if !k.active(ctx, "slurmd") {
|
||||
if err := k.systemctl(ctx, "enable", "--now", "slurmd"); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
st.SlurmdRunning = k.active(ctx, "slurmd")
|
||||
|
||||
env := k.clientEnv()
|
||||
out, err := k.host.Run(ctx, env, "scontrol", "show", "node", "-o", name)
|
||||
if err != nil {
|
||||
// Not in slurm.conf yet: the leader writes it once this node's
|
||||
// hardware is in.
|
||||
return nil
|
||||
}
|
||||
ns, err := slurm.ParseNodeState(string(out))
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
ours := ns.Reason == "" || strings.HasPrefix(ns.Reason, reasonPrefix)
|
||||
switch {
|
||||
case want.GetLent() && ns.Drained() && ours:
|
||||
k.log.Info("this node is lent to Slurm: resumed there")
|
||||
if out, err := k.host.Run(ctx, env, "scontrol", "update", "NodeName="+name, "State=RESUME"); err != nil {
|
||||
return fmt.Errorf("resuming the node: %v: %s", err, lastLine(out))
|
||||
}
|
||||
case !want.GetLent() && !strings.Contains(ns.State, "DRAIN"):
|
||||
k.log.Info("this node is llmgw's: drained in Slurm, its jobs left to finish")
|
||||
if out, err := k.host.Run(ctx, env, "scontrol", "update", "NodeName="+name, "State=DRAIN",
|
||||
"Reason="+ReasonServing); err != nil {
|
||||
return fmt.Errorf("draining the node: %v: %s", err, lastLine(out))
|
||||
}
|
||||
}
|
||||
if out, err := k.host.Run(ctx, env, "scontrol", "show", "node", "-o", name); err == nil {
|
||||
if ns2, err := slurm.ParseNodeState(string(out)); err == nil {
|
||||
ns = ns2
|
||||
}
|
||||
}
|
||||
st.State, st.Reason = ns.State, ns.Reason
|
||||
if out, err := k.host.Run(ctx, env, "squeue", "-h", "-w", name, "-o", "%i"); err == nil {
|
||||
st.Jobs = int32(len(strings.Fields(string(out)))) //nolint:gosec // a count of jobs
|
||||
} else {
|
||||
st.Jobs = -1
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// clientEnv points scontrol at slurmd's copy of the files on a node that is
|
||||
// not the controller, which has none of its own.
|
||||
func (k *Keeper) clientEnv() []string {
|
||||
if _, err := k.host.ReadFile(filepath.Join(slurm.ConfDir, "slurm.conf")); err == nil {
|
||||
return nil
|
||||
}
|
||||
return []string{"SLURM_CONF=" + filepath.Join(slurm.SpoolDir, "conf-cache", "slurm.conf")}
|
||||
}
|
||||
|
||||
// stop stops what llmgw started here, if it did.
|
||||
func (k *Keeper) stop(ctx context.Context) {
|
||||
if _, err := k.host.ReadFile(k.marker); err != nil {
|
||||
return
|
||||
}
|
||||
k.log.Info("this node has no part in Slurm any more: slurmd and slurmctld stopped; the packages stay")
|
||||
for _, u := range []string{"slurmd", "slurmctld"} {
|
||||
_ = k.systemctl(ctx, "disable", "--now", u)
|
||||
}
|
||||
_ = k.host.Remove(DropIn)
|
||||
_ = k.systemctl(ctx, "daemon-reload")
|
||||
_ = k.host.Remove(k.marker)
|
||||
k.files = ""
|
||||
}
|
||||
|
||||
// write writes a file when it differs, saying whether it did.
|
||||
func (k *Keeper) write(path string, b []byte, mode fs.FileMode, owner string) (bool, error) {
|
||||
if old, err := k.host.ReadFile(path); err == nil && bytes.Equal(old, b) {
|
||||
return false, nil
|
||||
}
|
||||
return true, k.host.WriteFile(path, b, mode, owner)
|
||||
}
|
||||
|
||||
// systemctl runs systemctl; a last argument "&&restart" restarts the unit
|
||||
// after.
|
||||
func (k *Keeper) systemctl(ctx context.Context, args ...string) error {
|
||||
restart := false
|
||||
if n := len(args); n > 0 && args[n-1] == "&&restart" {
|
||||
args, restart = args[:n-1], true
|
||||
}
|
||||
if out, err := k.host.Run(ctx, nil, "systemctl", args...); err != nil {
|
||||
return fmt.Errorf("systemctl %s: %v: %s", strings.Join(args, " "), err, lastLine(out))
|
||||
}
|
||||
if restart {
|
||||
u := args[len(args)-1]
|
||||
if out, err := k.host.Run(ctx, nil, "systemctl", "restart", u); err != nil {
|
||||
return fmt.Errorf("systemctl restart %s: %v: %s", u, err, lastLine(out))
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (k *Keeper) active(ctx context.Context, unit string) bool {
|
||||
out, err := k.host.Run(ctx, nil, "systemctl", "is-active", unit)
|
||||
return err == nil && strings.TrimSpace(string(out)) == "active"
|
||||
}
|
||||
|
||||
func lastLine(b []byte) string {
|
||||
s := strings.TrimSpace(string(b))
|
||||
if i := strings.LastIndexByte(s, '\n'); i >= 0 {
|
||||
s = s[i+1:]
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// Local is the machine the agent runs on.
|
||||
type Local struct{}
|
||||
|
||||
// Run runs a command, with env added to the agent's.
|
||||
func (Local) Run(ctx context.Context, env []string, name string, args ...string) ([]byte, error) {
|
||||
cmd := exec.CommandContext(ctx, name, args...)
|
||||
if len(env) > 0 {
|
||||
cmd.Env = append(os.Environ(), env...)
|
||||
}
|
||||
return cmd.CombinedOutput()
|
||||
}
|
||||
|
||||
// LookPath finds a command, in the sbin directories too: an agent started
|
||||
// by systemd has them, one started by hand may not.
|
||||
func (Local) LookPath(name string) (string, error) {
|
||||
if p, err := exec.LookPath(name); err == nil {
|
||||
return p, nil
|
||||
}
|
||||
for _, d := range []string{"/usr/sbin", "/sbin", "/usr/local/sbin"} {
|
||||
p := filepath.Join(d, name)
|
||||
if fi, err := os.Stat(p); err == nil && !fi.IsDir() {
|
||||
return p, nil
|
||||
}
|
||||
}
|
||||
return "", fmt.Errorf("%s is not installed", name)
|
||||
}
|
||||
|
||||
// ReadFile reads a file.
|
||||
func (Local) ReadFile(path string) ([]byte, error) { return os.ReadFile(path) } //nolint:gosec // our paths
|
||||
|
||||
// WriteFile writes a file whole, through a temporary one beside it.
|
||||
func (Local) WriteFile(path string, b []byte, mode fs.FileMode, owner string) error {
|
||||
if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { //nolint:gosec // system directories
|
||||
return err
|
||||
}
|
||||
tmp := path + ".llmgw-new"
|
||||
if err := os.WriteFile(tmp, b, mode); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := os.Chmod(tmp, mode); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := chown(tmp, owner); err != nil {
|
||||
_ = os.Remove(tmp)
|
||||
return err
|
||||
}
|
||||
return os.Rename(tmp, path)
|
||||
}
|
||||
|
||||
// MkdirAll makes a directory owned by owner.
|
||||
func (Local) MkdirAll(path string, mode fs.FileMode, owner string) error {
|
||||
if err := os.MkdirAll(path, mode); err != nil {
|
||||
return err
|
||||
}
|
||||
return chown(path, owner)
|
||||
}
|
||||
|
||||
// Remove removes a file, and is content when there is none.
|
||||
func (Local) Remove(path string) error {
|
||||
if err := os.Remove(path); err != nil && !errors.Is(err, fs.ErrNotExist) {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func chown(path, owner string) error {
|
||||
if owner == "" {
|
||||
return nil
|
||||
}
|
||||
u, err := user.Lookup(owner)
|
||||
if err != nil {
|
||||
return fmt.Errorf("user %s: %w (installed with its package)", owner, err)
|
||||
}
|
||||
uid, _ := strconv.Atoi(u.Uid)
|
||||
gid, _ := strconv.Atoi(u.Gid)
|
||||
return os.Chown(path, uid, gid)
|
||||
}
|
||||
@@ -0,0 +1,306 @@
|
||||
package slurmnode
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io/fs"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
||||
"gitea.mynodes.xyz/ilker/llmgw/internal/slurm"
|
||||
)
|
||||
|
||||
type fakeHost struct {
|
||||
files map[string][]byte
|
||||
owners map[string]string
|
||||
bins map[string]bool
|
||||
active map[string]bool
|
||||
ran []string
|
||||
outputs map[string]string
|
||||
fail map[string]bool
|
||||
state string
|
||||
reason string
|
||||
}
|
||||
|
||||
func newFake() *fakeHost {
|
||||
return &fakeHost{files: map[string][]byte{}, owners: map[string]string{}, bins: map[string]bool{},
|
||||
active: map[string]bool{}, outputs: map[string]string{
|
||||
"scontrol --version": "slurm 23.11.4\n",
|
||||
"scontrol ping": "Slurmctld(primary) at node06 is UP\n",
|
||||
"slurmd -C": "NodeName=cuda01 CPUs=16 RealMemory=64000\nUpTime=1-00:00:00\n",
|
||||
}, fail: map[string]bool{}, state: "IDLE+DRAIN", reason: slurm.ReasonNew}
|
||||
}
|
||||
|
||||
func (f *fakeHost) Run(_ context.Context, env []string, name string, args ...string) ([]byte, error) {
|
||||
c := strings.Join(append([]string{name}, args...), " ")
|
||||
f.ran = append(f.ran, c)
|
||||
if f.fail[c] {
|
||||
return []byte("it failed"), errors.New("exit status 1")
|
||||
}
|
||||
switch {
|
||||
case name == "systemctl" && len(args) >= 1:
|
||||
u := args[len(args)-1]
|
||||
switch args[0] {
|
||||
case "is-active":
|
||||
if f.active[u] {
|
||||
return []byte("active\n"), nil
|
||||
}
|
||||
return []byte("inactive\n"), errors.New("exit status 3")
|
||||
case "restart", "start":
|
||||
f.active[u] = true
|
||||
case "enable":
|
||||
if slices.Contains(args, "--now") {
|
||||
f.active[u] = true
|
||||
}
|
||||
case "disable":
|
||||
f.active[u] = false
|
||||
}
|
||||
return nil, nil
|
||||
case strings.HasPrefix(c, "scontrol show node"):
|
||||
return []byte("NodeName=cuda01 Arch=x86_64 State=" + f.state + " Partitions=gpu Reason=" + f.reason + " [root@now]\n"), nil
|
||||
case strings.HasPrefix(c, "scontrol update"):
|
||||
if slices.Contains(args, "State=RESUME") {
|
||||
f.state, f.reason = "IDLE", ""
|
||||
} else {
|
||||
f.state, f.reason = "IDLE+DRAIN", ReasonServing
|
||||
}
|
||||
return nil, nil
|
||||
case strings.HasPrefix(c, "squeue"):
|
||||
return []byte(f.outputs["squeue"]), nil
|
||||
}
|
||||
return []byte(f.outputs[c]), nil
|
||||
}
|
||||
|
||||
func (f *fakeHost) LookPath(name string) (string, error) {
|
||||
if f.bins[name] {
|
||||
return "/usr/sbin/" + name, nil
|
||||
}
|
||||
return "", errors.New("not found")
|
||||
}
|
||||
|
||||
func (f *fakeHost) ReadFile(path string) ([]byte, error) {
|
||||
b, ok := f.files[path]
|
||||
if !ok {
|
||||
return nil, fs.ErrNotExist
|
||||
}
|
||||
return b, nil
|
||||
}
|
||||
|
||||
func (f *fakeHost) WriteFile(path string, b []byte, _ fs.FileMode, owner string) error {
|
||||
f.files[path] = append([]byte(nil), b...)
|
||||
f.owners[path] = owner
|
||||
return nil
|
||||
}
|
||||
|
||||
func (f *fakeHost) MkdirAll(string, fs.FileMode, string) error { return nil }
|
||||
|
||||
func (f *fakeHost) Remove(path string) error { delete(f.files, path); return nil }
|
||||
|
||||
func (f *fakeHost) count(prefix string) int {
|
||||
n := 0
|
||||
for _, c := range f.ran {
|
||||
if strings.HasPrefix(c, prefix) {
|
||||
n++
|
||||
}
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
type fakeInst struct {
|
||||
host *fakeHost
|
||||
got map[string][]string
|
||||
}
|
||||
|
||||
func (i *fakeInst) InstallPackages(_ context.Context, by map[string][]string) error {
|
||||
i.got = by
|
||||
for _, b := range []string{"munged", "scontrol", "slurmd", "slurmctld"} {
|
||||
i.host.bins[b] = true
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
var key = []byte(strings.Repeat("k", 1024))
|
||||
|
||||
func nodeMsg(lent bool) *llmgwv1.Slurm {
|
||||
return &llmgwv1.Slurm{Node: true, Name: "cuda01", ConfServer: "10.0.0.146:6817", MungeKey: key, Lent: lent}
|
||||
}
|
||||
|
||||
func TestNodeInstallsAndStaysDrained(t *testing.T) {
|
||||
h := newFake()
|
||||
inst := &fakeInst{host: h}
|
||||
k := New(h, inst, "/var/lib/llmgw-agent/slurm", nil)
|
||||
st := k.Reconcile(context.Background(), nodeMsg(false))
|
||||
if st.GetError() != "" {
|
||||
t.Fatal(st.GetError())
|
||||
}
|
||||
if !slices.Equal(inst.got["deb"], []string{"munge", "slurm-client", "slurmd"}) ||
|
||||
!slices.Equal(inst.got["rpm"], []string{"munge", "slurm", "slurm-slurmd"}) {
|
||||
t.Fatalf("installed %v", inst.got)
|
||||
}
|
||||
if string(h.files[slurm.MungeKeyPath]) != string(key) || h.owners[slurm.MungeKeyPath] != "munge" {
|
||||
t.Fatal("the munge key was not written for munge")
|
||||
}
|
||||
unit := string(h.files[DropIn])
|
||||
if !strings.Contains(unit, "ExecStart=/usr/sbin/slurmd -D -s --conf-server 10.0.0.146:6817 -N cuda01\n") ||
|
||||
!strings.Contains(unit, "ConditionPathExists=\n") {
|
||||
t.Fatalf("drop-in:\n%s", unit)
|
||||
}
|
||||
if !st.GetSlurmdRunning() || st.GetVersion() != "23.11.4" || !strings.HasPrefix(st.GetHardware(), "NodeName=cuda01 CPUs=16") {
|
||||
t.Fatalf("status %v", st)
|
||||
}
|
||||
// Not lent: left drained, with llmgw's first reason.
|
||||
if h.count("scontrol update") != 0 || st.GetState() != "IDLE+DRAIN" {
|
||||
t.Fatalf("a node not lent was changed: %v, %v", h.ran, st)
|
||||
}
|
||||
if _, ok := h.files["/var/lib/llmgw-agent/slurm"]; !ok {
|
||||
t.Fatal("no marker")
|
||||
}
|
||||
|
||||
// Again, nothing changed: no install, no restart.
|
||||
h.ran = nil
|
||||
inst.got = nil
|
||||
k.Reconcile(context.Background(), nodeMsg(false))
|
||||
if inst.got != nil || h.count("systemctl restart") != 0 || h.count("systemctl enable") != 0 {
|
||||
t.Fatalf("a second pass did work: %v", h.ran)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLendAndReclaim(t *testing.T) {
|
||||
h := newFake()
|
||||
k := New(h, &fakeInst{host: h}, "/m", nil)
|
||||
st := k.Reconcile(context.Background(), nodeMsg(true))
|
||||
if h.count("scontrol update NodeName=cuda01 State=RESUME") != 1 || st.GetState() != "IDLE" {
|
||||
t.Fatalf("a lent node was not resumed: %v %v", h.ran, st)
|
||||
}
|
||||
h.outputs["squeue"] = "101\n102\n"
|
||||
st = k.Reconcile(context.Background(), nodeMsg(false))
|
||||
if h.count("scontrol update NodeName=cuda01 State=DRAIN Reason="+ReasonServing) != 1 {
|
||||
t.Fatalf("a node taken back was not drained: %v", h.ran)
|
||||
}
|
||||
if st.GetJobs() != 2 || !strings.Contains(st.GetState(), "DRAIN") {
|
||||
t.Fatalf("status %v", st)
|
||||
}
|
||||
// Drained already: not drained again.
|
||||
h.ran = nil
|
||||
k.Reconcile(context.Background(), nodeMsg(false))
|
||||
if h.count("scontrol update") != 0 {
|
||||
t.Fatalf("drained twice: %v", h.ran)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdminsDrainIsLeft(t *testing.T) {
|
||||
h := newFake()
|
||||
h.reason = "bad DIMM, ticket 42"
|
||||
k := New(h, &fakeInst{host: h}, "/m", nil)
|
||||
k.Reconcile(context.Background(), nodeMsg(true))
|
||||
if h.count("scontrol update") != 0 {
|
||||
t.Fatalf("llmgw resumed a node a Slurm administrator drained: %v", h.ran)
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientEnvOnANodeWithoutFiles(t *testing.T) {
|
||||
h := newFake()
|
||||
k := New(h, &fakeInst{host: h}, "/m", nil)
|
||||
if env := k.clientEnv(); len(env) != 1 || !strings.HasSuffix(env[0], "conf-cache/slurm.conf") {
|
||||
t.Fatalf("env %v", env)
|
||||
}
|
||||
h.files[filepath.Join(slurm.ConfDir, "slurm.conf")] = []byte("x")
|
||||
if env := k.clientEnv(); env != nil {
|
||||
t.Fatalf("env %v on the controller", env)
|
||||
}
|
||||
}
|
||||
|
||||
func TestController(t *testing.T) {
|
||||
h := newFake()
|
||||
inst := &fakeInst{host: h}
|
||||
k := New(h, inst, "/m", nil)
|
||||
files := map[string]string{"slurm.conf": "ClusterName=llmgw\n", "gres.conf": "#\n", "cgroup.conf": "#\n"}
|
||||
m := &llmgwv1.Slurm{Controller: true, MungeKey: key, Files: files}
|
||||
st := k.Reconcile(context.Background(), m)
|
||||
if st.GetError() != "" {
|
||||
t.Fatal(st.GetError())
|
||||
}
|
||||
if !slices.Contains(inst.got["deb"], "slurmctld") || slices.Contains(inst.got["deb"], "slurmd") {
|
||||
t.Fatalf("installed %v", inst.got)
|
||||
}
|
||||
if string(h.files["/etc/slurm/slurm.conf"]) != "ClusterName=llmgw\n" {
|
||||
t.Fatal("slurm.conf not written")
|
||||
}
|
||||
if h.count("systemctl restart slurmctld") != 1 || h.count("scontrol reconfigure") != 1 {
|
||||
t.Fatalf("slurmctld not started on its files: %v", h.ran)
|
||||
}
|
||||
if !st.GetSlurmctldRunning() || len(st.GetFilesSha256()) != 64 {
|
||||
t.Fatalf("status %v", st)
|
||||
}
|
||||
sum := st.GetFilesSha256()
|
||||
|
||||
h.ran = nil
|
||||
k.Reconcile(context.Background(), m)
|
||||
if h.count("systemctl restart") != 0 {
|
||||
t.Fatalf("restarted on the same files: %v", h.ran)
|
||||
}
|
||||
|
||||
files["gres.conf"] = "NodeName=cuda01 Name=gpu File=/dev/nvidia0\n"
|
||||
h.ran = nil
|
||||
st = k.Reconcile(context.Background(), m)
|
||||
if h.count("systemctl restart slurmctld") != 1 || st.GetFilesSha256() == sum {
|
||||
t.Fatalf("a changed file did not restart slurmctld: %v", h.ran)
|
||||
}
|
||||
|
||||
if st := k.Reconcile(context.Background(), &llmgwv1.Slurm{Controller: true, MungeKey: key,
|
||||
Files: map[string]string{"slurm.conf": "x", "../evil": "y"}}); !strings.Contains(st.GetError(), "named") {
|
||||
t.Fatalf("a path in a file name was taken: %v", st)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTooOld(t *testing.T) {
|
||||
h := newFake()
|
||||
h.outputs["scontrol --version"] = "slurm-wlm 21.08.5\n"
|
||||
k := New(h, &fakeInst{host: h}, "/m", nil)
|
||||
st := k.Reconcile(context.Background(), nodeMsg(false))
|
||||
if !strings.Contains(st.GetError(), "older than 22.05") {
|
||||
t.Fatalf("status %v", st)
|
||||
}
|
||||
if _, ok := h.files[slurm.MungeKeyPath]; ok {
|
||||
t.Fatal("went on with a Slurm too old")
|
||||
}
|
||||
}
|
||||
|
||||
func TestStop(t *testing.T) {
|
||||
h := newFake()
|
||||
k := New(h, &fakeInst{host: h}, "/m", nil)
|
||||
k.Reconcile(context.Background(), nodeMsg(false))
|
||||
h.ran = nil
|
||||
if st := k.Reconcile(context.Background(), &llmgwv1.Slurm{}); st != nil {
|
||||
t.Fatalf("status %v for no part", st)
|
||||
}
|
||||
if h.count("systemctl disable --now slurmd") != 1 {
|
||||
t.Fatalf("slurmd not stopped: %v", h.ran)
|
||||
}
|
||||
if _, ok := h.files[DropIn]; ok {
|
||||
t.Fatal("drop-in left")
|
||||
}
|
||||
// Never started here: nothing touched.
|
||||
h2 := newFake()
|
||||
k2 := New(h2, &fakeInst{host: h2}, "/m", nil)
|
||||
k2.Reconcile(context.Background(), &llmgwv1.Slurm{})
|
||||
if len(h2.ran) != 0 {
|
||||
t.Fatalf("a machine llmgw never ran Slurm on was touched: %v", h2.ran)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBadNames(t *testing.T) {
|
||||
h := newFake()
|
||||
k := New(h, &fakeInst{host: h}, "/m", nil)
|
||||
m := nodeMsg(false)
|
||||
m.Name = "cuda01 -c /etc/evil"
|
||||
if st := k.Reconcile(context.Background(), m); !strings.Contains(st.GetError(), "spaces") {
|
||||
t.Fatalf("status %v", st)
|
||||
}
|
||||
if st := k.Reconcile(context.Background(), &llmgwv1.Slurm{Node: true, Name: "a", ConfServer: "b"}); !strings.Contains(st.GetError(), "no key") {
|
||||
t.Fatalf("status %v", st)
|
||||
}
|
||||
}
|
||||
@@ -29,6 +29,10 @@ const (
|
||||
CtldPort = 6817
|
||||
)
|
||||
|
||||
// ReasonNew is why a node new to slurm.conf starts drained. Every reason
|
||||
// llmgw drains with starts with "llmgw", and only those are resumed by it.
|
||||
const ReasonNew = "llmgw: not lent yet"
|
||||
|
||||
// MinVersion is the oldest Slurm llmgw runs: the first to put jobs in
|
||||
// cgroup v2, which every distribution llmgw supports now mounts alone.
|
||||
var MinVersion = [2]int{22, 5}
|
||||
@@ -226,7 +230,9 @@ SlurmdTimeout=300
|
||||
}
|
||||
fmt.Fprintf(&b, " Gres=%s", strings.Join(gres, ","))
|
||||
}
|
||||
b.WriteString(" State=UNKNOWN\n")
|
||||
// Drained until llmgw lends it: a node new to the controller would
|
||||
// otherwise take jobs in the minute before its agent drains it.
|
||||
fmt.Fprintf(&b, " State=DRAIN Reason=%q\n", ReasonNew)
|
||||
}
|
||||
if len(names) > 0 {
|
||||
fmt.Fprintf(&b, "PartitionName=%s Nodes=%s Default=YES MaxTime=INFINITE State=UP\n",
|
||||
|
||||
@@ -79,8 +79,8 @@ func TestSlurmConf(t *testing.T) {
|
||||
"ClusterName=llmgw\n",
|
||||
"SlurmctldHost=node06(10.0.0.146)\n",
|
||||
"SlurmctldParameters=enable_configless\n",
|
||||
"NodeName=cuda01 NodeAddr=10.0.0.147 CPUs=16 RealMemory=62080 Gres=gpu:rtx3090:1 State=UNKNOWN\n",
|
||||
"NodeName=site2 NodeAddr=10.1.0.5 CPUs=32 Boards=1 SocketsPerBoard=2 CoresPerSocket=8 ThreadsPerCore=2 RealMemory=97000 Gres=gpu:t4:1,gpu:v100:2 State=UNKNOWN\n",
|
||||
"NodeName=cuda01 NodeAddr=10.0.0.147 CPUs=16 RealMemory=62080 Gres=gpu:rtx3090:1 State=DRAIN Reason=\"llmgw: not lent yet\"\n",
|
||||
"NodeName=site2 NodeAddr=10.1.0.5 CPUs=32 Boards=1 SocketsPerBoard=2 CoresPerSocket=8 ThreadsPerCore=2 RealMemory=97000 Gres=gpu:t4:1,gpu:v100:2 State=DRAIN Reason=\"llmgw: not lent yet\"\n",
|
||||
"PartitionName=gpu Nodes=cuda01,site2 Default=YES",
|
||||
} {
|
||||
if !strings.Contains(conf, want) {
|
||||
|
||||
Reference in new issue
Block a user