Use constrain_resources() in Tides

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

Hide file tree

Showing 2 changed files with 32 additions and 10 deletions.
diff --git a/compass/ocean/tests/tides/forward/forward.py b/compass/ocean/tests/tides/forward/forward.py
@@ -72,12 +72,24 @@ def setup(self):
         Set up the test case in the work directory, including downloading any
         dependencies
         """
-        self.ntasks = self.config.getint('tides', 'forward_ntasks')
-        self.min_tasks = self.config.getint('tides', 'forward_min_tasks')
-        self.threads = self.config.getint('tides', 'forward_threads')
+        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):
         """
         Run this step of the testcase
         """
         run_model(self)
+
+    def _get_resources(self):
+        # get the these properties from the config options
+        config = self.config
+        self.ntasks = config.getint('tides', 'forward_ntasks')
+        self.min_tasks = config.getint('tides', 'forward_min_tasks')
+        self.openmp_threads = config.getint('tides', 'forward_threads')
diff --git a/compass/ocean/tests/tides/init/initial_state.py b/compass/ocean/tests/tides/init/initial_state.py
@@ -85,14 +85,17 @@ def __init__(self, test_case, mesh):
 
     def setup(self):
         """
-        Set up the test case in the work directory, including downloading any
-        dependencies.
+        Get resources at setup from config options
         """
-        # get the these properties from the config options
-        config = self.config
-        self.ntasks = config.getint('tides', 'init_ntasks')
-        self.min_tasks = config.getint('tides', 'init_min_tasks')
-        self.threads = config.getint('tides', 'init_threads')
+        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):
         """
@@ -119,3 +122,10 @@ def run(self):
         init["layerThickness"][0, :, 0] = init["bottomDepth"][:]
         init["restingThickness"][:, 0] = init["bottomDepth"][:]
         init.close()
+
+    def _get_resources(self):
+        # get the these properties from the config options
+        config = self.config
+        self.ntasks = config.getint('tides', 'init_ntasks')
+        self.min_tasks = config.getint('tides', 'init_min_tasks')
+        self.openmp_threads = config.getint('tides', 'init_threads')