diff --git a/utils/python/CIME/aprun.py b/utils/python/CIME/aprun.py index 25b4f3c13150..a01d6fb15cb8 100755 --- a/utils/python/CIME/aprun.py +++ b/utils/python/CIME/aprun.py @@ -57,16 +57,14 @@ def _get_aprun_cmd_for_case_impl(ntasks, nthreads, rootpes, pstrids, c2 += 1 - logger.info("total tasks is: %s" % total_tasks) - # make sure all maxt values at least 1 for c1 in xrange(0, total_tasks): if maxt[c1] < 1: maxt[c1] = 1 # Compute task and thread settings for batch commands - tasks_per_node, task_count, thread_count, max_thread_count, aprun = \ - 0, 1, maxt[0], maxt[0], "aprun" + tasks_per_node, task_count, thread_count, max_thread_count, total_node_count, aprun = \ + 0, 1, maxt[0], maxt[0], 0, "aprun" for c1 in xrange(1, total_tasks): if maxt[c1] != thread_count: tasks_per_node = min(pes_per_node, max_tasks_per_node / thread_count) @@ -83,6 +81,9 @@ def _get_aprun_cmd_for_case_impl(ntasks, nthreads, rootpes, pstrids, aprun += " -n %d -N %d -d %d %s :" % (task_count, tasks_per_node, thread_count, run_exe) + node_count = int(math.ceil(float(task_count) / tasks_per_node)) + total_node_count += node_count + thread_count = maxt[c1] max_thread_count = max(max_thread_count, maxt[c1]) task_count = 1 @@ -99,6 +100,8 @@ def _get_aprun_cmd_for_case_impl(ntasks, nthreads, rootpes, pstrids, task_per_numa = int(math.ceil(tasks_per_node / 2.0)) + total_node_count += int(math.ceil(float(task_count) / tasks_per_node)) + # Special option for Titan with intel compiler if machine == "titan" and tasks_per_node > 1: aprun += " -S %d" % task_per_numa @@ -107,13 +110,13 @@ def _get_aprun_cmd_for_case_impl(ntasks, nthreads, rootpes, pstrids, aprun += " -n %d -N %d -d %d %s " % (task_count, tasks_per_node, thread_count, run_exe) - return aprun + return aprun, total_node_count ############################################################################### def get_aprun_cmd_for_case(case, run_exe): ############################################################################### """ - Given a case, construct and return the aprun command + Given a case, construct and return the aprun command and optimized node count """ models = case.get_values("COMP_CLASSES") ntasks, nthreads, rootpes, pstrids = [], [], [], [] diff --git a/utils/python/CIME/case.py b/utils/python/CIME/case.py index a4d49aaa5465..33d589e50a6b 100644 --- a/utils/python/CIME/case.py +++ b/utils/python/CIME/case.py @@ -132,7 +132,7 @@ def _initialize_derived_attributes(self): self.thread_count = env_mach_pes.get_max_thread_count(comp_classes) self.tasks_per_node = env_mach_pes.get_tasks_per_node(self.total_tasks, self.thread_count) logger.debug("total_tasks %s thread_count %s"%(self.total_tasks, self.thread_count)) - self.num_nodes = env_mach_pes.get_total_nodes(self.total_tasks, self.thread_count) + self.tasks_per_numa = int(math.ceil(self.tasks_per_node / 2.0)) smt_factor = max(1,int(self.get_value("MAX_TASKS_PER_NODE") / pes_per_node)) @@ -140,6 +140,11 @@ def _initialize_derived_attributes(self): threads_per_core = 1 if (threads_per_node <= pes_per_node) else smt_factor self.cores_per_task = self.thread_count / threads_per_core + if self.get_value("MACH") == "titan": + self.num_nodes = get_aprun_cmd_for_case(self, "acme.exe")[1] + else: + self.num_nodes = env_mach_pes.get_total_nodes(self.total_tasks, self.thread_count) + # Define __enter__ and __exit__ so that we can use this as a context manager # and force a flush on exit. def __enter__(self): @@ -1094,7 +1099,7 @@ def get_mpirun_cmd(self, job="case.run"): # special case for aprun if executable == "aprun": - return get_aprun_cmd_for_case(self, run_exe) + " " + run_misc_suffix + return get_aprun_cmd_for_case(self, run_exe)[0] + " " + run_misc_suffix else: mpi_arg_string = " ".join(args.values()) diff --git a/utils/python/CIME/case_setup.py b/utils/python/CIME/case_setup.py index 84723585b305..f65293eee594 100644 --- a/utils/python/CIME/case_setup.py +++ b/utils/python/CIME/case_setup.py @@ -163,7 +163,7 @@ def _case_setup_impl(case, caseroot, clean=False, test_mode=False, reset=False): # create batch files logger.info("Creating batch script case.run") env_batch = case.get_env("batch") - num_nodes = env_mach_pes.get_total_nodes(pestot, thread_count) + num_nodes = case.num_nodes tasks_per_node = env_mach_pes.get_tasks_per_node(pestot, thread_count) for job in env_batch.get_jobs(): input_batch_script = os.path.join(case.get_value("MACHDIR"), env_batch.get_value('template', subgroup=job))