diff --git a/nemo_run/core/execution/slurm.py b/nemo_run/core/execution/slurm.py index 35fa9ee3..163d733a 100644 --- a/nemo_run/core/execution/slurm.py +++ b/nemo_run/core/execution/slurm.py @@ -306,6 +306,9 @@ class ResourceRequest: time: str = "00:10:00" nodes: int = 1 ntasks_per_node: int = 1 + #: Total task count for the allocation, mapped to sbatch --ntasks. + #: Mutually exclusive with ntasks_per_node, which is dropped when this is set. + ntasks: Optional[int] = None cpus_per_task: Optional[int] = None cpus_per_gpu: Optional[int] = None gpus_per_node: Optional[int] = None @@ -430,6 +433,11 @@ def __post_init__(self): if self.wait_time_for_group_job < 0: self.wait_time_for_group_job = 0 + assert self.ntasks is None or not self.heterogeneous, ( + "ntasks cannot be combined with heterogeneous=True, " + "size each group via resource_group instead." + ) + def info(self) -> str: return f"{self.__class__.__qualname__} on {self.tunnel.key}" @@ -463,6 +471,10 @@ def srun( for arg in self.SRUN_ARGS if getattr(self, arg.replace("-", "_"), None) is not None } + if _arg_dict.get("ntasks") is not None: + # Same exclusion as sbatch: ntasks-per-node would cap an explicit ntasks. + _arg_dict.pop("ntasks-per-node", None) + _arg_dict["container-mounts"] = ",".join(self.container_mounts) if env_vars: _arg_dict["container-env"] = ",".join(list(env_vars.keys())) @@ -814,6 +826,11 @@ def materialize(self) -> str: k: v for k, v in args.items() if v is not None and k in SlurmExecutor.SBATCH_FLAGS } + # --ntasks-per-node acts as a per-node maximum when --ntasks is also present, so + # leaving its default of 1 in place would cap an explicit --ntasks request. + if parameters.get("ntasks") is not None: + parameters.pop("ntasks_per_node", None) + # rename and reformat parameters if "cpus_per_gpu" in parameters and "gpus_per_task" not in parameters: diff --git a/test/core/execution/test_slurm_templates.py b/test/core/execution/test_slurm_templates.py index 5457563d..8e9ab331 100644 --- a/test/core/execution/test_slurm_templates.py +++ b/test/core/execution/test_slurm_templates.py @@ -17,6 +17,7 @@ import os import re from pathlib import Path +from unittest.mock import MagicMock, PropertyMock, patch import pytest @@ -504,6 +505,44 @@ def test_dummy_batch_request_array( in sbatch_script ) + def test_dummy_batch_request_ntasks( + self, + dummy_slurm_request_with_artifact: tuple[SlurmBatchRequest, str], + ): + dummy_slurm_request, _ = dummy_slurm_request_with_artifact + dummy_slurm_request.executor.ntasks = 8 + + sbatch_script = dummy_slurm_request.materialize() + assert "#SBATCH --ntasks=8" in sbatch_script + assert "--ntasks-per-node" not in sbatch_script + + def test_dummy_batch_request_ntasks_per_node_default( + self, + dummy_slurm_request_with_artifact: tuple[SlurmBatchRequest, str], + ): + dummy_slurm_request, _ = dummy_slurm_request_with_artifact + + sbatch_script = dummy_slurm_request.materialize() + assert "#SBATCH --ntasks-per-node=1" in sbatch_script + assert "#SBATCH --ntasks=" not in sbatch_script + + def test_srun_ntasks_drops_ntasks_per_node(self): + # sbatch and srun must apply the same exclusion, or the two drift apart. + executor = SlurmExecutor(account="account", tunnel=LocalTunnel(job_dir="/tmp"), ntasks=8) + mock_slurm = MagicMock() + with patch.object( + SlurmExecutor, "slurm", new_callable=PropertyMock, return_value=mock_slurm + ): + executor.srun("echo hi", job_name="interactive") + + srun_cmd = mock_slurm.run.call_args.args[0] + assert "--ntasks=8" in srun_cmd + assert "--ntasks-per-node" not in srun_cmd + + def test_ntasks_rejected_for_heterogeneous(self): + with pytest.raises(AssertionError, match="ntasks cannot be combined with heterogeneous"): + SlurmExecutor(account="account", ntasks=8, heterogeneous=True) + def test_dummy_batch_additonal_params( self, dummy_slurm_request_with_artifact: tuple[SlurmBatchRequest, str],