Merge branch 'main' into feature/hooks-TrackOperation

flow/environments/umich.py

@@ -17,7 +17,11 @@ class GreatLakesEnvironment(DefaultSlurmEnvironment):
     _gpus_per_node = {"default": 2}
     _shared_partitions = {"standard", "gpu"}
 
-    mpi_cmd = "srun"
+    # For unknown reasons, srun fails to export environment variables such as
+    # PATH on Great Lakes unless explicitly requested to with --export=ALL.
+    # On Great Lakes, srun also fails to flush the buffer until the end of
+    # the job without explicitly setting -u.
+    mpi_cmd = "srun -u --export=ALL"
 
     @classmethod
     def add_args(cls, parser):

flow/templates/umich-greatlakes.sh

@@ -1,16 +1,18 @@
 {% extends "slurm.sh" %}
 {% set partition = partition|default('standard', true) %}
+{% set nranks = (operations|calc_tasks("nranks", parallel, force), 1) | max %}
 {% block tasks %}
     {% if resources.ngpu_tasks and 'gpu' not in partition and not force %}
         {% raise "Requesting GPUs requires a gpu partition!" %}
     {% endif %}
     {% if 'gpu' in partition and resources.ngpu_tasks == 0 and not force %}
         {% raise "Requesting gpu partition without GPUs!" %}
     {% endif %}
-#SBATCH --nodes={{ resources.num_nodes }}
-#SBATCH --ntasks={{ resources.ncpu_tasks }}
+#SBATCH --nodes={{ resources.num_nodes }}-{{ resources.num_nodes }}
+#SBATCH --ntasks={{ nranks }}
+#SBATCH --cpus-per-task={{ resources.ncpu_tasks // nranks}}
     {% if partition == 'gpu' %}
-#SBATCH --gpus={{ resources.ngpu_tasks }}
+#SBATCH --gpus-per-task={{ resources.ngpu_tasks // nranks }}
     {% endif %}
 {% endblock tasks %}
 {% block header %}