def get_initial_atoms_from_fwid(fwid):
'''
Function written by Kevin
This function will return the `ase.Atoms` object given to a rocket.
Args:
fwid Integer indicating the FireWorks ID that you're trying to get
the atoms object for
Returns
atoms `ase.Atoms` instance from the Firework you provided
'''
# Get the firework object, which will have all the information we'll need
try:
lpad = get_launchpad()
fw = lpad.get_fw_by_id(fwid)
# Close the Mongo connection for cleanliness' sake
finally:
lpad.fireworks.database.client.close()
# Get the Firework task that was meant to convert the original hexstring to
# a trajectory file. We'll get the original atoms from this task (in
# hexstring format). Note that over the course of our use, we have had
# different names for these FireWorks tasks, so we check for them all.
function_names_of_hex_encoders = set([
'vasp_functions.hex_to_file',
'fireworks_helper_scripts.atoms_hex_to_file',
'fireworks_helper_scripts.atomsHexToFile'
])
trajhexes = [
task['args'][1] for task in fw.spec['_tasks']
if task.get('func', '') in function_names_of_hex_encoders
]
# If there was not one task, then ignore it because it's probably some old
# firework that isn't formatted correctly for us
if len(trajhexes) != 1:
warnings.warn('%i does not have any trajhexes. Moving on.' % fwid,
RuntimeWarning)
return None
# If the trajhex is empty, then it probably bugged out. Ignore it.
elif not trajhexes[0]:
warnings.warn('%i has an empty trajhex. Moving on.' % fwid,
RuntimeWarning)
return None
# Use our vanilla trajhex opener
try:
atoms = decode_trajhex_to_atoms(trajhexes[0])
# Some hexes are weird and don't open properly for a bunch of reasons. Ignore them.
except:
warnings.warn('%i did not open properly. Moving on.' % fwid,
RuntimeWarning)
return None
return atoms
def get_state_by_launch_id(launch_id):
"""
Obtain the state, or job status, of the firework by the given launch id
Parameters:
launch_id (int): launch id of the firework
Returns:
state (string): job status (e.g. COMPLETED, FIZZLED, READY, DEFUSED)
"""
lpad = get_launchpad()
launch_fw = list(lpad.launches.find({"launch_id": launch_id}))[0]
fw_id = launch_fw["fw_id"]
fw = lpad.get_fw_by_id(fw_id)
state = fw.state
return state
def get_n_jobs_to_submit(user_name, quota=300):
'''
This function helps you figure out how many jobs you should submit by
calculating the difference between your quota of jobs in queue and the
number of jobs you actually have in queue.
Args:
user_name String indicating your user name in FireWorks
quota Integer indicating the number of jobs you want to
have running/ready in FireWorks at a given time.
Returns:
n_jobs_to_submit The integer difference between the number of jobs
you have running and the number you want running.
'''
lpad = get_launchpad()
query = {'name.user': user_name, 'state': {'$in': ['READY', 'RUNNING']}}
n_jobs_running = lpad.fireworks.find(query).count()
n_jobs_to_submit = quota - n_jobs_running
return n_jobs_to_submit
os.system("rsync -azqP --max-size=100M gilgamesh:%s /tmp/%s" % (launch.launch_dir, launch_id))
elif 'arjuna' in cluster:
os.system("rsync -azqP --max-size=100M arjuna:%s /tmp/%s" % (launch.launch_dir, launch_id))
elif 'Cori' in cluster or '/global/project/projectdirs/m2755/' in launch.launch_dir:
launch_dir = launch.launch_dir
if '/global/u2/z/zulissi' in launch_dir:
launch_dir = '/global/project/projectdirs/m2755/fireworks_zu/fireworks/' + launch_dir[30:]
os.system("rsync -azqP --max-size=100M %s /tmp/%s" % (launch_dir, launch_id))
else:
print('unknown cluster!: %s' % cluster)
print(launch.fw_id)
launch_dir_folder = launch.launch_dir.split('/')[-1]
os.system("(cd /tmp/%s/%s && tar -czf %s *)" % (launch_id, launch_dir_folder, backup_loc))
os.system("rm -r /tmp/%s" % launch_id)
lpad = get_launchpad()
all_fw_ids = lpad.get_fw_ids({"state": "COMPLETED"})
launch_list = []
for fwid in tqdm(all_fw_ids):
fw = lpad.get_fw_by_id(fwid)
for launch in fw.launches:
launch_list.append(launch.launch_id)
with multiprocess.Pool(4) as pool:
list(tqdm(pool.imap(fetch_tar_file, launch_list, chunksize=20), total=len(launch_list)))