fix: mpi task settings (#363)

This PR restores former behaviour regarding MPI allowing both `tasks`
and `tasks_per_node` to be set.

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

* **Bug Fixes**
* MPI warning now appears only when tasks_per_node, nodes, and tasks are
all missing.
* MPI submissions with missing or conflicting task counts now produce
clear errors to prevent ambiguous jobs.

* **Improvements**
* Enforces sensible minimums for MPI task settings and disallows
specifying both per-node and total task counts.
  * Non-MPI jobs no longer force total-task flags.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
Co-authored-by: Johannes Köster <[email protected]>

Author: 3 people

snakemake_executor_plugin_slurm/__init__.py

@@ -341,17 +341,6 @@ def run_job(self, job: JobExecutorInterface):
                 "- submitting without. This might or might not work on your cluster."
             )
 
-        # MPI job
-        if job.resources.get("mpi", False):
-            if not job.resources.get("tasks_per_node") and not job.resources.get(
-                "nodes"
-            ):
-                self.logger.warning(
-                    "MPI job detected, but no 'tasks_per_node' or 'nodes' "
-                    "specified. Assuming 'tasks_per_node=1'."
-                    "Probably not what you want."
-                )
-
         exec_job = self.format_job_exec(job)
 
         # and finally the job to execute with all the snakemake parameters

snakemake_executor_plugin_slurm/submit_string.py

@@ -1,3 +1,4 @@
+from snakemake_interface_common.exceptions import WorkflowError
 from snakemake_executor_plugin_slurm_jobstep import get_cpu_setting
 from types import SimpleNamespace
 import shlex
@@ -83,7 +84,36 @@ def get_submit_command(job, params):
     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') or 1}"
+        # if the job is a MPI job, ntasks will be set later
+        if not job.resources.get("mpi", False):
+            call += f" --ntasks={job.resources.get('tasks') or 1}"
+
+    # if the job is an MPI job, we need to have some task setting:
+    if job.resources.get("mpi", False):
+        if not job.resources.get("tasks_per_node") and not job.resources.get("tasks"):
+            raise WorkflowError(
+                "For MPI jobs, please specify either "
+                "'tasks_per_node' or 'tasks' (at least one is required)."
+            )
+        # raise an error if both task settings are used
+        if job.resources.get("tasks_per_node") and job.resources.get("tasks"):
+            raise WorkflowError(
+                "For MPI jobs, please specify either 'tasks_per_node' or 'tasks', "
+                "but not both."
+            )
+        if job.resources.get("tasks_per_node"):
+            if job.resources.get("tasks_per_node") <= 1:
+                raise WorkflowError(
+                    "For MPI jobs, 'tasks_per_node' must be greater than 1."
+                )
+            call += f" --ntasks-per-node={job.resources.tasks_per_node}"
+        elif job.resources.get("tasks"):
+            if job.resources.get("tasks") == 1:
+                raise WorkflowError("For MPI jobs, 'tasks' must be greater than 1.")
+            call += f" --ntasks={job.resources.tasks}"
+        # nodes CAN be set independently of tasks or tasks_per_node
+        # this is at a user's discretion. The nodes flag might already
+        # be set above, if the user specified it.
 
     # we need to set cpus-per-task OR cpus-per-gpu, the function
     # will return a string with the corresponding value