diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index ff9b257b56..b4e53146fb 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,29 +267,79 @@ 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"]): + 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) 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, - ) - elif dataclasses.is_dataclass(f.type): - self._deserialize_flyte_type(v, f.type) + 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 + ) + ), + uri=python_val.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 + ) + ), + uri=python_val.path, + ) + ) + ), + 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: @@ -326,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: @@ -817,18 +875,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..128aebab27 100644 --- a/flytekit/types/directory/types.py +++ b/flytekit/types/directory/types.py @@ -3,8 +3,12 @@ import os import pathlib import typing +from dataclasses import dataclass, field from pathlib import Path +from dataclasses_json import config, dataclass_json +from marshmallow import fields + from flytekit.core.context_manager import FlyteContext from flytekit.core.type_engine import TypeEngine, TypeTransformer from flytekit.models import types as _type_models @@ -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,16 @@ 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 + # 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 +135,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 +167,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 +179,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 +242,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..c2568e6a6f 100644 --- a/flytekit/types/file/file.py +++ b/flytekit/types/file/file.py @@ -3,6 +3,10 @@ import os import pathlib import typing +from dataclasses import dataclass, field + +from dataclasses_json import config, dataclass_json +from marshmallow import fields from flytekit.core.context_manager import FlyteContext from flytekit.core.type_engine import TypeEngine, TypeTransformer @@ -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,18 @@ 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 + # 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 +184,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 +216,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 +323,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..2041acdf8e 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,9 +31,10 @@ 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.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 @@ -456,7 +458,6 @@ def test_dataclass_transformer(): }, }, } - tf = DataclassTransformer() t = tf.get_literal_type(TestStruct) assert t is not None @@ -506,6 +507,74 @@ 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: TestInnerFileStruct + + f = FlyteFile("s3://tmp/file") + 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 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(): + @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: TestInnerFileStruct + + tempdir = tempfile.mkdtemp(prefix="flyte-") + f = FlyteDirectory(tempdir) + 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 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 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..fcebeb53d6 100644 --- a/tests/flytekit/unit/core/test_type_hints.py +++ b/tests/flytekit/unit/core/test_type_hints.py @@ -32,6 +32,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 +350,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 "/tmp/flyte/" in wf(path="s3://somewhere").path + + +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 "/tmp/flyte/" in wf(path="s3://somewhere").path + + def test_wf1_with_map(): @task def t1(a: int) -> int: @@ -1429,7 +1493,7 @@ class Foo(object): @dataclass class Bar(object): x: int - y: str + y: dict z: Foo @task @@ -1448,14 +1512,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