fix: gpu tasks are unset if <= 0 (#347)

previously, `--ntasks-per-gpu` was set, despite being < 1. This is a fix
for #342

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- Bug Fixes
- Simplified task flag behavior: task-related SBATCH flags (--ntasks /
--ntasks-per-gpu) are no longer conditionally injected in the previous
path, producing consistent resource sizing for GPU and non‑GPU jobs.
- Unset or falsy task values now reliably default to 1, preventing empty
or invalid task flags.
- CPU-related settings are now included in the generated SBATCH command.

- Tests
- Added tests validating presence or omission of --ntasks and
--ntasks-per-gpu across GPU/non‑GPU and unset scenarios.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>

Author: cmeesters

snakemake_executor_plugin_slurm/__init__.py

@@ -328,13 +328,6 @@ def run_job(self, job: JobExecutorInterface):
                 "- submitting without. This might or might not work on your cluster."
             )
 
-            # fixes #40 - set ntasks regardless of mpi, because
-            # SLURM v22.05 introduced the requirement for all jobs
-            gpu_job = job.resources.get("gpu") or "gpu" in job.resources.get("gres", "")
-            if gpu_job:
-                call += f" --ntasks-per-gpu={job.resources.get('tasks', 1)}"
-            else:
-                call += f" --ntasks={job.resources.get('tasks', 1)}"
         # MPI job
         if job.resources.get("mpi", False):
             if not job.resources.get("tasks_per_node") and not job.resources.get(

snakemake_executor_plugin_slurm/submit_string.py

@@ -55,15 +55,18 @@ def get_submit_command(job, params):
         # fixes #316 - allow unsetting of tasks per gpu
         # apparently, python's internal process manangement interfers with SLURM
         # e.g. for pytorch
-        ntasks_per_gpu = job.resources.get(
-            "tasks_per_gpu", job.resources.get("tasks", 1)
-        )
+        ntasks_per_gpu = job.resources.get("tasks_per_gpu")
+        if ntasks_per_gpu is None:
+            ntasks_per_gpu = job.resources.get("tasks")
+        if ntasks_per_gpu is None:
+            ntasks_per_gpu = 1
+
         if ntasks_per_gpu >= 1:
             call += f" --ntasks-per-gpu={ntasks_per_gpu}"
     else:
         # fixes #40 - set ntasks regardless of mpi, because
         # SLURM v22.05 will require it for all jobs
-        call += f" --ntasks={job.resources.get('tasks', 1)}"
+        call += f" --ntasks={job.resources.get('tasks') or 1}"
 
     # we need to set cpus-per-task OR cpus-per-gpu, the function
     # will return a string with the corresponding value

tests/tests.py

@@ -520,6 +520,124 @@ def test_empty_qos(self, mock_job):
             # Assert the qoes is included (even if empty)
             assert "--qos=''" in get_submit_command(job, params)
 
+    def test_taks(self, mock_job):
+        """Test that tasks are correctly included in the sbatch command."""
+        # Create a job with tasks
+        job = mock_job(tasks=4)
+        params = {
+            "run_uuid": "test_run",
+            "slurm_logfile": "test_logfile",
+            "comment_str": "test_comment",
+            "account": None,
+            "partition": None,
+            "workdir": ".",
+            "tasks": 4,
+        }
+
+        # Patch subprocess.Popen to capture the sbatch command
+        with patch("subprocess.Popen") as mock_popen:
+            # Configure the mock to return successful submission
+            process_mock = MagicMock()
+            process_mock.communicate.return_value = ("123", "")
+            process_mock.returncode = 0
+            mock_popen.return_value = process_mock
+
+        assert "--ntasks=4" in get_submit_command(job, params)
+
+    def test_gpu_tasks(self, mock_job):
+        """Test that GPU tasks are correctly included in the sbatch command."""
+        # Create a job with GPU tasks
+        job = mock_job(gpu=1, tasks_per_gpu=2)
+        params = {
+            "run_uuid": "test_run",
+            "slurm_logfile": "test_logfile",
+            "comment_str": "test_comment",
+            "account": None,
+            "partition": None,
+            "workdir": ".",
+            "tasks_per_gpu": 2,
+        }
+
+        # Patch subprocess.Popen to capture the sbatch command
+        with patch("subprocess.Popen") as mock_popen:
+            # Configure the mock to return successful submission
+            process_mock = MagicMock()
+            process_mock.communicate.return_value = ("123", "")
+            process_mock.returncode = 0
+            mock_popen.return_value = process_mock
+
+        assert "--ntasks-per-gpu=2" in get_submit_command(job, params)
+
+    def test_no_gpu_task(self, mock_job):
+        """Test that no GPU tasks are included when not specified."""
+        # Create a job without GPU tasks
+        job = mock_job(gpu=1, tasks_per_gpu=-1)
+        params = {
+            "run_uuid": "test_run",
+            "slurm_logfile": "test_logfile",
+            "comment_str": "test_comment",
+            "account": None,
+            "partition": None,
+            "workdir": ".",
+            "tasks_per_gpu": -1,
+        }
+
+        # Patch subprocess.Popen to capture the sbatch command
+        with patch("subprocess.Popen") as mock_popen:
+            # Configure the mock to return successful submission
+            process_mock = MagicMock()
+            process_mock.communicate.return_value = ("123", "")
+            process_mock.returncode = 0
+            mock_popen.return_value = process_mock
+
+        assert "--ntasks-per-gpu" not in get_submit_command(job, params)
+
+    def test_task_set_for_unset_tasks(self, mock_job):
+        """Test that tasks are set to 1 when unset."""
+        # Create a job without tasks
+        job = mock_job(tasks=None)
+        params = {
+            "run_uuid": "test_run",
+            "slurm_logfile": "test_logfile",
+            "comment_str": "test_comment",
+            "account": None,
+            "partition": None,
+            "workdir": ".",
+        }
+
+        # Patch subprocess.Popen to capture the sbatch command
+        with patch("subprocess.Popen") as mock_popen:
+            # Configure the mock to return successful submission
+            process_mock = MagicMock()
+            process_mock.communicate.return_value = ("123", "")
+            process_mock.returncode = 0
+            mock_popen.return_value = process_mock
+
+        assert "--ntasks=1" in get_submit_command(job, params)
+
+    def test_gpu_tasks_set_for_unset_tasks(self, mock_job):
+        """Test that GPU tasks are set to 1 when unset."""
+        # Create a job without GPU tasks
+        job = mock_job(gpu=1)
+        params = {
+            "run_uuid": "test_run",
+            "slurm_logfile": "test_logfile",
+            "comment_str": "test_comment",
+            "account": None,
+            "partition": None,
+            "workdir": ".",
+        }
+
+        # Patch subprocess.Popen to capture the sbatch command
+        with patch("subprocess.Popen") as mock_popen:
+            # Configure the mock to return successful submission
+            process_mock = MagicMock()
+            process_mock.communicate.return_value = ("123", "")
+            process_mock.returncode = 0
+            mock_popen.return_value = process_mock
+
+        assert "--ntasks-per-gpu=1" in get_submit_command(job, params)
+
 
 class TestWildcardsWithSlashes(snakemake.common.tests.TestWorkflowsLocalStorageBase):
     """