def test_basic_fw(self): test1 = ScriptTask.from_str("python -c 'print \"test1\"'", {'store_stdout': True}) fw = FireWork(test1) self.lp.add_wf(fw) launch_rocket(self.lp) time.sleep(3) self.assertEqual(self.lp.get_launch_by_id(1).action.stored_data['stdout'], 'test1\n')
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(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}) 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 test_postproc_exception(self): fw = Firework(MalformedAdditionTask()) self.lp.add_wf(fw) launch_rocket(self.lp, self.fworker) fw = self.lp.get_fw_by_id(1) self.assertEqual(fw.state, 'FIZZLED')
def test_parallel(self): n_procs = 10 self.mc.configure(wf_creator=wf_creator_accuracy, dimensions=self.dims_parallel, opt_label='test_parallel', maximize=True, **common_kwargs) for i in range(n_procs): # Assume the worst case, with n_procs forced duplicates launchpad.add_wf(wf_creator_accuracy([1, 5, 3])) launch_multiprocess(launchpad, None, 'INFO', 12, n_procs, 0) try: for _ in range(10): launch_rocket(launchpad) except ExhaustedSpaceError: pass self.assertEqual(125, self.c.count_documents({'y': {'$exists': 1}})) all_x, all_y = self.mc.fetch_matrices(include_reserved=True) self.assertEqual(len(all_x), 125) self.assertEqual(len(all_x), 125) all_x, all_y = self.mc.fetch_matrices(include_reserved=False) self.assertGreaterEqual(len(all_x), 114) self.assertGreaterEqual(len(all_y), 114)
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 test_tracker_failed_fw(self): """ Add a bad firetask to workflow and test the tracking """ self._teardown([self.dest1]) try: fts = [] for i in range(5,50): ft = ScriptTask.from_str('echo "' + str(i) + '" >> '+ self.dest1, {'store_stdout':True}) fts.append(ft) fts.append(ScriptTask.from_str('cat 4 >> ' + self.dest1)) for i in range(51,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=21, name='test_fw') self.lp.add_wf(fw) try: print("===========================================") print("Bad rocket launched. The failure below is OK") print("===========================================") launch_rocket(self.lp, self.fworker) except: pass self.assertEqual('48\n49',self.tracker1.track_file()) finally: self._teardown([self.dest1])
def rlaunch(): m_description = 'This program launches one or more Rockets. A Rocket grabs a job from the central database and ' \ 'runs it. The "single-shot" option launches a single Rocket, ' \ 'whereas the "rapidfire" option loops until all FireWorks are completed.' parser = ArgumentParser(description=m_description) subparsers = parser.add_subparsers(help='command', dest='command') single_parser = subparsers.add_parser('singleshot', help='launch a single Rocket') rapid_parser = subparsers.add_parser('rapidfire', help='launch multiple Rockets (loop until all FireWorks complete)') single_parser.add_argument('-f', '--fw_id', help='specific fw_id to run', default=None, type=int) single_parser.add_argument('--offline', help='run in offline mode (FW.json required)', action='store_true') rapid_parser.add_argument('--nlaunches', help='num_launches (int or "infinite"; default 0 is all jobs in DB)', default=0) rapid_parser.add_argument('--sleep', help='sleep time between loops (secs)', default=None, type=int) parser.add_argument('-l', '--launchpad_file', help='path to launchpad file', default=LAUNCHPAD_LOC) parser.add_argument('-w', '--fworker_file', help='path to fworker file', default=FWORKER_LOC) parser.add_argument('-c', '--config_dir', help='path to a directory containing the config file (used if -l, -w unspecified)', default=CONFIG_FILE_DIR) parser.add_argument('--loglvl', help='level to print log messages', default='INFO') parser.add_argument('-s', '--silencer', help='shortcut to mute log messages', action='store_true') args = parser.parse_args() signal.signal(signal.SIGINT, handle_interrupt) # graceful exist on ^C if not args.launchpad_file and os.path.exists(os.path.join(args.config_dir, 'my_launchpad.yaml')): args.launchpad_file = os.path.join(args.config_dir, 'my_launchpad.yaml') if not args.fworker_file and os.path.exists(os.path.join(args.config_dir, 'my_fworker.yaml')): args.fworker_file = os.path.join(args.config_dir, 'my_fworker.yaml') args.loglvl = 'CRITICAL' if args.silencer else args.loglvl if args.command == 'singleshot' and args.offline: launchpad = None else: launchpad = LaunchPad.from_file(args.launchpad_file) if args.launchpad_file else LaunchPad(strm_lvl=args.loglvl) if args.fworker_file: fworker = FWorker.from_file(args.fworker_file) else: fworker = FWorker() # prime addr lookups _log = get_fw_logger("rlaunch", stream_level="INFO") _log.info("Hostname/IP lookup (this will take a few seconds)") get_my_host() get_my_ip() if args.command == 'rapidfire': rapidfire(launchpad, fworker, None, args.nlaunches, -1, args.sleep, args.loglvl) else: launch_rocket(launchpad, fworker, args.fw_id, args.loglvl)
def test_add_lp_and_fw_id(self): fw1 = Firework([DummyLPTask()], spec={"_add_launchpad_and_fw_id": True}) self.lp.add_wf(fw1) launch_rocket(self.lp, self.fworker) self.assertEqual( self.lp.get_launch_by_id(1).action.stored_data['fw_id'], 1) self.assertIsNotNone( self.lp.get_launch_by_id(1).action.stored_data['host'])
def test_dupefinder(self): test1 = ScriptTask.from_str("python -c 'print(\"test1\")'", {'store_stdout': True}) fw = Firework(test1, {"_dupefinder": DupeFinderExact()}) self.lp.add_wf(fw) self.lp.add_wf(fw) launch_rocket(self.lp, self.fworker) launch_rocket(self.lp, self.fworker) self.assertEqual(self.lp.launches.count(), 1)
def test_basic_fw(self): test1 = ScriptTask.from_str("python -c 'print(\"test1\")'", {'store_stdout': True}) fw = Firework(test1) self.lp.add_wf(fw) launch_rocket(self.lp, self.fworker) self.assertEqual( self.lp.get_launch_by_id(1).action.stored_data['stdout'], 'test1\n')
def test_delete_fw(self): test1 = ScriptTask.from_str("python -c 'print(\"test1\")'", {'store_stdout': True}) fw = Firework(test1) self.lp.add_wf(fw) launch_rocket(self.lp, self.fworker) self.assertEqual(self.lp.get_launch_by_id(1).action.stored_data[ 'stdout'], 'test1\n') self.lp.delete_wf(fw.fw_id) self.assertRaises(ValueError, self.lp.get_fw_by_id, fw.fw_id) self.assertRaises(ValueError, self.lp.get_launch_by_id, 1)
def test_serializable_exception(self): error_test_dict = {'error': 'description', 'error_code': 1} fw = Firework(ExceptionTestTask(exc_details=error_test_dict)) self.lp.add_wf(fw) launch_rocket(self.lp, self.fworker) # Check data in the exception fw = self.lp.get_fw_by_id(1) launches = fw.launches self.assertEqual(launches[0].action.stored_data['_exception']['_details'], error_test_dict)
def test_full_opt_fw_writes_isif_setting_takes_effect(patch_pmg_psp_dir, launch_dir, lpad, fworker): fw = OptimizeFW(STRUCT, isif=7, job_type='full_opt_run', vasp_cmd=None) wf = Workflow([fw]) lpad.add_wf(wf) launch_rocket(lpad, fworker=fworker) incar = Incar.from_file(os.path.join(launch_dir, 'INCAR.gz')) desired_parameters = {'ISIF': 7} assert all([incar[k] == v for k, v in desired_parameters.items()])
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.assertEqual(self.lp.get_launch_by_id(1).action.stored_data['stdout'], 'test1\n') launch_rocket(self.lp) self.assertEqual(self.lp.get_launch_by_id(2).action.stored_data['stdout'], 'test2\n')
def test_fw_env(self): t = DummyFWEnvTask() fw = Firework(t) self.lp.add_wf(fw) launch_rocket(self.lp, self.fworker) self.assertEqual(self.lp.get_launch_by_id(1).action.stored_data['data'], "hello") self.lp.add_wf(fw) launch_rocket(self.lp, FWorker(env={"hello": "world"})) self.assertEqual(self.lp.get_launch_by_id(2).action.stored_data[ 'data'], "world")
def test_multi_fw(self): test1 = ScriptTask.from_str("python -c 'print(\"test1\")'", {'store_stdout': True}) test2 = ScriptTask.from_str("python -c 'print(\"test2\")'", {'store_stdout': True}) fw = Firework([test1, test2]) self.lp.add_wf(fw) launch_rocket(self.lp, self.fworker) self.assertEqual( self.lp.get_launch_by_id(1).action.stored_data['stdout'], "test2\n")
def test_multi_fw(self): test1 = ScriptTask.from_str("python -c 'print(\"test1\")'", {"store_stdout": True}) test2 = ScriptTask.from_str("python -c 'print(\"test2\")'", {"store_stdout": True}) fw = Firework([test1, test2]) self.lp.add_wf(fw) launch_rocket(self.lp, self.fworker) self.assertEqual( self.lp.get_launch_by_id(1).action.stored_data["stdout"], "test2\n")
def test_serializable_exception(self): error_test_dict = {"error": "description", "error_code": 1} fw = Firework(ExceptionTestTask(exc_details=error_test_dict)) self.lp.add_wf(fw) launch_rocket(self.lp, self.fworker) # Check data in the exception fw = self.lp.get_fw_by_id(1) launches = fw.launches self.assertEqual( launches[0].action.stored_data["_exception"]["_details"], error_test_dict)
def test_duplicate_delete_fw(self): test1 = ScriptTask.from_str("python -c 'print(\"test1\")'", {'store_stdout': True}) fw = Firework(test1, {"_dupefinder": DupeFinderExact()}) self.lp.add_wf(fw) self.lp.add_wf(fw) launch_rocket(self.lp, self.fworker) launch_rocket(self.lp, self.fworker) self.lp.delete_wf(2) self.assertRaises(ValueError, self.lp.get_fw_by_id, 2) self.assertEqual(self.lp.get_launch_by_id(1).action.stored_data[ 'stdout'], 'test1\n')
def test_multi_fw(self): test1 = ScriptTask.from_str("python -c 'print \"test1\"'", {'store_stdout': True}) test2 = ScriptTask.from_str("python -c 'print \"test2\"'", {'store_stdout': True}) fw = FireWork([test1, test2]) self.lp.add_wf(fw) launch_rocket(self.lp) time.sleep(3) self.assertEqual( self.lp.get_launch_by_id(1).action.stored_data['stdout'], 'test2\n')
def test_alternate_builtin_predictor(self): gp = "GaussianProcessRegressor" self.mc.configure(wf_creator=wf_creator_basic, dimensions=self.dims_basic, predictor=gp) config = self.c.find_one({"doctype": "config"}) self.assertEqual(config["predictor"], gp) launchpad.add_wf(wf_creator_basic([5, 11, 'blue'])) launch_rocket(launchpad) done = self.c.find_one({'y': {'$exists': 1, '$ne': 'reserved'}}) self.assertEqual(done["predictor"], f"{gp} with acquisition: Expected Improvement")
def test_add_wf_to_fws_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.add_wf_to_fws(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_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_archive_wf(self): # Run a firework before archiving Zeus launch_rocket(self.lp, self.fworker) # archive Workflow containing Zeus. self.lp.archive_wf(self.zeus_fw_id) # Ensure the states are sync wf = self.lp.get_wf_by_fw_id_lzyfw(self.zeus_fw_id) fws = wf.id_fw for fw_id in wf.fw_states: fw_state = fws[fw_id].state fw_cache_state = wf.fw_states[fw_id] self.assertEqual(fw_state, fw_cache_state)
def test_delete_wf(self): # Run a firework before deleting Zeus launch_rocket(self.lp, self.fworker) # Delete workflow containing Zeus. self.lp.delete_wf(self.zeus_fw_id) # Check if any fireworks and the workflow are available with self.assertRaises(ValueError): self.lp.get_wf_by_fw_id(self.zeus_fw_id) fw_ids = self.lp.get_fw_ids() self.assertFalse(fw_ids) wf_ids = self.lp.get_wf_ids() self.assertFalse(wf_ids)
def basic_fw_ex(): print('--- BASIC FIREWORK EXAMPLE ---') # setup launchpad = setup() # add FireWork firetask = ScriptTask.from_str('echo "howdy, your job launched successfully!"') firework = FireWork(firetask) launchpad.add_wf(firework) # launch Rocket launch_rocket(launchpad, FWorker())
def test_full_opt_fw_writes_correct_fw_for_UIS_in_set_constructor( patch_pmg_psp_dir, launch_dir, lpad, fworker): s = PRLRelaxSet(STRUCT, user_incar_settings={'ISIF': 4}) fw = OptimizeFW(STRUCT, vasp_input_set=s, job_type='full_opt_run', vasp_cmd=None) wf = Workflow([fw]) lpad.add_wf(wf) launch_rocket(lpad, fworker=fworker) incar = Incar.from_file(os.path.join(launch_dir, 'INCAR.gz')) desired_parameters = {'ISIF': 4} assert all([incar[k] == v for k, v in desired_parameters.items()])
def basic_fw_ex(): print("--- BASIC FIREWORK EXAMPLE ---") # setup launchpad = setup() # add Firework firetask = ScriptTask.from_str('echo "howdy, your job launched successfully!"') firework = Firework(firetask) launchpad.add_wf(firework) # launch Rocket launch_rocket(launchpad, FWorker())
def test_complex(self): self.lp.reset(password=None, require_password=False) self.lp.add_wf(wf_creator_complex([5, 11, 'blue'], self.lp)) for _ in range(10): launch_rocket(self.lp) col = self.db.test_complex self.assertEqual(col.count_documents({}), 4) self.assertEqual(col.count_documents({'y': 'reserved'}), 1) # should return one doc, for first WF self.assertEqual(col.count_documents({'index': 1}), 1) # loop 1 # should return one doc, for second WF self.assertEqual(col.count_documents({'index': 2}), 1) # loop 2
def rlaunch(): m_description = 'This program launches one or more Rockets. A Rocket grabs a job from the central database and ' \ 'runs it. The "single-shot" option launches a single Rocket, ' \ 'whereas the "rapidfire" option loops until all FireWorks are completed.' parser = ArgumentParser(description=m_description) subparsers = parser.add_subparsers(help='command', dest='command') single_parser = subparsers.add_parser('singleshot', help='launch a single Rocket') rapid_parser = subparsers.add_parser('rapidfire', help='launch multiple Rockets (loop until all FireWorks complete)') single_parser.add_argument('-f', '--fw_id', help='specific fw_id to run', default=None, type=int) single_parser.add_argument('--offline', help='run in offline mode (FW.json required)', action='store_true') rapid_parser.add_argument('--nlaunches', help='num_launches (int or "infinite"; default 0 is all jobs in DB)', default=0) rapid_parser.add_argument('--sleep', help='sleep time between loops (secs)', default=None, type=int) parser.add_argument('-l', '--launchpad_file', help='path to launchpad file', default=LAUNCHPAD_LOC) parser.add_argument('-w', '--fworker_file', help='path to fworker file', default=FWORKER_LOC) parser.add_argument('-c', '--config_dir', help='path to a directory containing the config file (used if -l, -w unspecified)', default=CONFIG_FILE_DIR) parser.add_argument('--loglvl', help='level to print log messages', default='INFO') parser.add_argument('-s', '--silencer', help='shortcut to mute log messages', action='store_true') args = parser.parse_args() if not args.launchpad_file and os.path.exists(os.path.join(args.config_dir, 'my_launchpad.yaml')): args.launchpad_file = os.path.join(args.config_dir, 'my_launchpad.yaml') if not args.fworker_file and os.path.exists(os.path.join(args.config_dir, 'my_fworker.yaml')): args.fworker_file = os.path.join(args.config_dir, 'my_fworker.yaml') args.loglvl = 'CRITICAL' if args.silencer else args.loglvl if args.command == 'singleshot' and args.offline: launchpad = None else: launchpad = LaunchPad.from_file(args.launchpad_file) if args.launchpad_file else LaunchPad(strm_lvl=args.loglvl) if args.fworker_file: fworker = FWorker.from_file(args.fworker_file) else: fworker = FWorker() if args.command == 'rapidfire': rapidfire(launchpad, fworker, None, args.nlaunches, -1, args.sleep, args.loglvl) else: launch_rocket(launchpad, fworker, args.fw_id, args.loglvl)
def test_complex(self): self.mc.configure(wf_creator=wf_creator_complex, dimensions=self.dims_basic, **common_kwargs) launchpad.add_wf(wf_creator_complex([5, 11, 'blue'])) for _ in range(10): launch_rocket(launchpad) self.assertEqual(self.c.count_documents({}), 5) self.assertEqual(self.c.count_documents({'y': 'reserved'}), 1) # should return one doc, for first WF self.assertEqual(self.c.count_documents({'index': 1}), 1) # loop 1 # should return one doc, for second WF self.assertEqual(self.c.count_documents({'index': 2}), 1) # loop 2
def test_full_opt_fw_writes_isif_setting_does_take_effects_with_VIS( patch_pmg_psp_dir, launch_dir, lpad, fworker): s = PRLRelaxSet(STRUCT) fw = PRLOptimizeFW(STRUCT, vasp_input_set=s, isif=5, job_type='full_opt_run', vasp_cmd=None) wf = Workflow([fw]) lpad.add_wf(wf) launch_rocket(lpad, fworker=fworker) incar = Incar.from_file(os.path.join(launch_dir, 'INCAR.gz')) desired_parameters = {'ISIF': 5} assert all([incar[k] == v for k, v in desired_parameters.items()])
def test__recover_completed(self): fw, launch_id = self.lp.reserve_fw(self.fworker, self.launch_dir) fw = self.lp.get_fw_by_id(1) with cd(self.launch_dir): setup_offline_job(self.lp, fw, launch_id) # launch rocket without launchpad to trigger offline mode launch_rocket(launchpad=None, fworker=self.fworker, fw_id=1) self.assertIsNone(self.lp.recover_offline(launch_id)) fw = self.lp.get_fw_by_id(launch_id) self.assertEqual(fw.state, 'COMPLETED')
def test_get_z(self): self.lp.reset(password=None, require_password=False) self.lp.add_wf(wf_creator_get_z([5, 11, 'blue'], self.lp)) for _ in range(2): launch_rocket(self.lp) col = self.db.test_get_z loop1 = col.find_one({'index': 1}) loop2 = col.find_one({'index': 2}) self.assertEqual(col.find({}).count(), 4) self.assertEqual(loop1['x'], [5, 11, 'blue']) self.assertEqual(loop1['z'], [25.0, 121.0]) self.assertEqual(loop2['x'], [3, 12.0, 'green']) self.assertEqual(loop2['z'], [9, 144.0])
def test_basic(self): self.lp.reset(password=None, require_password=False) self.lp.add_wf(wf_creator_basic([5, 11, 'blue'], self.lp)) launch_rocket(self.lp) col = self.db.test_basic manager = col.find_one({'y': {'$exists': 0}}) done = col.find_one({'y': {'$exists': 1, '$ne': 'reserved'}}) reserved = col.find_one({'y': 'reserved'}) self.assertEqual(col.count_documents({}), 3) self.assertEqual(manager['lock'], None) self.assertEqual(manager['queue'], []) self.assertEqual(done['x'], [5, 11, 'blue']) self.assertEqual(done['index'], 1)
def test_duplicates(self): self.lp.reset(password=None, require_password=False) self.lp.add_wf(wf_creator_duplicates([5, 11, 'blue'], self.lp)) for _ in range(2): launch_rocket(self.lp) col = self.db.test_duplicates self.assertEqual(col.count_documents({}), 4) self.assertEqual(col.count_documents({'y': 'reserved'}), 1) # should return one doc, for the first WF self.assertEqual(col.count_documents({'x': [5, 11, 'blue']}), 1) # should return one doc, for the second WF # no duplicates are in the db self.assertEqual(col.count_documents({'x': [3, 12, 'green']}), 1)
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_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_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_basic(self): self.mc.configure(wf_creator=wf_creator_basic, dimensions=self.dims_basic, **common_kwargs) launchpad.add_wf(wf_creator_basic([5, 11, 'blue'])) launch_rocket(launchpad) manager = self.c.find_one({'doctype': "manager"}) done = self.c.find_one({'y': {'$exists': 1, '$ne': 'reserved'}}) reserved = {'y': 'reserved'} self.assertEqual(self.c.count_documents({}), 4) self.assertEqual(manager['lock'], None) self.assertEqual(manager['queue'], []) self.assertEqual(done['x'], [5, 11, 'blue']) self.assertEqual(done['index'], 1) self.assertEqual(self.c.count_documents(reserved), 1)
def test_duplicate_delete_fw(self): test1 = ScriptTask.from_str("python -c 'print(\"test1\")'", {'store_stdout': True}) fw = Firework(test1, {"_dupefinder": DupeFinderExact()}) self.lp.add_wf(fw) self.lp.add_wf(fw) launch_rocket(self.lp, self.fworker) # TODO: if test keeps failing on Travis, add an explicit check of nlaunches>0 in the database here # this will ensure the first Rocket is actually in the DB launch_rocket(self.lp, self.fworker) run_id = self.lp.get_launch_by_id(1).fw_id del_id = 1 if run_id == 2 else 2 self.lp.delete_wf(del_id) self.assertRaises(ValueError, self.lp.get_fw_by_id, del_id) self.assertEqual(self.lp.get_launch_by_id(1).action.stored_data[ 'stdout'], 'test1\n')
def test_basic(self): self.mc.configure(wf_creator=wf_creator_basic, dimensions=self.dims_basic, **common_kwargs) launchpad.add_wf(wf_creator_basic([5, 11, "blue"])) launch_rocket(launchpad) manager = self.c.find_one({"doctype": "manager"}) done = self.c.find_one({"y": {"$exists": 1, "$ne": "reserved"}}) reserved = {"y": "reserved"} self.assertEqual(self.c.count_documents({}), 4) self.assertEqual(manager["lock"], None) self.assertEqual(manager["queue"], []) self.assertEqual(done["x"], [5, 11, "blue"]) self.assertEqual(done["index"], 1) self.assertEqual(self.c.count_documents(reserved), 1)
def test_archive_wf(self): # Run a firework before archiving Zeus launch_rocket(self.lp, self.fworker) # archive Workflow containing Zeus. Ensure all are archived self.lp.archive_wf(self.zeus_fw_id) archived_ids = set(self.lp.get_fw_ids({'state':'ARCHIVED'})) self.assertEqual(archived_ids, self.all_ids) # Try to launch the fireworks and check if any are launched rapidfire(self.lp, self.fworker,m_dir=MODULE_DIR) fws_completed = set(self.lp.get_fw_ids({'state':'COMPLETED'})) self.assertFalse(fws_completed) # Query for provenance fw = self.lp.get_fw_by_id(self.zeus_fw_id) self.assertEqual(fw.state,'ARCHIVED')
def test_fix_db_inconsistencies_completed(self): class RocketProcess(Process): def __init__(self, lpad, fworker, fw_id): super(self.__class__,self).__init__() self.lpad = lpad self.fworker = fworker self.fw_id = fw_id def run(self): launch_rocket(self.lpad, self.fworker, fw_id=self.fw_id) # Launch the slow firework in a separate process rp = RocketProcess(self.lp, self.fworker, fw_id=1) rp.start() time.sleep(1) launch_rocket(self.lp, self.fworker, fw_id=2) # wait for the slow to complete rp.join() fast_fw = self.lp.get_fw_by_id(2) if fast_fw.state == 'FIZZLED': stacktrace = self.lp.launches.find_one( {'fw_id': 2}, {'action.stored_data._exception._stacktrace': 1})['action']['stored_data']['_exception']['_stacktrace'] if 'SkipTest' in stacktrace: self.skipTest("The test didn't run correctly") self.assertEqual(fast_fw.state, 'RUNNING') child_fw = self.lp.get_fw_by_id(3) self.assertTrue("SlowAdditionTask" in child_fw.spec) self.assertFalse("WaitWFLockTask" in child_fw.spec) self.lp._refresh_wf(fw_id=2) child_fw = self.lp.get_fw_by_id(3) self.assertTrue("WaitWFLockTask" in child_fw.spec) fast_fw = self.lp.get_fw_by_id(2) self.assertEqual(fast_fw.state, 'COMPLETED')
def test_dupefinder(self): test1 = ScriptTask.from_str("python -c 'print(\"test1\")'", {'store_stdout': True}) fw = Firework(test1, {"_dupefinder": DupeFinderExact()}) self.lp.add_wf(fw) self.lp.add_wf(fw) launch_rocket(self.lp, self.fworker) # TODO: if test keeps failing on Travis, add an explicit check of nlaunches>0 in the database here # this will ensure the first Rocket is actually in the DB launch_rocket(self.lp, self.fworker) if self.lp.launches.count() > 1: print("TOO MANY LAUNCHES FOUND!") print("--------") for d in self.lp.launches.find(): print(d) print("--------") self.assertEqual(self.lp.launches.count(), 1)
def test_custom_predictor(self): self.mc.configure(wf_creator=wf_creator_basic, dimensions=self.dims_basic, predictor=custom_predictor) launchpad.add_wf(wf_creator_basic([5, 11, 'blue'])) launch_rocket(launchpad) manager = self.c.find_one({'doctype': "manager"}) done = self.c.find_one({'y': {'$exists': 1, '$ne': 'reserved'}}) reserved = self.c.find_one({'y': 'reserved'}) self.assertEqual(self.c.count_documents({}), 4) self.assertEqual(manager['lock'], None) self.assertEqual(manager['queue'], []) self.assertEqual(done['x'], [5, 11, 'blue']) self.assertEqual(done['x_new'], [3, 12, 'green']) self.assertEqual(done['index'], 1) self.assertEqual(reserved['x'], [3, 12, 'green'])
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_fix_db_inconsistencies_fizzled(self): class RocketProcess(Process): def __init__(self, lpad, fworker, fw_id): super(self.__class__, self).__init__() self.lpad = lpad self.fworker = fworker self.fw_id = fw_id def run(self): launch_rocket(self.lpad, self.fworker, fw_id=self.fw_id) self.lp.update_spec([2], {"fizzle": True}) # Launch the slow firework in a separate process rp = RocketProcess(self.lp, self.fworker, fw_id=1) rp.start() time.sleep(1) launch_rocket(self.lp, self.fworker, fw_id=2) # wait for the slow to complete rp.join() fast_fw = self.lp.get_fw_by_id(2) if fast_fw.state == "FIZZLED": stacktrace = self.lp.launches.find_one({"fw_id": 2}, {"action.stored_data._exception._stacktrace": 1})[ "action" ]["stored_data"]["_exception"]["_stacktrace"] if "SkipTest" in stacktrace: self.skipTest("The test didn't run correctly") self.assertEqual(fast_fw.state, "RUNNING") child_fw = self.lp.get_fw_by_id(3) self.assertTrue("SlowAdditionTask" in child_fw.spec) self.assertFalse("WaitWFLockTask" in child_fw.spec) self.lp._refresh_wf(fw_id=2) fast_fw = self.lp.get_fw_by_id(2) self.assertEqual(fast_fw.state, "FIZZLED")
def test_duplicates(self): self.mc.configure(wf_creator=wf_creator_basic, dimensions=self.dims_basic, duplicate_check=True, tolerances=[0, 1e-6, None], predictor=custom_predictor, acq=None ) launchpad.add_wf(wf_creator_basic([5, 11, 'blue'])) for _ in range(2): launch_rocket(launchpad) self.assertEqual(self.c.count_documents({}), 5) self.assertEqual(self.c.count_documents({'y': 'reserved'}), 1) # should return one doc, for the first WF self.assertEqual(self.c.count_documents({'x': [5, 11, 'blue']}), 1) # should return one doc, for the second WF # no duplicates are in the db self.assertEqual(self.c.count_documents({'x': [3, 12, 'green']}), 1)
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'])