[kjobctl] Move unmask_filename to templates.

cmd/experimental/kjobctl/pkg/builder/slurm_builder.go

@@ -57,39 +57,6 @@ const (
 	slurmEnvsPath          = "/slurm/env"
 	slurmSbatchEnvFilename = "sbatch.env"
 	slurmSlurmEnvFilename  = "slurm.env"
-
-	//# \\ - Do not process any of the replacement symbols.
-	//# %% - The character "%".
-	//# %A - Job array's master job allocation number (for now it is equivalent to SLURM_JOB_ID).
-	//# %a - Job array ID (index) number (SLURM_ARRAY_TASK_ID).
-	//# %j - job id of the running job (SLURM_JOB_ID).
-	//# %N - short hostname (pod name).
-	//# %n - node(pod) identifier relative to current job - index from K8S index job.
-	//# %t - task identifier (rank) relative to current job - It is array id position.
-	//# %u - username (from the client machine).
-	//# %x - job name.
-	unmaskFilenameFunction = `unmask_filename () {
-  replaced="$1"
-
-  if [[ "$replaced" == "\\"* ]]; then
-      replaced="${replaced//\\/}"
-      echo "${replaced}"
-      return 0
-  fi
-
-  replaced=$(echo "$replaced" | sed -E "s/(%)(%A)/\1\n\2/g;:a s/(^|[^\n])%A/\1$SLURM_ARRAY_JOB_ID/;ta;s/\n//g")
-  replaced=$(echo "$replaced" | sed -E "s/(%)(%a)/\1\n\2/g;:a s/(^|[^\n])%a/\1$SLURM_ARRAY_TASK_ID/;ta;s/\n//g")
-  replaced=$(echo "$replaced" | sed -E "s/(%)(%j)/\1\n\2/g;:a s/(^|[^\n])%j/\1$SLURM_JOB_ID/;ta;s/\n//g")
-  replaced=$(echo "$replaced" | sed -E "s/(%)(%N)/\1\n\2/g;:a s/(^|[^\n])%N/\1$HOSTNAME/;ta;s/\n//g")
-  replaced=$(echo "$replaced" | sed -E "s/(%)(%n)/\1\n\2/g;:a s/(^|[^\n])%n/\1$JOB_COMPLETION_INDEX/;ta;s/\n//g")
-  replaced=$(echo "$replaced" | sed -E "s/(%)(%t)/\1\n\2/g;:a s/(^|[^\n])%t/\1$SLURM_ARRAY_TASK_ID/;ta;s/\n//g")
-  replaced=$(echo "$replaced" | sed -E "s/(%)(%u)/\1\n\2/g;:a s/(^|[^\n])%u/\1$USER_ID/;ta;s/\n//g")
-  replaced=$(echo "$replaced" | sed -E "s/(%)(%x)/\1\n\2/g;:a s/(^|[^\n])%x/\1$SBATCH_JOB_NAME/;ta;s/\n//g")
-
-  replaced="${replaced//%%/%}"
-
-  echo "$replaced"
-}`
 )
 
 var (
@@ -570,7 +537,6 @@ type slurmEntrypointScript struct {
 	EnvsPath               string
 	SbatchEnvFilename      string
 	SlurmEnvFilename       string
-	UnmaskFilenameFunction string
 	ChangeDir              string
 	BuildEntrypointCommand string
 }
@@ -580,7 +546,6 @@ func (b *slurmBuilder) buildEntrypointScript() (string, error) {
 		EnvsPath:               slurmEnvsPath,
 		SbatchEnvFilename:      slurmSbatchEnvFilename,
 		SlurmEnvFilename:       slurmSlurmEnvFilename,
-		UnmaskFilenameFunction: unmaskFilenameFunction,
 		ChangeDir:              b.changeDir,
 		BuildEntrypointCommand: b.buildEntrypointCommand(),
 	}

cmd/experimental/kjobctl/pkg/builder/slurm_builder_test.go

@@ -17,12 +17,14 @@ limitations under the License.
 package builder
 
 import (
+	"bytes"
 	"context"
 	"errors"
 	"fmt"
 	"os"
 	"os/exec"
 	"testing"
+	"text/template"
 	"time"
 
 	"github.com/google/go-cmp/cmp"
@@ -84,6 +86,17 @@ func TestUnmaskFilenameFunction(t *testing.T) {
 				return
 			}
 
+			tmpl, err := template.ParseFS(slurmTemplates, "templates/*")
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			var unmaskFilenameFunction bytes.Buffer
+
+			if err := tmpl.ExecuteTemplate(&unmaskFilenameFunction, "slurm_unmask_filename.sh.tmpl", nil); err != nil {
+				t.Fatal(err)
+			}
+
 			script := fmt.Sprintf(`#!/usr/bin/bash
 
 set -o errexit
@@ -94,7 +107,7 @@ set -o pipefail
 
 unmask_filename "%s"
 
-`, unmaskFilenameFunction, tc.input)
+`, unmaskFilenameFunction.String(), tc.input)
 
 			if _, err := file.WriteString(script); err != nil {
 				t.Error(err)

cmd/experimental/kjobctl/pkg/builder/templates/slurm_entrypoint_script.sh.tmpl

@@ -15,7 +15,7 @@ source {{.EnvsPath}}/$JOB_CONTAINER_INDEX/{{.SbatchEnvFilename}}
 
 export $(cat {{.EnvsPath}}/$JOB_CONTAINER_INDEX/{{.SlurmEnvFilename}} | xargs)
 
-{{.UnmaskFilenameFunction}}
+{{template "slurm_unmask_filename.sh.tmpl" .}}
 
 input_file=$(unmask_filename "$SBATCH_INPUT")
 output_file=$(unmask_filename "$SBATCH_OUTPUT")

cmd/experimental/kjobctl/pkg/builder/templates/slurm_unmask_filename.sh.tmpl

@@ -0,0 +1,34 @@
+{{- /*
+\\ - Do not process any of the replacement symbols.
+%% - The character "%".
+%A - Job array's master job allocation number (for now it is equivalent to SLURM_JOB_ID).
+%a - Job array ID (index) number (SLURM_ARRAY_TASK_ID).
+%j - job id of the running job (SLURM_JOB_ID).
+%N - short hostname (pod name).
+%n - node(pod) identifier relative to current job - index from K8S index job.
+%t - task identifier (rank) relative to current job - It is array id position.
+%u - username (from the client machine).
+%x - job name.
+*/ -}}
+unmask_filename () {
+  replaced="$1"
+
+  if [[ "$replaced" == "\\"* ]]; then
+      replaced="${replaced//\\/}"
+      echo "${replaced}"
+      return 0
+  fi
+
+  replaced=$(echo "$replaced" | sed -E "s/(%)(%A)/\1\n\2/g;:a s/(^|[^\n])%A/\1$SLURM_ARRAY_JOB_ID/;ta;s/\n//g")
+  replaced=$(echo "$replaced" | sed -E "s/(%)(%a)/\1\n\2/g;:a s/(^|[^\n])%a/\1$SLURM_ARRAY_TASK_ID/;ta;s/\n//g")
+  replaced=$(echo "$replaced" | sed -E "s/(%)(%j)/\1\n\2/g;:a s/(^|[^\n])%j/\1$SLURM_JOB_ID/;ta;s/\n//g")
+  replaced=$(echo "$replaced" | sed -E "s/(%)(%N)/\1\n\2/g;:a s/(^|[^\n])%N/\1$HOSTNAME/;ta;s/\n//g")
+  replaced=$(echo "$replaced" | sed -E "s/(%)(%n)/\1\n\2/g;:a s/(^|[^\n])%n/\1$JOB_COMPLETION_INDEX/;ta;s/\n//g")
+  replaced=$(echo "$replaced" | sed -E "s/(%)(%t)/\1\n\2/g;:a s/(^|[^\n])%t/\1$SLURM_ARRAY_TASK_ID/;ta;s/\n//g")
+  replaced=$(echo "$replaced" | sed -E "s/(%)(%u)/\1\n\2/g;:a s/(^|[^\n])%u/\1$USER_ID/;ta;s/\n//g")
+  replaced=$(echo "$replaced" | sed -E "s/(%)(%x)/\1\n\2/g;:a s/(^|[^\n])%x/\1$SBATCH_JOB_NAME/;ta;s/\n//g")
+
+  replaced="${replaced//%%/%}"
+
+  echo "$replaced"
+}