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 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_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_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 run_task(self, fw_spec):
print('Running the Multiple Detour Task')
dt1 = Firework(
ScriptTask.from_str('echo "this is intermediate job 1"'))
dt2 = Firework(
ScriptTask.from_str('echo "this is intermediate job 2"'))
dt3 = Firework(
ScriptTask.from_str('echo "this is intermediate job 3"'))
return FWAction(detours=[dt1, dt2, dt3])
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 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 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_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 test_copyfilesfromcalcloc(self):
fw1 = Firework(
[
CopyVaspOutputs(calc_dir=self.plain_outdir),
PassCalcLocs(name="fw1")
],
name="fw1",
)
fw2 = Firework(
[
CopyVaspOutputs(calc_dir=self.relax2_outdir),
PassCalcLocs(name="fw2")
],
name="fw2",
)
fw3 = Firework(
[
CopyFilesFromCalcLoc(
calc_loc="fw1",
filenames=["POSCAR"],
name_prepend="",
name_append="_0",
),
CopyFilesFromCalcLoc(
calc_loc="fw2",
filenames=["POSCAR"],
name_prepend="",
name_append="_1",
),
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_0"))
self.assertTrue(
os.path.exists(
get_calc_loc("fw3", calc_locs)["path"] + "/POSCAR_1"))
def createVaspSRCFireworks(vasp_input_set, task_helper, task_type, control_procedure,
custodian_handlers, max_restarts, src_cleaning, task_index, spec,
setup_spec_update=None, run_spec_update=None):
# Make a full copy of the spec
if spec is None:
spec = {}
spec = copy.deepcopy(spec)
spec['_add_launchpad_and_fw_id'] = True
spec['_add_fworker'] = True
# Initialize the SRC task_index
if task_index is not None:
src_task_index = SRCTaskIndex.from_any(task_index)
else:
src_task_index = SRCTaskIndex.from_string(task_type)
spec['SRC_task_index'] = src_task_index
# SetupTask
setup_spec = copy.deepcopy(spec)
# Remove any initial queue_adapter_update from the spec
setup_spec.pop('queue_adapter_update', None)
setup_spec = set_short_single_core_to_spec(setup_spec)
setup_spec['_preserve_fworker'] = True
setup_spec['_pass_job_info'] = True
setup_spec.update({} if setup_spec_update is None else setup_spec_update)
setup_task = VaspSetupTask(vasp_input_set=vasp_input_set, deps=None, task_helper=task_helper, task_type=task_type)
setup_fw = Firework(setup_task, spec=setup_spec, name=src_task_index.setup_str)
# RunTask
run_spec = copy.deepcopy(spec)
run_spec['SRC_task_index'] = src_task_index
run_spec['_preserve_fworker'] = True
run_spec['_pass_job_info'] = True
run_spec.update({} if run_spec_update is None else run_spec_update)
run_task = VaspRunTask(control_procedure=control_procedure, task_helper=task_helper, task_type=task_type,
custodian_handlers=custodian_handlers)
run_fw = Firework(run_task, spec=run_spec, name=src_task_index.run_str)
# ControlTask
control_spec = copy.deepcopy(spec)
control_spec = set_short_single_core_to_spec(control_spec)
control_spec['SRC_task_index'] = src_task_index
control_spec['_allow_fizzled_parents'] = True
control_task = VaspControlTask(control_procedure=control_procedure, manager=None, max_restarts=max_restarts,
src_cleaning=src_cleaning, task_helper=task_helper)
control_fw = Firework(control_task, spec=control_spec, name=src_task_index.control_str)
links_dict = {setup_fw.fw_id: [run_fw.fw_id],
run_fw.fw_id: [control_fw.fw_id]}
return {'setup_fw': setup_fw, 'run_fw': run_fw, 'control_fw': control_fw, 'links_dict': links_dict,
'fws': [setup_fw, run_fw, control_fw]}
def test_parallel_fibadder(self):
# this is really testing to see if a Workflow can handle multiple FWs updating it at once
parent = Firework(ScriptTask.from_str("python -c 'print(\"test1\")'", {'store_stdout': True}))
fib1 = Firework(FibonacciAdderTask(), {'smaller': 0, 'larger': 1, 'stop_point': 30}, parents=[parent])
fib2 = Firework(FibonacciAdderTask(), {'smaller': 0, 'larger': 1, 'stop_point': 30}, parents=[parent])
fib3 = Firework(FibonacciAdderTask(), {'smaller': 0, 'larger': 1, 'stop_point': 30}, parents=[parent])
fib4 = Firework(FibonacciAdderTask(), {'smaller': 0, 'larger': 1, 'stop_point': 30}, parents=[parent])
wf = Workflow([parent, fib1, fib2, fib3, fib4])
self.lp.add_wf(wf)
p = Pool(NCORES_PARALLEL_TEST)
creds_array = [self.lp.to_dict()] * NCORES_PARALLEL_TEST
p.map(random_launch, creds_array)
def test_category_pt2(self):
task1 = ScriptTask.from_str('echo "Task 1"')
task2 = ScriptTask.from_str('echo "Task 2"')
fw1 = Firework(task1, fw_id=1, name='Task 1')
fw2 = Firework(task2, fw_id=2, name='Task 2')
self.lp.add_wf(Workflow([fw1, fw2]))
self.assertFalse(self.lp.run_exists(FWorker(category="dummy_category")))
self.assertTrue(self.lp.run_exists(FWorker(category="__none__")))
self.assertTrue(self.lp.run_exists(FWorker())) # can run any category
self.assertFalse(self.lp.run_exists(FWorker(category=["dummy_category",
"other category"])))
def test_org_wf(self):
test1 = ScriptTask.from_str("python -c 'print(\"test1\")'",
{'store_stdout': True})
test2 = ScriptTask.from_str("python -c 'print(\"test2\")'",
{'store_stdout': True})
fw1 = Firework(test1, fw_id=-1)
fw2 = Firework(test2, fw_id=-2)
wf = Workflow([fw1, fw2], {-1: -2})
self.lp.add_wf(wf)
launch_rocket(self.lp, self.fworker)
self.assertEqual(self.lp.get_launch_by_id(1).action.stored_data['stdout'],
'test1\n')
launch_rocket(self.lp, self.fworker)
self.assertEqual(self.lp.get_launch_by_id(2).action.stored_data['stdout'],
'test2\n')
def test_cleanupfiles_calc_loc(self):
# will test deleting files from previous folder
fw1 = Firework(
[
CreateFolder(folder_name="to_remove.relax0"),
CreateFolder(folder_name="to_remove.relax1"),
CreateFolder(folder_name="dont_remove.relax0"),
CreateFolder(folder_name="shouldnt_touch"),
PassCalcLocs(name="fw1"),
],
name="fw1",
)
fw2 = Firework(
[
DeleteFilesPrevFolder(files=["to_remove*", "dont_remove"],
calc_loc=True),
PassCalcLocs(name="fw2"),
],
name="fw2",
parents=fw1,
)
wf = Workflow([fw1, fw2])
self.lp.add_wf(wf)
rapidfire(self.lp)
fw2 = self.lp.get_fw_by_id(self.lp.get_fw_ids({"name": "fw2"})[0])
calc_locs = fw2.spec["calc_locs"]
self.assertTrue(
os.path.exists(
os.path.join(
get_calc_loc("fw1", calc_locs)["path"],
"dont_remove.relax0")))
self.assertTrue(
os.path.exists(
os.path.join(
get_calc_loc("fw1", calc_locs)["path"], "shouldnt_touch")))
self.assertFalse(
os.path.exists(
os.path.join(
get_calc_loc("fw1", calc_locs)["path"],
"to_remove.relax0")))
self.assertFalse(
os.path.exists(
os.path.join(
get_calc_loc("fw1", calc_locs)["path"],
"to_remove.relax1")))
def test_job_info(self):
fw1 = Firework([ScriptTask.from_str('echo "Testing job info"')],
spec={"_pass_job_info": True},
fw_id=1)
fw2 = Firework([DummyJobPassTask()],
parents=[fw1],
spec={
"_pass_job_info": True,
"target": 1
},
fw_id=2)
fw3 = Firework([DummyJobPassTask()],
parents=[fw2],
spec={"target": 2},
fw_id=3)
self.lp.add_wf(Workflow([fw1, fw2, fw3]))
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 '_job_info' not in modified_spec and cnt < 5:
print(modified_spec)
modified_spec = self.lp.get_fw_by_id(target_fw_id).spec
time.sleep(5)
cnt += 1
"""
self.assertIsNotNone(modified_spec['_job_info'])
self.assertIsNotNone(modified_spec['_job_info'][0]["launch_dir"])
self.assertEqual(modified_spec['_job_info'][0]['name'], 'Unnamed FW')
self.assertEqual(modified_spec['_job_info'][0]['fw_id'], 1)
launch_rocket(self.lp, self.fworker)
target_fw_id = self.lp.get_fw_ids({"spec.target": 2})[0]
modified_spec = self.lp.get_fw_by_id(target_fw_id).spec
"""
cnt = 0
while '_job_info' not in modified_spec and cnt < 5:
print(modified_spec)
modified_spec = self.lp.get_fw_by_id(target_fw_id).spec
time.sleep(5)
cnt += 1
"""
self.assertEqual(len(modified_spec['_job_info']), 2)
def test_tracker(self):
"""
Launch a workflow and track the files
"""
self._teardown([self.dest1])
try:
fts = []
for i in range(5, 100):
ft = ScriptTask.from_str(
'echo "' + str(i) + '" >> ' + self.dest1,
{"store_stdout": True})
fts.append(ft)
fw = Firework(fts,
spec={"_trackers": [self.tracker1]},
fw_id=20,
name="test_fw")
self.lp.add_wf(fw)
launch_rocket(self.lp, self.fworker)
# print (self.tracker1.track_file())
self.assertEqual("98\n99", self.tracker1.track_file())
finally:
self._teardown([self.dest1])
def test_basic_fw_offline(self):
test1 = ScriptTask.from_str("python -c 'print(\"test1\")'",
{'store_stdout': True})
fw = Firework(test1)
self.lp.add_wf(fw)
fw, launch_id = self.lp.reserve_fw(self.fworker, os.getcwd())
setup_offline_job(self.lp, fw, launch_id)
launch_rocket(None, self.fworker)
with open(os.path.join(os.getcwd(), "FW_offline.json")) as f:
fwo = json.load(f)
self.assertEquals(fwo["state"], "COMPLETED")
self.assertEquals(fwo["launch_id"], 1)
self.assertEquals(fwo["fwaction"], {'update_spec': {}, 'mod_spec': [], 'stored_data': {'returncode': 0, 'stdout': u'test1\n', 'all_returncodes': [0]}, 'exit': False, 'detours': [], 'additions': [], 'defuse_children': False, 'defuse_workflow': False})
with open(os.path.join(os.getcwd(), "FW_ping.json")) as f:
fwp = json.load(f)
self.assertIsNotNone(fwp["ping_time"])
l = self.lp.offline_runs.find_one({"completed": False, "deprecated": False}, {"launch_id": 1})
self.lp.recover_offline(l['launch_id'])
self.assertEqual(self.lp.get_launch_by_id(1).action.stored_data['stdout'], 'test1\n')
def create_autoparal_fw(cls, firetask, folder=None):
spec = {'_queueadapter': {'ntasks': 1, 'walltime': '00:10:00'}}
if folder:
spec['_launch_dir'] = folder + '_autoparal'
autoparal_task = AutoparalFireTask(firetask)
autoparal_fw = Firework(autoparal_task, spec=spec)
return autoparal_fw
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 counterpoise_correction_generation_fw(molname, charge, spin_multiplicity,
qm_method, fragments,
mission, priority=1,
parent_fwid=None,
additional_user_tags=None,
large=False):
fw_spec = dict()
fw_spec["user_tags"] = dict()
fw_spec["user_tags"]["molname"] = molname
fw_spec["user_tags"]["mission"] = mission
fw_spec["qm_method"] = qm_method
fw_spec["fragments"] = fragments
fw_spec["charge"] = charge
fw_spec["spin_multiplicity"] = spin_multiplicity
fw_spec["large"] = large
fw_spec["task_type"] = "counterpoise correction generation"
if priority:
fw_spec['_priority'] = priority
fw_spec["user_tags"].update(additional_user_tags)
fwid_base = 1
if parent_fwid:
if not (isinstance(parent_fwid, int) or isinstance(parent_fwid, list)):
raise ValueError("Parent FireWork ID must be integer or list")
parent_fwid = parent_fwid if isinstance(parent_fwid, list) \
else [parent_fwid]
fwid_base = max(parent_fwid) + 1
current_fwid = fwid_base
links_dict = dict()
fw_cp = Firework([CounterpoiseCorrectionGenerationTask()],
spec=fw_spec,
name=molname + " Counterpoise Correction Generation",
fw_id=current_fwid)
for p_fwid in parent_fwid:
links_dict[p_fwid] = current_fwid
return [fw_cp], links_dict
def test_archive(self):
p = PyTask(func="fireworks.tests.mongo_tests.throw_error",
args=["This should not happen"])
fw = Firework(p)
self.lp.add_wf(fw)
self.lp.archive_wf(fw.fw_id)
self.assertFalse(launch_rocket(self.lp, self.fworker))
def test_backgroundtask(self):
dest1 = os.path.join(os.path.dirname(os.path.abspath(__file__)),
'hello.txt')
self._teardown([dest1])
try:
test1 = ScriptTask.from_str(
"python -c 'print(\"testing background...\")'",
{'store_stdout': True})
bg_task1 = BackgroundTask(FileWriteTask(
{'files_to_write': [{
'filename': dest1,
'contents': 'hello'
}]}),
num_launches=1,
run_on_finish=True)
fw = Firework(test1, spec={'_background_tasks': [bg_task1]})
self.lp.add_wf(fw)
launch_rocket(self.lp, self.fworker)
with open(dest1) as f:
self.assertEqual(f.read(), 'hello')
finally:
self._teardown([dest1])
def test_walltime_handler(self):
# Simple test from a queue_adapter
f = open(os.path.join(test_dir, "fw_handlers", "walltime", "PBS_qadapter.json"))
dd = json.load(f)
f.close()
qad = QueueAdapter.from_dict(dd)
walltime_handler = WalltimeHandler(job_rundir=os.path.join(test_dir, "fw_handlers", "walltime"),
qout_file='queue.qout', qerr_file='queue.qerr', queue_adapter=qad,
max_timelimit=1200, timelimit_increase=600)
self.assertTrue(walltime_handler.check())
# Test with a given Firework
walltime_handler = WalltimeHandler(job_rundir='.',
qout_file='queue.qout', qerr_file='queue.qerr', queue_adapter=None,
max_timelimit=1200, timelimit_increase=600)
# Get the Firework from json file
f = open(os.path.join(test_dir, 'fw_handlers', 'walltime', 'sleep_fw.json'))
dd = json.load(f)
f.close()
fw_to_check = Firework.from_dict(dd)
# Hack the launch_dir so that it points to the directory where the queue.qout is
fw_to_check.launches[-1].launch_dir = os.path.join(test_dir, fw_to_check.launches[-1].launch_dir)
walltime_handler.src_setup(fw_spec={}, fw_to_check=fw_to_check)
self.assertTrue(walltime_handler.check())
actions = walltime_handler.correct()
self.assertEqual(actions, {'errors': ['WalltimeHandler'],
'actions': [{'action': {'_set': {'timelimit': 660.0}},
'object': {'source': 'fw_spec', 'key': 'qtk_queueadapter'},
'action_type': 'modify_object'}]})
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 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_fizzle(self):
p = PyTask(func="fireworks.tests.mongo_tests.throw_error", args=["Testing; this error is normal."])
fw = Firework(p)
self.lp.add_wf(fw)
self.assertTrue(launch_rocket(self.lp, self.fworker))
self.assertEqual(self.lp.get_fw_by_id(1).state, 'FIZZLED')
self.assertFalse(launch_rocket(self.lp, self.fworker))
def run(self):
"""
Run the rocket (check out a job from the database and execute it)
"""
all_stored_data = {} # combined stored data for *all* the Tasks
all_update_spec = {} # combined update_spec for *all* the Tasks
all_mod_spec = [] # combined mod_spec for *all* the Tasks
lp = self.launchpad
launch_dir = os.path.abspath(os.getcwd())
# check a FW job out of the launchpad
if lp:
m_fw, launch_id = lp.checkout_fw(self.fworker, launch_dir, self.fw_id)
else: # offline mode
m_fw = Firework.from_file(os.path.join(os.getcwd(), "FW.json"))
# set the run start time
with open('FW_offline.json', 'r+') as f:
d = json.loads(f.read())
d['started_on'] = datetime.utcnow().isoformat()
f.seek(0)
f.write(json.dumps(d))
f.truncate()
launch_id = None # we don't need this in offline mode...
if not m_fw:
print("No FireWorks are ready to run and match query! {}".format(self.fworker.query))
return False
if lp:
message = 'RUNNING fw_id: {} in directory: {}'.\
format(m_fw.fw_id, os.getcwd())
lp.log_message(logging.INFO, message)
# write FW.json and/or FW.yaml to the directory
if PRINT_FW_JSON:
m_fw.to_file('FW.json', indent=4)
if PRINT_FW_YAML:
m_fw.to_file('FW.yaml')
try:
if '_launch_dir' in m_fw.spec:
prev_dir = launch_dir
launch_dir = os.path.expandvars(m_fw.spec['_launch_dir'])
# thread-safe "mkdir -p"
try:
os.makedirs(launch_dir)
except OSError as exception:
if exception.errno != errno.EEXIST:
raise
os.chdir(launch_dir)
launch_dir = os.path.abspath(os.getcwd())
if lp:
lp.change_launch_dir(launch_id, launch_dir)
if not os.listdir(prev_dir) and REMOVE_USELESS_DIRS:
try:
os.rmdir(prev_dir)
except:
pass
if m_fw.spec.get('_recover_launch', None):
launch_to_recover = lp.get_launch_by_id(m_fw.spec['_recover_launch']['_launch_id'])
starting_task = launch_to_recover.action.stored_data.get('_exception', {}).get('_failed_task_n', 0)
recover_launch_dir = launch_to_recover.launch_dir
if lp:
lp.log_message(
logging.INFO,
'Recovering from task number {} in folder {}.'.format(starting_task, recover_launch_dir))
if m_fw.spec['_recover_launch']['_recover_mode'] == 'cp' and launch_dir != recover_launch_dir:
if lp:
lp.log_message(
logging.INFO,
'Copying data from recovery folder {} to folder {}.'.format(recover_launch_dir, launch_dir))
distutils.dir_util.copy_tree(recover_launch_dir, launch_dir, update=1)
else:
starting_task = 0
my_spec = dict(m_fw.spec) # make a copy of spec, don't override original
my_spec["_fw_env"] = self.fworker.env
# set up heartbeat (pinging the server that we're still alive)
ping_stop = start_ping_launch(lp, launch_id)
# start background tasks
btask_stops = []
if '_background_tasks' in my_spec:
for bt in my_spec['_background_tasks']:
btask_stops.append(start_background_task(bt, m_fw.spec))
# execute the FireTasks!
for t_counter, t in enumerate(m_fw.tasks[starting_task:], start=starting_task):
if lp:
lp.log_message(logging.INFO, "Task started: %s." % t.fw_name)
try:
m_action = t.run_task(my_spec)
except BaseException as e:
traceback.print_exc()
tb = traceback.format_exc()
stop_backgrounds(ping_stop, btask_stops)
do_ping(lp, launch_id) # one last ping, esp if there is a monitor
# If the exception is serializable, save its details
try:
exception_details = e.to_dict()
except AttributeError:
exception_details = None
except BaseException as e:
if lp:
lp.log_message(logging.WARNING, "Exception couldn't be serialized: %s " % e)
exception_details = None
try:
m_task = t.to_dict()
except:
m_task = None
m_action = FWAction(stored_data={'_message': 'runtime error during task', '_task': m_task,
'_exception': {'_stacktrace': tb, '_details': exception_details,
'_failed_task_n': t_counter}}, exit=True)
if lp:
lp.complete_launch(launch_id, m_action, 'FIZZLED')
else:
with open('FW_offline.json', 'r+') as f:
d = json.loads(f.read())
d['fwaction'] = m_action.to_dict()
d['state'] = 'FIZZLED'
f.seek(0)
f.write(json.dumps(d))
f.truncate()
return True
# read in a FWAction from a file, in case the task is not Python and cannot return it explicitly
if os.path.exists('FWAction.json'):
m_action = FWAction.from_file('FWAction.json')
elif os.path.exists('FWAction.yaml'):
m_action = FWAction.from_file('FWAction.yaml')
if not m_action:
m_action = FWAction()
# update the global stored data with the data to store and update from this particular Task
all_stored_data.update(m_action.stored_data)
all_update_spec.update(m_action.update_spec)
all_mod_spec.extend(m_action.mod_spec)
# update spec for next task as well
my_spec.update(m_action.update_spec)
for mod in m_action.mod_spec:
apply_mod(mod, my_spec)
if lp:
lp.log_message(logging.INFO, "Task completed: %s " % t.fw_name)
if m_action.skip_remaining_tasks:
break
# add job packing info if this is needed
if FWData().MULTIPROCESSING and STORE_PACKING_INFO:
all_stored_data['multiprocess_name'] = multiprocessing.current_process().name
# perform finishing operation
stop_backgrounds(ping_stop, btask_stops)
for b in btask_stops:
b.set()
do_ping(lp, launch_id) # one last ping, esp if there is a monitor
# last background monitors
if '_background_tasks' in my_spec:
for bt in my_spec['_background_tasks']:
if bt.run_on_finish:
for task in bt.tasks:
task.run_task(m_fw.spec)
m_action.stored_data = all_stored_data
m_action.mod_spec = all_mod_spec
m_action.update_spec = all_update_spec
if lp:
lp.complete_launch(launch_id, m_action, 'COMPLETED')
else:
with open('FW_offline.json', 'r+') as f:
d = json.loads(f.read())
d['fwaction'] = m_action.to_dict()
d['state'] = 'COMPLETED'
d['completed_on'] = datetime.utcnow().isoformat()
f.seek(0)
f.write(json.dumps(d))
f.truncate()
return True
except:
# problems while processing the results. high probability of malformed data.
traceback.print_exc()
stop_backgrounds(ping_stop, btask_stops)
# restore initial state to prevent the raise of further exceptions
if lp:
lp.restore_backup_data(launch_id, m_fw.fw_id)
do_ping(lp, launch_id) # one last ping, esp if there is a monitor
# the action produced by the task is discarded
m_action = FWAction(stored_data={'_message': 'runtime error during task', '_task': None,
'_exception': {'_stacktrace': traceback.format_exc(),
'_details': None}}, exit=True)
if lp:
lp.complete_launch(launch_id, m_action, 'FIZZLED')
else:
with open('FW_offline.json', 'r+') as f:
d = json.loads(f.read())
d['fwaction'] = m_action.to_dict()
d['state'] = 'FIZZLED'
f.seek(0)
f.write(json.dumps(d))
f.truncate()
return True
def run(self):
"""
Run the rocket (check out a job from the database and execute it)
"""
all_stored_data = {} # combined stored data for *all* the Tasks
all_update_spec = {} # combined update_spec for *all* the Tasks
all_mod_spec = [] # combined mod_spec for *all* the Tasks
lp = self.launchpad
launch_dir = os.path.abspath(os.getcwd())
# check a FW job out of the launchpad
if lp:
m_fw, launch_id = lp.checkout_fw(self.fworker, launch_dir, self.fw_id)
else: # offline mode
m_fw = Firework.from_file(os.path.join(os.getcwd(), "FW.json"))
# set the run start time
with open('FW_offline.json', 'r+') as f:
d = json.loads(f.read())
d['started_on'] = datetime.utcnow().isoformat()
f.seek(0)
f.write(json.dumps(d))
f.truncate()
launch_id = None # we don't need this in offline mode...
if not m_fw:
print("No FireWorks are ready to run and match query! {}".format(self.fworker.query))
return False
if '_launch_dir' in m_fw.spec:
prev_dir = launch_dir
os.chdir(m_fw.spec['_launch_dir'])
launch_dir = os.path.abspath(os.getcwd())
if lp:
lp.change_launch_dir(launch_id, launch_dir)
if not os.listdir(prev_dir) and REMOVE_USELESS_DIRS:
try:
os.rmdir(prev_dir)
except:
pass
if lp:
message = 'RUNNING fw_id: {} in directory: {}'.\
format(m_fw.fw_id, os.getcwd())
lp.log_message(logging.INFO, message)
# write FW.json and/or FW.yaml to the directory
if PRINT_FW_JSON:
m_fw.to_file('FW.json', indent=4)
if PRINT_FW_YAML:
m_fw.to_file('FW.yaml')
try:
my_spec = dict(m_fw.spec) # make a copy of spec, don't override original
my_spec["_fw_env"] = self.fworker.env
# set up heartbeat (pinging the server that we're still alive)
ping_stop = start_ping_launch(lp, launch_id)
# start background tasks
btask_stops = []
if '_background_tasks' in my_spec:
for bt in my_spec['_background_tasks']:
btask_stops.append(start_background_task(bt, m_fw.spec))
# execute the FireTasks!
for t in m_fw.tasks:
lp.log_message(logging.INFO, "Task started: %s." % t.fw_name)
m_action = t.run_task(my_spec)
# read in a FWAction from a file, in case the task is not Python and cannot return it explicitly
if os.path.exists('FWAction.json'):
m_action = FWAction.from_file('FWAction.json')
elif os.path.exists('FWAction.yaml'):
m_action = FWAction.from_file('FWAction.yaml')
if not m_action:
m_action = FWAction()
# update the global stored data with the data to store and update from this particular Task
all_stored_data.update(m_action.stored_data)
all_update_spec.update(m_action.update_spec)
all_mod_spec.extend(m_action.mod_spec)
# update spec for next task as well
my_spec.update(m_action.update_spec)
for mod in m_action.mod_spec:
apply_mod(mod, my_spec)
lp.log_message(logging.INFO, "Task completed: %s " % t.fw_name)
if m_action.skip_remaining_tasks:
break
# add job packing info if this is needed
if FWData().MULTIPROCESSING and STORE_PACKING_INFO:
all_stored_data['multiprocess_name'] = multiprocessing.current_process().name
# perform finishing operation
stop_backgrounds(ping_stop, btask_stops)
for b in btask_stops:
b.set()
do_ping(lp, launch_id) # one last ping, esp if there is a monitor
# last background monitors
if '_background_tasks' in my_spec:
for bt in my_spec['_background_tasks']:
if bt.run_on_finish:
for task in bt.tasks:
task.run_task(m_fw.spec)
m_action.stored_data = all_stored_data
m_action.mod_spec = all_mod_spec
m_action.update_spec = all_update_spec
if lp:
lp.complete_launch(launch_id, m_action, 'COMPLETED')
else:
with open('FW_offline.json', 'r+') as f:
d = json.loads(f.read())
d['fwaction'] = m_action.to_dict()
d['state'] = 'COMPLETED'
d['completed_on'] = datetime.utcnow().isoformat()
f.seek(0)
f.write(json.dumps(d))
f.truncate()
return True
except:
stop_backgrounds(ping_stop, btask_stops)
do_ping(lp, launch_id) # one last ping, esp if there is a monitor
traceback.print_exc()
try:
m_action = FWAction(stored_data={'_message': 'runtime error during task', '_task': t.to_dict(),
'_exception': traceback.format_exc()}, exit=True)
except:
m_action = FWAction(stored_data={'_message': 'runtime error during task', '_task': None,
'_exception': traceback.format_exc()}, exit=True)
if lp:
lp.complete_launch(launch_id, m_action, 'FIZZLED')
else:
with open('FW_offline.json', 'r+') as f:
d = json.loads(f.read())
d['fwaction'] = m_action.to_dict()
d['state'] = 'FIZZLED'
f.seek(0)
f.write(json.dumps(d))
f.truncate()
return True
def run(self, pdb_on_exception=False):
"""
Run the rocket (check out a job from the database and execute it)
Args:
pdb_on_exception (bool): whether to invoke the debugger on
a caught exception. Default False.
"""
all_stored_data = {} # combined stored data for *all* the Tasks
all_update_spec = {} # combined update_spec for *all* the Tasks
all_mod_spec = [] # combined mod_spec for *all* the Tasks
lp = self.launchpad
launch_dir = os.path.abspath(os.getcwd())
logdir = lp.get_logdir() if lp else None
l_logger = get_fw_logger('rocket.launcher', l_dir=logdir,
stream_level=ROCKET_STREAM_LOGLEVEL)
# check a FW job out of the launchpad
if lp:
m_fw, launch_id = lp.checkout_fw(self.fworker, launch_dir, self.fw_id)
else: # offline mode
m_fw = Firework.from_file(os.path.join(os.getcwd(), "FW.json"))
# set the run start time
fpath = zpath("FW_offline.json")
with zopen(fpath) as f_in:
d = json.loads(f_in.read())
d['started_on'] = datetime.utcnow().isoformat()
with zopen(fpath, "wt") as f_out:
f_out.write(json.dumps(d, ensure_ascii=False))
launch_id = None # we don't need this in offline mode...
if not m_fw:
print("No FireWorks are ready to run and match query! {}".format(self.fworker.query))
return False
final_state = None
ping_stop = None
btask_stops = []
try:
if '_launch_dir' in m_fw.spec and lp:
prev_dir = launch_dir
launch_dir = os.path.expandvars(m_fw.spec['_launch_dir'])
if not os.path.abspath(launch_dir):
launch_dir = os.path.normpath(os.path.join(os.getcwd(), launch_dir))
# thread-safe "mkdir -p"
try:
os.makedirs(launch_dir)
except OSError as exception:
if exception.errno != errno.EEXIST:
raise
os.chdir(launch_dir)
if not os.path.samefile(launch_dir, prev_dir):
lp.change_launch_dir(launch_id, launch_dir)
if not os.listdir(prev_dir) and REMOVE_USELESS_DIRS:
try:
os.rmdir(prev_dir)
except:
pass
recovery = m_fw.spec.get('_recovery', None)
if recovery:
recovery_dir = recovery.get('_prev_dir')
recovery_mode = recovery.get('_mode')
starting_task = recovery.get('_task_n')
all_stored_data.update(recovery.get('_all_stored_data'))
all_update_spec.update(recovery.get('_all_update_spec'))
all_mod_spec.extend(recovery.get('_all_mod_spec'))
if lp:
l_logger.log(
logging.INFO,
'Recovering from task number {} in folder {}.'.format(starting_task,
recovery_dir))
if recovery_mode == 'cp' and launch_dir != recovery_dir:
if lp:
l_logger.log(
logging.INFO,
'Copying data from recovery folder {} to folder {}.'.format(recovery_dir,
launch_dir))
distutils.dir_util.copy_tree(recovery_dir, launch_dir, update=1)
else:
starting_task = 0
files_in = m_fw.spec.get("_files_in", {})
prev_files = m_fw.spec.get("_files_prev", {})
for f in set(files_in.keys()).intersection(prev_files.keys()):
# We use zopen for the file objects for transparent handling
# of zipped files. shutil.copyfileobj does the actual copy
# in chunks that avoid memory issues.
with zopen(prev_files[f], "rb") as fin, zopen(files_in[f], "wb") as fout:
shutil.copyfileobj(fin, fout)
if lp:
message = 'RUNNING fw_id: {} in directory: {}'.\
format(m_fw.fw_id, os.getcwd())
l_logger.log(logging.INFO, message)
# write FW.json and/or FW.yaml to the directory
if PRINT_FW_JSON:
m_fw.to_file('FW.json', indent=4)
if PRINT_FW_YAML:
m_fw.to_file('FW.yaml')
my_spec = dict(m_fw.spec) # make a copy of spec, don't override original
my_spec["_fw_env"] = self.fworker.env
# set up heartbeat (pinging the server that we're still alive)
ping_stop = start_ping_launch(lp, launch_id)
# start background tasks
if '_background_tasks' in my_spec:
for bt in my_spec['_background_tasks']:
btask_stops.append(start_background_task(bt, m_fw.spec))
# execute the Firetasks!
for t_counter, t in enumerate(m_fw.tasks[starting_task:], start=starting_task):
checkpoint = {'_task_n': t_counter,
'_all_stored_data': all_stored_data,
'_all_update_spec': all_update_spec,
'_all_mod_spec': all_mod_spec}
Rocket.update_checkpoint(lp, launch_dir, launch_id, checkpoint)
if lp:
l_logger.log(logging.INFO, "Task started: %s." % t.fw_name)
if my_spec.get("_add_launchpad_and_fw_id"):
t.fw_id = m_fw.fw_id
if FWData().MULTIPROCESSING:
# hack because AutoProxy manager can't access attributes
t.launchpad = LaunchPad.from_dict(self.launchpad.to_dict())
else:
t.launchpad = self.launchpad
if my_spec.get("_add_fworker"):
t.fworker = self.fworker
try:
m_action = t.run_task(my_spec)
except BaseException as e:
traceback.print_exc()
tb = traceback.format_exc()
stop_backgrounds(ping_stop, btask_stops)
do_ping(lp, launch_id) # one last ping, esp if there is a monitor
# If the exception is serializable, save its details
if pdb_on_exception:
pdb.post_mortem()
try:
exception_details = e.to_dict()
except AttributeError:
exception_details = None
except BaseException as e:
if lp:
l_logger.log(logging.WARNING,
"Exception couldn't be serialized: %s " % e)
exception_details = None
try:
m_task = t.to_dict()
except:
m_task = None
m_action = FWAction(stored_data={'_message': 'runtime error during task',
'_task': m_task,
'_exception': {'_stacktrace': tb,
'_details': exception_details}},
exit=True)
m_action = self.decorate_fwaction(m_action, my_spec, m_fw, launch_dir)
if lp:
final_state = 'FIZZLED'
lp.complete_launch(launch_id, m_action, final_state)
else:
fpath = zpath("FW_offline.json")
with zopen(fpath) as f_in:
d = json.loads(f_in.read())
d['fwaction'] = m_action.to_dict()
d['state'] = 'FIZZLED'
d['completed_on'] = datetime.utcnow().isoformat()
with zopen(fpath, "wt") as f_out:
f_out.write(json.dumps(d, ensure_ascii=False))
return True
# read in a FWAction from a file, in case the task is not Python and cannot return
# it explicitly
if os.path.exists('FWAction.json'):
m_action = FWAction.from_file('FWAction.json')
elif os.path.exists('FWAction.yaml'):
m_action = FWAction.from_file('FWAction.yaml')
if not m_action:
m_action = FWAction()
# update the global stored data with the data to store and update from this
# particular Task
all_stored_data.update(m_action.stored_data)
all_update_spec.update(m_action.update_spec)
all_mod_spec.extend(m_action.mod_spec)
# update spec for next task as well
my_spec.update(m_action.update_spec)
for mod in m_action.mod_spec:
apply_mod(mod, my_spec)
if lp:
l_logger.log(logging.INFO, "Task completed: %s " % t.fw_name)
if m_action.skip_remaining_tasks:
break
# add job packing info if this is needed
if FWData().MULTIPROCESSING and STORE_PACKING_INFO:
all_stored_data['multiprocess_name'] = multiprocessing.current_process().name
# perform finishing operation
stop_backgrounds(ping_stop, btask_stops)
for b in btask_stops:
b.set()
do_ping(lp, launch_id) # one last ping, esp if there is a monitor
# last background monitors
if '_background_tasks' in my_spec:
for bt in my_spec['_background_tasks']:
if bt.run_on_finish:
for task in bt.tasks:
task.run_task(m_fw.spec)
m_action.stored_data = all_stored_data
m_action.mod_spec = all_mod_spec
m_action.update_spec = all_update_spec
m_action = self.decorate_fwaction(m_action, my_spec, m_fw, launch_dir)
if lp:
final_state = 'COMPLETED'
lp.complete_launch(launch_id, m_action, final_state)
else:
fpath = zpath("FW_offline.json")
with zopen(fpath) as f_in:
d = json.loads(f_in.read())
d['fwaction'] = m_action.to_dict()
d['state'] = 'COMPLETED'
d['completed_on'] = datetime.utcnow().isoformat()
with zopen(fpath, "wt") as f_out:
f_out.write(json.dumps(d, ensure_ascii=False))
return True
except LockedWorkflowError as e:
l_logger.log(logging.DEBUG, traceback.format_exc())
l_logger.log(logging.WARNING,
"Firework {} reached final state {} but couldn't complete the update of "
"the database. Reason: {}\nRefresh the WF to recover the result "
"(lpad admin refresh -i {}).".format(
self.fw_id, final_state, e, self.fw_id))
return True
except:
# problems while processing the results. high probability of malformed data.
traceback.print_exc()
stop_backgrounds(ping_stop, btask_stops)
# restore initial state to prevent the raise of further exceptions
if lp:
lp.restore_backup_data(launch_id, m_fw.fw_id)
do_ping(lp, launch_id) # one last ping, esp if there is a monitor
# the action produced by the task is discarded
m_action = FWAction(stored_data={'_message': 'runtime error during task', '_task': None,
'_exception': {'_stacktrace': traceback.format_exc(),
'_details': None}},
exit=True)
try:
m_action = self.decorate_fwaction(m_action, my_spec, m_fw, launch_dir)
except:
traceback.print_exc()
if lp:
try:
lp.complete_launch(launch_id, m_action, 'FIZZLED')
except LockedWorkflowError as e:
l_logger.log(logging.DEBUG, traceback.format_exc())
l_logger.log(logging.WARNING,
"Firework {} fizzled but couldn't complete the update of the database."
" Reason: {}\nRefresh the WF to recover the result "
"(lpad admin refresh -i {}).".format(
self.fw_id, final_state, e, self.fw_id))
return True
else:
fpath = zpath("FW_offline.json")
with zopen(fpath) as f_in:
d = json.loads(f_in.read())
d['fwaction'] = m_action.to_dict()
d['state'] = 'FIZZLED'
d['completed_on'] = datetime.utcnow().isoformat()
with zopen(fpath, "wt") as f_out:
f_out.write(json.dumps(d, ensure_ascii=False))
return True
