Add command-line args to the Step class and a ModelStep base class

compass/model.py

@@ -4,6 +4,208 @@
 from mpas_tools.logging import check_call
 from compass.parallel import run_command
 
+from compass.step import Step
+
+
+class ModelStep(Step):
+    """
+    Attributes
+    ----------
+
+    namelist : str
+        The name of the namelist file
+
+    streams : str
+        The name of the streams file
+
+    update_pio : bool
+        Whether to modify the namelist so the number of PIO tasks and the
+        stride between them is consistent with the number of nodes and
+        cores (one PIO task per node).
+
+    make_graph : bool
+        Whether to make a graph file from the given MPAS mesh file.  If
+        ``True``, ``mesh_filename`` must be given.
+
+    mesh_filename : str
+        The name of an MPAS mesh file to use to make the graph file
+
+    partition_graph : bool
+        Whether to partition the domain for the requested number of cores.
+        If so, the partitioning executable is taken from the ``partition``
+        option of the ``[executables]`` config section.
+
+    graph_filename : str
+        The name of the graph file to partition
+
+    """
+    def __init__(self, test_case, name, subdir=None, ntasks=None,
+                 min_tasks=None, openmp_threads=None, max_memory=None,
+                 cached=False, namelist=None, streams=None, update_pio=True,
+                 make_graph=False, mesh_filename=None, partition_graph=True,
+                 graph_filename='graph.info'):
+        """
+        Make a step for running the model
+
+        Parameters
+        ----------
+        test_case : compass.TestCase
+            The test case this step belongs to
+
+        name : str
+            The name of the step
+
+        subdir : str, optional
+            the subdirectory for the step.  The default is ``name``
+
+        ntasks : int, optional
+            the target number of tasks the step would ideally use.  If too
+            few cores are available on the system to accommodate the number of
+            tasks and the number of cores per task, the step will run on
+            fewer tasks as long as as this is not below ``min_tasks``
+
+        min_tasks : int, optional
+            the number of tasks the step requires.  If the system has too
+            few cores to accommodate the number of tasks and cores per task,
+            the step will fail
+
+        openmp_threads : int, optional
+            the number of OpenMP threads to use
+
+        max_memory : int, optional
+            the amount of memory that the step is allowed to use in MB.
+            This is currently just a placeholder for later use with task
+            parallelism
+
+        cached : bool, optional
+            Whether to get all of the outputs for the step from the database of
+            cached outputs for this MPAS core
+
+        namelist : str, optional
+            The name of the namelist file, default is ``namelist.<mpas_core>``
+
+        streams : str, optional
+            The name of the streams file, default is ``streams.<mpas_core>``
+
+        update_pio : bool, optional
+            Whether to modify the namelist so the number of PIO tasks and the
+            stride between them is consistent with the number of nodes and
+            cores (one PIO task per node).
+
+        make_graph : bool, optional
+            Whether to make a graph file from the given MPAS mesh file.  If
+            ``True``, ``mesh_filename`` must be given.
+
+        mesh_filename : str, optional
+            The name of an MPAS mesh file to use to make the graph file
+
+        partition_graph : bool, optional
+            Whether to partition the domain for the requested number of cores.
+            If so, the partitioning executable is taken from the ``partition``
+            option of the ``[executables]`` config section.
+
+        graph_filename : str, optional
+            The name of the graph file to partition
+        """
+        super().__init__(test_case=test_case, name=name, subdir=subdir,
+                         cpus_per_task=openmp_threads,
+                         min_cpus_per_task=openmp_threads, ntasks=ntasks,
+                         min_tasks=min_tasks, openmp_threads=openmp_threads,
+                         max_memory=max_memory, cached=cached)
+
+        mpas_core = test_case.mpas_core.name
+        if namelist is None:
+            namelist = 'namelist.{}'.format(mpas_core)
+
+        if streams is None:
+            streams = 'streams.{}'.format(mpas_core)
+
+        self.namelist = namelist
+        self.streams = streams
+        self.update_pio = update_pio
+        self.make_graph = make_graph
+        self.mesh_filename = mesh_filename
+        self.partition_graph = partition_graph
+        self.graph_filename = graph_filename
+
+        self.add_input_file(filename='<<<model>>>')
+
+    def setup(self):
+        """ Setup the command-line arguments """
+        config = self.config
+        model = config.get('executables', 'model')
+        model_basename = os.path.basename(model)
+        self.args = [f'./{model_basename}', '-n', self.namelist,
+                     '-s', self.streams]
+
+    def set_model_resources(self, ntasks=None, min_tasks=None,
+                            openmp_threads=None, max_memory=None):
+        """
+        Update the resources for the step.  This can be done within init,
+        ``setup()`` or ``runtime_setup()`` for the step that this step
+        belongs to, or init, ``configure()`` or ``run()`` for the test case
+        that this step belongs to.
+        Parameters
+        ----------
+        ntasks : int, optional
+            the number of tasks the step would ideally use.  If too few
+            cores are available on the system to accommodate the number of
+            tasks and the number of cores per task, the step will run on
+            fewer tasks as long as as this is not below ``min_tasks``
+
+        min_tasks : int, optional
+            the number of tasks the step requires.  If the system has too
+            few cores to accommodate the number of tasks and cores per task,
+            the step will fail
+
+        openmp_threads : int, optional
+            the number of OpenMP threads to use
+
+        max_memory : int, optional
+            the amount of memory that the step is allowed to use in MB.
+            This is currently just a placeholder for later use with task
+            parallelism
+        """
+        self.set_resources(cpus_per_task=openmp_threads,
+                           min_cpus_per_task=openmp_threads, ntasks=ntasks,
+                           min_tasks=min_tasks, openmp_threads=openmp_threads,
+                           max_memory=max_memory)
+
+    def runtime_setup(self):
+        """
+        Update PIO namelist options, make graph file, and partition graph file
+        (if any of these are requested)
+        """
+
+        namelist = self.namelist
+        cores = self.ntasks
+        config = self.config
+        logger = self.logger
+
+        if self.update_pio:
+            self.update_namelist_pio(namelist)
+
+        if self.make_graph:
+            make_graph_file(mesh_filename=self.mesh_filename,
+                            graph_filename=self.graph_filename)
+
+        if self.partition_graph:
+            partition(cores, config, logger, graph_file=self.graph_filename)
+
+    def process_inputs_and_outputs(self):
+        """
+        Process the model as an input, then call the parent class' version
+        """
+        for entry in self.input_data:
+            filename = entry['filename']
+
+            if filename == '<<<model>>>':
+                model = self.config.get('executables', 'model')
+                entry['filename'] = os.path.basename(model)
+                entry['target'] = os.path.abspath(model)
+
+        super().process_inputs_and_outputs()
+
 
 def run_model(step, update_pio=True, partition_graph=True,
               graph_file='graph.info', namelist=None, streams=None):

