feat!: improved GPU job support

snakemake_executor_plugin_slurm/__init__.py

@@ -114,6 +114,9 @@ class ExecutorSettings(ExecutorSettingsBase):
 # Required:
 # Implementation of your executor
 class Executor(RemoteExecutor):
+    gres_re = re.compile(r"^[a-zA-Z0-9_]+:([a-zA-Z0-9_]+:)?\d+$")
+    gpu_model_re = re.compile(r"^[a-zA-Z0-9_]+$")
+
     def __post_init__(self):
         # run check whether we are running in a SLURM job context
         self.warn_on_jobcontext()
@@ -219,6 +222,45 @@ def run_job(self, job: JobExecutorInterface):
         if self.workflow.executor_settings.requeue:
             call += " --requeue"
 
+        if job.resources.get("gres"):
+            # Validate GRES format (e.g., "gpu:1", "gpu:tesla:2")
+            gres = job.resources.gres
+            if not Executor.gres_re.match(gres):
+                raise WorkflowError(
+                    f"Invalid GRES format: {gres}. Expected format: "
+                    "'<name>:<number>' or '<name>:<type>:<number>'"
+                )
+            gres_string = f" --gres={job.resources.gres}"
+        if job.resources.get("gpu"):
+            # ensure that gres is not set, if gpu and gpu_model are set
+            if job.resources.get("gres"):
+                raise WorkflowError(
+                    "GRES and GPU are set. Please only set one of them."
+                )
+            # ensure that 'gpu' is an integer
+            if not isinstance(job.resources.gpu, int):
+                raise WorkflowError(
+                    "The 'gpu' resource must be an integer. "
+                    f"Got: {job.resources.gpu} ({type(job.resources.gpu)})."
+                )
+            gres_string = f" --gpus={job.resources.gpu}"
+        if job.resources.get("gpu_model") and job.resources.get("gpu"):
+            # validate GPU model format
+            if not Executor.gpu_model_re.match(job.resources.gpu_model):
+                raise WorkflowError(
+                    f"Invalid GPU model format: {job.resources.gpu_model}. "
+                    "Expected format: '<name>'"
+                )
+            gres_string = f" --gpus:{job.resources.gpu_model}:{job.resources.gpu}"
+        elif job.resources.get("gpu_model") and not job.resources.get("gpu"):
+            raise WorkflowError(
+                "GPU model is set, but no GPU number is given. "
+                "Please set 'gpu' as well."
+            )
+        call += (
+            gres_string if job.resources.get("gres") or job.resources.get("gpu") else ""
+        )
+
         if job.resources.get("clusters"):
             call += f" --clusters {job.resources.clusters}"
 
@@ -249,7 +291,11 @@ def run_job(self, job: JobExecutorInterface):
 
         # 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)}"
+        gpu_job = job.resources.get("gpus") 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(
@@ -260,8 +306,9 @@ def run_job(self, job: JobExecutorInterface):
                     "specified. Assuming 'tasks_per_node=1'."
                     "Probably not what you want."
                 )
-
-        call += f" --cpus-per-task={get_cpus_per_task(job)}"
+        # we need to set cpus-per-task OR cpus-per-gpu, the function
+        # will return a string with the corresponding value
+        call += f" {get_cpus_per_task(job, gpu_job)}"
 
         if job.resources.get("slurm_extra"):
             self.check_slurm_extra(job)