New RunLLM response generated: >>>*# Question*

from typing import Union
from dataclasses import dataclass
from flytekit import task, workflow

@dataclass
class A:
    a: int

@dataclass
class B:
    b: int

@task
def foo(inp: Union[A, B]):
    ...

@workflow
def wf():
    foo(inp=B(b=1))

if __name__ == "__main__":
    wf()

Our users want to do something like this and I also would have naively expected this should work but it actually doesn’t. Details in 🧵 Would be curious if somebody knows how to fix this 🙏 # Answer Flyte currently does not support the direct use of

Union

types in tasks and workflows. This limitation is due to the way Flyte handles type serialization and deserialization, which requires explicit type information that

Union

types do not provide. According to Flyte's GitHub issue #4505, Flyte's type system does not handle

Union

types well, especially when it comes to distinguishing between different types at runtime. This issue is compounded when dealing with complex types like dataclasses. To work around this limitation, you can use separate tasks for each type or use a more generic type like

Any

and handle type checking within the task. Here is an example of using separate tasks:

from dataclasses import dataclass
from flytekit import task, workflow

@dataclass
class A:
    a: int

@dataclass
class B:
    b: int

@task
def foo_a(inp: A):
    print(f"Processing A: {inp}")

@task
def foo_b(inp: B):
    print(f"Processing B: {inp}")

@workflow
def wf():
    foo_b(inp=B(b=1))

if __name__ == "__main__":
    wf()

Alternatively, using

Any

with type checks:

from typing import Any
from dataclasses import dataclass
from flytekit import task, workflow

@dataclass
class A:
    a: int

@dataclass
class B:
    b: int

@task
def foo(inp: Any):
    if isinstance(inp, A):
        print(f"Processing A: {inp}")
    elif isinstance(inp, B):
        print(f"Processing B: {inp}")
    else:
        raise ValueError("Unsupported type")

@workflow
def wf():
    foo(inp=B(b=1))

if __name__ == "__main__":
    wf()

Would you like more details on either approach? Sources: - Flyte GitHub Issue #4505 Slack Link https://flyte-org.slack.com/archives/CP2HDHKE1/p1718792181.448699 👍 0 button 👎 0 button 💡 Hint: Mention @runllm in the thread for followups.