Is it possible to decorate functions with @script at runtime?

[autumnkelsey]

Hi all!
I'm trying to decorate functions with the script decorator at runtime rather than compile time, but I keep running into exceptions.

The following code runs as expected--build the workflow and print out the YAML:

from hera.workflows import Steps, Workflow, script
from typing import Callable

@script(image="python:3.11")
def f1():
    print("function 1")

@script(image="python:3.11")
def f2():
    print("function 2")

def staticallyDecorate(*steps: Callable):
    with Workflow(entrypoint="steps", generate_name="idk") as w:
        with Steps(name="steps"):
            for step in steps:
                step()
    print(w.to_yaml())

staticallyDecorate(f1, f2)

But this code, which I believe should do the same thing, only at runtime, fails:

from hera.workflows import Steps, Workflow, script
from typing import Callable

def f1():
    print("function 1")

def f2():
    print("function 2")

def dynamicallyDecorate(*steps: Callable):
    with Workflow(entrypoint="steps", generate_name="idk") as w:
        with Steps(name="steps"):
            for step in steps:
                script(image="python:3.11", constructor="inline")(step)
    print(w.to_yaml())

dynamicallyDecorate(f1, f2)

The stack trace I get is:

Traceback (most recent call last):
  File "/Users/me/Repos/repo/./scripts/spawner.py", line 19, in <module>
    dynamicallyDecorate(f1, f2)
  File "/Users/me/Repos/repo/./scripts/spawner.py", line 16, in dynamicallyDecorate
    script(image="python:3.11", constructor="inline")(step)
  File "/Users/me/pythonvenv/lib/python3.12/site-packages/hera/workflows/script.py", line 771, in script_wrapper
    s = Script(name=name, source=source, **script_kwargs)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ame/pythonvenv/lib/python3.12/site-packages/pydantic/v1/main.py", line 349, in __init__
    __pydantic_self__._init_private_attributes()
  File "/Users/me/pythonvenv/lib/python3.12/site-packages/hera/shared/_global_config.py", line 181, in _init_private_attributes
    self.__hera_init__()
  File "/Users/me/pythonvenv/lib/python3.12/site-packages/hera/workflows/_context.py", line 31, in __hera_init__
    _context.add_sub_node(self)
  File "/Users/me/pythonvenv/lib/python3.12/site-packages/hera/workflows/_context.py", line 117, in add_sub_node
    pieces[-1]._add_sub(node)
  File "/Users/me/pythonvenv/lib/python3.12/site-packages/hera/workflows/steps.py", line 185, in _add_sub
    raise InvalidType(type(node))
hera.workflows.exceptions.InvalidType: <class 'hera.workflows.script.Script'>

Is there something special I have to do to invoke the script decorator at runtime?

[alicederyn]

Try calling the decorator outside of the with blocks?

def dynamicallyDecorate(*steps: Callable):
    decorated_steps = [
        script(image="python:3.11", constructor="inline")(step)
        for step in steps
    ]
    with Workflow(entrypoint="steps", generate_name="idk") as w:
        with Steps(name="steps"):
            for step in decorated_steps:
                step()
    print(w.to_yaml())

This works locally for me and produces the same yaml as staticallyDecorate

[autumnkelsey]

Thank you so much!! I feel so silly now lol. I had tried decorating the steps outside of the with blocks with

decorated_steps = map(lambda step: script(image="python:3.11", constructor="inline")(step), steps)

but I didn't realize the lambda was executed during iteration and not immediately 😅