def test_copy(self):
"""Test that we can produce a copy of a Workflow but that the copy
has unique fw_ids.
"""
fws = []
for i in range(5):
fw = Firework([PyTask(func="print", args=[i])], fw_id=i, name=i)
fws.append(fw)
wf = Workflow(fws, links_dict={0: [1, 2, 3], 1: [4], 2: [4]})
wf_copy = Workflow.from_wflow(wf)
# now we compare to the original to make sure dependencies are same.
# have to do gymnastics because ids will NOT be the same
# but names are retained
for fw in wf_copy.fws:
children = wf_copy.links.get(fw.fw_id, list())
orig_id = fw.name
orig_children = wf.links.get(orig_id, list())
for child_id, orig_child_id in zip(children, orig_children):
self.assertEqual(orig_child_id, wf_copy.id_fw[child_id].name)
def test_append_wf(self):
fw1 = Firework([UpdateSpecTask()])
fw2 = Firework([ModSpecTask()])
self.lp.add_wf(Workflow([fw1, fw2]))
self.assertEqual(self.lp.fireworks.count(), 2)
launch_rocket(self.lp, self.fworker)
launch_rocket(self.lp, self.fworker)
self.assertEqual(self.lp.launches.count(), 2)
self.assertEqual(self.lp.fireworks.count(), 3) # due to detour
new_wf = Workflow([Firework([ModSpecTask()])])
self.lp.append_wf(new_wf, [1, 2])
launch_rocket(self.lp, self.fworker) # launch detour
launch_rocket(self.lp, self.fworker) # launch new FW
launch_rocket(self.lp, self.fworker) # dummy launch
new_fw = self.lp.get_fw_by_id(4)
self.assertEqual(new_fw.spec['dummy1'], 1)
self.assertEqual(new_fw.spec['dummy2'], [True])
self.assertEqual(self.lp.launches.count(), 4)
self.assertEqual(self.lp.fireworks.count(), 4)
new_wf = Workflow([Firework([ModSpecTask()])])
self.lp.append_wf(new_wf, [4])
launch_rocket(self.lp, self.fworker) # launch new FW
new_fw = self.lp.get_fw_by_id(5)
self.assertEqual(new_fw.spec['dummy2'], [True])
new_wf = Workflow([Firework([ModSpecTask()])])
self.assertRaises(ValueError, self.lp.append_wf, new_wf, [4], detour=True)
def test_remove_root_fws(self):
fw4 = Firework(Task1(), parents=[self.fw2, self.fw3])
fws = [self.fw1, self.fw2, self.fw3, fw4]
wflow = Workflow(fws)
root_ids = wflow.root_fw_ids
children = []
for i in root_ids:
children.extend(wflow.links[i])
wflow.remove_fws(wflow.root_fw_ids)
self.assertEqual(sorted(wflow.root_fw_ids), sorted(children))
def test_remove_leaf_fws(self):
fw4 = Firework(Task1(), parents=[self.fw2, self.fw3])
fws = [self.fw1, self.fw2, self.fw3, fw4]
wflow = Workflow(fws)
leaf_ids = wflow.leaf_fw_ids
parents = []
for i in leaf_ids:
parents.extend(wflow.links.parent_links[i])
wflow.remove_fws(wflow.leaf_fw_ids)
self.assertEqual(wflow.leaf_fw_ids, parents)
def test_remove_leaf_fws(self):
fw4 = Firework(Task1(), parents=[self.fw2, self.fw3])
fws = [self.fw1, self.fw2, self.fw3, fw4]
wflow = Workflow(fws)
leaf_ids = wflow.leaf_fw_ids
parents = []
for i in leaf_ids:
parents.extend(wflow.links.parent_links[i])
wflow.remove_fws(wflow.leaf_fw_ids)
self.assertEqual(wflow.leaf_fw_ids, parents)
def test_remove_root_fws(self):
fw4 = Firework(Task1(), parents=[self.fw2, self.fw3])
fws = [self.fw1, self.fw2, self.fw3, fw4]
wflow = Workflow(fws)
root_ids = wflow.root_fw_ids
children = []
for i in root_ids:
children.extend(wflow.links[i])
wflow.remove_fws(wflow.root_fw_ids)
self.assertEqual(sorted(wflow.root_fw_ids), sorted(children))
def test_append_wf_detour(self):
fw1 = Firework([ModSpecTask()], fw_id=1)
fw2 = Firework([ModSpecTask()], fw_id=2, parents=[fw1])
self.lp.add_wf(Workflow([fw1, fw2]))
new_wf = Workflow([Firework([ModSpecTask()])])
self.lp.append_wf(new_wf, [1], detour=True)
launch_rocket(self.lp, self.fworker)
launch_rocket(self.lp, self.fworker)
self.assertEqual(self.lp.get_fw_by_id(2).spec['dummy2'], [True, True])
def test_fwconnector(self):
fw1 = Firework(ScriptTask.from_str('echo "1"'))
fw2 = Firework(ScriptTask.from_str('echo "1"'))
wf1 = Workflow([fw1, fw2], {fw1.fw_id: fw2.fw_id})
self.assertEqual(wf1.links, {fw1.fw_id: [fw2.fw_id], fw2.fw_id: []})
wf2 = Workflow([fw1, fw2], {fw1: fw2})
self.assertEqual(wf2.links, {fw1.fw_id: [fw2.fw_id], fw2.fw_id: []})
wf3 = Workflow([fw1, fw2])
self.assertEqual(wf3.links, {fw1.fw_id: [], fw2.fw_id: []})
def test_cleanupfiles_calc_dir(self):
# will test deleting from some folder specified by calc_dir
with ScratchDir(".", copy_from_current_on_enter=True) as d:
current_path = os.getcwd()
os.mkdir("to_remove.relax0")
os.mkdir("to_remove.relax1")
os.mkdir("dont_remove.relax0")
os.mkdir("shouldnt_touch")
fw1 = Firework(
[
DeleteFilesPrevFolder(files=["to_remove*", "dont_remove"],
calc_dir=current_path),
PassCalcLocs(name="fw1"),
],
name="fw1",
)
fw2 = Firework([PassCalcLocs(name="fw2")], name="fw2", parents=fw1)
wf = Workflow([fw1, fw2])
self.lp.add_wf(wf)
rapidfire(self.lp)
self.assertTrue(
os.path.exists(os.path.join(current_path,
"dont_remove.relax0")))
self.assertTrue(
os.path.exists(os.path.join(current_path, "shouldnt_touch")))
self.assertFalse(
os.path.exists(os.path.join(current_path, "to_remove.relax0")))
self.assertFalse(
os.path.exists(os.path.join(current_path, "to_remove.relax1")))
def vasp_jobs_scan_and_run(dir, vasp_cmd, label):
"""
Args:
dir: directory need to scan
vasp_cmd: vasp run command executed by subprocess.Popen, e.g. ['mpirun','vasp_std'] or ['srun','vasp_std']
label: a label for these jobs
"""
work_dirs = init.get_directories_NeedVaspRun(dir)
fws = []
njobs = 0
for work_dir in work_dirs:
queue = init.queue_setup(work_dir)
fw_name = queue['job_name']
ftask = VaspRun()
fw = Firework(
[ftask],
spec={
'vasp_cmd': vasp_cmd,
'_launch_dir': work_dir,
'_queueadapter': queue,
'_fworker': fworker
},
name=fw_name)
fws.append(fw)
njobs = njobs + 1
wf = Workflow(fws, name=label)
launchpad.add_wf(wf)
return njobs
def test_iter_and_len(self):
fws = [self.fw1, self.fw2, self.fw3]
wflow = Workflow(fws)
for idx, fw in enumerate(wflow):
self.assertEqual(fw, fws[idx])
assert len(wflow) == len(fws)
def add_wf(self, wf, reassign_all=True):
"""
:param wf: a Workflow object.
"""
m_timer.start("add_wf")
if isinstance(wf, FireWork):
wf = Workflow.from_FireWork(wf)
# sets the root FWs as READY
# prefer to wf.refresh() for speed reasons w/many root FWs
for fw_id in wf.root_fw_ids:
wf.id_fw[fw_id].state = 'READY'
# insert the FireWorks and get back mapping of old to new ids
old_new = self._upsert_fws(list(wf.id_fw.values()), reassign_all=reassign_all)
# update the Workflow with the new ids
wf._reassign_ids(old_new)
# insert the WFLinks
self.workflows.insert(wf.to_db_dict())
m_timer.stop("add_wf")
self.m_logger.info('Added a workflow. id_map: {}'.format(old_new))
return old_new
def submit_new_workflow(self):
# finds a submitted job, creates a workflow, and submits it to FireWorks
job = self.jobs.find_and_modify({'state': 'SUBMITTED'},
{'$set': {
'state': 'WAITING'
}})
if job:
submission_id = job['submission_id']
try:
firework = FireWork(
[StructurePredictionTask()],
spec={
'species': job['species'],
'threshold': job['threshold'],
'submission_id': submission_id
})
wf = Workflow([firework],
metadata={'submission_id': submission_id})
self.launchpad.add_wf(wf)
print 'ADDED WORKFLOW FOR {}'.format(job['species'])
except:
self.jobs.find_and_modify({'submission_id': submission_id},
{'$set': {
'state': 'ERROR'
}})
traceback.print_exc()
return submission_id
def test_files_in_out(self):
# create the Workflow that passes files_in and files_out
fw1 = Firework(
[ScriptTask.from_str('echo "This is the first FireWork" > test1')],
spec={"_files_out": {"fwtest1": "test1"}}, fw_id=1)
fw2 = Firework([ScriptTask.from_str('gzip hello')], fw_id=2,
parents=[fw1],
spec={"_files_in": {"fwtest1": "hello"},
"_files_out": {"fw2": "hello.gz"}})
fw3 = Firework([ScriptTask.from_str('cat fwtest.2')], fw_id=3,
parents=[fw2],
spec={"_files_in": {"fw2": "fwtest.2"}})
wf = Workflow([fw1, fw2, fw3],
{fw1: [fw2], fw2: [fw3]})
# store workflow and launch it locally
self.lp.add_wf(wf)
launch_rocket(self.lp, self.fworker)
self.assertTrue(os.path.exists("test1"))
launch_rocket(self.lp, self.fworker)
self.assertTrue(os.path.exists("hello.gz"))
launch_rocket(self.lp, self.fworker)
self.assertTrue(os.path.exists("fwtest.2"))
for f in ["test1", "hello.gz", "fwtest.2"]:
os.remove(f)
def test_early_exit(self):
os.chdir(MODULE_DIR)
script_text = "echo hello from process $PPID; sleep 2"
fw1 = Firework(ScriptTask.from_str(shell_cmd=script_text,
parameters={"stdout_file": "task.out"}),
fw_id=1)
fw2 = Firework(ScriptTask.from_str(shell_cmd=script_text,
parameters={"stdout_file": "task.out"}),
fw_id=2)
fw3 = Firework(ScriptTask.from_str(shell_cmd=script_text,
parameters={"stdout_file": "task.out"}),
fw_id=3)
fw4 = Firework(ScriptTask.from_str(shell_cmd=script_text,
parameters={"stdout_file": "task.out"}),
fw_id=4)
wf = Workflow([fw1, fw2, fw3, fw4], {1: [2, 3], 2: [4], 3: [4]})
self.lp.add_wf(wf)
launch_multiprocess(self.lp, FWorker(), 'DEBUG', 0, 2, sleep_time=0.5)
fw2 = self.lp.get_fw_by_id(2)
fw3 = self.lp.get_fw_by_id(3)
with open(os.path.join(fw2.launches[0].launch_dir, "task.out")) as f:
fw2_text = f.read()
with open(os.path.join(fw3.launches[0].launch_dir, "task.out")) as f:
fw3_text = f.read()
self.assertNotEqual(fw2_text, fw3_text)
def test_copy(self):
"""Test that we can produce a copy of a Workflow but that the copy
has unique fw_ids.
"""
fws = []
for i in range(5):
fw = Firework([PyTask(func="print", args=[i])], fw_id=i,
name=i)
fws.append(fw)
wf = Workflow(fws, links_dict={0: [1, 2, 3], 1: [4], 2: [4]})
wf_copy = Workflow.from_wflow(wf)
# now we compare to the original to make sure dependencies are same.
# have to do gymnastics because ids will NOT be the same
# but names are retained
for fw in wf_copy.fws:
children = wf_copy.links.get(fw.fw_id, list())
orig_id = fw.name
orig_children = wf.links.get(orig_id, list())
for child_id, orig_child_id in zip(children, orig_children):
self.assertEqual(orig_child_id, wf_copy.id_fw[child_id].name)
def add_wf(self, wf, reassign_all=True):
"""
:param wf: a Workflow object.
"""
m_timer.start("add_wf")
if isinstance(wf, FireWork):
wf = Workflow.from_FireWork(wf)
# sets the root FWs as READY
# prefer to wf.refresh() for speed reasons w/many root FWs
for fw_id in wf.root_fw_ids:
wf.id_fw[fw_id].state = 'READY'
# insert the FireWorks and get back mapping of old to new ids
old_new = self._upsert_fws(list(wf.id_fw.values()), reassign_all=reassign_all)
# update the Workflow with the new ids
wf._reassign_ids(old_new)
# insert the WFLinks
self.workflows.insert(wf.to_db_dict())
m_timer.stop("add_wf")
self.m_logger.info('Added a workflow. id_map: {}'.format(old_new))
return old_new
def test_init(self):
fws = []
for i in range(5):
fw = Firework([PyTask(func="print", args=[i])], fw_id=i)
fws.append(fw)
wf = Workflow(fws, links_dict={0: [1, 2, 3], 1: [4], 2: [4]})
self.assertIsInstance(wf, Workflow)
self.assertRaises(ValueError,
Workflow,
fws,
links_dict={
0: [1, 2, 3],
1: [4],
100: [4]
})
self.assertRaises(ValueError,
Workflow,
fws,
links_dict={
0: [1, 2, 3],
1: [4],
2: [100]
})
def create(self):
from fireworks.core.firework import Workflow
self.wf = Workflow(self.work_list,
self.connections,
name='VaspGWFWWorkFlow',
created_on=now())
print('creating workflow')
def test_preserve_fworker(self):
fw1 = Firework(
[ScriptTask.from_str('echo "Testing preserve FWorker"')],
spec={"_preserve_fworker": True},
fw_id=1)
fw2 = Firework(
[ScriptTask.from_str('echo "Testing preserve FWorker pt 2"')],
spec={"target": 1},
parents=[fw1],
fw_id=2)
self.lp.add_wf(Workflow([fw1, fw2]))
launch_rocket(self.lp, self.fworker)
target_fw_id = self.lp.get_fw_ids({"spec.target": 1})[0]
modified_spec = self.lp.get_fw_by_id(target_fw_id).spec
"""
cnt = 0
while '_fworker' not in modified_spec and cnt < 5:
modified_spec = self.lp.get_fw_by_id(target_fw_id).spec
print(modified_spec)
time.sleep(5)
cnt += 1
"""
self.assertIsNotNone(modified_spec['_fworker'])
def snl_to_wf_elastic(snl, parameters):
# parameters["user_vasp_settings"] specifies user defined incar/kpoints parameters
fws = []
connections = defaultdict(list)
parameters = parameters if parameters else {}
snl_priority = parameters.get('priority', 1)
priority = snl_priority * 2 # once we start a job, keep going!
f = Composition(snl.structure.composition.reduced_formula).alphabetical_formula
# add the SNL to the SNL DB and figure out duplicate group
tasks = [AddSNLTask()]
spec = {'task_type': 'Add to SNL database', 'snl': snl.as_dict(),
'_queueadapter': QA_DB, '_priority': snl_priority}
if 'snlgroup_id' in parameters and isinstance(snl, MPStructureNL):
spec['force_mpsnl'] = snl.as_dict()
spec['force_snlgroup_id'] = parameters['snlgroup_id']
del spec['snl']
fws.append(Firework(tasks, spec,
name=get_slug(f + '--' + spec['task_type']), fw_id=0))
connections[0] = [1]
parameters["exact_structure"] = True
# run GGA structure optimization for force convergence
spec = snl_to_wf._snl_to_spec(snl, parameters=parameters)
user_vasp_settings = parameters.get("user_vasp_settings")
spec = update_spec_force_convergence(spec, user_vasp_settings)
spec['run_tags'].append("origin")
spec['_priority'] = priority
spec['_queueadapter'] = QA_VASP
del spec['_dupefinder']
spec['task_type'] = "Vasp force convergence optimize structure (2x)"
tasks = [VaspWriterTask(), get_custodian_task(spec)]
fws.append(Firework(tasks, spec,
name=get_slug(f + '--' + spec['task_type']), fw_id=1))
# insert into DB - GGA structure optimization
spec = {'task_type': 'VASP db insertion', '_priority': priority,
'_allow_fizzled_parents': True, '_queueadapter': QA_DB,
'clean_task_doc':True, 'elastic_constant':"force_convergence"}
fws.append(Firework([VaspToDBTask()], spec,
name=get_slug(f + '--' + spec['task_type']), fw_id=2))
connections[1] = [2]
spec = {'task_type': 'Setup Deformed Struct Task', '_priority': priority,
'_queueadapter': QA_CONTROL}
fws.append(Firework([SetupDeformedStructTask()], spec,
name=get_slug(f + '--' + spec['task_type']),fw_id=3))
connections[2] = [3]
wf_meta = get_meta_from_structure(snl.structure)
wf_meta['run_version'] = 'May 2013 (1)'
if '_materialsproject' in snl.data and 'submission_id' in snl.data['_materialsproject']:
wf_meta['submission_id'] = snl.data['_materialsproject']['submission_id']
return Workflow(fws, connections, name=Composition(
snl.structure.composition.reduced_formula).alphabetical_formula, metadata=wf_meta)
def append_wf(args):
lp = get_lp(args)
lp.append_wf(
Workflow.from_file(args.wf_file),
args.fw_id,
detour=args.detour,
pull_spec_mods=args.pull_spec_mods
)
def __init__(self,
structure,
pseudos,
kppa=1000,
ecut=None,
pawecutdg=None,
accuracy="normal",
spin_mode="polarized",
smearing="fermi_dirac:0.1 eV",
charge=0.0,
scf_algorithm=None,
autoparal=False,
folder=None,
**extra_abivars):
"""
Args:
"""
abiinput = ion_ioncell_relax_input(structure=structure,
pseudos=pseudos,
kppa=kppa,
ecut=ecut,
pawecutdg=pawecutdg,
accuracy=accuracy,
spin_mode=spin_mode,
smearing=smearing,
charge=charge,
scf_algorithm=scf_algorithm)
abiinput.set_vars(**extra_abivars)
abitask = RelaxTask(StrategyWithInput(abiinput.split_datasets()[0]))
self.task = RelaxFWTask(abitask)
spec = {}
if folder:
spec['_launch_dir'] = os.path.join(folder, 'atomic_relax')
fw = Firework(self.task, spec=spec)
if autoparal:
autoparal_fw = self.create_autoparal_fw(self.task,
spec['_launch_dir'])
self.wf = Workflow([autoparal_fw, fw], {autoparal_fw: [fw]})
else:
self.wf = Workflow([fw])
def __init__(self, structure, pseudos):
abiinput = ebands_input(structure, pseudos).split_datasets()[0]
abitask = AbiFireTask(abiinput)
self.scf_fw = Firework(abitask)
self.wf = Workflow([self.scf_fw])
def test_parentconnector(self):
fw1 = Firework(ScriptTask.from_str('echo "1"'))
fw2 = Firework(ScriptTask.from_str('echo "1"'), parents=fw1)
fw3 = Firework(ScriptTask.from_str('echo "1"'), parents=[fw1, fw2])
self.assertEqual(
Workflow([fw1, fw2, fw3]).links, {fw1.fw_id: [fw2.fw_id, fw3.fw_id], fw2.fw_id: [fw3.fw_id], fw3.fw_id: []}
)
self.assertRaises(ValueError, Workflow, [fw1, fw3]) # can't make this
def __init__(self,
structure,
pseudos,
ksampling=1000,
relax_algo="atoms_only",
accuracy="normal",
spin_mode="polarized",
smearing="fermi_dirac:0.1 eV",
charge=0.0,
scf_algorithm=None,
autoparal=False,
max_restart=10,
folder=None,
**extra_abivars):
"""
Args:
"""
task = MultiStepRelaxStrategyFireTask(structure=structure,
pseudos=pseudos,
ksampling=ksampling,
relax_algo=relax_algo,
accuracy=accuracy,
spin_mode=spin_mode,
smearing=smearing,
charge=charge,
scf_algorithm=scf_algorithm,
deps={},
additional_steps=max_restart,
**extra_abivars)
spec = {}
if folder:
spec['_launch_dir'] = os.path.join(folder, 'relax')
fw = Firework(task, spec=spec)
# Create the workflow
if autoparal:
autoparal_fw = self.create_autoparal_fw(task, spec['_launch_dir'])
self.wf = Workflow([autoparal_fw, fw], {autoparal_fw: [fw]})
else:
self.wf = Workflow([fw])
def run_task(self, fw_spec):
# Read structure from previous relaxation
relaxed_struct = fw_spec['output']['crystal']
# Generate deformed structures
d_struct_set = DeformedStructureSet(relaxed_struct, ns=0.06)
wf=[]
for i, d_struct in enumerate(d_struct_set.def_structs):
fws=[]
connections={}
f = Composition(d_struct.formula).alphabetical_formula
snl = StructureNL(d_struct, 'Joseph Montoya <*****@*****.**>',
projects=["Elasticity"])
tasks = [AddSNLTask()]
snl_priority = fw_spec.get('priority', 1)
spec = {'task_type': 'Add Deformed Struct to SNL database',
'snl': snl.as_dict(),
'_queueadapter': QA_DB,
'_priority': snl_priority}
if 'snlgroup_id' in fw_spec and isinstance(snl, MPStructureNL):
spec['force_mpsnl'] = snl.as_dict()
spec['force_snlgroup_id'] = fw_spec['snlgroup_id']
del spec['snl']
fws.append(Firework(tasks, spec,
name=get_slug(f + '--' + spec['task_type']),
fw_id=-1000+i*10))
connections[-1000+i*10] = [-999+i*10]
spec = snl_to_wf._snl_to_spec(snl,
parameters={'exact_structure':True})
spec = update_spec_force_convergence(spec)
spec['deformation_matrix'] = d_struct_set.deformations[i].tolist()
spec['original_task_id'] = fw_spec["task_id"]
spec['_priority'] = fw_spec['_priority']*2
#Turn off dupefinder for deformed structure
del spec['_dupefinder']
spec['task_type'] = "Optimize deformed structure"
fws.append(Firework([VaspWriterTask(), SetupElastConstTask(),
get_custodian_task(spec)],
spec,
name=get_slug(f + '--' + spec['task_type']),
fw_id=-999+i*10))
priority = fw_spec['_priority']*3
spec = {'task_type': 'VASP db insertion',
'_priority': priority,
'_allow_fizzled_parents': True,
'_queueadapter': QA_DB,
'elastic_constant':"deformed_structure",
'clean_task_doc':True,
'deformation_matrix':d_struct_set.deformations[i].tolist(),
'original_task_id':fw_spec["task_id"]}
fws.append(Firework([VaspToDBTask(), AddElasticDataToDBTask()], spec,
name=get_slug(f + '--' + spec['task_type']),
fw_id=-998+i*10))
connections[-999+i*10] = [-998+i*10]
wf.append(Workflow(fws, connections))
return FWAction(additions=wf)
def test_getinterpolatedposcar(self):
nimages = 5
this_image = 1
autosort_tol = 0.5
fw1 = Firework([CopyVaspOutputs(calc_dir=self.static_outdir,
contcar_to_poscar=False,
additional_files=["CONTCAR"]),
PassCalcLocs(name="fw1")], name="fw1")
fw2 = Firework([CopyVaspOutputs(calc_dir=self.opt_outdir,
contcar_to_poscar=False,
additional_files=["CONTCAR"]),
PassCalcLocs(name="fw2")], name="fw2")
fw3 = Firework([GetInterpolatedPOSCAR(start="fw1",
end="fw2",
this_image=this_image,
nimages=nimages,
autosort_tol=autosort_tol),
PassCalcLocs(name="fw3")],
name="fw3", parents=[fw1, fw2])
fw4 = Firework([PassCalcLocs(name="fw4")], name="fw4", parents=fw3)
wf = Workflow([fw1, fw2, fw3, fw4])
self.lp.add_wf(wf)
rapidfire(self.lp)
fw4 = self.lp.get_fw_by_id(self.lp.get_fw_ids({"name": "fw4"})[0])
calc_locs = fw4.spec["calc_locs"]
self.assertTrue(os.path.exists(get_calc_loc("fw3", calc_locs)["path"] +
"/POSCAR"))
self.assertTrue(os.path.exists(get_calc_loc("fw3", calc_locs)["path"] +
"/interpolate/CONTCAR_0"))
self.assertTrue(os.path.exists(get_calc_loc("fw3", calc_locs)["path"] +
"/interpolate/CONTCAR_1"))
struct_start = Structure.from_file(get_calc_loc("fw3", calc_locs)["path"] +
"/interpolate/CONTCAR_0")
struct_end = Structure.from_file(get_calc_loc("fw3", calc_locs)["path"] +
"/interpolate/CONTCAR_1")
struct_inter = Structure.from_file(get_calc_loc("fw3", calc_locs)["path"] +
"/POSCAR")
structs = struct_start.interpolate(struct_end,
nimages,
interpolate_lattices=True,
autosort_tol=autosort_tol)
# Check x of 1st site.
self.assertAlmostEqual(structs[this_image][1].coords[0],
struct_inter[1].coords[0])
# Check c lattice parameter
self.assertAlmostEqual(structs[this_image].lattice.abc[0],
struct_inter.lattice.abc[0])
def test_offline_fw_passinfo(self):
fw1 = Firework([AdditionTask()], {"input_array": [1, 1]}, name="1")
fw2 = Firework([AdditionTask()], {"input_array": [2, 2]}, name="2")
fw3 = Firework([AdditionTask()], {"input_array": [3]},
parents=[fw1, fw2],
name="3")
wf = Workflow([fw1, fw2, fw3])
self.lp.add_wf(wf)
# make dirs for launching jobs
cur_dir = os.path.dirname(os.path.abspath(__file__))
os.mkdir(os.path.join(cur_dir, "launcher_1"))
os.mkdir(os.path.join(cur_dir, "launcher_2"))
os.mkdir(os.path.join(cur_dir, "launcher_3"))
# launch two parent jobs
os.chdir(os.path.join(cur_dir, "launcher_1"))
fw, launch_id = self.lp.reserve_fw(self.fworker, os.getcwd())
setup_offline_job(self.lp, fw, launch_id)
launch_rocket(None, self.fworker)
os.chdir(os.path.join(cur_dir, "launcher_2"))
fw, launch_id = self.lp.reserve_fw(self.fworker, os.getcwd())
setup_offline_job(self.lp, fw, launch_id)
launch_rocket(None, self.fworker)
# recover jobs
for l in self.lp.offline_runs.find(
{
"completed": False,
"deprecated": False
}, {"launch_id": 1}):
fw = self.lp.recover_offline(l['launch_id'])
# launch child job
os.chdir(os.path.join(cur_dir, "launcher_3"))
fw, launch_id = self.lp.reserve_fw(self.fworker, os.getcwd())
last_fw_id = fw.fw_id
setup_offline_job(self.lp, fw, launch_id)
launch_rocket(None, self.fworker)
# recover jobs
for l in self.lp.offline_runs.find(
{
"completed": False,
"deprecated": False
}, {"launch_id": 1}):
fw = self.lp.recover_offline(l['launch_id'])
# confirm the sum in the child job
child_fw = self.lp.get_fw_by_id(last_fw_id)
self.assertEqual(set(child_fw.spec['input_array']), set([2, 3, 4]))
self.assertEqual(child_fw.launches[0].action.stored_data["sum"], 9)
def add_wf(args):
lp = get_lp(args)
if args.dir:
files = []
for f in args.wf_file:
files.extend([os.path.join(f, i) for i in os.listdir(f)])
else:
files = args.wf_file
for f in files:
fwf = Workflow.from_file(f)
lp.add_wf(fwf)
def get_wf_by_fw_id(self, fw_id):
"""
Given a FireWork id, give back the Workflow containing that FireWork
:param fw_id:
:return: A Workflow object
"""
links_dict = self.workflows.find_one({'nodes': fw_id})
fws = map(self.get_fw_by_id, links_dict["nodes"])
return Workflow(fws, links_dict['links'], links_dict['name'],
links_dict['metadata'])
def test_multi_detour(self):
fw1 = Firework([MultipleDetourTask()], fw_id=1)
fw2 = Firework([ScriptTask.from_str('echo "DONE"')], parents=[fw1], fw_id=2)
self.lp.add_wf(Workflow([fw1, fw2]))
rapidfire(self.lp)
links = self.lp.get_wf_by_fw_id(1).links
self.assertEqual(set(links[1]), set([2, 3, 4, 5]))
self.assertEqual(set(links[2]), set([]))
self.assertEqual(set(links[3]), set([2]))
self.assertEqual(set(links[4]), set([2]))
self.assertEqual(set(links[5]), set([2]))
def add_scripts(args):
lp = get_lp(args)
args.names = args.names if args.names else [None] * len(args.scripts)
args.wf_name = args.wf_name if args.wf_name else args.names[0]
fws = []
links = {}
for idx, s in enumerate(args.scripts):
fws.append(Firework(ScriptTask({'script': s, 'use_shell': True}), name=args.names[idx], fw_id=idx))
if idx != 0:
links[idx-1] = idx
lp.add_wf(Workflow(fws, links, args.wf_name))
def test_spec_copy(self):
task1 = ScriptTask.from_str('echo "Task 1"')
task2 = ScriptTask.from_str('echo "Task 2"')
spec = {'_category': 'dummy_category'}
fw1 = Firework(task1, fw_id=1, name='Task 1', spec=spec)
fw2 = Firework(task2, fw_id=2, name='Task 2', spec=spec)
self.lp.add_wf(Workflow([fw1, fw2]))
self.assertEqual(self.lp.get_fw_by_id(1).tasks[0]['script'][0], 'echo "Task 1"')
self.assertEqual(self.lp.get_fw_by_id(2).tasks[0]['script'][0], 'echo "Task 2"')
def task_dict_to_wf(task_dict, launchpad):
fw_id = launchpad.get_new_fw_id()
l_id = launchpad.get_new_launch_id()
spec = {'task_type': task_dict['task_type'], 'run_tags': task_dict['run_tags'],
'vaspinputset_name': None, 'vasp': None, 'mpsnl': task_dict['snl'],
'snlgroup_id': task_dict['snlgroup_id']}
tasks = [DummyLegacyTask()]
launch_dir = task_dict['dir_name_full']
stored_data = {'error_list': []}
update_spec = {'prev_vasp_dir': task_dict['dir_name'],
'prev_task_type': spec['task_type'],
'mpsnl': spec['mpsnl'], 'snlgroup_id': spec['snlgroup_id'],
'run_tags': spec['run_tags']}
fwaction = FWAction(stored_data=stored_data, update_spec=update_spec)
if task_dict['completed_at']:
complete_date = datetime.datetime.strptime(task_dict['completed_at'], "%Y-%m-%d %H:%M:%S")
state_history = [{"created_on": complete_date, 'state': 'COMPLETED'}]
else:
state_history = []
launches = [Launch('COMPLETED', launch_dir, fworker=None, host=None, ip=None, action=fwaction,
state_history=state_history, launch_id=l_id, fw_id=fw_id)]
f = Composition(task_dict['pretty_formula']).alphabetical_formula
fw = Firework(tasks, spec, name=get_slug(f + '--' + spec['task_type']), launches=launches, state='COMPLETED', created_on=None,
fw_id=fw_id)
wf_meta = get_meta_from_structure(Structure.from_dict(task_dict['snl']))
wf_meta['run_version'] = 'preproduction (0)'
wf = Workflow.from_FireWork(fw, name=f, metadata=wf_meta)
launchpad.add_wf(wf, reassign_all=False)
launchpad._upsert_launch(launches[0])
print 'ADDED', fw_id
# return fw_id
return fw_id
def add_wf_dir(args):
lp = get_lp(args)
for filename in os.listdir(args.wf_dir):
fwf = Workflow.from_file(filename)
lp.add_wf(fwf)
def task_dict_to_wf(task_dict, launchpad):
fw_id = launchpad.get_new_fw_id()
l_id = launchpad.get_new_launch_id()
spec = {
"task_type": task_dict["task_type"],
"run_tags": task_dict["run_tags"],
"vaspinputset_name": None,
"vasp": None,
"mpsnl": task_dict["snl"],
"snlgroup_id": task_dict["snlgroup_id"],
}
tasks = [DummyLegacyTask()]
launch_dir = task_dict["dir_name_full"]
stored_data = {"error_list": []}
update_spec = {
"prev_vasp_dir": task_dict["dir_name"],
"prev_task_type": spec["task_type"],
"mpsnl": spec["mpsnl"],
"snlgroup_id": spec["snlgroup_id"],
"run_tags": spec["run_tags"],
}
fwaction = FWAction(stored_data=stored_data, update_spec=update_spec)
if task_dict["completed_at"]:
complete_date = datetime.datetime.strptime(task_dict["completed_at"], "%Y-%m-%d %H:%M:%S")
state_history = [{"created_on": complete_date, "state": "COMPLETED"}]
else:
state_history = []
launches = [
Launch(
"COMPLETED",
launch_dir,
fworker=None,
host=None,
ip=None,
action=fwaction,
state_history=state_history,
launch_id=l_id,
fw_id=fw_id,
)
]
f = Composition(task_dict["pretty_formula"]).alphabetical_formula
fw = Firework(
tasks,
spec,
name=get_slug(f + "--" + spec["task_type"]),
launches=launches,
state="COMPLETED",
created_on=None,
fw_id=fw_id,
)
wf_meta = get_meta_from_structure(Structure.from_dict(task_dict["snl"]))
wf_meta["run_version"] = "preproduction (0)"
wf = Workflow.from_FireWork(fw, name=f, metadata=wf_meta)
launchpad.add_wf(wf, reassign_all=False)
launchpad._upsert_launch(launches[0])
print "ADDED", fw_id
# return fw_id
return fw_id
