From 6b366013f531671548a011d708ee892ac47d7ecb Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Tue, 30 Nov 2021 22:57:00 +0800 Subject: [PATCH 1/6] Add support Flyte File and directory in dataclass Signed-off-by: Kevin Su --- flytekit/core/type_engine.py | 81 ++++++++++++++++---- flytekit/types/directory/types.py | 29 ++++--- flytekit/types/file/file.py | 37 +++++---- tests/flytekit/unit/core/test_type_engine.py | 65 +++++++++++++++- tests/flytekit/unit/core/test_type_hints.py | 78 ++++++++++++++++++- 5 files changed, 243 insertions(+), 47 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index ff9b257b56..1d7d0d6e9a 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -28,7 +28,16 @@ from flytekit.models import interface as _interface_models from flytekit.models import types as _type_models from flytekit.models.core import types as _core_types -from flytekit.models.literals import Literal, LiteralCollection, LiteralMap, Primitive, Scalar, Schema +from flytekit.models.literals import ( + Blob, + BlobMetadata, + Literal, + LiteralCollection, + LiteralMap, + Primitive, + Scalar, + Schema, +) from flytekit.models.types import LiteralType, SimpleType T = typing.TypeVar("T") @@ -258,27 +267,68 @@ def _serialize_flyte_type(self, python_val: T, python_type: Type[T]): """ If any field inside the dataclass is flyte type, we should use flyte type transformer for that field. """ - from flytekit.types.schema.types import FlyteSchema, FlyteSchemaTransformer + from flytekit.types.directory.types import FlyteDirectory + from flytekit.types.file import FlyteFile + from flytekit.types.schema.types import FlyteSchema for f in dataclasses.fields(python_type): v = python_val.__getattribute__(f.name) - if inspect.isclass(f.type) and issubclass(f.type, FlyteSchema): - FlyteSchemaTransformer().to_literal(FlyteContext.current_context(), v, f.type, None) + if inspect.isclass(f.type) and ( + issubclass(f.type, FlyteSchema) or issubclass(f.type, FlyteFile) or issubclass(f.type, FlyteDirectory) + ): + TypeEngine.to_literal(FlyteContext.current_context(), v, f.type, None) elif dataclasses.is_dataclass(f.type): self._serialize_flyte_type(v, f.type) def _deserialize_flyte_type(self, python_val: T, expected_python_type: Type["FlyteSchema"]): + from flytekit.types.directory.types import FlyteDirectory, FlyteDirToMultipartBlobTransformer + from flytekit.types.file.file import FlyteFile, FlyteFilePathTransformer from flytekit.types.schema.types import FlyteSchema, FlyteSchemaTransformer for f in dataclasses.fields(expected_python_type): v = python_val.__getattribute__(f.name) - if inspect.isclass(f.type) and issubclass(f.type, FlyteSchema): - t = FlyteSchemaTransformer() - t.to_python_value( - FlyteContext.current_context(), - Literal(scalar=Scalar(schema=Schema(v.remote_path, t._get_schema_type(f.type)))), - f.type, - ) + if inspect.isclass(f.type): + if issubclass(f.type, FlyteSchema): + t = FlyteSchemaTransformer() + t.to_python_value( + FlyteContext.current_context(), + Literal(scalar=Scalar(schema=Schema(v.remote_path, t._get_schema_type(f.type)))), + f.type, + ) + elif issubclass(f.type, FlyteFile): + FlyteFilePathTransformer().to_python_value( + FlyteContext.current_context(), + Literal( + scalar=Scalar( + blob=Blob( + metadata=BlobMetadata( + type=_core_types.BlobType( + format="", dimensionality=_core_types.BlobType.BlobDimensionality.SINGLE + ) + ), + uri=v.path, + ) + ) + ), + f.type, + ) + elif issubclass(f.type, FlyteDirectory): + FlyteDirToMultipartBlobTransformer().to_python_value( + FlyteContext.current_context(), + Literal( + scalar=Scalar( + blob=Blob( + metadata=BlobMetadata( + type=_core_types.BlobType( + format="", dimensionality=_core_types.BlobType.BlobDimensionality.MULTIPART + ) + ), + uri=v.path, + ) + ) + ), + f.type, + ) elif dataclasses.is_dataclass(f.type): self._deserialize_flyte_type(v, f.type) @@ -817,18 +867,21 @@ def convert_json_schema_to_python_class(schema: dict, schema_name) -> Type[datac """ attribute_list = [] for property_key, property_val in schema[schema_name]["properties"].items(): + property_type = property_val["type"] # Handle list if property_val["type"] == "array": attribute_list.append((property_key, typing.List[_get_element_type(property_val["items"])])) # Handle dataclass and dict - elif property_val["type"] == "object": - if "$ref" in property_val: + elif property_type == "object": + if property_val.get("$ref"): name = property_val["$ref"].split("/")[-1] attribute_list.append((property_key, convert_json_schema_to_python_class(schema, name))) - else: + elif property_val.get("additionalProperties"): attribute_list.append( (property_key, typing.Dict[str, _get_element_type(property_val["additionalProperties"])]) ) + else: + attribute_list.append((property_key, typing.Dict[str, _get_element_type(property_val)])) # Handle int, float, bool or str else: attribute_list.append([property_key, _get_element_type(property_val)]) diff --git a/flytekit/types/directory/types.py b/flytekit/types/directory/types.py index a0112c1eea..a5d7261336 100644 --- a/flytekit/types/directory/types.py +++ b/flytekit/types/directory/types.py @@ -3,9 +3,13 @@ import os import pathlib import typing +from dataclasses import dataclass, field from pathlib import Path -from flytekit.core.context_manager import FlyteContext +from dataclasses_json import config, dataclass_json +from marshmallow import fields + +from flytekit.core.context_manager import FlyteContext, FlyteContextManager from flytekit.core.type_engine import TypeEngine, TypeTransformer from flytekit.models import types as _type_models from flytekit.models.core import types as _core_types @@ -19,7 +23,10 @@ def noop(): ... +@dataclass_json +@dataclass class FlyteDirectory(os.PathLike, typing.Generic[T]): + path: typing.Union[str, os.PathLike] = field(default=None, metadata=config(mm_field=fields.String())) """ .. warning:: @@ -106,15 +113,17 @@ def t1(in1: FlyteDirectory["svg"]): field in the ``BlobType``. """ - def __init__(self, path: str, downloader: typing.Callable = None, remote_directory=None): + def __init__(self, path: typing.Union[str, os.PathLike], downloader: typing.Callable = None, remote_directory=None): """ :param path: The source path that users are expected to call open() on :param downloader: Optional function that can be passed that used to delay downloading of the actual fil until a user actually calls open(). :param remote_directory: If the user wants to return something and also specify where it should be uploaded to. """ - - self._path = path + ctx = FlyteContextManager.current_context() + # Make this field public, so that the dataclass transformer can set a value for it + # https://github.com/flyteorg/flytekit/blob/bcc8541bd6227b532f8462563fe8aac902242b21/flytekit/core/type_engine.py#L298 + self.path = path self._downloader = downloader or noop self._downloaded = False self._remote_directory = remote_directory @@ -127,7 +136,7 @@ def __fspath__(self): if not self._downloaded: self._downloader() self._downloaded = True - return self._path + return self.path @classmethod def extension(cls) -> str: @@ -159,10 +168,6 @@ def downloaded(self) -> bool: def remote_directory(self) -> typing.Optional[str]: return self._remote_directory - @property - def path(self) -> str: - return self._path - @property def remote_source(self) -> str: """ @@ -175,10 +180,10 @@ def download(self) -> str: return self.__fspath__() def __repr__(self): - return self._path + return self.path def __str__(self): - return self._path + return self.path class FlyteDirToMultipartBlobTransformer(TypeTransformer[FlyteDirectory]): @@ -238,7 +243,7 @@ def to_literal( return Literal(scalar=Scalar(blob=Blob(metadata=meta, uri=python_val._remote_source))) source_path = python_val.path - # If the user specified the remote_path to be False, that means no matter what, do not upload. Also if the + # If the user specified the remote_directory to be False, that means no matter what, do not upload. Also if the # path given is already a remote path, say https://www.google.com, the concept of uploading to the Flyte # blob store doesn't make sense. if python_val.remote_directory is False or ctx.file_access.is_remote(source_path): diff --git a/flytekit/types/file/file.py b/flytekit/types/file/file.py index 1fbcee049a..bd5894da7b 100644 --- a/flytekit/types/file/file.py +++ b/flytekit/types/file/file.py @@ -3,8 +3,12 @@ import os import pathlib import typing +from dataclasses import dataclass, field -from flytekit.core.context_manager import FlyteContext +from dataclasses_json import config, dataclass_json +from marshmallow import fields + +from flytekit.core.context_manager import FlyteContext, FlyteContextManager from flytekit.core.type_engine import TypeEngine, TypeTransformer from flytekit.loggers import logger from flytekit.models.core.types import BlobType @@ -19,7 +23,10 @@ def noop(): T = typing.TypeVar("T") +@dataclass_json +@dataclass class FlyteFile(os.PathLike, typing.Generic[T]): + path: typing.Union[str, os.PathLike] = field(default=None, metadata=config(mm_field=fields.String())) """ Since there is no native Python implementation of files and directories for the Flyte Blob type, (like how int exists for Flyte's Integer type) we need to create one so that users can express that their tasks take @@ -155,14 +162,19 @@ def extension(cls) -> str: return _SpecificFormatClass - def __init__(self, path: str, downloader: typing.Callable = noop, remote_path=None): + def __init__( + self, path: typing.Union[str, os.PathLike], downloader: typing.Callable = noop, remote_path: os.PathLike = None + ): """ :param path: The source path that users are expected to call open() on :param downloader: Optional function that can be passed that used to delay downloading of the actual fil until a user actually calls open(). :param remote_path: If the user wants to return something and also specify where it should be uploaded to. """ - self._path = path + ctx = FlyteContextManager.current_context() + # Make this field public, so that the dataclass transformer can set a value for it + # https://github.com/flyteorg/flytekit/blob/bcc8541bd6227b532f8462563fe8aac902242b21/flytekit/core/type_engine.py#L298 + self.path = path self._downloader = downloader self._downloaded = False self._remote_path = remote_path @@ -173,30 +185,26 @@ def __fspath__(self): if not self._downloaded: self._downloader() self._downloaded = True - return self._path + return self.path def __eq__(self, other): if isinstance(other, FlyteFile): return ( - self._path == other._path + self.path == other.path and self._remote_path == other._remote_path and self.extension() == other.extension() ) else: - return self._path == other + return self.path == other @property def downloaded(self) -> bool: return self._downloaded @property - def remote_path(self) -> typing.Optional[str]: + def remote_path(self) -> os.PathLike: return self._remote_path - @property - def path(self) -> str: - return self._path - @property def remote_source(self) -> str: """ @@ -209,10 +217,10 @@ def download(self) -> str: return self.__fspath__() def __repr__(self): - return self._path + return self.path def __str__(self): - return self._path + return self.path class FlyteFilePathTransformer(TypeTransformer[FlyteFile]): @@ -316,11 +324,10 @@ def to_literal( return Literal(scalar=Scalar(blob=Blob(metadata=meta, uri=source_path))) def to_python_value( - self, ctx: FlyteContext, lv: Literal, expected_python_type: typing.Union[typing.Type[FlyteFile]] + self, ctx: FlyteContext, lv: Literal, expected_python_type: typing.Union[typing.Type[FlyteFile], os.PathLike] ) -> FlyteFile: uri = lv.scalar.blob.uri - # In this condition, we still return a FlyteFile instance, but it's a simple one that has no downloading tricks # Using is instead of issubclass because FlyteFile does actually subclass it if expected_python_type is os.PathLike: diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 826c8df1fb..5ebaf587f2 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -1,5 +1,6 @@ import datetime import os +import tempfile import typing from dataclasses import asdict, dataclass from datetime import timedelta @@ -30,8 +31,9 @@ from flytekit.models.core.types import BlobType from flytekit.models.literals import Blob, BlobMetadata, Literal, LiteralCollection, LiteralMap, Primitive, Scalar from flytekit.models.types import LiteralType, SimpleType +from flytekit.types.directory import TensorboardLogs from flytekit.types.directory.types import FlyteDirectory -from flytekit.types.file import JPEGImageFile +from flytekit.types.file import JPEGImageFile, PNGImageFile from flytekit.types.file.file import FlyteFile, FlyteFilePathTransformer from flytekit.types.pickle import FlytePickle from flytekit.types.pickle.pickle import FlytePickleTransformer @@ -456,7 +458,6 @@ def test_dataclass_transformer(): }, }, } - tf = DataclassTransformer() t = tf.get_literal_type(TestStruct) assert t is not None @@ -506,6 +507,66 @@ def test_dataclass_int_preserving(): assert ot == o +def test_flyte_file_in_dataclass(): + @dataclass_json + @dataclass + class TestInnerFileStruct(object): + a: JPEGImageFile + b: typing.List[FlyteFile] + c: typing.Dict[str, FlyteFile] + + @dataclass_json + @dataclass + class TestFileStruct(object): + a: FlyteFile + b: typing.List[FlyteFile] + c: typing.Dict[str, FlyteFile] + d: TestInnerFileStruct + + f = FlyteFile("s3://tmp/file") + o = TestFileStruct( + a=f, b=[f], c={"hello": f}, d=TestInnerFileStruct(a=JPEGImageFile("s3://tmp/file.jpeg"), b=[f], c={"hello": f}) + ) + + ctx = FlyteContext.current_context() + tf = DataclassTransformer() + lt = tf.get_literal_type(TestFileStruct) + lv = tf.to_literal(ctx, o, TestFileStruct, lt) + ot = tf.to_python_value(ctx, lv=lv, expected_python_type=TestFileStruct) + assert o.a.path == ot.a.path + assert o == ot + + +def test_flyte_directory_in_dataclass(): + @dataclass_json + @dataclass + class TestInnerFileStruct(object): + a: TensorboardLogs + b: typing.List[FlyteDirectory] + c: typing.Dict[str, FlyteDirectory] + + @dataclass_json + @dataclass + class TestFileStruct(object): + a: FlyteDirectory + b: typing.List[FlyteDirectory] + c: typing.Dict[str, FlyteDirectory] + d: TestInnerFileStruct + + tempdir = tempfile.mkdtemp(prefix="flyte-") + f = FlyteDirectory(tempdir) + o = TestFileStruct( + a=f, b=[f], c={"hello": f}, d=TestInnerFileStruct(a=TensorboardLogs("s3://tensorboard"), b=[f], c={"hello": f}) + ) + + ctx = FlyteContext.current_context() + tf = DataclassTransformer() + lt = tf.get_literal_type(TestFileStruct) + lv = tf.to_literal(ctx, o, TestFileStruct, lt) + ot = tf.to_python_value(ctx, lv=lv, expected_python_type=TestFileStruct) + # assert o == ot + + # Enums should have string values class Color(Enum): RED = "red" diff --git a/tests/flytekit/unit/core/test_type_hints.py b/tests/flytekit/unit/core/test_type_hints.py index a0559377e7..5df2f8607f 100644 --- a/tests/flytekit/unit/core/test_type_hints.py +++ b/tests/flytekit/unit/core/test_type_hints.py @@ -18,7 +18,13 @@ from flytekit.common.translator import get_serializable from flytekit.core import context_manager, launch_plan, promise from flytekit.core.condition import conditional -from flytekit.core.context_manager import ExecutionState, FastSerializationSettings, Image, ImageConfig +from flytekit.core.context_manager import ( + ExecutionState, + FastSerializationSettings, + FlyteContextManager, + Image, + ImageConfig, +) from flytekit.core.data_persistence import FileAccessProvider from flytekit.core.node import Node from flytekit.core.promise import NodeOutput, Promise, VoidPromise @@ -32,6 +38,8 @@ from flytekit.models.interface import Parameter from flytekit.models.task import Resources as _resource_models from flytekit.models.types import LiteralType, SimpleType +from flytekit.types.directory import FlyteDirectory, TensorboardLogs +from flytekit.types.file import FlyteFile, PNGImageFile from flytekit.types.schema import FlyteSchema, SchemaOpenMode serialization_settings = context_manager.SerializationSettings( @@ -348,6 +356,68 @@ def test_user_demo_test(mock_sql): assert context_manager.FlyteContextManager.size() == 1 +def test_flyte_file_in_dataclass(): + @dataclass_json + @dataclass + class InnerFileStruct(object): + a: FlyteFile + b: PNGImageFile + + @dataclass_json + @dataclass + class FileStruct(object): + a: FlyteFile + b: InnerFileStruct + + @task + def t1(path: str) -> FileStruct: + file = FlyteFile(path) + fs = FileStruct(a=file, b=InnerFileStruct(a=file, b=PNGImageFile(path))) + return fs + + @task + def t2(fs: FileStruct) -> os.PathLike: + return fs.a.path + + @workflow + def wf(path: str) -> os.PathLike: + n1 = t1(path=path) + return t2(fs=n1) + + assert wf(path="s3://somewhere") == "s3://somewhere" + + +def test_flyte_directory_in_dataclass(): + @dataclass_json + @dataclass + class InnerFileStruct(object): + a: FlyteDirectory + b: TensorboardLogs + + @dataclass_json + @dataclass + class FileStruct(object): + a: FlyteDirectory + b: InnerFileStruct + + @task + def t1(path: str) -> FileStruct: + dir = FlyteDirectory(path) + fs = FileStruct(a=dir, b=InnerFileStruct(a=dir, b=TensorboardLogs(path))) + return fs + + @task + def t2(fs: FileStruct) -> os.PathLike: + return fs.a.path + + @workflow + def wf(path: str) -> os.PathLike: + n1 = t1(path=path) + return t2(fs=n1) + + assert wf(path="s3://somewhere") == "s3://somewhere" + + def test_wf1_with_map(): @task def t1(a: int) -> int: @@ -1429,7 +1499,7 @@ class Foo(object): @dataclass class Bar(object): x: int - y: str + y: dict z: Foo @task @@ -1448,14 +1518,14 @@ def t1() -> Foo: @task def t2() -> Bar: - return Bar(x=1, y="bar", z=Foo(x=1, y="foo", z={"hello": "world"})) + return Bar(x=1, y={"hello": "world"}, z=Foo(x=1, y="foo", z={"hello": "world"})) task_spec = get_serializable(OrderedDict(), serialization_settings, t2) pt_map = TypeEngine.guess_python_types(task_spec.template.interface.outputs) assert dataclasses.is_dataclass(pt_map["o0"]) output_lm = t2.dispatch_execute(ctx, _literal_models.LiteralMap(literals={})) - expected_struct.update({"x": 1, "y": "bar", "z": {"x": 1, "y": "foo", "z": {"hello": "world"}}}) + expected_struct.update({"x": 1, "y": {"hello": "world"}, "z": {"x": 1, "y": "foo", "z": {"hello": "world"}}}) assert output_lm.literals["o0"].scalar.generic == expected_struct From 2da054b0adaf393d341e9eb482bb0d05fa541587 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Wed, 1 Dec 2021 04:01:50 +0800 Subject: [PATCH 2/6] Fixed tests Signed-off-by: Kevin Su --- flytekit/types/directory/types.py | 3 +-- flytekit/types/file/file.py | 3 +-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/flytekit/types/directory/types.py b/flytekit/types/directory/types.py index a5d7261336..128aebab27 100644 --- a/flytekit/types/directory/types.py +++ b/flytekit/types/directory/types.py @@ -9,7 +9,7 @@ from dataclasses_json import config, dataclass_json from marshmallow import fields -from flytekit.core.context_manager import FlyteContext, FlyteContextManager +from flytekit.core.context_manager import FlyteContext from flytekit.core.type_engine import TypeEngine, TypeTransformer from flytekit.models import types as _type_models from flytekit.models.core import types as _core_types @@ -120,7 +120,6 @@ def __init__(self, path: typing.Union[str, os.PathLike], downloader: typing.Call until a user actually calls open(). :param remote_directory: If the user wants to return something and also specify where it should be uploaded to. """ - ctx = FlyteContextManager.current_context() # Make this field public, so that the dataclass transformer can set a value for it # https://github.com/flyteorg/flytekit/blob/bcc8541bd6227b532f8462563fe8aac902242b21/flytekit/core/type_engine.py#L298 self.path = path diff --git a/flytekit/types/file/file.py b/flytekit/types/file/file.py index bd5894da7b..c2568e6a6f 100644 --- a/flytekit/types/file/file.py +++ b/flytekit/types/file/file.py @@ -8,7 +8,7 @@ from dataclasses_json import config, dataclass_json from marshmallow import fields -from flytekit.core.context_manager import FlyteContext, FlyteContextManager +from flytekit.core.context_manager import FlyteContext from flytekit.core.type_engine import TypeEngine, TypeTransformer from flytekit.loggers import logger from flytekit.models.core.types import BlobType @@ -171,7 +171,6 @@ def __init__( until a user actually calls open(). :param remote_path: If the user wants to return something and also specify where it should be uploaded to. """ - ctx = FlyteContextManager.current_context() # Make this field public, so that the dataclass transformer can set a value for it # https://github.com/flyteorg/flytekit/blob/bcc8541bd6227b532f8462563fe8aac902242b21/flytekit/core/type_engine.py#L298 self.path = path From 2a998332ab388f7f27951d62e7f6aec6ee209d68 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Wed, 1 Dec 2021 04:12:20 +0800 Subject: [PATCH 3/6] Fixed tests Signed-off-by: Kevin Su --- tests/flytekit/unit/core/test_type_engine.py | 4 ++-- tests/flytekit/unit/core/test_type_hints.py | 8 +------- 2 files changed, 3 insertions(+), 9 deletions(-) diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 5ebaf587f2..b26df2af38 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -33,7 +33,7 @@ from flytekit.models.types import LiteralType, SimpleType from flytekit.types.directory import TensorboardLogs from flytekit.types.directory.types import FlyteDirectory -from flytekit.types.file import JPEGImageFile, PNGImageFile +from flytekit.types.file import JPEGImageFile from flytekit.types.file.file import FlyteFile, FlyteFilePathTransformer from flytekit.types.pickle import FlytePickle from flytekit.types.pickle.pickle import FlytePickleTransformer @@ -564,7 +564,7 @@ class TestFileStruct(object): lt = tf.get_literal_type(TestFileStruct) lv = tf.to_literal(ctx, o, TestFileStruct, lt) ot = tf.to_python_value(ctx, lv=lv, expected_python_type=TestFileStruct) - # assert o == ot + assert o == ot # Enums should have string values diff --git a/tests/flytekit/unit/core/test_type_hints.py b/tests/flytekit/unit/core/test_type_hints.py index 5df2f8607f..ee32727055 100644 --- a/tests/flytekit/unit/core/test_type_hints.py +++ b/tests/flytekit/unit/core/test_type_hints.py @@ -18,13 +18,7 @@ from flytekit.common.translator import get_serializable from flytekit.core import context_manager, launch_plan, promise from flytekit.core.condition import conditional -from flytekit.core.context_manager import ( - ExecutionState, - FastSerializationSettings, - FlyteContextManager, - Image, - ImageConfig, -) +from flytekit.core.context_manager import ExecutionState, FastSerializationSettings, Image, ImageConfig from flytekit.core.data_persistence import FileAccessProvider from flytekit.core.node import Node from flytekit.core.promise import NodeOutput, Promise, VoidPromise From dd6e79ab51339d0b907bebc6fc8ccc6a9809b6f4 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Sat, 4 Dec 2021 00:46:37 +0800 Subject: [PATCH 4/6] Updated Signed-off-by: Kevin Su --- flytekit/core/type_engine.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 1d7d0d6e9a..197b34da5f 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -290,13 +290,13 @@ def _deserialize_flyte_type(self, python_val: T, expected_python_type: Type["Fly if inspect.isclass(f.type): if issubclass(f.type, FlyteSchema): t = FlyteSchemaTransformer() - t.to_python_value( + v = t.to_python_value( FlyteContext.current_context(), Literal(scalar=Scalar(schema=Schema(v.remote_path, t._get_schema_type(f.type)))), f.type, ) elif issubclass(f.type, FlyteFile): - FlyteFilePathTransformer().to_python_value( + v = FlyteFilePathTransformer().to_python_value( FlyteContext.current_context(), Literal( scalar=Scalar( @@ -313,7 +313,7 @@ def _deserialize_flyte_type(self, python_val: T, expected_python_type: Type["Fly f.type, ) elif issubclass(f.type, FlyteDirectory): - FlyteDirToMultipartBlobTransformer().to_python_value( + v = FlyteDirToMultipartBlobTransformer().to_python_value( FlyteContext.current_context(), Literal( scalar=Scalar( From a614ba6801505acaeb901c37ec6d00070f817873 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Sat, 4 Dec 2021 03:20:19 +0800 Subject: [PATCH 5/6] Updated Signed-off-by: Kevin Su --- flytekit/core/type_engine.py | 98 +++++++++++--------- tests/flytekit/unit/core/test_type_engine.py | 40 ++++---- 2 files changed, 77 insertions(+), 61 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 197b34da5f..b4e53146fb 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -280,57 +280,66 @@ def _serialize_flyte_type(self, python_val: T, python_type: Type[T]): elif dataclasses.is_dataclass(f.type): self._serialize_flyte_type(v, f.type) - def _deserialize_flyte_type(self, python_val: T, expected_python_type: Type["FlyteSchema"]): + def _deserialize_flyte_type(self, python_val: T, expected_python_type: Type) -> T: from flytekit.types.directory.types import FlyteDirectory, FlyteDirToMultipartBlobTransformer from flytekit.types.file.file import FlyteFile, FlyteFilePathTransformer from flytekit.types.schema.types import FlyteSchema, FlyteSchemaTransformer - for f in dataclasses.fields(expected_python_type): - v = python_val.__getattribute__(f.name) - if inspect.isclass(f.type): - if issubclass(f.type, FlyteSchema): - t = FlyteSchemaTransformer() - v = t.to_python_value( - FlyteContext.current_context(), - Literal(scalar=Scalar(schema=Schema(v.remote_path, t._get_schema_type(f.type)))), - f.type, - ) - elif issubclass(f.type, FlyteFile): - v = FlyteFilePathTransformer().to_python_value( - FlyteContext.current_context(), - Literal( - scalar=Scalar( - blob=Blob( - metadata=BlobMetadata( - type=_core_types.BlobType( - format="", dimensionality=_core_types.BlobType.BlobDimensionality.SINGLE - ) - ), - uri=v.path, + if not dataclasses.is_dataclass(expected_python_type): + return python_val + + if issubclass(expected_python_type, FlyteSchema): + t = FlyteSchemaTransformer() + return t.to_python_value( + FlyteContext.current_context(), + Literal(scalar=Scalar(schema=Schema(python_val.remote_path, t._get_schema_type(expected_python_type)))), + expected_python_type, + ) + elif issubclass(expected_python_type, FlyteFile): + return FlyteFilePathTransformer().to_python_value( + FlyteContext.current_context(), + Literal( + scalar=Scalar( + blob=Blob( + metadata=BlobMetadata( + type=_core_types.BlobType( + format="", dimensionality=_core_types.BlobType.BlobDimensionality.SINGLE ) - ) - ), - f.type, + ), + uri=python_val.path, + ) ) - elif issubclass(f.type, FlyteDirectory): - v = FlyteDirToMultipartBlobTransformer().to_python_value( - FlyteContext.current_context(), - Literal( - scalar=Scalar( - blob=Blob( - metadata=BlobMetadata( - type=_core_types.BlobType( - format="", dimensionality=_core_types.BlobType.BlobDimensionality.MULTIPART - ) - ), - uri=v.path, + ), + expected_python_type, + ) + elif issubclass(expected_python_type, FlyteDirectory): + return FlyteDirToMultipartBlobTransformer().to_python_value( + FlyteContext.current_context(), + Literal( + scalar=Scalar( + blob=Blob( + metadata=BlobMetadata( + type=_core_types.BlobType( + format="", dimensionality=_core_types.BlobType.BlobDimensionality.MULTIPART ) - ) - ), - f.type, + ), + uri=python_val.path, + ) ) - elif dataclasses.is_dataclass(f.type): - self._deserialize_flyte_type(v, f.type) + ), + expected_python_type, + ) + else: + for f in dataclasses.fields(expected_python_type): + value = python_val.__getattribute__(f.name) + if hasattr(f.type, "__origin__") and f.type.__origin__ is list: + value = [self._deserialize_flyte_type(v, f.type.__args__[0]) for v in value] + elif hasattr(f.type, "__origin__") and f.type.__origin__ is dict: + value = {k: self._deserialize_flyte_type(v, f.type.__args__[1]) for k, v in value.items()} + else: + value = self._deserialize_flyte_type(value, f.type) + python_val.__setattr__(f.name, value) + return python_val def _fix_val_int(self, t: typing.Type, val: typing.Any) -> typing.Any: if t == int: @@ -376,8 +385,7 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: ) dc = cast(DataClassJsonMixin, expected_python_type).from_json(_json_format.MessageToJson(lv.scalar.generic)) - self._deserialize_flyte_type(dc, expected_python_type) - return self._fix_dataclass_int(expected_python_type, dc) + return self._fix_dataclass_int(expected_python_type, self._deserialize_flyte_type(dc, expected_python_type)) def guess_python_type(self, literal_type: LiteralType) -> Type[T]: if literal_type.simple == SimpleType.STRUCT: diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index b26df2af38..2041acdf8e 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -34,7 +34,7 @@ from flytekit.types.directory import TensorboardLogs from flytekit.types.directory.types import FlyteDirectory from flytekit.types.file import JPEGImageFile -from flytekit.types.file.file import FlyteFile, FlyteFilePathTransformer +from flytekit.types.file.file import FlyteFile, FlyteFilePathTransformer, noop from flytekit.types.pickle import FlytePickle from flytekit.types.pickle.pickle import FlytePickleTransformer from flytekit.types.schema import FlyteSchema @@ -519,22 +519,25 @@ class TestInnerFileStruct(object): @dataclass class TestFileStruct(object): a: FlyteFile - b: typing.List[FlyteFile] - c: typing.Dict[str, FlyteFile] - d: TestInnerFileStruct + b: TestInnerFileStruct f = FlyteFile("s3://tmp/file") - o = TestFileStruct( - a=f, b=[f], c={"hello": f}, d=TestInnerFileStruct(a=JPEGImageFile("s3://tmp/file.jpeg"), b=[f], c={"hello": f}) - ) + o = TestFileStruct(a=f, b=TestInnerFileStruct(a=JPEGImageFile("s3://tmp/file.jpeg"), b=[f], c={"hello": f})) ctx = FlyteContext.current_context() tf = DataclassTransformer() lt = tf.get_literal_type(TestFileStruct) lv = tf.to_literal(ctx, o, TestFileStruct, lt) ot = tf.to_python_value(ctx, lv=lv, expected_python_type=TestFileStruct) - assert o.a.path == ot.a.path - assert o == ot + assert ot.a._downloader is not noop + assert ot.b.a._downloader is not noop + assert ot.b.b[0]._downloader is not noop + assert ot.b.c["hello"]._downloader is not noop + + assert o.a.path == ot.a.remote_source + assert o.b.a.path == ot.b.a.remote_source + assert o.b.b[0].path == ot.b.b[0].remote_source + assert o.b.c["hello"].path == ot.b.c["hello"].remote_source def test_flyte_directory_in_dataclass(): @@ -549,22 +552,27 @@ class TestInnerFileStruct(object): @dataclass class TestFileStruct(object): a: FlyteDirectory - b: typing.List[FlyteDirectory] - c: typing.Dict[str, FlyteDirectory] - d: TestInnerFileStruct + b: TestInnerFileStruct tempdir = tempfile.mkdtemp(prefix="flyte-") f = FlyteDirectory(tempdir) - o = TestFileStruct( - a=f, b=[f], c={"hello": f}, d=TestInnerFileStruct(a=TensorboardLogs("s3://tensorboard"), b=[f], c={"hello": f}) - ) + o = TestFileStruct(a=f, b=TestInnerFileStruct(a=TensorboardLogs("s3://tensorboard"), b=[f], c={"hello": f})) ctx = FlyteContext.current_context() tf = DataclassTransformer() lt = tf.get_literal_type(TestFileStruct) lv = tf.to_literal(ctx, o, TestFileStruct, lt) ot = tf.to_python_value(ctx, lv=lv, expected_python_type=TestFileStruct) - assert o == ot + + assert ot.a._downloader is not noop + assert ot.b.a._downloader is not noop + assert ot.b.b[0]._downloader is not noop + assert ot.b.c["hello"]._downloader is not noop + + assert o.a.path == ot.a.path + assert o.b.a.path == ot.b.a.remote_source + assert o.b.b[0].path == ot.b.b[0].path + assert o.b.c["hello"].path == ot.b.c["hello"].path # Enums should have string values From 7288a2c07bc3b8abeb8f0f6b118d664a910c88f5 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Mon, 6 Dec 2021 15:29:08 +0800 Subject: [PATCH 6/6] Fixed tests Signed-off-by: Kevin Su --- tests/flytekit/unit/core/test_type_hints.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/flytekit/unit/core/test_type_hints.py b/tests/flytekit/unit/core/test_type_hints.py index ee32727055..fcebeb53d6 100644 --- a/tests/flytekit/unit/core/test_type_hints.py +++ b/tests/flytekit/unit/core/test_type_hints.py @@ -378,7 +378,7 @@ def wf(path: str) -> os.PathLike: n1 = t1(path=path) return t2(fs=n1) - assert wf(path="s3://somewhere") == "s3://somewhere" + assert "/tmp/flyte/" in wf(path="s3://somewhere").path def test_flyte_directory_in_dataclass(): @@ -409,7 +409,7 @@ def wf(path: str) -> os.PathLike: n1 = t1(path=path) return t2(fs=n1) - assert wf(path="s3://somewhere") == "s3://somewhere" + assert "/tmp/flyte/" in wf(path="s3://somewhere").path def test_wf1_with_map():