diff --git a/shelper/local.go b/shelper/local.go index cb38983..3107cad 100644 --- a/shelper/local.go +++ b/shelper/local.go @@ -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 @@ -33,6 +34,7 @@ type SlurmMetadataList struct { NumNodes []string ArrayJobID []string ArrayTaskID []string + Comment []string } func parseNewLineToList(input string) []string { diff --git a/shelper/slurm_helpers.go b/shelper/slurm_helpers.go index 11879a7..f87b842 100644 --- a/shelper/slurm_helpers.go +++ b/shelper/slurm_helpers.go @@ -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]) @@ -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 @@ -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), @@ -95,6 +97,7 @@ func GetGPUData(GPUToSlurm map[string]SlurmMetadata) SlurmMetadataList { Partition: setToSlice(allPartition), Account: setToSlice(allAccount), NumNodes: setToSlice(allNumNodes), + Comment: setToSlice(allComment), } } @@ -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], "(") @@ -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])...) } @@ -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 { @@ -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] @@ -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]) diff --git a/shelper/slurm_helpers_test.go b/shelper/slurm_helpers_test.go index 773815b..3b6f195 100644 --- a/shelper/slurm_helpers_test.go +++ b/shelper/slurm_helpers_test.go @@ -29,6 +29,7 @@ func TestGetSlurmDataFromSlurmLineAllGpus(t *testing.T) { Account: "test_account", Partition: "learn", NumNodes: "1", + Comment: "test_comment", }, "1": { User: "test_username", @@ -40,6 +41,7 @@ func TestGetSlurmDataFromSlurmLineAllGpus(t *testing.T) { Account: "test_account", Partition: "learn", NumNodes: "1", + Comment: "test_comment", }, "2": { User: "test_username", @@ -51,6 +53,7 @@ func TestGetSlurmDataFromSlurmLineAllGpus(t *testing.T) { Account: "test_account", Partition: "learn", NumNodes: "1", + Comment: "test_comment", }, "3": { User: "test_username", @@ -62,6 +65,7 @@ func TestGetSlurmDataFromSlurmLineAllGpus(t *testing.T) { Account: "test_account", Partition: "learn", NumNodes: "1", + Comment: "test_comment", }, "4": { User: "test_username", @@ -73,6 +77,7 @@ func TestGetSlurmDataFromSlurmLineAllGpus(t *testing.T) { Account: "test_account", Partition: "learn", NumNodes: "1", + Comment: "test_comment", }, "5": { User: "test_username", @@ -84,6 +89,7 @@ func TestGetSlurmDataFromSlurmLineAllGpus(t *testing.T) { Account: "test_account", Partition: "learn", NumNodes: "1", + Comment: "test_comment", }, "6": { User: "test_username", @@ -95,6 +101,7 @@ func TestGetSlurmDataFromSlurmLineAllGpus(t *testing.T) { Account: "test_account", Partition: "learn", NumNodes: "1", + Comment: "test_comment", }, "7": { User: "test_username", @@ -106,6 +113,7 @@ func TestGetSlurmDataFromSlurmLineAllGpus(t *testing.T) { Account: "test_account", Partition: "learn", NumNodes: "1", + Comment: "test_comment", }, } @@ -133,6 +141,7 @@ func TestGetSlurmDataFromSlurmLineSomeGpus(t *testing.T) { Account: "test_account", Partition: "learn", NumNodes: "1", + Comment: "test_comment", }, "1": { User: "test_username", @@ -144,6 +153,7 @@ func TestGetSlurmDataFromSlurmLineSomeGpus(t *testing.T) { Account: "test_account", Partition: "learn", NumNodes: "1", + Comment: "test_comment", }, "2": { User: "test_username", @@ -155,6 +165,7 @@ func TestGetSlurmDataFromSlurmLineSomeGpus(t *testing.T) { Account: "test_account", Partition: "learn", NumNodes: "1", + Comment: "test_comment", }, "3": { User: "test_username", @@ -166,6 +177,7 @@ func TestGetSlurmDataFromSlurmLineSomeGpus(t *testing.T) { Account: "test_account", Partition: "learn", NumNodes: "1", + Comment: "test_comment", }, } @@ -208,6 +220,7 @@ func TestGetSlurmDataFromSlurmLineUniqueEntries(t *testing.T) { Account: "test2_account", Partition: "learn", NumNodes: "1", + Comment: "test_comment", }, "1": { User: "test_username", @@ -219,6 +232,7 @@ func TestGetSlurmDataFromSlurmLineUniqueEntries(t *testing.T) { Account: "test2_account", Partition: "learn", NumNodes: "1", + Comment: "test_comment", }, "2": { User: "test_username", @@ -230,6 +244,7 @@ func TestGetSlurmDataFromSlurmLineUniqueEntries(t *testing.T) { Account: "test2_account", Partition: "learn", NumNodes: "1", + Comment: "test_comment", }, "3": { User: "test_username", @@ -241,6 +256,7 @@ func TestGetSlurmDataFromSlurmLineUniqueEntries(t *testing.T) { Account: "test2_account", Partition: "learn", NumNodes: "1", + Comment: "test_comment", }, "5": { User: "test_username_2", @@ -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", @@ -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", @@ -274,6 +292,7 @@ func TestGetSlurmDataFromSlurmLineUniqueEntries(t *testing.T) { Account: "test_account", Partition: "test", NumNodes: "3", + Comment: "test_comment_2", }, } AttributeGPU2SlurmMetadata(blocks, "node1751", GPU2Slurm) @@ -300,6 +319,7 @@ func TestGetSlurmDataFromSlurmLineMainArrayJob(t *testing.T) { Account: "test2_account", Partition: "learn", NumNodes: "1", + Comment: "test_comment", }, "1": { User: "test_username", @@ -311,6 +331,7 @@ func TestGetSlurmDataFromSlurmLineMainArrayJob(t *testing.T) { Account: "test2_account", Partition: "learn", NumNodes: "1", + Comment: "test_comment", }, "2": { User: "test_username", @@ -322,6 +343,7 @@ func TestGetSlurmDataFromSlurmLineMainArrayJob(t *testing.T) { Account: "test2_account", Partition: "learn", NumNodes: "1", + Comment: "test_comment", }, "3": { User: "test_username", @@ -333,6 +355,7 @@ func TestGetSlurmDataFromSlurmLineMainArrayJob(t *testing.T) { Account: "test2_account", Partition: "learn", NumNodes: "1", + Comment: "test_comment", }, "5": { User: "test_username_2", @@ -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", @@ -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", @@ -366,6 +391,7 @@ func TestGetSlurmDataFromSlurmLineMainArrayJob(t *testing.T) { Account: "test_account", Partition: "test", NumNodes: "3", + Comment: "test_comment_2", }, } AttributeGPU2SlurmMetadata(blocks, "node1751", GPU2Slurm) @@ -555,6 +581,7 @@ func TestGetGPUData(t *testing.T) { NumNodes: "1", ArrayJobID: "0", ArrayTaskID: "0", + Comment: "test_comment", }, "1": { JobID: "1234", @@ -566,6 +593,7 @@ func TestGetGPUData(t *testing.T) { NumNodes: "1", ArrayJobID: "10", ArrayTaskID: "10", + Comment: "test_comment", }, } expectedMetadata := SlurmMetadataList{ @@ -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) diff --git a/shelper/testdata/scontrol_out_all_gpus.txt b/shelper/testdata/scontrol_out_all_gpus.txt index 3bdf145..f329c31 100644 --- a/shelper/testdata/scontrol_out_all_gpus.txt +++ b/shelper/testdata/scontrol_out_all_gpus.txt @@ -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 diff --git a/shelper/testdata/scontrol_out_main_array_job.txt b/shelper/testdata/scontrol_out_main_array_job.txt index 6544560..8a8bd8a 100644 --- a/shelper/testdata/scontrol_out_main_array_job.txt +++ b/shelper/testdata/scontrol_out_main_array_job.txt @@ -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 @@ -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 diff --git a/shelper/testdata/scontrol_out_multi_node.txt b/shelper/testdata/scontrol_out_multi_node.txt index 86dcd32..5bfaceb 100644 --- a/shelper/testdata/scontrol_out_multi_node.txt +++ b/shelper/testdata/scontrol_out_multi_node.txt @@ -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 diff --git a/shelper/testdata/scontrol_out_no_gpus.txt b/shelper/testdata/scontrol_out_no_gpus.txt index 50876f8..31d36ca 100644 --- a/shelper/testdata/scontrol_out_no_gpus.txt +++ b/shelper/testdata/scontrol_out_no_gpus.txt @@ -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 diff --git a/shelper/testdata/scontrol_out_repeated_entry.txt b/shelper/testdata/scontrol_out_repeated_entry.txt index dd0fc72..bd91343 100644 --- a/shelper/testdata/scontrol_out_repeated_entry.txt +++ b/shelper/testdata/scontrol_out_repeated_entry.txt @@ -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 @@ -34,6 +35,7 @@ JobId=30214 ArrayJobId=30185 ArrayTaskId=28 ArrayTaskThrottle=100 JobName=demo_o 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 diff --git a/shelper/testdata/scontrol_out_some_gpus.txt b/shelper/testdata/scontrol_out_some_gpus.txt index 09864e4..9cf52e1 100644 --- a/shelper/testdata/scontrol_out_some_gpus.txt +++ b/shelper/testdata/scontrol_out_some_gpus.txt @@ -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 diff --git a/shelper/testdata/scontrol_out_unique_entries.txt b/shelper/testdata/scontrol_out_unique_entries.txt index fc29c38..3b8ff85 100644 --- a/shelper/testdata/scontrol_out_unique_entries.txt +++ b/shelper/testdata/scontrol_out_unique_entries.txt @@ -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=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 @@ -34,6 +35,7 @@ JobId=30214 ArrayJobId=30185 ArrayTaskId=28 ArrayTaskThrottle=100 JobName=demo_o 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 diff --git a/slurmprocessor/common.go b/slurmprocessor/common.go index 7f8e5d4..c4011cf 100644 --- a/slurmprocessor/common.go +++ b/slurmprocessor/common.go @@ -40,6 +40,9 @@ const SlurmAccount = "account" // SlurmNNodes is used as the attribute key for slurm number of nodes assigned to the job const SlurmNNodes = "num_nodes" +// SlurmComment is used as the attribute key for slurm job comment +const SlurmComment = "comment" + // gpuUUID is used as the attribute key for GPU UUID const gpuUUID = "uuid" @@ -93,6 +96,7 @@ func (spb *SlurmProcessorBase) AddSlurmMetadataStr(attributes pcommon.Map, slurm SlurmPartition: slurmMetadata.Partition, SlurmAccount: slurmMetadata.Account, SlurmNNodes: slurmMetadata.NumNodes, + SlurmComment: slurmMetadata.Comment, } for name, value := range columnMap { @@ -112,6 +116,7 @@ func (spb *SlurmProcessorBase) AddSlurmMetadataSlice(attributes pcommon.Map, all SlurmPartition: allGPUData.Partition, SlurmAccount: allGPUData.Account, SlurmNNodes: allGPUData.NumNodes, + SlurmComment: allGPUData.Comment, } for name, value := range columnMap {