add working dir param for sbj

chriskery · Nov 14, 2023 · 6defc76 · 6defc76
1 parent c10127b
commit 6defc76
Show file tree

Hide file tree

Showing 8 changed files with 289 additions and 267 deletions.
diff --git a/apis/kubecluster.org/v1alpha1/zz_generated.deepcopy.go b/apis/kubecluster.org/v1alpha1/zz_generated.deepcopy.go
diff --git a/manifests/samples/kubecluster.org_v1alpha1_slurmbridgejob.yaml b/manifests/samples/kubecluster.org_v1alpha1_slurmbridgejob.yaml
@@ -13,12 +13,16 @@ spec:
   sbatchScript: |
     #!/bin/sh
     #SBATCH --nodes=1
-    #SBATCH --array 1-3
     srun hostname
     hostname
-    for i in {1..5};do echo $i && sleep 5;done
-#  runAsUser: 1003
-#  runAsGroup: 0
+    pwd
+    for i in {1..5};do echo $i && sleep 15;done
+  runAsUser: 0
+  runAsGroup: 0
+  ntasks: 3
+  workingDir: /tmp
+  memPerCpu: 500
+  cpusPerTask: 1
   result:
     volume:
       name: myresult

diff --git a/pkg/slurm-agent/api/slurm.go b/pkg/slurm-agent/api/slurm.go
@@ -102,6 +102,7 @@ func (s *Slurm) SubmitJob(ctx context.Context, req *workload.SubmitJobRequest) (
 			Ntasks:        req.Ntasks,
 			Nodes:         req.Nodes,
 			JobName:       req.JobName,
+			WorkingDir:    req.WorkingDir,
 		})
 		if err != nil {
 			return nil, errors.Wrap(err, "could not submit sbatch script")

diff --git a/pkg/slurm-agent/slurm.go b/pkg/slurm-agent/slurm.go
@@ -160,6 +160,7 @@ type SbatchRequest struct {
 	Ntasks        int64  `protobuf:"varint,11,opt,name=ntasks,proto3" json:"ntasks,omitempty"`
 	Nodes         int64  `protobuf:"varint,12,opt,name=nodes,proto3" json:"nodes,omitempty"`
 	JobName       string `protobuf:"bytes,13,opt,name=job_name,json=jobName,proto3" json:"job_name,omitempty"`
+	WorkingDir    string `protobuf:"bytes,14,opt,name=working_dir,json=workingDir,proto3" json:"working_dir,omitempty"`
 }
 
 // SBatch submits batch job and returns job id if succeeded.
@@ -218,9 +219,12 @@ func getSbatchOpts(req *SbatchRequest) []string {
 	if req.NtasksPerNode > 0 {
 		opts = append(opts, "--ntasks-per-node="+strconv.Itoa(int(req.NtasksPerNode)))
 	}
-	if len(req.JobName) > 0 {
+	if req.JobName != "" {
 		opts = append(opts, "--job-name="+req.JobName)
 	}
+	if req.WorkingDir != "" {
+		opts = append(opts, "--chdir="+req.WorkingDir)
+	}
 	return opts
 }
 

diff --git a/pkg/slurm-bridge-operator/pod.go b/pkg/slurm-bridge-operator/pod.go
@@ -49,6 +49,7 @@ func (r *SlurmBridgeJobReconciler) newPodForSJ(sjb *v1alpha1.SlurmBridgeJob) (*c
 					Resources:       corev1.ResourceRequirements{Requests: resourceList, Limits: resourceList},
 					Command:         []string{sjb.Spec.SbatchScript},
 					SecurityContext: &corev1.SecurityContext{RunAsUser: sjb.Spec.RunAsUser},
+					WorkingDir:      sjb.Spec.WorkingDir,
 				},
 			},
 			RestartPolicy: corev1.RestartPolicyNever,

diff --git a/pkg/slurm-virtual-kubelet/provider.go b/pkg/slurm-virtual-kubelet/provider.go
@@ -69,6 +69,7 @@ func (s *SlurmVirtualKubeletProvider) newSubmitRequestForPod(pod *v1.Pod) *workl
 		submitRequest.RunAsUser = strconv.FormatInt(*pod.Spec.Containers[0].SecurityContext.RunAsUser, 10)
 	}
 	submitRequest.Script = pod.Spec.Containers[0].Command[0]
+	submitRequest.WorkingDir = pod.Spec.Containers[0].WorkingDir
 
 	labels := pod.GetLabels()
 	nodes, ok := labels[common.LabelsResourceRequestNodes]

diff --git a/pkg/workload/workload.pb.go b/pkg/workload/workload.pb.go
diff --git a/pkg/workload/workload.proto b/pkg/workload/workload.proto
@@ -78,6 +78,7 @@ message SubmitJobRequest {
   int64 ntasks = 11;
   int64 nodes = 12;
   string job_name = 13;
+  string working_dir = 14;
 }
 
 message SubmitJobResponse {