diff --git a/docker/Dockerfile b/docker/Dockerfile index e377f58..83f27ba 100644 --- a/docker/Dockerfile +++ b/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 "] \ No newline at end of file diff --git a/pkg/slurm/prepare.go b/pkg/slurm/prepare.go index 953cee4..0ed6638 100644 --- a/pkg/slurm/prepare.go +++ b/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 +