Add gpu_mig40 to Greatlakes. (#811)

* Add gpu_mig40 partition. * Generate test data for gpu_mig40 partition. * Request gpus on the gpu_mig40 partition. * Update test reference data. * Update change log.
glotzerlab · Feb 6, 2024 · 6c1bd8b · 6c1bd8b
1 parent 1d0f9fd
commit 6c1bd8b
Show file tree

Hide file tree

Showing 5 changed files with 10 additions and 5 deletions.
diff --git a/changelog.txt b/changelog.txt
@@ -11,6 +11,11 @@ Version 0.27
 [next] -- not yet released
 --------------------------
 
+Added
++++++
+
+- Support the ``gpu_mig40`` partition on UMich Great Lakes (#811).
+
 Fixed
 +++++
 

diff --git a/flow/environments/umich.py b/flow/environments/umich.py
@@ -14,8 +14,8 @@ class GreatLakesEnvironment(DefaultSlurmEnvironment):
     hostname_pattern = r"gl(-login)?[0-9]+\.arc-ts\.umich\.edu"
     template = "umich-greatlakes.sh"
     _partition_config = _PartitionConfig(
-        cpus_per_node={"default": 36, "gpu": 40},
-        gpus_per_node={"gpu": 2},
+        cpus_per_node={"default": 36, "gpu": 40, "gpu_mig40": 64},
+        gpus_per_node={"gpu": 2, "gpu_mig40": 8},
     )
     # For unknown reasons, srun fails to export environment variables such as
     # PATH on Great Lakes unless explicitly requested to with --export=ALL.
@@ -36,7 +36,7 @@ def add_args(cls, parser):
         super().add_args(parser)
         parser.add_argument(
             "--partition",
-            choices=("standard", "gpu", "largemem"),
+            choices=("standard", "gpu_mig40", "gpu", "largemem"),
             default="standard",
             help="Specify the partition to submit to. (default=standard)",
         )

diff --git a/flow/templates/umich-greatlakes.sh b/flow/templates/umich-greatlakes.sh
@@ -5,7 +5,7 @@
 #SBATCH --nodes={{ resources.num_nodes }}-{{ resources.num_nodes }}
 #SBATCH --ntasks={{ nranks }}
 #SBATCH --cpus-per-task={{ resources.ncpu_tasks // nranks}}
-    {% if partition == 'gpu' %}
+    {% if partition.startswith('gpu') %}
 #SBATCH --gpus-per-task={{ resources.ngpu_tasks // nranks }}
     {% endif %}
 {% endblock tasks %}

diff --git a/tests/generate_template_reference_data.py b/tests/generate_template_reference_data.py
@@ -76,7 +76,7 @@ def init(project):
         ],
         "environments.umich.GreatLakesEnvironment": [
             {
-                "partition": ["standard", "gpu"],
+                "partition": ["standard", "gpu", "gpu_mig40"],
             },
             {
                 "parallel": [False, True],

diff --git a/tests/template_reference_data.tar.gz b/tests/template_reference_data.tar.gz
-Original file line number
+Diff line change
@@ Expand Up / @@ -11,6 +11,11 @@ Version 0.27 @@
     [next] -- not yet released
     --------------------------
+    Added
+    +++++
+    - Support the ``gpu_mig40`` partition on UMich Great Lakes (#811).
     Fixed
     +++++
@@ Expand Down @@