"""Provides the TemplatedScriptTask.
Copyright 2021, Bill & Melinda Gates Foundation. All rights reserved.
"""
import copy
import inspect
import os
from contextlib import suppress
from dataclasses import dataclass, field
from functools import partial
from logging import getLogger, DEBUG
from typing import List, Callable, Type, Dict, Any, Union, TYPE_CHECKING
from jinja2 import Environment
from idmtools.assets import AssetCollection, Asset
from idmtools.entities import CommandLine
from idmtools.entities.itask import ITask
from idmtools.entities.iworkflow_item import IWorkflowItem
from idmtools.entities.simulation import Simulation
from idmtools.registry.task_specification import TaskSpecification
if TYPE_CHECKING: # pragma: no cover
from idmtools.entities.iplatform import IPlatform
logger = 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 templates
WINDOWS_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 scripts
LINUX_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()
class TemplatedScriptTask(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 script
script_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 contents
template: str = field(default=None, metadata={"md": True})
#: The template file. You can only use either template or template_file at once
template_file: str = field(default=None, metadata={"md": True})
#: Controls whether a template should be an experiment or a simulation level asset
template_is_common: bool = field(default=True, metadata={"md": True})
#: Template variables used for rendering the template
# Note: large amounts of variables will increase metadata size
variables: 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 line
extra_command_arguments: str = field(default='', metadata={"md": True})
#: Hooks to gather common assets
gather_common_asset_hooks: List[Callable[[ITask], AssetCollection]] = field(default_factory=list)
#: Hooks to gather transient assets
gather_transient_asset_hooks: List[Callable[[ITask], AssetCollection]] = field(default_factory=list)
def __post_init__(self):
"""Constructor."""
super().__post_init__()
if self.script_path is None and self.template_file is None:
raise ValueError("Either script name or template file is required")
elif self.script_path is None: # Get name from the file
self.script_path = os.path.basename(self.template_file)
if self.template is None and self.template_file is None:
raise ValueError("Either template or template_file is required")
if self.template_file and not os.path.exists(self.template_file):
raise FileNotFoundError(f"Could not find the file the template file {self.template_file}")
# is the template a common(experiment) asset?
if self.template_is_common:
if logger.isEnabledFor(DEBUG):
logger.debug("Adding common asset hook")
# add hook to render it to common asset hooks
hook = 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 asset
if logger.isEnabledFor(DEBUG):
logger.debug("Adding transient asset hook")
# create hook to render it to our transient assets
hook = partial(
TemplatedScriptTask._add_template_to_asset_collection,
asset_collection=self.transient_assets
)
# add the hook
self.gather_transient_asset_hooks.append(hook)
@staticmethod
def _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 environment
env = Environment()
# try to load template from string or file
if task.template:
template = env.from_string(task.template)
else:
template = env.get_template(task.template_file)
if logger.isEnabledFor(DEBUG):
logger.debug(f"Rendering Script template: {template}")
# render the template
result = template.render(vars=task.variables, env=os.environ)
if logger.isEnabledFor(DEBUG):
logger.debug(f"Rendered Template: {result}")
# create an asset
asset = Asset(filename=task.script_path, content=result)
asset_collection.add_asset(asset)
return asset_collection
[docs] def gather_common_assets(self) -> AssetCollection:
"""
Gather common(experiment-level) assets for task.
Returns:
AssetCollection containing common assets
"""
# TODO validate hooks have expected return type
ac = AssetCollection()
for x in self.gather_common_asset_hooks:
ac += x(self)
return ac
[docs] def gather_transient_assets(self) -> AssetCollection:
"""
Gather transient(experiment-level) assets for task.
Returns:
AssetCollection containing transient assets
"""
ac = AssetCollection()
for x in self.gather_transient_asset_hooks:
ac += x(self)
ac += self.transient_assets
if len(ac.assets) != 0:
self.transient_assets = ac
return ac
[docs] def reload_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 script
if simulation.parent.assets:
# prep new asset collection in case we have to remove our asset from the experiment level
new_assets = AssetCollection()
for _i, asset in enumerate(simulation.parent.assets.assets):
# is it our script?
if asset.filename != self.script_path and asset.absolute_path != self.script_path:
# nope keep it
new_assets.add_asset(asset)
# set filtered assets back to parent
simulation.parent.assets = new_assets
[docs] def pre_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?
if self.script_binary:
sn = self.script_binary + ' '
else:
sn = ''
if self.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 script
self.command = CommandLine.from_string(sn)
if self.path_sep != "/":
self.command.executable = self.command.executable.replace("/", self.path_sep)
self.command.is_windows = True
# set any extra arguments
if self.extra_command_arguments:
other_command = CommandLine.from_string(self.extra_command_arguments)
self.command._args.append(other_command.executable)
if other_command._options:
self.command._args += other_command._options
if other_command._args:
self.command._args += other_command._args
super().pre_creation(parent, platform)
[docs]@dataclass()
class ScriptWrapperTask(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."""
if self.template_script_task is None:
raise ValueError("Template Script Task is required")
if self.task is None:
raise ValueError("Task is required")
if isinstance(self.task, dict):
self.task = self.from_dict(self.task)
if isinstance(self.template_script_task, dict):
self.template_script_task = self.from_dict(self.template_script_task)
[docs] @staticmethod
def from_dict(task_dictionary: Dict[str, Any]):
"""Load the task from a dictionary."""
from idmtools.core.task_factory import TaskFactory
task_args = {k: v for k, v in task_dictionary.items() if k not in ['task_type']}
return TaskFactory().create(task_dictionary['task_type'], **task_args)
@property
def command(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))
return cmd
@command.setter
def command(self, value: Union[str, CommandLine]):
"""Set our command. because we are wrapping a task, we have to overload this."""
callers = []
with suppress(KeyError, IndexError):
callers = inspect.stack()[1:3]
if not callers or callers[0].filename == __file__ or callers[1].filename == __file__:
if isinstance(value, property):
self._command = None
elif isinstance(value, str):
self._command = CommandLine.from_string(value)
else:
self._command = value
else:
self.task.command = value
@property
def wrapped_task(self):
"""
Our task we are wrapping with a shell script.
Returns:
Our wrapped task
"""
return self.task
@wrapped_task.setter
def wrapped_task(self, value: ITask):
"""Set our wrapped task."""
return None if isinstance(value, property) else value
[docs] def gather_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())
return self.common_assets
[docs] def gather_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())
return self.transient_assets
[docs] def reload_from_simulation(self, simulation: Simulation):
"""
Reload from simulation.
Args:
simulation: simulation
Returns:
None
"""
pass
[docs] def pre_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] def post_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."""
if item not in self.__dict__:
return getattr(self.task, item)
else:
return super(ScriptWrapperTask, self).__getattr__(item)
[docs]def get_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`
"""
if variables is None:
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
)
return ScriptWrapperTask(template_script_task=template_task, task=task)
[docs]def get_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`
"""
return get_script_wrapper_task(task, wrapper_script_name, template_content, template_file, template_is_common,
variables, "\\")
[docs]def get_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`
"""
return get_script_wrapper_task(task, wrapper_script_name, template_content, template_file, template_is_common,
variables, "/")
[docs]class TemplatedScriptTaskSpecification(TaskSpecification):
"""
TemplatedScriptTaskSpecification provides the plugin specs for TemplatedScriptTask.
"""
[docs] def get(self, configuration: dict) -> TemplatedScriptTask:
"""
Get instance of TemplatedScriptTask with configuration.
Args:
configuration: configuration for TemplatedScriptTask
Returns:
TemplatedScriptTask with configuration
"""
return TemplatedScriptTask(**configuration)
[docs] def get_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] def get_example_urls(self) -> List[str]:
"""
Get example urls related to TemplatedScriptTask.
Returns:
List of urls that have examples related to CommandTask
"""
return []
[docs] def get_type(self) -> Type[TemplatedScriptTask]:
"""
Get task type provided by plugin.
Returns:
TemplatedScriptTask
"""
return TemplatedScriptTask
[docs] def get_version(self) -> str:
"""
Returns the version of the plugin.
Returns:
Plugin Version
"""
from idmtools_models import __version__
return __version__
[docs]class ScriptWrapperTaskSpecification(TaskSpecification):
"""
ScriptWrapperTaskSpecification defines the plugin specs for ScriptWrapperTask.
"""
[docs] def get(self, configuration: dict) -> ScriptWrapperTask:
"""
Get instance of ScriptWrapperTask with configuration.
Args:
configuration: configuration for ScriptWrapperTask
Returns:
TemplatedScriptTask with configuration
"""
return ScriptWrapperTask(**configuration)
[docs] def get_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] def get_example_urls(self) -> List[str]:
"""
Get example urls related to ScriptWrapperTask.
Returns:
List of urls that have examples related to CommandTask
"""
return []
[docs] def get_type(self) -> Type[ScriptWrapperTask]:
"""
Get task type provided by plugin.
Returns:
TemplatedScriptTask
"""
return ScriptWrapperTask
[docs] def get_version(self) -> str:
"""
Returns the version of the plugin.
Returns:
Plugin Version
"""
from idmtools_models import __version__
return __version__