Add gres settings to slurm job config

ncsa · Sep 2, 2023 · 5e72fd9 · 5e72fd9
1 parent b19eddf
commit 5e72fd9
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 8 deletions.
diff --git a/README.md b/README.md
@@ -21,14 +21,15 @@ Slurm `sbatch` command. It will tag the run with the Slurm JobID
 You can set values in a json file to control job submission. The supported
 properties in this file are:
 
-|Config File Setting|Use|
-|-------------------|---|
-|partition          | Which Slurm partition should the job run in? |
- |account            | What account name to run under |
-| gpus_per_node     | On GPU partitions how many GPUs to allocate per node |
-| mem               | Amount of memory to allocate to CPU jobs |
-| modules           | List of modules to load before starting job |
-| time              | Max CPU time job may run |
+|Config File Setting| Use                                                                                                            |
+|-------------------|----------------------------------------------------------------------------------------------------------------|
+|partition          | Which Slurm partition should the job run in?                                                                   |
+ |account            | What account name to run under                                                                                 |
+| gpus_per_node     | On GPU partitions how many GPUs to allocate per node                                                           |
+| gres              | SLURM Generic RESources requests                                                                               |
+| mem               | Amount of memory to allocate to CPU jobs                                                                       |
+| modules           | List of modules to load before starting job                                                                    |
+| time              | Max CPU time job may run                                                                                       |
 | sbatch-script-file | Name of batch file to be produced. Leave blank to have service generate a script file name based on the run ID |
 
 ## Development

diff --git a/mlflow_slurm/templates/sbatch_template.sh b/mlflow_slurm/templates/sbatch_template.sh
@@ -8,6 +8,9 @@
 {% if config.gpus_per_node %}
 #SBATCH --gpus-per-node={{ config.gpus_per_node }}
 {% endif %}
+{% if config.gres %}
+#SBATCH --gres={{ config.gres }}
+{% endif %}
 {% if config.mem %}
 #SBATCH --mem={{ config.mem }}
 {% endif %}