Repository: ESMCI/cime

Issue:
add logic for aprun for narwhal, warhawk, onyx
Add support for narwhal, warhawk, and onxy for use of aprun.

This should ultimately NOT be required for porting cime, it would be nice if whatever this piece of code is doing were generalized or fixed, so modifications per machine were not needed.

A few test cases were run on narwhal, warhawk, and onyx.

Test suite:
Test baseline:
Test namelist changes:
Test status: bit for bit

Fixes [CIME Github issue #]

User interface changes?: N

Update gh-pages html (Y/N)?: N



Discussion:
See #2024
thanks @jedwards4b.  All I know is that if I don't have this mod, the scripts fail.

```
Traceback (most recent call last):
  File "./case.setup", line 106, in <module>
    _main_func(__doc__)
  File "./case.setup", line 102, in _main_func
    case.case_setup(clean=clean, test_mode=test_mode, reset=reset, keep=keep)
  File "/p/home/apcraig/E3SM/E3SM.dod-port/cime/CIME/case/case_setup.py", line 483, in case_setup
    is_batch=is_batch,
  File "/p/home/apcraig/E3SM/E3SM.dod-port/cime/CIME/utils.py", line 2433, in run_and_log_case_status
    rv = func()
  File "/p/home/apcraig/E3SM/E3SM.dod-port/cime/CIME/case/case_setup.py", line 446, in <lambda>
    self, caseroot, clean=clean, test_mode=test_mode, reset=reset, keep=keep
  File "/p/home/apcraig/E3SM/E3SM.dod-port/cime/CIME/case/case_setup.py", line 372, in _case_setup_impl
    env_batch.make_all_batch_files(case)
  File "/p/home/apcraig/E3SM/E3SM.dod-port/cime/CIME/XML/env_batch.py", line 1361, in make_all_batch_files
    self.make_batch_script(input_batch_script, job, case)
  File "/p/home/apcraig/E3SM/E3SM.dod-port/cime/CIME/XML/env_batch.py", line 232, in make_batch_script
    overrides = self.get_job_overrides(job, case)
  File "/p/home/apcraig/E3SM/E3SM.dod-port/cime/CIME/XML/env_batch.py", line 224, in get_job_overrides
    overrides["mpirun"] = case.get_mpirun_cmd(job=job, overrides=overrides)
  File "/p/home/apcraig/E3SM/E3SM.dod-port/cime/CIME/case/case.py", line 2017, in get_mpirun_cmd
    self, run_exe, overrides=overrides
  File "/p/home/apcraig/E3SM/E3SM.dod-port/cime/CIME/aprun.py", line 177, in get_aprun_cmd_for_case
    run_exe,
  File "/p/home/apcraig/E3SM/E3SM.dod-port/cime/CIME/aprun.py", line 58, in _get_aprun_cmd_for_case_impl
    tt = rootpe + (ntask - 1) * pstrid + 1
TypeError: unsupported operand type(s) for -: 'str' and 'int'
```

It could be that some different settings in config_machines or elsewhere would help, but CIME is very complicated.  Does the above error suggest I can fix this another way?
Also, I cannot find aleph or the other machines in config_machines.xml.  I'm doing this in E3SM, I don't think they are there.
This error seems to indicate that the variable ntask is being interpreted as a string instead of int.   Do you have a PR for ccs_config that you can post as a draft so that I can look at it and try to spot the issue?
@jedwards4b, not sure what you're asking for, PR for css_config?
Right those are cesm machines.   I guess I'll extract myself and let you work with the e3sm team.
@apcraig point us to the branch in github that has those machine configs and we can look at it.
I'm going to create a PR for these machines soon in E3SM.  My dev branch is

https://github.com/apcraig/E3SM/tree/apcraig/e3sm/dod-port

Thanks!

Resolving patch:
diff --git a/CIME/XML/env_mach_specific.py b/CIME/XML/env_mach_specific.py
--- a/CIME/XML/env_mach_specific.py
+++ b/CIME/XML/env_mach_specific.py
@@ -581,16 +581,12 @@ def get_module_system_cmd_path(self, lang):
             else None
         )
 
-    def get_mpirun(self, case, attribs, job, exe_only=False, overrides=None):
-        """
-        Find best match, return (executable, {arg_name : text})
-        """
+    def _find_best_mpirun_match(self, attribs):
         mpirun_nodes = self.get_children("mpirun")
         best_match = None
         best_num_matched = -1
         default_match = None
         best_num_matched_default = -1
