Revert "revert src changes for later merge with #1138 after testing"

This reverts commit a56bed6.

Author: mandresm

src/esm_runscripts/slurm.py

@@ -63,11 +63,21 @@ def get_jobid():
         return os.environ.get("SLURM_JOB_ID")
 
     def prepare_launcher(self, config, cluster):
+        # which launcher are we using?
+        launcher = config["computer"].get("launcher",None)
+        # friendly check that you are using a launcher that we support
+        if launcher not in ["srun", "mpirun"]:
+            print(" The launcher %s is not compatible with ESM-Tools in SLURM " % (launcher,))
+            print(" Supported launchers for SLURM are srun and mpirun ")
+        
         # MA: not sure how this will play with heterogeneous parallelization
         if "multi_srun" in config["general"]:
             for run_type in list(config["general"]["multi_srun"]):
                 current_hostfile = self.path + "_" + run_type
-                write_one_hostfile(current_hostfile, config)
+                if launcher == "srun":
+                    write_one_hostfile_srun(current_hostfile, config)
+                elif launcher == "mpirun":
+                    write_one_hostfile_mpirun(current_hostfile, config)
 
         if config["computer"].get(
             "heterogeneous_parallelization", False
@@ -76,7 +86,11 @@ def prepare_launcher(self, config, cluster):
             config["general"]["batch"].het_par_launcher_lines(config, cluster)
         else:
             # Standard/old way of running jobs with slurm
-            self.write_one_hostfile(self.path, config)
+            if launcher == "srun":
+                self.write_one_hostfile_srun(self.path, config)
+            elif launcher == "mpirun":
+                # JK: Need to think about how to handle heterogeneous paralleisation here...
+                self.write_one_hostfile_mpirun(self.path, config)
 
             hostfile_in_work = (
                 config["general"]["work_dir"] + "/" + os.path.basename(self.path)
@@ -85,10 +99,11 @@ def prepare_launcher(self, config, cluster):
 
         return config
 
-    def write_one_hostfile(self, hostfile, config):
+    def write_one_hostfile_srun(self, hostfile, config):
         """
         Gathers previously prepared requirements
         (batch_system.calculate_requirements) and writes them to ``self.path``.
+        Suitable for srun
         """
 
         with open(hostfile, "w") as hostfile:
@@ -112,7 +127,50 @@ def write_one_hostfile(self, hostfile, config):
                 hostfile.write(
                     str(start_proc) + "-" + str(end_proc) + "  " + command + "\n"
                 )
+    
+    def write_one_hostfile_mpirun(self, hostfile, config):
+        """ 
+        Gathers previously prepared requirements
+        (batch_system.calculate_requirements) and writes them to ``self.path``.
+        Suitable for mpirun launcher
+        """
+        
+        # make an empty string which we will append commands to
+        mpirun_options = ""
 
+        for model in config["general"]["valid_model_names"]:
+            end_proc = config[model].get("end_proc", None)
+            start_proc = config[model].get("start_proc", None)
+            
+            # a model component like oasis3mct does not need cores
+            # since its technically a library
+            # So start_proc and end_proc will be None. Skip it
+            if start_proc == None or end_proc == None:
+                continue
+            
+            # number of cores needed
+            no_cpus = end_proc - start_proc + 1
+            
+            # check if execution_command or executable exist
+            if "execution_command" in config[model]:
+                command = "./" + config[model]["execution_command"]
+            elif "executable" in config[model]:
+                command = "./" + config[model]["executable"]
+            else:
+                print('warning: the executable or execution_command could not be detemined for %s' % (model,))
+                continue
+            
+            # the mpirun command is set here. 
+            mpirun_options += (
+                    " -np %d %s :" % (no_cpus, command)
+                )
+    
+        mpirun_options = mpirun_options[:-1]  # remove trailing ":"
+    
+        with open(hostfile, "w") as hostfile:
+            hostfile.write(mpirun_options)
+            
+    
     @staticmethod
     def get_job_state(jobid):
         """