Repository: ESMCI/cime

Issue:
Need support for TEST specifiers in config_pes.xml
Previously, it was possible to define test-specific PE layouts: e.g.
```
<pes MACH="cetus" TEST="^.+$">
<NTASKS_ATM>64</NTASKS_ATM> <NTHRDS_ATM>16</NTHRDS_ATM> <ROOTPE_ATM>0</ROOTPE_ATM>
...
```

The new PE XML appears to support only machine, compset, grid, pesize tags. Is there a specifier for tests?


Discussion:
You can use a modifier in the test name such as SMS_P64.f19_f19.FC5  here the P64 says set all NTASKS = 64.
Test names cannot be changed, they are part of a test suite. The whole test suite needs to run with 64x16 PEs.
It seems risky to have test-specific PE layouts that are applied in a hidden way like this: It means that it's impossible to test the relevant production PE layout, and it is not obvious to someone working with the tests that they're using a different PE layout from what would be used in production.

That said, I don't feel strongly enough about this to stop someone else from doing it, but I'd want to avoid doing so myself - with a preference for the method suggested by Jim Edwards.
Our case is we want to test non-production cases on Mira so that things like compiler bugs and software stack can be tested.  acme_developer is a bunch of simple compsets and small grids and, on cetus, they all need to be run at 64x16.  I guess the way in cime5 would be to pretend each compset/res is production and make a cetus-specific entry in the appropriate config_pes.xml.  That's a lot more typing then what cime2 required.
Another thought:  if you want a single pe-layout for all tests in a run of a test-suite, the right place to say that is on the command line to create_test.  @jgfouca ?
@rljacob that might be possible. This could potentially implemented in the XML_PHASE of the tests without having to modify the test names as JimE suggested.
If you want a separate command line option for this in create_test it should follow the same internal code path that the _P name modifier does now - I think that's possible without too much trouble.

