Repository: ESMCI/cime

Issue:
Allow ability to select batch queue by num_nodes instead of tasks
Tasks is not necessarily a good metric on queue selection since the number of tasks per node will vary based on the number of threads.


Discussion:

Resolving patch:
diff --git a/scripts/lib/CIME/XML/env_batch.py b/scripts/lib/CIME/XML/env_batch.py
--- a/scripts/lib/CIME/XML/env_batch.py
+++ b/scripts/lib/CIME/XML/env_batch.py
@@ -164,11 +164,11 @@ def make_batch_script(self, input_template, job, case, total_tasks, tasks_per_no
         self.thread_count = thread_count
         task_count = self.get_value("task_count", subgroup=job)
 
-        if task_count == "default":
+        if task_count is None:
             self.total_tasks = total_tasks
             self.num_nodes = num_nodes
         else:
-            self.total_tasks = task_count
+            self.total_tasks = int(task_count)
             self.num_nodes = int(math.ceil(float(task_count)/float(tasks_per_node)))
 
         self.pedocumentation = ""
@@ -185,7 +185,7 @@ def make_batch_script(self, input_template, job, case, total_tasks, tasks_per_no
             fd.write(output_text)
         os.chmod(job, os.stat(job).st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)
 
-    def set_job_defaults(self, batch_jobs, pesize=None, walltime=None, force_queue=None, allow_walltime_override=False):
+    def set_job_defaults(self, batch_jobs, pesize=None, num_nodes=None, tasks_per_node=None, walltime=None, force_queue=None, allow_walltime_override=False):
         if self.batchtype is None:
             self.batchtype = self.get_batch_system_type()
 
@@ -193,24 +193,27 @@ def set_job_defaults(self, batch_jobs, pesize=None, walltime=None, force_queue=N
             return
 
         for job, jsect in batch_jobs:
-            task_count = jsect["task_count"]
-            if task_count is None or task_count == "default":
+            task_count = jsect["task_count"] if "task_count" in jsect else None
+            if task_count is None:
                 task_count = pesize
+                node_count = num_nodes
             else:
-                task_count = int(task_count)
+                expect(tasks_per_node is not None, "Must provide tasks_per_node for custom task_count job '{}'".format(job))
+                task_count = task_count
+                node_count = int(math.ceil(float(task_count)/float(tasks_per_node)))
 
             if force_queue:
-                if not self.queue_meets_spec(force_queue, task_count, walltime=walltime, job=job):
+                if not self.queue_meets_spec(force_queue, task_count, node_count, walltime=walltime, job=job):
                     logger.warning("WARNING: User-requested queue '{}' does not meet requirements for job '{}'".format(force_queue, job))
                 queue = force_queue
             else:
-                queue = self.select_best_queue(task_count, walltime=walltime, job=job)
+                queue = self.select_best_queue(task_count, node_count, walltime=walltime, job=job)
                 if queue is None and walltime is not None:
                     # Try to see if walltime was the holdup
-                    queue = self.select_best_queue(task_count, walltime=None, job=job)
+                    queue = self.select_best_queue(task_count, node_count, walltime=None, job=job)
                     if queue is not None:
                         # It was, override the walltime if a test, otherwise just warn the user
-                        new_walltime = self.get_queue_specs(queue)[3]
+                        new_walltime = self._get_queue_specs(queue)[5]
                         expect(new_walltime is not None, "Should never make it here")
                         logger.warning("WARNING: Requested walltime '{}' could not be matched by any queue".format(walltime))
                         if allow_walltime_override:
@@ -223,16 +226,16 @@ def set_job_defaults(self, batch_jobs, pesize=None, walltime=None, force_queue=N
                     logger.warning("WARNING: No queue on this system met the requirements for this job. Falling back to defaults")
                     default_queue_node = self.get_default_queue()
                     queue = default_queue_node.text
-                    walltime = self.get_queue_specs(queue)[3]
+                    walltime = self._get_queue_specs(queue)[5]
 
             if walltime is None:
                 # Figure out walltime
-                specs = self.get_queue_specs(queue)
+                specs = self._get_queue_specs(queue)
                 if specs is None:
                     # Queue is unknown, use specs from default queue
                     walltime = self.get_default_queue().get("walltimemax")
                 else:
-                    walltime = specs[3]
+                    walltime = specs[5]
 
                 walltime = self._default_walltime if walltime is None else walltime # last-chance fallback
 
@@ -485,23 +488,22 @@ def get_job_id(self, output):
         jobid = search_match.group(1)
         return jobid
 
-    def queue_meets_spec(self, queue, num_pes, walltime=None, job=None):
-        specs = self.get_queue_specs(queue)
+    def queue_meets_spec(self, queue, num_pes, num_nodes, walltime=None, job=None):
+        specs = self._get_queue_specs(queue)
         if specs is None:
             logger.warning("WARNING: queue '{}' is unknown to this system".format(queue))
             return True
 
-        jobmin, jobmax, jobname, walltimemax, strict = specs
+        jobmin, jobmax, nodemin, nodemax, jobname, walltimemax, strict = specs
 
         # A job name match automatically meets spec
         if job is not None and jobname is not None:
             return jobname == job
 
-        if jobmin is not None and num_pes < int(jobmin):
-            return False
-
-        if jobmax is not None and num_pes > int(jobmax):
-            return False
+        for minval, maxval, val in [(jobmin, jobmax, num_pes), (nodemin, nodemax, num_nodes)]:
+            if (minval is not None and val < int(minval)) or \
+               (maxval is not None and val > int(maxval)):
+                return False
 
         if walltime is not None and walltimemax is not None and strict:
             walltime_s = convert_to_seconds(walltime)
@@ -511,7 +513,7 @@ def queue_meets_spec(self, queue, num_pes, walltime=None, job=None):
 
         return True
 
-    def select_best_queue(self, num_pes, walltime=None, job=None):
+    def select_best_queue(self, num_pes, num_nodes, walltime=None, job=None):
         # Make sure to check default queue first.
         all_queues = []
         all_queues.append( self.get_default_queue())
@@ -519,12 +521,12 @@ def select_best_queue(self, num_pes, walltime=None, job=None):
         for queue in all_queues:
             if queue is not None:
                 qname = queue.text
-                if self.queue_meets_spec(qname, num_pes, walltime=walltime, job=job):
+                if self.queue_meets_spec(qname, num_pes, num_nodes, walltime=walltime, job=job):
                     return qname
 
         return None
 
-    def get_queue_specs(self, queue):
+    def _get_queue_specs(self, queue):
         """
         Get queue specifications by name.
 
@@ -534,11 +536,13 @@ def get_queue_specs(self, queue):
             if queue_node.text == queue:
                 jobmin = queue_node.get("jobmin")
                 jobmax = queue_node.get("jobmax")
+                nodemin = queue_node.get("nodemin")
+                nodemax = queue_node.get("nodemax")
                 jobname = queue_node.get("jobname")
                 walltimemax = queue_node.get("walltimemax")
                 strict = queue_node.get("strict") == "true"
 
-                return jobmin, jobmax, jobname, walltimemax, strict
+                return jobmin, jobmax, nodemin, nodemax, jobname, walltimemax, strict
 
         return None
 
diff --git a/scripts/lib/CIME/case.py b/scripts/lib/CIME/case.py
--- a/scripts/lib/CIME/case.py
+++ b/scripts/lib/CIME/case.py
@@ -653,6 +653,7 @@ def _setup_mach_pes(self, pecount, multi_driver, ninst, machine_name, mpilib):
             new_mach_pes_obj = EnvMachPes(infile=self._pesfile, components=self._component_classes)
             self.update_env(new_mach_pes_obj, "mach_pes")
             return new_mach_pes_obj.get_value("TOTALPES")
+
         pesobj = Pes(self._pesfile)
 
         match1 = re.match('(.+)x([0-9]+)', "" if pecount is None else pecount)
@@ -845,20 +846,6 @@ def configure(self, compset_name, grid_name, machine_name=None,
         if multi_driver and ninst>1:
             logger.info(" Driver/Coupler has %s instances" % ninst)
 
-        #--------------------------------------------
-        # batch system
-        #--------------------------------------------
-        env_batch = self.get_env("batch")
-
-        batch_system_type = machobj.get_value("BATCH_SYSTEM")
-        batch = Batch(batch_system=batch_system_type, machine=machine_name)
-        bjobs = batch.get_batch_jobs()
-
-        env_batch.set_batch_system(batch, batch_system_type=batch_system_type)
-        env_batch.create_job_groups(bjobs)
-        env_batch.set_job_defaults(bjobs, pesize=pesize, walltime=walltime, force_queue=queue, allow_walltime_override=test)
-        self.schedule_rewrite(env_batch)
-
         #--------------------------------------------
         # archiving system
         #--------------------------------------------
@@ -931,6 +918,20 @@ def configure(self, compset_name, grid_name, machine_name=None,
 
         self.initialize_derived_attributes()
 
+        #--------------------------------------------
+        # batch system (must come after initialize_derived_attributes)
+        #--------------------------------------------
+        env_batch = self.get_env("batch")
+
+        batch_system_type = machobj.get_value("BATCH_SYSTEM")
+        batch = Batch(batch_system=batch_system_type, machine=machine_name)
+        bjobs = batch.get_batch_jobs()
+
+        env_batch.set_batch_system(batch, batch_system_type=batch_system_type)
+        env_batch.create_job_groups(bjobs)
+        env_batch.set_job_defaults(bjobs, pesize=pesize, num_nodes=self.num_nodes, tasks_per_node=self.tasks_per_node, walltime=walltime, force_queue=queue, allow_walltime_override=test)
+        self.schedule_rewrite(env_batch)
+
         # Make sure that parallel IO is not specified if total_tasks==1
         if self.total_tasks == 1:
             for compclass in self._component_classes:


Test patch: