102-handle-meshsh-script-generation-coming-from-vk

docker/Dockerfile

@@ -16,7 +16,8 @@ ENV TIMEZONE=America/New_York
 # Install system packages
 RUN DEBIAN_FRONTEND=noninteractive \
   dnf install -y epel-release \
-  && dnf install -y slurm-slurmd rpm-build slurm-slurmctld apptainer sshpass \
+  && dnf install -y --allowerasing slurm-slurmd rpm-build slurm-slurmctld apptainer sshpass slirp4netns iproute curl tar xz make gcc ca-certificates \
+  && update-ca-trust \
   && dnf clean all
 
 RUN slurm-setuser -y
@@ -52,4 +53,4 @@ ENV SLURMCONFIGPATH=/root/SlurmConfig.yaml
 
 RUN mkdir -p /cvmfs/grid.cern.ch/etc/grid-security
 
-ENTRYPOINT ["/bin/sh", "-c", "/etc/startup.sh && SHARED_FS=true /sidecar/slurm-sidecar "]
+ENTRYPOINT ["/bin/sh", "-c", "/etc/startup.sh && SHARED_FS=true /sidecar/slurm-sidecar "]

pkg/slurm/prepare.go

@@ -61,6 +61,66 @@ type FlavorResolution struct {
 	SlurmFlags    []string
 }
 
+func extractHeredoc(content, marker string) (string, error) {
+	// Find the start of the heredoc
+	startPattern := fmt.Sprintf("cat <<'%s'", marker)
+	startIdx := strings.Index(content, startPattern)
+	if startIdx == -1 {
+		return "", fmt.Errorf("heredoc start marker not found")
+	}
+
+	// Find the line after the cat command (start of actual content)
+	contentStart := strings.Index(content[startIdx:], "\n")
+	if contentStart == -1 {
+		return "", fmt.Errorf("invalid heredoc format")
+	}
+	contentStart += startIdx + 1
+
+	// Find the end marker
+	endMarker := "\n" + marker
+	endIdx := strings.Index(content[contentStart:], endMarker)
+	if endIdx == -1 {
+		return "", fmt.Errorf("heredoc end marker not found")
+	}
+
+	// Extract the content between start and end markers
+	return content[contentStart : contentStart+endIdx], nil
+}
+
+func removeHeredoc(content, marker string) string {
+	// Find the start of the heredoc
+	startPattern := fmt.Sprintf("cat <<'%s'", marker)
+	startIdx := strings.Index(content, startPattern)
+	if startIdx == -1 {
+		return content // No heredoc found, return as-is
+	}
+
+	// Find the line after the cat command (start of actual content)
+	contentStart := strings.Index(content[startIdx:], "\n")
+	if contentStart == -1 {
+		return content // Invalid heredoc format
+	}
+	contentStart += startIdx + 1
+
+	// Find the end marker
+	endMarker := "\n" + marker
+	endIdx := strings.Index(content[contentStart:], endMarker)
+	if endIdx == -1 {
+		return content // Heredoc end marker not found
+	}
+
+	// Calculate the actual end position (after the end marker)
+	heredocEnd := contentStart + endIdx + len(endMarker)
+
+	// Skip trailing newline if present
+	if heredocEnd < len(content) && content[heredocEnd] == '\n' {
+		heredocEnd++
+	}
+
+	// Remove the heredoc block and return
+	return content[:startIdx] + content[heredocEnd:]
+}
+
 // stringToHex encodes the provided str string into a hex string and removes all trailing redundant zeroes to keep the output more compact
 func stringToHex(str string) string {
 	var buffer bytes.Buffer
@@ -955,7 +1015,38 @@ func produceSLURMScript(
 	}
 
 	if preExecAnnotations, ok := metadata.Annotations["slurm-job.vk.io/pre-exec"]; ok {
-		prefix += "\n" + preExecAnnotations
+		// Check if pre-exec contains a heredoc that creates mesh.sh
+		if strings.Contains(preExecAnnotations, "cat <<'EOFMESH' > $TMPDIR/mesh.sh") {
+			// Extract the heredoc content
+			meshScript, err := extractHeredoc(preExecAnnotations, "EOFMESH")
+			if err == nil && meshScript != "" {
+
+				meshPath := filepath.Join(path, "mesh.sh")
+				err := os.WriteFile(meshPath, []byte(meshScript), 0755)
+				if err != nil {
+					prefix += "\n" + preExecAnnotations
+				} else {
+					// wrote mesh.sh, now add pre-exec without the mesh.sh heredoc
+					preExecWithoutHeredoc := removeHeredoc(preExecAnnotations, "EOFMESH")
+					prefix += "\n" + preExecWithoutHeredoc + "\n" + fmt.Sprintf(" %s", meshPath)
+				}
+
+				err = os.Chmod(path+"/mesh.sh", 0774)
+				if err != nil {
+					log.G(Ctx).Error("Unable to chmod file ", path, "/job.sh")
+					log.G(Ctx).Error(err)
+					return "", err
+				} else {
+					log.G(Ctx).Debug("--- Created with correct permission file ", path, "/job.sh")
+				}
+			} else {
+				// Could not extract heredoc, include as-is
+				prefix += "\n" + preExecAnnotations
+			}
+		} else {
+			// No heredoc pattern, include pre-exec as-is
+			prefix += "\n" + preExecAnnotations
+		}
 	}
 
 	sbatch_macros := "#!" + config.BashPath +