Override Dataclass Serialization/Deserialization Behavior for FlyteTypes by mashumaro

[Future-Outlier]

Why are the changes needed?

In the dataclass transformer, we call _serialize_flyte_types before converting our dataclass to a JSON string and _deserialize_flyte_type after converting the JSON string back to dataclasses.
This symmetric behavior is implemented to apply certain modifications, allowing Flyte Types to function as desired.
We must move these codes to the Flyte Types's method.

1. Use the mashumaro SerializableType interface so we can define our serialization/deserialization in my class's method.
2. remove from json and to json in the dataclass Transformer so that we can use mashumaro SerializableType instead of dataclases_json's method.

Note: pydantic will call serialize method from here to serialize values from flyte.
Whether this change is introduced or not, it will not affect the behavior of pydantic now.

What changes were proposed in this pull request?

1. Use mashumaro SerializableType interface to define our serialization and deserialization for FlyteTypes.
   We move code in function serialize flyte types and deserialize flyte types the _serialize and deserialize function in Flyte Types.
2. Remove from_json and to_jsonin the dataclass Transformer, since we have to use mashumaro's to serialize and deserialize, we can't use dataclasess_json's method.

Flyte types: FlyteSchema, FlyteFile, FlyteDirectory and StructuredDataset

Note: I didn't make any changes in FlyteSchema because it doesn't do any additional steps in the _serialize_flyte_types, so there's no reason to implement it.

How was this patch tested?

local execution, remote execution and unit tests.
For remote execution, we have to use old image to test backward compatibility.

Scenario	Upstream Image	Downstream Image
1	Current Branch	Current Branch
2	Current Branch	ghcr.io/flyteorg/flytekit:py3.9-1.13.0
3	ghcr.io/flyteorg/flytekit:py3.9-1.13.0	Current Branch

Setup process

python example

from dataclasses import dataclass
from typing import Annotated, List, Dict, Optional
from flytekit import kwtypes, task, workflow
from flytekit import ImageSpec
from flytekit.types.directory import FlyteDirectory
from flytekit.types.file import FlyteFile
from flytekit.types.structured import StructuredDataset
from dataclasses_json import DataClassJsonMixin, dataclass_json
from mashumaro.mixins.json import DataClassJSONMixin
from flytekit.types.schema import FlyteSchema
# from flytekitplugins import envd


key = "3c03cac15c381431611d5e5000018864af94145b"
flytekit_dev_version = f"https://github.com/flyteorg/flytekit.git@{key}"
image = ImageSpec(
    # registry="localhost:30000",
    registry="futureoutlier",
    apt_packages=["git"],
    packages=[
        f"git+{flytekit_dev_version}",
    ],
    builder="default",
    # builder="envd",
)
# image.run_command = ["""echo "Acquire::http::Pipeline-Depth 0; \n Acquire::http::No-Cache true; \n Acquire::BrokenProxy    true;" > /etc/apt/apt.conf.d/99fixbadproxy"""]

# image = "futureoutlier/flytekit:0712-0347"

# old_image = ImageSpec(registry="futureoutlier",
#                       base_image="ghcr.io/flyteorg/flytekit:py3.9-1.9.0",
#                         apt_packages=["git"],
#                         packages=["mashumaro"])

# old_image = "ghcr.io/flyteorg/flytekit:py3.9-1.9.0"


@dataclass
class FlyteTypes:
    flytefile: FlyteFile
    flytedir: FlyteDirectory
    structured_dataset: StructuredDataset
    fs: FlyteSchema

@dataclass
class NestedFlyteTypes:
    flytefile: FlyteFile
    flytedir: FlyteDirectory
    structured_dataset: StructuredDataset
    fs: FlyteSchema
    flyte_types: FlyteTypes
    list_flyte_types: List[FlyteTypes]
    dict_flyte_types: Dict[str, FlyteTypes]
    optional_flyte_types: Optional[FlyteTypes] = None

@dataclass
class FlyteTypesWithDataClassJson(DataClassJsonMixin):
    flytefile: FlyteFile
    flytedir: FlyteDirectory
    structured_dataset: StructuredDataset
    fs: FlyteSchema

@dataclass
class NestedFlyteTypesWithDataClassJson(DataClassJsonMixin):
    flytefile: FlyteFile
    flytedir: FlyteDirectory
    structured_dataset: StructuredDataset
    fs: FlyteSchema
    flyte_types: FlyteTypesWithDataClassJson
    list_flyte_types: List[FlyteTypesWithDataClassJson]
    dict_flyte_types: Dict[str, FlyteTypesWithDataClassJson]
    optional_flyte_types: Optional[FlyteTypesWithDataClassJson] = None

