Template set

The hera.workflows.template_set module provides the TemplateSet class.

The TemplateSet class lets you create a collection of templates unattached to a Workflow.

Warning

This class is only for use with experimental Workflow decorators. Read the guide here.

The TemplateSet class can be used to arrange templates across modules in a package.

Source code in src/hera/workflows/template_set.py

class TemplateSet(
    TemplateDecoratorFuncsMixin,
):
    """The TemplateSet class can be used to arrange templates across modules in a package."""

    templates: List[Union[_ModelTemplate, Templatable]] = []

    def _add_sub(self, node: Any):
        """Adds the given node (expected to satisfy the `Templatable` protocol) to the context."""
        if not isinstance(node, (Templatable, _ModelTemplate)):
            raise InvalidType(type(node))
        self.templates.append(node)

templates

templates: List[Union[Template, Templatable]] = []

script

script(**script_kwargs) -> Callable

A decorator that wraps a function into a Script object.

Using this decorator users can define a function that will be executed as a script in a container. Once the Script is returned users can use it as they generally use a Script e.g. as a callable inside a DAG or Steps. Note that invoking the function will result in the template associated with the script to be added to the workflow context, so users do not have to worry about that.

Parameters:

Name	Type	Description	Default
**script_kwargs		Keyword arguments to be passed to the Script object.	{}

Returns:

Type	Description
Callable	Function wrapper that holds a Script and allows the function to be called to create a Step or Task if in a Steps or DAG context.

Source code in src/hera/workflows/_meta_mixins.py

@_add_type_hints(Script)  # type: ignore
def script(self, **script_kwargs) -> Callable:
    """A decorator that wraps a function into a Script object.

    Using this decorator users can define a function that will be executed as a script in a container. Once the
    `Script` is returned users can use it as they generally use a `Script` e.g. as a callable inside a DAG or Steps.
    Note that invoking the function will result in the template associated with the script to be added to the
    workflow context, so users do not have to worry about that.

    Args:
        **script_kwargs: Keyword arguments to be passed to the Script object.

    Returns:
        Function wrapper that holds a `Script` and allows the function to be called to create a Step or Task if in a Steps or DAG context.
    """
    self._check_if_enabled("script")

    from hera.workflows.script import RunnerScriptConstructor, Script

    def script_decorator(func: Callable[FuncIns, FuncR]) -> Callable:
        """The internal decorator function.

        Parameters
        ----------
        func: Callable
            Function to wrap.

        Returns:
        -------
        Callable
            Callable that calls the `Script` object `__call__` method when in a Steps or DAG context,
            otherwise calls function itself.
        """
        # instance methods are wrapped in `staticmethod`. Hera can capture that type and extract the underlying
        # function for remote submission since it does not depend on any class or instance attributes, so it is
        # submittable
        if isinstance(func, staticmethod):
            source: Callable = func.__func__
        else:
            source = func

        # take the client-provided `name` if it is submitted, pop the name for otherwise there will be two
        # kwargs called `name`
        # otherwise populate the `name` from the function name
        name = script_kwargs.pop("name", source.__name__.replace("_", "-"))

        if "constructor" not in script_kwargs and "constructor" not in global_config._get_class_defaults(Script):
            script_kwargs["constructor"] = RunnerScriptConstructor()

        signature = inspect.signature(func)
        func_inputs = signature.parameters
        if len(func_inputs) > 1:
            raise SyntaxError("script decorator must be used with a single `Input` arg, or no args.")

        if len(func_inputs) == 1:
            func_input = list(func_inputs.values())[0].annotation
            if not issubclass(func_input, (InputV1, InputV2)):
                raise SyntaxError("script decorator must be used with a single `Input` arg, or no args.")

        # Open (Workflow) context to add `Script` object automatically
        with self:
            script_template = Script(name=name, source=source, **script_kwargs)

        if not isinstance(script_template.constructor, RunnerScriptConstructor):
            raise ValueError(f"Script '{name}' must use RunnerScriptConstructor")

        @functools.wraps(func)
        def script_call_wrapper(*args, **kwargs) -> Union[FuncR, Step, Task, None]:
            """Invokes a CallableTemplateMixin's `__call__` method using the given SubNode (Step or Task) args/kwargs."""
            if _context.declaring:
                subnode_name = kwargs.pop("name", None)
                if not subnode_name:
                    try:
                        # ignore decorator function assignment
                        subnode_name = varname()
                    except ImproperUseError:
                        # Template is being used without variable assignment (so use function name or provided name)
                        subnode_name = script_template.name  # type: ignore

                    assert isinstance(subnode_name, str)
                    subnode_name = subnode_name.replace("_", "-")

                return self._create_subnode(subnode_name, func, script_template, *args, **kwargs)

            if _context.pieces:
                return script_template.__call__(*args, **kwargs)

            # Do not allow kwargs
            return func(*args)  # type: ignore

        # Set the wrapped function to the original function so that we can use it later
        script_call_wrapper.wrapped_function = func  # type: ignore
        # Set the template name to the inferred name
        script_call_wrapper.template_name = name  # type: ignore

        return script_call_wrapper

    return script_decorator

container

container(**container_kwargs) -> Callable

Source code in src/hera/workflows/_meta_mixins.py

@_add_type_hints(Container)  # type: ignore
def container(self, **container_kwargs) -> Callable:
    self._check_if_enabled("container")
    from hera.workflows.container import Container

    def container_decorator(func: Callable[FuncIns, FuncR]) -> Callable:
        name = container_kwargs.pop("name", func.__name__.replace("_", "-"))

        signature = inspect.signature(func)
        func_inputs = signature.parameters
        inputs = []
        if len(func_inputs) >= 1:
            input_arg = list(func_inputs.values())[0].annotation
            if issubclass(input_arg, (InputV1, InputV2)):
                inputs = input_arg._get_inputs(add_missing_path=True)

        func_return = signature.return_annotation
        outputs = []
        if func_return and issubclass(func_return, (OutputV1, OutputV2)):
            outputs = func_return._get_outputs(add_missing_path=True)

        # Open context to add `Container` object automatically
        with self:
            container_template = Container(
                name=name,
                inputs=inputs,
                outputs=outputs,
                **container_kwargs,
            )

        @functools.wraps(func)
        def container_call_wrapper(*args, **kwargs) -> Union[FuncR, Step, Task, None]:
            """Invokes a CallableTemplateMixin's `__call__` method using the given SubNode (Step or Task) args/kwargs."""
            if _context.declaring:
                subnode_name = kwargs.pop("name", None)
                if not subnode_name:
                    try:
                        # ignore decorator function assignment
                        subnode_name = varname()
                    except ImproperUseError:
                        # Template is being used without variable assignment (so use function name or provided name)
                        subnode_name = container_template.name  # type: ignore

                    assert isinstance(subnode_name, str)
                    subnode_name = subnode_name.replace("_", "-")

                return self._create_subnode(subnode_name, func, container_template, *args, **kwargs)

            if _context.pieces:
                return container_template.__call__(*args, **kwargs)

            # Do not allow kwargs
            return func(*args)  # type: ignore

        # Set the template name to the inferred name
        container_call_wrapper.template_name = name  # type: ignore

        return container_call_wrapper

    return container_decorator

dag

dag(**dag_kwargs) -> Callable

Source code in src/hera/workflows/_meta_mixins.py

@_add_type_hints(DAG)  # type: ignore
def dag(self, **dag_kwargs) -> Callable:
    self._check_if_enabled("dag")

    from hera.workflows.dag import DAG

    return self._construct_invocator_decorator(DAG, **dag_kwargs)

steps

steps(**steps_kwargs) -> Callable

Source code in src/hera/workflows/_meta_mixins.py

@_add_type_hints(Steps)  # type: ignore
def steps(self, **steps_kwargs) -> Callable:
    self._check_if_enabled("steps")

    from hera.workflows.steps import Steps

    return self._construct_invocator_decorator(Steps, **steps_kwargs)

Comments