Use constrain_resources() in SphereTransport

MPAS-Dev · Dec 18, 2022 · 4fe5542 · 4fe5542
1 parent da9bb77
commit 4fe5542
Show file tree

Hide file tree

Showing 8 changed files with 64 additions and 68 deletions.
diff --git a/compass/ocean/tests/sphere_transport/correlated_tracers_2d/__init__.py b/compass/ocean/tests/sphere_transport/correlated_tracers_2d/__init__.py
@@ -57,23 +57,6 @@ def configure(self):
 
         self.update_cores()
 
-    def run(self):
-        """
-        Run each step of the testcase
-        """
-        config = self.config
-        for resolution in self.resolutions:
-            ntasks = config.getint('correlated_tracers_2d',
-                                   f'QU{resolution}_ntasks')
-            min_tasks = config.getint('correlated_tracers_2d',
-                                      f'QU{resolution}_min_tasks')
-            step = self.steps[f'QU{resolution}_forward']
-            step.ntasks = ntasks
-            step.min_tasks = min_tasks
-
-        # run the step
-        super().run()
-
     def update_cores(self):
         """ Update the number of cores and min_tasks for each forward step """
 

diff --git a/compass/ocean/tests/sphere_transport/correlated_tracers_2d/forward.py b/compass/ocean/tests/sphere_transport/correlated_tracers_2d/forward.py
@@ -62,6 +62,14 @@ def setup(self):
                                    config.get(
                                        'correlated_tracers_2d',
                                        'time_integrator')})
+        self._get_resources()
+
+    def constrain_resources(self, available_cores):
+        """
+        Update resources at runtime from config options
+        """
+        self._get_resources()
+        super().constrain_resources(available_cores)
 
     def run(self):
         """
@@ -100,3 +108,11 @@ def get_timestep_str(self):
             dtminutes = dt / timedelta(minutes=1)
             dtstr = "00:" + str(int(dtminutes))[:2].zfill(2) + ":00"
         return dtstr
+
+    def _get_resources(self):
+        resolution = self.resolution
+        config = self.config
+        self.ntasks = config.getint('correlated_tracers_2d',
+                                    f'QU{resolution}_ntasks')
+        self.min_tasks = config.getint('correlated_tracers_2d',
+                                       f'QU{resolution}_min_tasks')
diff --git a/compass/ocean/tests/sphere_transport/divergent_2d/__init__.py b/compass/ocean/tests/sphere_transport/divergent_2d/__init__.py
@@ -53,23 +53,6 @@ def configure(self):
 
         self.update_cores()
 
-    def run(self):
-        """
-        Run each step of the testcase
-        """
-        config = self.config
-        for resolution in self.resolutions:
-            ntasks = config.getint('divergent_2d',
-                                   f'QU{resolution}_ntasks')
-            min_tasks = config.getint('divergent_2d',
-                                      f'QU{resolution}_min_tasks')
-            step = self.steps[f'QU{resolution}_forward']
-            step.ntasks = ntasks
-            step.min_tasks = min_tasks
-
-        # run the step
-        super().run()
-
     def update_cores(self):
         """ Update the number of cores and min_tasks for each forward step """
 

diff --git a/compass/ocean/tests/sphere_transport/divergent_2d/forward.py b/compass/ocean/tests/sphere_transport/divergent_2d/forward.py
@@ -59,6 +59,14 @@ def setup(self):
         self.add_namelist_options({'config_dt': dtstr,
                                    'config_time_integrator': config.get(
                                        'divergent_2d', 'time_integrator')})
+        self._get_resources()
+
+    def constrain_resources(self, available_cores):
+        """
+        Update resources at runtime from config options
+        """
+        self._get_resources()
+        super().constrain_resources(available_cores)
 
     def run(self):
         """
@@ -97,3 +105,11 @@ def get_timestep_str(self):
             dtminutes = dt / timedelta(minutes=1)
             dtstr = "00:" + str(int(dtminutes))[:2].zfill(2) + ":00"
         return dtstr
