Cleaned a lot the cluster interface. Now it should be much easier to rewrite it in a new way

Author: mesonepigreco

Modules/Cluster.py

@@ -162,7 +162,7 @@ def __init__(self, hostname=None, pwd=None, extra_options="", workdir = "",
         self.workdir=r""
         self.submit_command="sbatch --wait"
         self.submit_name="SBATCH"
-        self.terminal="#!/bin/bash"
+        self.terminal="/bin/bash"
         self.v_nodes="-N "
         self.use_nodes = True
         self.v_cpu="-n "
@@ -384,7 +384,182 @@ def CheckCommunication(self):
             return False
 
         return True
-    
+
+    def create_submission_script(self, labels):
+        """
+        CREATE THE SUBMISSION SCRIPT
+        ===========================================
+
+        This is a function that is general and does not depend on the specific
+        calculator. It is usefull to create the header of the submission script.
+
+        Parameters
+        ----------
+            labels : list 
+                It is a list of the labels of the calculations to be done.
+        
+        Returns
+        -------
+            submission_header : string
+                The text of the submission header.
+        """
+
+        # prepare the submission script
+        submission = "#!" + self.terminal + "\n"
+        
+        
+
+        # Add the submission options
+        if self.use_nodes:
+            submission += "#%s %s%d\n" % (self.submit_name, self.v_nodes, self.n_nodes)
+        if self.use_cpu:
+            submission += "#%s %s%d\n" % (self.submit_name, self.v_cpu, self.n_cpu)
+        if self.use_time:
+            submission += "#%s %s%s\n" % (self.submit_name, self.v_time, self.time)
+        if self.use_account:
+            submission += "#%s %s%s\n" % (self.submit_name, self.v_account, self.account_name)
+        if self.use_memory:
+            submission += "#%s %s%s\n" % (self.submit_name, self.v_memory, self.ram)
+        if self.use_partition:
+            submission += "#%s %s%s\n" % (self.submit_name, self.v_partition, self.partition_name)
+
+        # Append the additional parameters
+        for add_parameter in self.custom_params:
+            if self.custom_params[add_parameter] is None:
+                submission += "#{} --{}\n".format(self.submit_name, add_parameter)
+            else:
+                submission += "#{} --{}={}\n".format(self.submit_name, add_parameter, self.custom_params[add_parameter])
+
+        
+        # Add the set -x option
+        if self.add_set_minus_x:
+            submission += "set -x\n"
+        
+        # Add the loading of the modules
+        submission += self.load_modules + "\n"
+        
+        # Go to the working directory
+        submission += "cd " + self.workdir + "\n"
+
+        # If any, apply the extra text before and after the calculation
+        other_input = ""
+        other_output = ""
+        if (self.additional_script_parameters is not None):
+            other_input, other_output = self.additional_script_parameters(labels)
+
+        submission += other_input
+        
+        # Use the xargs trick
+        #submission += "xargs -d " + r"'\n'" + " -L1 -P%d -a %s -- bash -c\n" % (n_togheder, 
+        for i, lbl in enumerate(labels):
+            submission += self.get_execution_command(lbl)
+
+
+        submission += other_output
+
+        return submission
+
+    def get_execution_command(self, label):
+        """
+        GET THE EXECUTION COMMAND
+        =========================
+
+        Return the command used in the submission script to actually execute the calculation.
+
+        Parameters
+        ----------
+            label : string
+                The label of the calculation
+
+        Returns
+        -------
+            commnad : string
+                The command to be appended to the submission script
+        """
+
+        # Get the MPI command replacing NPROC
+        new_mpicmd  = self.mpi_cmd.replace("NPROC", str(self.n_cpu))
+        
+        # Replace the NPOOL variable and the PREFIX in the binary
+        binary = self.binary.replace("NPOOL", str(self.n_pool)).replace("PREFIX", label)
+
+
+        tmt_str = ""
+        if self.use_timeout:
+            tmt_str = "timeout %d " % self.timeout
+        return "%s%s %s\n" % (tmt_str, new_mpicmd, binary)
+
+    def prepare_input_file(self, structure, calc, label):
+        """
+        PREPARE THE INPUT FILE
+        ======================
+
+        This is specific for quantum espresso and must be inherit and replaced for 
+        other calculators.
+
+        This crates the input file and copy it in the working directory.
+
+        Parameters
+        ----------
+            structure : CellConstructor.Structure.Structure
+                The atomic structure on which to run the calculation
+            calc : the ASE or CellConstructor calculator.
+                In this case, it works with quantum espresso
+            label : string
+                The unique name of this calculation
+        """
+
+        # Prepare the input file
+        atm = structure.get_ase_atoms()
+        atm.set_calculator(calc)
+        ase.io.write("%s/%s.pwi"% (self.local_workdir, label),
+                        atm, **calc.parameters)
+        
+        
+            
+        # First of all clean eventually input/output file of this very same calculation
+        cmd = self.sshcmd + " %s 'rm -f %s/%s%s %s/%s%s'" % (self.hostname, 
+                                                                self.workdir, label, ".pwi",
+                                                                self.workdir, label, ".pwo")
+        self.ExecuteCMD(cmd, False)
+#            cp_res = os.system(cmd + " > /dev/null")
+#            if cp_res != 0:
+#                print "Error while executing:", cmd
+#                print "Return code:", cp_res
+#                sys.stderr.write(cmd + ": exit with code " + str(cp_res) + "\n")
+#            
+        # Copy the file into the cluster
+        cmd = self.scpcmd + " %s/%s%s %s:%s/" % (self.local_workdir, label, 
+                                                ".pwi", self.hostname, 
+                                                self.workdir)
+        cp_res = self.ExecuteCMD(cmd, False)
+        if not cp_res:
+            print ("Error while executing:", cmd)
+            print ("Return code:", cp_res)
+            sys.stderr.write(cmd + ": exit with code " + str(cp_res) + "\n")
+        return cp_res
+        #cp_res = os.system(cmd + " > /dev/null")
+        
+
+    def submit(self, script_location):
+        """
+        SUBMIT THE CALCULATION
+        ======================
+
+        Submit the calculation
+        """
+        
+        cmd = "{ssh} {host} '{submit_cmd} {script}'"
+        if self.use_active_shell:
+            cmd = "{ssh} {host} -t '{shell} --login -c \"{submit_cmd} {script}\"'".format(shell = self.terminal)
+        
+        
+        cmd = cmd.format(ssh = self.sshcmd, host = self.hostname, 
+                         submit_cmd = self.submit_name, script = script_location) 
+
+        return cmd
+
+
     def batch_submission(self, list_of_structures, calc, indices, 
                          in_extension, out_extension,
                          label = "ESP", n_togheder=1):
@@ -434,7 +609,6 @@ def batch_submission(self, list_of_structures, calc, indices,
         # Prepare the input atoms
         app_list = ""
         new_ncpu = self.n_cpu * n_togheder
-        new_mpicmd  = self.mpi_cmd.replace("NPROC", str(self.n_cpu))
         results = [None] * N_structs
         submitted = []
         submission_labels = []
@@ -443,44 +617,10 @@ def batch_submission(self, list_of_structures, calc, indices,
             lbl = label + "_" + str(indices[i])
             submission_labels.append(lbl)
 
-            atm = list_of_structures[i].get_ase_atoms()
-            atm.set_calculator(calc)
-            ase.io.write("%s/%s%s"% (self.local_workdir, lbl, in_extension),
-                         atm,**calc.parameters)
-            
-            
-            # Add the file in the applist
-            binary = self.binary.replace("NPOOL", str(self.n_pool)).replace("PREFIX", lbl)
-            
-                
-            # First of all clean eventually input/output file of this very same calculation
-            cmd = self.sshcmd + " %s 'rm -f %s/%s%s %s/%s%s'" % (self.hostname, 
-                                                                 self.workdir, lbl, in_extension,
-                                                                 self.workdir, lbl, out_extension)
-            self.ExecuteCMD(cmd, False)
-#            cp_res = os.system(cmd + " > /dev/null")
-#            if cp_res != 0:
-#                print "Error while executing:", cmd
-#                print "Return code:", cp_res
-#                sys.stderr.write(cmd + ": exit with code " + str(cp_res) + "\n")
-#            
-            # Copy the file into the cluster
-            cmd = self.scpcmd + " %s/%s%s %s:%s/" % (self.local_workdir, lbl, 
-                                                    in_extension, self.hostname, 
-                                                    self.workdir)
-            cp_res = self.ExecuteCMD(cmd, False)
-            
-            #cp_res = os.system(cmd + " > /dev/null")
-            if not cp_res:
-                print ("Error while executing:", cmd)
-                print ("Return code:", cp_res)
-                sys.stderr.write(cmd + ": exit with code " + str(cp_res) + "\n")
+            # Create the input file and copy it into the cluster
+            if not self.prepare_input_file(list_of_structures[i], calc, lbl):
                 continue
-            
-            tmt_str = ""
-            if self.use_timeout:
-                tmt_str = "timeout %d " % self.timeout
-            app_list += "%s%s %s\n" % (tmt_str, new_mpicmd, binary)
+
             submitted.append(i)
 
         # Save the app list and copy it to the destination
@@ -500,55 +640,7 @@ def batch_submission(self, list_of_structures, calc, indices,
 #            sys.stderr.write(cmd + ": exit with code " + str(cp_res) + "\n")
 #            return results #[None] * N_structs
 
-        
-        # prepare the submission script
-        submission = self.terminal + "\n"
-        
-        # Add the submission options
-        if self.use_nodes:
-            submission += "#%s %s%d\n" % (self.submit_name, self.v_nodes, self.n_nodes)
-        if self.use_cpu:
-            submission += "#%s %s%d\n" % (self.submit_name, self.v_cpu, new_ncpu)
-        if self.use_time:
-            submission += "#%s %s%s\n" % (self.submit_name, self.v_time, self.time)
-        if self.use_account:
-            submission += "#%s %s%s\n" % (self.submit_name, self.v_account, self.account_name)
-        if self.use_memory:
-            submission += "#%s %s%s\n" % (self.submit_name, self.v_memory, self.ram)
-        if self.use_partition:
-            submission += "#%s %s%s\n" % (self.submit_name, self.v_partition, self.partition_name)
-
-        # Append the additional parameters
-        for add_parameter in self.custom_params:
-            if self.custom_params[add_parameter] is None:
-                submission += "#{} --{}\n".format(self.submit_name, add_parameter)
-            else:
-                submission += "#{} --{}={}\n".format(self.submit_name, add_parameter, self.custom_params[add_parameter])
-
-        
-        # Add the set -x option
-        if self.add_set_minus_x:
-            submission += "set -x\n"
-        
-        # Add the loading of the modules
-        submission += self.load_modules + "\n"
-        
-        # Go to the working directory
-        submission += "cd " + self.workdir + "\n"
-
-        # If any, apply the extra text before and after the calculation
-        other_input = ""
-        other_output = ""
-        if self.additional_script_parameters is not None:
-            other_input, other_output = self.additional_script_parameters(submission_labels)
-
-        submission += other_input
-        
-        # Use the xargs trick
-        #submission += "xargs -d " + r"'\n'" + " -L1 -P%d -a %s -- bash -c\n" % (n_togheder, 
-        submission += app_list 
-
-        submission += other_output
+        submission = self.create_submission_script(submission_labels)
 
         # Copy the submission script
         sub_fpath = "%s/%s.sh" % (self.local_workdir, label + "_" + str(indices[0]))