feat: Support LSF job arrays (#491)

* fix: Support LSF job arrays

* This fix includes the necessary changes in cgroup regex and bjobs handling to account job arrays in LSF

Signed-off-by: Mahendra Paipuri <mahendra.paipuri@gmail.com>

* style: Fix lint issues

Signed-off-by: Mahendra Paipuri <mahendra.paipuri@gmail.com>

---------

Signed-off-by: Mahendra Paipuri <mahendra.paipuri@gmail.com>

Author: mahendrapaipuri

pkg/collector/cgroup.go

@@ -98,9 +98,11 @@ var (
 /*
 	For v1 possibilities are /cpuacct/lsf/cluster_name/job.42.34071.1775289779
 							 /memory/lsf/cluster_name/job.42.34071.1775289779
+							 /memory/lsf/cluster_name/job.42[2].34071.1775289779 for job arrays
 
 	For v2 possibilities are /lsf/cluster1/job.9.30675.1775295256
 							 /lsf/cluster1/job.9.30675.1775295256/leaf
+							 /lsf/cluster1/job.9[1].30675.1775295256/leaf for job arrays
 
 	Seems like there are no cgroups for each MPI task created by LSF. At least I am not
 	able to find proper settings to get them created
@@ -112,8 +114,8 @@ var (
 	/bin/bash /home/usr1/.lsbatch/1775296755.13.shell
 */
 var (
-	lsfCgroupV1PathRegex = regexp.MustCompile(`^.*/lsf/(?:(?P<cluster_name>.*?))/job\.(?P<id>[0-9]+)\.(?:(?P<int_id>[0-9]+))\.(?:(?P<ts>[0-9]+))(?:.*$)`)
-	lsfCgroupV2PathRegex = regexp.MustCompile(`^.*/lsf/(?:(?P<cluster_name>.*?))/job\.(?P<id>[0-9]+)\.(?:(?P<int_id>[0-9]+))\.(?:(?P<ts>[0-9]+))/leaf(?:.*$)`)
+	lsfCgroupV1PathRegex = regexp.MustCompile(`^.*/lsf/(?:(?P<cluster_name>.*?))/job\.(?P<id>[0-9]+(?:\[[0-9]+\])?)\.(?:(?P<int_id>[0-9]+))\.(?:(?P<ts>[0-9]+))(?:.*$)`)
+	lsfCgroupV2PathRegex = regexp.MustCompile(`^.*/lsf/(?:(?P<cluster_name>.*?))/job\.(?P<id>[0-9]+(?:\[[0-9]+\])?)\.(?:(?P<int_id>[0-9]+))\.(?:(?P<ts>[0-9]+))/leaf(?:.*$)`)
 	lsfIgnoreProcsRegex  = regexp.MustCompile("(.*)/.lsbatch/(.*)")
 )
 

pkg/collector/helper.go

@@ -488,7 +488,7 @@ func readProcEnvirons(data any) error {
 		var targetEnvVarsToFetch []string
 
 		for _, targetEnvVar := range d.targetEnvVars {
-			if _, ok := d.targetEnvVarValues[targetEnvVar]; !ok {
+			if val, ok := d.targetEnvVarValues[targetEnvVar]; !ok || val == "" {
 				targetEnvVarsToFetch = append(targetEnvVarsToFetch, targetEnvVar)
 			}
 		}
@@ -510,8 +510,12 @@ func readProcEnvirons(data any) error {
 		// When env var entry found, get all necessary env vars
 		for _, env := range environments {
 			for _, targetEnvVar := range targetEnvVarsToFetch {
-				if strings.Contains(env, targetEnvVar) {
-					d.targetEnvVarValues[targetEnvVar] = strings.Split(env, "=")[1]
+				// Split env var name and value and compare the name with targetEnvVar.
+				// This is because if we have targetEnvVar as CUDA_VISIBLE_DEVICES and
+				// env is CUDA_VISIBLE_DEVICES_SOMEPREFIX, we will match targetEnvVar with
+				// CUDA_VISIBLE_DEVICES_SOMEPREFIX which is not we want.
+				if p := strings.Split(env, "="); len(p) == 2 && p[0] == targetEnvVar {
+					d.targetEnvVarValues[targetEnvVar] = p[1]
 				}
 			}
 		}

pkg/collector/helper_test.go

@@ -255,7 +255,6 @@ func TestReadProcEnvirons(t *testing.T) {
 				targetEnvVars: []string{"CUDA_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES1"},
 			},
 			expected: map[string]string{
-				"CUDA_VISIBLE_DEVICES":  "MIG-GPU-956348bc-d43d-23ed-53d4-857749fa2b67/1/0,MIG-GPU-feba7e40-d724-01ff-b00f-3a439a28a6c7/1/0",
 				"CUDA_VISIBLE_DEVICES1": "MIG-GPU-956348bc-d43d-23ed-53d4-857749fa2b67/1/0,MIG-GPU-feba7e40-d724-01ff-b00f-3a439a28a6c7/1/0",
 			},
 		},
@@ -266,7 +265,6 @@ func TestReadProcEnvirons(t *testing.T) {
 				targetEnvVars: []string{"CUDA_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES1"},
 			},
 			expected: map[string]string{
-				"CUDA_VISIBLE_DEVICES":  "MIG-GPU-956348bc-d43d-23ed-53d4-857749fa2b67/1/0,MIG-GPU-feba7e40-d724-01ff-b00f-3a439a28a6c7/1/0",
 				"CUDA_VISIBLE_DEVICES1": "MIG-GPU-956348bc-d43d-23ed-53d4-857749fa2b67/1/0,MIG-GPU-feba7e40-d724-01ff-b00f-3a439a28a6c7/1/0",
 			},
 		},

pkg/collector/lsf.go

@@ -50,9 +50,11 @@ const (
 	lsfReadProcCtx = "lsf_read_procs"
 )
 
-// Regex to get physical GPU UUID and GPU_I_ID, GPU_C_ID from MIG device IDs.
 var (
+	// Regex to get physical GPU UUID and GPU_I_ID, GPU_C_ID from MIG device IDs.
 	migDeviceIDRegex = regexp.MustCompile("^MIG-(?P<GPU_UUID>.*?)/(?P<GPU_I_ID>[0-9]+?)/(?P<GPU_C_ID>[0-9]+)$")
+	// Regex to extract job index from job ID.
+	jobIndexRegex = regexp.MustCompile(`^(?P<id>[0-9]+)(?:\[(?P<index>[0-9]+)\])?$`)
 )
 
 // Cache interval.
@@ -68,6 +70,70 @@ type LSFJobRecord struct {
 	GPUSlot      string `json:"GPU_ALLOC"`
 }
 
+// UnmarshalJSON unmarshals byte array into LSFJobRecord.
+func (r *LSFJobRecord) UnmarshalJSON(b []byte) error {
+	// Define a temporary type to avoid infinite looping
+	type LSFJobRecordTmp LSFJobRecord
+
+	type tmp struct {
+		LSFJobRecordTmp
+
+		Index string `json:"JOBINDEX"`
+	}
+
+	var s tmp
+
+	err := json.Unmarshal(b, &s)
+	if err != nil {
+		return err
+	}
+
+	*r = LSFJobRecord(s.LSFJobRecordTmp)
+
+	// If Index exists, attach it to job ID. So, if id is 6 and index is 2,
+	// the final id will be "6[2]" as appaears in cgroup paths
+	if s.Index != "" {
+		r.ID = fmt.Sprintf("%s[%s]", r.ID, s.Index)
+	}
+
+	return nil
+}
+
+// MarshalJSON marshals LSFJobRecord into byte array.
+func (r LSFJobRecord) MarshalJSON() ([]byte, error) {
+	// Define a temporary type to avoid infinite looping
+	type LSFJobRecordTmp LSFJobRecord
+
+	// Another tmp struct that contains job index
+	type tmp struct {
+		LSFJobRecordTmp
+
+		Index string `json:"JOBINDEX"`
+	}
+
+	var s tmp
+
+	s.LSFJobRecordTmp = LSFJobRecordTmp(r)
+
+	// If jobID has index inside it, extract it
+	match := jobIndexRegex.FindStringSubmatch(r.ID)
+	// If no matches found, return
+	if len(match) > 0 {
+		// Get index of the job
+		for i, name := range jobIndexRegex.SubexpNames() {
+			if name == "id" {
+				s.ID = strings.TrimSpace(match[i])
+			}
+
+			if name == "index" {
+				s.Index = strings.TrimSpace(match[i])
+			}
+		}
+	}
+
+	return json.Marshal(s)
+}
+
 // LSFJobsList contains list of all job records.
 type LSFJobsList struct {
 	Command string         `json:"COMMAND"`
@@ -379,7 +445,7 @@ func (c *lsfCollector) jobResources(cgroups []cgroup) {
 	// Get job resources from bjobs command
 	cmdOut, err := osexec.ExecuteContext(
 		ctx, c.bjobsExePath,
-		[]string{"-u", "all", "-m", c.hostname, "-json", "-o", "jobid alloc_slot nalloc_slot gpu_alloc"},
+		[]string{"-u", "all", "-m", c.hostname, "-json", "-o", "jobid jobindex alloc_slot nalloc_slot gpu_alloc"},
 		nil,
 	)
 	if err != nil {
@@ -563,7 +629,8 @@ func (c *lsfCollector) jobGPUInstanceResources(uuid string, procs []procfs.Proc,
 	dataPtr := &readProcSecurityCtxData{
 		procs: procs,
 		ignoreProc: func(envs []string) bool {
-			return !slices.Contains(envs, "LSB_JOBID="+uuid)
+			// LSB_JOBID WILL NOT HAVE job index. LSB_BATCH_JID is more realiable
+			return !slices.Contains(envs, "LSB_BATCH_JID="+uuid)
 		},
 		// In the case of AMD GPUs, GPU_DEVICE_ORDINAL and GPU_DEVICE_ORDINAL1 are
 		// exported with "real" GPU minor numbers. But at the same time, bjobs

pkg/collector/lsf_test.go

@@ -145,10 +145,10 @@ func TestLSFJobDevices(t *testing.T) {
 }
 
 func writeLSFEnvironFile(procFS string, jobid string, val string) error {
-	envs := []string{"LSB_JOBID=" + jobid, "CUDA_VISIBLE_DEVICES=" + val}
+	envs := []string{"LSB_BATCH_JID=" + jobid, "CUDA_VISIBLE_DEVICES=" + val}
 
 	return os.WriteFile(
-		procFS+"/"+jobid+"/environ",
+		procFS+"/"+strings.ReplaceAll(strings.ReplaceAll(jobid, "[", ""), "]", "")+"/environ",
 		[]byte(strings.Join(envs, "\000")+"\000"),
 		0o600,
 	)
@@ -443,20 +443,20 @@ func TestLSFJobDevicesCaching(t *testing.T) {
 		require.NoError(t, err)
 	}
 
-	for i := 19; i < 40; i++ {
-		dir := fmt.Sprintf("%s/cpuacct/lsf/cluster1/job.%d.12345.123443", cgroupsPath, i)
+	for i := 1; i <= 20; i++ {
+		dir := fmt.Sprintf("%s/cpuacct/lsf/cluster1/job.19[%d].12345.123443", cgroupsPath, i)
 
 		err = os.MkdirAll(dir, 0o750)
 		require.NoError(t, err)
 
 		err = os.WriteFile(
 			dir+"/cgroup.procs",
-			fmt.Appendf(nil, "%d\n", i),
+			fmt.Appendf(nil, "19%d\n", i),
 			0o600,
 		)
 		require.NoError(t, err)
 
-		procDir := fmt.Sprintf("%s/%d", procFS, i)
+		procDir := fmt.Sprintf("%s/19%d", procFS, i)
 
 		err = os.MkdirAll(procDir, 0o750)
 		require.NoError(t, err)
@@ -539,32 +539,32 @@ func TestLSFJobDevicesCaching(t *testing.T) {
 		NumJobs: 4,
 		Records: []LSFJobRecord{
 			{
-				ID:        "19",
+				ID:        "19[1]",
 				AllocSlot: "testhost-1.example.com:testhost-1.example.com:testhost-1.example.com:testhost-1.example.com",
 				GPUSlot:   "testhost-1.example.com:0,0,0,0",
 			},
 			{
-				ID:        "20",
+				ID:        "19[2]",
 				AllocSlot: "testhost-1.example.com:testhost-1.example.com:testhost-1.example.com:testhost-1.example.com",
 				GPUSlot:   "testhost-1.example.com:2,2,2,2",
 			},
 			{
-				ID:        "21",
+				ID:        "19[3]",
 				AllocSlot: "testhost-1.example.com:testhost-1.example.com:testhost-2.example.com:testhost-2.example.com",
 				GPUSlot:   `testhost-1.example.com:1:1\/0;testhost-2.example.com:0:1\/0`,
 			},
 			{
-				ID:        "22",
+				ID:        "19[4]",
 				AllocSlot: "testhost-1.example.com:testhost-1.example.com",
 				GPUSlot:   "testhost-1.example.com:0,2",
 			},
 			{
-				ID:        "23",
+				ID:        "19[5]",
 				AllocSlot: "testhost-1.example.com:testhost-1.example.com",
 				GPUSlot:   `testhost-1.example.com:3:13\/1`,
 			},
 			{
-				ID:        "24",
+				ID:        "19[6]",
 				AllocSlot: "testhost-1.example.com:testhost-1.example.com",
 				GPUSlot:   `testhost-1.example.com:3:13\/1`,
 			},
@@ -610,10 +610,10 @@ func TestLSFJobDevicesCaching(t *testing.T) {
 
 	// New expected jobs
 	expected = map[string][]ComputeUnit{
-		"0": {{UUID: "19", NumShares: 1}, {UUID: "22", NumShares: 1}},
-		"1": {{UUID: "21", NumShares: 1}},
-		"2": {{UUID: "20", NumShares: 1}, {UUID: "22", NumShares: 1}},
-		"5": {{UUID: "23", NumShares: 1}, {UUID: "24", NumShares: 1}},
+		"0": {{UUID: "19[1]", NumShares: 1}, {UUID: "19[4]", NumShares: 1}},
+		"1": {{UUID: "19[3]", NumShares: 1}},
+		"2": {{UUID: "19[2]", NumShares: 1}, {UUID: "19[4]", NumShares: 1}},
+		"5": {{UUID: "19[5]", NumShares: 1}, {UUID: "19[6]", NumShares: 1}},
 	}
 
 	for _, gpu := range c.gpuSMI.Devices {

pkg/collector/testdata/bjobs

@@ -6,18 +6,21 @@ echo -e '{
   "RECORDS":[
     {
       "JOBID":"1009248",
+      "JOBINDEX":"",
       "ALLOC_SLOT":"testhost-1.example.com:testhost-1.example.com:testhost-1.example.com:testhost-1.example.com:testhost-1.example.com:testhost-1.example.com",
       "NALLOC_SLOT":"6",
       "GPU_ALLOC":"testhost-1.example.com:2:1\/0,3:1\/0,2:1\/0,3:1\/0,2:1\/0,3:1\/0"
     },
     {
       "JOBID":"1009249",
+      "JOBINDEX":"",
       "ALLOC_SLOT":"testhost-1.example.com:testhost-1.example.com:testhost-1.example.com:testhost-1.example.com:testhost-2.example.com:testhost-2.example.com:testhost-2.example.com:testhost-2.example.com",
       "NALLOC_SLOT":"8",
       "GPU_ALLOC":"testhost-1.example.com:0,0,1,1;testhost-2.example.com:0,0,1,1"
     },
     {
       "JOBID":"1009250",
+      "JOBINDEX":"",
       "ALLOC_SLOT":"testhost-1.example.com:testhost-1.example.com:testhost-1.example.com:testhost-1.example.com",
       "NALLOC_SLOT":"4",
       "GPU_ALLOC":"testhost-1.example.com:1,4,1,4"

pkg/collector/testdata/proc.ttar

@@ -9920,7 +9920,7 @@ SymlinkTo: /usr/bin
 # ttar - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
 Path: proc/46281/environ
 Lines: 1
-PATH=/go/bin:/usr/local/go/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/binNULLBYTEHOSTNAME=cd24e11f73a5NULLBYTETERM=xtermNULLBYTEGOLANG_VERSION=1.12.5NULLBYTEGOPATH=/goNULLBYTEHOME=/rootNULLBYTEENABLE_PROFILING=1NULLBYTESLURM_JOB_GPUS=2NULLBYTESLURM_JOB_ID=1009248NULLBYTESLURM_SHARDS_ON_NODE=5NULLBYTELSB_JOBID=1009248NULLBYTECUDA_VISIBLE_DEVICES1=MIG-GPU-956348bc-d43d-23ed-53d4-857749fa2b67/1/0,MIG-GPU-feba7e40-d724-01ff-b00f-3a439a28a6c7/1/0NULLBYTE
+PATH=/go/bin:/usr/local/go/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/binNULLBYTEHOSTNAME=cd24e11f73a5NULLBYTETERM=xtermNULLBYTEGOLANG_VERSION=1.12.5NULLBYTEGOPATH=/goNULLBYTEHOME=/rootNULLBYTEENABLE_PROFILING=1NULLBYTESLURM_JOB_GPUS=2NULLBYTESLURM_JOB_ID=1009248NULLBYTESLURM_SHARDS_ON_NODE=5NULLBYTELSB_JOBID=1009248NULLBYTECUDA_VISIBLE_DEVICES1=MIG-GPU-956348bc-d43d-23ed-53d4-857749fa2b67/1/0,MIG-GPU-feba7e40-d724-01ff-b00f-3a439a28a6c7/1/0NULLBYTECUDA_VISIBLE_DEVICES1_ORIG=NULLBYTELSB_BATCH_JID=1009248NULLBYTE
 Mode: 644
 # ttar - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
 Path: proc/46281/exe