Extend to more tasks

Author: cbegeman

polaris/ocean/ice_shelf/__init__.py

@@ -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_filename,
                     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
@@ -91,7 +91,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/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
 
@@ -70,15 +70,14 @@ def __init__(self, component, resolution, name='forward', subdir=None,
         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
 
@@ -28,7 +28,8 @@ def __init__(self, component, resolution, name, indir):
         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/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_target='culled_graph.info',
+                         graph_target=f'{init.path}/culled_graph.info',
                          output_filename='output.nc',
                          validate_vars=['layerThickness', 'normalVelocity'])