@dataclass
class FlyteTypesWithMashumaro(DataClassJSONMixin):
    flytefile: FlyteFile
    flytedir: FlyteDirectory
    structured_dataset: StructuredDataset
    fs: FlyteSchema
@dataclass
class NestedFlyteTypesWithMashumaro(DataClassJSONMixin):
    flytefile: FlyteFile
    flytedir: FlyteDirectory
    structured_dataset: StructuredDataset
    flyte_types: FlyteTypesWithMashumaro
    fs: FlyteSchema
    list_flyte_types: List[FlyteTypesWithMashumaro]
    dict_flyte_types: Dict[str, FlyteTypesWithMashumaro]
    optional_flyte_types: Optional[FlyteTypesWithMashumaro] = None

@task(container_image=image)
# @task
def pass_and_return_flyte_types(a: NestedFlyteTypes,
                                b: NestedFlyteTypesWithDataClassJson,
                                c: NestedFlyteTypesWithMashumaro) -> (NestedFlyteTypes, NestedFlyteTypesWithDataClassJson, NestedFlyteTypesWithMashumaro):
    return a,b,c


@workflow
def wf(a: NestedFlyteTypes = NestedFlyteTypes(
            flytefile=FlyteFile(path="gs://opta-gcp-dogfood-gcp/92/f643fd4a06a9b4b57858-n0-0/5dc8c51e8ad277b3829be47547287bbf/example.txt"),
            flytedir=FlyteDirectory(path="gs://opta-gcp-dogfood-gcp/92/f643fd4a06a9b4b57858-n0-0/5dc8c51e8ad277b3829be47547287bbf"),
            structured_dataset=StructuredDataset(uri="gs://opta-gcp-dogfood-gcp/92/f643fd4a06a9b4b57858-n0-0/5dc8c51e8ad277b3829be47547287bbf/example.txt", file_format="txt"),
            fs=FlyteSchema(remote_path="s3://my-bucket/my-path"),
            flyte_types=FlyteTypes(
                flytefile=FlyteFile(
                    path="gs://opta-gcp-dogfood-gcp/92/f643fd4a06a9b4b57858-n0-0/5dc8c51e8ad277b3829be47547287bbf/example.txt"),
                flytedir=FlyteDirectory(
                    path="gs://opta-gcp-dogfood-gcp/92/f643fd4a06a9b4b57858-n0-0/5dc8c51e8ad277b3829be47547287bbf"),
                structured_dataset=StructuredDataset(
                    uri="gs://opta-gcp-dogfood-gcp/92/f643fd4a06a9b4b57858-n0-0/5dc8c51e8ad277b3829be47547287bbf/example.txt",
                    file_format="txt"),
                fs=FlyteSchema(remote_path="s3://my-bucket/my-path"),
            ),
            list_flyte_types=[FlyteTypes(
                flytefile=FlyteFile(
                    path="gs://opta-gcp-dogfood-gcp/92/f643fd4a06a9b4b57858-n0-0/5dc8c51e8ad277b3829be47547287bbf/example.txt"),
                flytedir=FlyteDirectory(
                    path="gs://opta-gcp-dogfood-gcp/92/f643fd4a06a9b4b57858-n0-0/5dc8c51e8ad277b3829be47547287bbf"),
                structured_dataset=StructuredDataset(
                    uri="gs://opta-gcp-dogfood-gcp/92/f643fd4a06a9b4b57858-n0-0/5dc8c51e8ad277b3829be47547287bbf/example.txt",
                    file_format="txt"),
                fs=FlyteSchema(remote_path="s3://my-bucket/my-path"),
            )],
            dict_flyte_types={"a": FlyteTypes(
                flytefile=FlyteFile(
                    path="gs://opta-gcp-dogfood-gcp/92/f643fd4a06a9b4b57858-n0-0/5dc8c51e8ad277b3829be47547287bbf/example.txt"),
                flytedir=FlyteDirectory(
                    path="gs://opta-gcp-dogfood-gcp/92/f643fd4a06a9b4b57858-n0-0/5dc8c51e8ad277b3829be47547287bbf"),
                structured_dataset=StructuredDataset(
                    uri="gs://opta-gcp-dogfood-gcp/92/f643fd4a06a9b4b57858-n0-0/5dc8c51e8ad277b3829be47547287bbf/example.txt",
                    file_format="txt"),
                fs=FlyteSchema(remote_path="s3://my-bucket/my-path"),
            )},
        )
       ,b: NestedFlyteTypesWithDataClassJson = NestedFlyteTypesWithDataClassJson(
            flytefile=FlyteFile(path="gs://opta-gcp-dogfood-gcp/92/f643fd4a06a9b4b57858-n0-0/5dc8c51e8ad277b3829be47547287bbf/example.txt"),
            flytedir=FlyteDirectory(path="gs://opta-gcp-dogfood-gcp/92/f643fd4a06a9b4b57858-n0-0/5dc8c51e8ad277b3829be47547287bbf"),
            structured_dataset=StructuredDataset(uri="gs://opta-gcp-dogfood-gcp/92/f643fd4a06a9b4b57858-n0-0/5dc8c51e8ad277b3829be47547287bbf/example.txt", file_format="txt"),
            fs=FlyteSchema(remote_path="s3://my-bucket/my-path"),
            flyte_types=FlyteTypesWithDataClassJson(
                flytefile=FlyteFile(path="gs://opta-gcp-dogfood-gcp/92/f643fd4a06a9b4b57858-n0-0/5dc8c51e8ad277b3829be47547287bbf/example.txt"),
                flytedir=FlyteDirectory(path="gs://opta-gcp-dogfood-gcp/92/f643fd4a06a9b4b57858-n0-0/5dc8c51e8ad277b3829be47547287bbf"),
                structured_dataset=StructuredDataset(uri="gs://opta-gcp-dogfood-gcp/92/f643fd4a06a9b4b57858-n0-0/5dc8c51e8ad277b3829be47547287bbf/example.txt", file_format="txt"),
                fs=FlyteSchema(remote_path="s3://my-bucket/my-path"),
            ),
            list_flyte_types=[FlyteTypesWithDataClassJson(
                flytefile=FlyteFile(path="gs://opta-gcp-dogfood-gcp/92/f643fd4a06a9b4b57858-n0-0/5dc8c51e8ad277b3829be47547287bbf/example.txt"),
                flytedir=FlyteDirectory(path="gs://opta-gcp-dogfood-gcp/92/f643fd4a06a9b4b57858-n0-0/5dc8c51e8ad277b3829be47547287bbf"),
                structured_dataset=StructuredDataset(uri="gs://opta-gcp-dogfood-gcp/92/f643fd4a06a9b4b57858-n0-0/5dc8c51e8ad277b3829be47547287bbf/example.txt", file_format="txt"),
                fs=FlyteSchema(remote_path="s3://my-bucket/my-path"),
            )],
            dict_flyte_types={"a": FlyteTypesWithDataClassJson(
                flytefile=FlyteFile(path="gs://opta-gcp-dogfood-gcp/92/f643fd4a06a9b4b57858-n0-0/5dc8c51e8ad277b3829be47547287bbf/example.txt"),
                flytedir=FlyteDirectory(path="gs://opta-gcp-dogfood-gcp/92/f643fd4a06a9b4b57858-n0-0/5dc8c51e8ad277b3829be47547287bbf"),
                structured_dataset=StructuredDataset(uri="gs://opta-gcp-dogfood-gcp/92/f643fd4a06a9b4b57858-n0-0/5dc8c51e8ad277b3829be47547287bbf/example.txt", file_format="txt"),
                fs=FlyteSchema(remote_path="s3://my-bucket/my-path"),
            )},
        ),
        c: NestedFlyteTypesWithMashumaro = NestedFlyteTypesWithMashumaro(
            flytefile=FlyteFile(
                path="gs://opta-gcp-dogfood-gcp/92/f643fd4a06a9b4b57858-n0-0/5dc8c51e8ad277b3829be47547287bbf/example.txt"),
            flytedir=FlyteDirectory(
                path="gs://opta-gcp-dogfood-gcp/92/f643fd4a06a9b4b57858-n0-0/5dc8c51e8ad277b3829be47547287bbf"),
            structured_dataset=StructuredDataset(
                uri="gs://opta-gcp-dogfood-gcp/92/f643fd4a06a9b4b57858-n0-0/5dc8c51e8ad277b3829be47547287bbf/example.txt",
                file_format="txt"),
            fs=FlyteSchema(remote_path="s3://my-bucket/my-path"),
            flyte_types=FlyteTypesWithMashumaro(
                flytefile=FlyteFile(
                    path="gs://opta-gcp-dogfood-gcp/92/f643fd4a06a9b4b57858-n0-0/5dc8c51e8ad277b3829be47547287bbf/example.txt"),
                flytedir=FlyteDirectory(
                    path="gs://opta-gcp-dogfood-gcp/92/f643fd4a06a9b4b57858-n0-0/5dc8c51e8ad277b3829be47547287bbf"),
                structured_dataset=StructuredDataset(
                    uri="gs://opta-gcp-dogfood-gcp/92/f643fd4a06a9b4b57858-n0-0/5dc8c51e8ad277b3829be47547287bbf/example.txt",
                    file_format="txt"),
                fs=FlyteSchema(remote_path="s3://my-bucket/my-path"),
            ),
            list_flyte_types=[FlyteTypesWithMashumaro(
                flytefile=FlyteFile(
                    path="gs://opta-gcp-dogfood-gcp/92/f643fd4a06a9b4b57858-n0-0/5dc8c51e8ad277b3829be47547287bbf/example.txt"),
                flytedir=FlyteDirectory(
                    path="gs://opta-gcp-dogfood-gcp/92/f643fd4a06a9b4b57858-n0-0/5dc8c51e8ad277b3829be47547287bbf"),
                structured_dataset=StructuredDataset(
                    uri="gs://opta-gcp-dogfood-gcp/92/f643fd4a06a9b4b57858-n0-0/5dc8c51e8ad277b3829be47547287bbf/example.txt",
                    file_format="txt"),
                fs=FlyteSchema(remote_path="s3://my-bucket/my-path"),
            )],
            dict_flyte_types={"a": FlyteTypesWithMashumaro(
                flytefile=FlyteFile(
                    path="gs://opta-gcp-dogfood-gcp/92/f643fd4a06a9b4b57858-n0-0/5dc8c51e8ad277b3829be47547287bbf/example.txt"),
                flytedir=FlyteDirectory(
                    path="gs://opta-gcp-dogfood-gcp/92/f643fd4a06a9b4b57858-n0-0/5dc8c51e8ad277b3829be47547287bbf"),
                structured_dataset=StructuredDataset(
                    uri="gs://opta-gcp-dogfood-gcp/92/f643fd4a06a9b4b57858-n0-0/5dc8c51e8ad277b3829be47547287bbf/example.txt",
                    file_format="txt"),
                fs=FlyteSchema(remote_path="s3://my-bucket/my-path"),
            )},
        )):
    a1,b1,c1 = pass_and_return_flyte_types(a,b,c)
    a2,b2,c2 = pass_and_return_flyte_types(a1,b1,c1).with_overrides(container_image="ghcr.io/flyteorg/flytekit:py3.9-1.13.0")
    a3,b3,c3 = pass_and_return_flyte_types(a2,b2,c2)
    # with 1.9.0 flytekit
    a4, b4, c4 = pass_and_return_flyte_types(a3,b3,c3).with_overrides(container_image="futureoutlier/flytekit:iV4YIwLdioU0jkKL11SUIw")
    pass_and_return_flyte_types(a4, b4, c4)

