feat!: improved GPU job support

docs/further.md

@@ -118,8 +118,9 @@ You can use the following specifications:
 | `--cpus-per-task`       | `cpus_per_task`      | number of cpus per task (in case of SMP, rather use `threads`)   |
 | `--nodes` | `nodes`    | number of nodes                       |
 | `--clusters` | `clusters` | comma separated string of clusters |
+| `--gres`  | `gres` | generic resource specification, e.g. `gres=gpu:1`, allows for arbitray strings after the `=` |
 
-Each of these can be part of a rule, e.g.:
+Each of these can be part of a rule, e.g.: 
 
 ``` python
 rule:

snakemake_executor_plugin_slurm/__init__.py

@@ -83,6 +83,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()
@@ -158,6 +161,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" --gres=gpu:{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" --gres=gpu:{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}"