compass/ocean/tests/baroclinic_channel/forward.py

@@ -1,8 +1,7 @@
-from compass.model import run_model
-from compass.step import Step
+from compass.model import ModelStep
 
 
-class Forward(Step):
+class Forward(ModelStep):
     """
     A step for performing forward MPAS-Ocean runs as part of baroclinic
     channel test cases.
@@ -58,7 +57,7 @@ def __init__(self, test_case, resolution, name='forward', subdir=None,
                                'namelist.{}.forward'.format(resolution))
         if nu is not None:
             # update the viscosity to the requested value
-            options = {'config_mom_del2': '{}'.format(nu)}
+            options = {'config_mom_del2': f'{nu}'}
             self.add_namelist_options(options)
 
         # make sure output is double precision
@@ -72,14 +71,4 @@ def __init__(self, test_case, resolution, name='forward', subdir=None,
         self.add_input_file(filename='graph.info',
                             target='../initial_state/culled_graph.info')
 
-        self.add_model_as_input()
-
         self.add_output_file(filename='output.nc')
-
-    # no setup() is needed
-
-    def run(self):
-        """
-        Run this step of the test case
-        """
-        run_model(self)

compass/ocean/tests/drying_slope/default/__init__.py

@@ -1,4 +1,5 @@
 from compass.testcase import TestCase
+from compass.ocean.tests.drying_slope.mesh import Mesh
 from compass.ocean.tests.drying_slope.initial_state import InitialState
 from compass.ocean.tests.drying_slope.forward import Forward
 from compass.ocean.tests.drying_slope.viz import Viz
@@ -14,8 +15,8 @@ class Default(TestCase):
     resolution : float
         The resolution of the test case in km
 
-    coord_type : str
-        The type of vertical coordinate (``sigma``, ``single_layer``, etc.)
+    coord_type : {'sigma', 'single_layer'}
+        The type of vertical coordinate
     """
 
     def __init__(self, test_group, resolution, coord_type):
@@ -30,8 +31,8 @@ def __init__(self, test_group, resolution, coord_type):
         resolution : float
             The resolution of the test case in km
 
-        coord_type : str
-            The type of vertical coordinate (``sigma``, ``single_layer``)
+        coord_type : {'sigma', 'single_layer'}
+            The type of vertical coordinate
         """
         name = 'default'
 
@@ -44,6 +45,8 @@ def __init__(self, test_group, resolution, coord_type):
         subdir = f'{res_name}/{coord_type}/{name}'
         super().__init__(test_group=test_group, name=name,
                          subdir=subdir)
+
+        self.add_step(Mesh(test_case=self, coord_type=coord_type))
         self.add_step(InitialState(test_case=self, coord_type=coord_type))
         if coord_type == 'single_layer':
             self.add_step(Forward(test_case=self, resolution=resolution,

compass/ocean/tests/drying_slope/forward.py

@@ -1,8 +1,7 @@
-from compass.model import run_model
-from compass.step import Step
+from compass.model import ModelStep
 
 
-class Forward(Step):
+class Forward(ModelStep):
     """
     A step for performing forward MPAS-Ocean runs as part of drying slope
     test cases.
@@ -70,28 +69,17 @@ def __init__(self, test_case, resolution, name='forward', subdir=None,
         self.add_streams_file('compass.ocean.tests.drying_slope',
                               'streams.forward')
 
-        input_path = '../initial_state'
         self.add_input_file(filename='mesh.nc',
-                            target=f'{input_path}/culled_mesh.nc')
+                            target=f'../mesh/culled_mesh.nc')
+
+        self.add_input_file(filename='graph.info',
+                            target=f'../mesh/culled_graph.info')
 
+        input_path = '../initial_state'
         self.add_input_file(filename='init.nc',
                             target=f'{input_path}/ocean.nc')
 
         self.add_input_file(filename='forcing.nc',
                             target=f'{input_path}/init_mode_forcing_data.nc')
 
-        self.add_input_file(filename='graph.info',
-                            target=f'{input_path}/culled_graph.info')
-
-        self.add_model_as_input()
-
         self.add_output_file(filename='output.nc')
-
-    # no setup() is needed
-
-    def run(self):
-        """
-        Run this step of the test case
-        """
-
-        run_model(self)