def get_slurm_queues():
'''
Get a list of all available queues to submit a job to.
**Returns**
avail_queues: *list, str*
A list of available queues by name.
'''
sinfo_path = which("sinfo")
assert sinfo_path is not None,\
"Error - Unable to find sinfo in PATH."
p = subprocess.Popen([sinfo_path],
shell=False,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE)
all_queues = p.stdout.read().decode("utf-8").strip()
if all_queues == '':
close_pipes(p)
return []
all_queues = all_queues.split("\n")[1:]
all_queues = [q.split()[0] for q in all_queues if q.split()[1] == 'up']
all_queues = list(set(all_queues))
close_pipes(p)
return [q if "*" not in q else q.replace("*", "") for q in all_queues]
def get_nbs_queues():
'''
Get a list of all available queues to submit a job to.
**Returns**
avail_queues: *list, str*
A list of available queues by name.
'''
qlist_path = which("qlist")
p = subprocess.Popen([qlist_path],
shell=False,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE)
all_queues = p.stdout.read().decode("utf-8").strip().split('\n')[:-1]
all_queues = [a.split() for a in all_queues]
all_queues = [a[0] for a in all_queues if len(a) > 1]
close_pipes(p)
return [
a.lower() for a in all_queues
if a.lower() not in ["queue", "name", ""]
]
def get_ovito_obj(version="2.9.0"):
'''
This function returns the ovito object. Note, currently the code below
only works on version 2.9.0.
'''
ovito_path = which("ovitos")
# Determine version
ovito_pipe = subprocess.Popen([ovito_path, "-v"],
shell=False,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE)
stdout = str(ovito_pipe.stdout.read().decode("utf-8").strip())
assert "Ovito" in stdout,\
"Error - Unable to access Ovito. Please ensure it is in your PATH \
environment variable!"
assert version in stdout,\
"Error - Incorrect Ovito version! It should be %s, but is %s."\
% (version, stdout.strip().split()[-1])
close_pipes(ovito_pipe)
return ovito_path
def submit_job(name, job_to_submit, **kwargs):
'''
Code to submit a simulation to the specified queue and queueing system.
**Parameters**
name: *str*
Name of the job to be submitted to the queue.
job_to_submit: *str*
String holding code you wish to submit.
queue: *str, optional*
What queue to run the simulation on (queueing system dependent).
walltime: *str, optional*
How long to post the job on the queue for in d-h:m:s where d are
days, h are hours, m are minutes, and s are seconds. Default is
for 30 minutes (00:30:00).
cpus_per_task: *int, optional*
How many processors to run the simulation on. Note, the actual
number of cores mpirun will use is nprocs * ntasks.
ntasks: *int, optional*
How many processors to run the simulation on. Note, the actual
number of cores mpirun will use is nprocs * ntasks.
nodes: *int, optional*
How many nodes to run the simulation on.
sub_flag: *str, optional*
Additional strings/flags/arguments to add at the end when we
submit a job using sbatch. That is: sbatch demo.slurm sub_flag.
unique_name: *bool, optional*
Whether to force the requirement of a unique name or not. NOTE! If
you submit simulations from the same folder, ensure that this is
True lest you have a redundancy problem! To overcome said issue,
you can set redundancy to True as well (but only if the simulation
is truly redundant).
outfile_name: *str, optional*
Whether to give a unique output file name, or one based on the sim
name.procs
allocation: *str, optional*
The SLURM allocation to submit the job to.
jobarray: *str, optional*
If specified, instead of indicating a range for job arrays, we
will use these specific values. For example,
jobarray=1,2,4,5 would submit jobs, but skip the 3rd
index by name.
gpu: *int, optional*
How many GPUs to use, if submitting to a GPU node.
redundancy: *bool, optional*
With redundancy on, if the job is submitted and unique_name is on,
then if another job of the same name is running, a pointer to that
job will instead be returned.
**Returns**
job_obj: :class:`squid.jobs.slurm.Job`
A Job object.
'''
# Store the defaults
params = {
"queue": "shared",
"ntasks": 1,
"cpus_per_task": 1,
"nodes": 1,
"walltime": "00:30:00",
"sub_flag": "",
"unique_name": True,
"redundancy": False,
"allocation": None,
"gpu": None,
"jobarray": None,
"sandbox": None,
"outfile_name": None,
"email": None,
}
AVAIL_GPU_PARTS = ["unlimited", "gpuk80", "gpup100", "debugger"]
## We used to check to see if any pointless values were being passed, but
## we found it was easier to simply just ignore those and not crash.
## Ensure we are passing only the above
#for key, value in kwargs.items():
# assert key in params,\
# "Error - Unknown variable (%s) passed to slurm.submit_job." % key
params.update(kwargs)
# Ensure variables of correct types
param_types = {
"queue": lambda s: str(s).strip(),
"ntasks": int,
"cpus_per_task": int,
"nodes": int,
"unique_name": bool,
"redundancy": bool,
"walltime": lambda s: str(s).strip(),
"sub_flag": lambda s: str(s).strip()
}
for k, f in param_types.items():
params[k] = f(params[k])
# Ensure default values make sense
# Check Queue
slurm_queues = get_slurm_queues()
assert params["queue"] in slurm_queues,\
"Error - Invalid queue (%s) requested. Options: %s"\
% (params["queue"], ", ".join(slurm_queues))
# Check ntasks and nodes
if params["cpus_per_task"] * params["ntasks"] > 24 * params["nodes"]:
print("Warning - You requested %d tasks and %d cpus_per_task. This \
equates to %d nodes on marcc; however, you only requested %d nodes." %
(params["cpus_per_task"], params["ntasks"],
(params["cpus_per_task"] * params["ntasks"] - 1) // 24 + 1,
params["nodes"]))
print("\tWill adjust nodes accordingly...")
params["nodes"] =\
(params["cpus_per_task"] * params["ntasks"] - 1) // 24 + 1
# We need to remove gpu nodes from available nodes on SLURM/MARCC
gpu_flag_slurm = "#SBATCH --exclude=gpu004,gpu005"
# However, if we want to submit to GPU, handle accordingly
if params["gpu"] is not None:
msg = "Error - queue (%s) not available with gpus. Choose one: %s"\
% (params["queue"], ', '.join(AVAIL_GPU_PARTS))
assert params["queue"] in AVAIL_GPU_PARTS, msg
gpu_flag_slurm = "#SBATCH --gres=gpu:%d" % params["gpu"]
# On MARCC we need gpu tasks, and 6 cores per task
params["ntasks"] = params["gpu"]
params["cpus_per_task"] = 6
if params["allocation"] is None:
allocation = ""
else:
allocation = "#SBATCH --account=" + params["allocation"]
# Generate your script
jobarray_outfile_append = ""
job_array_script = ""
if params["jobarray"] is not None:
if isinstance(params["jobarray"], str):
job_array_script = "#SBATCH --array=%s"\
% simplify_numerical_array(params["jobarray"])
else:
job_array_script = "#SBATCH --array=%d-%d"\
% tuple(params["jobarray"])
jobarray_outfile_append = ".a%a"
if params["outfile_name"] is None:
params["outfile_name"] = name + jobarray_outfile_append + ".o%j"
generic_script = '''#!/bin/sh
#SBATCH --job-name="''' + name + '''"
#SBATCH --output="''' + params["outfile_name"] + '''"
#SBATCH --nodes=''' + str(params["nodes"]) + '''
#SBATCH --ntasks=''' + str(params["ntasks"]) + ('''
#SBATCH --cpus-per-task=''' + str(
params["cpus_per_task"]) if params["cpus_per_task"] > 1 else "") + '''
#SBATCH --partition=''' + params["queue"] + '''
#SBATCH --time=''' + params["walltime"] + '''
''' + allocation + '''
''' + gpu_flag_slurm + '''
''' + job_array_script + '''
'''
# Take care of sandboxing if needed
if params["sandbox"] is not None:
raise Exception("Sandbox not implemented in slurm.")
# Add in your script now
generic_script = generic_script +\
"\ndate\n" +\
job_to_submit +\
"\ndate\n\n"
f = open(name + '.slurm', 'w')
f.write(generic_script)
f.close()
# Get a list of all jobs
all_jobs = get_job("RUNNING", detail=0) + get_job("PENDING", detail=0)
job_exists = name in all_jobs
if params["redundancy"] and job_exists:
try:
job_info = get_job(name, detail=1)[0]
except IndexError:
# Job finished in process of submitting and redundancy call
job_to_return = Job(name)
# Attach the redundancy flag
job_to_return.redundancy = True
return job_to_return
job_to_return = Job(job_info[0], job_id=job_info[-1])
job_to_return.redundancy = True
return job_to_return
elif params["unique_name"] and job_exists:
raise Exception("Job with name %s already exists in the queue!" % name)
# Submit job
cmd = 'sbatch %s.slurm %s' % (name, params["sub_flag"])
job_pipe = subprocess.Popen(cmd.split(),
shell=False,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE)
# Get the error message
job_err = job_pipe.stderr.read().decode("utf-8")
# If we figure out redundancy, add it here
# CODE FORE REDUNDANCY IN SLURM
job_id = job_pipe.stdout.read().decode("utf-8")
if "Submitted batch job" not in job_id:
print("\nFailed to submit the job!")
print("--------------- JOB OUTPUT ---------------")
print(job_id)
print("---------------- JOB ERROR ---------------")
print(job_err)
print("---------------------------------")
sys.stdout.flush()
raise Exception()
job_id = job_id.split()[-1].strip()
close_pipes(job_pipe)
return Job(name, job_id=job_id)
def get_orca_obj(parallel=True):
'''
This function will find the orca executable and the corresponding openmpi
executable. It will handle errors accordingly.
**Parameters**
parallel: *bool, optional*
Whether we guarantee the relevant parallel openmpi is setup (True)
or not (False).
**Returns**
orca_path: *str*
The path to the orca executable.
'''
# This is to ensure we read in ORCA correctly
orca_string_id = "An Ab Initio, DFT and Semiempirical electronic structure"
# This is to find the version
version_string_id = "Program Version"
orca_path = which("orca")
# Determine orca version
orca_pipe = subprocess.Popen(
[orca_path, "FAKE_FILE"], shell=False,
stdout=subprocess.PIPE, stderr=subprocess.PIPE)
stdout = str(orca_pipe.stdout.read().decode("utf-8").strip())
# stderr = str(p.stderr.read().decode("utf-8").strip())
assert orca_string_id in stdout,\
"Error - Unable to access Orca. Please ensure it is in your PATH \
environment variable!"
assert version_string_id in stdout,\
"Error - Unable to assess Orca version!"
orca_version = stdout.split(version_string_id)[1].strip().split()[0]
# If running in parallel, ensure we have the correct version of openmpi
ompi_pipe = None
if parallel:
ompi_version_should_be = {
"4.1.2": "3.1",
"4.2.0": "3.1"
}
assert orca_version in ompi_version_should_be,\
"Error - Please contact squid dev. We do not have stored the \
required openmpi version for Orca %s" % orca_version
# Find openmpi
ompi_path = which("mpiexec")
ompi_pipe = subprocess.Popen(
[ompi_path, "--V"], shell=False,
stdout=subprocess.PIPE, stderr=subprocess.PIPE)
stdout = str(ompi_pipe.stdout.read().decode("utf-8").strip())
# stderr = str(p.stderr.read().decode("utf-8").strip())
# Simple check for openmpi
assert "open" in stdout.lower(),\
"Error - Unable to access openmpi. Please ensure it is in your \
PATH environment variable!"
ompi_version = stdout.strip().split("\n")[0].split()[-1]
ompi_version_held = ompi_version_should_be[orca_version]
assert ompi_version.startswith(ompi_version_held),\
"Error - Incorrect openmpi version for the loaded orca version. \
Should be openmpi %s (found %s) for orca %s."\
% (ompi_version_held, ompi_version, orca_version)
close_pipes(orca_pipe)
close_pipes(ompi_pipe)
return orca_path
def get_job(s_flag, detail=0):
'''
Get a list of all jobs currently on your queue. From this, only return
the values that have s_flag in them. The *detail* variable can be used
to specify how much information you want returned.
**Parameters**
s_flag: *str*
A string to parse out job information with.
detail: *int, optional*
The amount of information you want returned.
**Returns**
all_jobs: *list*
Depending on *detail*, you get the following:
- *details* =0: *list, str*
List of all jobs on the queue.
- *details* =1: *list, tuple, str*
List of all jobs on the queue as:
(job name, time run, job status)
- *details* =2: *list, tuple, str*
List of all jobs on the queue as:
(job name,
time run,
job status,
queue,
number of processors)
'''
detail = int(detail)
main_detail = detail
if detail <= 0:
detail = 1
# Get input from jlist as a string
jlist_path = which("jlist")
qlist_path = which("qlist")
jshow_path = which("jshow")
assert jlist_path is not None,\
"Error - Cannot find jlist."
assert qlist_path is not None,\
"Error - Cannot find qlist."
assert jshow_path is not None,\
"Error - Cannot find jshow."
p = subprocess.Popen([jlist_path],
shell=False,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE)
output = p.stdout.read().decode("utf-8")
# Get data from string
pattern = getpass.getuser() +\
'''[\s]+([\S]+)[\s]+([\S]+)[\s]+([\S]+)'''
info = re.findall(pattern, output)
# Get a list of names
names = []
for a in info:
names.append(a[0])
if len(names) > 0:
out_ids = output.split("\n")
out_ids = [
x.split()[0] for x in out_ids
if len(x.split()) > 0 and is_numeric(x.split()[0])
]
info = [tuple(list(i) + [j]) for i, j in zip(info, out_ids)]
# If user wants more information
all_jobs = None
if detail == 3:
close_pipes(p)
all_jobs = [i[-1] for i in info]
elif detail == 2:
for i, a in enumerate(info):
p = subprocess.Popen([jshow_path, a[0]], stdout=subprocess.PIPE)
s = p.stdout.read().decode("utf-8")
serv = s[s.find('Queue name:'):].split()[2].strip()
threads = 1
if "Slot Reservations" in s:
threads = s[s.find('Slot Reservations'):].split()[4]
threads = threads.strip()
info[i] = info[i] + (
serv,
threads,
)
close_pipes(p)
all_jobs = info
if all_jobs is None:
# Return appropriate information
close_pipes(p)
if detail == 1:
all_jobs = info
else:
all_jobs = names
job_indices = [i for i, j in enumerate(all_jobs) if s_flag in " ".join(j)]
chosen_jobs = [all_jobs[i] for i in job_indices]
if main_detail == 0:
return [j[0] for j in chosen_jobs]
else:
return chosen_jobs
def submit_job(name, job_to_submit, **kwargs):
'''
Code to submit a simulation to the specified queue and queueing system.
**Parameters**
name: *str*
Name of the job to be submitted to the queue.
job_to_submit: *str*
String holding code you wish to submit.
queue: *str, optional*
What queue to run the simulation on (queueing system dependent).
walltime: *str, optional*
How long to post the job on the queue for in d-h:m:s where d are
days, h are hours, m are minutes, and s are seconds. Default is
for 30 minutes (00:30:00).
nprocs: *int, optional*
How many processors to run the simulation on. Note, the actual
number of cores mpirun will use is procs * ntasks.
sub_flag: *str, optional*
Additional strings/flags/arguments to add at the end when we
submit a job using jsub. That is: jsub demo.nbs sub_flag.
unique_name: *bool, optional*
Whether to force the requirement of a unique name or not. NOTE! If
you submit simulations from the same folder, ensure that this is
True lest you have a redundancy problem! To overcome said issue,
you can set redundancy to True as well (but only if the simulation
is truly redundant).
outfile_name: *str, optional*
Whether to give a unique output file name, or one based on the sim
name.procs
xhosts: *str* or *list, str, optional*
Which cpu to submit the job to.
email: *str, optional*
An email address for sending job information to.
priority: *int, optional*
What priority to give the submitted job.
sandbox: *bool, optional*
Whether to sandbox the job or not.
redundancy: *bool, optional*
With redundancy on, if the job is submitted and unique_name is on,
then if another job of the same name is running, a pointer to that
job will instead be returned.
**Returns**
job_obj: :class:`squid.jobs.nbs.Job`
A Job object.
'''
# Store the defaults
params = {
"queue": "shared",
"nprocs": 1,
"sub_flag": "",
"unique_name": True,
"redundancy": False,
"outfile_name": None,
"xhosts": None,
"email": None,
"priority": None,
"sandbox": True,
"walltime": None,
}
## Ensure we are passing only the above
#for key, value in kwargs.items():
# assert key in params,\
# "Error - Unknown variable (%s) passed to nbs.submit_job." % key
params.update(kwargs)
if params["walltime"] is not None:
print("Warning - Walltime is not handled in NBS yet. Your job will \
have the default time of the given queue.")
# Ensure variables of correct types
param_types = {
"queue": lambda s: str(s).strip(),
"nprocs": int,
"sub_flag": lambda s: str(s).strip(),
"unique_name": bool,
"redundancy": bool,
"sandbox": bool,
}
for k, f in param_types.items():
params[k] = f(params[k])
# Ensure default values make sense
# Check Queue
nbs_queues = get_nbs_queues()
assert params["queue"] in nbs_queues,\
"Error - Invalid queue (%s) requested. Options: %s"\
% (params["queue"], ", ".join(nbs_queues))
jsub_path = which("jsub")
assert jsub_path is not None,\
"Error - Unable to find jsub path!"
# Deal with variables accordingly
if params["xhosts"] is not None:
if isinstance(params["xhosts"], str):
xhosts = "##NBS-xhost: \"%s\"" % params["xhosts"]
elif is_array(params["xhosts"]):
xhosts = "##NBS-xhost: " +\
", ".join(map(lambda x: '"' + x + '"', params["xhosts"]))
else:
raise Exception("xhosts has been passed oddly!")
else:
xhosts = ""
# Generate your script
generic_script = '''#!/bin/sh
##NBS-name: ''' + name + '''
##NBS-nproc: ''' + params["nprocs"] + '''
##NBS-queue: ''' + params["queue"] + '''
''' + ["", "##NBS-unique: yes"][int(params["unique_name"])] + '''
''' + xhosts
# If emailing, set here
if params["email"] is not None:
generic_script += "##NBS-email: " + params["email"] + "\n"
# If priority is set, add it
if params["priority"] is not None:
if int(params["priority"]) > 255:
params["priority"] = 255
if int(params["priority"]) < 1:
params["priority"] = 1
generic_script += "##NBS-priority: " + str(params["priority"]) + "\n"
# Take care of sandboxing if needed
if params["sandbox"] is not None:
generic_script = generic_script + '''
##NBS-fdisk: 8192
##NBS-fswap: 8192
##NBS-sandbox: yes
##NBS-tmp-sandbox: yes
'''
for sb_in in params["sandbox"][0]:
generic_script = generic_script + '''
##NBS-input: ''' + sb_in
for sb_out in params["sandbox"][1]:
generic_script = generic_script + '''
##NBS-output: ''' + sb_out + ''' -overwrite'''
generic_script = generic_script + "\n\n"
# Add in your script now
generic_script = generic_script + "\ndate\n" + job_to_submit + "\ndate\n\n"
f = open(name + '.nbs', 'w')
f.write(generic_script)
f.close()
# Submit job
cmd = "%s %s.nbs %s" % (jsub_path, name, params["sub_flag"])
job_pipe = subprocess.Popen(cmd.split(),
shell=False,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE)
job_err = job_pipe.stderr.read().decode("utf-8")
if params["redundancy"] and "+notunique:" in job_err:
try:
job_info = get_job(name, detail=1)[0]
except IndexError:
# Job finished in process of submitting and redundancy call
job_to_return = Job(name)
# Attach the redundancy flag
job_to_return.redundancy = True
return job_to_return
job_to_return = Job(job_info[0], job_id=job_info[-1])
job_to_return.redundancy = True
close_pipes(job_pipe)
return job_to_return
elif "+notunique:" in job_err:
raise Exception("Job with name %s already exists in the queue!" % name)
job_id_str = job_pipe.stdout.read().decode("utf-8")
if "submitted to queue" not in job_id_str:
print("\nFailed to submit the job!")
print("--------------- JOB OUTPUT ---------------")
print(job_id_str)
print("---------------- JOB ERROR ---------------")
print(job_err)
print("---------------------------------")
sys.stdout.flush()
raise Exception()
try:
job_id = job_id_str.split("submitted to queue")[0].split()[-1][2:-1]
except IndexError:
print("ERROR - job_id_str is:")
print(job_id_str)
print("Defaulting to None, should still work... FIX!")
job_id = None
close_pipes(job_pipe)
return Job(name, job_id=job_id)
def get_lmp_obj(parallel=True):
'''
This function will find the lmp executable and a corresponding mpi
executable. It will handle errors accordingly.
**Parameters**
parallel: *bool, optional*
Whether to get corresponding mpiexec info or not.
**Returns**
lmp_path: *str*
Path to a lammps executable.
mpi_path: *str*
Path to an mpi executable.
'''
# If running in parallel, ensure we have mpi
mpi_path = None
if parallel:
mpi_path = which("mpiexec")
p = subprocess.Popen([mpi_path, "-h"],
shell=False,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE)
stdout = str(p.stdout.read().decode("utf-8").strip())
close_pipes(p)
# Simple check for openmpi
assert "mpi" in stdout.lower(),\
"Error - Unable to access mpiexec. Please ensure it is in your \
PATH environment variable!"
# First, look for lmp_X in order of common names
lmp_path = None
common_names = ["lmp_mpi", "lmp_serial", "lmp_smrff"]
lmp_string_id = "Large-scale Atomic/Molecular Massively Parallel Simulator"
for name in common_names:
if which(name) is not None:
# Check stuff
if mpi_path is not None:
cmd = [mpi_path, "-n", "1", which(name), "-h"]
else:
cmd = [which(name), "-h"]
p = subprocess.Popen(cmd,
shell=False,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE)
stdout = str(p.stdout.read().decode("utf-8").strip())
if lmp_string_id not in stdout:
close_pipes(p)
continue
else:
# If it works, then save it
close_pipes(p)
lmp_path = which(name)
break
assert lmp_path is not None,\
"Error - Unable to find lmp executable. Please ensure it is \
in your PATH environment variable!"
return lmp_path, mpi_path