Add gpus-per-node to job scripts and resources

E3SM-Project · Jul 12, 2024 · 5fb5ac1 · 5fb5ac1
1 parent 2c61650
commit 5fb5ac1
Show file tree

Hide file tree

Showing 5 changed files with 24 additions and 5 deletions.
diff --git a/polaris/job/__init__.py b/polaris/job/__init__.py
@@ -42,7 +42,7 @@ def write_job_script(config, machine, target_cores, min_cores, work_dir,
     cores = np.sqrt(target_cores * min_cores)
     nodes = int(np.ceil(cores / cores_per_node))
 
-    partition, qos, constraint, wall_time = get_slurm_options(
+    partition, qos, constraint, gpus_per_node, wall_time = get_slurm_options(
         config, machine, nodes)
 
     job_name = config.get('job', 'job_name')
@@ -58,7 +58,7 @@ 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)
+                           gpus_per_node=gpus_per_node, suite=suite)
     text = clean_up_whitespace(text)
     if suite == '':
         script_filename = 'job_script.sh'
@@ -95,6 +95,9 @@ def get_slurm_options(config, machine, nodes):
     constraint : str
         Slurm constraint
 
+    gpus_per_node : str
+        The numer of GPUs per node (if any)
+
     wall_time : str
         Slurm wall time
     """
@@ -131,9 +134,14 @@ def get_slurm_options(config, machine, nodes):
         else:
             constraint = ''
 
+    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')
 
-    return partition, qos, constraint, wall_time
+    return partition, qos, constraint, gpus_per_node, wall_time
 
 
 def clean_up_whitespace(text):

diff --git a/polaris/job/job_script.template b/polaris/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_polaris_env.sh
 polaris serial {{suite}}
diff --git a/polaris/parallel.py b/polaris/parallel.py
@@ -77,6 +77,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
 
 

diff --git a/utils/omega/ctest/job_script.template b/utils/omega/ctest/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 %}
 
 cd {{ build_dir }}
 ./omega_ctest.sh
diff --git a/utils/omega/ctest/omega_ctest.py b/utils/omega/ctest/omega_ctest.py
@@ -120,7 +120,7 @@ def write_job_script(config, machine, compiler, submit):
 
     nodes = 1
 
-    partition, qos, constraint, _ = get_slurm_options(
+    partition, qos, constraint, gpus_per_node, _ = get_slurm_options(
         config, machine, nodes)
 
     wall_time = '0:15:00'
@@ -156,7 +156,7 @@ def write_job_script(config, machine, compiler, submit):
     script = template.render(job_name=job_name, account=account,
                              nodes=f'{nodes}', wall_time=wall_time, qos=qos,
                              partition=partition, constraint=constraint,
-                             build_dir=build_dir)
+                             gpus_per_node=gpus_per_node, build_dir=build_dir)
     script = clean_up_whitespace(script)
 
     build_omega_dir = os.path.abspath('build_omega')