Is there support for tasks with inputs that use Un...
# ask-the-community
l
Is there support for tasks with inputs that use Unions?
Union[str, flytekit.types.directory.FlyteDirectory]
. Trying to create a task that could take in a str or a flytedirectory but I get
Copy code
raise ValueError(f"Expected a directory. {source_path} is not a directory")
ValueError: Expected a directory. X is not a directory

AssertionError: Failed to Bind variable input for function
k
Should work
Cc @Kevin Su
l
Yea it does to seem work; I tested it with a small task, I think for this task its not working bc I have a decorator around it to do some custom io and calling
func(*args, **kwargs)
inside the decorator might be breaking something?
k
Ohh read the doc on decorators
You have to ensure you pass the type info Correctly
l
Copy code
def mount(job_type):
    def inner(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
I am using the the
@wraps
, seems like I need to use partial. I'll play around with it
I think there might be something special with flytedirectory and decorators? I'm not really sure but here's a minimal example that is failing
Copy code
from functools import partial, wraps
def validate_output(fn=None, *, floor=0):
    @wraps(fn)
    def wrapper(*args, **kwargs):
        out = fn(*args, **kwargs)
        if out <= floor:
            raise ValueError(
                f"output of task {fn.__name__} must be a positive number, found {out}"
            )
        return out

    if fn is None:
        return partial(validate_output, floor=floor)

    return wrapper

@task
@validate_output(floor=0)
def test_task(input: Union[str, flytekit.types.directory.FlyteDirectory]) ->int :
    print(input)
    return 5

if __name__ == "__main__":
    result = test_task(
        input="some str",
    )
I just copied the decorator from the docs and it's still trying to resolve the input as a directory instead of as a string
cc @Ketan (kumare3) @Kevin Su
k
@Laura Lin Sorry, there is a small bug in type engine. I just fixed in this PR. you can try this by running
pip install git+<https://github.com/flyteorg/flytekit@flyte_dir_union>
l
yes, it seems to work with the code snippet I shared above! thanks so much ❤️
113 Views