class VaspFirework():
"""
This is an interface to Fireworks for VaspInputInterface workflows
Required Params:
-vasp_task (obj): FireTask Object used to create a Firework.
2 FireTasks are automatically created upon constructing
the VaspFirework object from a WriteVaspInputTask:
[<WriteVaspInputTask()>, <RunCustodianTask()>].
Optional Params:
-name (str): Name given to the Firework
"""
def __init__(self,vasp_task=None,name='vaspfw',handlers=None, handler_params=None, config_file=None):
self.name = name
self.handlers=handlers if handlers else []
self.handler_params=handler_params if handler_params else {}
if config_file:
config_dict = loadfn(config_file)
elif os.path.exists(os.path.join(os.environ['HOME'], 'vasp_interface_defaults.yaml')):
config_dict = loadfn(os.path.join(os.environ['HOME'], 'vasp_interface_defaults.yaml'))
else:
config_dict = {}
if config_dict:
self.custodian_opts = config_dict.get('CUSTODIAN_PARAMS', {})
if self.custodian_opts.get('handlers', []):
self.handlers.extend(self.custodian_opts.get('handlers', []))
self.handler_params.update(self.custodian_opts.get('handler_params', {}))
self.tasks=[vasp_task.input,RunCustodianTask(handlers=self.handlers,
handler_params=self.handler_params)] if isinstance(vasp_task,
VaspInputInterface) else [vasp_task]
self.Firework=Firework(self.tasks,name=self.name)
# Try to establish connection with Launchpad
try:
self.LaunchPad=LaunchPad.from_file(os.path.join(os.environ["HOME"], ".fireworks", "my_launchpad.yaml"))
except:
self.LaunchPad = None
def add_task(self, task):
'''
Function used to add another FireTask to the Firework. If
given task is a VaspInputInterface Object, it will create the WriteInputTask
as well as add another RunCustodianTask() to the FireTask list.
'''
if isinstance(task, VaspInputInterface):
self.tasks.extend([task.input,
RunCustodianTask(handlers=self.handlers,
handler_params=self.handler_params)])
else:
self.tasks.append(task)
self.Firework=Firework(self.tasks,name=self.name)
def to_file(self,file_name):
self.Firework.to_file(file_name)
def add_fw_to_launchpad(self):
if self.LaunchPad:
self.Workflow=Workflow([self.Firework])
self.LaunchPad.add_wf(self.Workflow)
else:
print("No connection to LaunchPad. \n"
"Use 'to_file(<filename>)' to write a yaml file\n"
"to manually add Firework to LaunchPad later.\n")
def copy_files_from_previous(self, *args, **kwargs):
'''
Function used to automatically insert VASPTransferTask
between WriteVaspInputTask() and RunCustodianTask() in the
FireTask list. This is used to copy files from a previous
VASP Firework Calculation or from a known directory.
Known directories require the full path of the directory.
This function should not be used to copy files from a
previous Firetask within the same Firework.
Parameters:
-mode (str): Mode to use in transferring files:
move, mv, copy, cp, copy2, copytree, copyfile
Default: 'copy'
-ignore_errors (bool): Whether to quit upon error.
Default: True
-dir (str): Directory to transfer files from if not copying
directly from a previous Firework
Default: None
Example:
fw2.copy_files_from_previous('file1', 'file2', ... , 'filen', mode='copy', ignore_errors=False)
Example to copy files from directory other than previous Firework:
fw2.copy_files_from_previous('file1', 'file2', ... , 'filen',
dir='/path/to/dir', mode='copy', ignore_errors=False)
'''
mode = kwargs.get('mode', 'copy')
ignore_errors = kwargs.get('ignore_errors', True)
files = args[0] if isinstance(args[0], list) else [k for k in args]
self.dir = kwargs.get('dir', None)
if self.dir:
self.transfer=VaspTransferTask(files=files, dir=self.dir,
mode=mode, ignore_errors=ignore_errors)
else:
self.transfer=VaspTransferTask(files=files, mode=mode, ignore_errors=ignore_errors)
if isinstance(self.tasks[-1], RunCustodianTask):
self.tasks.pop(-1)
self.add_task(self.transfer)
self.add_task(RunCustodianTask(handlers=self.handlers,
handler_params=self.handler_params))
else:
self.add_task(self.transfer)
def add_handler(self, handler, **kwargs):
'''
Member function to add handler and handler options to
all Fireworks in the VaspFirework.tasks list.
Example (assuming fw1 is the firework object):
fw1.add_handler('WalltimeHandler', wall_time=3600, buffer_time=300)
fw1.add_handler('FrozenJobErrorHandler')
'''
for i,j in enumerate(self.Firework.tasks):
if isinstance(j, RunCustodianTask):
cust_handlers = j.get('handlers', [])
cust_params = j.get('handler_params', {})
if handler not in cust_handlers:
j['handlers'].append(handler)
j['handler_params'].update(kwargs)
self.Firework.spec['_tasks']=[t.to_dict() for t in
self.Firework.tasks] #re-initialize FW.spec
from fireworks import Firework, LaunchPad, ScriptTask
from fireworks.core.rocket_launcher import launch_rocket, rapidfire
# set up the LaunchPad and reset it
launchpad = LaunchPad(strm_lvl='WARNING') # set messaging lowest level to WARNING
launchpad.reset('', require_password=False)
# create the Firework consisting of a single task
firetask = ScriptTask.from_str('cd /projects/development/LDRDSANS/fireworks/localtest; ./test_cluster.sh')
firework = Firework(firetask)
fw_yaml = firework.to_file("my_firework.yaml") # save to yaml file, and get the string representation
fw_json = firework.to_file("my_firework.json") # save to json file, and get the string representation
# store workflow and launch it locally
launchpad.add_wf(firework)
launch_rocket(launchpad) # same as "rlaunch singleshot"
#rapidfire(launchpad, FWorker(), strm_lvl='WARNING') # same as "rlaunch rapidfire"
# loading from file
# any class in FireWorks that subclasses FWSerializable has methods from_file and to_file
#firework = Firework.from_file("fw_test.yaml")
#fw_yaml = Firework.from_file("fw_test.json")