Resolving patch:
diff --git a/utils/python/CIME/XML/pes.py b/utils/python/CIME/XML/pes.py
--- a/utils/python/CIME/XML/pes.py
+++ b/utils/python/CIME/XML/pes.py
@@ -87,13 +87,13 @@ def _find_matches(self, grid_nodes, grid, compset, machine, pesize_opts, overrid
                                         logger.info("vid is %s"%vid)
                                         if "ntasks" in vid:
                                             for child in node:
-                                                pes_ntasks[child.tag.upper()] = child.text
+                                                pes_ntasks[child.tag.upper()] = int(child.text)
                                         elif "nthrds" in vid:
                                             for child in node:
-                                                pes_nthrds[child.tag.upper()] = child.text
+                                                pes_nthrds[child.tag.upper()] = int(child.text)
                                         elif "rootpe" in vid:
                                             for child in node:
-                                                pes_rootpe[child.tag.upper()] = child.text
+                                                pes_rootpe[child.tag.upper()] = int(child.text)
                                     # if the value is already upper case its something else we are trying to set
                                         elif vid == node.tag:
                                             other_settings[vid] = node.text
@@ -119,13 +119,13 @@ def _find_matches(self, grid_nodes, grid, compset, machine, pesize_opts, overrid
                 logger.debug("vid is %s"%vid)
                 if "ntasks" in vid:
                     for child in node:
-                        pes_ntasks[child.tag.upper()] = child.text
+                        pes_ntasks[child.tag.upper()] = int(child.text)
                 elif "nthrds" in vid:
                     for child in node:
-                        pes_nthrds[child.tag.upper()] = child.text
+                        pes_nthrds[child.tag.upper()] = int(child.text)
                 elif "rootpe" in vid:
                     for child in node:
-                        pes_rootpe[child.tag.upper()] = child.text
+                        pes_rootpe[child.tag.upper()] = int(child.text)
             # if the value is already upper case its something else we are trying to set
                 elif vid == node.tag:
                     other_settings[vid] = node.text
diff --git a/utils/python/CIME/case.py b/utils/python/CIME/case.py
--- a/utils/python/CIME/case.py
+++ b/utils/python/CIME/case.py
@@ -652,50 +652,72 @@ def configure(self, compset_name, grid_name, machine_name=None,
         #--------------------------------------------
         # pe layout
         #--------------------------------------------
-        match1 = re.match('([0-9]+)x([0-9]+)', "" if pecount is None else pecount)
+        match1 = re.match('(.+)x([0-9]+)', "" if pecount is None else pecount)
         match2 = re.match('([0-9]+)', "" if pecount is None else pecount)
+
         pes_ntasks = {}
         pes_nthrds = {}
         pes_rootpe = {}
+        other      = {}
+
+        pesobj = Pes(self._pesfile)
+
+        force_tasks = None
+        force_thrds = None
+
         if match1:
             opti_tasks = match1.group(1)
-            opti_thrds = match1.group(2)
+            if opti_tasks.isdigit():
+                force_tasks = int(opti_tasks)
+            else:
+                pes_ntasks = pesobj.find_pes_layout(self._gridname, self._compsetname, machine_name,
+                                                    pesize_opts=opti_tasks, mpilib=mpilib)[0]
+            force_thrds = int(match1.group(2))
         elif match2:
-            opti_tasks = match2.group(1)
-            opti_thrds = 1
+            force_tasks = int(match2.group(1))
+            pes_nthrds = pesobj.find_pes_layout(self._gridname, self._compsetname, machine_name, mpilib=mpilib)[1]
+        else:
+            pes_ntasks, pes_nthrds, pes_rootpe, other = pesobj.find_pes_layout(self._gridname, self._compsetname,
+                                                                               machine_name, pesize_opts=pecount, mpilib=mpilib)
 
-        other = {}
         if match1 or match2:
             for component_class in self._component_classes:
-                string_ = "NTASKS_" + component_class
-                pes_ntasks[string_] = opti_tasks
-                string_ = "NTHRDS_" + component_class
-                pes_nthrds[string_] = opti_thrds
+                if force_tasks is not None:
+                    string_ = "NTASKS_" + component_class
+                    pes_ntasks[string_] = force_tasks
+
+                if force_thrds is not None:
+                    string_ = "NTHRDS_" + component_class
+                    pes_nthrds[string_] = force_thrds
+
+                # Always default to zero rootpe if user forced procs and or threads
                 string_ = "ROOTPE_" + component_class
                 pes_rootpe[string_] = 0
-        else:
-            pesobj = Pes(self._pesfile)
 
-            pes_ntasks, pes_nthrds, pes_rootpe, other = pesobj.find_pes_layout(self._gridname, self._compsetname,
-                                                                    machine_name, pesize_opts=pecount, mpilib=mpilib)
         mach_pes_obj = self.get_env("mach_pes")
-        totaltasks = {}
+
+
         if other is not None:
             for key, value in other.items():
                 self.set_value(key, value)
-        for key, value in pes_ntasks.items():
-            totaltasks[key[-3:]] = int(value)
-            mach_pes_obj.set_value(key,int(value))
-        for key, value in pes_rootpe.items():
-            totaltasks[key[-3:]] += int(value)
-            mach_pes_obj.set_value(key,int(value))
-        for key, value in pes_nthrds.items():
-            totaltasks[key[-3:]] *= int(value)
-            mach_pes_obj.set_value(key,int(value))
+
+        totaltasks = []
+        for comp_class in self._component_classes:
+            ntasks_str, nthrds_str, rootpe_str = "NTASKS_%s" % comp_class, "NTHRDS_%s" % comp_class, "ROOTPE_%s" % comp_class
+
+            ntasks = pes_ntasks[ntasks_str] if ntasks_str in pes_ntasks else 1
+            nthrds = pes_nthrds[nthrds_str] if nthrds_str in pes_nthrds else 1
+            rootpe = pes_rootpe[rootpe_str] if rootpe_str in pes_rootpe else 0
+
+            totaltasks.append( (ntasks + rootpe) * nthrds )
+
+            mach_pes_obj.set_value(ntasks_str, ntasks)
+            mach_pes_obj.set_value(nthrds_str, nthrds)
+            mach_pes_obj.set_value(rootpe_str, rootpe)
 
         maxval = 1
         pes_per_node = self.get_value("PES_PER_NODE")
-        for key, val in totaltasks.items():
+        for val in totaltasks:
             if val < 0:
                 val = -1*val*pes_per_node
             if val > maxval:
@@ -715,8 +737,6 @@ def configure(self, compset_name, grid_name, machine_name=None,
             if compclass not in pes_nthrds.keys():
                 mach_pes_obj.set_value(compclass,1)
 
-
-
         #--------------------------------------------
         # batch system
         #--------------------------------------------


Test patch:
diff --git a/scripts/create_test b/scripts/create_test
--- a/scripts/create_test
+++ b/scripts/create_test
@@ -211,6 +211,12 @@ OR
     parser.add_argument("--wait", action="store_true",
                         help="On batch systems, wait for submitted jobs to complete")
 
+    parser.add_argument("--force-procs", type=int, default=None,
+                        help="For all tests to run with this number of processors")
+
+    parser.add_argument("--force-threads", type=int, default=None,
+                        help="For all tests to run with this number of threads")
+
     args = parser.parse_args(args[1:])
 
     CIME.utils.handle_standard_logging_options(args)
@@ -335,7 +341,7 @@ OR
         args.test_root, args.baseline_root, args.clean, baseline_cmp_name, baseline_gen_name, \
         args.namelists_only, args.project, args.test_id, args.parallel_jobs, args.walltime, \
         args.single_submit, args.proc_pool, args.use_existing, args.save_timing, args.queue, \
-        args.allow_baseline_overwrite, args.output_root, args.wait
+        args.allow_baseline_overwrite, args.output_root, args.wait, args.force_procs, args.force_threads
 
 ###############################################################################
 def single_submit_impl(machine_name, test_id, proc_pool, project, args, job_cost_map, wall_time, test_root):
@@ -440,7 +446,8 @@ def single_submit_impl(machine_name, test_id, proc_pool, project, args, job_cost
 # pragma pylint: disable=protected-access
 def create_test(test_names, test_data, compiler, machine_name, no_run, no_build, no_setup, no_batch, test_root,
                 baseline_root, clean, baseline_cmp_name, baseline_gen_name, namelists_only, project, test_id, parallel_jobs,
-                walltime, single_submit, proc_pool, use_existing, save_timing, queue, allow_baseline_overwrite, output_root, wait):
+                walltime, single_submit, proc_pool, use_existing, save_timing, queue, allow_baseline_overwrite, output_root, wait,
+                force_procs, force_threads):
 ###############################################################################
     impl = TestScheduler(test_names, test_data=test_data,
                          no_run=no_run, no_build=no_build, no_setup=no_setup, no_batch=no_batch,
@@ -452,7 +459,7 @@ def create_test(test_names, test_data, compiler, machine_name, no_run, no_build,
                          project=project, parallel_jobs=parallel_jobs, walltime=walltime,
                          proc_pool=proc_pool, use_existing=use_existing, save_timing=save_timing,
                          queue=queue, allow_baseline_overwrite=allow_baseline_overwrite,
-                         output_root=output_root)
+                         output_root=output_root, force_procs=force_procs, force_threads=force_threads)
 
     success = impl.run_tests(wait=wait)
 
@@ -493,13 +500,13 @@ def _main_func(description):
     test_names, test_data, compiler, machine_name, no_run, no_build, no_setup, no_batch, \
     test_root, baseline_root, clean, baseline_cmp_name, baseline_gen_name, namelists_only, \
     project, test_id, parallel_jobs, walltime, single_submit, proc_pool, use_existing, \
-    save_timing, queue, allow_baseline_overwrite, output_root, wait \
+    save_timing, queue, allow_baseline_overwrite, output_root, wait, force_procs, force_threads \
         = parse_command_line(sys.argv, description)
 
     sys.exit(create_test(test_names, test_data, compiler, machine_name, no_run, no_build, no_setup, no_batch, test_root,
                          baseline_root, clean, baseline_cmp_name, baseline_gen_name, namelists_only,
                          project, test_id, parallel_jobs, walltime, single_submit, proc_pool, use_existing, save_timing,
-                         queue, allow_baseline_overwrite, output_root, wait))
+                         queue, allow_baseline_overwrite, output_root, wait, force_procs, force_threads))
 
 ###############################################################################
 
diff --git a/utils/python/CIME/test_scheduler.py b/utils/python/CIME/test_scheduler.py
--- a/utils/python/CIME/test_scheduler.py
+++ b/utils/python/CIME/test_scheduler.py
@@ -44,14 +44,18 @@ def __init__(self, test_names, test_data=None,
                  project=None, parallel_jobs=None,
                  walltime=None, proc_pool=None,
                  use_existing=False, save_timing=False, queue=None,
-                 allow_baseline_overwrite=False, output_root=None):
+                 allow_baseline_overwrite=False, output_root=None,
+                 force_procs=None, force_threads=None):
     ###########################################################################
-        self._cime_root  = CIME.utils.get_cime_root()
-        self._cime_model = CIME.utils.get_model()
+        self._cime_root     = CIME.utils.get_cime_root()
+        self._cime_model    = CIME.utils.get_model()
+        self._save_timing   = save_timing
+        self._queue         = queue
+        self._test_data     = {} if test_data is None else test_data # Format:  {test_name -> {data_name -> data}}
+        self._force_procs   = force_procs
+        self._force_threads = force_threads
+
         self._allow_baseline_overwrite  = allow_baseline_overwrite
-        self._save_timing = save_timing
-        self._queue       = queue
-        self._test_data   = {} if test_data is None else test_data # Format:  {test_name -> {data_name -> data}}
 
         self._machobj = Machines(machine=machine_name)
 
@@ -312,6 +316,13 @@ def _create_newcase_phase(self, test):
         if self._output_root is not None:
             create_newcase_cmd += " --output-root %s " % self._output_root
 
+        if self._force_procs or self._force_threads:
+            pecount_str = "M" if self._force_procs is None else str(self._force_procs)
+            if self._force_threads is not None:
+                pecount_str += "x%d" % self._force_threads
+
+            create_newcase_cmd += " --pecount %s" % pecount_str
+
         if test_mods is not None:
             files = Files()
             (component,modspath) = test_mods.split('/',1)
diff --git a/utils/python/tests/scripts_regression_tests.py b/utils/python/tests/scripts_regression_tests.py
--- a/utils/python/tests/scripts_regression_tests.py
+++ b/utils/python/tests/scripts_regression_tests.py
@@ -1182,6 +1182,21 @@ def test_cime_case_mpi_serial(self):
             # Serial cases should be using 1 task
             self.assertEqual(case.get_value("TOTALPES"), 1)
 
+    ###########################################################################
+    def test_cime_case_force_pecount(self):
+    ###########################################################################
+        run_cmd_assert_result(self, "%s/create_test TESTRUNPASS_Mmpi-serial.f19_g16_rx1.A -t %s --no-build --test-root %s --output-root %s --force-procs 16 --force-threads 8"
+                              % (SCRIPT_DIR, self._baseline_name, self._testroot, self._testroot))
+
+        casedir = os.path.join(self._testroot,
+                               "%s.%s" % (CIME.utils.get_full_test_name("TESTRUNPASS_Mmpi-serial.f19_g16_rx1.A", machine=self._machine, compiler=self._compiler), self._baseline_name))
+        self.assertTrue(os.path.isdir(casedir), msg="Missing casedir '%s'" % casedir)
+
+        with Case(casedir, read_only=True) as case:
+            self.assertEqual(case.get_value("NTASKS_CPL"), 16)
+
+            self.assertEqual(case.get_value("NTHRDS_CPL"), 8)
+
 ###############################################################################
 class X_TestSingleSubmit(TestCreateTestCommon):
 ###############################################################################