[Core feature] Allow dictionaries to be passed to execute workflows that take dataclasses

[jasonlai1218]

Tracking issue

flyteorg/flyte#4098

Docs link

Describe your changes

• Import the json module
• Add a check for dictionary types in the DataclassTransformer class
• Raise a TypeTransformerFailedError if a key in the dictionary does not match any field in the dataclass
• Convert the dictionary to a JSON string using json.dumps()

Check all the applicable boxes

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

Setup Process

1. write a data_class2.py

import typing
from dataclasses import dataclass
from dataclasses_json import dataclass_json
from flytekit import workflow, task


@dataclass_json
@dataclass
class Data:
    a: int = 10
    b: str = "default_b"
    c: typing.Optional[str] = "default_c"


@task
def concat(x: list[Data]):
    for data in x:
        print(f"{data.a}_{data.b}_{data.c}")


@workflow
def my_wf(x: list[Data]) -> int:
    concat(x=x)
    return 0

2. register data_class2.py
   execute cmd: pyflyte register -p flytesnacks --domain development --version 4 data_class2.py

3. write a remote_wf2.py

from flytekit.remote import FlyteRemote
from flytekit.configuration import Config


if __name__ == "__main__":
    # FlyteRemote object is the main entrypoint to API
    remote = FlyteRemote(
        config=Config.for_endpoint(endpoint="localhost:30080", insecure=True),
        default_project="flytesnacks",
        default_domain="development",
    )

    # Get Task
    flyte_workflow = remote.fetch_workflow(name="data_class2.my_wf", version="4")

    # Run Task
    execution = remote.execute(
        flyte_workflow,
        inputs={'x': [{"a": 10, "b": "hello"}, {"a": 100, "b": "hello2"}]}
    )

and execute the program

Screenshots

Successful screenshot

---

---

FlyteFile

1. write a flytefile_data_class.py

import csv
import os
from collections import defaultdict
from dataclasses import dataclass
from typing import List

from dataclasses_json import dataclass_json

import flytekit
from flytekit import task, workflow
from flytekit.types.file import FlyteFile


@dataclass_json
@dataclass
class Data:
    csv_url: FlyteFile
    columns_to_normalize: List[str]
    column_names: List[str]
    output_location: str


@task
def normalize_columns(data: Data) -> FlyteFile:
    # read the data from the raw csv file
    parsed_data = defaultdict(list)
    with open(data.csv_url, newline="\n") as input_file:
        reader = csv.DictReader(input_file, fieldnames=data.column_names)
        next(reader)  # Skip header
        for row in reader:
            for column in data.columns_to_normalize:
                parsed_data[column].append(float(row[column].strip()))

    # normalize the data
    normalized_data = defaultdict(list)
    for colname, values in parsed_data.items():
        mean = sum(values) / len(values)
        std = (sum([(x - mean) ** 2 for x in values]) / len(values)) ** 0.5
        normalized_data[colname] = [(x - mean) / std for x in values]

    # write to local path
    out_path = os.path.join(
        flytekit.current_context().working_directory,
        f"normalized-{os.path.basename(data.csv_url.path).rsplit('.')[0]}.csv",
    )
    with open(out_path, mode="w") as output_file:
        writer = csv.DictWriter(output_file, fieldnames=data.columns_to_normalize)
        writer.writeheader()
        for row in zip(*normalized_data.values()):
            writer.writerow({k: row[i] for i, k in enumerate(data.columns_to_normalize)})

    if data.output_location:
        return FlyteFile(path=out_path, remote_path=data.output_location)
    else:
        return FlyteFile(path=out_path)


@workflow
def normalize_csv_file(
    data: Data,
) -> FlyteFile:
    return normalize_columns(data=data)

2. register flytefile_data_class.py
   pyflyte register -p flytesnacks --domain development --version 1 flytefile_data_class.py

3. write a flytefile_remote_wf.py

from flytekit.remote import FlyteRemote
from flytekit.configuration import Config

