Slurm's jobs in llmgw's accounting: projects as accounts, jobs priced
CI / test (push) Failing after 7m24s

With slurm.users llmgw, each project is a Slurm account, its organization's
and team's above it, with the project's members and its team's and
organization's on it; slurm.conf holds a job to its user's accounts. The
controller's agent keeps them with sacctmgr, only what it made, and reports
the jobs that ended in the last day; the leader records each once, its
GPU-hours (core-hours without cards) priced at auth.gpu_hour_cost into its
project's usage as model slurm/PARTITION.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Ilker ManapandClaude Opus 5.5 committed 2026-11-15 11:54:01 +01:00
1 parent 8547528fb7
commit fcda4d3f85
22 files changed
+1390 -76

No files matched your search

+10 -1
View File
@@ -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
+7
View File
@@ -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
+283 -51
View File
@@ -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,
},
+30
View File
@@ -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;
}
+4
View File
@@ -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'."
}
}
},
+11
View File
@@ -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
+214
View File
@@ -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)
}
+83
View File
@@ -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"))
}
}
+14
View File
@@ -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 {
+3
View File
@@ -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)
+83
View File
@@ -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
}
+126
View File
@@ -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
+57
View File
@@ -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):
}
}
+193
View File
@@ -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
}
+89
View File
@@ -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)
}
}
+7
View File
@@ -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
+1
View File
@@ -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
@@ -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);
@@ -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);
+65
View File
@@ -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
}
+41
View File
@@ -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)
}
}
+30 -24
View File
@@ -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.