diff --git a/src/madengine/deployment/slurm.py b/src/madengine/deployment/slurm.py index af99c08d..088a3fb2 100644 --- a/src/madengine/deployment/slurm.py +++ b/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"), diff --git a/src/madengine/deployment/templates/slurm/job.sh.j2 b/src/madengine/deployment/templates/slurm/job.sh.j2 index 3b236d9f..4cdddba0 100644 --- a/src/madengine/deployment/templates/slurm/job.sh.j2 +++ b/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" diff --git a/src/madengine/orchestration/run_orchestrator.py b/src/madengine/orchestration/run_orchestrator.py index 296af1ee..148c953c 100644 --- a/src/madengine/orchestration/run_orchestrator.py +++ b/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]") diff --git a/tests/unit/test_slurm_job_template.py b/tests/unit/test_slurm_job_template.py new file mode 100644 index 00000000..6ce19628 --- /dev/null +++ b/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