Fixup partition file presence

docs/developers_guide/framework/model.md

@@ -20,7 +20,9 @@ method,  you can set number of tasks, threads, etc. determined from the
 `openmp_threads` attributes.  These resources need to be set at construction or
 in the  {ref}`dev-step-setup` method (i.e. before calling {ref}`dev-step-run`) 
 so that  the polaris framework  can ensure that the required resources are 
-available.
+available. If the graph partition file has been constructed prior to the ocean
+model step, the path to the graph file should be provided in the
+`graph_target` argument to {py:meth}`polaris.ocean.OceanModelStep.__init__()`.
 
 (dev-model-yaml-namelists-and-streams)=
 

docs/developers_guide/ocean/api.md

@@ -215,6 +215,7 @@
 
    init.Init
    init.Init.run
+   init.Init.setup
 
    viz.Viz
    viz.Viz.run

docs/developers_guide/organization/steps.md

@@ -244,7 +244,6 @@ class SetupMesh(Step):
             self.add_input_file(filename='mpas_grid.nc',
                                 target='dome_varres_grid.nc', database='')
 
-        self.add_output_file(filename='graph.info')
         self.add_output_file(filename='landice_grid.nc')
 ```
 
@@ -366,8 +365,6 @@ class Forward(OceanModelStep):
 
         self.add_input_file(filename='initial_state.nc',
                             target='../../init/initial_state.nc')
-        self.add_input_file(filename='graph.info',
-                            target='../../init/culled_graph.info')
 
         self.add_yaml_file('polaris.ocean.tasks.baroclinic_channel',
                            'forward.yaml')
@@ -803,9 +800,6 @@ def __init__(self, component, mesh, init):
     self.add_input_file(
         filename='forcing_data.nc',
         work_dir_target=f'{init.path}/init/init_mode_forcing_data.nc')
-    self.add_input_file(
-        filename='graph.info',
-        work_dir_target=f'{mesh_path}/culled_graph.info')
 ```
 
 (dev-step-input-polaris)=

docs/tutorials/dev_add_test_group.md

@@ -1024,8 +1024,7 @@ class Init(Step):
 
         self.resolution = resolution
 
-        for file in ['base_mesh.nc', 'culled_mesh.nc', 'culled_graph.info',
-                     'initial_state.nc']:
+        for file in ['base_mesh.nc', 'culled_mesh.nc', 'initial_state.nc']:
             self.add_output_file(file)
 ```
 
@@ -1286,8 +1285,6 @@ class Forward(OceanModelStep):
 
         self.add_input_file(filename='initial_state.nc',
                             target='../init/initial_state.nc')
-        self.add_input_file(filename='graph.info',
-                            target='../init/culled_graph.info')
 
         self.add_output_file(filename='output.nc')
 ```

polaris/ocean/convergence/forward.py

@@ -21,7 +21,7 @@ class ConvergenceForward(OceanModelStep):
 
     def __init__(self, component, name, subdir, resolution, mesh, init,
                  package, yaml_filename='forward.yaml', options=None,
-                 graph_filename='graph.info', output_filename='output.nc',
+                 graph_target=None, output_filename='output.nc',
                  validate_vars=None):
         """
         Create a new step
@@ -50,6 +50,10 @@ def __init__(self, component, name, subdir, resolution, mesh, init,
             A nested dictionary of options and value for each ``config_model``
             to replace model config options with new values
 
+        graph_target : str, optional
+            The graph file name (relative to the base work directory).
+            If none, it will be created.
+
         output_filename : str, optional
             The output file that will be written out at the end of the forward
             run
@@ -58,7 +62,7 @@ def __init__(self, component, name, subdir, resolution, mesh, init,
             A list of variables to validate against a baseline if requested
         """
         super().__init__(component=component, name=name, subdir=subdir,
-                         openmp_threads=1)
+                         openmp_threads=1, graph_target=graph_target)
 
         self.resolution = resolution
         self.package = package
