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_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_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_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_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_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_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_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 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_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_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__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_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_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'])
