Skip to content

Commit e7e4fd2

Browse files
mesonepigrecomesonepigreco
authored
Merge branch 'new_cluster' into new_cluster_v2
2 parents 3785d10 + 856342e commit e7e4fd2

1 file changed

Lines changed: 205 additions & 102 deletions

File tree

‎Modules/Cluster.py‎

Lines changed: 205 additions & 102 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import subprocess
55
import threading
66
import copy
7+
import time, datetime
78

89
__DIFFLIB__ = False
910
try:
@@ -163,7 +164,7 @@ def __init__(self, hostname=None, pwd=None, extra_options="", workdir = "",
163164
self.workdir=r""
164165
self.submit_command="sbatch --wait"
165166
self.submit_name="SBATCH"
166-
self.terminal="#!/bin/bash"
167+
self.terminal="/bin/bash"
167168
self.v_nodes="-N "
168169
self.use_nodes = True
169170
self.v_cpu="-n "
@@ -240,6 +241,7 @@ def __init__(self, hostname=None, pwd=None, extra_options="", workdir = "",
240241
self.additional_script_parameters = None
241242

242243

244+
243245
# Allow to setup additional custom extra parameters
244246
self.custom_params = {}
245247

@@ -396,7 +398,201 @@ def CheckCommunication(self):
396398
return False
397399

398400
return True
399-
401+
402+
def create_submission_script(self, labels):
403+
"""
404+
CREATE THE SUBMISSION SCRIPT
405+
===========================================
406+
407+
This is a function that is general and does not depend on the specific
408+
calculator. It is usefull to create the header of the submission script.
409+
410+
Parameters
411+
----------
412+
labels : list
413+
It is a list of the labels of the calculations to be done.
414+
415+
Returns
416+
-------
417+
submission_header : string
418+
The text of the submission header.
419+
"""
420+
421+
# prepare the submission script
422+
submission = "#!" + self.terminal + "\n"
423+
424+
425+
426+
# Add the submission options
427+
if self.use_nodes:
428+
submission += "#%s %s%d\n" % (self.submit_name, self.v_nodes, self.n_nodes)
429+
if self.use_cpu:
430+
submission += "#%s %s%d\n" % (self.submit_name, self.v_cpu, self.n_cpu)
431+
if self.use_time:
432+
submission += "#%s %s%s\n" % (self.submit_name, self.v_time, self.time)
433+
if self.use_account:
434+
submission += "#%s %s%s\n" % (self.submit_name, self.v_account, self.account_name)
435+
if self.use_memory:
436+
submission += "#%s %s%s\n" % (self.submit_name, self.v_memory, self.ram)
437+
if self.use_partition:
438+
submission += "#%s %s%s\n" % (self.submit_name, self.v_partition, self.partition_name)
439+
440+
# Append the additional parameters
441+
for add_parameter in self.custom_params:
442+
if self.custom_params[add_parameter] is None:
443+
submission += "#{} --{}\n".format(self.submit_name, add_parameter)
444+
else:
445+
submission += "#{} --{}={}\n".format(self.submit_name, add_parameter, self.custom_params[add_parameter])
446+
447+
448+
# Add the set -x option
449+
if self.add_set_minus_x:
450+
submission += "set -x\n"
451+
452+
# Add the loading of the modules
453+
submission += self.load_modules + "\n"
454+
455+
# Go to the working directory
456+
submission += "cd " + self.workdir + "\n"
457+
458+
# If any, apply the extra text before and after the calculation
459+
other_input = ""
460+
other_output = ""
461+
if (self.additional_script_parameters is not None):
462+
other_input, other_output = self.additional_script_parameters(labels)
463+
464+
submission += other_input
465+
466+
# Use the xargs trick
467+
#submission += "xargs -d " + r"'\n'" + " -L1 -P%d -a %s -- bash -c\n" % (n_togheder,
468+
for i, lbl in enumerate(labels):
469+
submission += self.get_execution_command(lbl)
470+
471+
472+
submission += other_output
473+
474+
return submission
475+
476+
def get_execution_command(self, label):
477+
"""
478+
GET THE EXECUTION COMMAND
479+
=========================
480+
481+
Return the command used in the submission script to actually execute the calculation.
482+
483+
Parameters
484+
----------
485+
label : string
486+
The label of the calculation
487+
488+
Returns
489+
-------
490+
commnad : string
491+
The command to be appended to the submission script
492+
"""
493+
494+
# Get the MPI command replacing NPROC
495+
new_mpicmd = self.mpi_cmd.replace("NPROC", str(self.n_cpu))
496+
497+
# Replace the NPOOL variable and the PREFIX in the binary
498+
binary = self.binary.replace("NPOOL", str(self.n_pool)).replace("PREFIX", label)
499+
500+
501+
tmt_str = ""
502+
if self.use_timeout:
503+
tmt_str = "timeout %d " % self.timeout
504+
return "%s%s %s\n" % (tmt_str, new_mpicmd, binary)
505+
506+
def prepare_input_file(self, structure, calc, label):
507+
"""
508+
PREPARE THE INPUT FILE
509+
======================
510+
511+
This is specific for quantum espresso and must be inherit and replaced for
512+
other calculators.
513+
514+
This crates the input file and copy it in the working directory.
515+
516+
Parameters
517+
----------
518+
structure : CellConstructor.Structure.Structure
519+
The atomic structure on which to run the calculation
520+
calc : the ASE or CellConstructor calculator.
521+
In this case, it works with quantum espresso
522+
label : string
523+
The unique name of this calculation
524+
"""
525+
526+
# Prepare the input file
527+
atm = structure.get_ase_atoms()
528+
atm.set_calculator(calc)
529+
ase.io.write("%s/%s.pwi"% (self.local_workdir, label),
530+
atm, **calc.parameters)
531+
532+
533+
534+
# First of all clean eventually input/output file of this very same calculation
535+
cmd = self.sshcmd + " %s 'rm -f %s/%s%s %s/%s%s'" % (self.hostname,
536+
self.workdir, label, ".pwi",
537+
self.workdir, label, ".pwo")
538+
self.ExecuteCMD(cmd, False)
539+
# cp_res = os.system(cmd + " > /dev/null")
540+
# if cp_res != 0:
541+
# print "Error while executing:", cmd
542+
# print "Return code:", cp_res
543+
# sys.stderr.write(cmd + ": exit with code " + str(cp_res) + "\n")
544+
#
545+
# Copy the file into the cluster
546+
cmd = self.scpcmd + " %s/%s%s %s:%s/" % (self.local_workdir, label,
547+
".pwi", self.hostname,
548+
self.workdir)
549+
cp_res = self.ExecuteCMD(cmd, False)
550+
if not cp_res:
551+
print ("Error while executing:", cmd)
552+
print ("Return code:", cp_res)
553+
sys.stderr.write(cmd + ": exit with code " + str(cp_res) + "\n")
554+
return cp_res
555+
#cp_res = os.system(cmd + " > /dev/null")
556+
557+
558+
def submit(self, script_location):
559+
"""
560+
SUBMIT THE CALCULATION
561+
======================
562+
563+
Submit the calculation. Compose the command into a cmd variable, then submit it through:
564+
565+
.. code ::
566+
567+
return self.ExecuteCMD(cmd, True, return_output=True)
568+
569+
570+
571+
Parameters
572+
----------
573+
script_localtion : string
574+
Path to the submission script inside the cluster.
575+
576+
Results
577+
-------
578+
success : bool
579+
Result of the execution of the submission command.
580+
It is what returned from self.ExecuteCMD(cmd, False)
581+
"""
582+
583+
cmd = "{ssh} {host} '{submit_cmd} {script}'"
584+
if self.use_active_shell:
585+
cmd = "{ssh} {host} -t '{shell} --login -c \"{submit_cmd} {script}\"'".format(shell = self.terminal)
586+
587+
588+
cmd = cmd.format(ssh = self.sshcmd, host = self.hostname,
589+
submit_cmd = self.submit_name, script = script_location)
590+
591+
#cmd = self.sshcmd + " %s '%s %s/%s.sh'" % (self.hostname, self.submit_command,
592+
# self.workdir, label+ "_" + str(indices[0]))
593+
594+
return self.ExecuteCMD(cmd, True, return_output=True)
595+
400596
def batch_submission(self, list_of_structures, calc, indices,
401597
in_extension, out_extension,
402598
label = "ESP", n_togheder=1):
@@ -446,7 +642,6 @@ def batch_submission(self, list_of_structures, calc, indices,
446642
# Prepare the input atoms
447643
app_list = ""
448644
new_ncpu = self.n_cpu * n_togheder
449-
new_mpicmd = self.mpi_cmd.replace("NPROC", str(self.n_cpu))
450645
results = [None] * N_structs
451646
submitted = []
452647
submission_labels = []
@@ -455,44 +650,10 @@ def batch_submission(self, list_of_structures, calc, indices,
455650
lbl = label + "_" + str(indices[i])
456651
submission_labels.append(lbl)
457652

458-
atm = list_of_structures[i].get_ase_atoms()
459-
atm.set_calculator(calc)
460-
ase.io.write("%s/%s%s"% (self.local_workdir, lbl, in_extension),
461-
atm,**calc.parameters)
462-
463-
464-
# Add the file in the applist
465-
binary = self.binary.replace("NPOOL", str(self.n_pool)).replace("PREFIX", lbl)
466-
467-
468-
# First of all clean eventually input/output file of this very same calculation
469-
cmd = self.sshcmd + " %s 'rm -f %s/%s%s %s/%s%s'" % (self.hostname,
470-
self.workdir, lbl, in_extension,
471-
self.workdir, lbl, out_extension)
472-
self.ExecuteCMD(cmd, False)
473-
# cp_res = os.system(cmd + " > /dev/null")
474-
# if cp_res != 0:
475-
# print "Error while executing:", cmd
476-
# print "Return code:", cp_res
477-
# sys.stderr.write(cmd + ": exit with code " + str(cp_res) + "\n")
478-
#
479-
# Copy the file into the cluster
480-
cmd = self.scpcmd + " %s/%s%s %s:%s/" % (self.local_workdir, lbl,
481-
in_extension, self.hostname,
482-
self.workdir)
483-
cp_res = self.ExecuteCMD(cmd, False)
484-
485-
#cp_res = os.system(cmd + " > /dev/null")
486-
if not cp_res:
487-
print ("Error while executing:", cmd)
488-
print ("Return code:", cp_res)
489-
sys.stderr.write(cmd + ": exit with code " + str(cp_res) + "\n")
653+
# Create the input file and copy it into the cluster
654+
if not self.prepare_input_file(list_of_structures[i], calc, lbl):
490655
continue
491-
492-
tmt_str = ""
493-
if self.use_timeout:
494-
tmt_str = "timeout %d " % self.timeout
495-
app_list += "%s%s %s\n" % (tmt_str, new_mpicmd, binary)
656+
496657
submitted.append(i)
497658

498659
# Save the app list and copy it to the destination
@@ -512,59 +673,8 @@ def batch_submission(self, list_of_structures, calc, indices,
512673
# sys.stderr.write(cmd + ": exit with code " + str(cp_res) + "\n")
513674
# return results #[None] * N_structs
514675

515-
516-
# prepare the submission script
517-
submission = self.terminal + "\n"
518-
519-
# Add the submission options
520-
if self.use_nodes:
521-
submission += "#%s %s%d\n" % (self.submit_name, self.v_nodes, self.n_nodes)
522-
if self.use_cpu:
523-
submission += "#%s %s%d\n" % (self.submit_name, self.v_cpu, new_ncpu)
524-
if self.use_time:
525-
submission += "#%s %s%s\n" % (self.submit_name, self.v_time, self.time)
526-
if self.use_account:
527-
submission += "#%s %s%s\n" % (self.submit_name, self.v_account, self.account_name)
528-
if self.use_memory:
529-
submission += "#%s %s%s\n" % (self.submit_name, self.v_memory, self.ram)
530-
if self.use_partition:
531-
submission += "#%s %s%s\n" % (self.submit_name, self.v_partition, self.partition_name)
532-
533-
# Append the additional parameters
534-
for add_parameter in self.custom_params:
535-
adder_string = "--{}".format(add_parameter)
536-
if add_parameter.startswith("-"):
537-
adder_string = add_parameter
538-
539-
if self.custom_params[add_parameter] is None:
540-
submission += "#{} {}\n".format(self.submit_name, adder_string)
541-
else:
542-
submission += "#{} {}={}\n".format(self.submit_name, adder_string, self.custom_params[add_parameter])
543-
544-
545-
# Add the set -x option
546-
if self.add_set_minus_x:
547-
submission += "set -x\n"
548-
549-
# Add the loading of the modules
550-
submission += self.load_modules + "\n"
551-
552-
# Go to the working directory
553-
submission += "cd " + self.workdir + "\n"
554-
555-
# If any, apply the extra text before and after the calculation
556-
other_input = ""
557-
other_output = ""
558-
if self.additional_script_parameters is not None:
559-
other_input, other_output = self.additional_script_parameters(submission_labels)
560676

561-
submission += other_input
562-
563-
# Use the xargs trick
564-
#submission += "xargs -d " + r"'\n'" + " -L1 -P%d -a %s -- bash -c\n" % (n_togheder,
565-
submission += app_list
566-
567-
submission += other_output
677+
submission = self.create_submission_script(submission_labels)
568678

569679
# Copy the submission script
570680
sub_fpath = "%s/%s.sh" % (self.local_workdir, label + "_" + str(indices[0]))
@@ -582,11 +692,10 @@ def batch_submission(self, list_of_structures, calc, indices,
582692

583693

584694
# Run the simulation
585-
cmd = self.sshcmd + " %s '%s %s/%s.sh'" % (self.hostname, self.submit_command,
586-
self.workdir, label+ "_" + str(indices[0]))
587-
status, submission_output = self.ExecuteCMD(cmd, True, return_output = True)
588695

589-
# If the command for submission is non blocking, we need to check periodically wether the calculation has been completed
696+
sub_script_loc = os.path.join(self.workdir, label + "_" + str(indices[0]) + ".sh")
697+
cp_res, submission_output = self.submit(sub_script_loc)
698+
590699
if self.nonblocking_command:
591700
job_id = self.get_job_id_from_submission_output(submission_output)
592701

@@ -598,12 +707,6 @@ def batch_submission(self, list_of_structures, calc, indices,
598707
while not self.check_job_finished(job_id):
599708
time.sleep(self.check_timeout)
600709

601-
# cp_res = os.system(cmd + " > /dev/null")
602-
# if cp_res != 0:
603-
# print "Error while executing:", cmd
604-
# print "Return code:", cp_res
605-
# sys.stderr.write(cmd + ": exit with code " + str(cp_res))
606-
#
607710

608711
# Collect the output back
609712
for i in submitted:

0 commit comments

Comments
 (0)