diff --git a/BACKLOG.md b/BACKLOG.md index c7502b1..4122a16 100644 --- a/BACKLOG.md +++ b/BACKLOG.md @@ -2107,7 +2107,16 @@ lent, partition`); the leader tells each named node's agent (root) what to be, t controller ~20 s later, sbatch ran on n1 as ada. Disabled in AD: suspended in 56 s, the Linux user gone in 89 s, ssh refused, home kept. Enabled: back in 32 s, same UID, key kept. Taken out of ml-team: out of the team in 11 s; put back: in again in 56 s. -- P6, Slurm's jobs in llmgw's accounting (written 2026-11-12, not started). The tenancy tree to Slurm: +- P6, done 2026-11-15 but for budgets: each project a Slurm account (org and team accounts above it, + ids keyed by kind: the first install's org, team and project share the id "default"), its members and its + team's and org's on it; AccountingStorageEnforce=associations; the controller's agent keeps them with + sacctmgr (only what it made; a user's default account moved before an association goes, the user taken out + when none is left) and reports the jobs that ended in the last day; the leader records each once + (slurm_jobs, keyed cluster/id) with a usage row, model slurm/PARTITION, priced at gpu_hour_cost or + core_hour_cost. Tried on the VMs: accounts made, ayse's job on her project's account ran and another + project's was refused, both jobs in deneme's usage at 345 and 0 micro-units. Left: budgets (a job's cost + is not counted against a budget, which binds keys; GrpTRESMins from a project's budget would hold its jobs). +- P6 as written 2026-11-12: The tenancy tree to Slurm: each project an account (`sacctmgr add account`, parent its team and organization), each member a user of it, kept by the leader from the store; `sbatch --account=PROJECT`. The records back: the leader reads `sacct --json` for jobs ended since the last read (through the controller's agent) into a diff --git a/CHANGELOG.md b/CHANGELOG.md index 1971ba7..d1fb8ed 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -65,6 +65,13 @@ APIs change between minors without a deprecation period. there, their session put in the job (pam_slurm_adopt), so nobody uses a node's cards outside Slurm, nor on a node llmgw serves from. Root and the machines' own users are not held to it. +- **Slurm's jobs in llmgw's accounting.** With `slurm.users: llmgw` each + project is a Slurm account (organizations and teams above it), its members + and its team's and organization's on it, and a job must run on one of its + user's (`AccountingStorageEnforce=associations`). Each job that ends is + recorded once, its GPU-hours (or core-hours) priced at `auth.gpu_hour_cost` + into its project's usage as model `slurm/PARTITION`. Accounts a Slurm + administrator made are not touched. - **Prices from what the cluster costs.** `auth.gpu_hour_cost` (and `core_hour_cost`) is what an hour of a GPU (a core) costs; the leader keeps the hours each model's replicas hold, and `llmgw ctl prices suggest` prices diff --git a/api/gen/llmgw/v1/agent.pb.go b/api/gen/llmgw/v1/agent.pb.go index 6e4945a..32ef177 100644 --- a/api/gen/llmgw/v1/agent.pb.go +++ b/api/gen/llmgw/v1/agent.pb.go @@ -5317,9 +5317,14 @@ type Slurm struct { MakeHomes bool `protobuf:"varint,11,opt,name=make_homes,json=makeHomes,proto3" json:"make_homes,omitempty"` // containers has the machine install pyxis and enroot, so jobs may run // in containers (srun --container-image), and set them up. - Containers bool `protobuf:"varint,12,opt,name=containers,proto3" json:"containers,omitempty"` - unknownFields protoimpl.UnknownFields - sizeCache protoimpl.SizeCache + Containers bool `protobuf:"varint,12,opt,name=containers,proto3" json:"containers,omitempty"` + // manage_accounts has the controller keep Slurm's accounts as accounts + // says: one for each organization, team and project, nested, the users + // on each project's. Only those llmgw made are changed or taken away. + ManageAccounts bool `protobuf:"varint,13,opt,name=manage_accounts,json=manageAccounts,proto3" json:"manage_accounts,omitempty"` + Accounts []*SlurmAccount `protobuf:"bytes,14,rep,name=accounts,proto3" json:"accounts,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache } func (x *Slurm) Reset() { @@ -5436,6 +5441,199 @@ func (x *Slurm) GetContainers() bool { return false } +func (x *Slurm) GetManageAccounts() bool { + if x != nil { + return x.ManageAccounts + } + return false +} + +func (x *Slurm) GetAccounts() []*SlurmAccount { + if x != nil { + return x.Accounts + } + return nil +} + +// SlurmAccount is an account in Slurm's accounting, parents before their +// children. +type SlurmAccount struct { + state protoimpl.MessageState `protogen:"open.v1"` + Name string `protobuf:"bytes,1,opt,name=name,proto3" json:"name,omitempty"` + Parent string `protobuf:"bytes,2,opt,name=parent,proto3" json:"parent,omitempty"` + Users []string `protobuf:"bytes,3,rep,name=users,proto3" json:"users,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *SlurmAccount) Reset() { + *x = SlurmAccount{} + mi := &file_llmgw_v1_agent_proto_msgTypes[49] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *SlurmAccount) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*SlurmAccount) ProtoMessage() {} + +func (x *SlurmAccount) ProtoReflect() protoreflect.Message { + mi := &file_llmgw_v1_agent_proto_msgTypes[49] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use SlurmAccount.ProtoReflect.Descriptor instead. +func (*SlurmAccount) Descriptor() ([]byte, []int) { + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{49} +} + +func (x *SlurmAccount) GetName() string { + if x != nil { + return x.Name + } + return "" +} + +func (x *SlurmAccount) GetParent() string { + if x != nil { + return x.Parent + } + return "" +} + +func (x *SlurmAccount) GetUsers() []string { + if x != nil { + return x.Users + } + return nil +} + +// SlurmJob is a job that ended, as sacct has it, for llmgw's accounting. +type SlurmJob struct { + state protoimpl.MessageState `protogen:"open.v1"` + Id string `protobuf:"bytes,1,opt,name=id,proto3" json:"id,omitempty"` + User string `protobuf:"bytes,2,opt,name=user,proto3" json:"user,omitempty"` + Account string `protobuf:"bytes,3,opt,name=account,proto3" json:"account,omitempty"` + Partition string `protobuf:"bytes,4,opt,name=partition,proto3" json:"partition,omitempty"` + State string `protobuf:"bytes,5,opt,name=state,proto3" json:"state,omitempty"` + StartUnix int64 `protobuf:"varint,6,opt,name=start_unix,json=startUnix,proto3" json:"start_unix,omitempty"` + EndUnix int64 `protobuf:"varint,7,opt,name=end_unix,json=endUnix,proto3" json:"end_unix,omitempty"` + ElapsedSeconds int64 `protobuf:"varint,8,opt,name=elapsed_seconds,json=elapsedSeconds,proto3" json:"elapsed_seconds,omitempty"` + Cpus int32 `protobuf:"varint,9,opt,name=cpus,proto3" json:"cpus,omitempty"` + Gpus int32 `protobuf:"varint,10,opt,name=gpus,proto3" json:"gpus,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *SlurmJob) Reset() { + *x = SlurmJob{} + mi := &file_llmgw_v1_agent_proto_msgTypes[50] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *SlurmJob) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*SlurmJob) ProtoMessage() {} + +func (x *SlurmJob) ProtoReflect() protoreflect.Message { + mi := &file_llmgw_v1_agent_proto_msgTypes[50] + 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 SlurmJob.ProtoReflect.Descriptor instead. +func (*SlurmJob) Descriptor() ([]byte, []int) { + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{50} +} + +func (x *SlurmJob) GetId() string { + if x != nil { + return x.Id + } + return "" +} + +func (x *SlurmJob) GetUser() string { + if x != nil { + return x.User + } + return "" +} + +func (x *SlurmJob) GetAccount() string { + if x != nil { + return x.Account + } + return "" +} + +func (x *SlurmJob) GetPartition() string { + if x != nil { + return x.Partition + } + return "" +} + +func (x *SlurmJob) GetState() string { + if x != nil { + return x.State + } + return "" +} + +func (x *SlurmJob) GetStartUnix() int64 { + if x != nil { + return x.StartUnix + } + return 0 +} + +func (x *SlurmJob) GetEndUnix() int64 { + if x != nil { + return x.EndUnix + } + return 0 +} + +func (x *SlurmJob) GetElapsedSeconds() int64 { + if x != nil { + return x.ElapsedSeconds + } + return 0 +} + +func (x *SlurmJob) GetCpus() int32 { + if x != nil { + return x.Cpus + } + return 0 +} + +func (x *SlurmJob) GetGpus() int32 { + if x != nil { + return x.Gpus + } + return 0 +} + // SlurmShared is where the Slurm cluster's shared files come from. type SlurmShared struct { state protoimpl.MessageState `protogen:"open.v1"` @@ -5454,7 +5652,7 @@ type SlurmShared struct { func (x *SlurmShared) Reset() { *x = SlurmShared{} - mi := &file_llmgw_v1_agent_proto_msgTypes[49] + mi := &file_llmgw_v1_agent_proto_msgTypes[51] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -5466,7 +5664,7 @@ func (x *SlurmShared) String() string { func (*SlurmShared) ProtoMessage() {} func (x *SlurmShared) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[49] + mi := &file_llmgw_v1_agent_proto_msgTypes[51] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -5479,7 +5677,7 @@ func (x *SlurmShared) ProtoReflect() protoreflect.Message { // Deprecated: Use SlurmShared.ProtoReflect.Descriptor instead. func (*SlurmShared) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{49} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{51} } func (x *SlurmShared) GetPath() string { @@ -5527,7 +5725,7 @@ type SlurmUser struct { func (x *SlurmUser) Reset() { *x = SlurmUser{} - mi := &file_llmgw_v1_agent_proto_msgTypes[50] + mi := &file_llmgw_v1_agent_proto_msgTypes[52] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -5539,7 +5737,7 @@ func (x *SlurmUser) String() string { func (*SlurmUser) ProtoMessage() {} func (x *SlurmUser) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[50] + mi := &file_llmgw_v1_agent_proto_msgTypes[52] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -5552,7 +5750,7 @@ func (x *SlurmUser) ProtoReflect() protoreflect.Message { // Deprecated: Use SlurmUser.ProtoReflect.Descriptor instead. func (*SlurmUser) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{50} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{52} } func (x *SlurmUser) GetLogin() string { @@ -5623,14 +5821,17 @@ type SlurmStatus struct { Users int32 `protobuf:"varint,12,opt,name=users,proto3" json:"users,omitempty"` // containers is what this machine has for jobs in containers, as // "pyxis 0.24.0, enroot 4.2.1"; empty while it has not got both. - Containers string `protobuf:"bytes,13,opt,name=containers,proto3" json:"containers,omitempty"` + Containers string `protobuf:"bytes,13,opt,name=containers,proto3" json:"containers,omitempty"` + // ended are, on the controller, the jobs that ended in the last day; the + // leader records each once. + Ended []*SlurmJob `protobuf:"bytes,14,rep,name=ended,proto3" json:"ended,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } func (x *SlurmStatus) Reset() { *x = SlurmStatus{} - mi := &file_llmgw_v1_agent_proto_msgTypes[51] + mi := &file_llmgw_v1_agent_proto_msgTypes[53] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -5642,7 +5843,7 @@ func (x *SlurmStatus) String() string { func (*SlurmStatus) ProtoMessage() {} func (x *SlurmStatus) ProtoReflect() protoreflect.Message { - mi := &file_llmgw_v1_agent_proto_msgTypes[51] + mi := &file_llmgw_v1_agent_proto_msgTypes[53] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -5655,7 +5856,7 @@ func (x *SlurmStatus) ProtoReflect() protoreflect.Message { // Deprecated: Use SlurmStatus.ProtoReflect.Descriptor instead. func (*SlurmStatus) Descriptor() ([]byte, []int) { - return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{51} + return file_llmgw_v1_agent_proto_rawDescGZIP(), []int{53} } func (x *SlurmStatus) GetVersion() string { @@ -5749,6 +5950,13 @@ func (x *SlurmStatus) GetContainers() string { return "" } +func (x *SlurmStatus) GetEnded() []*SlurmJob { + if x != nil { + return x.Ended + } + return nil +} + var File_llmgw_v1_agent_proto protoreflect.FileDescriptor const file_llmgw_v1_agent_proto_rawDesc = "" + @@ -6254,7 +6462,7 @@ const file_llmgw_v1_agent_proto_rawDesc = "" + "\x04kind\x18\x02 \x01(\tR\x04kind\x12\x10\n" + "\x03seq\x18\x03 \x01(\x05R\x03seq\x12\x12\n" + "\x04data\x18\x04 \x01(\fR\x04data\x12\x12\n" + - "\x04last\x18\x05 \x01(\bR\x04last\"\xc9\x03\n" + + "\x04last\x18\x05 \x01(\bR\x04last\"\xa6\x04\n" + "\x05Slurm\x12\x1e\n" + "\n" + "controller\x18\x01 \x01(\bR\n" + @@ -6274,11 +6482,30 @@ const file_llmgw_v1_agent_proto_rawDesc = "" + "make_homes\x18\v \x01(\bR\tmakeHomes\x12\x1e\n" + "\n" + "containers\x18\f \x01(\bR\n" + - "containers\x1a8\n" + + "containers\x12'\n" + + "\x0fmanage_accounts\x18\r \x01(\bR\x0emanageAccounts\x122\n" + + "\baccounts\x18\x0e \x03(\v2\x16.llmgw.v1.SlurmAccountR\baccounts\x1a8\n" + "\n" + "FilesEntry\x12\x10\n" + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + - "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"c\n" + + "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"P\n" + + "\fSlurmAccount\x12\x12\n" + + "\x04name\x18\x01 \x01(\tR\x04name\x12\x16\n" + + "\x06parent\x18\x02 \x01(\tR\x06parent\x12\x14\n" + + "\x05users\x18\x03 \x03(\tR\x05users\"\x87\x02\n" + + "\bSlurmJob\x12\x0e\n" + + "\x02id\x18\x01 \x01(\tR\x02id\x12\x12\n" + + "\x04user\x18\x02 \x01(\tR\x04user\x12\x18\n" + + "\aaccount\x18\x03 \x01(\tR\aaccount\x12\x1c\n" + + "\tpartition\x18\x04 \x01(\tR\tpartition\x12\x14\n" + + "\x05state\x18\x05 \x01(\tR\x05state\x12\x1d\n" + + "\n" + + "start_unix\x18\x06 \x01(\x03R\tstartUnix\x12\x19\n" + + "\bend_unix\x18\a \x01(\x03R\aendUnix\x12'\n" + + "\x0felapsed_seconds\x18\b \x01(\x03R\x0eelapsedSeconds\x12\x12\n" + + "\x04cpus\x18\t \x01(\x05R\x04cpus\x12\x12\n" + + "\x04gpus\x18\n" + + " \x01(\x05R\x04gpus\"c\n" + "\vSlurmShared\x12\x12\n" + "\x04path\x18\x01 \x01(\tR\x04path\x12\x10\n" + "\x03nfs\x18\x02 \x01(\tR\x03nfs\x12\x14\n" + @@ -6289,7 +6516,7 @@ const file_llmgw_v1_agent_proto_rawDesc = "" + "\x03uid\x18\x02 \x01(\rR\x03uid\x12\x12\n" + "\x04name\x18\x03 \x01(\tR\x04name\x12\x12\n" + "\x04home\x18\x04 \x01(\tR\x04home\x12\x19\n" + - "\bssh_keys\x18\x05 \x03(\tR\asshKeys\"\x8b\x03\n" + + "\bssh_keys\x18\x05 \x03(\tR\asshKeys\"\xb5\x03\n" + "\vSlurmStatus\x12\x18\n" + "\aversion\x18\x01 \x01(\tR\aversion\x12\x1a\n" + "\bhardware\x18\x02 \x01(\tR\bhardware\x12%\n" + @@ -6306,7 +6533,8 @@ const file_llmgw_v1_agent_proto_rawDesc = "" + "\x05users\x18\f \x01(\x05R\x05users\x12\x1e\n" + "\n" + "containers\x18\r \x01(\tR\n" + - "containers2O\n" + + "containers\x12(\n" + + "\x05ended\x18\x0e \x03(\v2\x12.llmgw.v1.SlurmJobR\x05ended2O\n" + "\fAgentService\x12?\n" + "\aConnect\x12\x16.llmgw.v1.AgentMessage\x1a\x18.llmgw.v1.ControlMessage(\x010\x01B\x93\x01\n" + "\fcom.llmgw.v1B\n" + @@ -6325,7 +6553,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, 61) +var file_llmgw_v1_agent_proto_msgTypes = make([]protoimpl.MessageInfo, 63) var file_llmgw_v1_agent_proto_goTypes = []any{ (ShardStatus_State)(0), // 0: llmgw.v1.ShardStatus.State (*AgentMessage)(nil), // 1: llmgw.v1.AgentMessage @@ -6377,18 +6605,20 @@ var file_llmgw_v1_agent_proto_goTypes = []any{ (*TrainStatus)(nil), // 47: llmgw.v1.TrainStatus (*TrainArtifact)(nil), // 48: llmgw.v1.TrainArtifact (*Slurm)(nil), // 49: llmgw.v1.Slurm - (*SlurmShared)(nil), // 50: llmgw.v1.SlurmShared - (*SlurmUser)(nil), // 51: llmgw.v1.SlurmUser - (*SlurmStatus)(nil), // 52: llmgw.v1.SlurmStatus - nil, // 53: llmgw.v1.ControlMessage.TraceEntry - nil, // 54: llmgw.v1.Register.LabelsEntry - nil, // 55: llmgw.v1.Register.EngineVersionsEntry - nil, // 56: llmgw.v1.GPU.NicPathsEntry - nil, // 57: llmgw.v1.IBPort.CountersEntry - nil, // 58: llmgw.v1.Heartbeat.EngineVersionsEntry - nil, // 59: llmgw.v1.EngineWant.OfflineSha256Entry - nil, // 60: llmgw.v1.EngineProfile.EnvEntry - nil, // 61: llmgw.v1.Slurm.FilesEntry + (*SlurmAccount)(nil), // 50: llmgw.v1.SlurmAccount + (*SlurmJob)(nil), // 51: llmgw.v1.SlurmJob + (*SlurmShared)(nil), // 52: llmgw.v1.SlurmShared + (*SlurmUser)(nil), // 53: llmgw.v1.SlurmUser + (*SlurmStatus)(nil), // 54: llmgw.v1.SlurmStatus + nil, // 55: llmgw.v1.ControlMessage.TraceEntry + nil, // 56: llmgw.v1.Register.LabelsEntry + nil, // 57: llmgw.v1.Register.EngineVersionsEntry + nil, // 58: llmgw.v1.GPU.NicPathsEntry + nil, // 59: llmgw.v1.IBPort.CountersEntry + nil, // 60: llmgw.v1.Heartbeat.EngineVersionsEntry + nil, // 61: llmgw.v1.EngineWant.OfflineSha256Entry + nil, // 62: llmgw.v1.EngineProfile.EnvEntry + nil, // 63: llmgw.v1.Slurm.FilesEntry } var file_llmgw_v1_agent_proto_depIdxs = []int32{ 8, // 0: llmgw.v1.AgentMessage.register:type_name -> llmgw.v1.Register @@ -6401,7 +6631,7 @@ var file_llmgw_v1_agent_proto_depIdxs = []int32{ 2, // 7: llmgw.v1.AgentMessage.kv_events:type_name -> llmgw.v1.KVEvents 47, // 8: llmgw.v1.AgentMessage.train_status:type_name -> llmgw.v1.TrainStatus 48, // 9: llmgw.v1.AgentMessage.train_artifact:type_name -> llmgw.v1.TrainArtifact - 52, // 10: llmgw.v1.AgentMessage.slurm_status:type_name -> llmgw.v1.SlurmStatus + 54, // 10: llmgw.v1.AgentMessage.slurm_status:type_name -> llmgw.v1.SlurmStatus 30, // 11: llmgw.v1.ControlMessage.start:type_name -> llmgw.v1.StartShard 33, // 12: llmgw.v1.ControlMessage.stop:type_name -> llmgw.v1.StopShard 34, // 13: llmgw.v1.ControlMessage.drain:type_name -> llmgw.v1.DrainShard @@ -6419,11 +6649,11 @@ var file_llmgw_v1_agent_proto_depIdxs = []int32{ 45, // 25: llmgw.v1.ControlMessage.train_start:type_name -> llmgw.v1.TrainStart 46, // 26: llmgw.v1.ControlMessage.train_cancel:type_name -> llmgw.v1.TrainCancel 49, // 27: llmgw.v1.ControlMessage.slurm:type_name -> llmgw.v1.Slurm - 53, // 28: llmgw.v1.ControlMessage.trace:type_name -> llmgw.v1.ControlMessage.TraceEntry - 54, // 29: llmgw.v1.Register.labels:type_name -> llmgw.v1.Register.LabelsEntry + 55, // 28: llmgw.v1.ControlMessage.trace:type_name -> llmgw.v1.ControlMessage.TraceEntry + 56, // 29: llmgw.v1.Register.labels:type_name -> llmgw.v1.Register.LabelsEntry 12, // 30: llmgw.v1.Register.gpus:type_name -> llmgw.v1.GPU 14, // 31: llmgw.v1.Register.fabrics:type_name -> llmgw.v1.FabricAttachment - 55, // 32: llmgw.v1.Register.engine_versions:type_name -> llmgw.v1.Register.EngineVersionsEntry + 57, // 32: llmgw.v1.Register.engine_versions:type_name -> llmgw.v1.Register.EngineVersionsEntry 11, // 33: llmgw.v1.Register.running_shards:type_name -> llmgw.v1.RunningShard 15, // 34: llmgw.v1.Register.weights:type_name -> llmgw.v1.LocalWeights 10, // 35: llmgw.v1.Register.storage:type_name -> llmgw.v1.Storage @@ -6431,15 +6661,15 @@ var file_llmgw_v1_agent_proto_depIdxs = []int32{ 21, // 37: llmgw.v1.Register.imex:type_name -> llmgw.v1.IMEX 0, // 38: llmgw.v1.RunningShard.state:type_name -> llmgw.v1.ShardStatus.State 27, // 39: llmgw.v1.RunningShard.compute:type_name -> llmgw.v1.Compute - 56, // 40: llmgw.v1.GPU.nic_paths:type_name -> llmgw.v1.GPU.NicPathsEntry + 58, // 40: llmgw.v1.GPU.nic_paths:type_name -> llmgw.v1.GPU.NicPathsEntry 13, // 41: llmgw.v1.GPU.mig:type_name -> llmgw.v1.MIGDevice 18, // 42: llmgw.v1.NodeLoad.shards:type_name -> llmgw.v1.ShardLoad 17, // 43: llmgw.v1.NodeLoad.gpus:type_name -> llmgw.v1.GPUSensor - 57, // 44: llmgw.v1.IBPort.counters:type_name -> llmgw.v1.IBPort.CountersEntry + 59, // 44: llmgw.v1.IBPort.counters:type_name -> llmgw.v1.IBPort.CountersEntry 20, // 45: llmgw.v1.IBPort.subnet_managers:type_name -> llmgw.v1.SubnetManager 10, // 46: llmgw.v1.Heartbeat.storage:type_name -> llmgw.v1.Storage 15, // 47: llmgw.v1.Heartbeat.weights:type_name -> llmgw.v1.LocalWeights - 58, // 48: llmgw.v1.Heartbeat.engine_versions:type_name -> llmgw.v1.Heartbeat.EngineVersionsEntry + 60, // 48: llmgw.v1.Heartbeat.engine_versions:type_name -> llmgw.v1.Heartbeat.EngineVersionsEntry 26, // 49: llmgw.v1.Heartbeat.provision:type_name -> llmgw.v1.ProvisionStatus 16, // 50: llmgw.v1.Heartbeat.load:type_name -> llmgw.v1.NodeLoad 19, // 51: llmgw.v1.Heartbeat.ib_ports:type_name -> llmgw.v1.IBPort @@ -6447,23 +6677,25 @@ var file_llmgw_v1_agent_proto_depIdxs = []int32{ 22, // 53: llmgw.v1.Heartbeat.gpu_fabric:type_name -> llmgw.v1.GPUFabricState 21, // 54: llmgw.v1.Heartbeat.imex:type_name -> llmgw.v1.IMEX 25, // 55: llmgw.v1.Provision.engines:type_name -> llmgw.v1.EngineWant - 59, // 56: llmgw.v1.EngineWant.offline_sha256:type_name -> llmgw.v1.EngineWant.OfflineSha256Entry + 61, // 56: llmgw.v1.EngineWant.offline_sha256:type_name -> llmgw.v1.EngineWant.OfflineSha256Entry 0, // 57: llmgw.v1.ShardStatus.state:type_name -> llmgw.v1.ShardStatus.State 29, // 58: llmgw.v1.ShardStatus.fetch:type_name -> llmgw.v1.FetchProgress 27, // 59: llmgw.v1.ShardStatus.compute:type_name -> llmgw.v1.Compute 31, // 60: llmgw.v1.StartShard.profile:type_name -> llmgw.v1.EngineProfile - 60, // 61: llmgw.v1.EngineProfile.env:type_name -> llmgw.v1.EngineProfile.EnvEntry + 62, // 61: llmgw.v1.EngineProfile.env:type_name -> llmgw.v1.EngineProfile.EnvEntry 44, // 62: llmgw.v1.EngineProfile.adapters:type_name -> llmgw.v1.LoraAdapter - 61, // 63: llmgw.v1.Slurm.files:type_name -> llmgw.v1.Slurm.FilesEntry - 50, // 64: llmgw.v1.Slurm.shared:type_name -> llmgw.v1.SlurmShared - 51, // 65: llmgw.v1.Slurm.users:type_name -> llmgw.v1.SlurmUser - 1, // 66: llmgw.v1.AgentService.Connect:input_type -> llmgw.v1.AgentMessage - 3, // 67: llmgw.v1.AgentService.Connect:output_type -> llmgw.v1.ControlMessage - 67, // [67:68] is the sub-list for method output_type - 66, // [66:67] is the sub-list for method input_type - 66, // [66:66] is the sub-list for extension type_name - 66, // [66:66] is the sub-list for extension extendee - 0, // [0:66] is the sub-list for field type_name + 63, // 63: llmgw.v1.Slurm.files:type_name -> llmgw.v1.Slurm.FilesEntry + 52, // 64: llmgw.v1.Slurm.shared:type_name -> llmgw.v1.SlurmShared + 53, // 65: llmgw.v1.Slurm.users:type_name -> llmgw.v1.SlurmUser + 50, // 66: llmgw.v1.Slurm.accounts:type_name -> llmgw.v1.SlurmAccount + 51, // 67: llmgw.v1.SlurmStatus.ended:type_name -> llmgw.v1.SlurmJob + 1, // 68: llmgw.v1.AgentService.Connect:input_type -> llmgw.v1.AgentMessage + 3, // 69: llmgw.v1.AgentService.Connect:output_type -> llmgw.v1.ControlMessage + 69, // [69:70] is the sub-list for method output_type + 68, // [68:69] is the sub-list for method input_type + 68, // [68:68] is the sub-list for extension type_name + 68, // [68:68] is the sub-list for extension extendee + 0, // [0:68] is the sub-list for field type_name } func init() { file_llmgw_v1_agent_proto_init() } @@ -6509,7 +6741,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: 61, + NumMessages: 63, NumExtensions: 0, NumServices: 1, }, diff --git a/api/proto/llmgw/v1/agent.proto b/api/proto/llmgw/v1/agent.proto index 04964b5..f092da6 100644 --- a/api/proto/llmgw/v1/agent.proto +++ b/api/proto/llmgw/v1/agent.proto @@ -1023,6 +1023,33 @@ message Slurm { // containers has the machine install pyxis and enroot, so jobs may run // in containers (srun --container-image), and set them up. bool containers = 12; + // manage_accounts has the controller keep Slurm's accounts as accounts + // says: one for each organization, team and project, nested, the users + // on each project's. Only those llmgw made are changed or taken away. + bool manage_accounts = 13; + repeated SlurmAccount accounts = 14; +} + +// SlurmAccount is an account in Slurm's accounting, parents before their +// children. +message SlurmAccount { + string name = 1; + string parent = 2; + repeated string users = 3; +} + +// SlurmJob is a job that ended, as sacct has it, for llmgw's accounting. +message SlurmJob { + string id = 1; + string user = 2; + string account = 3; + string partition = 4; + string state = 5; + int64 start_unix = 6; + int64 end_unix = 7; + int64 elapsed_seconds = 8; + int32 cpus = 9; + int32 gpus = 10; } // SlurmShared is where the Slurm cluster's shared files come from. @@ -1083,4 +1110,7 @@ message SlurmStatus { // containers is what this machine has for jobs in containers, as // "pyxis 0.24.0, enroot 4.2.1"; empty while it has not got both. string containers = 13; + // ended are, on the controller, the jobs that ended in the last day; the + // leader records each once. + repeated SlurmJob ended = 14; } diff --git a/docs/admin-api.openapi.json b/docs/admin-api.openapi.json index 5aee1aa..4187bf5 100644 --- a/docs/admin-api.openapi.json +++ b/docs/admin-api.openapi.json @@ -12573,6 +12573,10 @@ "containers": { "type": "string", "description": "Containers is whether jobs may run in containers: \"on\" once every machine has pyxis and enroot, \"installing\" until then, empty when slurm.containers is off." + }, + "accounts": { + "type": "integer", + "description": "Accounts is how many of Slurm's accounts are llmgw's projects'." } } }, diff --git a/docs/install.md b/docs/install.md index 58cc941..041c9c6 100644 --- a/docs/install.md +++ b/docs/install.md @@ -217,6 +217,17 @@ agents install); root and the machines' own users log in as before. A site whose directory (sssd against Active Directory or LDAP) passes `--slurm-users site`, or later `llmgw ctl slurm users off`. +With llmgw's users, Slurm's accounts are llmgw's projects: an account for +each organization, team and project, nested (`acme`, `acme-ml`, +`acme-ml-training`), and on each project's account its members and its +team's and organization's. A job runs on one of its user's accounts and no +other (`sbatch --account=acme-ml-training`; without it, the user's first), +and when it ends its cards' hours (a job without cards, its cores') are +priced at `auth.gpu_hour_cost` (`core_hour_cost`) into the project's usage, +where it shows as model `slurm/PARTITION` beside the project's requests. +Accounts a Slurm administrator made are left alone. A job's cost does not +yet count against a budget, which binds API keys only. + Jobs may run in containers: `llmgw ctl slurm containers on` has the agents install enroot and pyxis on every Slurm machine (NVIDIA's releases, at versions llmgw pins and checks; pyxis is built against each machine's own diff --git a/internal/agent/slurmnode/accounts.go b/internal/agent/slurmnode/accounts.go new file mode 100644 index 0000000..478c68f --- /dev/null +++ b/internal/agent/slurmnode/accounts.go @@ -0,0 +1,214 @@ +package slurmnode + +import ( + "context" + "errors" + "fmt" + "slices" + "strings" + + llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1" + "gitea.mynodes.xyz/ilker/llmgw/internal/slurm" +) + +// accountsMarker lists what llmgw made in Slurm's accounting, a line each: +// "account NAME" and "user LOGIN NAME". Nothing else is changed or taken +// away: an account a Slurm administrator made is theirs. +func (k *Keeper) accountsMarker() string { return k.marker + "-accounts" } + +// maxEnded is the most ended jobs sent at once: a day's on a busy cluster +// would otherwise be a status message of megabytes. +const maxEnded = 2000 + +// accounts keeps Slurm's accounts and their users as want says, on the +// controller. +func (k *Keeper) accounts(ctx context.Context, want *llmgwv1.Slurm) error { + made := k.madeAccounts() + if !want.GetManageAccounts() { + if len(made) == 0 { + return nil + } + return k.applyAccounts(ctx, nil, made) + } + return k.applyAccounts(ctx, want.GetAccounts(), made) +} + +func (k *Keeper) madeAccounts() []string { + b, err := k.host.ReadFile(k.accountsMarker()) + if err != nil { + return nil + } + var out []string + for _, l := range strings.Split(strings.TrimSpace(string(b)), "\n") { + if l != "" { + out = append(out, l) + } + } + return out +} + +func (k *Keeper) sacctmgr(ctx context.Context, args ...string) error { + out, err := k.host.Run(ctx, nil, "sacctmgr", append([]string{"-i"}, args...)...) + if err != nil { + return fmt.Errorf("sacctmgr %s: %v: %s", strings.Join(args, " "), err, lastLine(out)) + } + return nil +} + +func (k *Keeper) applyAccounts(ctx context.Context, want []*llmgwv1.SlurmAccount, made []string) error { + out, err := k.host.Run(ctx, nil, "sacctmgr", slurm.AssocFormat...) + if err != nil { + return fmt.Errorf("reading Slurm's accounts: %v: %s", err, lastLine(out)) + } + have := map[string]bool{} // "account NAME", "user LOGIN NAME" + for _, a := range slurm.ParseAssociations(string(out)) { + if a.User == "" { + have["account "+a.Account] = true + } else { + have["user "+a.User+" "+a.Account] = true + } + } + keep := map[string]bool{} + var errs []error + ours := slices.Clone(made) + note := func(line string) { + if !slices.Contains(ours, line) { + ours = append(ours, line) + } + } + for _, a := range want { + name, parent := a.GetName(), a.GetParent() + if !validAccount(name) || (parent != "root" && !validAccount(parent)) { + errs = append(errs, fmt.Errorf("the leader sent an account %q under %q", name, parent)) + continue + } + line := "account " + name + keep[line] = true + if !have[line] { + k.log.Info("a Slurm account is made, as llmgw's projects say", "account", name, "parent", parent) + if err := k.sacctmgr(ctx, "add", "account", name, "parent="+parent, "description=llmgw"); err != nil { + errs = append(errs, err) + continue + } + note(line) + } + for _, u := range a.GetUsers() { + if !slurm.ValidLogin(u) { + continue + } + ul := "user " + u + " " + name + keep[ul] = true + if have[ul] { + continue + } + if err := k.sacctmgr(ctx, "add", "user", u, "account="+name); err != nil { + errs = append(errs, err) + continue + } + note(ul) + } + } + // Taken away: what llmgw made and no longer wants, users before their + // accounts, children before their parents (made later, listed later). + var left []string + for i := len(ours) - 1; i >= 0; i-- { + line := ours[i] + if keep[line] || !strings.HasPrefix(line, "user ") { + continue + } + f := strings.Fields(line) + if !have[line] { + continue + } + if err := k.removeAssoc(ctx, f[1], f[2], have, keep); err != nil { + errs = append(errs, err) + left = append(left, line) + } + } + for i := len(ours) - 1; i >= 0; i-- { + line := ours[i] + if keep[line] || !strings.HasPrefix(line, "account ") { + continue + } + if have[line] { + k.log.Info("a Slurm account llmgw made is taken away; its jobs' records stay", "account", line[len("account "):]) + if err := k.sacctmgr(ctx, "remove", "account", line[len("account "):]); err != nil { + errs = append(errs, err) + left = append(left, line) + } + } + } + var lines []string + for _, l := range ours { + if keep[l] || slices.Contains(left, l) { + lines = append(lines, l) + } + } + if !slices.Equal(lines, made) { + var err error + if len(lines) == 0 { + err = k.host.Remove(k.accountsMarker()) + } else { + _, err = k.write(k.accountsMarker(), []byte(strings.Join(lines, "\n")+"\n"), 0o644, "") + } + errs = append(errs, err) + } + return errors.Join(errs...) +} + +func validAccount(s string) bool { + return s != "" && len(s) <= 64 && strings.Trim(s, "abcdefghijklmnopqrstuvwxyz0123456789_-") == "" +} + +// ended are the jobs that ended in the last day, newest last. +func (k *Keeper) ended(ctx context.Context) ([]*llmgwv1.SlurmJob, error) { + out, err := k.host.Run(ctx, nil, "sacct", slurm.SacctArgs...) + if err != nil { + return nil, fmt.Errorf("sacct: %v: %s", err, lastLine(out)) + } + jobs, err := slurm.ParseJobs(string(out)) + if err != nil { + return nil, err + } + if len(jobs) > maxEnded { + jobs = jobs[len(jobs)-maxEnded:] + } + ps := make([]*llmgwv1.SlurmJob, 0, len(jobs)) + for _, j := range jobs { + ps = append(ps, &llmgwv1.SlurmJob{Id: j.ID, User: j.User, Account: j.Account, Partition: j.Partition, + State: j.State, StartUnix: j.Start.Unix(), EndUnix: j.End.Unix(), ElapsedSeconds: int64(j.Elapsed.Seconds()), + Cpus: int32(j.CPUs), Gpus: int32(j.GPUs)}) //nolint:gosec // counts of a node's + } + return ps, nil +} + +// removeAssoc takes login off account. Slurm refuses to take away a +// user's default account, so another of theirs is made the default first, +// one llmgw keeps if there is one; with none left the user goes from the +// accounting entirely. +func (k *Keeper) removeAssoc(ctx context.Context, login, account string, have, keep map[string]bool) error { + var others []string + for line := range have { + f := strings.Fields(line) + if len(f) == 3 && f[0] == "user" && f[1] == login && f[2] != account { + others = append(others, f[2]) + } + } + if len(others) == 0 { + return k.sacctmgr(ctx, "remove", "user", login) + } + slices.SortFunc(others, func(a, b string) int { + ka, kb := keep["user "+login+" "+a], keep["user "+login+" "+b] + switch { + case ka && !kb: + return -1 + case kb && !ka: + return 1 + } + return strings.Compare(a, b) + }) + if err := k.sacctmgr(ctx, "modify", "user", login, "set", "defaultaccount="+others[0]); err != nil { + return err + } + return k.sacctmgr(ctx, "remove", "user", login, "account="+account) +} diff --git a/internal/agent/slurmnode/accounts_test.go b/internal/agent/slurmnode/accounts_test.go new file mode 100644 index 0000000..611259c --- /dev/null +++ b/internal/agent/slurmnode/accounts_test.go @@ -0,0 +1,83 @@ +package slurmnode + +import ( + "context" + "slices" + "strings" + "testing" + + llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1" + "gitea.mynodes.xyz/ilker/llmgw/internal/slurm" +) + +func TestAccountsMadeAndTakenAwayOnlyIfLlmgwsOwn(t *testing.T) { + h := newFake() + k := New(h, nil, "/var/lib/llmgw-agent/slurm", nil) + assoc := "sacctmgr " + strings.Join(slurm.AssocFormat, " ") + // root, and an account a Slurm administrator made. + h.outputs[assoc] = "root||\nroot|root|\nphysics||root\n" + m := &llmgwv1.Slurm{ManageAccounts: true, Accounts: []*llmgwv1.SlurmAccount{ + {Name: "test", Parent: "root"}, {Name: "test-arge", Parent: "test"}, + {Name: "test-arge-deneme", Parent: "test-arge", Users: []string{"ada", "ayse"}}, + }} + if err := k.accounts(context.Background(), m); err != nil { + t.Fatal(err) + } + for _, want := range []string{"sacctmgr -i add account test parent=root description=llmgw", + "sacctmgr -i add account test-arge-deneme parent=test-arge description=llmgw", + "sacctmgr -i add user ayse account=test-arge-deneme"} { + if !slices.Contains(h.ran, want) { + t.Errorf("did not run %q:\n%s", want, strings.Join(h.ran, "\n")) + } + } + if !strings.Contains(string(h.files[k.accountsMarker()]), "user ada test-arge-deneme") { + t.Errorf("marker %q", h.files[k.accountsMarker()]) + } + + // Now there; ayse leaves the project, and the team goes. + h.outputs[assoc] = "root||\nphysics||root\ntest||root\ntest-arge||test\ntest-arge-deneme||test-arge\n" + + "test-arge-deneme|ada|test-arge\ntest-arge-deneme|ayse|test-arge\n" + h.ran = nil + m.Accounts = []*llmgwv1.SlurmAccount{{Name: "test", Parent: "root"}} + if err := k.accounts(context.Background(), m); err != nil { + t.Fatal(err) + } + // Neither has another account: they go from the accounting. + want := []string{"sacctmgr -i remove user ayse", "sacctmgr -i remove user ada", + "sacctmgr -i remove account test-arge-deneme", "sacctmgr -i remove account test-arge"} + var got []string + for _, c := range h.ran { + if strings.Contains(c, "remove") || strings.Contains(c, " add ") { + got = append(got, c) + } + } + if !slices.Equal(got, want) { + t.Errorf("ran:\n%s\nwant:\n%s", strings.Join(got, "\n"), strings.Join(want, "\n")) + } + if slices.ContainsFunc(h.ran, func(c string) bool { return strings.Contains(c, "physics") }) { + t.Error("touched an account llmgw did not make") + } + + // One still on another account: that one is made the default first. + h.outputs[assoc] = "root||\nphysics||root\nphysics|ada|\ntest||root\nx||test\nx|ada|\n" + h.files[k.accountsMarker()] = []byte("account test\naccount x\nuser ada x\n") + h.ran = nil + m.Accounts = []*llmgwv1.SlurmAccount{{Name: "test", Parent: "root"}} + if err := k.accounts(context.Background(), m); err != nil { + t.Fatal(err) + } + if i, j := slices.Index(h.ran, "sacctmgr -i modify user ada set defaultaccount=physics"), + slices.Index(h.ran, "sacctmgr -i remove user ada account=x"); i < 0 || j < i { + t.Errorf("ran:\n%s", strings.Join(h.ran, "\n")) + } + + // Off: what llmgw made goes. + h.outputs[assoc] = "root||\nphysics||root\ntest||root\n" + h.ran = nil + if err := k.accounts(context.Background(), &llmgwv1.Slurm{}); err != nil { + t.Fatal(err) + } + if !slices.Contains(h.ran, "sacctmgr -i remove account test") || h.files[k.accountsMarker()] != nil { + t.Errorf("not taken away:\n%s", strings.Join(h.ran, "\n")) + } +} diff --git a/internal/agent/slurmnode/slurmnode.go b/internal/agent/slurmnode/slurmnode.go index a00b533..7c5a8f5 100644 --- a/internal/agent/slurmnode/slurmnode.go +++ b/internal/agent/slurmnode/slurmnode.go @@ -267,6 +267,20 @@ func (k *Keeper) Reconcile(ctx context.Context, want *llmgwv1.Slurm) *llmgwv1.Sl return fail("slurmctld", err) } st.SlurmctldRunning = k.active(ctx, "slurmctld") + // The accounts, and the jobs that ended, need the accounting. + if st.SlurmdbdRunning && st.SlurmctldRunning { + if err := k.accounts(ctx, want); err != nil { + k.log.Warn("keeping Slurm's accounts", "err", err) + if acctErr == nil { + acctErr = fmt.Errorf("accounts: %w", err) + } + } + if ended, err := k.ended(ctx); err != nil { + k.log.Warn("reading the jobs that ended", "err", err) + } else { + st.Ended = ended + } + } } if want.GetNode() { if err := k.node(ctx, want, st); err != nil { diff --git a/internal/cli/ctl_slurm.go b/internal/cli/ctl_slurm.go index e4aa9a7..ca728f1 100644 --- a/internal/cli/ctl_slurm.go +++ b/internal/cli/ctl_slurm.go @@ -247,6 +247,9 @@ func newSlurmStatusCmd(client func() *ctlClient, asJSON *bool) *cobra.Command { fmt.Fprintf(w, "shared files: %s\n", sharedLine(st.Shared)) if st.Users != "" { fmt.Fprintf(w, "users: %d, made from llmgw's accounts\n", st.UserCount) + if st.Accounts > 0 { + fmt.Fprintf(w, "accounts: %d, one for each project (sbatch --account=ORG-TEAM-PROJECT)\n", st.Accounts) + } } if line := containersLine(st); line != "" { fmt.Fprintf(w, "containers: %s\n", line) diff --git a/internal/cli/server.go b/internal/cli/server.go index 5be1fb8..faa5059 100644 --- a/internal/cli/server.go +++ b/internal/cli/server.go @@ -10,6 +10,7 @@ import ( "fmt" "log/slog" "maps" + "math" "net" "net/http" "os" @@ -487,6 +488,15 @@ func newServerCmd() *cobra.Command { return out, err }) api.SetSlurmUsers(st) + // Slurm's accounts from the projects, and the jobs + // that ended priced as the cluster's hours are. + slurmSvc.SetAccounts(func(ctx context.Context) ([]slurm.Account, error) { + return slurmAccounts(ctx, st) + }) + slurmSvc.SetJobs(func(ctx context.Context, jobs []slurmsvc.JobRecord) error { + au := current.Load().Auth + return recordSlurmJobs(ctx, st, current.Load().Slurm.Cluster(), au.GPUHourCost, au.CoreHourCost, jobs, log) + }) agentSvc.SetSlurm(slurmSvc) go slurmSvc.Run(ctx) api.SetSlurm(slurmSvc) @@ -1447,3 +1457,76 @@ func accountingPrice(p config.Price) accounting.Price { return accounting.Price{PromptPer1k: p.PromptPer1k, CompletionPer1k: p.CompletionPer1k, CachedPer1k: p.CachedPer1k, PerImage: p.PerImage, PerMesh: p.PerMesh} } + +// slurmAccounts are Slurm's accounts from the store's projects, members and +// Slurm users. +func slurmAccounts(ctx context.Context, st *storepkg.DB) ([]slurm.Account, error) { + ps, err := st.SlurmProjects(ctx) + if err != nil { + return nil, err + } + ms, err := st.AllMemberships(ctx) + if err != nil { + return nil, err + } + us, err := st.SlurmUsers(ctx) + if err != nil { + return nil, err + } + logins := make(map[string]string, len(us)) + for _, u := range us { + logins[u.UserID] = u.Login + } + tps := make([]slurm.TenancyProject, 0, len(ps)) + for _, p := range ps { + tps = append(tps, slurm.TenancyProject{ID: p.ID, Name: p.Name, TeamID: p.TeamID, TeamName: p.TeamName, + OrgID: p.OrgID, OrgName: p.OrgName}) + } + tms := make([]slurm.TenancyMember, 0, len(ms)) + for _, m := range ms { + tms = append(tms, slurm.TenancyMember{UserID: m.UserID, Scope: string(m.Scope), ScopeID: m.ScopeID}) + } + return slurm.Accounts(tps, tms, logins), nil +} + +// recordSlurmJobs records the jobs that ended, each with its project (from +// its account) and person (from its login), priced as the hours the +// cluster's models hold are: a GPU job its cards' hours, any other its +// cores'. +func recordSlurmJobs(ctx context.Context, st *storepkg.DB, cluster string, gpuHour, coreHour float64, + jobs []slurmsvc.JobRecord, log *slog.Logger) error { + as, err := slurmAccounts(ctx, st) + if err != nil { + return err + } + project := map[string]string{} + for _, a := range as { + if a.Project != "" { + project[a.Name] = a.Project + } + } + us, err := st.SlurmUsers(ctx) + if err != nil { + return err + } + person := map[string]string{} + for _, u := range us { + person[u.Login] = u.UserID + } + rows := make([]storepkg.SlurmJob, 0, len(jobs)) + for _, j := range jobs { + secs := j.Elapsed.Seconds() + cost := float64(j.CPUs) * secs / 3600 * coreHour + if j.GPUs > 0 { + cost = float64(j.GPUs) * secs / 3600 * gpuHour + } + rows = append(rows, storepkg.SlurmJob{Key: cluster + "/" + j.ID, UserID: person[j.Login], ProjectID: project[j.Account], + Login: j.Login, Account: j.Account, Partition: j.Partition, State: j.State, Start: j.Start, End: j.End, + ElapsedSeconds: int64(secs), CPUs: int64(j.CPUs), GPUs: int64(j.GPUs), CostMicros: int64(math.Ceil(cost * 1e6))}) + } + n, err := st.RecordSlurmJobs(ctx, rows) + if n > 0 { + log.Info("Slurm jobs that ended recorded, with their cost", "jobs", n) + } + return err +} diff --git a/internal/control/slurmsvc/service.go b/internal/control/slurmsvc/service.go index a7b3b6b..3de4ea1 100644 --- a/internal/control/slurmsvc/service.go +++ b/internal/control/slurmsvc/service.go @@ -116,6 +116,8 @@ type Status struct { // machine has pyxis and enroot, "installing" until then, empty when // slurm.containers is off. Containers string `json:"containers,omitempty"` + // Accounts is how many of Slurm's accounts are llmgw's projects'. + Accounts int `json:"accounts,omitempty"` } // User is a Linux user made from an llmgw account. @@ -166,6 +168,53 @@ type Service struct { users []User usersAt time.Time usersOK bool + // accountList reads Slurm's accounts from the projects, with the users; + // accountsOK is false until it has been read, and the controller is + // told nothing of accounts until then: none would take every one away. + accountList func(ctx context.Context) ([]slurm.Account, error) + accounts []slurm.Account + accountsOK bool + // jobSink records the jobs that ended; recorded are those it took, + // recording whether a recording is under way. + jobSink func(ctx context.Context, jobs []JobRecord) error + recorded map[string]time.Time + recording bool +} + +// JobRecord is a job that ended, for llmgw's accounting. +type JobRecord struct { + ID, Login, Account, Partition, State string + Start, End time.Time + Elapsed time.Duration + CPUs, GPUs int +} + +// SetAccounts has Slurm's accounts read with list, beside the users. +func (s *Service) SetAccounts(list func(ctx context.Context) ([]slurm.Account, error)) { + s.mu.Lock() + s.accountList = list + s.mu.Unlock() +} + +// SetJobs has the jobs that ended recorded with record; each is given to it +// once it took it without an error. +func (s *Service) SetJobs(record func(ctx context.Context, jobs []JobRecord) error) { + s.mu.Lock() + s.jobSink = record + s.mu.Unlock() +} + +// AccountsOf are the accounts a login may run jobs on. +func (s *Service) AccountsOf(login string) []string { + s.mu.Lock() + defer s.mu.Unlock() + var out []string + for _, a := range s.accounts { + if slices.Contains(a.Users, login) { + out = append(out, a.Name) + } + } + return out } // UsersEvery is how often the users are read again without being asked. @@ -199,6 +248,64 @@ func (s *Service) readUsers(ctx context.Context, set config.SlurmSettings) { return } s.users, s.usersOK = us, true + if s.accountList == nil { + return + } + as, err := s.accountList(ctx) + if err != nil { + s.log.Warn("could not read Slurm's accounts from the projects; the controller keeps the last ones", "err", err) + return + } + s.accounts, s.accountsOK = as, true +} + +// takeJobs records the jobs a controller says ended that are not recorded +// yet, away from the agent's stream. +func (s *Service) takeJobs(ended []*llmgwv1.SlurmJob) { + s.mu.Lock() + if s.jobSink == nil || s.recording { + s.mu.Unlock() + return + } + if s.recorded == nil { + s.recorded = map[string]time.Time{} + } + for id, at := range s.recorded { + if time.Since(at) > 3*24*time.Hour { + delete(s.recorded, id) + } + } + var jobs []JobRecord + for _, j := range ended { + if _, ok := s.recorded[j.GetId()]; ok || j.GetId() == "" { + continue + } + jobs = append(jobs, JobRecord{ID: j.GetId(), Login: j.GetUser(), Account: j.GetAccount(), Partition: j.GetPartition(), + State: j.GetState(), Start: time.Unix(j.GetStartUnix(), 0), End: time.Unix(j.GetEndUnix(), 0), + Elapsed: time.Duration(j.GetElapsedSeconds()) * time.Second, CPUs: int(j.GetCpus()), GPUs: int(j.GetGpus())}) + } + if len(jobs) == 0 { + s.mu.Unlock() + return + } + sink := s.jobSink + s.recording = true + s.mu.Unlock() + go func() { + ctx, cancel := context.WithTimeout(context.Background(), time.Minute) + defer cancel() + err := sink(ctx, jobs) + s.mu.Lock() + defer s.mu.Unlock() + s.recording = false + if err != nil { + s.log.Warn("could not record the Slurm jobs that ended; tried again with the next status", "jobs", len(jobs), "err", err) + return + } + for _, j := range jobs { + s.recorded[j.ID] = time.Now() + } + }() } // HoldFor is how long a new leader waits for a node it neither has heard @@ -276,7 +383,11 @@ func (s *Service) OnSlurmStatus(node string, st *llmgwv1.SlurmStatus) { s.mu.Lock() prev := s.status[node] s.status[node] = r + isCtl := s.settings().Controller == node s.mu.Unlock() + if isCtl && len(st.GetEnded()) > 0 { + s.takeJobs(st.GetEnded()) + } if r.Error != "" && (prev == nil || prev.Error != r.Error) { s.log.Warn("a node could not keep its part in Slurm", "node", node, "err", r.Error) } @@ -517,6 +628,7 @@ func (s *Service) cluster(set config.SlurmSettings) (slurm.Cluster, []string) { return c, nil } c.ControllerAddr = s.controllerAddr(set) + c.EnforceAccounts = set.Users == config.UsersLlmgw && s.accountList != nil var waiting []string for _, node := range set.Nodes { sn, ok := s.reported(node) @@ -582,6 +694,15 @@ func (s *Service) message(set config.SlurmSettings, node string, lent bool, key if !s.sharedPart(set, node, m) || !s.usersPart(set, node, m) { return nil } + if isCtl && set.Users == config.UsersLlmgw && s.accountList != nil { + if !s.accountsOK { + return nil + } + m.ManageAccounts = true + for _, a := range s.accounts { + m.Accounts = append(m.Accounts, &llmgwv1.SlurmAccount{Name: a.Name, Parent: a.Parent, Users: slices.Clone(a.Users)}) + } + } return m } @@ -675,6 +796,11 @@ func (s *Service) Status() Status { Partition: set.PartitionName(), Ctld: s.status[set.Controller], Shared: set.Shared, Users: set.Users} if set.Users == config.UsersLlmgw { out.UserCount = len(s.users) + for _, a := range s.accounts { + if a.Project != "" { + out.Accounts++ + } + } } if !set.Enabled { return out diff --git a/internal/control/slurmsvc/service_test.go b/internal/control/slurmsvc/service_test.go index 9bc4434..30f750b 100644 --- a/internal/control/slurmsvc/service_test.go +++ b/internal/control/slurmsvc/service_test.go @@ -11,6 +11,7 @@ import ( "gitea.mynodes.xyz/ilker/llmgw/internal/agent/discovery" "gitea.mynodes.xyz/ilker/llmgw/internal/config" "gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory" + "gitea.mynodes.xyz/ilker/llmgw/internal/slurm" ) type fakeInv struct { @@ -323,3 +324,59 @@ func TestPyxisIsLoadedOnlyOnceEveryMachineHasIt(t *testing.T) { t.Fatal("pyxis still loaded with containers off") } } + +func TestTheControllerKeepsTheProjectsAccountsAndJobsAreRecordedOnce(t *testing.T) { + r := newRig() + r.set.Users = config.UsersLlmgw + r.svc.SetUsers(func(context.Context) ([]User, error) { return []User{{Login: "ayse", UID: 200001}}, nil }) + read := 0 + r.svc.SetAccounts(func(context.Context) ([]slurm.Account, error) { + read++ + return []slurm.Account{{Name: "test", Parent: "root"}, + {Name: "test-arge-deneme", Parent: "test", Project: "p1", Users: []string{"ayse"}}}, nil + }) + got := make(chan []JobRecord, 4) + r.svc.SetJobs(func(_ context.Context, jobs []JobRecord) error { got <- jobs; return nil }) + r.report("cuda01", &llmgwv1.SlurmStatus{Hardware: hw}) + r.tick() + ctl := r.agents.last("node06") + if !ctl.GetManageAccounts() || len(ctl.GetAccounts()) != 2 || ctl.GetAccounts()[1].GetUsers()[0] != "ayse" { + t.Fatalf("controller told %v", ctl.GetAccounts()) + } + if r.agents.last("cuda01").GetManageAccounts() { + t.Error("a node was told the accounts") + } + if !strings.Contains(ctl.GetFiles()["slurm.conf"], "AccountingStorageEnforce=associations") { + t.Error("jobs not held to accounts") + } + if r.svc.Status().Accounts != 1 || len(r.svc.AccountsOf("ayse")) != 1 { + t.Errorf("status %d, ayse's %v", r.svc.Status().Accounts, r.svc.AccountsOf("ayse")) + } + ended := []*llmgwv1.SlurmJob{{Id: "41", User: "ayse", Account: "test-arge-deneme", ElapsedSeconds: 60, Gpus: 1}} + r.report("node06", &llmgwv1.SlurmStatus{Ended: ended}) + select { + case jobs := <-got: + if len(jobs) != 1 || jobs[0].Login != "ayse" || jobs[0].GPUs != 1 { + t.Fatalf("jobs %+v", jobs) + } + case <-time.After(5 * time.Second): + t.Fatal("not recorded") + } + // The next status says it again: not recorded twice. + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + r.svc.mu.Lock() + busy := r.svc.recording + r.svc.mu.Unlock() + if !busy { + break + } + time.Sleep(10 * time.Millisecond) + } + r.report("node06", &llmgwv1.SlurmStatus{Ended: ended}) + select { + case jobs := <-got: + t.Fatalf("recorded again: %+v", jobs) + case <-time.After(200 * time.Millisecond): + } +} diff --git a/internal/slurm/accounts.go b/internal/slurm/accounts.go new file mode 100644 index 0000000..90347b2 --- /dev/null +++ b/internal/slurm/accounts.go @@ -0,0 +1,193 @@ +package slurm + +import ( + "crypto/sha256" + "encoding/hex" + "fmt" + "slices" + "strconv" + "strings" + "time" +) + +// Slurm's accounts from llmgw's tree: an account for each organization, +// team and project, nested as they are, and on each project's the users +// that may use it, its members and its team's and organization's. A job +// is charged to the project its --account names (`sbatch --account +// acme-ml-training`), and what it used to that project's spend. + +// Account is one Slurm account llmgw keeps. +type Account struct { + Name string + Parent string // "root" at the top + // Project is the llmgw project it is for; empty for an organization's + // or a team's. + Project string + // Users are the logins that may run jobs on it. + Users []string +} + +// TenancyProject is a project with where it is. +type TenancyProject struct { + ID, Name, TeamID, TeamName, OrgID, OrgName string +} + +// TenancyMember is a membership: a person in an organization, team or +// project ("org", "team", "project"). +type TenancyMember struct { + UserID, Scope, ScopeID string +} + +// Accounts are the accounts for projects, with logins by user id. Parents +// come before their children, as sacctmgr needs them. +func Accounts(projects []TenancyProject, members []TenancyMember, logins map[string]string) []Account { + // By kind and id: an organization, its team and its project may share + // an id (the first install's are all "default"). + names := map[string]string{} // kind/id -> name + taken := map[string]string{} // name -> kind/id + name := func(id string, parts ...string) string { + if n, ok := names[id]; ok { + return n + } + n := AccountName(parts...) + if other, ok := taken[n]; ok && other != id { + // Two places whose names come out the same: told apart by id. + h := sha256.Sum256([]byte(id)) + n = AccountName(append(parts, hex.EncodeToString(h[:3]))...) + } + names[id], taken[n] = n, id + return n + } + in := map[string][]string{} // kind/id -> logins + for _, m := range members { + if l, ok := logins[m.UserID]; ok { + in[m.Scope+"/"+m.ScopeID] = append(in[m.Scope+"/"+m.ScopeID], l) + } + } + ps := slices.Clone(projects) + slices.SortFunc(ps, func(a, b TenancyProject) int { + return strings.Compare(a.OrgName+"\x00"+a.TeamName+"\x00"+a.Name, b.OrgName+"\x00"+b.TeamName+"\x00"+b.Name) + }) + var out []Account + seen := map[string]bool{} + for _, p := range ps { + org := name("org/"+p.OrgID, p.OrgName) + team := name("team/"+p.TeamID, p.OrgName, p.TeamName) + proj := name("project/"+p.ID, p.OrgName, p.TeamName, p.Name) + if !seen[org] { + out = append(out, Account{Name: org, Parent: "root"}) + seen[org] = true + } + if !seen[team] { + out = append(out, Account{Name: team, Parent: org}) + seen[team] = true + } + var users []string + for _, id := range []string{"project/" + p.ID, "team/" + p.TeamID, "org/" + p.OrgID} { + users = append(users, in[id]...) + } + slices.Sort(users) + out = append(out, Account{Name: proj, Parent: team, Project: p.ID, Users: slices.Compact(users)}) + } + return out +} + +// AccountName is a Slurm account name from names: lower case, ASCII, +// joined with "-", at most 64 characters. +func AccountName(parts ...string) string { + var segs []string + for _, p := range parts { + var b strings.Builder + for _, r := range strings.ToLower(p) { + switch { + case r >= 'a' && r <= 'z', r >= '0' && r <= '9', r == '_': + b.WriteRune(r) + case strings.ContainsRune("çğıöşüâîû", r): + b.WriteString(map[rune]string{'ç': "c", 'ğ': "g", 'ı': "i", 'ö': "o", 'ş': "s", 'ü': "u", + 'â': "a", 'î': "i", 'û': "u"}[r]) + default: + b.WriteByte('_') + } + } + if s := strings.Trim(b.String(), "_"); s != "" { + segs = append(segs, s) + } + } + n := strings.Join(segs, "-") + if n == "" { + n = "x" + } + if len(n) > 64 { + h := sha256.Sum256([]byte(n)) + n = n[:57] + "-" + hex.EncodeToString(h[:3]) + } + return n +} + +// Association is one line of `sacctmgr show assoc format=account,user,parentname`. +type Association struct { + Account, User, Parent string +} + +// AssocFormat is what ParseAssociations reads. +var AssocFormat = []string{"show", "assoc", "format=account,user,parentname", "-n", "-P"} + +// ParseAssociations reads sacctmgr's parsable lines. +func ParseAssociations(out string) []Association { + var as []Association + for _, line := range strings.Split(strings.TrimSpace(out), "\n") { + f := strings.Split(line, "|") + if len(f) < 3 || f[0] == "" { + continue + } + as = append(as, Association{Account: f[0], User: f[1], Parent: f[2]}) + } + return as +} + +// Job is a finished job, as sacct has it. +type Job struct { + ID, User, Account, Partition, State string + Start, End time.Time + Elapsed time.Duration + CPUs, GPUs int +} + +// SacctArgs are sacct's for the jobs that ended in the last day. +var SacctArgs = []string{"-a", "-X", "-n", "-P", "--noconvert", "-S", "now-1days", "-E", "now", + "-s", "CD,F,TO,CA,NF,OOM,PR,DL,BF", + "-o", "JobIDRaw,User,Account,Partition,Start,End,ElapsedRaw,AllocTRES,State"} + +// ParseJobs reads sacct's lines; a job that has not ended is left out. +func ParseJobs(out string) ([]Job, error) { + var jobs []Job + for _, line := range strings.Split(strings.TrimSpace(out), "\n") { + if line == "" { + continue + } + f := strings.Split(line, "|") + if len(f) < 9 { + return nil, fmt.Errorf("sacct said %q, not 9 fields", line) + } + end, err := time.ParseInLocation("2006-01-02T15:04:05", f[5], time.Local) + if err != nil { + continue // Unknown: still running, or never started + } + start, _ := time.ParseInLocation("2006-01-02T15:04:05", f[4], time.Local) + secs, _ := strconv.ParseInt(f[6], 10, 64) + j := Job{ID: f[0], User: f[1], Account: f[2], Partition: f[3], Start: start, End: end, + Elapsed: time.Duration(secs) * time.Second, State: strings.Fields(f[8] + " ")[0]} + for _, t := range strings.Split(f[7], ",") { + k, v, _ := strings.Cut(t, "=") + n, _ := strconv.Atoi(v) + switch k { + case "cpu": + j.CPUs = n + case "gres/gpu": + j.GPUs = n + } + } + jobs = append(jobs, j) + } + return jobs, nil +} diff --git a/internal/slurm/accounts_test.go b/internal/slurm/accounts_test.go new file mode 100644 index 0000000..ad40e7c --- /dev/null +++ b/internal/slurm/accounts_test.go @@ -0,0 +1,89 @@ +package slurm + +import ( + "slices" + "testing" + "time" +) + +func TestAccountsFollowTheTree(t *testing.T) { + projects := []TenancyProject{ + {ID: "p1", Name: "deneme", TeamID: "t1", TeamName: "arge", OrgID: "o1", OrgName: "Test"}, + {ID: "p2", Name: "Görüntü İşleme", TeamID: "t1", TeamName: "arge", OrgID: "o1", OrgName: "Test"}, + {ID: "p3", Name: "x", TeamID: "t2", TeamName: "ops", OrgID: "o2", OrgName: "Other"}, + } + members := []TenancyMember{ + {UserID: "u-ayse", Scope: "project", ScopeID: "p1"}, + {UserID: "u-ada", Scope: "team", ScopeID: "t1"}, + {UserID: "u-bob", Scope: "org", ScopeID: "o2"}, + {UserID: "u-nologin", Scope: "project", ScopeID: "p1"}, + } + logins := map[string]string{"u-ayse": "ayse", "u-ada": "ada", "u-bob": "bob"} + got := Accounts(projects, members, logins) + var names []string + by := map[string]Account{} + for _, a := range got { + names = append(names, a.Name) + by[a.Name] = a + } + // Parents first, each before its children. + want := []string{"other", "other-ops", "other-ops-x", "test", "test-arge", "test-arge-goruntu_isleme", "test-arge-deneme"} + if !slices.Equal(names, want) { + t.Fatalf("names %v", names) + } + if a := by["test-arge-deneme"]; a.Parent != "test-arge" || a.Project != "p1" || !slices.Equal(a.Users, []string{"ada", "ayse"}) { + t.Errorf("deneme %+v", a) + } + if a := by["other-ops-x"]; !slices.Equal(a.Users, []string{"bob"}) { + t.Errorf("x %+v", a) + } + if a := by["test"]; a.Parent != "root" || a.Project != "" || len(a.Users) != 0 { + t.Errorf("org %+v", a) + } + if n := AccountName("Test", "arge", "Görüntü İşleme"); n != "test-arge-goruntu_isleme" { + t.Errorf("name %q", n) + } + // The first install's tree: one id for all three. + def := Accounts([]TenancyProject{{ID: "default", Name: "default", TeamID: "default", TeamName: "default", + OrgID: "default", OrgName: "default"}}, []TenancyMember{{UserID: "u", Scope: "project", ScopeID: "default"}}, + map[string]string{"u": "admin2"}) + if len(def) != 3 || def[2].Name != "default-default-default" || def[2].Parent != "default-default" || + def[1].Parent != "default" || !slices.Equal(def[2].Users, []string{"admin2"}) { + t.Errorf("default tree %+v", def) + } + // Two that come out the same are told apart. + same := Accounts([]TenancyProject{ + {ID: "a", Name: "ml", TeamID: "t", TeamName: "x", OrgID: "o", OrgName: "A"}, + {ID: "b", Name: "ML", TeamID: "t", TeamName: "x", OrgID: "o", OrgName: "A"}, + }, nil, nil) + if same[2].Name == same[3].Name { + t.Errorf("same name twice: %+v", same) + } +} + +func TestParseJobs(t *testing.T) { + out := "41|ayse|test-arge-deneme|gpu|2026-11-15T10:00:00|2026-11-15T10:30:00|1800|billing=4,cpu=4,gres/gpu=1,mem=16G,node=1|COMPLETED\n" + + "42|ada|test-arge-deneme|gpu|2026-11-15T10:00:00|Unknown|0|cpu=2|RUNNING\n" + + "43|root|root|gpu|2026-11-15T09:00:00|2026-11-15T09:00:05|5|cpu=1,mem=1G|CANCELLED by 0\n" + jobs, err := ParseJobs(out) + if err != nil { + t.Fatal(err) + } + if len(jobs) != 2 { + t.Fatalf("jobs %+v", jobs) + } + j := jobs[0] + if j.ID != "41" || j.User != "ayse" || j.GPUs != 1 || j.CPUs != 4 || j.Elapsed != 30*time.Minute || j.State != "COMPLETED" { + t.Errorf("job %+v", j) + } + if jobs[1].State != "CANCELLED" { + t.Errorf("state %q", jobs[1].State) + } + if _, err := ParseJobs("short|line\n"); err == nil { + t.Error("a short line taken") + } + as := ParseAssociations("root||\ntest|ayse|test-arge\n") + if len(as) != 2 || as[1] != (Association{"test", "ayse", "test-arge"}) { + t.Errorf("assoc %+v", as) + } +} diff --git a/internal/slurm/slurm.go b/internal/slurm/slurm.go index 5612701..62ceca7 100644 --- a/internal/slurm/slurm.go +++ b/internal/slurm/slurm.go @@ -131,6 +131,10 @@ type Cluster struct { Controller string ControllerAddr string Partition string + // EnforceAccounts refuses a job whose user has no account in Slurm's + // accounting: with llmgw's users, each is on the accounts of the + // projects they are in, and a job is charged to one of them. + EnforceAccounts bool // Nodes without hardware (not reported yet) are left out of the files // until they are. Nodes []Node @@ -199,6 +203,9 @@ func (c Cluster) SlurmConf() string { // The jobs' records, in slurmdbd beside slurmctld (DbdConf). fmt.Fprintf(&b, "AccountingStorageType=accounting_storage/slurmdbd\nAccountingStorageHost=%s\nAccountingStoragePort=%d\n", c.ControllerAddr, DbdPort) + if c.EnforceAccounts { + b.WriteString("# A job runs on an account of a project its user is in.\nAccountingStorageEnforce=associations\n") + } b.WriteString(`SlurmctldParameters=enable_configless SlurmctldPort=6817 SlurmdPort=6818 diff --git a/internal/store/copy.go b/internal/store/copy.go index 47d9240..1103497 100644 --- a/internal/store/copy.go +++ b/internal/store/copy.go @@ -45,6 +45,7 @@ var copyOrder = []string{ "usage_1h", "usage_1d", "held_1h", + "slurm_jobs", "models", // No foreign key to models, deliberately: a measurement outlives the // catalog entry it was taken for, and losing it when somebody removes a diff --git a/internal/store/migrations/postgres/0054_slurm_jobs.sql b/internal/store/migrations/postgres/0054_slurm_jobs.sql new file mode 100644 index 0000000..e9e14a7 --- /dev/null +++ b/internal/store/migrations/postgres/0054_slurm_jobs.sql @@ -0,0 +1,20 @@ +-- Ported from the SQLite migration of the same number. +-- The Slurm jobs that ended, each recorded once (job_key is the cluster's +-- name and the job's id), with what it used and cost; a usage row (model +-- "slurm") carries the cost into the usage pages with the requests'. +CREATE TABLE slurm_jobs ( + job_key TEXT PRIMARY KEY, + user_id TEXT NOT NULL DEFAULT '', + project_id TEXT NOT NULL DEFAULT '', + login TEXT NOT NULL DEFAULT '', + account TEXT NOT NULL DEFAULT '', + partition TEXT NOT NULL DEFAULT '', + state TEXT NOT NULL DEFAULT '', + started_at BIGINT NOT NULL DEFAULT 0, + ended_at BIGINT NOT NULL DEFAULT 0, + elapsed_seconds BIGINT NOT NULL DEFAULT 0, + cpus BIGINT NOT NULL DEFAULT 0, + gpus BIGINT NOT NULL DEFAULT 0, + cost_micros BIGINT NOT NULL DEFAULT 0 +); +CREATE INDEX slurm_jobs_by_project ON slurm_jobs(project_id, ended_at); diff --git a/internal/store/migrations/sqlite/0054_slurm_jobs.sql b/internal/store/migrations/sqlite/0054_slurm_jobs.sql new file mode 100644 index 0000000..61d92f9 --- /dev/null +++ b/internal/store/migrations/sqlite/0054_slurm_jobs.sql @@ -0,0 +1,19 @@ +-- The Slurm jobs that ended, each recorded once (job_key is the cluster's +-- name and the job's id), with what it used and cost; a usage row (model +-- "slurm") carries the cost into the usage pages with the requests'. +CREATE TABLE slurm_jobs ( + job_key TEXT PRIMARY KEY, + user_id TEXT NOT NULL DEFAULT '', + project_id TEXT NOT NULL DEFAULT '', + login TEXT NOT NULL DEFAULT '', + account TEXT NOT NULL DEFAULT '', + partition TEXT NOT NULL DEFAULT '', + state TEXT NOT NULL DEFAULT '', + started_at BIGINT NOT NULL DEFAULT 0, + ended_at BIGINT NOT NULL DEFAULT 0, + elapsed_seconds BIGINT NOT NULL DEFAULT 0, + cpus BIGINT NOT NULL DEFAULT 0, + gpus BIGINT NOT NULL DEFAULT 0, + cost_micros BIGINT NOT NULL DEFAULT 0 +); +CREATE INDEX slurm_jobs_by_project ON slurm_jobs(project_id, ended_at); diff --git a/internal/store/slurm_jobs.go b/internal/store/slurm_jobs.go new file mode 100644 index 0000000..69fe596 --- /dev/null +++ b/internal/store/slurm_jobs.go @@ -0,0 +1,65 @@ +package store + +import ( + "context" + "database/sql" + "time" +) + +// SlurmProject is a project with where it is, for Slurm's accounts. +type SlurmProject struct { + ID, Name, TeamID, TeamName, OrgID, OrgName string +} + +// SlurmProjects are every project, with its team and organization. +func (s *DB) SlurmProjects(ctx context.Context) ([]SlurmProject, error) { + return queryAll(ctx, s.db, ` + SELECT p.id, p.name, t.id, t.name, o.id, o.name + FROM projects p JOIN teams t ON t.id = p.team_id JOIN organizations o ON o.id = t.org_id + ORDER BY o.name, t.name, p.name`, nil, func(rows *sql.Rows) (SlurmProject, error) { + var p SlurmProject + return p, rows.Scan(&p.ID, &p.Name, &p.TeamID, &p.TeamName, &p.OrgID, &p.OrgName) + }) +} + +// SlurmJob is a job that ended, as recorded: once, by its cluster's id. +type SlurmJob struct { + // Key is the cluster's name and the job's id: unique among jobs. + Key string + UserID, ProjectID string + Login, Account, Partition string + State string + Start, End time.Time + ElapsedSeconds, CPUs, GPUs int64 + CostMicros int64 +} + +// RecordSlurmJobs records jobs not recorded before, each with a usage row +// (model "slurm") so its cost is in the usage pages and the exports with +// the requests'. Answers how many were new. +func (s *DB) RecordSlurmJobs(ctx context.Context, jobs []SlurmJob) (int, error) { + n := 0 + err := s.inTx(ctx, func(tx conn) error { + var usage []Usage + for _, j := range jobs { + res, err := tx.ExecContext(ctx, + `INSERT INTO slurm_jobs (job_key, user_id, project_id, login, account, partition, state, + started_at, ended_at, elapsed_seconds, cpus, gpus, cost_micros) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(job_key) DO NOTHING`, + j.Key, j.UserID, j.ProjectID, j.Login, j.Account, j.Partition, j.State, + millis(j.Start), millis(j.End), j.ElapsedSeconds, j.CPUs, j.GPUs, j.CostMicros) + if err != nil { + return err + } + if k, _ := res.RowsAffected(); k == 0 { + continue + } + n++ + usage = append(usage, Usage{At: j.End, ProjectID: j.ProjectID, UserID: j.UserID, Model: "slurm", + ServedModel: "slurm/" + j.Partition, CostMicros: j.CostMicros, Status: 200, + LatencyMs: j.ElapsedSeconds * 1000, RequestID: "slurm-" + j.Key}) + } + return insertUsage(ctx, tx, usage) + }) + return n, err +} diff --git a/internal/store/slurm_jobs_test.go b/internal/store/slurm_jobs_test.go new file mode 100644 index 0000000..6b621c7 --- /dev/null +++ b/internal/store/slurm_jobs_test.go @@ -0,0 +1,41 @@ +package store + +import ( + "testing" + "time" +) + +func TestASlurmJobIsRecordedOnce(t *testing.T) { + s := capacityStore(t) + ctx := t.Context() + end := time.Now().Add(-time.Hour).Truncate(time.Second) + j := SlurmJob{Key: "llmgw/41", ProjectID: "p1", UserID: "u1", Login: "ayse", Account: "test-arge-deneme", + Partition: "gpu", State: "COMPLETED", Start: end.Add(-30 * time.Minute), End: end, ElapsedSeconds: 1800, + CPUs: 4, GPUs: 1, CostMicros: 675000} + for range 2 { + n, err := s.RecordSlurmJobs(ctx, []SlurmJob{j}) + if err != nil { + t.Fatal(err) + } + _ = n + } + byModel, err := s.UsageByModel(ctx, end.Add(-24*time.Hour)) + if err != nil { + t.Fatal(err) + } + var found bool + for _, m := range byModel { + if m.Model == "slurm/gpu" { + found = true + if m.Requests != 1 || m.CostMicros != 675000 { + t.Errorf("recorded twice or wrong: %+v", m) + } + } + } + if !found { + t.Fatalf("no usage row: %+v", byModel) + } + if n, _ := s.RecordSlurmJobs(ctx, []SlurmJob{j}); n != 0 { + t.Errorf("recorded again: %d", n) + } +} diff --git a/internal/store/sqlite.go b/internal/store/sqlite.go index a77787e..4c0c9f4 100644 --- a/internal/store/sqlite.go +++ b/internal/store/sqlite.go @@ -447,32 +447,38 @@ func (s *DB) WriteUsage(ctx context.Context, records []Usage) error { if len(records) == 0 { return nil } - return s.inTx(ctx, func(tx conn) error { - stmt, err := tx.PrepareContext(ctx, - `INSERT INTO usage (ts, key_id, project_id, model, served_model, - prompt_tokens, completion_tokens, cost_micros, status, latency_ms, user_id, - request_id, replica, queue_ms, ttft_ms, - guard_action, guard_policy, guard_category, guard_rule, guard_ms, end_user, tags, cache_hit, images, meshes) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`) - if err != nil { + return s.inTx(ctx, func(tx conn) error { return insertUsage(ctx, tx, records) }) +} + +// insertUsage writes records, and their sums, in tx. +func insertUsage(ctx context.Context, tx conn, records []Usage) error { + if len(records) == 0 { + return nil + } + stmt, err := tx.PrepareContext(ctx, + `INSERT INTO usage (ts, key_id, project_id, model, served_model, + prompt_tokens, completion_tokens, cost_micros, status, latency_ms, user_id, + request_id, replica, queue_ms, ttft_ms, + guard_action, guard_policy, guard_category, guard_rule, guard_ms, end_user, tags, cache_hit, images, meshes) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`) + if err != nil { + return err + } + defer stmt.Close() + for _, r := range records { + hit := 0 + if r.CacheHit { + hit = 1 + } + if _, err := stmt.ExecContext(ctx, millis(r.At), r.KeyID, r.ProjectID, + r.Model, r.ServedModel, r.PromptTokens, r.CompletionTokens, + r.CostMicros, r.Status, r.LatencyMs, r.UserID, + r.RequestID, r.Replica, r.QueueMs, r.TTFTMs, + r.GuardAction, r.GuardPolicy, r.GuardCategory, r.GuardRule, r.GuardMs, r.EndUser, tagsJSON(r.Tags), hit, r.Images, r.Meshes); err != nil { return err } - defer stmt.Close() - for _, r := range records { - hit := 0 - if r.CacheHit { - hit = 1 - } - if _, err := stmt.ExecContext(ctx, millis(r.At), r.KeyID, r.ProjectID, - r.Model, r.ServedModel, r.PromptTokens, r.CompletionTokens, - r.CostMicros, r.Status, r.LatencyMs, r.UserID, - r.RequestID, r.Replica, r.QueueMs, r.TTFTMs, - r.GuardAction, r.GuardPolicy, r.GuardCategory, r.GuardRule, r.GuardMs, r.EndUser, tagsJSON(r.Tags), hit, r.Images, r.Meshes); err != nil { - return err - } - } - return writeRollup(ctx, tx, records) - }) + } + return writeRollup(ctx, tx, records) } // SpendSince totals cost per key since a moment.