diff --git a/api/gen/llmgw/v1/agent.pb.go b/api/gen/llmgw/v1/agent.pb.go index 70eafcb..d9f42a6 100644 --- a/api/gen/llmgw/v1/agent.pb.go +++ b/api/gen/llmgw/v1/agent.pb.go @@ -106,6 +106,8 @@ type AgentMessage struct { // *AgentMessage_FabricTest // *AgentMessage_Upgrade // *AgentMessage_KvEvents + // *AgentMessage_TrainStatus + // *AgentMessage_TrainArtifact Msg isAgentMessage_Msg `protobuf_oneof:"msg"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache @@ -220,6 +222,24 @@ func (x *AgentMessage) GetKvEvents() *KVEvents { return nil } +func (x *AgentMessage) GetTrainStatus() *TrainStatus { + if x != nil { + if x, ok := x.Msg.(*AgentMessage_TrainStatus); ok { + return x.TrainStatus + } + } + return nil +} + +func (x *AgentMessage) GetTrainArtifact() *TrainArtifact { + if x != nil { + if x, ok := x.Msg.(*AgentMessage_TrainArtifact); ok { + return x.TrainArtifact + } + } + return nil +} + type isAgentMessage_Msg interface { isAgentMessage_Msg() } @@ -256,6 +276,14 @@ type AgentMessage_KvEvents struct { KvEvents *KVEvents `protobuf:"bytes,8,opt,name=kv_events,json=kvEvents,proto3,oneof"` } +type AgentMessage_TrainStatus struct { + TrainStatus *TrainStatus `protobuf:"bytes,9,opt,name=train_status,json=trainStatus,proto3,oneof"` +} + +type AgentMessage_TrainArtifact struct { + TrainArtifact *TrainArtifact `protobuf:"bytes,10,opt,name=train_artifact,json=trainArtifact,proto3,oneof"` +} + func (*AgentMessage_Register) isAgentMessage_Msg() {} func (*AgentMessage_Heartbeat) isAgentMessage_Msg() {} @@ -272,6 +300,10 @@ func (*AgentMessage_Upgrade) isAgentMessage_Msg() {} func (*AgentMessage_KvEvents) isAgentMessage_Msg() {} +func (*AgentMessage_TrainStatus) isAgentMessage_Msg() {} + +func (*AgentMessage_TrainArtifact) 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). @@ -392,6 +424,8 @@ type ControlMessage struct { // *ControlMessage_Firewall // *ControlMessage_RestartAgent // *ControlMessage_AgentSettings + // *ControlMessage_TrainStart + // *ControlMessage_TrainCancel 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 @@ -566,6 +600,24 @@ func (x *ControlMessage) GetAgentSettings() *AgentSettings { return nil } +func (x *ControlMessage) GetTrainStart() *TrainStart { + if x != nil { + if x, ok := x.Msg.(*ControlMessage_TrainStart); ok { + return x.TrainStart + } + } + return nil +} + +func (x *ControlMessage) GetTrainCancel() *TrainCancel { + if x != nil { + if x, ok := x.Msg.(*ControlMessage_TrainCancel); ok { + return x.TrainCancel + } + } + return nil +} + func (x *ControlMessage) GetTrace() map[string]string { if x != nil { return x.Trace @@ -633,6 +685,14 @@ type ControlMessage_AgentSettings struct { AgentSettings *AgentSettings `protobuf:"bytes,15,opt,name=agent_settings,json=agentSettings,proto3,oneof"` } +type ControlMessage_TrainStart struct { + TrainStart *TrainStart `protobuf:"bytes,16,opt,name=train_start,json=trainStart,proto3,oneof"` +} + +type ControlMessage_TrainCancel struct { + TrainCancel *TrainCancel `protobuf:"bytes,17,opt,name=train_cancel,json=trainCancel,proto3,oneof"` +} + func (*ControlMessage_Start) isControlMessage_Msg() {} func (*ControlMessage_Stop) isControlMessage_Msg() {} @@ -661,6 +721,10 @@ func (*ControlMessage_RestartAgent) isControlMessage_Msg() {} func (*ControlMessage_AgentSettings) isControlMessage_Msg() {} +func (*ControlMessage_TrainStart) isControlMessage_Msg() {} + +func (*ControlMessage_TrainCancel) 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 { @@ -4699,11 +4763,335 @@ func (x *LoraAdapter) GetSource() string { return "" } +// TrainStart has a node train a LoRA adapter (BACKLOG F4): on a base model's +// weights from its registry, on question/answer pairs, with a Python trainer +// the agent carries, in an environment it makes over an installed engine's. +// One at a time per node. +type TrainStart struct { + state protoimpl.MessageState `protogen:"open.v1"` + JobId string `protobuf:"bytes,1,opt,name=job_id,json=jobId,proto3" json:"job_id,omitempty"` + // base_repo and base_revision name the base model's weights on its + // registry (Hugging Face): safetensors, not the GGUF a catalog may serve. + BaseRepo string `protobuf:"bytes,2,opt,name=base_repo,json=baseRepo,proto3" json:"base_repo,omitempty"` + BaseRevision string `protobuf:"bytes,3,opt,name=base_revision,json=baseRevision,proto3" json:"base_revision,omitempty"` + // data is the training pairs, NDJSON of {"q", "a"}. + Data []byte `protobuf:"bytes,4,opt,name=data,proto3" json:"data,omitempty"` + // params is the trainer's settings as JSON: rank, alpha, epochs, lr, + // batch, max_len. + Params string `protobuf:"bytes,5,opt,name=params,proto3" json:"params,omitempty"` + // gpu_memory_mib is what the job needs of a card; the agent starts it + // only when one has that much free, which lets it run beside an engine. + GpuMemoryMib int32 `protobuf:"varint,6,opt,name=gpu_memory_mib,json=gpuMemoryMib,proto3" json:"gpu_memory_mib,omitempty"` + // gguf asks for the adapter in llama.cpp's format too. + Gguf bool `protobuf:"varint,7,opt,name=gguf,proto3" json:"gguf,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *TrainStart) Reset() { + *x = TrainStart{} + mi := &file_llmgw_v1_agent_proto_msgTypes[44] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *TrainStart) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*TrainStart) ProtoMessage() {} + +func (x *TrainStart) ProtoReflect() protoreflect.Message { + mi := &file_llmgw_v1_agent_proto_msgTypes[44] + 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 TrainStart.ProtoReflect.Descriptor instead. +func (*TrainStart) Descriptor() ([]byte, []int) { + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{44} +} + +func (x *TrainStart) GetJobId() string { + if x != nil { + return x.JobId + } + return "" +} + +func (x *TrainStart) GetBaseRepo() string { + if x != nil { + return x.BaseRepo + } + return "" +} + +func (x *TrainStart) GetBaseRevision() string { + if x != nil { + return x.BaseRevision + } + return "" +} + +func (x *TrainStart) GetData() []byte { + if x != nil { + return x.Data + } + return nil +} + +func (x *TrainStart) GetParams() string { + if x != nil { + return x.Params + } + return "" +} + +func (x *TrainStart) GetGpuMemoryMib() int32 { + if x != nil { + return x.GpuMemoryMib + } + return 0 +} + +func (x *TrainStart) GetGguf() bool { + if x != nil { + return x.Gguf + } + return false +} + +// TrainCancel stops a training job. +type TrainCancel struct { + state protoimpl.MessageState `protogen:"open.v1"` + JobId string `protobuf:"bytes,1,opt,name=job_id,json=jobId,proto3" json:"job_id,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *TrainCancel) Reset() { + *x = TrainCancel{} + mi := &file_llmgw_v1_agent_proto_msgTypes[45] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *TrainCancel) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*TrainCancel) ProtoMessage() {} + +func (x *TrainCancel) ProtoReflect() protoreflect.Message { + mi := &file_llmgw_v1_agent_proto_msgTypes[45] + 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 TrainCancel.ProtoReflect.Descriptor instead. +func (*TrainCancel) Descriptor() ([]byte, []int) { + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{45} +} + +func (x *TrainCancel) GetJobId() string { + if x != nil { + return x.JobId + } + return "" +} + +// TrainStatus is how a training job is going, sent as it changes. +type TrainStatus struct { + state protoimpl.MessageState `protogen:"open.v1"` + JobId string `protobuf:"bytes,1,opt,name=job_id,json=jobId,proto3" json:"job_id,omitempty"` + // state is preparing, running, done or failed. + State string `protobuf:"bytes,2,opt,name=state,proto3" json:"state,omitempty"` + Step int32 `protobuf:"varint,3,opt,name=step,proto3" json:"step,omitempty"` + Steps int32 `protobuf:"varint,4,opt,name=steps,proto3" json:"steps,omitempty"` + Loss float64 `protobuf:"fixed64,5,opt,name=loss,proto3" json:"loss,omitempty"` + // message is what is being done, or why it failed. + Message string `protobuf:"bytes,6,opt,name=message,proto3" json:"message,omitempty"` + // peak_mib is the most the job held of the card. + PeakMib int32 `protobuf:"varint,7,opt,name=peak_mib,json=peakMib,proto3" json:"peak_mib,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *TrainStatus) Reset() { + *x = TrainStatus{} + mi := &file_llmgw_v1_agent_proto_msgTypes[46] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *TrainStatus) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*TrainStatus) ProtoMessage() {} + +func (x *TrainStatus) ProtoReflect() protoreflect.Message { + mi := &file_llmgw_v1_agent_proto_msgTypes[46] + 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 TrainStatus.ProtoReflect.Descriptor instead. +func (*TrainStatus) Descriptor() ([]byte, []int) { + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{46} +} + +func (x *TrainStatus) GetJobId() string { + if x != nil { + return x.JobId + } + return "" +} + +func (x *TrainStatus) GetState() string { + if x != nil { + return x.State + } + return "" +} + +func (x *TrainStatus) GetStep() int32 { + if x != nil { + return x.Step + } + return 0 +} + +func (x *TrainStatus) GetSteps() int32 { + if x != nil { + return x.Steps + } + return 0 +} + +func (x *TrainStatus) GetLoss() float64 { + if x != nil { + return x.Loss + } + return 0 +} + +func (x *TrainStatus) GetMessage() string { + if x != nil { + return x.Message + } + return "" +} + +func (x *TrainStatus) GetPeakMib() int32 { + if x != nil { + return x.PeakMib + } + return 0 +} + +// TrainArtifact is a piece of what a training job made, sent before its +// final status: a gzipped tar of the PEFT adapter (kind peft), or the GGUF +// adapter (kind gguf). +type TrainArtifact struct { + state protoimpl.MessageState `protogen:"open.v1"` + JobId string `protobuf:"bytes,1,opt,name=job_id,json=jobId,proto3" json:"job_id,omitempty"` + Kind string `protobuf:"bytes,2,opt,name=kind,proto3" json:"kind,omitempty"` + Seq int32 `protobuf:"varint,3,opt,name=seq,proto3" json:"seq,omitempty"` + Data []byte `protobuf:"bytes,4,opt,name=data,proto3" json:"data,omitempty"` + Last bool `protobuf:"varint,5,opt,name=last,proto3" json:"last,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *TrainArtifact) Reset() { + *x = TrainArtifact{} + mi := &file_llmgw_v1_agent_proto_msgTypes[47] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *TrainArtifact) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*TrainArtifact) ProtoMessage() {} + +func (x *TrainArtifact) ProtoReflect() protoreflect.Message { + mi := &file_llmgw_v1_agent_proto_msgTypes[47] + 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 TrainArtifact.ProtoReflect.Descriptor instead. +func (*TrainArtifact) Descriptor() ([]byte, []int) { + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{47} +} + +func (x *TrainArtifact) GetJobId() string { + if x != nil { + return x.JobId + } + return "" +} + +func (x *TrainArtifact) GetKind() string { + if x != nil { + return x.Kind + } + return "" +} + +func (x *TrainArtifact) GetSeq() int32 { + if x != nil { + return x.Seq + } + return 0 +} + +func (x *TrainArtifact) GetData() []byte { + if x != nil { + return x.Data + } + return nil +} + +func (x *TrainArtifact) GetLast() bool { + if x != nil { + return x.Last + } + return false +} + var File_llmgw_v1_agent_proto protoreflect.FileDescriptor const file_llmgw_v1_agent_proto_rawDesc = "" + "\n" + - "\x14llmgw/v1/agent.proto\x12\bllmgw.v1\"\xb9\x03\n" + + "\x14llmgw/v1/agent.proto\x12\bllmgw.v1\"\xb7\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" + @@ -4714,7 +5102,10 @@ const file_llmgw_v1_agent_proto_rawDesc = "" + "\vfabric_test\x18\x06 \x01(\v2\x19.llmgw.v1.FabricTestReplyH\x00R\n" + "fabricTest\x122\n" + "\aupgrade\x18\a \x01(\v2\x16.llmgw.v1.UpgradeReplyH\x00R\aupgrade\x121\n" + - "\tkv_events\x18\b \x01(\v2\x12.llmgw.v1.KVEventsH\x00R\bkvEventsB\x05\n" + + "\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" + "\x03msg\"\xb5\x01\n" + "\bKVEvents\x12\x1d\n" + "\n" + @@ -4725,7 +5116,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\"\x83\a\n" + + "\aproblem\x18\x06 \x01(\tR\aproblem\"\xf8\a\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" + @@ -4743,7 +5134,10 @@ const file_llmgw_v1_agent_proto_rawDesc = "" + "\x0econtrol_planes\x18\v \x01(\v2\x17.llmgw.v1.ControlPlanesH\x00R\rcontrolPlanes\x120\n" + "\bfirewall\x18\r \x01(\v2\x12.llmgw.v1.FirewallH\x00R\bfirewall\x12=\n" + "\rrestart_agent\x18\x0e \x01(\v2\x16.llmgw.v1.RestartAgentH\x00R\frestartAgent\x12@\n" + - "\x0eagent_settings\x18\x0f \x01(\v2\x17.llmgw.v1.AgentSettingsH\x00R\ragentSettings\x129\n" + + "\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" + "\x05trace\x18\f \x03(\v2#.llmgw.v1.ControlMessage.TraceEntryR\x05trace\x1a8\n" + "\n" + "TraceEntry\x12\x10\n" + @@ -5150,7 +5544,32 @@ const file_llmgw_v1_agent_proto_rawDesc = "" + "\x05error\x18\x03 \x01(\tR\x05error\"9\n" + "\vLoraAdapter\x12\x12\n" + "\x04name\x18\x01 \x01(\tR\x04name\x12\x16\n" + - "\x06source\x18\x02 \x01(\tR\x06source2O\n" + + "\x06source\x18\x02 \x01(\tR\x06source\"\xcb\x01\n" + + "\n" + + "TrainStart\x12\x15\n" + + "\x06job_id\x18\x01 \x01(\tR\x05jobId\x12\x1b\n" + + "\tbase_repo\x18\x02 \x01(\tR\bbaseRepo\x12#\n" + + "\rbase_revision\x18\x03 \x01(\tR\fbaseRevision\x12\x12\n" + + "\x04data\x18\x04 \x01(\fR\x04data\x12\x16\n" + + "\x06params\x18\x05 \x01(\tR\x06params\x12$\n" + + "\x0egpu_memory_mib\x18\x06 \x01(\x05R\fgpuMemoryMib\x12\x12\n" + + "\x04gguf\x18\a \x01(\bR\x04gguf\"$\n" + + "\vTrainCancel\x12\x15\n" + + "\x06job_id\x18\x01 \x01(\tR\x05jobId\"\xad\x01\n" + + "\vTrainStatus\x12\x15\n" + + "\x06job_id\x18\x01 \x01(\tR\x05jobId\x12\x14\n" + + "\x05state\x18\x02 \x01(\tR\x05state\x12\x12\n" + + "\x04step\x18\x03 \x01(\x05R\x04step\x12\x14\n" + + "\x05steps\x18\x04 \x01(\x05R\x05steps\x12\x12\n" + + "\x04loss\x18\x05 \x01(\x01R\x04loss\x12\x18\n" + + "\amessage\x18\x06 \x01(\tR\amessage\x12\x19\n" + + "\bpeak_mib\x18\a \x01(\x05R\apeakMib\"t\n" + + "\rTrainArtifact\x12\x15\n" + + "\x06job_id\x18\x01 \x01(\tR\x05jobId\x12\x12\n" + + "\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" + "\fAgentService\x12?\n" + "\aConnect\x12\x16.llmgw.v1.AgentMessage\x1a\x18.llmgw.v1.ControlMessage(\x010\x01B\x93\x01\n" + "\fcom.llmgw.v1B\n" + @@ -5169,7 +5588,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, 52) +var file_llmgw_v1_agent_proto_msgTypes = make([]protoimpl.MessageInfo, 56) var file_llmgw_v1_agent_proto_goTypes = []any{ (ShardStatus_State)(0), // 0: llmgw.v1.ShardStatus.State (*AgentMessage)(nil), // 1: llmgw.v1.AgentMessage @@ -5216,14 +5635,18 @@ var file_llmgw_v1_agent_proto_goTypes = []any{ (*UpgradeAgent)(nil), // 42: llmgw.v1.UpgradeAgent (*UpgradeReply)(nil), // 43: llmgw.v1.UpgradeReply (*LoraAdapter)(nil), // 44: llmgw.v1.LoraAdapter - nil, // 45: llmgw.v1.ControlMessage.TraceEntry - nil, // 46: llmgw.v1.Register.LabelsEntry - nil, // 47: llmgw.v1.Register.EngineVersionsEntry - nil, // 48: llmgw.v1.GPU.NicPathsEntry - nil, // 49: llmgw.v1.IBPort.CountersEntry - nil, // 50: llmgw.v1.Heartbeat.EngineVersionsEntry - nil, // 51: llmgw.v1.EngineWant.OfflineSha256Entry - nil, // 52: llmgw.v1.EngineProfile.EnvEntry + (*TrainStart)(nil), // 45: llmgw.v1.TrainStart + (*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 } var file_llmgw_v1_agent_proto_depIdxs = []int32{ 8, // 0: llmgw.v1.AgentMessage.register:type_name -> llmgw.v1.Register @@ -5234,62 +5657,66 @@ var file_llmgw_v1_agent_proto_depIdxs = []int32{ 41, // 5: llmgw.v1.AgentMessage.fabric_test:type_name -> llmgw.v1.FabricTestReply 43, // 6: llmgw.v1.AgentMessage.upgrade:type_name -> llmgw.v1.UpgradeReply 2, // 7: llmgw.v1.AgentMessage.kv_events:type_name -> llmgw.v1.KVEvents - 30, // 8: llmgw.v1.ControlMessage.start:type_name -> llmgw.v1.StartShard - 33, // 9: llmgw.v1.ControlMessage.stop:type_name -> llmgw.v1.StopShard - 34, // 10: llmgw.v1.ControlMessage.drain:type_name -> llmgw.v1.DrainShard - 35, // 11: llmgw.v1.ControlMessage.ack:type_name -> llmgw.v1.Ack - 32, // 12: llmgw.v1.ControlMessage.remove_weights:type_name -> llmgw.v1.RemoveWeights - 24, // 13: llmgw.v1.ControlMessage.provision:type_name -> llmgw.v1.Provision - 36, // 14: llmgw.v1.ControlMessage.logs:type_name -> llmgw.v1.LogsRequest - 38, // 15: llmgw.v1.ControlMessage.fabric_scan:type_name -> llmgw.v1.FabricScanRequest - 40, // 16: llmgw.v1.ControlMessage.fabric_test:type_name -> llmgw.v1.FabricTestRequest - 42, // 17: llmgw.v1.ControlMessage.upgrade:type_name -> llmgw.v1.UpgradeAgent - 7, // 18: llmgw.v1.ControlMessage.control_planes:type_name -> llmgw.v1.ControlPlanes - 5, // 19: llmgw.v1.ControlMessage.firewall:type_name -> llmgw.v1.Firewall - 4, // 20: llmgw.v1.ControlMessage.restart_agent:type_name -> llmgw.v1.RestartAgent - 6, // 21: llmgw.v1.ControlMessage.agent_settings:type_name -> llmgw.v1.AgentSettings - 45, // 22: llmgw.v1.ControlMessage.trace:type_name -> llmgw.v1.ControlMessage.TraceEntry - 46, // 23: llmgw.v1.Register.labels:type_name -> llmgw.v1.Register.LabelsEntry - 12, // 24: llmgw.v1.Register.gpus:type_name -> llmgw.v1.GPU - 14, // 25: llmgw.v1.Register.fabrics:type_name -> llmgw.v1.FabricAttachment - 47, // 26: llmgw.v1.Register.engine_versions:type_name -> llmgw.v1.Register.EngineVersionsEntry - 11, // 27: llmgw.v1.Register.running_shards:type_name -> llmgw.v1.RunningShard - 15, // 28: llmgw.v1.Register.weights:type_name -> llmgw.v1.LocalWeights - 10, // 29: llmgw.v1.Register.storage:type_name -> llmgw.v1.Storage - 9, // 30: llmgw.v1.Register.host:type_name -> llmgw.v1.Host - 21, // 31: llmgw.v1.Register.imex:type_name -> llmgw.v1.IMEX - 0, // 32: llmgw.v1.RunningShard.state:type_name -> llmgw.v1.ShardStatus.State - 27, // 33: llmgw.v1.RunningShard.compute:type_name -> llmgw.v1.Compute - 48, // 34: llmgw.v1.GPU.nic_paths:type_name -> llmgw.v1.GPU.NicPathsEntry - 13, // 35: llmgw.v1.GPU.mig:type_name -> llmgw.v1.MIGDevice - 18, // 36: llmgw.v1.NodeLoad.shards:type_name -> llmgw.v1.ShardLoad - 17, // 37: llmgw.v1.NodeLoad.gpus:type_name -> llmgw.v1.GPUSensor - 49, // 38: llmgw.v1.IBPort.counters:type_name -> llmgw.v1.IBPort.CountersEntry - 20, // 39: llmgw.v1.IBPort.subnet_managers:type_name -> llmgw.v1.SubnetManager - 10, // 40: llmgw.v1.Heartbeat.storage:type_name -> llmgw.v1.Storage - 15, // 41: llmgw.v1.Heartbeat.weights:type_name -> llmgw.v1.LocalWeights - 50, // 42: llmgw.v1.Heartbeat.engine_versions:type_name -> llmgw.v1.Heartbeat.EngineVersionsEntry - 26, // 43: llmgw.v1.Heartbeat.provision:type_name -> llmgw.v1.ProvisionStatus - 16, // 44: llmgw.v1.Heartbeat.load:type_name -> llmgw.v1.NodeLoad - 19, // 45: llmgw.v1.Heartbeat.ib_ports:type_name -> llmgw.v1.IBPort - 14, // 46: llmgw.v1.Heartbeat.fabrics:type_name -> llmgw.v1.FabricAttachment - 22, // 47: llmgw.v1.Heartbeat.gpu_fabric:type_name -> llmgw.v1.GPUFabricState - 21, // 48: llmgw.v1.Heartbeat.imex:type_name -> llmgw.v1.IMEX - 25, // 49: llmgw.v1.Provision.engines:type_name -> llmgw.v1.EngineWant - 51, // 50: llmgw.v1.EngineWant.offline_sha256:type_name -> llmgw.v1.EngineWant.OfflineSha256Entry - 0, // 51: llmgw.v1.ShardStatus.state:type_name -> llmgw.v1.ShardStatus.State - 29, // 52: llmgw.v1.ShardStatus.fetch:type_name -> llmgw.v1.FetchProgress - 27, // 53: llmgw.v1.ShardStatus.compute:type_name -> llmgw.v1.Compute - 31, // 54: llmgw.v1.StartShard.profile:type_name -> llmgw.v1.EngineProfile - 52, // 55: llmgw.v1.EngineProfile.env:type_name -> llmgw.v1.EngineProfile.EnvEntry - 44, // 56: llmgw.v1.EngineProfile.adapters:type_name -> llmgw.v1.LoraAdapter - 1, // 57: llmgw.v1.AgentService.Connect:input_type -> llmgw.v1.AgentMessage - 3, // 58: llmgw.v1.AgentService.Connect:output_type -> llmgw.v1.ControlMessage - 58, // [58:59] is the sub-list for method output_type - 57, // [57:58] is the sub-list for method input_type - 57, // [57:57] is the sub-list for extension type_name - 57, // [57:57] is the sub-list for extension extendee - 0, // [0:57] is the sub-list for field type_name + 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 } func init() { file_llmgw_v1_agent_proto_init() } @@ -5306,6 +5733,8 @@ func file_llmgw_v1_agent_proto_init() { (*AgentMessage_FabricTest)(nil), (*AgentMessage_Upgrade)(nil), (*AgentMessage_KvEvents)(nil), + (*AgentMessage_TrainStatus)(nil), + (*AgentMessage_TrainArtifact)(nil), } file_llmgw_v1_agent_proto_msgTypes[2].OneofWrappers = []any{ (*ControlMessage_Start)(nil), @@ -5322,6 +5751,8 @@ func file_llmgw_v1_agent_proto_init() { (*ControlMessage_Firewall)(nil), (*ControlMessage_RestartAgent)(nil), (*ControlMessage_AgentSettings)(nil), + (*ControlMessage_TrainStart)(nil), + (*ControlMessage_TrainCancel)(nil), } type x struct{} out := protoimpl.TypeBuilder{ @@ -5329,7 +5760,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: 52, + NumMessages: 56, NumExtensions: 0, NumServices: 1, }, diff --git a/api/proto/llmgw/v1/agent.proto b/api/proto/llmgw/v1/agent.proto index 749af28..8be40c7 100644 --- a/api/proto/llmgw/v1/agent.proto +++ b/api/proto/llmgw/v1/agent.proto @@ -28,6 +28,8 @@ message AgentMessage { FabricTestReply fabric_test = 6; UpgradeReply upgrade = 7; KVEvents kv_events = 8; + TrainStatus train_status = 9; + TrainArtifact train_artifact = 10; } } @@ -74,6 +76,8 @@ message ControlMessage { Firewall firewall = 13; RestartAgent restart_agent = 14; AgentSettings agent_settings = 15; + TrainStart train_start = 16; + TrainCancel train_cancel = 17; } // 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 @@ -884,3 +888,55 @@ message LoraAdapter { string name = 1; string source = 2; } + +// TrainStart has a node train a LoRA adapter (BACKLOG F4): on a base model's +// weights from its registry, on question/answer pairs, with a Python trainer +// the agent carries, in an environment it makes over an installed engine's. +// One at a time per node. +message TrainStart { + string job_id = 1; + // base_repo and base_revision name the base model's weights on its + // registry (Hugging Face): safetensors, not the GGUF a catalog may serve. + string base_repo = 2; + string base_revision = 3; + // data is the training pairs, NDJSON of {"q", "a"}. + bytes data = 4; + // params is the trainer's settings as JSON: rank, alpha, epochs, lr, + // batch, max_len. + string params = 5; + // gpu_memory_mib is what the job needs of a card; the agent starts it + // only when one has that much free, which lets it run beside an engine. + int32 gpu_memory_mib = 6; + // gguf asks for the adapter in llama.cpp's format too. + bool gguf = 7; +} + +// TrainCancel stops a training job. +message TrainCancel { + string job_id = 1; +} + +// TrainStatus is how a training job is going, sent as it changes. +message TrainStatus { + string job_id = 1; + // state is preparing, running, done or failed. + string state = 2; + int32 step = 3; + int32 steps = 4; + double loss = 5; + // message is what is being done, or why it failed. + string message = 6; + // peak_mib is the most the job held of the card. + int32 peak_mib = 7; +} + +// TrainArtifact is a piece of what a training job made, sent before its +// final status: a gzipped tar of the PEFT adapter (kind peft), or the GGUF +// adapter (kind gguf). +message TrainArtifact { + string job_id = 1; + string kind = 2; + int32 seq = 3; + bytes data = 4; + bool last = 5; +} diff --git a/docs/admin-api.openapi.json b/docs/admin-api.openapi.json index 8b24cfb..79d7829 100644 --- a/docs/admin-api.openapi.json +++ b/docs/admin-api.openapi.json @@ -95,6 +95,9 @@ { "name": "summary" }, + { + "name": "train" + }, { "name": "updates" }, @@ -3689,6 +3692,254 @@ } } }, + "/api/v1/train/jobs": { + "post": { + "operationId": "postTrainJobs", + "summary": "Queue a LoRA adapter's training on a node", + "description": "The node trains when its agent is there, nothing else trains on it, and a card has gpu_memory_mib free; the pairs are NDJSON of {\"q\", \"a\"}. An adapter learns how answers sound far better than what they say.", + "tags": [ + "train" + ], + "parameters": [ + { + "$ref": "#/components/parameters/Actor" + } + ], + "requestBody": { + "required": true, + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/TrainRequest" + } + } + } + }, + "responses": { + "201": { + "description": "OK", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/TrainJobView" + } + } + } + }, + "401": { + "$ref": "#/components/responses/Unauthorized" + }, + "default": { + "$ref": "#/components/responses/Error" + } + } + }, + "get": { + "operationId": "getTrainJobs", + "summary": "Training jobs, newest first", + "tags": [ + "train" + ], + "parameters": [ + { + "name": "limit", + "in": "query", + "description": "How many; 50 when absent.", + "schema": { + "type": "integer" + } + } + ], + "responses": { + "200": { + "description": "OK", + "content": { + "application/json": { + "schema": { + "type": "array", + "items": { + "$ref": "#/components/schemas/TrainJobView" + } + } + } + } + }, + "401": { + "$ref": "#/components/responses/Unauthorized" + }, + "default": { + "$ref": "#/components/responses/Error" + } + } + } + }, + "/api/v1/train/jobs/{id}": { + "get": { + "operationId": "getTrainJobsId", + "summary": "One training job: how it goes, and what it made", + "tags": [ + "train" + ], + "parameters": [ + { + "name": "id", + "in": "path", + "required": true, + "description": "The job's ID, as starting it answered.", + "schema": { + "type": "string" + } + } + ], + "responses": { + "200": { + "description": "OK", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/TrainJobView" + } + } + } + }, + "401": { + "$ref": "#/components/responses/Unauthorized" + }, + "default": { + "$ref": "#/components/responses/Error" + } + } + }, + "delete": { + "operationId": "deleteTrainJobsId", + "summary": "Remove an ended training job and what it made", + "tags": [ + "train" + ], + "parameters": [ + { + "name": "id", + "in": "path", + "required": true, + "description": "The job's ID, as starting it answered.", + "schema": { + "type": "string" + } + }, + { + "$ref": "#/components/parameters/Actor" + } + ], + "responses": { + "200": { + "description": "OK", + "content": { + "application/json": { + "schema": { + "type": "object", + "additionalProperties": {} + } + } + } + }, + "401": { + "$ref": "#/components/responses/Unauthorized" + }, + "default": { + "$ref": "#/components/responses/Error" + } + } + } + }, + "/api/v1/train/jobs/{id}/artifacts/{kind}": { + "get": { + "operationId": "getTrainJobsIdArtifactsKind", + "summary": "What a training job made: peft (a gzipped tar) or gguf", + "tags": [ + "train" + ], + "parameters": [ + { + "name": "id", + "in": "path", + "required": true, + "description": "The job's ID, as starting it answered.", + "schema": { + "type": "string" + } + }, + { + "name": "kind", + "in": "path", + "required": true, + "description": "What the job made: peft, a gzipped tar of the PEFT adapter, or gguf, the adapter for llama.cpp.", + "schema": { + "type": "string" + } + } + ], + "responses": { + "200": { + "description": "OK", + "content": { + "application/json": { + "schema": { + "type": "object", + "additionalProperties": {} + } + } + } + }, + "401": { + "$ref": "#/components/responses/Unauthorized" + }, + "default": { + "$ref": "#/components/responses/Error" + } + } + } + }, + "/api/v1/train/jobs/{id}/cancel": { + "post": { + "operationId": "postTrainJobsIdCancel", + "summary": "Stop a training job", + "tags": [ + "train" + ], + "parameters": [ + { + "name": "id", + "in": "path", + "required": true, + "description": "The job's ID, as starting it answered.", + "schema": { + "type": "string" + } + }, + { + "$ref": "#/components/parameters/Actor" + } + ], + "responses": { + "200": { + "description": "OK", + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/TrainJobView" + } + } + } + }, + "401": { + "$ref": "#/components/responses/Unauthorized" + }, + "default": { + "$ref": "#/components/responses/Error" + } + } + } + }, "/api/v1/updates": { "get": { "operationId": "getUpdates", @@ -6816,6 +7067,123 @@ } } }, + "TrainJobView": { + "type": "object", + "description": "TrainJobView is a job as listed.", + "properties": { + "id": { + "type": "string" + }, + "name": { + "type": "string" + }, + "model": { + "type": "string" + }, + "base_repo": { + "type": "string" + }, + "base_revision": { + "type": "string" + }, + "node": { + "type": "string" + }, + "gpu_memory_mib": { + "type": "integer" + }, + "params": { + "type": "string" + }, + "gguf": { + "type": "boolean" + }, + "pairs": { + "type": "integer" + }, + "state": { + "type": "string" + }, + "step": { + "type": "integer" + }, + "steps": { + "type": "integer" + }, + "loss": { + "type": "number" + }, + "message": { + "type": "string" + }, + "peak_mib": { + "type": "integer" + }, + "created_at": { + "type": "string", + "format": "date-time" + }, + "started_at": { + "type": "string", + "format": "date-time" + }, + "finished_at": { + "type": "string", + "format": "date-time" + }, + "created_by": { + "type": "string" + }, + "artifacts": { + "type": "array", + "items": { + "type": "string" + }, + "description": "Artifacts are what it made: peft (a gzipped tar), gguf." + } + } + }, + "TrainRequest": { + "type": "object", + "description": "TrainRequest queues a job.", + "properties": { + "name": { + "type": "string", + "description": "Name is the adapter's: what requests will ask for it by." + }, + "model": { + "type": "string", + "description": "Model is the catalog model the adapter is for." + }, + "base_repo": { + "type": "string", + "description": "BaseRepo and BaseRevision are the weights trained on: the model's own on its registry, in safetensors (Hugging Face), not a GGUF." + }, + "base_revision": { + "type": "string" + }, + "node": { + "type": "string", + "description": "Node trains it; GPUMemoryMiB is what it needs of a card, which can be what an engine there leaves free." + }, + "gpu_memory_mib": { + "type": "integer" + }, + "data": { + "type": "string", + "description": "Data is the pairs, NDJSON of {\"q\", \"a\"} (llmgw ctl kb dataset)." + }, + "params": { + "type": "object", + "additionalProperties": {}, + "description": "Params tune the trainer: rank, alpha, dropout, epochs, lr, batch, max_len, seed." + }, + "gguf": { + "type": "boolean", + "description": "GGUF also makes the adapter in llama.cpp's format; when absent, it does when the model is served by llama.cpp." + } + } + }, "UpgradeRequest": { "type": "object", "description": "UpgradeRequest names the release a node should run.", diff --git a/internal/agent/agent.go b/internal/agent/agent.go index a5828ec..34ef9aa 100644 --- a/internal/agent/agent.go +++ b/internal/agent/agent.go @@ -34,6 +34,7 @@ import ( "gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines" "gitea.mynodes.xyz/ilker/llmgw/internal/agent/provision" "gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor" + "gitea.mynodes.xyz/ilker/llmgw/internal/agent/train" "gitea.mynodes.xyz/ilker/llmgw/internal/agent/weights" "gitea.mynodes.xyz/ilker/llmgw/internal/control/agentsvc" "gitea.mynodes.xyz/ilker/llmgw/internal/firewall" @@ -140,6 +141,9 @@ type Agent struct { // settingsCheck says whether the cluster's settings for this agent // would load beside its agent.yaml; see SetSettingsCheck. settingsCheck func([]byte) error + + // trainer runs LoRA training jobs, one at a time (BACKLOG F4). + trainer *train.Runner } // connectHTTPClient is the HTTP client connect uses; an interface so tests can @@ -186,6 +190,12 @@ func New(cfg Config, httpClient connectHTTPClient, log *slog.Logger) (*Agent, er kvSubs: map[string]*kvevents.Subscriber{}, } a.prov = &provision.Installer{Root: cfg.EnginesDir, Log: log, OnStatus: a.provisionStatus} + runAs := cfg.Superv.RunAs + if runAs == "" { + runAs = supervisor.DefaultRunAs + } + a.trainer = train.New(train.Env{EnginesDir: cfg.EnginesDir, WorkDir: filepath.Join(cfg.Superv.WorkDir, "trainer"), + RunAs: runAs, Log: log}) return a, nil } @@ -886,6 +896,10 @@ func (a *Agent) handle(ctx context.Context, msg *llmgwv1.ControlMessage, out cha go a.applyFirewall(m.Firewall) case *llmgwv1.ControlMessage_AgentSettings: go a.applySettings(m.AgentSettings) + case *llmgwv1.ControlMessage_TrainStart: + go a.trainJob(ctx, m.TrainStart, out) + case *llmgwv1.ControlMessage_TrainCancel: + a.trainer.Cancel(m.TrainCancel.GetJobId()) case *llmgwv1.ControlMessage_RestartAgent: go func() { a.log.Info("restarting, as the control plane asks; the engines keep serving", diff --git a/internal/agent/supervisor/runas.go b/internal/agent/supervisor/runas.go index 94e6992..e217dc4 100644 --- a/internal/agent/supervisor/runas.go +++ b/internal/agent/supervisor/runas.go @@ -3,6 +3,7 @@ package supervisor import ( "fmt" "os" + "os/exec" "os/user" "path/filepath" "strconv" @@ -90,3 +91,25 @@ func credentialFor(u *user.User, workDir string) (*runAs, error) { func (r *runAs) env() []string { return []string{"HOME=" + r.home, "XDG_CACHE_HOME=" + filepath.Join(r.home, ".cache"), "USER=" + r.name, "LOGNAME=" + r.name} } + +// AsEngineUser has cmd run as the engine user, as an engine would, with its +// home under workDir: for work the agent does that is not an engine but +// reads what an engine reads — training an adapter. It answers that home +// and a way to hand a path to the user, or why it runs as this process. +func AsEngineUser(cmd *exec.Cmd, name, workDir string) (home string, give func(path string) error, why string) { + r, why := resolveRunAs(name, workDir) + if r == nil { + return "", func(string) error { return nil }, why + } + cmd.SysProcAttr = &syscall.SysProcAttr{Credential: r.cred, Setpgid: true} + cmd.Env = append(cmd.Env, r.env()...) + uid, gid := int(r.cred.Uid), int(r.cred.Gid) + return r.home, func(path string) error { + return filepath.WalkDir(path, func(p string, _ os.DirEntry, err error) error { + if err != nil { + return err + } + return os.Lchown(p, uid, gid) + }) + }, "" +} diff --git a/internal/agent/train/runner.go b/internal/agent/train/runner.go new file mode 100644 index 0000000..0601b94 --- /dev/null +++ b/internal/agent/train/runner.go @@ -0,0 +1,574 @@ +// Package train runs LoRA training jobs on a node (BACKLOG F4): with an +// installed engine's Python — its torch and transformers, already matched to +// the node's CUDA — and a small overlay of what training needs besides, so a +// node that serves can train without a second multi-gigabyte environment. +package train + +import ( + "archive/tar" + "bufio" + "bytes" + "compress/gzip" + "context" + "crypto/sha256" + _ "embed" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "io" + "log/slog" + "net/http" + "os" + "os/exec" + "path/filepath" + "sort" + "strconv" + "strings" + "sync" + "syscall" + "time" + + "gitea.mynodes.xyz/ilker/llmgw/internal/agent/supervisor" +) + +//go:embed trainer.py +var trainerPy []byte + +// overlayPackages are what training needs besides an engine's Python, pinned +// and installed without their dependencies: those are the engine's own, and +// letting an installer resolve them pulled a second torch. +var overlayPackages = []string{"peft==0.21.2", "accelerate==1.15.0"} + +// Job is one training job, as the control plane sent it. +type Job struct { + ID string + BaseRepo string + BaseRevision string + // Data is NDJSON of {"q", "a"}. + Data []byte + // Params is the trainer's settings as JSON. + Params string + GPUMemoryMiB int + GGUF bool +} + +// Status is how a job is going. +type Status struct { + State string + Message string + Step int + Steps int + Loss float64 + PeakMiB int +} + +// Env is where a node keeps what training needs. +type Env struct { + // EnginesDir is where engines are installed (the agent's engines_dir). + EnginesDir string + // WorkDir is the trainer's own directory: its overlay, the llama.cpp + // conversion code, and each job's files while it runs. + WorkDir string + // RunAs is the user jobs run as: the engines' user. + RunAs string + Log *slog.Logger + // Python, UV, ConvDir and GPUs replace what would be found, for tests. + Python, UV, ConvDir string + GPUs func(ctx context.Context) ([]GPU, error) +} + +// GPU is one card and how much of it is free. +type GPU struct { + Index int + FreeMiB int +} + +// Runner runs one job at a time. +type Runner struct { + env Env + mu sync.Mutex + cancel context.CancelFunc + job string +} + +// New is a runner. +func New(env Env) *Runner { + if env.Log == nil { + env.Log = slog.New(slog.DiscardHandler) + } + return &Runner{env: env} +} + +// Cancel stops a job if it is the one running. +func (r *Runner) Cancel(id string) { + r.mu.Lock() + defer r.mu.Unlock() + if r.job == id && r.cancel != nil { + r.cancel() + } +} + +// ErrBusy is a node already training. +var ErrBusy = errors.New("this node is already training") + +// Run runs a job to its end, telling status as it goes and artifact each +// piece of what it made. It answers the error it failed with. +func (r *Runner) Run(ctx context.Context, job Job, status func(Status), artifact func(kind string, seq int, data []byte, last bool) error) error { + r.mu.Lock() + if r.job != "" { + r.mu.Unlock() + return ErrBusy + } + ctx, cancel := context.WithCancel(ctx) + r.job, r.cancel = job.ID, cancel + r.mu.Unlock() + defer func() { + cancel() + r.mu.Lock() + r.job, r.cancel = "", nil + r.mu.Unlock() + }() + + status(Status{State: "preparing", Message: "preparing the trainer"}) + py, err := r.python() + if err != nil { + return err + } + overlay, err := r.overlay(ctx, py) + if err != nil { + return fmt.Errorf("installing %s: %w", strings.Join(overlayPackages, ", "), err) + } + conv := "" + if job.GGUF { + if conv, err = r.conversion(ctx); err != nil { + return fmt.Errorf("fetching llama.cpp's conversion code: %w", err) + } + } + gpu, err := r.pickGPU(ctx, job.GPUMemoryMiB) + if err != nil { + return err + } + + dir := filepath.Join(r.env.WorkDir, "jobs", job.ID) + if err := os.MkdirAll(dir, 0o755); err != nil { + return err + } + defer os.RemoveAll(dir) + out := filepath.Join(dir, "adapter") + var params map[string]any + if job.Params != "" { + if err := json.Unmarshal([]byte(job.Params), ¶ms); err != nil { + return fmt.Errorf("the job's params: %w", err) + } + } + spec, _ := json.Marshal(map[string]any{"base_repo": job.BaseRepo, "base_revision": job.BaseRevision, + "data": filepath.Join(dir, "data.jsonl"), "out": out, "params": params}) + for name, b := range map[string][]byte{"data.jsonl": job.Data, "job.json": spec, "trainer.py": trainerPy} { + if err := os.WriteFile(filepath.Join(dir, name), b, 0o644); err != nil { + return err + } + } + + var baseDir string + err = r.python3(ctx, py, dir, []string{"PYTHONPATH=" + overlay, "CUDA_VISIBLE_DEVICES=" + strconv.Itoa(gpu)}, + []string{filepath.Join(dir, "trainer.py"), filepath.Join(dir, "job.json")}, + func(line map[string]any) { + s := Status{State: str(line["state"]), Message: str(line["message"]), Step: num(line["step"]), + Steps: num(line["steps"]), PeakMiB: num(line["peak_mib"])} + if v, ok := line["loss"].(float64); ok { + s.Loss = v + } + if b := str(line["base_dir"]); b != "" { + baseDir = b + } + if s.State == "trained" { + s.State = "running" + s.Message = "trained; packing the adapter" + } + status(s) + }) + if err != nil { + return err + } + + if job.GGUF { + status(Status{State: "running", Message: "converting the adapter to GGUF for llama.cpp"}) + gguf := filepath.Join(dir, "adapter.gguf") + err := r.python3(ctx, py, dir, []string{"PYTHONPATH=" + strings.Join([]string{filepath.Join(conv, "gguf-py"), conv, overlay}, ":")}, + []string{filepath.Join(conv, "convert_lora_to_gguf.py"), "--base", baseDir, "--outtype", "f16", "--outfile", gguf, out}, nil) + if err != nil { + return fmt.Errorf("converting to GGUF: %w", err) + } + if err := sendFile(gguf, "gguf", artifact); err != nil { + return err + } + } + var tgz bytes.Buffer + if err := tarGz(&tgz, out); err != nil { + return err + } + return sendBytes(tgz.Bytes(), "peft", artifact) +} + +func str(v any) string { + s, _ := v.(string) + return s +} + +func num(v any) int { + f, _ := v.(float64) + return int(f) +} + +// python3 runs a Python script as the engines' user, passing each JSON line +// it prints to line, and answers the last of what it wrote to stderr when it +// fails. +func (r *Runner) python3(ctx context.Context, py, dir string, env, args []string, line func(map[string]any)) error { + cmd := exec.CommandContext(ctx, py, args...) + cmd.Dir = dir + cmd.Env = append(os.Environ(), "PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True", "TOKENIZERS_PARALLELISM=false", + "HF_HUB_DISABLE_PROGRESS_BARS=1", "PYTHONUNBUFFERED=1") + cmd.Env = append(cmd.Env, env...) + home, give, why := supervisor.AsEngineUser(cmd, r.env.RunAs, r.env.WorkDir) + if why != "" { + r.env.Log.Warn(why) + } + if home != "" { + cmd.Env = append(cmd.Env, "HF_HOME="+filepath.Join(home, ".cache", "huggingface")) + } else { + cmd.Env = append(cmd.Env, "HF_HOME="+filepath.Join(r.env.WorkDir, "hf")) + } + if err := give(dir); err != nil { + return fmt.Errorf("handing %s to %s: %w", dir, r.env.RunAs, err) + } + // The whole process group goes when the job is cancelled: torch starts + // workers of its own. + if cmd.SysProcAttr == nil { + cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true} + } + cmd.Cancel = func() error { return syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL) } + stdout, err := cmd.StdoutPipe() + if err != nil { + return err + } + var stderr tailBuffer + cmd.Stderr = &stderr + if err := cmd.Start(); err != nil { + return err + } + sc := bufio.NewScanner(stdout) + sc.Buffer(make([]byte, 64<<10), 1<<20) + for sc.Scan() { + var m map[string]any + if line != nil && json.Unmarshal(sc.Bytes(), &m) == nil { + line(m) + } + } + err = cmd.Wait() + if ctx.Err() != nil { + return errors.New("cancelled") + } + if err != nil { + return fmt.Errorf("%v: %s", err, stderr.last()) + } + return nil +} + +// tailBuffer keeps the end of what a process wrote: the reason it failed is +// at the end, and the rest is progress bars. +type tailBuffer struct{ b []byte } + +func (t *tailBuffer) Write(p []byte) (int, error) { + t.b = append(t.b, p...) + if len(t.b) > 16<<10 { + t.b = t.b[len(t.b)-16<<10:] + } + return len(p), nil +} + +func (t *tailBuffer) last() string { + lines := strings.Split(strings.TrimSpace(string(t.b)), "\n") + if len(lines) > 3 { + lines = lines[len(lines)-3:] + } + return strings.Join(lines, " | ") +} + +// python is the newest installed vLLM's interpreter: its torch is built for +// this node's CUDA. +func (r *Runner) python() (string, error) { + if r.env.Python != "" { + return r.env.Python, nil + } + found, _ := filepath.Glob(filepath.Join(r.env.EnginesDir, "vllm", "*", "bin", "python3")) + if len(found) == 0 { + return "", errors.New("training needs vLLM installed on this node (its torch is what trains); " + + "llmgw ctl engine-build, or serving any vLLM model here, installs it") + } + sort.Strings(found) + return found[len(found)-1], nil +} + +func (r *Runner) uv() (string, error) { + if r.env.UV != "" { + return r.env.UV, nil + } + var found []string + _ = filepath.WalkDir(filepath.Join(r.env.EnginesDir, "uv"), func(p string, d os.DirEntry, err error) error { + if err == nil && !d.IsDir() && d.Name() == "uv" { + found = append(found, p) + } + return nil + }) + if len(found) == 0 { + return "", errors.New("no uv under " + filepath.Join(r.env.EnginesDir, "uv")) + } + sort.Strings(found) + return found[len(found)-1], nil +} + +// overlay installs the pinned packages once per Python and set of pins. +func (r *Runner) overlay(ctx context.Context, py string) (string, error) { + sum := sha256.Sum256([]byte(py + "\x00" + strings.Join(overlayPackages, " "))) + dir := filepath.Join(r.env.WorkDir, "overlay-"+hex.EncodeToString(sum[:6])) + if _, err := os.Stat(filepath.Join(dir, ".done")); err == nil { + return dir, nil + } + uv, err := r.uv() + if err != nil { + return "", err + } + tmp := dir + ".new" + _ = os.RemoveAll(tmp) + args := append([]string{"pip", "install", "--quiet", "--python", py, "--no-deps", "--target", tmp}, overlayPackages...) + cmd := exec.CommandContext(ctx, uv, args...) + if out, err := cmd.CombinedOutput(); err != nil { + return "", fmt.Errorf("%v: %s", err, strings.TrimSpace(string(out))) + } + _ = os.WriteFile(filepath.Join(tmp, ".done"), nil, 0o644) + _ = os.RemoveAll(dir) + if err := os.Rename(tmp, dir); err != nil { + return "", err + } + // Readable by the engines' user, which runs the trainer. + _ = filepath.WalkDir(dir, func(p string, d os.DirEntry, err error) error { + if err == nil { + mode := os.FileMode(0o644) + if d.IsDir() { + mode = 0o755 + } + _ = os.Chmod(p, mode) + } + return nil + }) + return dir, nil +} + +// llamaCppTag is the release of the newest llama.cpp installed: its +// conversion code is the one whose GGUF that llama.cpp reads. +func (r *Runner) llamaCppTag() (string, error) { + found, _ := filepath.Glob(filepath.Join(r.env.EnginesDir, "llamacpp", "b*")) + if len(found) == 0 { + return "", errors.New("a GGUF adapter needs llama.cpp installed on this node, to match its release") + } + sort.Strings(found) + tag, _, _ := strings.Cut(filepath.Base(found[len(found)-1]), "-") + return tag, nil +} + +// conversionSource is where llama.cpp's source is fetched from, by tag. +var conversionSource = "https://codeload.github.com/ggml-org/llama.cpp/tar.gz/refs/tags/" + +// conversion fetches llama.cpp's LoRA conversion code for the installed +// release, once. +func (r *Runner) conversion(ctx context.Context) (string, error) { + if r.env.ConvDir != "" { + return r.env.ConvDir, nil + } + tag, err := r.llamaCppTag() + if err != nil { + return "", err + } + dir := filepath.Join(r.env.WorkDir, "llama.cpp-"+tag) + if _, err := os.Stat(filepath.Join(dir, "convert_lora_to_gguf.py")); err == nil { + return dir, nil + } + req, err := http.NewRequestWithContext(ctx, http.MethodGet, conversionSource+tag, nil) + if err != nil { + return "", err + } + res, err := http.DefaultClient.Do(req) + if err != nil { + return "", err + } + defer res.Body.Close() + if res.StatusCode != http.StatusOK { + return "", fmt.Errorf("%s answered %s", conversionSource+tag, res.Status) + } + tmp := dir + ".new" + _ = os.RemoveAll(tmp) + if err := untarConversion(res.Body, tmp); err != nil { + return "", err + } + _ = os.RemoveAll(dir) + return dir, os.Rename(tmp, dir) +} + +// untarConversion keeps, of llama.cpp's source, the conversion scripts and +// the Python packages they import. +func untarConversion(r io.Reader, dir string) error { + gz, err := gzip.NewReader(r) + if err != nil { + return err + } + tr := tar.NewReader(gz) + kept := 0 + for { + h, err := tr.Next() + if errors.Is(err, io.EOF) { + break + } + if err != nil { + return err + } + _, rel, ok := strings.Cut(h.Name, "/") + if !ok || h.Typeflag != tar.TypeReg || strings.Contains(rel, "..") { + continue + } + if rel != "convert_lora_to_gguf.py" && rel != "convert_hf_to_gguf.py" && + !strings.HasPrefix(rel, "conversion/") && !strings.HasPrefix(rel, "gguf-py/") { + continue + } + p := filepath.Join(dir, filepath.FromSlash(rel)) + if err := os.MkdirAll(filepath.Dir(p), 0o755); err != nil { + return err + } + f, err := os.OpenFile(p, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o644) + if err != nil { + return err + } + _, err = io.Copy(f, io.LimitReader(tr, 64<<20)) + f.Close() + if err != nil { + return err + } + kept++ + } + if kept == 0 { + return errors.New("no conversion code in llama.cpp's source") + } + return nil +} + +// pickGPU is the card with the most memory free, if that is enough. +func (r *Runner) pickGPU(ctx context.Context, need int) (int, error) { + list := r.env.GPUs + if list == nil { + list = nvidiaSMI + } + gpus, err := list(ctx) + if err != nil { + return 0, fmt.Errorf("reading the cards' free memory: %w", err) + } + if len(gpus) == 0 { + return 0, errors.New("this node has no GPU to train on") + } + best := gpus[0] + for _, g := range gpus[1:] { + if g.FreeMiB > best.FreeMiB { + best = g + } + } + if best.FreeMiB < need { + return 0, fmt.Errorf("no card here has %d MiB free (the most free is GPU %d with %d MiB); "+ + "a smaller base, or a card freed for it, would do", need, best.Index, best.FreeMiB) + } + return best.Index, nil +} + +func nvidiaSMI(ctx context.Context) ([]GPU, error) { + ctx, cancel := context.WithTimeout(ctx, 15*time.Second) + defer cancel() + out, err := exec.CommandContext(ctx, "nvidia-smi", "--query-gpu=index,memory.free", "--format=csv,noheader,nounits").Output() + if err != nil { + return nil, err + } + return parseGPUs(string(out)), nil +} + +func parseGPUs(s string) []GPU { + var out []GPU + for _, line := range strings.Split(strings.TrimSpace(s), "\n") { + a, b, ok := strings.Cut(line, ",") + if !ok { + continue + } + i, err1 := strconv.Atoi(strings.TrimSpace(a)) + f, err2 := strconv.Atoi(strings.TrimSpace(b)) + if err1 == nil && err2 == nil { + out = append(out, GPU{Index: i, FreeMiB: f}) + } + } + return out +} + +// tarGz packs a directory. +func tarGz(w io.Writer, dir string) error { + gz := gzip.NewWriter(w) + tw := tar.NewWriter(gz) + err := filepath.WalkDir(dir, func(p string, d os.DirEntry, err error) error { + if err != nil || d.IsDir() { + return err + } + rel, _ := filepath.Rel(dir, p) + info, err := d.Info() + if err != nil { + return err + } + if err := tw.WriteHeader(&tar.Header{Name: filepath.ToSlash(rel), Mode: 0o644, Size: info.Size(), ModTime: info.ModTime()}); err != nil { + return err + } + f, err := os.Open(p) + if err != nil { + return err + } + defer f.Close() + _, err = io.Copy(tw, f) + return err + }) + if err != nil { + return err + } + if err := tw.Close(); err != nil { + return err + } + return gz.Close() +} + +// artifactChunk is how much of an artifact goes in one message. +const artifactChunk = 1 << 20 + +func sendFile(path, kind string, send func(string, int, []byte, bool) error) error { + b, err := os.ReadFile(path) + if err != nil { + return err + } + return sendBytes(b, kind, send) +} + +func sendBytes(b []byte, kind string, send func(string, int, []byte, bool) error) error { + seq := 0 + for { + n := min(artifactChunk, len(b)) + last := n == len(b) + if err := send(kind, seq, b[:n], last); err != nil { + return err + } + if last { + return nil + } + b = b[n:] + seq++ + } +} diff --git a/internal/agent/train/runner_test.go b/internal/agent/train/runner_test.go new file mode 100644 index 0000000..d6f0788 --- /dev/null +++ b/internal/agent/train/runner_test.go @@ -0,0 +1,164 @@ +package train + +import ( + "archive/tar" + "bytes" + "compress/gzip" + "context" + "io" + "os" + "path/filepath" + "strings" + "sync" + "testing" +) + +// fakePython stands in for an engine's interpreter: as the trainer it +// prints progress and writes an adapter; as the converter it writes a GGUF; +// told to fail, it fails. +const fakePython = `#!/bin/sh +case "$1" in +*trainer.py) + out=$(sed -n 's/.*"out": *"\([^"]*\)".*/\1/p' "$2") + grep -q FAIL "$(dirname "$2")/data.jsonl" && { echo "CUDA out of memory" >&2; exit 1; } + echo '{"state":"preparing","message":"fetching base"}' + echo '{"state":"running","step":5,"steps":10,"loss":0.9,"peak_mib":1900}' + mkdir -p "$out"; echo '{"r":16}' > "$out/adapter_config.json"; echo weights > "$out/adapter_model.safetensors" + echo '{"state":"trained","step":10,"steps":10,"peak_mib":2100,"base_dir":"/base"}' + ;; +*convert_lora_to_gguf.py) + while [ $# -gt 0 ]; do [ "$1" = "--outfile" ] && echo GGUF > "$2"; shift; done + ;; +esac +` + +func runner(t *testing.T, free int) *Runner { + t.Helper() + dir := t.TempDir() + py := filepath.Join(dir, "python3") + _ = os.WriteFile(py, []byte(fakePython), 0o755) + uv := filepath.Join(dir, "uv") + _ = os.WriteFile(uv, []byte("#!/bin/sh\nwhile [ $# -gt 0 ]; do [ \"$1\" = \"--target\" ] && mkdir -p \"$2\"; shift; done\n"), 0o755) + return New(Env{WorkDir: filepath.Join(dir, "work"), Python: py, UV: uv, ConvDir: dir, + GPUs: func(context.Context) ([]GPU, error) { return []GPU{{0, 500}, {1, free}}, nil }}) +} + +type sink struct { + mu sync.Mutex + statuses []Status + parts map[string][]byte + lasts int +} + +func (s *sink) status(st Status) { s.mu.Lock(); s.statuses = append(s.statuses, st); s.mu.Unlock() } +func (s *sink) artifact(kind string, _ int, b []byte, last bool) error { + s.mu.Lock() + defer s.mu.Unlock() + if s.parts == nil { + s.parts = map[string][]byte{} + } + s.parts[kind] = append(s.parts[kind], b...) + if last { + s.lasts++ + } + return nil +} + +func TestAJobTrainsOnTheFreestCardAndSendsWhatItMade(t *testing.T) { + r := runner(t, 3000) + s := &sink{} + err := r.Run(context.Background(), Job{ID: "j1", BaseRepo: "Qwen/x", Data: []byte(`{"q":"a","a":"b"}`), GPUMemoryMiB: 2500, GGUF: true}, + s.status, s.artifact) + if err != nil { + t.Fatal(err) + } + var running bool + for _, st := range s.statuses { + if st.State == "running" && st.Step == 5 && st.Steps == 10 && st.PeakMiB == 1900 { + running = true + } + } + if !running { + t.Errorf("statuses %+v", s.statuses) + } + if string(s.parts["gguf"]) != "GGUF\n" || s.lasts != 2 { + t.Errorf("gguf %q, lasts %d", s.parts["gguf"], s.lasts) + } + files := map[string]string{} + gz, err := gzip.NewReader(bytes.NewReader(s.parts["peft"])) + if err != nil { + t.Fatal(err) + } + tr := tar.NewReader(gz) + for { + h, err := tr.Next() + if err != nil { + break + } + b, _ := io.ReadAll(tr) + files[h.Name] = string(b) + } + if files["adapter_model.safetensors"] != "weights\n" || files["adapter_config.json"] == "" { + t.Errorf("peft archive %v", files) + } + // Nothing of the job is left on the node. + if entries, _ := os.ReadDir(filepath.Join(r.env.WorkDir, "jobs")); len(entries) != 0 { + t.Errorf("left behind: %v", entries) + } +} + +func TestAJobWaitsForNoCardItDoesNotFit(t *testing.T) { + r := runner(t, 1000) + err := r.Run(context.Background(), Job{ID: "j1", Data: []byte(`{"q":"a","a":"b"}`), GPUMemoryMiB: 2500}, (&sink{}).status, (&sink{}).artifact) + if err == nil || !strings.Contains(err.Error(), "no card here has 2500 MiB free") { + t.Errorf("got %v", err) + } +} + +func TestATrainerThatFailsSaysWhy(t *testing.T) { + r := runner(t, 3000) + err := r.Run(context.Background(), Job{ID: "j1", Data: []byte(`{"q":"FAIL","a":"b"}`)}, (&sink{}).status, (&sink{}).artifact) + if err == nil || !strings.Contains(err.Error(), "CUDA out of memory") { + t.Errorf("got %v", err) + } +} + +func TestOneJobAtATime(t *testing.T) { + r := runner(t, 3000) + r.job = "other" + if err := r.Run(context.Background(), Job{ID: "j2"}, (&sink{}).status, (&sink{}).artifact); err != ErrBusy { + t.Errorf("got %v", err) + } +} + +func TestCardsAreReadFromNvidiaSMI(t *testing.T) { + got := parseGPUs("0, 2310\n1, 80000\nbad line\n") + if len(got) != 2 || got[1].FreeMiB != 80000 || got[0].Index != 0 { + t.Errorf("parsed %+v", got) + } +} + +func TestOnlyTheConversionCodeIsKeptOfLlamaCpp(t *testing.T) { + var buf bytes.Buffer + gz := gzip.NewWriter(&buf) + tw := tar.NewWriter(gz) + for _, name := range []string{"llama.cpp-b1/convert_lora_to_gguf.py", "llama.cpp-b1/conversion/qwen.py", + "llama.cpp-b1/gguf-py/gguf/__init__.py", "llama.cpp-b1/src/llama.cpp", "llama.cpp-b1/../evil"} { + _ = tw.WriteHeader(&tar.Header{Name: name, Mode: 0o644, Size: 1, Typeflag: tar.TypeReg}) + _, _ = tw.Write([]byte("x")) + } + _ = tw.Close() + _ = gz.Close() + dir := filepath.Join(t.TempDir(), "conv") + if err := untarConversion(&buf, dir); err != nil { + t.Fatal(err) + } + for _, want := range []string{"convert_lora_to_gguf.py", "conversion/qwen.py", "gguf-py/gguf/__init__.py"} { + if _, err := os.Stat(filepath.Join(dir, want)); err != nil { + t.Errorf("%s not kept", want) + } + } + if _, err := os.Stat(filepath.Join(dir, "src")); err == nil { + t.Error("the rest of llama.cpp was kept") + } +} diff --git a/internal/agent/train/trainer.py b/internal/agent/train/trainer.py new file mode 100644 index 0000000..3d83a51 --- /dev/null +++ b/internal/agent/train/trainer.py @@ -0,0 +1,101 @@ +# llmgw's LoRA trainer (BACKLOG F4), run by the node agent with an installed +# engine's Python (its torch and transformers) and a small overlay (peft, +# accelerate). Reads a job, trains, writes the PEFT adapter, and says how it +# goes as one JSON object a line on standard output. +import json, math, os, random, sys, threading, subprocess, time + +def say(**kw): + print(json.dumps(kw), flush=True) + +job = json.load(open(sys.argv[1], encoding="utf-8")) +p = {"rank": 16, "alpha": 32, "dropout": 0.05, "epochs": 6, "lr": 2e-4, "batch": 4, "max_len": 512, "seed": 0} +p.update(job.get("params") or {}) + +say(state="preparing", message="loading libraries") +import torch +from huggingface_hub import snapshot_download +from transformers import AutoModelForCausalLM, AutoTokenizer +from peft import LoraConfig, get_peft_model + +random.seed(p["seed"]); torch.manual_seed(p["seed"]) +if not torch.cuda.is_available(): + raise SystemExit("no CUDA device is visible to the trainer") + +say(state="preparing", message="fetching " + job["base_repo"]) +base_dir = snapshot_download(job["base_repo"], revision=job.get("base_revision") or None, + allow_patterns=["*.json", "*.safetensors", "*.txt", "*.model", "*.tiktoken"]) +say(state="preparing", message="loading the base model", base_dir=base_dir) + +# bf16 where the card has it; fp16 with a loss scaler where it does not +# (Turing and older). +bf16 = torch.cuda.is_bf16_supported() +dtype = torch.bfloat16 if bf16 else torch.float16 +tok = AutoTokenizer.from_pretrained(base_dir) +if tok.pad_token_id is None: + tok.pad_token = tok.eos_token +model = AutoModelForCausalLM.from_pretrained(base_dir, dtype=dtype).to("cuda") + +rows = [json.loads(l) for l in open(job["data"], encoding="utf-8") if l.strip()] +rows = [r for r in rows if r.get("q") and r.get("a")] +if not rows: + raise SystemExit("no {q, a} pairs in the data") + +def encode(r): + prompt = tok.apply_chat_template([{"role": "user", "content": r["q"]}], add_generation_prompt=True, tokenize=False) + full = tok.apply_chat_template([{"role": "user", "content": r["q"]}, {"role": "assistant", "content": r["a"]}], tokenize=False) + pi = tok(prompt, add_special_tokens=False)["input_ids"] + fi = tok(full, add_special_tokens=False)["input_ids"][: p["max_len"]] + return fi, ([-100] * len(pi) + fi[len(pi):])[: len(fi)] + +data = [encode(r) for r in rows] + +model.gradient_checkpointing_enable() +model.enable_input_require_grads() +model = get_peft_model(model, LoraConfig(r=int(p["rank"]), lora_alpha=int(p["alpha"]), lora_dropout=float(p["dropout"]), + task_type="CAUSAL_LM", target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"])) +trainable = [x for x in model.parameters() if x.requires_grad] +for x in trainable: # fp32 trainable weights, so a loss scaler can unscale them + x.data = x.data.float() + +peak = {"mib": 0} +stop = threading.Event() +def watch(): + while not stop.is_set(): + try: + out = subprocess.run(["nvidia-smi", "--query-compute-apps=pid,used_memory", "--format=csv,noheader,nounits"], + capture_output=True, text=True, timeout=5).stdout + for line in out.splitlines(): + pid, mib = [v.strip() for v in line.split(",")] + if int(pid) == os.getpid(): + peak["mib"] = max(peak["mib"], int(mib)) + except Exception: + pass + time.sleep(0.5) +threading.Thread(target=watch, daemon=True).start() + +opt = torch.optim.AdamW(trainable, lr=float(p["lr"]), weight_decay=0.0) +scaler = torch.amp.GradScaler("cuda", enabled=not bf16) +bs, epochs = int(p["batch"]), int(p["epochs"]) +steps = epochs * math.ceil(len(data) / bs) +warm = max(1, steps // 10) +sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: min(1.0, (s + 1) / warm) * max(0.0, 1 - s / steps)) +model.train(); step = 0; last = time.time() +for ep in range(epochs): + random.shuffle(data) + for i in range(0, len(data), bs): + batch = data[i:i + bs]; n = max(len(x[0]) for x in batch) + ids = torch.tensor([x[0] + [tok.pad_token_id] * (n - len(x[0])) for x in batch], device="cuda") + lab = torch.tensor([x[1] + [-100] * (n - len(x[1])) for x in batch], device="cuda") + att = torch.tensor([[1] * len(x[0]) + [0] * (n - len(x[0])) for x in batch], device="cuda") + with torch.autocast("cuda", dtype=dtype): + loss = model(input_ids=ids, attention_mask=att, labels=lab).loss + scaler.scale(loss).backward(); scaler.unscale_(opt) + torch.nn.utils.clip_grad_norm_(trainable, 1.0) + scaler.step(opt); scaler.update(); opt.zero_grad(set_to_none=True); sched.step(); step += 1 + if time.time() - last > 2 or step == steps: + last = time.time() + say(state="running", step=step, steps=steps, loss=round(loss.item(), 4), peak_mib=peak["mib"]) + +model.save_pretrained(job["out"]) +stop.set() +say(state="trained", step=step, steps=steps, peak_mib=peak["mib"], base_dir=base_dir) diff --git a/internal/agent/trainjob.go b/internal/agent/trainjob.go new file mode 100644 index 0000000..893f7ba --- /dev/null +++ b/internal/agent/trainjob.go @@ -0,0 +1,46 @@ +package agent + +import ( + "context" + + llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1" + "gitea.mynodes.xyz/ilker/llmgw/internal/agent/train" +) + +// trainJob runs a training job the control plane sent, telling it how the +// job goes and sending what it made. A stream that ends ends the job: the +// control plane, which sees the node go, marks it failed, and it is started +// again from the beginning rather than resumed. +func (a *Agent) trainJob(ctx context.Context, m *llmgwv1.TrainStart, out chan<- *llmgwv1.AgentMessage) { + id := m.GetJobId() + send := func(msg *llmgwv1.AgentMessage) bool { + select { + case out <- msg: + return true + case <-ctx.Done(): + return false + } + } + status := func(s train.Status) { + send(&llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_TrainStatus{TrainStatus: &llmgwv1.TrainStatus{ + JobId: id, State: s.State, Message: s.Message, Step: int32(s.Step), Steps: int32(s.Steps), //nolint:gosec // counts + Loss: s.Loss, PeakMib: int32(s.PeakMiB)}}}) //nolint:gosec // MiB + } + artifact := func(kind string, seq int, data []byte, last bool) error { + if !send(&llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_TrainArtifact{TrainArtifact: &llmgwv1.TrainArtifact{ + JobId: id, Kind: kind, Seq: int32(seq), Data: data, Last: last}}}) { //nolint:gosec // chunk count + return ctx.Err() + } + return nil + } + a.log.Info("training an adapter", "job", id, "base", m.GetBaseRepo()) + err := a.trainer.Run(ctx, train.Job{ID: id, BaseRepo: m.GetBaseRepo(), BaseRevision: m.GetBaseRevision(), + Data: m.GetData(), Params: m.GetParams(), GPUMemoryMiB: int(m.GetGpuMemoryMib()), GGUF: m.GetGguf()}, status, artifact) + if err != nil { + a.log.Warn("training failed", "job", id, "err", err) + status(train.Status{State: "failed", Message: err.Error()}) + return + } + a.log.Info("trained an adapter", "job", id) + status(train.Status{State: "done"}) +} diff --git a/internal/cli/ctl.go b/internal/cli/ctl.go index 0f51844..d699094 100644 --- a/internal/cli/ctl.go +++ b/internal/cli/ctl.go @@ -186,6 +186,7 @@ func newCtlCmd() *cobra.Command { newQuotaCmd(client, &asJSON), newGuardCmd(client, &asJSON), newKBCmd(client, &asJSON), + newTrainCmd(client, &asJSON), newConfigCmd(client, &asJSON), newAuditCmd(client, &asJSON), newNotifyCmd(client, &asJSON), diff --git a/internal/cli/server.go b/internal/cli/server.go index 36df353..a747364 100644 --- a/internal/cli/server.go +++ b/internal/cli/server.go @@ -49,6 +49,7 @@ import ( "gitea.mynodes.xyz/ilker/llmgw/internal/control/fleetmetrics" "gitea.mynodes.xyz/ilker/llmgw/internal/control/modelsvc" "gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler" + "gitea.mynodes.xyz/ilker/llmgw/internal/control/train" "gitea.mynodes.xyz/ilker/llmgw/internal/guard" "gitea.mynodes.xyz/ilker/llmgw/internal/hf" "gitea.mynodes.xyz/ilker/llmgw/internal/kvevents" @@ -387,6 +388,13 @@ func newServerCmd() *cobra.Command { // Knowledge bases, with the cluster's own models to embed a // probe (a collection's dimension) and an operator's search. api.SetKB(&adminapi.KBAdmin{Store: st, Embed: embedder(rt.Replay)}) + // LoRA adapters trained on the nodes. + if agentSvc != nil { + trainer := train.New(st, agentSvc, log) + agentSvc.SetTrain(trainer) + go trainer.Run(ctx) + api.SetTrain(&adminapi.TrainAdmin{Store: st, Service: trainer}) + } // The cluster's settings, each change checked against this // machine's own keys before it is kept. api.SetSettings(&adminapi.Settings{Store: st, LocalKeys: config.LocalKeys(), diff --git a/internal/cli/train.go b/internal/cli/train.go new file mode 100644 index 0000000..dccfbb2 --- /dev/null +++ b/internal/cli/train.go @@ -0,0 +1,248 @@ +package cli + +import ( + "encoding/json" + "fmt" + "io" + "os" + "strings" + "time" + + "github.com/spf13/cobra" + + "gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi" + "gitea.mynodes.xyz/ilker/llmgw/internal/store" +) + +func newTrainCmd(client func() *ctlClient, asJSON *bool) *cobra.Command { + c := &cobra.Command{ + Use: "train", + Short: "Train LoRA adapters on the nodes, on question/answer pairs", + Long: "A node trains an adapter for a catalog model on pairs (llmgw ctl kb dataset\n" + + "writes them from a knowledge base; read them through first). It trains with\n" + + "its installed vLLM's torch and a small overlay, as the engines' user, on a card\n" + + "with --gpu-memory free — beside an engine, if it leaves that much. What it\n" + + "makes is kept by the control plane: llmgw ctl train download, or publish it\n" + + "as the model's adapter.\n\n" + + "An adapter learns how answers sound — words, format, tone — far better than\n" + + "what they say. Facts belong to the knowledge base.", + } + + var model, base, rev, node, data string + var gpuMem int + var epochs, rank int + var lr float64 + var noGGUF bool + start := &cobra.Command{ + Use: "start NAME --model MODEL --base-repo OWNER/NAME --node NODE --data FILE", + Short: "Queue an adapter's training on a node", + Example: " llmgw ctl train start policies-tone --model qwen2.5-0.5b-instruct \\\n" + + " --base-repo Qwen/Qwen2.5-0.5B-Instruct --node cuda01 --gpu-memory 2300 --data policies.jsonl", + Args: cobra.ExactArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + b, err := os.ReadFile(data) + if err != nil { + return err + } + req := adminapi.TrainRequest{Name: args[0], Model: model, BaseRepo: base, BaseRevision: rev, Node: node, + GPUMemoryMiB: gpuMem, Data: string(b), Params: map[string]any{}} + if cmd.Flags().Changed("epochs") { + req.Params["epochs"] = epochs + } + if cmd.Flags().Changed("rank") { + req.Params["rank"] = rank + } + if cmd.Flags().Changed("lr") { + req.Params["lr"] = lr + } + if noGGUF { + f := false + req.GGUF = &f + } + body, _ := json.Marshal(req) + var v adminapi.TrainJobView + if err := client().patient(2*time.Minute).do(cmd.Context(), "POST", "/api/v1/train/jobs", json.RawMessage(body), &v); err != nil { + return err + } + fmt.Fprintf(cmd.OutOrStdout(), "queued %s: %s for %s on %s, %d pairs\n", v.ID, v.Name, v.Model, v.Node, v.Pairs) + fmt.Fprintf(cmd.OutOrStdout(), "follow it with: llmgw ctl train show %s --follow\n", v.ID) + return nil + }, + } + f := start.Flags() + f.StringVar(&model, "model", "", "the catalog model the adapter is for") + f.StringVar(&base, "base-repo", "", "the weights to train on, as safetensors on Hugging Face (owner/name)") + f.StringVar(&rev, "base-revision", "", "a revision of them (default: the newest)") + f.StringVar(&node, "node", "", "the node that trains") + f.IntVar(&gpuMem, "gpu-memory", 4096, "MiB of a card the job needs free") + f.StringVar(&data, "data", "", "the pairs, NDJSON of {\"q\", \"a\"}") + f.IntVar(&epochs, "epochs", 6, "passes over the pairs") + f.IntVar(&rank, "rank", 16, "the adapter's rank") + f.Float64Var(&lr, "lr", 2e-4, "the learning rate") + f.BoolVar(&noGGUF, "no-gguf", false, "skip the llama.cpp GGUF, even for a model llama.cpp serves") + for _, req := range []string{"model", "base-repo", "node", "data"} { + _ = start.MarkFlagRequired(req) + } + c.AddCommand(start) + + c.AddCommand(&cobra.Command{ + Use: "list", + Short: "Training jobs, newest first", + Args: cobra.NoArgs, + RunE: func(cmd *cobra.Command, _ []string) error { + var jobs []adminapi.TrainJobView + if err := client().get(cmd.Context(), "/api/v1/train/jobs", &jobs); err != nil { + return err + } + if *asJSON { + return json.NewEncoder(cmd.OutOrStdout()).Encode(jobs) + } + if len(jobs) == 0 { + fmt.Fprintln(cmd.OutOrStdout(), "No training jobs (llmgw ctl train start).") + return nil + } + w := tab() + fmt.Fprintln(w, "JOB\tADAPTER\tMODEL\tNODE\tSTATE\tPROGRESS\tLOSS\tPEAK\tCREATED") + for _, j := range jobs { + fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\n", j.ID, j.Name, j.Model, j.Node, j.State, + trainProgress(j), trainLoss(j), trainPeak(j), j.CreatedAt.Local().Format("01-02 15:04")) + } + return w.Flush() + }, + }) + + var follow bool + show := &cobra.Command{ + Use: "show JOB", + Short: "How a training job goes, and what it made", + Args: cobra.ExactArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + out := cmd.OutOrStdout() + last := "" + for { + var j adminapi.TrainJobView + if err := client().get(cmd.Context(), "/api/v1/train/jobs/"+args[0], &j); err != nil { + return err + } + if *asJSON { + return json.NewEncoder(out).Encode(j) + } + line := fmt.Sprintf("%s %s %s loss %s peak %s %s", j.State, trainProgress(j), j.Node, trainLoss(j), trainPeak(j), j.Message) + if !follow { + fmt.Fprintf(out, "%s: adapter %s for %s, on %s (base %s), %d pairs\n", j.ID, j.Name, j.Model, j.Node, j.BaseRepo, j.Pairs) + fmt.Fprintln(out, line) + if len(j.Artifacts) > 0 { + fmt.Fprintf(out, "made: %s (llmgw ctl train download %s --kind KIND)\n", strings.Join(j.Artifacts, ", "), j.ID) + } + return nil + } + if line != last { + fmt.Fprintf(out, "%s %s\n", time.Now().Format("15:04:05"), line) + last = line + } + if store.TrainOver(j.State) { + if j.State != store.TrainDone { + return fmt.Errorf("%s %s", j.ID, j.State) + } + return nil + } + select { + case <-cmd.Context().Done(): + return nil + case <-time.After(3 * time.Second): + } + } + }, + } + show.Flags().BoolVarP(&follow, "follow", "f", false, "follow it until it ends") + c.AddCommand(show) + + c.AddCommand(&cobra.Command{ + Use: "cancel JOB", + Short: "Stop a training job", + Args: cobra.ExactArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + if err := client().do(cmd.Context(), "POST", "/api/v1/train/jobs/"+args[0]+"/cancel", nil, nil); err != nil { + return err + } + fmt.Fprintf(cmd.OutOrStdout(), "%s cancelled\n", args[0]) + return nil + }, + }) + c.AddCommand(&cobra.Command{ + Use: "rm JOB", + Short: "Remove an ended training job and what it made", + Args: cobra.ExactArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + if err := client().do(cmd.Context(), "DELETE", "/api/v1/train/jobs/"+args[0], nil, nil); err != nil { + return err + } + fmt.Fprintf(cmd.OutOrStdout(), "%s removed\n", args[0]) + return nil + }, + }) + + var kind, outFile string + download := &cobra.Command{ + Use: "download JOB", + Short: "Save what a training job made: the PEFT adapter (peft) or the GGUF (gguf)", + Args: cobra.ExactArgs(1), + RunE: func(cmd *cobra.Command, args []string) error { + if outFile == "" { + outFile = args[0] + ".tar.gz" + if kind == "gguf" { + outFile = args[0] + ".gguf" + } + } + f, err := os.Create(outFile) + if err != nil { + return err + } + defer f.Close() + cw := &countWriter{w: f} + if err := client().patient(10*time.Minute).do(cmd.Context(), "GET", + "/api/v1/train/jobs/"+args[0]+"/artifacts/"+kind, nil, io.Writer(cw)); err != nil { + _ = os.Remove(outFile) + return err + } + fmt.Fprintf(cmd.OutOrStdout(), "%s: %s\n", outFile, humanBytes(cw.n)) + return nil + }, + } + download.Flags().StringVar(&kind, "kind", "peft", "peft (a gzipped tar of the PEFT adapter) or gguf") + download.Flags().StringVarP(&outFile, "out", "o", "", "the file to write") + c.AddCommand(download) + return c +} + +type countWriter struct { + w io.Writer + n int64 +} + +func (c *countWriter) Write(p []byte) (int, error) { + n, err := c.w.Write(p) + c.n += int64(n) + return n, err +} + +func trainProgress(j adminapi.TrainJobView) string { + if j.Steps == 0 { + return "-" + } + return fmt.Sprintf("%d/%d", j.Step, j.Steps) +} + +func trainLoss(j adminapi.TrainJobView) string { + if j.Loss == 0 { + return "-" + } + return fmt.Sprintf("%.3f", j.Loss) +} + +func trainPeak(j adminapi.TrainJobView) string { + if j.PeakMiB == 0 { + return "-" + } + return fmt.Sprintf("%d MiB", j.PeakMiB) +} diff --git a/internal/control/adminapi/api.go b/internal/control/adminapi/api.go index ddd921a..821ed52 100644 --- a/internal/control/adminapi/api.go +++ b/internal/control/adminapi/api.go @@ -70,7 +70,9 @@ type API struct { // guard is the guardrail policies' library and attachments. guard *GuardAdmin // kb is knowledge bases; nil without a store. - kb *KBAdmin + kb *KBAdmin + // train is adapter training; nil without a store and agents. + train *TrainAdmin settings *Settings // restartAgents tells every connected agent to start again. restartAgents func(reason string) []string @@ -117,6 +119,7 @@ func (a *API) Register(mux Mux) { mux.HandleFunc("POST /api/v1/models/{name}/resources", a.setResources) a.registerQuotas(mux) a.registerKB(mux) + a.registerTrain(mux) a.registerGuard(mux) a.registerSettings(mux) mux.HandleFunc("POST /api/v1/agents/restart", a.restartAgentsHandler) diff --git a/internal/control/adminapi/routes.go b/internal/control/adminapi/routes.go index 72ca856..287b266 100644 --- a/internal/control/adminapi/routes.go +++ b/internal/control/adminapi/routes.go @@ -67,6 +67,7 @@ var PathParams = map[string]string{ "federation/{peer}": "The federation peer's name, as `llmgw ctl federation list` lists it.", "holds/{domain}": "The domain as level:name, rack:r12-07, as `llmgw ctl facility` lists it.", "collections/{id}": "The collection's ID, its name when that is unique, or ORG/NAME.", + "artifacts/{kind}": "What the job made: peft, a gzipped tar of the PEFT adapter, or gguf, the adapter for llama.cpp.", } // Routes is every endpoint Register mounts. @@ -295,6 +296,23 @@ var Routes = []Route{ Description: "No tenancy applies: this is for checking what a collection holds and finds.", Request: KBSearchRequest{}, Response: []KBHitView{}}, + // LoRA adapters trained on the nodes. + {Method: "POST", Path: "/api/v1/train/jobs", Summary: "Queue a LoRA adapter's training on a node", + Description: "The node trains when its agent is there, nothing else trains on it, and a card has gpu_memory_mib free; " + + "the pairs are NDJSON of {\"q\", \"a\"}. An adapter learns how answers sound far better than what they say.", + Request: TrainRequest{}, Response: TrainJobView{}, Status: http.StatusCreated}, + {Method: "GET", Path: "/api/v1/train/jobs", Summary: "Training jobs, newest first", + Query: []Param{{Name: "limit", Description: "How many; 50 when absent.", Type: "integer"}}, + Response: []TrainJobView{}}, + {Method: "GET", Path: "/api/v1/train/jobs/{id}", Summary: "One training job: how it goes, and what it made", + Response: TrainJobView{}}, + {Method: "POST", Path: "/api/v1/train/jobs/{id}/cancel", Summary: "Stop a training job", + Response: TrainJobView{}}, + {Method: "DELETE", Path: "/api/v1/train/jobs/{id}", Summary: "Remove an ended training job and what it made", + Response: map[string]any{}}, + {Method: "GET", Path: "/api/v1/train/jobs/{id}/artifacts/{kind}", Summary: "What a training job made: peft (a gzipped tar) or gguf", + Response: map[string]any{}}, + {Method: "POST", Path: "/api/v1/agents/restart", Summary: "Have every node's agent start again; their engines keep serving", Response: AgentRestartView{}}, diff --git a/internal/control/adminapi/train.go b/internal/control/adminapi/train.go new file mode 100644 index 0000000..7bd5d2b --- /dev/null +++ b/internal/control/adminapi/train.go @@ -0,0 +1,338 @@ +package adminapi + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "regexp" + "strconv" + "strings" + "time" + + "gitea.mynodes.xyz/ilker/llmgw/internal/control/train" + "gitea.mynodes.xyz/ilker/llmgw/internal/store" +) + +// Training LoRA adapters (BACKLOG F4), administered: a job names the model +// the adapter is for, the weights to train on, the node and how much of a +// card it needs, and the pairs; the node trains when it can, and what it +// made is kept here to download or publish. + +// TrainStore is where jobs are kept. +type TrainStore interface { + CreateTrainJob(ctx context.Context, j store.TrainJob) error + TrainJobs(ctx context.Context, limit int) ([]store.TrainJob, error) + TrainJob(ctx context.Context, id string) (store.TrainJob, error) + TrainArtifact(ctx context.Context, job, kind string) ([]byte, error) + TrainArtifactKinds(ctx context.Context, job string) ([]string, error) + DeleteTrainJob(ctx context.Context, id string) error +} + +// TrainAdmin is what the training endpoints need. +type TrainAdmin struct { + Store TrainStore + Service *train.Service +} + +// SetTrain enables the training endpoints. +func (a *API) SetTrain(t *TrainAdmin) { a.train = t } + +// TrainRequest queues a job. +type TrainRequest struct { + // Name is the adapter's: what requests will ask for it by. + Name string `json:"name"` + // Model is the catalog model the adapter is for. + Model string `json:"model"` + // BaseRepo and BaseRevision are the weights trained on: the model's own + // on its registry, in safetensors (Hugging Face), not a GGUF. + BaseRepo string `json:"base_repo"` + BaseRevision string `json:"base_revision,omitempty"` + // Node trains it; GPUMemoryMiB is what it needs of a card, which can be + // what an engine there leaves free. + Node string `json:"node"` + GPUMemoryMiB int `json:"gpu_memory_mib"` + // Data is the pairs, NDJSON of {"q", "a"} (llmgw ctl kb dataset). + Data string `json:"data"` + // Params tune the trainer: rank, alpha, dropout, epochs, lr, batch, + // max_len, seed. + Params map[string]any `json:"params,omitempty"` + // GGUF also makes the adapter in llama.cpp's format; when absent, it + // does when the model is served by llama.cpp. + GGUF *bool `json:"gguf,omitempty"` +} + +// TrainJobView is a job as listed. +type TrainJobView struct { + ID string `json:"id"` + Name string `json:"name"` + Model string `json:"model"` + BaseRepo string `json:"base_repo"` + BaseRevision string `json:"base_revision,omitempty"` + Node string `json:"node"` + GPUMemoryMiB int `json:"gpu_memory_mib"` + Params string `json:"params,omitempty"` + GGUF bool `json:"gguf"` + Pairs int `json:"pairs"` + State string `json:"state"` + Step int `json:"step"` + Steps int `json:"steps"` + Loss float64 `json:"loss"` + Message string `json:"message,omitempty"` + PeakMiB int `json:"peak_mib"` + CreatedAt time.Time `json:"created_at"` + StartedAt time.Time `json:"started_at,omitzero"` + FinishedAt time.Time `json:"finished_at,omitzero"` + CreatedBy string `json:"created_by,omitempty"` + // Artifacts are what it made: peft (a gzipped tar), gguf. + Artifacts []string `json:"artifacts,omitempty"` +} + +func trainView(j store.TrainJob) TrainJobView { + return TrainJobView{ID: j.ID, Name: j.Name, Model: j.Model, BaseRepo: j.BaseRepo, BaseRevision: j.BaseRevision, + Node: j.Node, GPUMemoryMiB: j.GPUMemoryMiB, Params: j.Params, GGUF: j.GGUF, Pairs: j.Pairs, State: j.State, + Step: j.Step, Steps: j.Steps, Loss: j.Loss, Message: j.Message, PeakMiB: j.PeakMiB, CreatedAt: j.CreatedAt, + StartedAt: j.StartedAt, FinishedAt: j.FinishedAt, CreatedBy: j.CreatedBy} +} + +func (a *API) registerTrain(mux Mux) { + mux.HandleFunc("POST /api/v1/train/jobs", a.trainCreate) + mux.HandleFunc("GET /api/v1/train/jobs", a.trainList) + mux.HandleFunc("GET /api/v1/train/jobs/{id}", a.trainGet) + mux.HandleFunc("POST /api/v1/train/jobs/{id}/cancel", a.trainCancel) + mux.HandleFunc("DELETE /api/v1/train/jobs/{id}", a.trainDelete) + mux.HandleFunc("GET /api/v1/train/jobs/{id}/artifacts/{kind}", a.trainArtifact) +} + +func (a *API) trainReady(w http.ResponseWriter) bool { + if a.train == nil { + writeError(w, http.StatusNotImplemented, "training needs the store and the nodes' agents, which this control plane runs without") + return false + } + return true +} + +// trainMaxData bounds a job's pairs. +const trainMaxData = 64 << 20 + +var adapterName = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9._-]{0,99}$`) + +// trainParams are the trainer's settings, with what each may be. +var trainParams = map[string][2]float64{ + "rank": {1, 256}, "alpha": {1, 1024}, "dropout": {0, 0.9}, "epochs": {1, 100}, + "lr": {1e-7, 1e-2}, "batch": {1, 128}, "max_len": {16, 32768}, "seed": {0, 1 << 31}, +} + +func (a *API) trainCreate(w http.ResponseWriter, r *http.Request) { + if !a.trainReady(w) { + return + } + var in TrainRequest + if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, trainMaxData+1<<20)).Decode(&in); err != nil { + writeError(w, http.StatusBadRequest, err.Error()) + return + } + bad := func(format string, args ...any) { + writeError(w, http.StatusBadRequest, fmt.Sprintf(format, args...)) + } + if !adapterName.MatchString(in.Name) { + bad("an adapter's name is letters, digits, . _ and -, up to 100: %q is not", in.Name) + return + } + m, ok := a.ctl.Catalog().Resolve(in.Model) + if !ok { + bad("no model %s in the catalog", in.Model) + return + } + if _, _, _, taken := a.ctl.Catalog().ResolveAdapter(in.Name); taken || in.Name == m.Name { + bad("%s is already a model's or an adapter's name", in.Name) + return + } + if strings.Count(in.BaseRepo, "/") != 1 || strings.ContainsAny(in.BaseRepo, " :") { + bad("base_repo is the weights' registry repository, owner/name: %q is not", in.BaseRepo) + return + } + if _, ok := a.ctl.Inventory().Get(in.Node); !ok { + bad("no node %s", in.Node) + return + } + if in.GPUMemoryMiB <= 0 { + bad("gpu_memory_mib is what the job needs of a card, in MiB") + return + } + pairs, err := countPairs(in.Data) + if err != nil { + bad("%v", err) + return + } + for k, v := range in.Params { + lim, known := trainParams[k] + f, isNum := v.(float64) + if !known || !isNum || f < lim[0] || f > lim[1] { + bad("params.%s: the trainer takes rank, alpha, dropout, epochs, lr, batch, max_len and seed, each a number in its range", k) + return + } + } + params := "" + if len(in.Params) > 0 { + b, _ := json.Marshal(in.Params) + params = string(b) + } + gguf := false + for _, p := range m.Profiles { + if p.Engine == "llamacpp" { + gguf = true + } + } + if in.GGUF != nil { + gguf = *in.GGUF + } + j := store.TrainJob{ID: store.NewID("train-"), Name: in.Name, Model: m.Name, BaseRepo: in.BaseRepo, + BaseRevision: in.BaseRevision, Node: in.Node, GPUMemoryMiB: in.GPUMemoryMiB, Params: params, GGUF: gguf, + Data: []byte(in.Data), Pairs: pairs, CreatedBy: actorOf(r)} + if err := a.train.Store.CreateTrainJob(r.Context(), j); err != nil { + writeError(w, http.StatusInternalServerError, err.Error()) + return + } + // At once rather than at the next look, when the node is here. + go a.train.Service.Dispatch(context.WithoutCancel(r.Context())) + j.State = store.TrainQueued + writeJSON(w, http.StatusCreated, trainView(j)) +} + +// countPairs checks a job's data: NDJSON of {"q", "a"}, at least one. +func countPairs(data string) (int, error) { + if len(data) > trainMaxData { + return 0, fmt.Errorf("the data is over %d MiB", trainMaxData>>20) + } + n := 0 + sc := bufio.NewScanner(strings.NewReader(data)) + sc.Buffer(make([]byte, 64<<10), 1<<20) + line := 0 + for sc.Scan() { + line++ + b := bytes.TrimSpace(sc.Bytes()) + if len(b) == 0 { + continue + } + var p struct{ Q, A string } + if err := json.Unmarshal(b, &p); err != nil { + return 0, fmt.Errorf("data line %d is not JSON: %v", line, err) + } + if strings.TrimSpace(p.Q) == "" || strings.TrimSpace(p.A) == "" { + return 0, fmt.Errorf("data line %d has no q or no a", line) + } + n++ + } + if err := sc.Err(); err != nil { + return 0, err + } + if n == 0 { + return 0, errors.New("the data holds no pairs: NDJSON of {\"q\", \"a\"}, as llmgw ctl kb dataset writes") + } + return n, nil +} + +func (a *API) trainList(w http.ResponseWriter, r *http.Request) { + if !a.trainReady(w) { + return + } + limit := 50 + if n, err := strconv.Atoi(r.URL.Query().Get("limit")); err == nil && n > 0 && n <= 500 { + limit = n + } + jobs, err := a.train.Store.TrainJobs(r.Context(), limit) + if err != nil { + writeError(w, http.StatusInternalServerError, err.Error()) + return + } + out := []TrainJobView{} + for _, j := range jobs { + out = append(out, trainView(j)) + } + writeJSON(w, http.StatusOK, out) +} + +func (a *API) trainJob(w http.ResponseWriter, r *http.Request) (store.TrainJob, bool) { + if !a.trainReady(w) { + return store.TrainJob{}, false + } + j, err := a.train.Store.TrainJob(r.Context(), r.PathValue("id")) + if errors.Is(err, store.ErrNotFound) { + writeError(w, http.StatusNotFound, "no training job "+r.PathValue("id")) + return j, false + } + if err != nil { + writeError(w, http.StatusInternalServerError, err.Error()) + return j, false + } + return j, true +} + +func (a *API) trainGet(w http.ResponseWriter, r *http.Request) { + j, ok := a.trainJob(w, r) + if !ok { + return + } + v := trainView(j) + v.Artifacts, _ = a.train.Store.TrainArtifactKinds(r.Context(), j.ID) + writeJSON(w, http.StatusOK, v) +} + +func (a *API) trainCancel(w http.ResponseWriter, r *http.Request) { + j, ok := a.trainJob(w, r) + if !ok { + return + } + if err := a.train.Service.Cancel(r.Context(), j.ID); err != nil { + writeError(w, http.StatusInternalServerError, err.Error()) + return + } + j, _ = a.train.Store.TrainJob(r.Context(), j.ID) + writeJSON(w, http.StatusOK, trainView(j)) +} + +func (a *API) trainDelete(w http.ResponseWriter, r *http.Request) { + j, ok := a.trainJob(w, r) + if !ok { + return + } + if !store.TrainOver(j.State) { + writeError(w, http.StatusConflict, "the job has not ended: cancel it first") + return + } + if err := a.train.Store.DeleteTrainJob(r.Context(), j.ID); err != nil { + writeError(w, http.StatusInternalServerError, err.Error()) + return + } + writeJSON(w, http.StatusOK, map[string]any{"id": j.ID, "deleted": true}) +} + +func (a *API) trainArtifact(w http.ResponseWriter, r *http.Request) { + j, ok := a.trainJob(w, r) + if !ok { + return + } + kind := r.PathValue("kind") + b, err := a.train.Store.TrainArtifact(r.Context(), j.ID, kind) + if errors.Is(err, store.ErrNotFound) { + writeError(w, http.StatusNotFound, "the job made no "+kind) + return + } + if err != nil { + writeError(w, http.StatusInternalServerError, err.Error()) + return + } + name := j.Name + ".tar.gz" + w.Header().Set("Content-Type", "application/gzip") + if kind == "gguf" { + name = j.Name + ".gguf" + w.Header().Set("Content-Type", "application/octet-stream") + } + w.Header().Set("Content-Disposition", `attachment; filename="`+name+`"`) + w.Header().Set("Content-Length", strconv.Itoa(len(b))) + _, _ = w.Write(b) +} diff --git a/internal/control/adminapi/train_test.go b/internal/control/adminapi/train_test.go new file mode 100644 index 0000000..10cb97c --- /dev/null +++ b/internal/control/adminapi/train_test.go @@ -0,0 +1,86 @@ +package adminapi + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "path/filepath" + "strings" + "testing" + + llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1" + "gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog" + "gitea.mynodes.xyz/ilker/llmgw/internal/control/train" + "gitea.mynodes.xyz/ilker/llmgw/internal/store" +) + +type noAgents struct{} + +func (noAgents) Send(string, *llmgwv1.ControlMessage) error { return nil } +func (noAgents) Connected(string) bool { return false } + +func TestATrainingJobIsCheckedBeforeItIsQueued(t *testing.T) { + models := append(testModels(), catalog.Model{Name: "small", Source: "/s", + Profiles: []catalog.Profile{{Name: "cpu", Engine: "llamacpp"}}}) + a, _ := setup(t, models, gpuNode("cuda01", 1)) + s, err := store.OpenSQLite(context.Background(), filepath.Join(t.TempDir(), "s.db")) + if err != nil { + t.Fatal(err) + } + defer s.Close() + a.SetTrain(&TrainAdmin{Store: s, Service: train.New(s, noAgents{}, nil)}) + mux := http.NewServeMux() + a.Register(mux) + post := func(body string) (int, string) { + rec := httptest.NewRecorder() + mux.ServeHTTP(rec, httptest.NewRequest("POST", "/api/v1/train/jobs", strings.NewReader(body))) + return rec.Code, rec.Body.String() + } + good := map[string]any{"name": "tone", "model": "small", "base_repo": "Qwen/Qwen2.5-0.5B-Instruct", "node": "cuda01", + "gpu_memory_mib": 2300, "data": "{\"q\":\"a\",\"a\":\"b\"}\n{\"q\":\"c\",\"a\":\"d\"}\n", "params": map[string]any{"epochs": 3}} + with := func(k string, v any) string { + m := map[string]any{} + for kk, vv := range good { + m[kk] = vv + } + m[k] = v + b, _ := json.Marshal(m) + return string(b) + } + for _, c := range []struct { + field string + value any + says string + }{ + {"name", "bad name", "letters, digits"}, + {"model", "nope", "no model nope"}, + {"name", "small", "already a model"}, + {"base_repo", "Qwen", "owner/name"}, + {"node", "gone", "no node gone"}, + {"data", "{\"q\":\"a\"}\n", "no q or no a"}, + {"data", "", "no pairs"}, + {"params", map[string]any{"epochs": 0}, "params.epochs"}, + {"params", map[string]any{"momentum": 0.9}, "params.momentum"}, + } { + if code, body := post(with(c.field, c.value)); code != http.StatusBadRequest || !strings.Contains(body, c.says) { + t.Errorf("%s=%v: %d %s", c.field, c.value, code, body) + } + } + code, body := post(with("name", "tone")) + var v TrainJobView + _ = json.Unmarshal([]byte(body), &v) + if code != http.StatusCreated || v.Pairs != 2 || !v.GGUF || v.State != "queued" || v.Params != `{"epochs":3}` { + t.Fatalf("%d %s", code, body) + } + rec := httptest.NewRecorder() + mux.ServeHTTP(rec, httptest.NewRequest("DELETE", "/api/v1/train/jobs/"+v.ID, nil)) + if rec.Code != http.StatusConflict { + t.Errorf("deleting a queued job: %d", rec.Code) + } + rec = httptest.NewRecorder() + mux.ServeHTTP(rec, httptest.NewRequest("GET", "/api/v1/train/jobs/"+v.ID+"/artifacts/peft", nil)) + if rec.Code != http.StatusNotFound { + t.Errorf("an artifact not made: %d", rec.Code) + } +} diff --git a/internal/control/agentsvc/service.go b/internal/control/agentsvc/service.go index 496c069..7bfeb3f 100644 --- a/internal/control/agentsvc/service.go +++ b/internal/control/agentsvc/service.go @@ -82,6 +82,8 @@ type Service struct { // agentSettings is what each node's agent is told to run on; see // SetAgentSettings. agentSettings func(node string) string + // train is told how training jobs go; see SetTrain. + train TrainHandler // logWaits are the log requests awaiting a reply, by request id. logMu sync.Mutex @@ -368,6 +370,9 @@ func (s *Service) readLoop(ctx context.Context, msg, err := stream.Receive() if err != nil { c.close() + if h := s.trainHandler(); h != nil { + h.OnNodeGone(nodeID) + } if ctx.Err() != nil || errors.Is(err, io.EOF) { s.log.Info("agent disconnected", "node", nodeID) return @@ -505,6 +510,16 @@ func (s *Service) handle(nodeID string, msg *llmgwv1.AgentMessage) { ch <- m.Upgrade } + case *llmgwv1.AgentMessage_TrainStatus: + if h := s.trainHandler(); h != nil { + h.OnTrainStatus(nodeID, m.TrainStatus) + } + + case *llmgwv1.AgentMessage_TrainArtifact: + if h := s.trainHandler(); h != nil { + h.OnTrainArtifact(nodeID, m.TrainArtifact) + } + case *llmgwv1.AgentMessage_Logs: s.logMu.Lock() ch, ok := s.logWaits[m.Logs.GetRequestId()] diff --git a/internal/control/agentsvc/train.go b/internal/control/agentsvc/train.go new file mode 100644 index 0000000..e9a32e1 --- /dev/null +++ b/internal/control/agentsvc/train.go @@ -0,0 +1,35 @@ +package agentsvc + +import ( + "slices" + + llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1" +) + +// TrainHandler is told how training jobs on the nodes go (BACKLOG F4). +type TrainHandler interface { + OnTrainStatus(node string, st *llmgwv1.TrainStatus) + OnTrainArtifact(node string, a *llmgwv1.TrainArtifact) + // OnNodeGone is a node's stream ending: a job running there ended with + // it. + OnNodeGone(node string) +} + +// SetTrain gives the service what training jobs' messages go to. +func (s *Service) SetTrain(h TrainHandler) { + s.mu.Lock() + s.train = h + s.mu.Unlock() +} + +func (s *Service) trainHandler() TrainHandler { + s.mu.RLock() + defer s.mu.RUnlock() + return s.train +} + +// Send gives a connected node's agent a command. +func (s *Service) Send(node string, msg *llmgwv1.ControlMessage) error { return s.send(node, msg) } + +// Connected is whether a node's agent holds a stream to this control plane. +func (s *Service) Connected(node string) bool { return slices.Contains(s.ConnectedNodes(), node) } diff --git a/internal/control/train/service.go b/internal/control/train/service.go new file mode 100644 index 0000000..b2a646d --- /dev/null +++ b/internal/control/train/service.go @@ -0,0 +1,196 @@ +// Package train runs LoRA training jobs from the control plane (BACKLOG F4): +// queued in the store, sent to the node each names when its agent is there +// and has nothing else training, followed as the agent reports, and what it +// makes kept in the store. +package train + +import ( + "context" + "log/slog" + "sync" + "time" + + llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1" + "gitea.mynodes.xyz/ilker/llmgw/internal/store" +) + +// Store is where jobs are kept. +type Store interface { + QueuedTrainJobs(ctx context.Context) ([]store.TrainJob, error) + ActiveTrainJobs(ctx context.Context) ([]store.TrainJob, error) + TrainJob(ctx context.Context, id string) (store.TrainJob, error) + UpdateTrainJob(ctx context.Context, j store.TrainJob) error + PutTrainArtifact(ctx context.Context, job, kind string, data []byte) error +} + +// Agents is how jobs reach the nodes. +type Agents interface { + Send(node string, msg *llmgwv1.ControlMessage) error + Connected(node string) bool +} + +// Service sends queued jobs and records how they go. +type Service struct { + store Store + agents Agents + log *slog.Logger + + mu sync.Mutex + // pieces are artifacts arriving, by job and kind. + pieces map[[2]string][]byte +} + +// New is a service. +func New(s Store, a Agents, log *slog.Logger) *Service { + if log == nil { + log = slog.New(slog.DiscardHandler) + } + return &Service{store: s, agents: a, log: log, pieces: map[[2]string][]byte{}} +} + +// Every is how often queued jobs are looked at. +var Every = 5 * time.Second + +// Run sends queued jobs until ctx ends. +func (s *Service) Run(ctx context.Context) { + t := time.NewTicker(Every) + defer t.Stop() + for { + s.Dispatch(ctx) + select { + case <-ctx.Done(): + return + case <-t.C: + } + } +} + +// Dispatch sends each queued job whose node is here and not already +// training, one per node. +func (s *Service) Dispatch(ctx context.Context) { + queued, err := s.store.QueuedTrainJobs(ctx) + if err != nil || len(queued) == 0 { + return + } + active, err := s.store.ActiveTrainJobs(ctx) + if err != nil { + return + } + busy := map[string]bool{} + for _, j := range active { + busy[j.Node] = true + } + for _, q := range queued { + if busy[q.Node] || !s.agents.Connected(q.Node) { + continue + } + j, err := s.store.TrainJob(ctx, q.ID) + if err != nil { + continue + } + err = s.agents.Send(j.Node, &llmgwv1.ControlMessage{Msg: &llmgwv1.ControlMessage_TrainStart{TrainStart: &llmgwv1.TrainStart{ + JobId: j.ID, BaseRepo: j.BaseRepo, BaseRevision: j.BaseRevision, Data: j.Data, Params: j.Params, + GpuMemoryMib: int32(j.GPUMemoryMiB), Gguf: j.GGUF}}}) //nolint:gosec // MiB + if err != nil { + s.log.Warn("could not send a training job to its node", "job", j.ID, "node", j.Node, "err", err) + continue + } + j.State, j.StartedAt, j.Message = store.TrainStarting, time.Now(), "sent to "+j.Node + _ = s.store.UpdateTrainJob(ctx, j) + busy[j.Node] = true + s.log.Info("training job sent", "job", j.ID, "node", j.Node, "adapter", j.Name) + } +} + +// OnTrainStatus records how a job goes, as its node says. +func (s *Service) OnTrainStatus(node string, st *llmgwv1.TrainStatus) { + ctx := context.Background() + j, err := s.store.TrainJob(ctx, st.GetJobId()) + if err != nil || j.Node != node || store.TrainOver(j.State) { + return + } + j.State = st.GetState() + if m := st.GetMessage(); m != "" || j.State == store.TrainDone { + j.Message = m + } + if st.GetSteps() > 0 { + j.Step, j.Steps = int(st.GetStep()), int(st.GetSteps()) + } + if st.GetLoss() > 0 { + j.Loss = st.GetLoss() + } + if p := int(st.GetPeakMib()); p > j.PeakMiB { + j.PeakMiB = p + } + if store.TrainOver(j.State) { + j.FinishedAt = time.Now() + s.mu.Lock() + for k := range s.pieces { + if k[0] == j.ID { + delete(s.pieces, k) + } + } + s.mu.Unlock() + s.log.Info("training job ended", "job", j.ID, "state", j.State, "message", j.Message) + } + if err := s.store.UpdateTrainJob(ctx, j); err != nil { + s.log.Warn("could not record a training job's progress", "job", j.ID, "err", err) + } +} + +// OnTrainArtifact gathers what a job made, and keeps it when it is whole. +func (s *Service) OnTrainArtifact(node string, a *llmgwv1.TrainArtifact) { + key := [2]string{a.GetJobId(), a.GetKind()} + s.mu.Lock() + if a.GetSeq() == 0 { + s.pieces[key] = nil + } + s.pieces[key] = append(s.pieces[key], a.GetData()...) + whole := s.pieces[key] + if a.GetLast() { + delete(s.pieces, key) + } + s.mu.Unlock() + if !a.GetLast() { + return + } + if err := s.store.PutTrainArtifact(context.Background(), key[0], key[1], whole); err != nil { + s.log.Warn("could not keep what a training job made", "job", key[0], "kind", key[1], "err", err) + } +} + +// OnNodeGone fails the jobs a node was running: its agent's stream ended, +// and the job with it. +func (s *Service) OnNodeGone(node string) { + ctx := context.Background() + active, err := s.store.ActiveTrainJobs(ctx) + if err != nil { + return + } + for _, j := range active { + if j.Node != node { + continue + } + j.State, j.FinishedAt = store.TrainFailed, time.Now() + j.Message = "the node's agent went away while it trained; start the job again" + _ = s.store.UpdateTrainJob(ctx, j) + } +} + +// Cancel stops a job: one still queued is not sent, one running is stopped +// on its node. +func (s *Service) Cancel(ctx context.Context, id string) error { + j, err := s.store.TrainJob(ctx, id) + if err != nil { + return err + } + if store.TrainOver(j.State) { + return nil + } + if j.State != store.TrainQueued && s.agents.Connected(j.Node) { + _ = s.agents.Send(j.Node, &llmgwv1.ControlMessage{Msg: &llmgwv1.ControlMessage_TrainCancel{ + TrainCancel: &llmgwv1.TrainCancel{JobId: id}}}) + } + j.State, j.FinishedAt, j.Message = store.TrainCancelled, time.Now(), "cancelled" + return s.store.UpdateTrainJob(ctx, j) +} diff --git a/internal/control/train/service_test.go b/internal/control/train/service_test.go new file mode 100644 index 0000000..5a7a8f0 --- /dev/null +++ b/internal/control/train/service_test.go @@ -0,0 +1,124 @@ +package train + +import ( + "context" + "path/filepath" + "sync" + "testing" + + llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1" + "gitea.mynodes.xyz/ilker/llmgw/internal/store" +) + +type fakeAgents struct { + mu sync.Mutex + connected map[string]bool + sent []*llmgwv1.ControlMessage +} + +func (f *fakeAgents) Send(_ string, m *llmgwv1.ControlMessage) error { + f.mu.Lock() + defer f.mu.Unlock() + f.sent = append(f.sent, m) + return nil +} +func (f *fakeAgents) Connected(n string) bool { + f.mu.Lock() + defer f.mu.Unlock() + return f.connected[n] +} + +func setup(t *testing.T) (*Service, *store.DB, *fakeAgents) { + t.Helper() + s, err := store.OpenSQLite(context.Background(), filepath.Join(t.TempDir(), "s.db")) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = s.Close() }) + a := &fakeAgents{connected: map[string]bool{}} + return New(s, a, nil), s, a +} + +func job(t *testing.T, s *store.DB, id, node string) { + t.Helper() + if err := s.CreateTrainJob(context.Background(), store.TrainJob{ID: id, Name: "a-" + id, Model: "m", BaseRepo: "o/n", + Node: node, GPUMemoryMiB: 2000, Data: []byte(`{"q":"x","a":"y"}`), Pairs: 1, GGUF: true}); err != nil { + t.Fatal(err) + } +} + +func TestAJobGoesToItsNodeWhenItIsHereAndFree(t *testing.T) { + ctx := context.Background() + svc, s, a := setup(t) + job(t, s, "j1", "cuda01") + job(t, s, "j2", "cuda01") + svc.Dispatch(ctx) + if len(a.sent) != 0 { + t.Fatal("sent to a node that is not here") + } + a.connected["cuda01"] = true + svc.Dispatch(ctx) + svc.Dispatch(ctx) + if len(a.sent) != 1 { + t.Fatalf("sent %d jobs to one node", len(a.sent)) + } + start := a.sent[0].GetTrainStart() + if start.GetJobId() != "j1" || start.GetGpuMemoryMib() != 2000 || !start.GetGguf() || string(start.GetData()) == "" { + t.Errorf("sent %+v", start) + } + if j, _ := s.TrainJob(ctx, "j1"); j.State != store.TrainStarting || j.StartedAt.IsZero() { + t.Errorf("j1 %+v", j) + } + + // Progress, then what it made in pieces, then done: the next is sent. + svc.OnTrainStatus("cuda01", &llmgwv1.TrainStatus{JobId: "j1", State: "running", Step: 5, Steps: 10, Loss: 0.8, PeakMib: 1900}) + svc.OnTrainArtifact("cuda01", &llmgwv1.TrainArtifact{JobId: "j1", Kind: "gguf", Seq: 0, Data: []byte("GG")}) + svc.OnTrainArtifact("cuda01", &llmgwv1.TrainArtifact{JobId: "j1", Kind: "gguf", Seq: 1, Data: []byte("UF"), Last: true}) + svc.OnTrainStatus("cuda01", &llmgwv1.TrainStatus{JobId: "j1", State: "done"}) + j, _ := s.TrainJob(ctx, "j1") + if j.State != store.TrainDone || j.Step != 5 || j.PeakMiB != 1900 || j.Loss != 0.8 || j.FinishedAt.IsZero() { + t.Errorf("j1 after %+v", j) + } + if b, _ := s.TrainArtifact(ctx, "j1", "gguf"); string(b) != "GGUF" { + t.Errorf("artifact %q", b) + } + // A late message does not bring an ended job back. + svc.OnTrainStatus("cuda01", &llmgwv1.TrainStatus{JobId: "j1", State: "running"}) + if j, _ := s.TrainJob(ctx, "j1"); j.State != store.TrainDone { + t.Errorf("revived: %s", j.State) + } + // Nor does another node speak for it. + svc.Dispatch(ctx) + if len(a.sent) != 2 || a.sent[1].GetTrainStart().GetJobId() != "j2" { + t.Fatalf("j2 not sent: %d", len(a.sent)) + } + svc.OnTrainStatus("node01", &llmgwv1.TrainStatus{JobId: "j2", State: "done"}) + if j, _ := s.TrainJob(ctx, "j2"); j.State == store.TrainDone { + t.Error("another node ended the job") + } + // The node goes: its job fails, saying why. + svc.OnNodeGone("cuda01") + if j, _ := s.TrainJob(ctx, "j2"); j.State != store.TrainFailed || j.Message == "" { + t.Errorf("after the node went: %+v", j) + } +} + +func TestCancellingAJob(t *testing.T) { + ctx := context.Background() + svc, s, a := setup(t) + job(t, s, "j1", "cuda01") + if err := svc.Cancel(ctx, "j1"); err != nil { + t.Fatal(err) + } + a.connected["cuda01"] = true + svc.Dispatch(ctx) + if len(a.sent) != 0 { + t.Error("a cancelled job was sent") + } + job(t, s, "j2", "cuda01") + svc.Dispatch(ctx) + _ = svc.Cancel(ctx, "j2") + if len(a.sent) != 2 || a.sent[1].GetTrainCancel().GetJobId() != "j2" { + t.Errorf("the node was not told: %+v", a.sent) + } +} diff --git a/internal/store/copy.go b/internal/store/copy.go index 96767f8..d0eb849 100644 --- a/internal/store/copy.go +++ b/internal/store/copy.go @@ -73,6 +73,9 @@ var copyOrder = []string{ "eval_runs", "eval_cases", "federation_peers", + // Before their artifacts, which point at them. + "train_jobs", + "train_artifacts", } // notCopied are tables deliberately left behind. diff --git a/internal/store/migrations/postgres/0037_train_jobs.sql b/internal/store/migrations/postgres/0037_train_jobs.sql new file mode 100644 index 0000000..ccf7b40 --- /dev/null +++ b/internal/store/migrations/postgres/0037_train_jobs.sql @@ -0,0 +1,42 @@ +-- Training jobs (BACKLOG F4): a LoRA adapter trained on a node, from a base +-- model's weights and question/answer pairs, and what it made. The pairs +-- and the adapter are kept here, in the store every control plane reads: +-- they are megabytes, and a job's record is worth little without them. +CREATE TABLE train_jobs ( + id TEXT PRIMARY KEY, + -- name is the adapter's: what clients will ask for it by. + name TEXT NOT NULL, + -- model is the catalog model the adapter is for; base_repo and + -- base_revision the weights it is trained on. + model TEXT NOT NULL, + base_repo TEXT NOT NULL, + base_revision TEXT NOT NULL DEFAULT '', + node TEXT NOT NULL, + gpu_memory_mib INTEGER NOT NULL, + params TEXT NOT NULL DEFAULT '', + gguf INTEGER NOT NULL DEFAULT 0, + data BYTEA NOT NULL, + pairs INTEGER NOT NULL, + -- state is queued, starting, preparing, running, done, failed or + -- cancelled. + state TEXT NOT NULL, + step INTEGER NOT NULL DEFAULT 0, + steps INTEGER NOT NULL DEFAULT 0, + loss DOUBLE PRECISION NOT NULL DEFAULT 0, + message TEXT NOT NULL DEFAULT '', + peak_mib INTEGER NOT NULL DEFAULT 0, + created_at BIGINT NOT NULL, + started_at BIGINT NOT NULL DEFAULT 0, + finished_at BIGINT NOT NULL DEFAULT 0, + created_by TEXT NOT NULL DEFAULT '' +); +CREATE INDEX train_jobs_by_state ON train_jobs (state); + +-- What a job made: the PEFT adapter as a gzipped tar, and for llama.cpp the +-- adapter as GGUF. +CREATE TABLE train_artifacts ( + job_id TEXT NOT NULL REFERENCES train_jobs(id) ON DELETE CASCADE, + kind TEXT NOT NULL, + data BYTEA NOT NULL, + PRIMARY KEY (job_id, kind) +); diff --git a/internal/store/migrations/sqlite/0037_train_jobs.sql b/internal/store/migrations/sqlite/0037_train_jobs.sql new file mode 100644 index 0000000..09b4234 --- /dev/null +++ b/internal/store/migrations/sqlite/0037_train_jobs.sql @@ -0,0 +1,42 @@ +-- Training jobs (BACKLOG F4): a LoRA adapter trained on a node, from a base +-- model's weights and question/answer pairs, and what it made. The pairs +-- and the adapter are kept here, in the store every control plane reads: +-- they are megabytes, and a job's record is worth little without them. +CREATE TABLE train_jobs ( + id TEXT PRIMARY KEY, + -- name is the adapter's: what clients will ask for it by. + name TEXT NOT NULL, + -- model is the catalog model the adapter is for; base_repo and + -- base_revision the weights it is trained on. + model TEXT NOT NULL, + base_repo TEXT NOT NULL, + base_revision TEXT NOT NULL DEFAULT '', + node TEXT NOT NULL, + gpu_memory_mib INTEGER NOT NULL, + params TEXT NOT NULL DEFAULT '', + gguf INTEGER NOT NULL DEFAULT 0, + data BLOB NOT NULL, + pairs INTEGER NOT NULL, + -- state is queued, starting, preparing, running, done, failed or + -- cancelled. + state TEXT NOT NULL, + step INTEGER NOT NULL DEFAULT 0, + steps INTEGER NOT NULL DEFAULT 0, + loss REAL NOT NULL DEFAULT 0, + message TEXT NOT NULL DEFAULT '', + peak_mib INTEGER NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + started_at INTEGER NOT NULL DEFAULT 0, + finished_at INTEGER NOT NULL DEFAULT 0, + created_by TEXT NOT NULL DEFAULT '' +); +CREATE INDEX train_jobs_by_state ON train_jobs (state); + +-- What a job made: the PEFT adapter as a gzipped tar, and for llama.cpp the +-- adapter as GGUF. +CREATE TABLE train_artifacts ( + job_id TEXT NOT NULL REFERENCES train_jobs(id) ON DELETE CASCADE, + kind TEXT NOT NULL, + data BLOB NOT NULL, + PRIMARY KEY (job_id, kind) +); diff --git a/internal/store/train.go b/internal/store/train.go new file mode 100644 index 0000000..6f0c7f5 --- /dev/null +++ b/internal/store/train.go @@ -0,0 +1,177 @@ +package store + +import ( + "context" + "database/sql" + "errors" + "time" +) + +// TrainJob is a LoRA adapter trained on a node (BACKLOG F4). +type TrainJob struct { + ID string + Name string + Model string + BaseRepo string + BaseRevision string + Node string + GPUMemoryMiB int + Params string + GGUF bool + // Data is the pairs, NDJSON; Pairs how many. + Data []byte + Pairs int + + State string + Step int + Steps int + Loss float64 + Message string + PeakMiB int + CreatedAt time.Time + StartedAt time.Time + FinishedAt time.Time + CreatedBy string +} + +// The states of a training job. +const ( + TrainQueued = "queued" + TrainStarting = "starting" + TrainDone = "done" + TrainFailed = "failed" + TrainCancelled = "cancelled" +) + +// TrainOver is whether a state is one a job ends in. +func TrainOver(state string) bool { + return state == TrainDone || state == TrainFailed || state == TrainCancelled +} + +// CreateTrainJob queues a job. +func (s *DB) CreateTrainJob(ctx context.Context, j TrainJob) error { + if j.CreatedAt.IsZero() { + j.CreatedAt = time.Now() + } + gguf := 0 + if j.GGUF { + gguf = 1 + } + _, err := s.db.ExecContext(ctx, `INSERT INTO train_jobs (id, name, model, base_repo, base_revision, node, gpu_memory_mib, + params, gguf, data, pairs, state, created_at, created_by) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + j.ID, j.Name, j.Model, j.BaseRepo, j.BaseRevision, j.Node, j.GPUMemoryMiB, j.Params, gguf, j.Data, j.Pairs, + TrainQueued, millis(j.CreatedAt), j.CreatedBy) + return err +} + +const trainCols = `id, name, model, base_repo, base_revision, node, gpu_memory_mib, params, gguf, pairs, state, step, steps, + loss, message, peak_mib, created_at, started_at, finished_at, created_by` + +func scanTrainJob(sc scanner) (TrainJob, error) { + var j TrainJob + var gguf int + var created, started, finished int64 + err := sc.Scan(&j.ID, &j.Name, &j.Model, &j.BaseRepo, &j.BaseRevision, &j.Node, &j.GPUMemoryMiB, &j.Params, &gguf, + &j.Pairs, &j.State, &j.Step, &j.Steps, &j.Loss, &j.Message, &j.PeakMiB, &created, &started, &finished, &j.CreatedBy) + j.GGUF = gguf != 0 + j.CreatedAt = fromMillis(created) + if started > 0 { + j.StartedAt = fromMillis(started) + } + if finished > 0 { + j.FinishedAt = fromMillis(finished) + } + return j, err +} + +// TrainJobs lists jobs, newest first. +func (s *DB) TrainJobs(ctx context.Context, limit int) ([]TrainJob, error) { + return queryAll(ctx, s.db, `SELECT `+trainCols+` FROM train_jobs ORDER BY created_at DESC LIMIT ?`, []any{limit}, + func(r *sql.Rows) (TrainJob, error) { return scanTrainJob(r) }) +} + +// TrainJob reads one job, with its data. +func (s *DB) TrainJob(ctx context.Context, id string) (TrainJob, error) { + j, err := scanTrainJob(s.db.QueryRowContext(ctx, `SELECT `+trainCols+` FROM train_jobs WHERE id = ?`, id)) + if errors.Is(err, sql.ErrNoRows) { + return TrainJob{}, ErrNotFound + } + if err != nil { + return j, err + } + err = s.db.QueryRowContext(ctx, `SELECT data FROM train_jobs WHERE id = ?`, id).Scan(&j.Data) + return j, err +} + +// QueuedTrainJobs are the jobs waiting for their node, oldest first. +func (s *DB) QueuedTrainJobs(ctx context.Context) ([]TrainJob, error) { + return queryAll(ctx, s.db, `SELECT `+trainCols+` FROM train_jobs WHERE state = ? ORDER BY created_at`, []any{TrainQueued}, + func(r *sql.Rows) (TrainJob, error) { return scanTrainJob(r) }) +} + +// ActiveTrainJobs are the jobs sent to a node and not ended. +func (s *DB) ActiveTrainJobs(ctx context.Context) ([]TrainJob, error) { + return queryAll(ctx, s.db, `SELECT `+trainCols+` FROM train_jobs WHERE state NOT IN (?, ?, ?, ?)`, + []any{TrainQueued, TrainDone, TrainFailed, TrainCancelled}, + func(r *sql.Rows) (TrainJob, error) { return scanTrainJob(r) }) +} + +// UpdateTrainJob records how a job goes; a job that has ended is not +// changed again, so a late message cannot bring one back. +func (s *DB) UpdateTrainJob(ctx context.Context, j TrainJob) error { + var started, finished int64 + if !j.StartedAt.IsZero() { + started = millis(j.StartedAt) + } + if !j.FinishedAt.IsZero() { + finished = millis(j.FinishedAt) + } + _, err := s.db.ExecContext(ctx, `UPDATE train_jobs SET state = ?, step = ?, steps = ?, loss = ?, message = ?, + peak_mib = ?, started_at = ?, finished_at = ? WHERE id = ? AND state NOT IN (?, ?, ?)`, + j.State, j.Step, j.Steps, j.Loss, j.Message, j.PeakMiB, started, finished, j.ID, + TrainDone, TrainFailed, TrainCancelled) + return err +} + +// PutTrainArtifact keeps what a job made. +func (s *DB) PutTrainArtifact(ctx context.Context, job, kind string, data []byte) error { + _, err := s.db.ExecContext(ctx, `INSERT INTO train_artifacts (job_id, kind, data) VALUES (?, ?, ?) + ON CONFLICT (job_id, kind) DO UPDATE SET data = excluded.data`, job, kind, data) + return err +} + +// TrainArtifact reads what a job made, of one kind. +func (s *DB) TrainArtifact(ctx context.Context, job, kind string) ([]byte, error) { + var b []byte + err := s.db.QueryRowContext(ctx, `SELECT data FROM train_artifacts WHERE job_id = ? AND kind = ?`, job, kind).Scan(&b) + if errors.Is(err, sql.ErrNoRows) { + return nil, ErrNotFound + } + return b, err +} + +// TrainArtifactKinds lists what a job made. +func (s *DB) TrainArtifactKinds(ctx context.Context, job string) ([]string, error) { + return queryAll(ctx, s.db, `SELECT kind FROM train_artifacts WHERE job_id = ? ORDER BY kind`, []any{job}, + func(r *sql.Rows) (string, error) { + var k string + return k, r.Scan(&k) + }) +} + +// DeleteTrainJob removes a job and what it made. +func (s *DB) DeleteTrainJob(ctx context.Context, id string) error { + return s.inTx(ctx, func(c conn) error { + if _, err := c.ExecContext(ctx, `DELETE FROM train_artifacts WHERE job_id = ?`, id); err != nil { + return err + } + res, err := c.ExecContext(ctx, `DELETE FROM train_jobs WHERE id = ?`, id) + if err != nil { + return err + } + if n, _ := res.RowsAffected(); n == 0 { + return ErrNotFound + } + return nil + }) +}