+
+    def _get_resources(self):
+        resolution = self.resolution
+        config = self.config
+        self.ntasks = config.getint('divergent_2d',
+                                    f'QU{resolution}_ntasks')
+        self.min_tasks = config.getint('divergent_2d',
+                                       f'QU{resolution}_min_tasks')
diff --git a/compass/ocean/tests/sphere_transport/nondivergent_2d/__init__.py b/compass/ocean/tests/sphere_transport/nondivergent_2d/__init__.py
@@ -55,23 +55,6 @@ def configure(self):
 
         self.update_cores()
 
-    def run(self):
-        """
-        Run each step of the testcase
-        """
-        config = self.config
-        for resolution in self.resolutions:
-            ntasks = config.getint('nondivergent_2d',
-                                   f'QU{resolution}_ntasks')
-            min_tasks = config.getint('nondivergent_2d',
-                                      f'QU{resolution}_min_tasks')
-            step = self.steps[f'QU{resolution}_forward']
-            step.ntasks = ntasks
-            step.min_tasks = min_tasks
-
-        # run the step
-        super().run()
-
     def update_cores(self):
         """ Update the number of cores and min_tasks for each forward step """
 

diff --git a/compass/ocean/tests/sphere_transport/nondivergent_2d/forward.py b/compass/ocean/tests/sphere_transport/nondivergent_2d/forward.py
@@ -60,6 +60,14 @@ def setup(self):
         self.add_namelist_options({'config_dt': dtstr,
                                    'config_time_integrator': config.get(
                                        'nondivergent_2d', 'time_integrator')})
+        self._get_resources()
+
+    def constrain_resources(self, available_cores):
+        """
+        Update resources at runtime from config options
+        """
+        self._get_resources()
+        super().constrain_resources(available_cores)
 
     def run(self):
         """
@@ -98,3 +106,11 @@ def get_timestep_str(self):
             dtminutes = dt / timedelta(minutes=1)
             dtstr = "00:" + str(int(dtminutes))[:2].zfill(2) + ":00"
         return dtstr
+
+    def _get_resources(self):
+        resolution = self.resolution
+        config = self.config
+        self.ntasks = config.getint('nondivergent_2d',
+                                    f'QU{resolution}_ntasks')
+        self.min_tasks = config.getint('nondivergent_2d',
+                                       f'QU{resolution}_min_tasks')
diff --git a/compass/ocean/tests/sphere_transport/rotation_2d/__init__.py b/compass/ocean/tests/sphere_transport/rotation_2d/__init__.py
@@ -52,23 +52,6 @@ def configure(self):
 
         self.update_cores()
 
-    def run(self):
-        """
-        Run each step of the testcase
-        """
-        config = self.config
-        for resolution in self.resolutions:
-            ntasks = config.getint('rotation_2d',
-                                   f'QU{resolution}_ntasks')
-            min_tasks = config.getint('rotation_2d',
-                                      f'QU{resolution}_min_tasks')
-            step = self.steps[f'QU{resolution}_forward']
-            step.ntasks = ntasks
-            step.min_tasks = min_tasks
-
-        # run the step
-        super().run()
-
     def update_cores(self):
         """ Update the number of cores and min_tasks for each forward step """
 

diff --git a/compass/ocean/tests/sphere_transport/rotation_2d/forward.py b/compass/ocean/tests/sphere_transport/rotation_2d/forward.py
@@ -59,6 +59,14 @@ def setup(self):
         self.add_namelist_options({'config_dt': dtstr,
                                    'config_time_integrator': config.get(
                                        'rotation_2d', 'time_integrator')})
+        self._get_resources()
+
+    def constrain_resources(self, available_cores):
+        """
+        Update resources at runtime from config options
+        """
+        self._get_resources()
+        super().constrain_resources(available_cores)
 
     def run(self):
         """
@@ -97,3 +105,11 @@ def get_timestep_str(self):
             dtminutes = dt / timedelta(minutes=1)
             dtstr = "00:" + str(int(dtminutes))[:2].zfill(2) + ":00"
         return dtstr
+
+    def _get_resources(self):
+        resolution = self.resolution
+        config = self.config
+        self.ntasks = config.getint('rotation_2d',
+                                    f'QU{resolution}_ntasks')
+        self.min_tasks = config.getint('rotation_2d',
+                                       f'QU{resolution}_min_tasks')