if __name__ == "__main__":
    # FlyteRemote object is the main entrypoint to API
    remote = FlyteRemote(
        config=Config.for_endpoint(endpoint="localhost:30080", insecure=True),
        default_project="flytesnacks",
        default_domain="development",
    )

    # Get Task
    flyte_workflow = remote.fetch_workflow(name="flytefile_data_class.normalize_csv_file", version="1")

    # Run Task
    execution = remote.execute(
        flyte_workflow,
        inputs={
            'data': {
                "csv_url": {"path": "https://people.sc.fsu.edu/~jburkardt/data/csv/biostats.csv"},
                "column_names": ["Name", "Sex", "Age", "Heights (in)", "Weight (lbs)"],
                "columns_to_normalize": ["Age"],
                "output_location": ""
            },
        }
    )

and execute the program

Screenshots

Successful screenshot

---

StructuredDataset

1. write a structureddataset_data_class.py

from dataclasses_json import dataclass_json
from flytekit.types.structured import StructuredDataset

from flytekit import task, workflow, kwtypes
import pandas as pd


@dataclass_json
@dataclass
class Data:
    structured_dataset: StructuredDataset


col = kwtypes(Age=int)


@task
def get_subset_pandas_df(data: Data) -> StructuredDataset:
    df = data.structured_dataset.open(pd.DataFrame).all()
    df = pd.concat([df, pd.DataFrame([[30]], columns=["Age"])])
    return StructuredDataset(dataframe=df)


@task
def print_df(df: pd.DataFrame):
    print(df)


@task
def save_df(df: pd.DataFrame, path: str):
    df.to_csv(path, index=False)


@workflow
def simple_sd_wf(data: Data) -> StructuredDataset:
    # pandas_df = generate_pandas_df(a=19)
    output = get_subset_pandas_df(data=data)
    print_df(df=output.dataframe)
    save_df(df=output.dataframe, path="test3.csv")
    return output.dataframe

2. register flytefile_data_class.py
   pyflyte register -p flytesnacks --domain development --version 1 structureddataset_data_class.py

3. upload parquet file to minio
   

4. write a structureddataset_remote_wf.py

from flytekit.remote import FlyteRemote
from flytekit.configuration import Config


if __name__ == "__main__":
    # FlyteRemote object is the main entrypoint to API
    remote = FlyteRemote(
        config=Config.for_endpoint(endpoint="localhost:30080", insecure=True),
        default_project="flytesnacks",
        default_domain="development",
    )

    # Get Task
    flyte_workflow = remote.fetch_workflow(name="structureddataset_data_class.simple_sd_wf", version="1")

    # Run Task
    execution = remote.execute(
        flyte_workflow,
        inputs={
            'data': {
                "structured_dataset": {
                    "uri": "s3://test/part-00000-8cbb2a62-33cc-475a-8895-b06e75aa1db6-c000.snappy.parquet",
                    "file_format": "parquet",
                }
            },
        }
    )

and execute the program

Screenshots

Successful screenshot

---

If the user configures the settings incorrectly, a screenshot of the error should be displayed

Scenario 1: Parameter type error

---

error message:
TypeTransformerFailedError: Type of Val '<class 'int'>' is not an instance of 
<class 'str'>

---

Scenario 2: Wrong field name

error message:
TypeTransformerFailedError: The original fields are missing the following keys 
from the dataclass fields: ['b']

---

TypeTransformerFailedError: The original fields are missing the following keys 
from the dataclass fields: ['b']

---

TypeTransformerFailedError: The original fields have the following extra keys 
that are not in dataclass fields: ['d']

Note to reviewers

Related PRs

[codecov]

Codecov Report

Attention: 1 lines in your changes are missing coverage. Please review.

Comparison is base (ad328a7) 85.70% compared to head (e4dbac9) 84.31%.

Files	Patch %	Lines
flytekit/core/type_engine.py	97.29%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2013      +/-   ##
==========================================
- Coverage   85.70%   84.31%   -1.40%     
==========================================
  Files         313      313              
  Lines       23364    23389      +25     
  Branches     3489     3499      +10     
==========================================
- Hits        20025    19720     -305     
- Misses       2733     3089     +356     
+ Partials      606      580      -26     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[pingsutw]

Thanks @jasonlai1218. Could you resolve the merge conflict?

[jasonlai1218]

updated