Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions shelper/local.go
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ type SlurmMetadata struct {
NumNodes string `json:"NumNodes"`
ArrayJobID string `json:"ArrayJobID"`
ArrayTaskID string `json:"ArrayTaskID"`
Comment string `json:"Comment"`
}

// SlurmMetadataList is a struct that contains metadata about a slurm job
Expand All @@ -33,6 +34,7 @@ type SlurmMetadataList struct {
NumNodes []string
ArrayJobID []string
ArrayTaskID []string
Comment []string
}

func parseNewLineToList(input string) []string {
Expand Down
32 changes: 20 additions & 12 deletions shelper/slurm_helpers.go
Original file line number Diff line number Diff line change
Expand Up @@ -22,9 +22,9 @@ func parseGRES(gresOut string) []string {
}

indexString := strings.SplitN(indicesKey[1], ":", 2)[1]
indices := strings.Split(indexString, ",")
indices := strings.SplitSeq(indexString, ",")

for _, index := range indices {
for index := range indices {
if strings.Contains(index, "-") {
indexRange := strings.Split(index, "-")
st, err1 := strconv.Atoi(indexRange[0])
Expand Down Expand Up @@ -74,6 +74,7 @@ func GetGPUData(GPUToSlurm map[string]SlurmMetadata) SlurmMetadataList {
allPartition := make(map[string]bool)
allAccount := make(map[string]bool)
allNumNodes := make(map[string]bool)
allComment := make(map[string]bool)
for _, value := range GPUToSlurm {
allJobID[value.JobID] = true
allJobName[value.JobName] = true
Expand All @@ -84,6 +85,7 @@ func GetGPUData(GPUToSlurm map[string]SlurmMetadata) SlurmMetadataList {
allPartition[value.Partition] = true
allAccount[value.Account] = true
allNumNodes[value.NumNodes] = true
allComment[value.Comment] = true
}
return SlurmMetadataList{
JobID: setToSlice(allJobID),
Expand All @@ -95,6 +97,7 @@ func GetGPUData(GPUToSlurm map[string]SlurmMetadata) SlurmMetadataList {
Partition: setToSlice(allPartition),
Account: setToSlice(allAccount),
NumNodes: setToSlice(allNumNodes),
Comment: setToSlice(allComment),
}
}

Expand All @@ -118,12 +121,13 @@ func AttributeGPU2SlurmMetadata(jobMetadata []string, hostname string, GPU2Slurm
allAccount := make(map[string]bool)
allPartition := make(map[string]bool)
allNumNodes := make(map[string]bool)
allComment := make(map[string]bool)

lines := strings.Split(jm, "\n")
for _, line := range lines {
field := strings.Fields(line)
lines := strings.SplitSeq(jm, "\n")
for line := range lines {
field := strings.FieldsSeq(line)

for _, data := range field {
for data := range field {
parts := strings.SplitN(data, "=", 2)
if parts[0] == "UserId" {
end := strings.Index(parts[1], "(")
Expand Down Expand Up @@ -161,6 +165,9 @@ func AttributeGPU2SlurmMetadata(jobMetadata []string, hostname string, GPU2Slurm
if parts[0] == "NumNodes" {
allNumNodes[parts[1]] = true
}
if parts[0] == "Comment" {
allComment[parts[1]] = true
}
if parts[0] == "GRES" {
gresIndex = append(gresIndex, parseGRES(parts[1])...)
}
Expand All @@ -177,6 +184,7 @@ func AttributeGPU2SlurmMetadata(jobMetadata []string, hostname string, GPU2Slurm
Account: stringifySet(allAccount),
Partition: stringifySet(allPartition),
NumNodes: stringifySet(allNumNodes),
Comment: stringifySet(allComment),
}

for _, gpu := range gresIndex {
Expand All @@ -187,11 +195,11 @@ func AttributeGPU2SlurmMetadata(jobMetadata []string, hostname string, GPU2Slurm

// GetHostList takes a slurm job metadata string and returns the hostlist
func GetHostList(jobMetadata string) string {
lines := strings.Split(jobMetadata, "\n")
lines := strings.SplitSeq(jobMetadata, "\n")

for _, line := range lines {
field := strings.Fields(line)
for _, data := range field {
for line := range lines {
field := strings.FieldsSeq(line)
for data := range field {
parts := strings.SplitN(data, "=", 2)
if parts[0] == "NodeList" {
return parts[1]
Expand Down Expand Up @@ -259,8 +267,8 @@ func hostnameMatchesGroup(hostname string, group string) bool {
return false
}

ranges := strings.Split(rangesStr, ",")
for _, r := range ranges {
ranges := strings.SplitSeq(rangesStr, ",")
for r := range ranges {
if strings.Contains(r, "-") {
bounds := strings.Split(r, "-")
start, err := strconv.Atoi(bounds[0])
Expand Down
29 changes: 29 additions & 0 deletions shelper/slurm_helpers_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ func TestGetSlurmDataFromSlurmLineAllGpus(t *testing.T) {
Account: "test_account",
Partition: "learn",
NumNodes: "1",
Comment: "test_comment",
},
"1": {
User: "test_username",
Expand All @@ -40,6 +41,7 @@ func TestGetSlurmDataFromSlurmLineAllGpus(t *testing.T) {
Account: "test_account",
Partition: "learn",
NumNodes: "1",
Comment: "test_comment",
},
"2": {
User: "test_username",
Expand All @@ -51,6 +53,7 @@ func TestGetSlurmDataFromSlurmLineAllGpus(t *testing.T) {
Account: "test_account",
Partition: "learn",
NumNodes: "1",
Comment: "test_comment",
},
"3": {
User: "test_username",
Expand All @@ -62,6 +65,7 @@ func TestGetSlurmDataFromSlurmLineAllGpus(t *testing.T) {
Account: "test_account",
Partition: "learn",
NumNodes: "1",
Comment: "test_comment",
},
"4": {
User: "test_username",
Expand All @@ -73,6 +77,7 @@ func TestGetSlurmDataFromSlurmLineAllGpus(t *testing.T) {
Account: "test_account",
Partition: "learn",
NumNodes: "1",
Comment: "test_comment",
},
"5": {
User: "test_username",
Expand All @@ -84,6 +89,7 @@ func TestGetSlurmDataFromSlurmLineAllGpus(t *testing.T) {
Account: "test_account",
Partition: "learn",
NumNodes: "1",
Comment: "test_comment",
},
"6": {
User: "test_username",
Expand All @@ -95,6 +101,7 @@ func TestGetSlurmDataFromSlurmLineAllGpus(t *testing.T) {
Account: "test_account",
Partition: "learn",
NumNodes: "1",
Comment: "test_comment",
},
"7": {
User: "test_username",
Expand All @@ -106,6 +113,7 @@ func TestGetSlurmDataFromSlurmLineAllGpus(t *testing.T) {
Account: "test_account",
Partition: "learn",
NumNodes: "1",
Comment: "test_comment",
},
}

Expand Down Expand Up @@ -133,6 +141,7 @@ func TestGetSlurmDataFromSlurmLineSomeGpus(t *testing.T) {
Account: "test_account",
Partition: "learn",
NumNodes: "1",
Comment: "test_comment",
},
"1": {
User: "test_username",
Expand All @@ -144,6 +153,7 @@ func TestGetSlurmDataFromSlurmLineSomeGpus(t *testing.T) {
Account: "test_account",
Partition: "learn",
NumNodes: "1",
Comment: "test_comment",
},
"2": {
User: "test_username",
Expand All @@ -155,6 +165,7 @@ func TestGetSlurmDataFromSlurmLineSomeGpus(t *testing.T) {
Account: "test_account",
Partition: "learn",
NumNodes: "1",
Comment: "test_comment",
},
"3": {
User: "test_username",
Expand All @@ -166,6 +177,7 @@ func TestGetSlurmDataFromSlurmLineSomeGpus(t *testing.T) {
Account: "test_account",
Partition: "learn",
NumNodes: "1",
Comment: "test_comment",
},
}

Expand Down Expand Up @@ -208,6 +220,7 @@ func TestGetSlurmDataFromSlurmLineUniqueEntries(t *testing.T) {
Account: "test2_account",
Partition: "learn",
NumNodes: "1",
Comment: "test_comment",
},
"1": {
User: "test_username",
Expand All @@ -219,6 +232,7 @@ func TestGetSlurmDataFromSlurmLineUniqueEntries(t *testing.T) {
Account: "test2_account",
Partition: "learn",
NumNodes: "1",
Comment: "test_comment",
},
"2": {
User: "test_username",
Expand All @@ -230,6 +244,7 @@ func TestGetSlurmDataFromSlurmLineUniqueEntries(t *testing.T) {
Account: "test2_account",
Partition: "learn",
NumNodes: "1",
Comment: "test_comment",
},
"3": {
User: "test_username",
Expand All @@ -241,6 +256,7 @@ func TestGetSlurmDataFromSlurmLineUniqueEntries(t *testing.T) {
Account: "test2_account",
Partition: "learn",
NumNodes: "1",
Comment: "test_comment",
},
"5": {
User: "test_username_2",
Expand All @@ -252,6 +268,7 @@ func TestGetSlurmDataFromSlurmLineUniqueEntries(t *testing.T) {
Account: "test_account",
Partition: "test",
NumNodes: "3",
Comment: "test_comment_2",
},
"6": {
User: "test_username_2",
Expand All @@ -263,6 +280,7 @@ func TestGetSlurmDataFromSlurmLineUniqueEntries(t *testing.T) {
Account: "test_account",
Partition: "test",
NumNodes: "3",
Comment: "test_comment_2",
},
"7": {
User: "test_username_2",
Expand All @@ -274,6 +292,7 @@ func TestGetSlurmDataFromSlurmLineUniqueEntries(t *testing.T) {
Account: "test_account",
Partition: "test",
NumNodes: "3",
Comment: "test_comment_2",
},
}
AttributeGPU2SlurmMetadata(blocks, "node1751", GPU2Slurm)
Expand All @@ -300,6 +319,7 @@ func TestGetSlurmDataFromSlurmLineMainArrayJob(t *testing.T) {
Account: "test2_account",
Partition: "learn",
NumNodes: "1",
Comment: "test_comment",
},
"1": {
User: "test_username",
Expand All @@ -311,6 +331,7 @@ func TestGetSlurmDataFromSlurmLineMainArrayJob(t *testing.T) {
Account: "test2_account",
Partition: "learn",
NumNodes: "1",
Comment: "test_comment",
},
"2": {
User: "test_username",
Expand All @@ -322,6 +343,7 @@ func TestGetSlurmDataFromSlurmLineMainArrayJob(t *testing.T) {
Account: "test2_account",
Partition: "learn",
NumNodes: "1",
Comment: "test_comment",
},
"3": {
User: "test_username",
Expand All @@ -333,6 +355,7 @@ func TestGetSlurmDataFromSlurmLineMainArrayJob(t *testing.T) {
Account: "test2_account",
Partition: "learn",
NumNodes: "1",
Comment: "test_comment",
},
"5": {
User: "test_username_2",
Expand All @@ -344,6 +367,7 @@ func TestGetSlurmDataFromSlurmLineMainArrayJob(t *testing.T) {
Account: "test_account",
Partition: "test",
NumNodes: "3",
Comment: "test_comment_2",
},
"6": {
User: "test_username_2",
Expand All @@ -355,6 +379,7 @@ func TestGetSlurmDataFromSlurmLineMainArrayJob(t *testing.T) {
Account: "test_account",
Partition: "test",
NumNodes: "3",
Comment: "test_comment_2",
},
"7": {
User: "test_username_2",
Expand All @@ -366,6 +391,7 @@ func TestGetSlurmDataFromSlurmLineMainArrayJob(t *testing.T) {
Account: "test_account",
Partition: "test",
NumNodes: "3",
Comment: "test_comment_2",
},
}
AttributeGPU2SlurmMetadata(blocks, "node1751", GPU2Slurm)
Expand Down Expand Up @@ -555,6 +581,7 @@ func TestGetGPUData(t *testing.T) {
NumNodes: "1",
ArrayJobID: "0",
ArrayTaskID: "0",
Comment: "test_comment",
},
"1": {
JobID: "1234",
Expand All @@ -566,6 +593,7 @@ func TestGetGPUData(t *testing.T) {
NumNodes: "1",
ArrayJobID: "10",
ArrayTaskID: "10",
Comment: "test_comment",
},
}
expectedMetadata := SlurmMetadataList{
Expand All @@ -578,6 +606,7 @@ func TestGetGPUData(t *testing.T) {
NumNodes: []string{"1"},
ArrayJobID: []string{"0", "10"},
ArrayTaskID: []string{"0", "10"},
Comment: []string{"test_comment"},
}

metadata := GetGPUData(GPUToSlurm)
Expand Down
1 change: 1 addition & 0 deletions shelper/testdata/scontrol_out_all_gpus.txt
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
JobId=30214 ArrayJobId=30185 ArrayTaskId=28 ArrayTaskThrottle=100 JobName=demo_ods
UserId=test_username(1472000020) GroupId=test_username(1472000020) MCS_label=N/A
Priority=1000000 Nice=0 Account=test_account QOS=normal
Comment=test_comment
JobState=RUNNING Reason=None Dependency=(null)
Requeue=1 Restarts=0 BatchFlag=1 Reboot=0 ExitCode=0:0
DerivedExitCode=0:0
Expand Down
2 changes: 2 additions & 0 deletions shelper/testdata/scontrol_out_main_array_job.txt
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
JobId=30214 ArrayJobId=30185 ArrayTaskId=28-59%100 ArrayTaskThrottle=100 JobName=demo_ods
UserId=test_username(1472000020) GroupId=test_username(1472000020) MCS_label=N/A
Priority=1000000 Nice=0 Account=test2_account QOS=normal
Comment=test_comment
JobState=RUNNING Reason=None Dependency=(null)
Requeue=1 Restarts=0 BatchFlag=1 Reboot=0 ExitCode=0:0
DerivedExitCode=0:0
Expand Down Expand Up @@ -34,6 +35,7 @@ JobId=30214 ArrayJobId=30185 ArrayTaskId=28-59%100 ArrayTaskThrottle=100 JobName
JobId=31214 ArrayJobId=31185 ArrayTaskId=128 ArrayTaskThrottle=100 JobName=demo_ods2
UserId=test_username_2(1472000020) GroupId=test_username_2(1472000020) MCS_label=N/A
Priority=1000000 Nice=0 Account=test_account QOS=dev
Comment=test_comment_2
JobState=RUNNING Reason=None Dependency=(null)
Requeue=1 Restarts=0 BatchFlag=1 Reboot=0 ExitCode=0:0
DerivedExitCode=0:0
Expand Down
1 change: 1 addition & 0 deletions shelper/testdata/scontrol_out_multi_node.txt
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
JobId=8401413 JobName=slurm_job_name
UserId=test_username(1472000020) GroupId=test_username(1472000020) MCS_label=N/A
Priority=1002612 Nice=0 Account=test_account QOS=test_qos
Comment=test_comment
JobState=RUNNING Reason=None Dependency=(null)
Requeue=1 Restarts=0 BatchFlag=1 Reboot=0 ExitCode=0:0
DerivedExitCode=0:0
Expand Down
1 change: 1 addition & 0 deletions shelper/testdata/scontrol_out_no_gpus.txt
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
JobId=30214 ArrayJobId=30185 ArrayTaskId=28 ArrayTaskThrottle=100 JobName=demo_ods
UserId=test_username(1472000020) GroupId=test_username(1472000020) MCS_label=N/A
Priority=1000000 Nice=0 Account=test_account QOS=normal
Comment=test_comment
JobState=RUNNING Reason=None Dependency=(null)
Requeue=1 Restarts=0 BatchFlag=1 Reboot=0 ExitCode=0:0
DerivedExitCode=0:0
Expand Down
Loading
Loading