-        args = []
         for mpirun_node in mpirun_nodes:
             xml_attribs = self.attrib(mpirun_node)
             all_match = True
@@ -638,14 +634,82 @@ def get_mpirun(self, case, attribs, job, exe_only=False, overrides=None):
             and attribs["mpilib"] == "mpi-serial"
             and best_match is None
         ):
-            return "", [], None, None
+            raise ValueError()
 
         expect(
             best_match is not None or default_match is not None,
             "Could not find a matching MPI for attributes: {}".format(attribs),
         )
 
-        the_match = best_match if best_match is not None else default_match
+        return best_match if best_match is not None else default_match
+
+    def get_aprun_mode(self, attribs):
+        default_mode = "default"
+        valid_modes = ("ignore", "default", "override")
+
+        try:
+            the_match = self._find_best_mpirun_match(attribs)
+        except ValueError:
+            return default_mode
+
+        mode_node = self.get_children("aprun_mode", root=the_match)
+
+        if len(mode_node) == 0:
+            return default_mode
+
+        expect(len(mode_node) == 1, 'Found multiple "aprun_mode" elements.')
+
+        # should have only one element to select from
+        mode = self.text(mode_node[0])
+
+        expect(
+            mode in valid_modes,
+            f"Value {mode!r} for \"aprun_mode\" is not valid, options are {', '.join(valid_modes)!r}",
+        )
+
+        return mode
+
+    def get_aprun_args(self, case, attribs, job, overrides=None):
+        args = {}
+
+        try:
+            the_match = self._find_best_mpirun_match(attribs)
+        except ValueError:
+            return None
+
+        arg_node = self.get_optional_child("arguments", root=the_match)
+
+        if arg_node:
+            arg_nodes = self.get_children("arg", root=arg_node)
+
+            for arg_node in arg_nodes:
+                position = self.get(arg_node, "position")
+
+                if position is None:
+                    position = "per"
+
+                arg_value = transform_vars(
+                    self.text(arg_node),
+                    case=case,
+                    subgroup=job,
+                    overrides=overrides,
+                    default=self.get(arg_node, "default"),
+                )
+
+                args[arg_value] = dict(position=position)
+
+        return args
+
+    def get_mpirun(self, case, attribs, job, exe_only=False, overrides=None):
+        """
+        Find best match, return (executable, {arg_name : text})
+        """
+        args = []
+
+        try:
+            the_match = self._find_best_mpirun_match(attribs)
+        except ValueError:
+            return "", [], None, None
 
         # Now that we know the best match, compute the arguments
         if not exe_only:
diff --git a/CIME/aprun.py b/CIME/aprun.py
--- a/CIME/aprun.py
+++ b/CIME/aprun.py
@@ -22,6 +22,7 @@ def _get_aprun_cmd_for_case_impl(
     compiler,
     machine,
     run_exe,
+    extra_args,
 ):
     ###############################################################################
     """
@@ -38,19 +39,22 @@ def _get_aprun_cmd_for_case_impl(
     >>> compiler = "pgi"
     >>> machine = "titan"
     >>> run_exe = "e3sm.exe"
-    >>> _get_aprun_cmd_for_case_impl(ntasks, nthreads, rootpes, pstrids, max_tasks_per_node, max_mpitasks_per_node, pio_numtasks, pio_async_interface, compiler, machine, run_exe)
-    (' -S 4 -n 680 -N 8 -d 2 e3sm.exe : -S 2 -n 128 -N 4 -d 4 e3sm.exe ', 117, 808, 4, 4)
+    >>> _get_aprun_cmd_for_case_impl(ntasks, nthreads, rootpes, pstrids, max_tasks_per_node, max_mpitasks_per_node, pio_numtasks, pio_async_interface, compiler, machine, run_exe, None)
+    ('  -S 4 -n 680 -N 8 -d 2  e3sm.exe : -S 2 -n 128 -N 4 -d 4  e3sm.exe ', 117, 808, 4, 4)
     >>> compiler = "intel"
-    >>> _get_aprun_cmd_for_case_impl(ntasks, nthreads, rootpes, pstrids, max_tasks_per_node, max_mpitasks_per_node, pio_numtasks, pio_async_interface, compiler, machine, run_exe)
-    (' -S 4 -cc numa_node -n 680 -N 8 -d 2 e3sm.exe : -S 2 -cc numa_node -n 128 -N 4 -d 4 e3sm.exe ', 117, 808, 4, 4)
+    >>> _get_aprun_cmd_for_case_impl(ntasks, nthreads, rootpes, pstrids, max_tasks_per_node, max_mpitasks_per_node, pio_numtasks, pio_async_interface, compiler, machine, run_exe, None)
+    ('  -S 4 -cc numa_node -n 680 -N 8 -d 2  e3sm.exe : -S 2 -cc numa_node -n 128 -N 4 -d 4  e3sm.exe ', 117, 808, 4, 4)
 
     >>> ntasks = [64, 64, 64, 64, 64, 64, 64, 64, 1]
     >>> nthreads = [1, 1, 1, 1, 1, 1, 1, 1, 1]
     >>> rootpes = [0, 0, 0, 0, 0, 0, 0, 0, 0]
     >>> pstrids = [1, 1, 1, 1, 1, 1, 1, 1, 1]
-    >>> _get_aprun_cmd_for_case_impl(ntasks, nthreads, rootpes, pstrids, max_tasks_per_node, max_mpitasks_per_node, pio_numtasks, pio_async_interface, compiler, machine, run_exe)
-    (' -S 8 -cc numa_node -n 64 -N 16 -d 1 e3sm.exe ', 4, 64, 16, 1)
+    >>> _get_aprun_cmd_for_case_impl(ntasks, nthreads, rootpes, pstrids, max_tasks_per_node, max_mpitasks_per_node, pio_numtasks, pio_async_interface, compiler, machine, run_exe, None)
+    ('  -S 8 -cc numa_node -n 64 -N 16 -d 1  e3sm.exe ', 4, 64, 16, 1)
     """
