Add --chdir flag (kubernetes-sigs#3182)

KunWuLuan · Oct 3, 2024 · c499188 · c499188
1 parent 0a18cb6
commit c499188
Show file tree

Hide file tree

Showing 5 changed files with 19 additions and 1 deletion.
diff --git a/cmd/experimental/kjobctl/pkg/builder/builder.go b/cmd/experimental/kjobctl/pkg/builder/builder.go
@@ -113,6 +113,7 @@ type Builder struct {
 	ignoreUnknown            bool
 	skipLocalQueueValidation bool
 	skipPriorityValidation   bool
+	changeDir                string
 
 	profile       *v1alpha1.ApplicationProfile
 	mode          *v1alpha1.SupportedMode
@@ -270,6 +271,11 @@ func (b *Builder) WithIgnoreUnknown(ignoreUnknown bool) *Builder {
 	return b
 }
 
+func (b *Builder) WithChangeDir(chdir string) *Builder {
+	b.changeDir = chdir
+	return b
+}
+
 func (b *Builder) WithSkipLocalQueueValidation(skip bool) *Builder {
 	b.skipLocalQueueValidation = skip
 	return b

diff --git a/cmd/experimental/kjobctl/pkg/builder/slurm_builder.go b/cmd/experimental/kjobctl/pkg/builder/slurm_builder.go
@@ -560,6 +560,7 @@ type slurmEntrypointScript struct {
 	SbatchEnvFilename      string
 	SlurmEnvFilename       string
 	UnmaskFilenameFunction string
+	ChangeDir              string
 	BuildEntrypointCommand string
 }
 
@@ -569,6 +570,7 @@ func (b *slurmBuilder) buildEntrypointScript() (string, error) {
 		SbatchEnvFilename:      slurmSbatchEnvFilename,
 		SlurmEnvFilename:       slurmSlurmEnvFilename,
 		UnmaskFilenameFunction: unmaskFilenameFunction,
+		ChangeDir:              b.changeDir,
 		BuildEntrypointCommand: b.buildEntrypointCommand(),
 	}
 

diff --git a/cmd/experimental/kjobctl/pkg/builder/templates/slurm_entrypoint_script.sh.tmpl b/cmd/experimental/kjobctl/pkg/builder/templates/slurm_entrypoint_script.sh.tmpl
@@ -20,5 +20,7 @@ export $(cat {{.EnvsPath}}/$JOB_CONTAINER_INDEX/{{.SlurmEnvFilename}} | xargs)
 input_file=$(unmask_filename "$SBATCH_INPUT")
 output_file=$(unmask_filename "$SBATCH_OUTPUT")
 error_path=$(unmask_filename "$SBATCH_ERROR")
-
+{{if .ChangeDir }}
+cd {{.ChangeDir}}
+{{end}}
 {{.BuildEntrypointCommand}}
diff --git a/cmd/experimental/kjobctl/pkg/cmd/create/create.go b/cmd/experimental/kjobctl/pkg/cmd/create/create.go
@@ -59,6 +59,7 @@ const (
 	initImageFlagName                = "init-image"
 	skipLocalQueueValidationFlagName = "skip-localqueue-validation"
 	skipPriorityValidationFlagName   = "skip-priority-validation"
+	changeDirFlagName                = "chdir"
 
 	commandFlagName     = string(v1alpha1.CmdFlag)
 	parallelismFlagName = string(v1alpha1.ParallelismFlag)
@@ -162,6 +163,7 @@ type CreateOptions struct {
 	InitImage            string
 	PodRunningTimeout    time.Duration
 	RemoveInteractivePod bool
+	ChangeDir            string
 
 	SlurmFlagSet *pflag.FlagSet
 
@@ -383,6 +385,8 @@ The minimum index value is 0. The maximum index value is 2147483647.`)
 				"Local queue name.")
 			o.SlurmFlagSet.StringVar(&o.Priority, priorityFlagName, "",
 				"Apply priority for the entire workload.")
+			o.SlurmFlagSet.StringVarP(&o.ChangeDir, changeDirFlagName, "D", "",
+				"Change directory before executing the script.")
 		},
 	},
 }
@@ -623,6 +627,7 @@ func (o *CreateOptions) Run(ctx context.Context, clientGetter util.ClientGetter,
 		WithIgnoreUnknown(o.IgnoreUnknown).
 		WithSkipLocalQueueValidation(o.SkipLocalQueueValidation).
 		WithSkipPriorityValidation(o.SkipPriorityValidation).
+		WithChangeDir(o.ChangeDir).
 		Do(ctx)
 	if err != nil {
 		return err

diff --git a/cmd/experimental/kjobctl/pkg/cmd/create/create_test.go b/cmd/experimental/kjobctl/pkg/cmd/create/create_test.go
@@ -913,6 +913,7 @@ error_path=$(unmask_filename "$SBATCH_ERROR")
 					"--input", "/slurm/input.txt",
 					"--job-name", "job-name",
 					"--partition", "lq1",
+					"--chdir", "/mydir",
 					tc.tempFile,
 				}
 			},
@@ -1118,6 +1119,8 @@ input_file=$(unmask_filename "$SBATCH_INPUT")
 output_file=$(unmask_filename "$SBATCH_OUTPUT")
 error_path=$(unmask_filename "$SBATCH_ERROR")
 
+cd /mydir
+
 /slurm/scripts/script <$input_file 1>$output_file 2>$error_file
 `,
 							}).