fix(slurm): make the generated job script portable across clusters

src/madengine/deployment/slurm.py

@@ -13,6 +13,7 @@
 
 import os
 import shlex
+import shutil
 import subprocess
 import time
 from pathlib import Path
@@ -77,6 +78,8 @@ def __init__(self, config: DeploymentConfig):
         self.time_limit = self.slurm_config.get("time", "24:00:00")
         self.output_dir = Path(self.slurm_config.get("output_dir", "./slurm_results"))
         self.reservation = self.slurm_config.get("reservation", None)
+        # Some clusters expose no GPU GRES, so sbatch rejects --gpus-per-node.
+        self.skip_gpus_directive = self.slurm_config.get("skip_gpus_directive", False)
 
         # Setup Jinja2 template engine
         template_dir = Path(__file__).parent / "templates" / "slurm"
@@ -226,6 +229,22 @@ def validate(self) -> bool:
         self.console.print("[green]✓ SLURM environment validated[/green]")
         return True
 
+    @staticmethod
+    def _submission_bin_dir() -> Optional[str]:
+        """
+        Directory the madengine console script was resolved from at submission time.
+
+        A batch job is not guaranteed to inherit the submitter's PATH: a site can
+        default sbatch to --export=NONE, and `module load` can rewrite PATH before
+        the job body runs. Passing the directory into the job script lets it put
+        the same madengine back on PATH instead of relying on inheritance.
+
+        Returns:
+            Optional[str]: absolute directory, or None if madengine is not on PATH
+        """
+        cli_path = shutil.which("madengine")
+        return str(Path(cli_path).resolve().parent) if cli_path else None
+
     def _validate_cli_availability(self) -> bool:
         """
         Validate madengine is available before job submission.
@@ -241,7 +260,9 @@ def _validate_cli_availability(self) -> bool:
                 ["madengine", "--version"],
                 capture_output=True,
                 text=True,
-                timeout=5,
+                # A cold import off shared/NFS storage can take far longer than a
+                # local one, so this only guards against a hung interpreter.
+                timeout=600,
                 check=False
             )
             if result.returncode == 0:
@@ -456,7 +477,10 @@ def _prepare_slurm_multi_script(self, model_info: Dict, docker_image_name: str =
             f"#SBATCH --partition={self.partition}",
             f"#SBATCH --nodes={self.nodes}",
             f"#SBATCH --ntasks={self.nodes}",
-            f"#SBATCH --gpus-per-node={self.gpus_per_node}",
+        ]
+        if not self.skip_gpus_directive:
+            script_lines.append(f"#SBATCH --gpus-per-node={self.gpus_per_node}")
+        script_lines += [
             f"#SBATCH --time={self.time_limit}",
         ]
         # Honour user-configured exclusivity (defaults to True to match the standard SLURM template).
@@ -669,6 +693,7 @@ def debug(self, msg):
             "partition": self.partition,
             "nodes": self.nodes,
             "gpus_per_node": resolved_gpus_per_node,  # Use resolved GPU count
+            "skip_gpus_directive": self.skip_gpus_directive,
             "time_limit": self.time_limit,
             "output_dir": str(self.output_dir),
             "master_port": master_port,
@@ -682,6 +707,7 @@ def debug(self, msg):
             "qos": self.slurm_config.get("qos"),
             "account": self.slurm_config.get("account"),
             "modules": self.slurm_config.get("modules", []),
+            "submission_bin_dir": self._submission_bin_dir(),
             "env_vars": self.config.additional_context.get("env_vars", {}),
             "shared_workspace": self.slurm_config.get("shared_workspace"),
             "shared_data": self.config.additional_context.get("shared_data"),

src/madengine/deployment/templates/slurm/job.sh.j2

@@ -6,7 +6,8 @@
 #SBATCH --nodes={{ nodes }}
 #SBATCH --ntasks={{ nodes }}
 #SBATCH --ntasks-per-node=1
-#SBATCH --gpus-per-node={{ gpus_per_node }}
+{% if not skip_gpus_directive %}#SBATCH --gpus-per-node={{ gpus_per_node }}
+{% endif %}
 #SBATCH --time={{ time_limit }}
 {% if reservation %}
 #SBATCH --reservation={{ reservation }}
@@ -41,6 +42,19 @@
 module load {{ module }}
 {% endfor %}
 
+# =============================================================================
+# PATH
+# =============================================================================
+# A batch job does not reliably inherit the submitter's PATH: a site can default
+# sbatch to --export=NONE, and the module loads above can rewrite it. Without
+# this the madengine console script is missing on the compute node even though
+# the pre-submission check found it. Re-add the per-user bin directory and the
+# directory madengine itself was resolved from at submission time.
+export PATH="$HOME/.local/bin:$PATH"
+{% if submission_bin_dir %}
+export PATH="{{ submission_bin_dir }}:$PATH"
+{% endif %}
+
 # =============================================================================
 # Environment Setup (Standard ML Environment Variables)
 # =============================================================================
@@ -141,7 +155,15 @@ echo "Submission directory: {{ manifest_file | dirname }}"
 # Single-node: Prefer shared storage (submission dir), with local fallback if needed
 # Check if submission directory is on shared filesystem
 SUBMIT_DIR={{ manifest_file | dirname }}
-if df -T "$SUBMIT_DIR" 2>/dev/null | grep -qE '\bnfs\b|\blustre\b|\bgpfs\b|\bceph\b'; then
+# Ask df for the filesystem type alone. Matching the whole df -T line reads the mount
+# point too, so a local disk under a path such as /mnt/nfs-scratch answered yes and the
+# job then trusted node-local storage to be shared. --output=fstype is GNU coreutils
+# 8.21 and up; the awk form is the fallback for anything older.
+SUBMIT_FSTYPE=$(df --output=fstype "$SUBMIT_DIR" 2>/dev/null | tail -n 1)
+if [ -z "$SUBMIT_FSTYPE" ]; then
+    SUBMIT_FSTYPE=$(df -T "$SUBMIT_DIR" 2>/dev/null | awk 'NR > 1 { print $2; exit }')
+fi
+if printf '%s' "$SUBMIT_FSTYPE" | grep -qE '^(nfs[0-9]*|lustre|gpfs|ceph|beegfs|panfs)$'; then
     # Submission directory is on shared storage - use it directly (best practice)
     WORKSPACE=$SUBMIT_DIR
     WORKSPACE_TYPE="shared-nfs"

src/madengine/orchestration/run_orchestrator.py

@@ -764,14 +764,17 @@ def _show_node_info(self):
         self.console.sh("echo 'MAD Run Models'")
 
         host_os = self.context.ctx.get("host_os", "")
+        # This is purely informational, but a package manager can block forever on
+        # an interactive prompt (e.g. yum asking to import a repo GPG key) with no
+        # tty to answer it, so every query is capped.
         if "HOST_UBUNTU" in host_os:
-            print(self.console.sh("apt show rocm-libs -a", canFail=True))
+            print(self.console.sh("timeout 10 apt show rocm-libs -a", canFail=True))
         elif "HOST_CENTOS" in host_os:
-            print(self.console.sh("yum info rocm-libs", canFail=True))
+            print(self.console.sh("timeout 10 yum info rocm-libs", canFail=True))
         elif "HOST_SLES" in host_os:
-            print(self.console.sh("zypper info rocm-libs", canFail=True))
+            print(self.console.sh("timeout 10 zypper info rocm-libs", canFail=True))
         elif "HOST_AZURE" in host_os:
-            print(self.console.sh("tdnf info rocm-libs", canFail=True))
+            print(self.console.sh("timeout 10 tdnf info rocm-libs", canFail=True))
         else:
             self.rich_console.print("[yellow]Warning: Unable to detect host OS[/yellow]")
 

tests/unit/test_slurm_job_template.py

@@ -0,0 +1,201 @@
+#!/usr/bin/env python3
+"""
+Unit tests for the generated SLURM job script (`job.sh.j2`).
+
+Locks in the portability contract points that clusters keep re-discovering
+downstream (see ROCm/rocm-systems#9055, which patched madengine's source
+rather than filing them):
+
+1. The job script puts madengine back on PATH itself instead of assuming the
+   batch environment inherited the submitter's PATH.
+2. The shared-filesystem probe recognizes `nfs4`, which is what `df -T`
+   reports on most modern NFS mounts.
+3. `slurm.skip_gpus_directive` removes `#SBATCH --gpus-per-node`, which a
+   cluster advertising no GPU GRES rejects outright.
+
+Copyright (c) Advanced Micro Devices, Inc. All rights reserved.
+"""
+
+import json
+import re
+from pathlib import Path
+from unittest.mock import patch
+
+import pytest
+
+from madengine.deployment.base import DeploymentConfig
+from madengine.deployment.slurm import SlurmDeployment
+
+
+MODEL_ENTRY = {
+    "name": "dummy_torchrun_multinode",
+    "url": "",
+    "dockerfile": "docker/dummy",
+    "scripts": "scripts/dummy/run.sh",
+    "n_gpus": "8",
+    "owner": "mad.support@amd.com",
+    "training_precision": "",
+    "tags": ["pyt", "training"],
+    "timeout": -1,
+    "args": "",
+}
+
+
+def _build_deployment(
+    tmp_path: Path,
+    slurm_overrides: dict = None,
+    distributed_overrides: dict = None,
+) -> SlurmDeployment:
+    """SlurmDeployment over a minimal torchrun manifest, output_dir under tmp_path."""
+    manifest = {
+        "built_images": {"dummy-image": {"docker_image": "dummy:latest"}},
+        "built_models": {"dummy-image": MODEL_ENTRY},
+        "context": {
+            "docker_env_vars": {},
+            "docker_mounts": {},
+            "docker_build_arg": {},
+            "gpu_vendor": "AMD",
+            "guest_os": "UBUNTU",
+            "docker_gpus": "all",
+        },
+    }
+    manifest_path = tmp_path / "build_manifest.json"
+    manifest_path.write_text(json.dumps(manifest))
+
+    slurm_config = {
+        "partition": "test-partition",
+        "nodes": 2,
+        "gpus_per_node": 8,
+        "time": "01:00:00",
+        "output_dir": str(tmp_path / "slurm_output"),
+        "exclusive": True,
+    }
+    slurm_config.update(slurm_overrides or {})
+
+    distributed_config = {
+        "launcher": "torchrun",
+        "nnodes": 2,
+        "nproc_per_node": 8,
+        "backend": "nccl",
+        "port": 29500,
+    }
+    distributed_config.update(distributed_overrides or {})
+
+    cfg = DeploymentConfig(
+        target="slurm",
+        manifest_file=str(manifest_path),
+        additional_context={
+            "deploy": "slurm",
+            "gpu_vendor": "AMD",
+            "guest_os": "UBUNTU",
+            "slurm": slurm_config,
+            "distributed": distributed_config,
+        },
+    )
+    return SlurmDeployment(cfg)
+
+
+def _render(deployment: SlurmDeployment) -> str:
+    """Render job.sh.j2 exactly as prepare() does, without submitting anything."""
+    context = deployment._prepare_template_context(MODEL_ENTRY)
+    return deployment.jinja_env.get_template("job.sh.j2").render(**context)
+
+
+# ---------------------------------------------------------------------------
+# 1. PATH is re-established inside the job
+
+class TestJobScriptPath:
+    """The job script must not depend on the submitter's PATH being inherited."""
+
+    def test_user_bin_dir_is_prepended(self, tmp_path):
+        script = _render(_build_deployment(tmp_path))
+        assert 'export PATH="$HOME/.local/bin:$PATH"' in script
+
+    def test_submission_bin_dir_is_prepended(self, tmp_path):
+        with patch("madengine.deployment.slurm.shutil.which", return_value="/opt/venv/bin/madengine"):
+            script = _render(_build_deployment(tmp_path))
+        assert 'export PATH="/opt/venv/bin:$PATH"' in script
+
+    def test_no_empty_export_when_cli_not_on_path(self, tmp_path):
+        """madengine missing at submission time must not render an empty PATH entry."""
+        with patch("madengine.deployment.slurm.shutil.which", return_value=None):
+            script = _render(_build_deployment(tmp_path))
+        assert 'export PATH=":$PATH"' not in script
+        assert 'export PATH="$HOME/.local/bin:$PATH"' in script
+
+    def test_path_is_set_before_madengine_is_looked_up(self, tmp_path):
+        """The export is useless if it lands after `command -v madengine`."""
+        with patch("madengine.deployment.slurm.shutil.which", return_value="/opt/venv/bin/madengine"):
+            script = _render(_build_deployment(tmp_path))
+        assert script.index('export PATH="/opt/venv/bin:$PATH"') < script.index("command -v madengine")
+
+
+# ---------------------------------------------------------------------------
+# 2. Shared-filesystem probe
+
+class TestSharedFilesystemProbe:
+    """`df -T` reports nfs4 on modern mounts; the probe must not miss it.
+
+    And it must read the filesystem type, nothing else: the mount point travels on the
+    same `df -T` line, so a local disk under a path such as /mnt/nfs-scratch used to answer
+    yes and the job then trusted node-local storage to be visible from every node.
+    """
+
+    @staticmethod
+    def _probe_pattern(script: str) -> str:
+        match = re.search(r"SUBMIT_FSTYPE\"?\s*\|\s*grep -qE '([^']+)'", script)
+        assert match, "shared-filesystem probe not found in rendered script"
+        return match.group(1)
+
+    @pytest.mark.parametrize("fstype,expected", [
+        ("nfs", True),
+        ("nfs3", True),
+        ("nfs4", True),
+        ("lustre", True),
+        ("gpfs", True),
+        ("ceph", True),
+        ("beegfs", True),
+        ("panfs", True),
+        ("ext4", False),
+        ("xfs", False),
+        ("overlay", False),
+        ("tmpfs", False),
+    ])
+    def test_probe_matches_shared_filesystems(self, tmp_path, fstype, expected):
+        # The probe only exists on the single-node branch of the template.
+        deployment = _build_deployment(tmp_path, {"nodes": 1}, {"nnodes": 1})
+        pattern = self._probe_pattern(_render(deployment))
+        assert bool(re.search(pattern, fstype)) is expected
+
+    def test_the_probe_reads_the_fstype_column_only(self, tmp_path):
+        script = _render(_build_deployment(tmp_path, {"nodes": 1}, {"nnodes": 1}))
+        assert 'df --output=fstype "$SUBMIT_DIR"' in script
+        assert 'df -T "$SUBMIT_DIR" 2>/dev/null | grep' not in script
+
+    def test_a_mount_point_that_says_nfs_does_not_make_a_disk_shared(self, tmp_path):
+        """/mnt/nfs-scratch on ext4 is local, whatever its name suggests."""
+        script = _render(_build_deployment(tmp_path, {"nodes": 1}, {"nnodes": 1}))
+        pattern = self._probe_pattern(script)
+        df_line = "/dev/nvme0n1p2 ext4 104857600 50106368 54751232 48% /mnt/nfs-scratch"
+        assert re.search(pattern, df_line) is None
+        assert re.search(pattern, "ext4") is None
+
+    def test_there_is_a_fallback_for_df_without_output(self, tmp_path):
+        """--output is coreutils 8.21; older df still has to be read correctly."""
+        script = _render(_build_deployment(tmp_path, {"nodes": 1}, {"nnodes": 1}))
+        assert "awk 'NR > 1 { print $2; exit }'" in script
+
+
+# ---------------------------------------------------------------------------
+# 3. GPU GRES directive opt-out
+
+class TestGpusPerNodeDirective:
+    """A cluster with GresTypes=(null) rejects any job carrying --gpus-per-node."""
+
+    def test_directive_present_by_default(self, tmp_path):
+        script = _render(_build_deployment(tmp_path))
+        assert "#SBATCH --gpus-per-node=8" in script
+
+    def test_directive_omitted_when_opted_out(self, tmp_path):
+        script = _render(_build_deployment(tmp_path, {"skip_gpus_directive": True}))
+        assert "--gpus-per-node" not in script