+    if extra_args is None:
+        extra_args = {}
+
     max_tasks_per_node = 1 if max_tasks_per_node < 1 else max_tasks_per_node
 
     total_tasks = 0
@@ -78,6 +82,12 @@ def _get_aprun_cmd_for_case_impl(
         if maxt[c1] < 1:
             maxt[c1] = 1
 
+    global_flags = " ".join(
+        [x for x, y in extra_args.items() if y["position"] == "global"]
+    )
+
+    per_flags = " ".join([x for x, y in extra_args.items() if y["position"] == "per"])
+
     # Compute task and thread settings for batch commands
     (
         tasks_per_node,
@@ -88,7 +98,7 @@ def _get_aprun_cmd_for_case_impl(
         total_node_count,
         total_task_count,
         aprun_args,
-    ) = (0, max_mpitasks_per_node, 1, maxt[0], maxt[0], 0, 0, "")
+    ) = (0, max_mpitasks_per_node, 1, maxt[0], maxt[0], 0, 0, f" {global_flags}")
     c1list = list(range(1, total_tasks))
     c1list.append(None)
     for c1 in c1list:
@@ -107,10 +117,11 @@ def _get_aprun_cmd_for_case_impl(
                 if compiler == "intel":
                     aprun_args += " -cc numa_node"
 
-            aprun_args += " -n {:d} -N {:d} -d {:d} {} {}".format(
+            aprun_args += " -n {:d} -N {:d} -d {:d} {} {} {}".format(
                 task_count,
                 tasks_per_node,
                 thread_count,
+                per_flags,
                 run_exe,
                 "" if c1 is None else ":",
             )
@@ -140,7 +151,7 @@ def _get_aprun_cmd_for_case_impl(
 
 
 ###############################################################################
-def get_aprun_cmd_for_case(case, run_exe, overrides=None):
+def get_aprun_cmd_for_case(case, run_exe, overrides=None, extra_args=None):
     ###############################################################################
     """
     Given a case, construct and return the aprun command and optimized node count
@@ -156,6 +167,10 @@ def get_aprun_cmd_for_case(case, run_exe, overrides=None):
             the_list.append(case.get_value("_".join([item_name, model])))
     max_tasks_per_node = case.get_value("MAX_TASKS_PER_NODE")
     if overrides:
+        overrides = {
+            x: y if isinstance(y, int) or y is None else int(y)
+            for x, y in overrides.items()
+        }
         if "max_tasks_per_node" in overrides:
             max_tasks_per_node = overrides["max_tasks_per_node"]
         if "total_tasks" in overrides:
@@ -175,4 +190,5 @@ def get_aprun_cmd_for_case(case, run_exe, overrides=None):
         case.get_value("COMPILER"),
         case.get_value("MACH"),
         run_exe,
+        extra_args,
     )
diff --git a/CIME/case/case.py b/CIME/case/case.py
--- a/CIME/case/case.py
+++ b/CIME/case/case.py
@@ -2011,15 +2011,25 @@ def get_mpirun_cmd(self, job=None, allow_unresolved_envvars=True, overrides=None
             )
             run_misc_suffix = custom_run_misc_suffix
 
+        aprun_mode = env_mach_specific.get_aprun_mode(mpi_attribs)
+
         # special case for aprun
         if (
             executable is not None
             and "aprun" in executable
-            and not "theta" in self.get_value("MACH")
+            and aprun_mode != "ignore"
+            # and not "theta" in self.get_value("MACH")
         ):
-            aprun_args, num_nodes = get_aprun_cmd_for_case(
-                self, run_exe, overrides=overrides
-            )[0:2]
+            extra_args = env_mach_specific.get_aprun_args(
+                self, mpi_attribs, job, overrides=overrides
+            )
+
+            aprun_args, num_nodes, _, _, _ = get_aprun_cmd_for_case(
+                self,
+                run_exe,
+                overrides=overrides,
+                extra_args=extra_args,
+            )
             if job in ("case.run", "case.test"):
                 expect(
                     (num_nodes + self.spare_nodes) == self.num_nodes,


Test patch:
diff --git a/CIME/tests/test_unit_aprun.py b/CIME/tests/test_unit_aprun.py
new file mode 100644
--- /dev/null
+++ b/CIME/tests/test_unit_aprun.py
@@ -0,0 +1,128 @@
+import unittest
+from unittest import mock
+
+from CIME import aprun
+
+# NTASKS, NTHRDS, ROOTPE, PSTRID
+DEFAULT_COMP_ATTRS = [
+    512,
+    2,
+    0,
+    1,
+    675,
+    2,
+    0,
+    1,
+    168,
+    2,
+    512,
+    1,
+    512,
+    2,
+    0,
+    1,
+    128,
+    4,
+    680,
+    1,
+    168,
+    2,
+    512,
+    1,
+    168,
+    2,
+    512,
+    1,
+    512,
+    2,
+    0,
+    1,
+    1,
+    1,
+    0,
+    1,
+]
+
+# MAX_TASKS_PER_NODE, MAX_MPITASKS_PER_NODE, PIO_NUMTASKS, PIO_ASYNC_INTERFACE, COMPILER, MACH
+DEFAULT_ARGS = [
+    16,
+    16,
+    -1,
+    False,
+    "gnu",
+    "docker",
+]
+
+
+class TestUnitAprun(unittest.TestCase):
+    def test_aprun_extra_args(self):
+        case = mock.MagicMock()
+
+        case.get_values.return_value = [
+            "CPL",
+            "ATM",
+            "LND",
+            "ICE",
+            "OCN",
+            "ROF",
+            "GLC",
+            "WAV",
+            "IAC",
+        ]
+
+        case.get_value.side_effect = DEFAULT_COMP_ATTRS + DEFAULT_ARGS
+
+        extra_args = {
+            "-e DEBUG=true": {"position": "global"},
+            "-j 20": {"position": "per"},
+        }
+
+        (
+            aprun_args,
+            total_node_count,
+            total_task_count,
+            min_tasks_per_node,
+            max_thread_count,
+        ) = aprun.get_aprun_cmd_for_case(case, "e3sm.exe", extra_args=extra_args)
+
+        assert (
+            aprun_args
+            == " -e DEBUG=true -n 680 -N 8 -d 2 -j 20 e3sm.exe : -n 128 -N 4 -d 4 -j 20 e3sm.exe "
+        )
+        assert total_node_count == 117
+        assert total_task_count == 808
+        assert min_tasks_per_node == 4
+        assert max_thread_count == 4
+
+    def test_aprun(self):
+        case = mock.MagicMock()
+
+        case.get_values.return_value = [
+            "CPL",
+            "ATM",
+            "LND",
+            "ICE",
+            "OCN",
+            "ROF",
+            "GLC",
+            "WAV",
+            "IAC",
+        ]
+
+        case.get_value.side_effect = DEFAULT_COMP_ATTRS + DEFAULT_ARGS
+
+        (
+            aprun_args,
+            total_node_count,
+            total_task_count,
+            min_tasks_per_node,
+            max_thread_count,
+        ) = aprun.get_aprun_cmd_for_case(case, "e3sm.exe")
+
+        assert (
+            aprun_args == "  -n 680 -N 8 -d 2  e3sm.exe : -n 128 -N 4 -d 4  e3sm.exe "
+        )
+        assert total_node_count == 117
+        assert total_task_count == 808
+        assert min_tasks_per_node == 4
+        assert max_thread_count == 4
diff --git a/CIME/tests/test_unit_xml_env_mach_specific.py b/CIME/tests/test_unit_xml_env_mach_specific.py
--- a/CIME/tests/test_unit_xml_env_mach_specific.py
+++ b/CIME/tests/test_unit_xml_env_mach_specific.py
@@ -4,12 +4,322 @@
 import tempfile
 from unittest import mock
 
+from CIME import utils
 from CIME.XML.env_mach_specific import EnvMachSpecific
 
 # pylint: disable=unused-argument
 
 
 class TestXMLEnvMachSpecific(unittest.TestCase):
+    def test_aprun_get_args(self):
+        with tempfile.NamedTemporaryFile() as temp:
+            temp.write(
+                b"""<?xml version="1.0"?>
+<file id="env_mach_specific.xml" version="2.0">
+  <header>
+    These variables control the machine dependent environment including
+    the paths to compilers and libraries external to cime such as netcdf,
+    environment variables for use in the running job should also be set here.
+    </header>
+  <group id="compliant_values">
+    <entry id="run_exe" value="${EXEROOT}/e3sm.exe ">
+      <type>char</type>
+      <desc>executable name</desc>
+    </entry>
+    <entry id="run_misc_suffix" value=" &gt;&gt; e3sm.log.$LID 2&gt;&amp;1 ">
+      <type>char</type>
+      <desc>redirect for job output</desc>
+    </entry>
+  </group>
+  <module_system type="none"/>
+  <environment_variables>
+    <env name="OMPI_ALLOW_RUN_AS_ROOT">1</env>
+    <env name="OMPI_ALLOW_RUN_AS_ROOT_CONFIRM">1</env>
+  </environment_variables>
+  <mpirun mpilib="openmpi">
+    <aprun_mode>override</aprun_mode>
+    <executable>aprun</executable>
+    <arguments>
+      <arg name="default_per">-j 10</arg>
+      <arg name="ntasks" position="global">-n {{ total_tasks }}</arg>
+      <arg name="oversubscribe" position="per">--oversubscribe</arg>
+    </arguments>
+  </mpirun>
+</file>
+"""
+            )
+            temp.seek(0)
+
+            mach_specific = EnvMachSpecific(infile=temp.name)
+
+            attribs = {"compiler": "gnu", "mpilib": "openmpi", "threaded": False}
+
+            case = mock.MagicMock()
+
+            type(case).total_tasks = mock.PropertyMock(return_value=4)
+
+            extra_args = mach_specific.get_aprun_args(case, attribs, "case.run")
+
+            expected_args = {
+                "-j 10": {"position": "per"},
+                "--oversubscribe": {"position": "per"},
+                "-n 4": {"position": "global"},
+            }
+
+            assert extra_args == expected_args
+
+    def test_get_aprun_mode_not_valid(self):
+        with tempfile.NamedTemporaryFile() as temp:
+            temp.write(
+                b"""<?xml version="1.0"?>
+<file id="env_mach_specific.xml" version="2.0">
+  <header>
+    These variables control the machine dependent environment including
+    the paths to compilers and libraries external to cime such as netcdf,
+    environment variables for use in the running job should also be set here.
+    </header>
+  <group id="compliant_values">
+    <entry id="run_exe" value="${EXEROOT}/e3sm.exe ">
+      <type>char</type>
+      <desc>executable name</desc>
+    </entry>
+    <entry id="run_misc_suffix" value=" &gt;&gt; e3sm.log.$LID 2&gt;&amp;1 ">
+      <type>char</type>
+      <desc>redirect for job output</desc>
+    </entry>
+  </group>
+  <module_system type="none"/>
+  <environment_variables>
+    <env name="OMPI_ALLOW_RUN_AS_ROOT">1</env>
+    <env name="OMPI_ALLOW_RUN_AS_ROOT_CONFIRM">1</env>
+  </environment_variables>
+  <mpirun mpilib="openmpi">
+    <aprun_mode>custom</aprun_mode>
+    <executable>aprun</executable>
+    <arguments>
+      <arg name="ntasks">-n {{ total_tasks }}</arg>
+      <arg name="oversubscribe">--oversubscribe</arg>
+    </arguments>
+  </mpirun>
+</file>
+"""
+            )
+            temp.seek(0)
+
+            mach_specific = EnvMachSpecific(infile=temp.name)
+
+            attribs = {"compiler": "gnu", "mpilib": "openmpi", "threaded": False}
+
+            with self.assertRaises(utils.CIMEError) as e:
+                mach_specific.get_aprun_mode(attribs)
+
+            assert (
+                str(e.exception)
+                == "ERROR: Value 'custom' for \"aprun_mode\" is not valid, options are 'ignore, default, override'"
+            )
+
+    def test_get_aprun_mode_user_defined(self):
+        with tempfile.NamedTemporaryFile() as temp:
+            temp.write(
+                b"""<?xml version="1.0"?>
+<file id="env_mach_specific.xml" version="2.0">
+  <header>
+    These variables control the machine dependent environment including
+    the paths to compilers and libraries external to cime such as netcdf,
+    environment variables for use in the running job should also be set here.
+    </header>
+  <group id="compliant_values">
+    <entry id="run_exe" value="${EXEROOT}/e3sm.exe ">
+      <type>char</type>
+      <desc>executable name</desc>
+    </entry>
+    <entry id="run_misc_suffix" value=" &gt;&gt; e3sm.log.$LID 2&gt;&amp;1 ">
+      <type>char</type>
+      <desc>redirect for job output</desc>
+    </entry>
+  </group>
+  <module_system type="none"/>
+  <environment_variables>
+    <env name="OMPI_ALLOW_RUN_AS_ROOT">1</env>
+    <env name="OMPI_ALLOW_RUN_AS_ROOT_CONFIRM">1</env>
+  </environment_variables>
+  <mpirun mpilib="openmpi">
+    <aprun_mode>default</aprun_mode>
+    <executable>aprun</executable>
+    <arguments>
+      <arg name="ntasks">-n {{ total_tasks }}</arg>
+      <arg name="oversubscribe">--oversubscribe</arg>
+    </arguments>
+  </mpirun>
+</file>
+"""
+            )
+            temp.seek(0)
+
+            mach_specific = EnvMachSpecific(infile=temp.name)
+
+            attribs = {"compiler": "gnu", "mpilib": "openmpi", "threaded": False}
+
+            aprun_mode = mach_specific.get_aprun_mode(attribs)
+
+            assert aprun_mode == "default"
+
+    def test_get_aprun_mode_default(self):
+        with tempfile.NamedTemporaryFile() as temp:
+            temp.write(
+                b"""<?xml version="1.0"?>
+<file id="env_mach_specific.xml" version="2.0">
+  <header>
+    These variables control the machine dependent environment including
+    the paths to compilers and libraries external to cime such as netcdf,
+    environment variables for use in the running job should also be set here.
+    </header>
+  <group id="compliant_values">
+    <entry id="run_exe" value="${EXEROOT}/e3sm.exe ">
+      <type>char</type>
+      <desc>executable name</desc>
+    </entry>
+    <entry id="run_misc_suffix" value=" &gt;&gt; e3sm.log.$LID 2&gt;&amp;1 ">
+      <type>char</type>
+      <desc>redirect for job output</desc>
+    </entry>
+  </group>
+  <module_system type="none"/>
+  <environment_variables>
+    <env name="OMPI_ALLOW_RUN_AS_ROOT">1</env>
+    <env name="OMPI_ALLOW_RUN_AS_ROOT_CONFIRM">1</env>
+  </environment_variables>
+  <mpirun mpilib="openmpi">
+    <executable>aprun</executable>
+    <arguments>
+      <arg name="ntasks">-n {{ total_tasks }}</arg>
+      <arg name="oversubscribe">--oversubscribe</arg>
+    </arguments>
+  </mpirun>
+</file>
+"""
+            )
+            temp.seek(0)
+
+            mach_specific = EnvMachSpecific(infile=temp.name)
+
+            attribs = {"compiler": "gnu", "mpilib": "openmpi", "threaded": False}
+
+            aprun_mode = mach_specific.get_aprun_mode(attribs)
+
+            assert aprun_mode == "default"
+
+    def test_find_best_mpirun_match(self):
+        with tempfile.NamedTemporaryFile() as temp:
+            temp.write(
+                b"""<?xml version="1.0"?>
+<file id="env_mach_specific.xml" version="2.0">
+  <header>
+    These variables control the machine dependent environment including
+    the paths to compilers and libraries external to cime such as netcdf,
+    environment variables for use in the running job should also be set here.
+    </header>
+  <group id="compliant_values">
+    <entry id="run_exe" value="${EXEROOT}/e3sm.exe ">
+      <type>char</type>
+      <desc>executable name</desc>
+    </entry>
+    <entry id="run_misc_suffix" value=" &gt;&gt; e3sm.log.$LID 2&gt;&amp;1 ">
+      <type>char</type>
+      <desc>redirect for job output</desc>
+    </entry>
+  </group>
+  <module_system type="none"/>
+  <environment_variables>
+    <env name="OMPI_ALLOW_RUN_AS_ROOT">1</env>
+    <env name="OMPI_ALLOW_RUN_AS_ROOT_CONFIRM">1</env>
+  </environment_variables>
+  <mpirun mpilib="openmpi">
+    <executable>aprun</executable>
+    <arguments>
+      <arg name="ntasks">-n {{ total_tasks }}</arg>
+      <arg name="oversubscribe">--oversubscribe</arg>
+    </arguments>
+  </mpirun>
+  <mpirun mpilib="openmpi" compiler="gnu">
+    <executable>srun</executable>
+  </mpirun>
+</file>
+"""
+            )
+            temp.seek(0)
+
+            mach_specific = EnvMachSpecific(infile=temp.name)
+
+            mock_case = mock.MagicMock()
+
+            type(mock_case).total_tasks = mock.PropertyMock(return_value=4)
+
+            attribs = {"compiler": "gnu", "mpilib": "openmpi", "threaded": False}
+
+            executable, args, run_exe, run_misc_suffix = mach_specific.get_mpirun(
+                mock_case, attribs, "case.run"
+            )
+
+            assert executable == "srun"
+            assert args == []
+            assert run_exe is None
+            assert run_misc_suffix is None
+
+    def test_get_mpirun(self):
+        with tempfile.NamedTemporaryFile() as temp:
+            temp.write(
+                b"""<?xml version="1.0"?>
+<file id="env_mach_specific.xml" version="2.0">
+  <header>
+    These variables control the machine dependent environment including
+    the paths to compilers and libraries external to cime such as netcdf,
+    environment variables for use in the running job should also be set here.
+    </header>
+  <group id="compliant_values">
+    <entry id="run_exe" value="${EXEROOT}/e3sm.exe ">
+      <type>char</type>
+      <desc>executable name</desc>
+    </entry>
+    <entry id="run_misc_suffix" value=" &gt;&gt; e3sm.log.$LID 2&gt;&amp;1 ">
+      <type>char</type>
+      <desc>redirect for job output</desc>
+    </entry>
+  </group>
+  <module_system type="none"/>
+  <environment_variables>
+    <env name="OMPI_ALLOW_RUN_AS_ROOT">1</env>
+    <env name="OMPI_ALLOW_RUN_AS_ROOT_CONFIRM">1</env>
+  </environment_variables>
+  <mpirun mpilib="openmpi">
+    <executable>aprun</executable>
+    <arguments>
+      <arg name="ntasks">-n {{ total_tasks }}</arg>
+      <arg name="oversubscribe">--oversubscribe</arg>
+    </arguments>
+  </mpirun>
+</file>
+"""
+            )
+            temp.seek(0)
+
+            mach_specific = EnvMachSpecific(infile=temp.name)
+
+            mock_case = mock.MagicMock()
+
+            type(mock_case).total_tasks = mock.PropertyMock(return_value=4)
+
+            attribs = {"compiler": "gnu", "mpilib": "openmpi", "threaded": False}
+
+            executable, args, run_exe, run_misc_suffix = mach_specific.get_mpirun(
+                mock_case, attribs, "case.run"
+            )
+
+            assert executable == "aprun"
+            assert args == ["-n 4", "--oversubscribe"]
+            assert run_exe is None
+            assert run_misc_suffix is None
+
     @mock.patch("CIME.XML.env_mach_specific.EnvMachSpecific.get_optional_child")
     @mock.patch("CIME.XML.env_mach_specific.EnvMachSpecific.text")
     @mock.patch.dict("os.environ", {"TEST_VALUE": "/testexec"})