Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

check #2486

Merged
Merged

check #2486

Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions flytekit/core/array_node_map_task.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
from flytekit.models.interface import Variable
from flytekit.models.task import Container, K8sPod, Sql, Task
from flytekit.tools.module_loader import load_object_from_module
from flytekit.types.pickle.pickle import FlytePickleTransformer


class ArrayNodeMapTask(PythonTask):
Expand Down Expand Up @@ -55,6 +56,11 @@ def __init__(
else:
actual_task = python_function_task

for _, v in actual_task.python_interface.inputs.items():
transformer = TypeEngine.get_transformer(v)
if isinstance(transformer, FlytePickleTransformer):
raise ValueError("Pickle transformers are not supported in map tasks.")

# TODO: add support for other Flyte entities
if not (isinstance(actual_task, PythonFunctionTask) or isinstance(actual_task, PythonInstanceTask)):
raise ValueError("Only PythonFunctionTask and PythonInstanceTask are supported in map tasks.")
Expand Down Expand Up @@ -224,6 +230,7 @@ def _literal_map_to_python_input(
inputs_interface = self._run_task.python_interface.inputs
for k in self.interface.inputs.keys():
v = literal_map.literals[k]

if k not in self.bound_inputs:
# assert that v.collection is not None
if not v.collection or not isinstance(v.collection.literals, list):
Expand Down
9 changes: 9 additions & 0 deletions tests/flytekit/unit/core/test_array_node_map_task.py
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,15 @@ def say_hello(name: str) -> str:
assert res.literals["o0"].scalar.primitive.string_value == "hello earth!"


def test_map_task_with_pickle():
@task
def say_hello(name: typing.Any) -> str:
return f"hello {name}!"

with pytest.raises(ValueError, match="Pickle transformers are not supported in map tasks"):
map_task(say_hello)(name=["abc", "def"])


def test_serialization(serialization_settings):
@task
def t1(a: int) -> int:
Expand Down
Loading