From c499188c51f46df658f92bfc8491c1622d35b6db Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Irving=20Mondrag=C3=B3n?=
 <IrvingMg@users.noreply.github.com>
Date: Thu, 3 Oct 2024 09:30:27 +0200
Subject: [PATCH] Add --chdir flag (#3182)

---
 cmd/experimental/kjobctl/pkg/builder/builder.go             | 6 ++++++
 cmd/experimental/kjobctl/pkg/builder/slurm_builder.go       | 2 ++
 .../pkg/builder/templates/slurm_entrypoint_script.sh.tmpl   | 4 +++-
 cmd/experimental/kjobctl/pkg/cmd/create/create.go           | 5 +++++
 cmd/experimental/kjobctl/pkg/cmd/create/create_test.go      | 3 +++
 5 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/cmd/experimental/kjobctl/pkg/builder/builder.go b/cmd/experimental/kjobctl/pkg/builder/builder.go
index efa1b70b0c..121d4536cb 100644
--- 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
index 74af928352..199ed992df 100644
--- 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
index a86fef0428..7c2a3f2cfe 100644
--- 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
index acff6dce73..68948edabd 100644
--- 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
index d40f6d7110..b98bf8668c 100644
--- 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
 `,
 							}).