[BUG] flytekit 1.6.0 task function calls do not type check #3682
Comments
tekumara
added
bug
eapolinario
added
flytekit
|
Verifying that this is still an issue with from flytekit import task, workflow
@task
def task_get_a() -> int:
return 1
@workflow
def my_workflow() -> int:
a: int = task_get_a()
return aAnd here's the error from Trying this with |
|
Returning a tuple of values also causes Here's an example workflow and task: from typing import Tuple
from flytekit import task, workflow
@task
def task_get_a_b() -> Tuple[int, int]:
return 1, 2
@workflow
def my_workflow() -> int:
a: int
b: int
a, b = task_get_a_b()
return aAnd here are the errors reported by |
|
And one more example using a dataclass: from dataclasses import dataclass
from dataclasses_json import DataClassJsonMixin
from flytekit import task, workflow
@dataclass
class MyDataClass(DataClassJsonMixin):
a: int
b: int
@task
def task_get_mydataclass() -> MyDataClass:
return MyDataClass(a=1, b=2)
@workflow
def my_workflow() -> bool:
my_dataclass: MyDataClass = task_get_mydataclass()
if my_dataclass.a == 1 and my_dataclass.b == 2:
return True
return FalseAnd the error from |
|
We also ran into this issue, blocking us from type checking the flyte parts of our code base. It seems like I don't really have a good idea how to solve this, as the |
|
After digging deeper into this, I assume there is currently no solution do solve this in flytekit. Internally we do use a workaround, which I will elaborate at the end. Example"""A simple example on the typing issues in flytekit."""
from typing import Callable, ParamSpec, TypeVar, Union
from flytekit import task, workflow
from flytekit.core.promise import Promise
@task
def add_one(a: int) -> int:
return a + 1
@workflow
def my_workflow() -> int:
# Add one should be able to receive an int
start: int = 0
res = add_one(a=start)
# The return type of the decorated task is no longer
# an int but a Promise
assert not isinstance(res, int)
# Thus the decorated add_one takes either a Promise or an int
add_one(a=res)
return resProblem analysisTo me it seems like the correct types to cover both runtime and registration time view on the workflow can't be expressed in the python type system, as it does not have a Map operator on types. A task decorator would need to change a function signature from T = TypeVar("T")
U = TypeVar("U")
V = TypeVar("V")
def python_func(arg1: T, arg2: U) -> V: ...to something like def decorated_python_func(arg1: Promise[T] | T, arg2: Promise[U] | U) -> Promise[V]: ...To annotate the types in the FuncIn = ParamSpec("FuncIn")
FuncOut = TypeVar("FuncOut")
ToDecorateFunction = Callable[FuncIn, FuncOut]where each type This is not yet supported in python's type system: python/typing#1163 WorkaroundOur company already internally replaced from typing import Callable, Optional, ParamSpec, TypeVar, Union, overload
from flytekit.core.task import task as flytekit_task
FuncOut = TypeVar("FuncOut")
FuncIn = ParamSpec("FuncIn")
@overload
def our_task(
_task_function: None = ...,
**kwargs,
) -> Callable[[Callable[FuncIn, FuncOut]], Callable[FuncIn, FuncOut]]: ...
@overload
def our_task(
_task_function: Callable[FuncIn, FuncOut],
**kwargs,
) -> Callable[FuncIn, FuncOut]: ...
def our_task(
_task_function: Optional[Callable[FuncIn, FuncOut]] = None,
**kwargs,
) -> Union[Callable[FuncIn, FuncOut], Callable[[Callable[FuncIn, FuncOut]], Callable[FuncIn, FuncOut]]]:
# Do some other stuff here
# ...
# Then actually call the flytekit task
def wrapped(_func: Callable[FuncIn, FuncOut]) -> Callable[FuncIn, FuncOut]:
# As we want to type check the runtime
return flytekit_task(_task_function=target_fn(_func), **kwargs) # type: ignore
if _task_function:
return wrapped(_task_function)
else:
return wrappedThis can be done in a similar way for other It allows us to type check workflows for their runtime view but also causes issues one needs to be aware of: from flytekit import workflow, Resources
@our_task
def add_one(a: int) -> int:
return a + 1
@workflow
def my_workflow() -> int:
start: int = 0
# res will be an int in pyright as we use `our_task`
res = add_one(a=start)
# IMPORTANT!
# The pyright types will be wrong but allow type checking, i.e.
# this type assertion will fail:
# assert isinstance(res, int)
# as res is actually still a Promise:
print(type(res)) # <class 'flytekit.core.promise.Promise'>
# When using anything related to `flytekit` Promises, as e.g. `with_overrides`, we need a small workaround:
res2 = add_one(a=res) # res2 will still be an int in pyright
# With overrides is called separately and needs type: ignore as pyright sees res2 as an int and not a Promise
res2.with_overrides(requests=Resources(cpu="0.1", mem="500Mi")) # type: ignore
return res2 |
|
I wasn't aware that we can use We could use this trick to provide a different type signature during at type-checking time. |
|
Indeed, we could offer the runtime types with this trick. Might be still a better developer experience? |
Describe the bug
In flyteorg/flytekit#1631 the
taskreturn type was changed to effectivelyTuple[Promise] | Promise | VoidPromise | Tuple[Unknown, ...] | None.This breaks usage of tasks in type checkers like pyright / vscode.
Expected behavior
Task function calls still type check correctly, as they did in flytekit 1.5.0
Additional context to reproduce
This type checks without errors in flytekit 1.5.0:
With flytekit 1.6.0 and pyright / vscode:
Screenshots
No response
Are you sure this issue hasn't been raised already?
Have you read the Code of Conduct?
The text was updated successfully, but these errors were encountered: