Source code for idmtools_models.templated_script_task
"""Provides the TemplatedScriptTask.Copyright 2021, Bill & Melinda Gates Foundation. All rights reserved."""importcopyimportinspectimportosfromcontextlibimportsuppressfromdataclassesimportdataclass,fieldfromfunctoolsimportpartialfromloggingimportgetLogger,DEBUGfromtypingimportList,Callable,Type,Dict,Any,Union,TYPE_CHECKINGfromjinja2importEnvironmentfromidmtools.assetsimportAssetCollection,Assetfromidmtools.entitiesimportCommandLinefromidmtools.entities.itaskimportITaskfromidmtools.entities.iworkflow_itemimportIWorkflowItemfromidmtools.entities.simulationimportSimulationfromidmtools.registry.task_specificationimportTaskSpecificationifTYPE_CHECKING:# pragma: no coverfromidmtools.entities.iplatformimportIPlatformlogger=getLogger(__name__)LINUX_DICT_TO_ENVIRONMENT="""{% for key, value in vars.items() %}export {{key}}="{{value}}"{% endfor %}"""WINDOWS_DICT_TO_ENVIRONMENT="""{% for key, value in vars.items() %}set {{key}}="{{value}}"{% endfor %}"""# Define our common windows templatesWINDOWS_BASE_WRAPPER="""echo Running %*%*"""WINDOWS_PYTHON_PATH_WRAPPER="""set PYTHONPATH=%cd%\\Assets\\site-packages\\;%cd%\\Assets\\;%PYTHONPATH%{}""".format(WINDOWS_BASE_WRAPPER)WINDOWS_DICT_TO_ENVIRONMENT=WINDOWS_DICT_TO_ENVIRONMENT+WINDOWS_BASE_WRAPPER# Define our linux common scriptsLINUX_BASE_WRAPPER="""echo Running args $@\"$@\""""LINUX_PYTHON_PATH_WRAPPER="""export PYTHONPATH=$(pwd)/Assets/site-packages:$(pwd)/Assets/:$PYTHONPATH{}""".format(LINUX_BASE_WRAPPER)LINUX_DICT_TO_ENVIRONMENT=LINUX_DICT_TO_ENVIRONMENT+LINUX_BASE_WRAPPER
[docs]@dataclass()classTemplatedScriptTask(ITask):""" Defines a task to run a script using a template. Best suited to shell scripts. Examples: In this example, we add modify the Python Path using TemplatedScriptTask and LINUX_PYTHON_PATH_WRAPPER .. literalinclude:: ../../examples/cookbook/python/python-path/python-path.py In this example, we modify environment variable using TemplatedScriptTask and LINUX_DICT_TO_ENVIRONMENT .. literalinclude:: ../../examples/cookbook/environment/variables/environment-vars.py """#: Name of scriptscript_path:str=field(default=None,metadata={"md":True})#: If platform requires path to script executing binary(ie /bin/bash)script_binary:str=field(default=None,metadata={"md":True})#: The template contentstemplate:str=field(default=None,metadata={"md":True})#: The template file. You can only use either template or template_file at oncetemplate_file:str=field(default=None,metadata={"md":True})#: Controls whether a template should be an experiment or a simulation level assettemplate_is_common:bool=field(default=True,metadata={"md":True})#: Template variables used for rendering the template# Note: large amounts of variables will increase metadata sizevariables:Dict[str,Any]=field(default_factory=dict,metadata={"md":True})#: Platform Path Separator. For Windows execution platforms, use \, otherwise use the default of /path_sep:str=field(default='/',metadata={"md":True})#: Extra arguments to add to the command lineextra_command_arguments:str=field(default='',metadata={"md":True})#: Hooks to gather common assetsgather_common_asset_hooks:List[Callable[[ITask],AssetCollection]]=field(default_factory=list)#: Hooks to gather transient assetsgather_transient_asset_hooks:List[Callable[[ITask],AssetCollection]]=field(default_factory=list)def__post_init__(self):"""Constructor."""super().__post_init__()ifself.script_pathisNoneandself.template_fileisNone:raiseValueError("Either script name or template file is required")elifself.script_pathisNone:# Get name from the fileself.script_path=os.path.basename(self.template_file)ifself.templateisNoneandself.template_fileisNone:raiseValueError("Either template or template_file is required")ifself.template_fileandnotos.path.exists(self.template_file):raiseFileNotFoundError(f"Could not find the file the template file {self.template_file}")# is the template a common(experiment) asset?ifself.template_is_common:iflogger.isEnabledFor(DEBUG):logger.debug("Adding common asset hook")# add hook to render it to common asset hookshook=partial(TemplatedScriptTask._add_template_to_asset_collection,asset_collection=self.common_assets)self.gather_common_asset_hooks.append(hook)else:# it must be a simulation level assetiflogger.isEnabledFor(DEBUG):logger.debug("Adding transient asset hook")# create hook to render it to our transient assetshook=partial(TemplatedScriptTask._add_template_to_asset_collection,asset_collection=self.transient_assets)# add the hookself.gather_transient_asset_hooks.append(hook)@staticmethoddef_add_template_to_asset_collection(task:'TemplatedScriptTask',asset_collection:AssetCollection)-> \
AssetCollection:""" Add our rendered template to the asset collection. Args: task: Task to add asset_collection:Asset collection Returns: Asset Collection with template added """# setup our jinja environmentenv=Environment()# try to load template from string or fileiftask.template:template=env.from_string(task.template)else:template=env.get_template(task.template_file)iflogger.isEnabledFor(DEBUG):logger.debug(f"Rendering Script template: {template}")# render the templateresult=template.render(vars=task.variables,env=os.environ)iflogger.isEnabledFor(DEBUG):logger.debug(f"Rendered Template: {result}")# create an assetasset=Asset(filename=task.script_path,content=result)asset_collection.add_asset(asset)returnasset_collection
[docs]defgather_common_assets(self)->AssetCollection:""" Gather common(experiment-level) assets for task. Returns: AssetCollection containing common assets """# TODO validate hooks have expected return typeac=AssetCollection()forxinself.gather_common_asset_hooks:ac+=x(self)returnac
[docs]defreload_from_simulation(self,simulation:Simulation):""" Reload a templated script task. When reloading, you will only have the rendered template available. Args: simulation: Returns: None """# check experiment level assets for our scriptifsimulation.parent.assets:# prep new asset collection in case we have to remove our asset from the experiment levelnew_assets=AssetCollection()for_i,assetinenumerate(simulation.parent.assets.assets):# is it our script?ifasset.filename!=self.script_pathandasset.absolute_path!=self.script_path:# nope keep itnew_assets.add_asset(asset)# set filtered assets back to parentsimulation.parent.assets=new_assets
[docs]defpre_creation(self,parent:Union[Simulation,IWorkflowItem],platform:'IPlatform'):""" Before creating simulation, we need to set our command line. Args: parent: Parent object platform: Platform item is being ran on Returns: """# are we experiment or simulation level asset?ifself.script_binary:sn=self.script_binary+' 'else:sn=''ifself.template_is_common:sn+=platform.join_path(platform.common_asset_path,self.script_path)else:sn+=self.script_path# set the command line to the rendered scriptself.command=CommandLine.from_string(sn)ifself.path_sep!="/":self.command.executable=self.command.executable.replace("/",self.path_sep)self.command.is_windows=True# set any extra argumentsifself.extra_command_arguments:other_command=CommandLine.from_string(self.extra_command_arguments)self.command._args.append(other_command.executable)ifother_command._options:self.command._args+=other_command._optionsifother_command._args:self.command._args+=other_command._argssuper().pre_creation(parent,platform)
[docs]@dataclass()classScriptWrapperTask(ITask):""" Allows you to wrap a script with another script. See Also: :py:class:`idmtools_models.templated_script_task.TemplatedScriptTask` Raises: ValueError if the template Script Task is not defined """template_script_task:TemplatedScriptTask=field(default=None)task:ITask=field(default=None)def__post_init__(self):"""Constructor."""ifself.template_script_taskisNone:raiseValueError("Template Script Task is required")ifself.taskisNone:raiseValueError("Task is required")ifisinstance(self.task,dict):self.task=self.from_dict(self.task)ifisinstance(self.template_script_task,dict):self.template_script_task=self.from_dict(self.template_script_task)
[docs]@staticmethoddeffrom_dict(task_dictionary:Dict[str,Any]):"""Load the task from a dictionary."""fromidmtools.core.task_factoryimportTaskFactorytask_args={k:vfork,vintask_dictionary.items()ifknotin['task_type']}returnTaskFactory().create(task_dictionary['task_type'],**task_args)
@propertydefcommand(self):"""Our task property. Again, we have to overload this because of wrapping a task."""cmd=copy.deepcopy(self.template_script_task.command)cmd.add_raw_argument(str(self.task.command))returncmd@command.setterdefcommand(self,value:Union[str,CommandLine]):"""Set our command. because we are wrapping a task, we have to overload this."""callers=[]withsuppress(KeyError,IndexError):callers=inspect.stack()[1:3]ifnotcallersorcallers[0].filename==__file__orcallers[1].filename==__file__:ifisinstance(value,property):self._command=Noneelifisinstance(value,str):self._command=CommandLine.from_string(value)else:self._command=valueelse:self.task.command=value@propertydefwrapped_task(self):""" Our task we are wrapping with a shell script. Returns: Our wrapped task """returnself.task@wrapped_task.setterdefwrapped_task(self,value:ITask):"""Set our wrapped task."""returnNoneifisinstance(value,property)elsevalue
[docs]defgather_common_assets(self):""" Gather all the common assets. Returns: Common assets(Experiment Assets) """self.common_assets.add_assets(self.template_script_task.gather_common_assets())self.common_assets.add_assets(self.task.gather_common_assets())returnself.common_assets
[docs]defgather_transient_assets(self)->AssetCollection:""" Gather all the transient assets. Returns: Transient Assets(Simulation level assets) """self.transient_assets.add_assets(self.template_script_task.gather_transient_assets())self.transient_assets.add_assets(self.task.gather_transient_assets())returnself.transient_assets
[docs]defreload_from_simulation(self,simulation:Simulation):""" Reload from simulation. Args: simulation: simulation Returns: None """pass
[docs]defpre_creation(self,parent:Union[Simulation,IWorkflowItem],platform:'IPlatform'):""" Before creation, create the true command by adding the wrapper name. Here we call both our wrapped task and our template_script_task pre_creation Args: parent: Parent Task platform: Platform Templated Task is executing on Returns: None """self.task.pre_creation(parent,platform)self.template_script_task.pre_creation(parent,platform)
[docs]defpost_creation(self,parent:Union[Simulation,IWorkflowItem],platform:'IPlatform'):""" Post creation of task. Here we call both our wrapped task and our template_script_task post_creation Args: parent: Parent of task platform: Platform we are running on Returns: None """self.task.post_creation(parent,platform)self.template_script_task.post_creation(parent,platform)
def__getattr__(self,item):"""Proxy get attr to child except for task item and items not in our object."""ifitemnotinself.__dict__:returngetattr(self.task,item)else:returnsuper(ScriptWrapperTask,self).__getattr__(item)
[docs]defget_script_wrapper_task(task:ITask,wrapper_script_name:str,template_content:str=None,template_file:str=None,template_is_common:bool=True,variables:Dict[str,Any]=None,path_sep:str='/')->ScriptWrapperTask:""" Convenience function that will wrap a task for you with some defaults. Args: task: Task to wrap wrapper_script_name: Wrapper script name template_content: Template Content template_file: Template File template_is_common: Is the template experiment level variables: Variables path_sep: Path sep(Window or Linux) Returns: ScriptWrapperTask wrapping the task See Also: :class:`idmtools_models.templated_script_task.TemplatedScriptTask` :func:`idmtools_models.templated_script_task.get_script_wrapper_windows_task` :func:`idmtools_models.templated_script_task.get_script_wrapper_unix_task` """ifvariablesisNone:variables=dict()template_task=TemplatedScriptTask(script_path=wrapper_script_name,template_file=template_file,template=template_content,template_is_common=template_is_common,variables=variables,path_sep=path_sep)returnScriptWrapperTask(template_script_task=template_task,task=task)
[docs]defget_script_wrapper_windows_task(task:ITask,wrapper_script_name:str='wrapper.bat',template_content:str=WINDOWS_DICT_TO_ENVIRONMENT,template_file:str=None,template_is_common:bool=True,variables:Dict[str,Any]=None)->ScriptWrapperTask:""" Get wrapper script task for windows platforms. The default content wraps a another task with a batch script that exports the variables to the run environment defined in variables. To modify python path, use WINDOWS_PYTHON_PATH_WRAPPER You can adapt this script to modify any pre-scripts you need or call others scripts in succession Args: task: Task to wrap wrapper_script_name: Wrapper script name(defaults to wrapper.bat) template_content: Template Content. template_file: Template File template_is_common: Is the template experiment level variables: Variables for template Returns: ScriptWrapperTask See Also:: :class:`idmtools_models.templated_script_task.TemplatedScriptTask` :func:`idmtools_models.templated_script_task.get_script_wrapper_task` :func:`idmtools_models.templated_script_task.get_script_wrapper_unix_task` """returnget_script_wrapper_task(task,wrapper_script_name,template_content,template_file,template_is_common,variables,"\\")
[docs]defget_script_wrapper_unix_task(task:ITask,wrapper_script_name:str='wrapper.sh',template_content:str=LINUX_DICT_TO_ENVIRONMENT,template_file:str=None,template_is_common:bool=True,variables:Dict[str,Any]=None):""" Get wrapper script task for unix platforms. The default content wraps a another task with a bash script that exports the variables to the run environment defined in variables. To modify python path, you can use LINUX_PYTHON_PATH_WRAPPER You can adapt this script to modify any pre-scripts you need or call others scripts in succession Args: task: Task to wrap wrapper_script_name: Wrapper script name(defaults to wrapper.sh) template_content: Template Content template_file: Template File template_is_common: Is the template experiment level variables: Variables for template Returns: ScriptWrapperTask See Also: :class:`idmtools_models.templated_script_task.TemplatedScriptTask` :func:`idmtools_models.templated_script_task.get_script_wrapper_task` :func:`idmtools_models.templated_script_task.get_script_wrapper_windows_task` """returnget_script_wrapper_task(task,wrapper_script_name,template_content,template_file,template_is_common,variables,"/")
[docs]classTemplatedScriptTaskSpecification(TaskSpecification):""" TemplatedScriptTaskSpecification provides the plugin specs for TemplatedScriptTask. """
[docs]defget(self,configuration:dict)->TemplatedScriptTask:""" Get instance of TemplatedScriptTask with configuration. Args: configuration: configuration for TemplatedScriptTask Returns: TemplatedScriptTask with configuration """returnTemplatedScriptTask(**configuration)
[docs]defget_description(self)->str:""" Get description of plugin. Returns: Plugin description """return"Defines a general command that provides user hooks. Intended for use in advanced scenarios"
[docs]defget_example_urls(self)->List[str]:""" Get example urls related to TemplatedScriptTask. Returns: List of urls that have examples related to CommandTask """return[]
[docs]defget_type(self)->Type[TemplatedScriptTask]:""" Get task type provided by plugin. Returns: TemplatedScriptTask """returnTemplatedScriptTask
[docs]defget_version(self)->str:""" Returns the version of the plugin. Returns: Plugin Version """fromidmtools_modelsimport__version__return__version__
[docs]classScriptWrapperTaskSpecification(TaskSpecification):""" ScriptWrapperTaskSpecification defines the plugin specs for ScriptWrapperTask. """
[docs]defget(self,configuration:dict)->ScriptWrapperTask:""" Get instance of ScriptWrapperTask with configuration. Args: configuration: configuration for ScriptWrapperTask Returns: TemplatedScriptTask with configuration """returnScriptWrapperTask(**configuration)
[docs]defget_description(self)->str:""" Get description of plugin. Returns: Plugin description """return"Defines a general command that provides user hooks. Intended for use in advanced scenarios"
[docs]defget_example_urls(self)->List[str]:""" Get example urls related to ScriptWrapperTask. Returns: List of urls that have examples related to CommandTask """return[]
[docs]defget_type(self)->Type[ScriptWrapperTask]:""" Get task type provided by plugin. Returns: TemplatedScriptTask """returnScriptWrapperTask
[docs]defget_version(self)->str:""" Returns the version of the plugin. Returns: Plugin Version """fromidmtools_modelsimport__version__return__version__