Merge pull request #110 from sylabs/support-job-array

Support job array

Author: pisarukv

cmd/job-companion/main.go

@@ -95,30 +95,49 @@ func runBatch(c api.WorkloadManagerClient, batch string, cOps *collectOptions) e
 	if err != nil {
 		return err
 	}
-	info := infoResp.Info
+	info := infoResp.Info[0] // response always contains at leas one element
 
 	log.Printf("JobID: %d", jobID)
 
 	ctx, cancelTailLogs := context.WithCancel(context.Background())
-	tailLogsDone := tailLogs(ctx, c, info.StdOut)
+	var tailLogsDone chan struct{}
+	// we are not tailing logs for JobArrays
+	// since there is no a correct solution how we can print multiple parallel running job outputs
+	// without affecting each other
+	if info.ArrayId == "" {
+		tailLogsDone = tailLogs(ctx, c, info.StdOut)
+	}
+
+	stopLogs := func() {
+		cancelTailLogs()
+
+		// if tail logs done is nil that means that job is a JobArray
+		// and we are not tailing logs for such jobs
+		if tailLogsDone == nil {
+			return
+		}
+
+		<-tailLogsDone // need to wail till all logs will be printed, not to ruin formatting
+	}
 
 	for {
 		time.Sleep(1 * time.Second)
 
 		infoResp, err = c.JobInfo(context.Background(), &api.JobInfoRequest{JobId: jobID})
 		if err != nil {
-			cancelTailLogs()
+			// waits till logs read reaches EOF
+			stopLogs()
 			return err
 		}
-		info = infoResp.Info
+		info = infoResp.Info[0]
 
 		state := info.Status
 		if state == api.JobStatus_COMPLETED ||
 			state == api.JobStatus_FAILED ||
 			state == api.JobStatus_CANCELLED {
 
-			cancelTailLogs()
-			<-tailLogsDone // need to wail till all logs will be printed, not to ruin formatting
+			// waits till logs read reaches EOF
+			stopLogs()
 
 			switch state {
 			case api.JobStatus_FAILED:

internal/red-box/api/slurm.go

@@ -75,6 +75,10 @@ func (a *Slurm) JobInfo(ctx context.Context, r *api.JobInfoRequest) (*api.JobInf
 		return nil, errors.Wrap(err, "can't convert slurm info into proto info")
 	}
 
+	if len(pInfo) == 0 {
+		return nil, errors.New("job info slice is empty, probably invalid scontrol output")
+	}
+
 	return &api.JobInfoResponse{Info: pInfo}, nil
 }
 
@@ -223,60 +227,65 @@ func mapSStepsToProtoSteps(ss []*slurm.JobStepInfo) ([]*api.JobStepInfo, error)
 	return pSteps, nil
 }
 
-func mapSInfoToProtoInfo(si *slurm.JobInfo) (*api.JobInfo, error) {
-	var submitTime *timestamp.Timestamp
-	if si.SubmitTime != nil {
-		pt, err := ptypes.TimestampProto(*si.SubmitTime)
-		if err != nil {
-			return nil, errors.Wrap(err, "can't convert submit go time to proto time")
+func mapSInfoToProtoInfo(si []*slurm.JobInfo) ([]*api.JobInfo, error) {
+	pInfs := make([]*api.JobInfo, len(si))
+	for i, inf := range si {
+		var submitTime *timestamp.Timestamp
+		if inf.SubmitTime != nil {
+			pt, err := ptypes.TimestampProto(*inf.SubmitTime)
+			if err != nil {
+				return nil, errors.Wrap(err, "can't convert submit go time to proto time")
+			}
+
+			submitTime = pt
 		}
 
-		submitTime = pt
-	}
+		var startTime *timestamp.Timestamp
+		if inf.StartTime != nil {
+			pt, err := ptypes.TimestampProto(*inf.StartTime)
+			if err != nil {
+				return nil, errors.Wrap(err, "can't convert start go time to proto time")
+			}
 
-	var startTime *timestamp.Timestamp
-	if si.StartTime != nil {
-		pt, err := ptypes.TimestampProto(*si.StartTime)
-		if err != nil {
-			return nil, errors.Wrap(err, "can't convert start go time to proto time")
+			startTime = pt
 		}
 
-		startTime = pt
-	}
-
-	var runTime *duration.Duration
-	if si.RunTime != nil {
-		runTime = ptypes.DurationProto(*si.RunTime)
-	}
+		var runTime *duration.Duration
+		if inf.RunTime != nil {
+			runTime = ptypes.DurationProto(*inf.RunTime)
+		}
 
-	var timeLimit *duration.Duration
-	if si.TimeLimit != nil {
-		timeLimit = ptypes.DurationProto(*si.TimeLimit)
-	}
+		var timeLimit *duration.Duration
+		if inf.TimeLimit != nil {
+			timeLimit = ptypes.DurationProto(*inf.TimeLimit)
+		}
 
-	status, ok := api.JobStatus_value[si.State]
-	if !ok {
-		status = int32(api.JobStatus_UNKNOWN)
-	}
+		status, ok := api.JobStatus_value[inf.State]
+		if !ok {
+			status = int32(api.JobStatus_UNKNOWN)
+		}
 
-	pi := api.JobInfo{
-		Id:         si.ID,
-		UserId:     si.UserID,
-		Name:       si.Name,
-		ExitCode:   si.ExitCode,
-		Status:     api.JobStatus(status),
-		SubmitTime: submitTime,
-		StartTime:  startTime,
-		RunTime:    runTime,
-		TimeLimit:  timeLimit,
-		WorkingDir: si.WorkDir,
-		StdOut:     si.StdOut,
-		StdErr:     si.StdErr,
-		Partition:  si.Partition,
-		NodeList:   si.NodeList,
-		BatchHost:  si.BatchHost,
-		NumNodes:   si.NumNodes,
+		pi := api.JobInfo{
+			Id:         inf.ID,
+			UserId:     inf.UserID,
+			Name:       inf.Name,
+			ExitCode:   inf.ExitCode,
+			Status:     api.JobStatus(status),
+			SubmitTime: submitTime,
+			StartTime:  startTime,
+			RunTime:    runTime,
+			TimeLimit:  timeLimit,
+			WorkingDir: inf.WorkDir,
+			StdOut:     inf.StdOut,
+			StdErr:     inf.StdErr,
+			Partition:  inf.Partition,
+			NodeList:   inf.NodeList,
+			BatchHost:  inf.BatchHost,
+			NumNodes:   inf.NumNodes,
+			ArrayId:    inf.ArrayJobID,
+		}
+		pInfs[i] = &pi
 	}
 
-	return &pi, nil
+	return pInfs, nil
 }

internal/red-box/api/slurm_test.go

@@ -41,9 +41,12 @@ func Test_mapSInfoToProtoInfo(t *testing.T) {
 		NodeList:   "node1",
 		BatchHost:  "host1",
 		NumNodes:   "2",
+		ArrayJobID: "111",
 	}
-	pi, err := mapSInfoToProtoInfo(&testInfo)
+	pinfs, err := mapSInfoToProtoInfo([]*slurm.JobInfo{&testInfo})
 	require.NoError(t, err)
+	require.Len(t, pinfs, 1)
+	pi := pinfs[0]
 
 	require.EqualValues(t, testInfo.ID, pi.Id)
 	require.EqualValues(t, testInfo.UserID, pi.UserId)
@@ -63,6 +66,7 @@ func Test_mapSInfoToProtoInfo(t *testing.T) {
 	require.EqualValues(t, testInfo.NodeList, pi.NodeList)
 	require.EqualValues(t, testInfo.BatchHost, pi.BatchHost)
 	require.EqualValues(t, testInfo.NumNodes, pi.NumNodes)
+	require.EqualValues(t, testInfo.ArrayJobID, pi.ArrayId)
 }
 
 func Test_mapSStepsToProtoSteps(t *testing.T) {

pkg/slurm/slurm.go

@@ -65,10 +65,11 @@ func NewClient() (*Client, error) {
 	return &Client{}, nil
 }
 
-// JobInfo contains information about a single Slurm job.
+// JobInfo contains information about a Slurm job.
 type JobInfo struct {
 	ID         string         `json:"id" slurm:"JobId"`
 	UserID     string         `json:"user_id" slurm:"UserId"`
+	ArrayJobID string         `json:"array_job_id" slurm:"ArrayJobId"`
 	Name       string         `json:"name" slurm:"JobName"`
 	ExitCode   string         `json:"exit_code" slurm:"ExitCode"`
 	State      string         `json:"state" slurm:"JobState"`
@@ -145,7 +146,7 @@ func (*Client) Tail(path string) (io.ReadCloser, error) {
 	return tr, nil
 }
 
-func (*Client) SJobInfo(jobID int64) (*JobInfo, error) {
+func (*Client) SJobInfo(jobID int64) ([]*JobInfo, error) {
 	cmd := exec.Command(scontrolBinaryName, "show", "jobid", strconv.FormatInt(jobID, 10))
 
 	out, err := cmd.Output()
@@ -188,24 +189,32 @@ func (*Client) SJobSteps(jobID int64) ([]*JobStepInfo, error) {
 	return jInfo, nil
 }
 
-func JobInfoFromScontrolResponse(r string) (*JobInfo, error) {
-	rFields := strings.Fields(r)
-	slurmFields := make(map[string]string)
-	for _, f := range rFields {
-		s := strings.Split(f, "=")
-		if len(s) != 2 {
-			// just skipping empty fields
-			continue
+func JobInfoFromScontrolResponse(r string) ([]*JobInfo, error) {
+	r = strings.TrimSpace(r)
+	rawInfos := strings.Split(r, "\n\n")
+
+	infos := make([]*JobInfo, len(rawInfos))
+
+	for i, raw := range rawInfos {
+		rFields := strings.Fields(raw)
+		slurmFields := make(map[string]string)
+		for _, f := range rFields {
+			s := strings.Split(f, "=")
+			if len(s) != 2 {
+				// just skipping empty fields
+				continue
+			}
+			slurmFields[s[0]] = s[1]
 		}
-		slurmFields[s[0]] = s[1]
-	}
 
-	ji := &JobInfo{}
-	if err := ji.fillFromSlurmFields(slurmFields); err != nil {
-		return nil, err
+		ji := &JobInfo{}
+		if err := ji.fillFromSlurmFields(slurmFields); err != nil {
+			return nil, err
+		}
+		infos[i] = ji
 	}
 
-	return ji, nil
+	return infos, nil
 }
 
 // ParseSacctResponse is a helper that parses sacct output and