if __name__ == "__main__":
    from flytekit.core.type_engine import TypeEngine
    a = TypeEngine.to_literal_type(NestedFlyteTypes)
    b = TypeEngine.to_literal_type(NestedFlyteTypesWithDataClassJson)
    c = TypeEngine.to_literal_type(NestedFlyteTypesWithMashumaro)
    print(a == b)

    # print(TypeEngine.to_literal_type(NestedFlyteTypes))
    # print(TypeEngine.to_literal_type(NestedFlyteTypesWithDataClassJson))
    # print(TypeEngine.to_literal_type(NestedFlyteTypesWithMashumaro))
    wf()


"""
1. flytefile -> flytefile
2. flytedir -> flytedir
3. structured dataset -> structured dataset
4. nested cases
"""

# if __name__ == "__main__":
"""
def wf(local_file: FlyteFile = None, local_dir: FlyteDirectory = None, remote_file: FlyteFile = None,
   remote_dir: FlyteDirectory = None, sd: StructuredDataset = None):
"""
# wf(local_file=create_local_file(), local_dir=create_dir(),
#    remote_file=create_remote_file(), remote_dir=create_remote_dir(),
#    sd=generate_sd())

Screenshots

sandbox

demo hosted

https://demo.hosted.unionai.cloud/console/projects/flytesnacks/domains/development/executions/f806fea899824435da55/nodes

dogfood

https://dogfood-gcp.cloud-staging.union.ai/console/projects/flytesnacks/domains/development/executions/f8e1b629a560c441ebdb/nodes

local execution

Check all the applicable boxes

• I updated the documentation accordingly.
• All new and existing tests passed.
• All commits are signed-off.

Related PRs

Docs link