@@ -75,9 +79,6 @@ def __init__(self, component, name, subdir, resolution, mesh, init,
         self.add_input_file(
             filename='init.nc',
             work_dir_target=f'{init.path}/initial_state.nc')
-        self.add_input_file(
-            filename='graph.info',
-            work_dir_target=f'{mesh.path}/{graph_filename}')
 
         self.add_output_file(filename=output_filename,
                              validate_vars=validate_vars)

polaris/ocean/ice_shelf/__init__.py

@@ -50,7 +50,7 @@ def __init__(self, component, min_resolution, name, subdir, sshdir=None):
         self.component = component
         self.min_resolution = min_resolution
 
-    def setup_ssh_adjustment_steps(self, mesh_filename, graph_filename,
+    def setup_ssh_adjustment_steps(self, mesh_filename, graph_target,
                                    init_filename, config, config_filename,
                                    ForwardStep, package=None,
                                    yaml_filename='ssh_forward.yaml',
@@ -67,7 +67,7 @@ def setup_ssh_adjustment_steps(self, mesh_filename, graph_filename,
         mesh_filename : str
             the mesh filename (relative to the base work directory)
 
-        graph_filename : str
+        graph_target: str
             the graph filename (relative to the base work directory)
 
         init_filename : str
@@ -114,7 +114,7 @@ def setup_ssh_adjustment_steps(self, mesh_filename, graph_filename,
             else:
                 ssh_forward = ForwardStep(
                     component=component, min_resolution=min_resolution,
-                    indir=indir, graph_filename=graph_filename,
+                    indir=indir, graph_target=graph_target,
                     init_filename=current_init_filename, name=name,
                     package=package, yaml_filename=yaml_filename,
                     yaml_replacements=yaml_replacements)

polaris/ocean/ice_shelf/ssh_forward.py

@@ -22,7 +22,7 @@ class SshForward(OceanModelStep):
         file
     """
     def __init__(self, component, min_resolution, init_filename,
-                 graph_filename, name='ssh_forward', subdir=None,
+                 graph_target, name='ssh_forward', subdir=None,
                  package=None, yaml_filename='ssh_forward.yaml',
                  yaml_replacements=None, indir=None, ntasks=None,
                  min_tasks=None, openmp_threads=1):
@@ -41,7 +41,7 @@ def __init__(self, component, min_resolution, init_filename,
             the initial condition filename (relative to the base work
             directory)
 
-        graph_filename : str
+        graph_target: str
             the graph filename (relative to the base work directory)
 
         name : str, optional
@@ -78,7 +78,8 @@ def __init__(self, component, min_resolution, init_filename,
         """
         super().__init__(component=component, name=name, subdir=subdir,
                          indir=f'{indir}/ssh_adjustment', ntasks=ntasks,
-                         min_tasks=min_tasks, openmp_threads=openmp_threads)
+                         min_tasks=min_tasks, openmp_threads=openmp_threads,
+                         graph_target=graph_target)
 
         self.min_resolution = min_resolution
         self.package = package
@@ -91,7 +92,7 @@ def __init__(self, component, min_resolution, init_filename,
         self.add_input_file(filename='init.nc',
                             work_dir_target=init_filename)
         self.add_input_file(filename='graph.info',
-                            work_dir_target=graph_filename)
+                            work_dir_target=graph_target)
 
         self.add_output_file(
             filename='output.nc',

polaris/ocean/model/ocean_model_step.py

@@ -20,12 +20,15 @@ class OceanModelStep(ModelStep):
     map : dict
         A nested dictionary that maps from MPAS-Ocean to Omega model config
         options
+
+    graph_target : str
+        The name of the graph partition file to link to (relative to the base
+        working directory)
     """
     def __init__(self, component, name, subdir=None, indir=None, ntasks=None,
                  min_tasks=None, openmp_threads=None, max_memory=None,
                  cached=False, yaml=None, update_pio=True, make_graph=False,
-                 mesh_filename=None, partition_graph=True,
-                 graph_filename='graph.info'):
+                 mesh_filename=None, partition_graph=True, graph_target=None):
         """
         Make a step for running the model
 
@@ -84,20 +87,25 @@ def __init__(self, component, name, subdir=None, indir=None, ntasks=None,
             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
+        graph_target : str, optional
+            The graph file name (relative to the base work directory).
+            If none, it will be created.
         """
+        if graph_target is None:
+            self.make_graph = True
+
         super().__init__(
             component=component, name=name, subdir=subdir, indir=indir,
             ntasks=ntasks, min_tasks=min_tasks, openmp_threads=openmp_threads,
             max_memory=max_memory, cached=cached, yaml=yaml,
             update_pio=update_pio, make_graph=make_graph,
             mesh_filename=mesh_filename, partition_graph=partition_graph,
-            graph_filename=graph_filename)
+            graph_filename='graph.info')
 
         self.dynamic_ntasks = (ntasks is None and min_tasks is None)
 
         self.map: Union[None, List[Dict[str, Dict[str, str]]]] = None
+        self.graph_target = graph_target
 
     def setup(self):
         """
@@ -112,9 +120,14 @@ def setup(self):
             self.config_models = ['ocean', 'Omega']
             self.yaml = 'omega.yml'
             self._read_map()
+            self.partition_graph = False
         elif model == 'mpas-ocean':
             self.config_models = ['ocean', 'mpas-ocean']
             self.make_yaml = False
+            self.add_input_file(
+                filename='graph.info',
+                work_dir_target=self.graph_target)
+
         else:
             raise ValueError(f'Unexpected ocean model: {model}')
 

polaris/ocean/tasks/baroclinic_channel/decomp/__init__.py

@@ -39,7 +39,8 @@ def __init__(self, component, resolution, indir, init):
             self.add_step(Forward(
                 component=component, name=name, indir=self.subdir,
                 ntasks=procs, min_tasks=procs, openmp_threads=1,
-                resolution=resolution, run_time_steps=3))
+                resolution=resolution, run_time_steps=3,
+                graph_target=f'{init.path}/culled_graph.info'))
             subdirs.append(name)
         self.add_step(Validate(component=component, step_subdirs=subdirs,
                                indir=self.subdir))

polaris/ocean/tasks/baroclinic_channel/default/__init__.py

@@ -34,7 +34,8 @@ def __init__(self, component, resolution, indir, init):
         self.add_step(
             Forward(component=component, indir=self.subdir, ntasks=None,
                     min_tasks=None, openmp_threads=1, resolution=resolution,
-                    run_time_steps=3))
+                    run_time_steps=3,
+                    graph_target=f'{init.path}/culled_graph.info'))
 
         self.add_step(
             Viz(component=component, indir=self.subdir))

polaris/ocean/tasks/baroclinic_channel/forward.py

@@ -25,7 +25,7 @@ class Forward(OceanModelStep):
     """
     def __init__(self, component, resolution, name='forward', subdir=None,
                  indir=None, ntasks=None, min_tasks=None, openmp_threads=1,
-                 nu=None, run_time_steps=None):
+                 nu=None, run_time_steps=None, graph_target='graph.info'):
         """
         Create a new task
 
@@ -65,20 +65,23 @@ def __init__(self, component, resolution, name='forward', subdir=None,
 
         run_time_steps : int, optional
             Number of time steps to run for
+
+        graph_target : str, optional
+            The graph file name (relative to the base work directory).
+            If none, it will be created.
         """
         self.resolution = resolution
         self.run_time_steps = run_time_steps
         super().__init__(component=component, name=name, subdir=subdir,
                          indir=indir, ntasks=ntasks, min_tasks=min_tasks,
-                         openmp_threads=openmp_threads)
+                         openmp_threads=openmp_threads,
+                         graph_target=graph_target)
 
         # make sure output is double precision
         self.add_yaml_file('polaris.ocean.config', 'output.yaml')
 
         self.add_input_file(filename='initial_state.nc',
                             target='../../init/initial_state.nc')
-        self.add_input_file(filename='graph.info',
-                            target='../../init/culled_graph.info')
 
         self.add_yaml_file('polaris.ocean.tasks.baroclinic_channel',
                            'forward.yaml')

polaris/ocean/tasks/baroclinic_channel/restart/__init__.py

@@ -35,11 +35,12 @@ def __init__(self, component, resolution, indir, init):
         self.add_step(init, symlink='init')
 
         full = RestartStep(component=component, resolution=resolution,
-                           name='full_run', indir=self.subdir)
+                           name='full_run', indir=self.subdir, init=init)
         self.add_step(full)
 
         restart = RestartStep(component=component, resolution=resolution,
-                              name='restart_run', indir=self.subdir)
+                              name='restart_run', indir=self.subdir,
+                              init=init)
         restart.add_dependency(full, full.name)
         self.add_step(restart)
 

polaris/ocean/tasks/baroclinic_channel/restart/restart_step.py

@@ -7,7 +7,7 @@ class RestartStep(Forward):
     """
     A forward model step in the restart test case
     """
-    def __init__(self, component, resolution, name, indir):
+    def __init__(self, component, resolution, name, indir, init):
         """
         Create a new test case
 
@@ -24,11 +24,15 @@ def __init__(self, component, resolution, name, indir):
 
         indir : str
             the directory the step is in, to which ``name`` will be appended
+
+        init: polaris.ocean.tasks.baroclinic_channel.init.Init
+            the initial state step
         """
         self.resolution = resolution
         super().__init__(component=component, name=name, indir=indir, ntasks=4,
                          min_tasks=4, openmp_threads=1,
-                         resolution=resolution)
+                         resolution=resolution,
+                         graph_target=f'{init.path}/culled_graph.info')
 
     def dynamic_model_config(self, at_setup):
         """

polaris/ocean/tasks/baroclinic_channel/threads/__init__.py

@@ -38,7 +38,8 @@ def __init__(self, component, resolution, indir, init):
             self.add_step(Forward(
                 component=component, name=name, indir=self.subdir, ntasks=4,
                 min_tasks=4, openmp_threads=openmp_threads,
-                resolution=resolution, run_time_steps=3))
+                resolution=resolution, run_time_steps=3,
+                graph_target=f'{init.path}/culled_graph.info'))
             subdirs.append(name)
         self.add_step(Validate(component=component, step_subdirs=subdirs,
                                indir=self.subdir))

polaris/ocean/tasks/cosine_bell/forward.py

@@ -38,4 +38,5 @@ def __init__(self, component, name, subdir, resolution, mesh, init):
                          init=init, package=package,
                          yaml_filename='forward.yaml',
                          output_filename='output.nc',
-                         validate_vars=validate_vars)
+                         validate_vars=validate_vars,
+                         graph_target=f'{init.path}/graph.info')

