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