Slurm's jobs in llmgw's accounting: projects as accounts, jobs priced
CI / test (push) Failing after 7m24s
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:
1 parent
8547528fb7
commit
fcda4d3f85
22 files changed
+1390
-76
No files matched your search
+10
-1
@@ -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
|
||||
|
||||
@@ -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
@@ -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,
|
||||
},
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
@@ -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'."
|
||||
}
|
||||
}
|
||||
},
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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"))
|
||||
}
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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):
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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);
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
@@ -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.
|
||||
|
||||
Reference in new issue
Block a user