polaris/ocean/tasks/geostrophic/forward.py

@@ -41,4 +41,5 @@ def __init__(self, component, name, subdir, resolution, mesh, init):
                          init=init, package=package,
                          yaml_filename='forward.yaml',
                          output_filename='output.nc',
-                         validate_vars=validate_vars)
+                         validate_vars=validate_vars,
+                         graph_target=f'{init.path}/graph.info')

polaris/ocean/tasks/ice_shelf_2d/default/__init__.py

@@ -73,7 +73,7 @@ def __init__(self, component, resolution, indir, init, config,
 
         last_adjust_step = self.setup_ssh_adjustment_steps(
             mesh_filename=f'{init.path}/culled_mesh.nc',
-            graph_filename=f'{init.path}/culled_graph.info',
+            graph_target=f'{init.path}/culled_graph.info',
             init_filename=f'{init.path}/output.nc',
             config=config, config_filename='ice_shelf_2d.cfg',
             ForwardStep=SshForward,

polaris/ocean/tasks/ice_shelf_2d/forward.py

@@ -69,7 +69,8 @@ def __init__(self, component, resolution, mesh, init,
             name = 'restart'
         super().__init__(component=component, name=name, subdir=subdir,
                          indir=indir, ntasks=ntasks, min_tasks=min_tasks,
-                         openmp_threads=openmp_threads)
+                         openmp_threads=openmp_threads,
+                         graph_target=f'{mesh.path}/culled_graph.info')
 
         self.resolution = resolution
         self.do_restart = do_restart
@@ -80,8 +81,6 @@ def __init__(self, component, resolution, mesh, init,
 
         self.add_input_file(filename='init.nc',
                             work_dir_target=f'{init.path}/output.nc')
-        self.add_input_file(filename='graph.info',
-                            work_dir_target=f'{mesh.path}/culled_graph.info')
 
         self.add_output_file(
             filename='output.nc',

polaris/ocean/tasks/inertial_gravity_wave/forward.py

@@ -37,7 +37,7 @@ def __init__(self, component, name, resolution, subdir, init):
                          resolution=resolution, mesh=init, init=init,
                          package='polaris.ocean.tasks.inertial_gravity_wave',
                          yaml_filename='forward.yaml',
-                         graph_filename='culled_graph.info',
+                         graph_target=f'{init.path}/culled_graph.info',
                          output_filename='output.nc',
                          validate_vars=['layerThickness', 'normalVelocity'])
 

polaris/ocean/tasks/internal_wave/default/__init__.py

@@ -32,8 +32,8 @@ def __init__(self, component, indir, init, vadv_method='standard'):
         self.add_step(init, symlink='init')
 
         self.add_step(
-            Forward(component=component, indir=self.subdir, ntasks=None,
-                    min_tasks=None, openmp_threads=1,
+            Forward(component=component, init=init, indir=self.subdir,
+                    ntasks=None, min_tasks=None, openmp_threads=1,
                     run_time_steps=3, vadv_method=vadv_method))
 
         self.add_step(