Merge pull request #61 from SSCHAcode/new_cluster_interface

New cluster interface

Author: mesonepigreco

Modules/Cluster.py

@@ -13,6 +13,7 @@
     pass
 
 import numpy as np
+import time, datetime
 
 from ase.units import Rydberg, Bohr
 import ase, ase.io
@@ -178,6 +179,10 @@ def __init__(self, hostname=None, pwd=None, extra_options="", workdir = "",
         self.timeout = 1000
         self.use_timeout = False
 
+        # Check the status of the job every TOT seconds
+        self.check_timeout = 300
+        self.nonblocking_command = False # True if you use a different version of slurm that does not accept blocking commands
+
         # This is the number of configurations to be computed for each jub submitted
         # This times the self.batch_size is the total amount of configurations submitted toghether
         self.job_number = 1
@@ -275,7 +280,7 @@ def __setattr__(self, name, value):
 
 
 
-    def ExecuteCMD(self, cmd, raise_error = True, return_output = False):
+    def ExecuteCMD(self, cmd, raise_error = True, return_output = False, on_cluster = False):
         """
         EXECUTE THE CMD ON THE CLUSTER
         ==============================
@@ -292,6 +297,8 @@ def ExecuteCMD(self, cmd, raise_error = True, return_output = False):
             return_output : bool, optional
                 If True (default False) the output of the command is 
                 returned as second value.
+            on_cluster : bool
+                If true, the command is executed directly on the cluster through ssh
                 
         Returns
         -------
@@ -301,6 +308,9 @@ def ExecuteCMD(self, cmd, raise_error = True, return_output = False):
             output : string
                 Returned only if return_output is True
         """
+
+        if on_cluster:
+            cmd = self.sshcmd + " {} '{}'".format(self.hostname, cmd)
 
         success = False
         output = ""
@@ -333,6 +343,8 @@ def ExecuteCMD(self, cmd, raise_error = True, return_output = False):
         if return_output:
             return success, output
         return success
+    
+    
 
 
     def set_timeout(self, timeout):
@@ -520,10 +532,14 @@ def batch_submission(self, list_of_structures, calc, indices,
 
         # Append the additional parameters
         for add_parameter in self.custom_params:
+            adder_string = "--{}".format(add_parameter)
+            if add_parameter.startswith("-"):
+                adder_string = add_parameter
+
             if self.custom_params[add_parameter] is None:
-                submission += "#{} --{}\n".format(self.submit_name, add_parameter)
+                submission += "#{} {}\n".format(self.submit_name, adder_string)
             else:
-                submission += "#{} --{}={}\n".format(self.submit_name, add_parameter, self.custom_params[add_parameter])
+                submission += "#{} {}={}\n".format(self.submit_name, adder_string, self.custom_params[add_parameter])
 
 
         # Add the set -x option
@@ -561,14 +577,27 @@ def batch_submission(self, list_of_structures, calc, indices,
         if not cp_res:
             print ("Error while executing:", cmd)
             print ("Return code:", cp_res)
-            sys.stderr.write(cmd + ": exit with code " + str(cp_res))
+            sys.stderr.write(cmd + ": exit with code " + str(cp_res) + "\n")
             return results#[None] * N_structs
 
 
         # Run the simulation
         cmd = self.sshcmd + " %s '%s %s/%s.sh'" % (self.hostname, self.submit_command, 
                                                    self.workdir, label+ "_" + str(indices[0]))
-        cp_res = self.ExecuteCMD(cmd, False)
+        status, submission_output = self.ExecuteCMD(cmd, True, return_output = True)
+
+        # If the command for submission is non blocking, we need to check periodically wether the calculation has been completed
+        if self.nonblocking_command:
+            job_id = self.get_job_id_from_submission_output(submission_output)
+
+            now = datetime.datetime.now()
+            sys.stderr.write("{}/{}/{} - {}:{} | submitted job id {} ({})\n".format(now.year, now.month, now.day, now.hour, now.minute, now.second, job_id, submission_output))
+            sys.stderr.flush()
+            time.sleep(self.check_timeout)
+
+            while not self.check_job_finished(job_id):
+                time.sleep(self.check_timeout)
+
 #        cp_res = os.system(cmd + " > /dev/null")
 #        if cp_res != 0:
 #            print "Error while executing:", cmd
@@ -602,7 +631,58 @@ def batch_submission(self, list_of_structures, calc, indices,
                 pass
 
         return results
-        
+
+    def get_job_id_from_submission_output(self, output):
+        """
+        GET THE JOB ID
+
+        Retreive the job id from the output of the submission. 
+        This depends on the software employed. It works for slurm.
+
+        Returns None if the output contains an error
+        """
+
+        try:
+            id = output.split()[-1]
+            return id
+        except:
+            print("Error, expected a standard output, but the result of the submission was: {}".format(output))
+            return None
+        
+    def check_job_finished(self, job_id, verbose = True):
+        """
+        Check if the job identified by the job_id is finished
+
+        Parameters
+        ----------
+            job_id : string
+                The string that identifies uniquely the job
+        """
+
+        status, output = self.ExecuteCMD("squeue -u $USER", False, return_output = True, on_cluster = True, )
+        lines = output.split("\n")
+        if len(lines):
+            for l in lines:
+                data = l.strip().split()
+                if data[0] == job_id:
+                    if verbose:
+                        now = datetime.datetime.now()
+                        sys.stderr.write("{}/{}/{} - {}:{}:{} | job {} still running\n".format(now.year, now.month, now.day, now.hour, now.minute, now.second, job_id))
+                        sys.stderr.flush()
+                    return False
+            
+            # If I'm here it means I did not find the job, but the command returned at least 1 line (so it was correctly executed).
+            if verbose:
+                now = datetime.datetime.now()
+                sys.stderr.write("{}/{}/{} - {}:{}:{} | job {} finished\n".format(now.year, now.month, now.day, now.hour, now.minute, now.second, job_id))
+                sys.stderr.flush()
+            return True
+        if verbose:
+            now = datetime.datetime.now()
+            sys.stderr.write("{}/{}/{} - {}:{}:{} | error while interrogating the cluster for job {}\n".format(now.year, now.month, now.day, now.hour, now.minute, now.second, job_id))
+            sys.stderr.flush()
+        return False
+
 
     def run_atoms(self, ase_calc, ase_atoms, label="ESP", 
                   in_extension = ".pwi", out_extension=".pwo",

Modules/Ensemble.py

@@ -82,6 +82,7 @@
     Bohr = 1/__A_TO_BOHR__
     __RyToK__ = 157887.32400374097
 
+__GPa__ = 14710.50763554043
 
 __DEBUG_RHO__ = False
 
@@ -817,6 +818,73 @@ def save_enhanced_xyz(self, filename, append_mode = True, stress_key = "virial",
         # Force other processors to wait for the master
         CC.Settings.barrier()
 
+    def save_raw(self, root_directory, type_dict = None):
+        """
+        Save the ensemble as a set of raw files.
+
+        This is the default format for training with deepmd
+
+        Parameters
+        ----------
+            filename : string
+                The directory on which to save the ensemble. If it does not exist, it is create.
+                NOTE: this will overwrite any other ensemble saved in raw format in that directory
+            type_dict : dict
+                The dictionary between integers and atomic types. If not provided, it is generated on the spot and returned.
+        
+        Returns
+        -------
+            type_dict : dict
+                The dictionary of the parameters
+        """
+        nat = self.current_dyn.structure.N_atoms * np.prod(self.current_dyn.GetSupercell())
+
+        if type_dict is None:
+            atm = np.unique(self.current_dyn.structure.atoms)
+            type_dict = {x : i for i, x in enumerate(atm)}
+        
+        inv_dict = {i : x for x, i in type_dict.items()}
+
+
+        # Save only if the current processor is the master
+        if Parallel.am_i_the_master():
+            if not os.path.exists(root_directory):
+                os.makedirs(root_directory)
+            
+            if not os.path.isdir(root_directory):
+                raise IOError("Error, save_raw expects a directory, but '{}' is not a directory.".format(root_directory))
+
+            # Save the energies
+            np.savetxt(os.path.join(root_directory, "energy.raw"), self.energies * Rydberg)
+
+            # Save the positions
+            np.savetxt(os.path.join(root_directory, "coord.raw"), self.xats.reshape((self.N, 3 * nat)))
+
+            # Save the box
+            np.savetxt(os.path.join(root_directory, "box.raw"), np.tile(self.current_dyn.structure.unit_cell.ravel(), (self.N, 1)))
+
+            # Save the forces
+            np.savetxt(os.path.join(root_directory, "force.raw"), self.forces.reshape((self.N, 3*nat)) * Rydberg)
+
+            # Save the stress
+            np.savetxt(os.path.join(root_directory, "virial.raw"), self.stresses.reshape((self.N, 9)) * __GPa__ * 10000)
+
+            # Save the types
+            ss = self.current_dyn.structure.generate_supercell(self.current_dyn.GetSupercell())
+
+            with open(os.path.join(root_directory, "type_map.raw"), "w") as fp:
+                line = " ".join([inv_dict[x] for x in np.arange(len(type_dict))])
+                fp.write(line + "\n")
+
+            with open(os.path.join(root_directory, "type.raw"), "w") as fp:
+                line = " ".join([str(type_dict[x]) for x in ss.atoms])
+                fp.write(line + "\n")
+
+                
+        
+        # Force other processors to wait for the master
+        CC.Settings.barrier()
+
 
 
 

tests/test_simple_relax/test_relax_other.py

@@ -69,7 +69,7 @@ def test_simple_relax(verbose = False):
     minim.min_step_struc = 0.5
     minim.meaningful_factor = 1e-10
     minim.init()
-    minim.run()
+    minim.run(verbose = 0)
     minim.finalize()
 
     # Check the differences in the atomic positions