Add ntasks to SlurmExecutor (#381)

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:

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],