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

[Core feature] Flytekit should support unsafe mode for types #2419

Merged
merged 19 commits into from
Nov 6, 2024
Merged
Prev Previous commit
revert None support
Signed-off-by: Mecoli1219 <michaellai901026@gmail.com>
  • Loading branch information
Mecoli1219 committed Nov 6, 2024

Unverified

This commit is not signed, but one or more authors requires that any commit attributed to them is signed.
commit 43f8948dfa5fa20eb26fcab07bded8f9511d59ce
16 changes: 6 additions & 10 deletions flytekit/types/pickle/pickle.py
Original file line number Diff line number Diff line change
@@ -1,13 +1,13 @@
import os
import typing
from typing import Optional, Type
from typing import Type

import cloudpickle

from flytekit.core.context_manager import FlyteContext, FlyteContextManager
from flytekit.core.type_engine import AsyncTypeTransformer, TypeEngine
from flytekit.models.core import types as _core_types
from flytekit.models.literals import Blob, BlobMetadata, Literal, Scalar, Void
from flytekit.models.literals import Blob, BlobMetadata, Literal, Scalar
from flytekit.models.types import LiteralType

T = typing.TypeVar("T")
@@ -73,9 +73,7 @@ def assert_type(self, t: Type[T], v: T):
# Every type can serialize to pickle, so we don't need to check the type here.
...

async def async_to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[T]) -> Optional[T]:
if lv.scalar.blob is None:
return None
async def async_to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[T]) -> T:
uri = lv.scalar.blob.uri
return await FlytePickle.from_pickle(uri)

@@ -87,11 +85,10 @@ async def async_to_literal(
expected: LiteralType,
) -> Literal:
if python_val is None:
return Literal(scalar=Scalar(none_type=Void()))
raise AssertionError("Cannot pickle None Value.")
meta = BlobMetadata(
type=_core_types.BlobType(
format=self.PYTHON_PICKLE_FORMAT,
dimensionality=_core_types.BlobType.BlobDimensionality.SINGLE,
format=self.PYTHON_PICKLE_FORMAT, dimensionality=_core_types.BlobType.BlobDimensionality.SINGLE
)
)
remote_path = await FlytePickle.to_pickle(ctx, python_val)
@@ -110,8 +107,7 @@ def guess_python_type(self, literal_type: LiteralType) -> typing.Type[FlytePickl
def get_literal_type(self, t: Type[T]) -> LiteralType:
lt = LiteralType(
blob=_core_types.BlobType(
format=self.PYTHON_PICKLE_FORMAT,
dimensionality=_core_types.BlobType.BlobDimensionality.SINGLE,
format=self.PYTHON_PICKLE_FORMAT, dimensionality=_core_types.BlobType.BlobDimensionality.SINGLE
)
)
lt.metadata = {"python_class_name": str(t)}
5 changes: 3 additions & 2 deletions tests/flytekit/unit/core/test_generice_idl_type_engine.py
Original file line number Diff line number Diff line change
@@ -2067,8 +2067,9 @@ def __init__(self, number: int):
pv = transformer.to_python_value(ctx, lv, expected_python_type=gt)
assert Foo(1).number == pv.number

lt = TypeEngine.to_literal_type(typing.Optional[typing.Any])
TypeEngine.to_literal(ctx, None, FlytePickle, lt)
with pytest.raises(AssertionError, match="Cannot pickle None Value"):
lt = TypeEngine.to_literal_type(typing.Optional[typing.Any])
TypeEngine.to_literal(ctx, None, FlytePickle, lt)

with pytest.raises(
AssertionError,
6 changes: 3 additions & 3 deletions tests/flytekit/unit/core/test_type_engine.py
Original file line number Diff line number Diff line change
@@ -2066,9 +2066,9 @@ def __init__(self, number: int):
pv = transformer.to_python_value(ctx, lv, expected_python_type=gt)
assert Foo(1).number == pv.number

lt = TypeEngine.to_literal_type(typing.Optional[typing.Any])
lv = TypeEngine.to_literal(ctx, None, FlytePickle, lt)
assert lv.scalar.none_type == Void()
with pytest.raises(AssertionError, match="Cannot pickle None Value"):
lt = TypeEngine.to_literal_type(typing.Optional[typing.Any])
TypeEngine.to_literal(ctx, None, FlytePickle, lt)

with pytest.raises(
AssertionError,
21 changes: 10 additions & 11 deletions tests/flytekit/unit/core/test_type_hints.py
Original file line number Diff line number Diff line change
@@ -2234,7 +2234,6 @@ def wf1_with_pickle_untyped(a) -> int:

assert wf1_with_pickle_untyped(a=1) == 2
assert wf1_with_pickle_untyped(a="1") == 0
assert wf1_with_pickle_untyped(a=None) == 0

with pytest.raises(FlyteMissingTypeException):

@@ -2247,13 +2246,13 @@ def test_pickle_untyped_wf_and_task():
@task(pickle_untyped=True)
def t1(a):
if type(a) != int:
return None
return "t1"
return a + 1

@task(pickle_untyped=True)
def t2(a):
if type(a) != int:
return None
return "t2"
return a + 2

@workflow(pickle_untyped=True)
@@ -2262,20 +2261,20 @@ def wf1_with_pickle_untyped(a):
return t2(a=a1)

assert wf1_with_pickle_untyped(a=1) == 4
assert wf1_with_pickle_untyped(a="1") is None
assert wf1_with_pickle_untyped(a="1") == "t2"


def test_wf_with_pickle_untyped_and_safe_tasks():
def test_wf_with_pickle_untyped_and_regular_tasks():
@task(pickle_untyped=True)
def t1(a):
if type(a) != int:
return None
return "t1"
return a + 1

@task
def t2(a: typing.Any) -> typing.Any:
if type(a) != int:
return None
return "t2"
return a + 2

@workflow(pickle_untyped=True)
@@ -2284,23 +2283,23 @@ def wf1_with_pickle_untyped(a):
return t2(a=a1)

assert wf1_with_pickle_untyped(a=1) == 4
assert wf1_with_pickle_untyped(a="1") is None
assert wf1_with_pickle_untyped(a="1") == "t2"

@workflow(pickle_untyped=True)
def wf2_with_pickle_untyped(a):
a1 = t2(a=a)
return t1(a=a1)

assert wf2_with_pickle_untyped(a=1) == 4
assert wf2_with_pickle_untyped(a="1") is None
assert wf2_with_pickle_untyped(a="1") == "t1"


def test_pickle_untyped_task_with_specified_input():
@task(pickle_untyped=True)
def t1(a, b: typing.Any):
if type(a) != int:
if type(b) != int:
return None
return "t1"
else:
return b
elif type(b) != int:
@@ -2315,4 +2314,4 @@ def wf1_with_pickle_untyped(a: typing.Any, b):
assert wf1_with_pickle_untyped(a=1, b=2) == 3
assert wf1_with_pickle_untyped(a="1", b=2) == 2
assert wf1_with_pickle_untyped(a=1, b="2") == 1
assert wf1_with_pickle_untyped(a="1", b="2") is None
assert wf1_with_pickle_untyped(a="1", b="2") == "t1"