Add gpus-per-node to job scripts and resources

compass/job/__init__.py

@@ -88,6 +88,12 @@ def write_job_script(config, machine, target_cores, min_cores, work_dir,
             job_name = 'compass'
         else:
             job_name = f'compass_{suite}'
+
+    if config.has_option('parallel', 'gpus_per_node'):
+        gpus_per_node = config.get('parallel', 'gpus_per_node')
+    else:
+        gpus_per_node = ''
+
     wall_time = config.get('job', 'wall_time')
 
     template = Template(resources.read_text(
@@ -96,7 +102,8 @@ def write_job_script(config, machine, target_cores, min_cores, work_dir,
     text = template.render(job_name=job_name, account=account,
                            nodes=f'{nodes}', wall_time=wall_time, qos=qos,
                            partition=partition, constraint=constraint,
-                           suite=suite, pre_run_commands=pre_run_commands,
+                           gpus_per_node=gpus_per_node, suite=suite,
+                           pre_run_commands=pre_run_commands,
                            post_run_commands=post_run_commands)
     text = _clean_up_whitespace(text)
     if suite == '':

compass/job/job_script.template

@@ -16,6 +16,9 @@
 {% if constraint != '' -%}
 #SBATCH  --constraint={{ constraint }}
 {%- endif %}
+{% if gpus_per_node != '' -%}
+#SBATCH  --gpus-per-node={{ gpus_per_node }}
+{%- endif %}
 
 source load_compass_env.sh
 {{ pre_run_commands }}

compass/parallel.py

@@ -66,6 +66,11 @@ def get_available_parallel_resources(config):
         cores_per_node=cores_per_node,
         mpi_allowed=mpi_allowed
     )
+
+    if config.has_option('parallel', 'gpus_per_node'):
+        available_resources['gpus_per_node'] = \
+            config.getint('parallel', 'gpus_per_node')
+
     return available_resources