Pydantic Transformer V2 #2792

Future-Outlier · 2024-10-08T10:08:11Z

Tracking issue

How to test it by others

git clone https://github.com/flyteorg/flytekit
gh pr checkout 2792
make setup-global-uv
cd plugins/flytekit-pydantic-v2 && pip install -e .
test a workflow example

Not Sure

Which pydantic version should we use as the lower bound?

This case will fail in the Flyte Console

@dataclass
class DC:
    a: Union[bool, str, int]
    b: Union[bool, str, int]

@task(container_image=image)
def add(dc1: DC, dc2: DC) -> Union[bool, int, str]:
    return dc1.a + dc2.b  # type: ignore

# input from flyte console to generate generic protobuf struct
# "{\"a\": 1, \"b\": 2}",
@workflow
def wf(dc: DC) -> Union[bool, int, str]:
    return add(dc1=dc, dc2=dc)

file tree structure

The file tree structure is the same as flytekit-pydantic

Why didn't integrate with pydantic v1 BaseModel? (make you run v1 and v2 BaseModel at the same time together)

This is an issue from pydantic.
pydantic/pydantic#9919
If this is fixed, then we can support both pydantic v1 and v2 at the same time.

story:
Kevin and I wanted to support v1 and v2 at the same time before, but after knowing that this would take lots of time, we asked Ketan for advice, then he said that if users want it, then we can try to support it or tell users to support it.

Why are the changes needed?

why from_generic_idl?
flyteconsole input to handle flytetypes.

when handling the input below, and attribute access to a flyte type, we need to teach flyte types how to convert a protobuf struct to flyte types.
Take FlyteFile as an example.

lifecycle

json str -> protobuf struct -> attribute access flyte type -> send to downstream input

class DC(BaseModel):
    ff: FlyteFile = Field(default_factory=lambda: FlyteFile("s3://my-s3-bucket/example.txt"))

@workflow
def wf(dc: DC) -> DC:
    t_ff(dc.ff)
    return t_args(dc=dc)

# console input: {"ff":{"path":"s3://my-s3-bucket/example.txt"}}

why _check_and_covert_int in the int transformer?
flyteconsole input to handle float issue.
It will be needed when in the following example.

json str -> protbuf struct -> attribute access and get float(due to javascript problem) -> convert float to int in flytekit

class DC(BaseModel):
    a: int = -1

@workflow
def wf(dc: DC):
    t_int(input_int=dc.a)

why basemodel -> json str -> dict obj -> msgpack bytes?
For enum class.

I've tried basemodel -> dict obj -> msgpack bytes first.
To make this happen, you need to call the function BaseModel.model_dump, but this function can't interpret Enum.
However, BaseModel.model_dump_json can.

What are @model_serializer and @model_validator(mode="after")?
You can understand them as _serialize and _deserialize in FlyteTypes, which use SerializableType to customize the serialize/deserialize behavior for flyte types.

Related PRs: #2554

What changes were proposed in this pull request?

attribute access (primitives and flyte types) (datetime not sure)
flyte types
nested cases
dataclasses.dataclass in pydantic.BaseModel
pydantic.dataclass in pydantic.BaseModel
pydantic.BaseModel in pydantic.BaseModel

note: we don't support pydantic BaseModel has a dataclass with FlyteTypes.
We support pydantic BaseModel has a dataclass with primitive types.

@dataclass
class dc:
    ff: FlyteFile

class DC(BaseModel):
    inner_dc: dc

# This is not supported
# ============================
@dataclass
class dc:
    a: int

class DC(BaseModel):
    inner_dc: dc

# This is supported
# ============================

How was this patch tested?

Example code.
(nested cases, flyte types, and attribute access.)

from pydantic import BaseModel, Field
from typing import Dict, List, Optional

from flytekit.types.schema import FlyteSchema
from flytekit.types.structured import StructuredDataset
from flytekit.types.file import FlyteFile
from flytekit.types.directory import FlyteDirectory
from flytekit import task, workflow, ImageSpec, kwtypes
from enum import Enum
import os
import pandas as pd

flytekit_hash = "fb82dd521615039f626c78489b2e83259d7db2a5"
flytekit = f"git+https://github.com/flyteorg/flytekit.git@{flytekit_hash}"
pydantic_plugin = f"git+https://github.com/flyteorg/flytekit.git@{flytekit_hash}#subdirectory=plugins/flytekit-pydantic-v2"

# Define custom image for the task
image = ImageSpec(packages=[
                            flytekit,
                            pydantic_plugin,
                            "pandas",
                            "pyarrow"],
                            apt_packages=["git"],
                            registry="localhost:30000",
                         )

class Status(Enum):
    PENDING = "pending"
    APPROVED = "approved"
    REJECTED = "rejected"

class InnerBM(BaseModel):
    a: int = -1
    b: float = 2.1
    c: str = "Hello, Flyte"
    d: bool = False
    e: List[int] = Field(default_factory=lambda: [0, 1, 2, -1, -2])
    f: List[FlyteFile] = Field(default_factory=lambda: [FlyteFile("s3://my-s3-bucket/example.txt")])
    g: List[List[int]] = Field(default_factory=lambda: [[0], [1], [-1]])
    h: List[Dict[int, bool]] = Field(default_factory=lambda: [{0: False}, {1: True}, {-1: True}])
    i: Dict[int, bool] = Field(default_factory=lambda: {0: False, 1: True, -1: False})
    j: Dict[int, FlyteFile] = Field(default_factory=lambda: {0: FlyteFile("s3://my-s3-bucket/example.txt"),
                                                             1: FlyteFile("s3://my-s3-bucket/example.txt"),
                                                             -1: FlyteFile("s3://my-s3-bucket/example.txt")})
    k: Dict[int, List[int]] = Field(default_factory=lambda: {0: [0, 1, -1]})
    l: Dict[int, Dict[int, int]] = Field(default_factory=lambda: {1: {-1: 0}})
    m: dict = Field(default_factory=lambda: {"key": "value"})
    n: FlyteFile = Field(default_factory=lambda: FlyteFile("s3://my-s3-bucket/example.txt"))
    o: FlyteDirectory = Field(default_factory=lambda: FlyteDirectory("s3://my-s3-bucket/s3_flyte_dir"))
    enum_status: Status = Status.PENDING
    sd: StructuredDataset = Field(default_factory=lambda: StructuredDataset(
        uri="s3://my-s3-bucket/data/uk/ahlg7qw7q5m4np7vwdqm-n0-0/7f31035fdf92510e40ee9340f9e5bf34",
        file_format="parquet"))
    fsc: FlyteSchema = Field(default_factory=lambda: FlyteSchema(
        remote_path="s3://my-s3-bucket/data/uk/ahlg7qw7q5m4np7vwdqm-n0-0/ab3aef21302d0529daef8c43825c3fdf"))

class BM(BaseModel):
    a: int = -1
    b: float = 2.1
    c: str = "Hello, Flyte"
    d: bool = False
    e: List[int] = Field(default_factory=lambda: [0, 1, 2, -1, -2])
    f: List[FlyteFile] = Field(default_factory=lambda: [FlyteFile("s3://my-s3-bucket/example.txt")])
    g: List[List[int]] = Field(default_factory=lambda: [[0], [1], [-1]])
    h: List[Dict[int, bool]] = Field(default_factory=lambda: [{0: False}, {1: True}, {-1: True}])
    i: Dict[int, bool] = Field(default_factory=lambda: {0: False, 1: True, -1: False})
    j: Dict[int, FlyteFile] = Field(default_factory=lambda: {0: FlyteFile("s3://my-s3-bucket/example.txt"),
                                                             1: FlyteFile("s3://my-s3-bucket/example.txt"),
                                                             -1: FlyteFile("s3://my-s3-bucket/example.txt")})
    k: Dict[int, List[int]] = Field(default_factory=lambda: {0: [0, 1, -1]})
    l: Dict[int, Dict[int, int]] = Field(default_factory=lambda: {1: {-1: 0}})
    m: dict = Field(default_factory=lambda: {"key": "value"})
    n: FlyteFile = Field(default_factory=lambda: FlyteFile("s3://my-s3-bucket/example.txt"))
    o: FlyteDirectory = Field(default_factory=lambda: FlyteDirectory("s3://my-s3-bucket/s3_flyte_dir"))
    inner_bm: InnerBM = Field(default_factory=lambda: InnerBM())
    enum_status: Status = Status.PENDING
    sd: StructuredDataset = Field(default_factory=lambda: StructuredDataset(
        uri="s3://my-s3-bucket/data/uk/ahlg7qw7q5m4np7vwdqm-n0-0/7f31035fdf92510e40ee9340f9e5bf34",
        file_format="parquet"))
    fsc: FlyteSchema = Field(default_factory=lambda: FlyteSchema(remote_path="s3://my-s3-bucket/data/uk/ahlg7qw7q5m4np7vwdqm-n0-0/ab3aef21302d0529daef8c43825c3fdf"))

@task(container_image=image)
def t_bm(bm: BM) -> BM:
    return bm

@task(container_image=image)
def t_inner(inner_bm: InnerBM):
    assert isinstance(inner_bm, InnerBM)

    expected_file_content = "Default content"

    # f: List[FlyteFile]
    for ff in inner_bm.f:
        assert isinstance(ff, FlyteFile)
        with open(ff, "r") as f:
            assert f.read() == expected_file_content
    # j: Dict[int, FlyteFile]
    for _, ff in inner_bm.j.items():
        assert isinstance(ff, FlyteFile)
        with open(ff, "r") as f:
            assert f.read() == expected_file_content
    # n: FlyteFile
    assert isinstance(inner_bm.n, FlyteFile)
    with open(inner_bm.n, "r") as f:
        assert f.read() == expected_file_content
    # o: FlyteDirectory
    assert isinstance(inner_bm.o, FlyteDirectory)
    assert not inner_bm.o.downloaded
    with open(os.path.join(inner_bm.o, "example.txt"), "r") as fh:
        assert fh.read() == expected_file_content
    assert inner_bm.o.downloaded
    print("Test InnerBM Successfully Passed")
    # enum: Status
    assert inner_bm.enum_status == Status.PENDING


@task(container_image=image)
def t_test_all_attributes(a: int, b: float, c: str, d: bool, e: List[int], f: List[FlyteFile], g: List[List[int]],
                          h: List[Dict[int, bool]], i: Dict[int, bool], j: Dict[int, FlyteFile],
                          k: Dict[int, List[int]], l: Dict[int, Dict[int, int]], m: dict,
                          n: FlyteFile, o: FlyteDirectory,
                          enum_status: Status,
                          sd: StructuredDataset,
                          fsc: FlyteSchema,
                          ):
    # Strict type checks for simple types
    assert isinstance(a, int), f"a is not int, it's {type(a)}"
    assert a == -1
    assert isinstance(b, float), f"b is not float, it's {type(b)}"
    assert isinstance(c, str), f"c is not str, it's {type(c)}"
    assert isinstance(d, bool), f"d is not bool, it's {type(d)}"

    # Strict type checks for List[int]
    assert isinstance(e, list) and all(isinstance(i, int) for i in e), "e is not List[int]"

    # Strict type checks for List[FlyteFile]
    assert isinstance(f, list) and all(isinstance(i, FlyteFile) for i in f), "f is not List[FlyteFile]"

    # Strict type checks for List[List[int]]
    assert isinstance(g, list) and all(
        isinstance(i, list) and all(isinstance(j, int) for j in i) for i in g), "g is not List[List[int]]"

    # Strict type checks for List[Dict[int, bool]]
    assert isinstance(h, list) and all(
        isinstance(i, dict) and all(isinstance(k, int) and isinstance(v, bool) for k, v in i.items()) for i in h
    ), "h is not List[Dict[int, bool]]"

    # Strict type checks for Dict[int, bool]
    assert isinstance(i, dict) and all(
        isinstance(k, int) and isinstance(v, bool) for k, v in i.items()), "i is not Dict[int, bool]"

    # Strict type checks for Dict[int, FlyteFile]
    assert isinstance(j, dict) and all(
        isinstance(k, int) and isinstance(v, FlyteFile) for k, v in j.items()), "j is not Dict[int, FlyteFile]"

    # Strict type checks for Dict[int, List[int]]
    assert isinstance(k, dict) and all(
        isinstance(k, int) and isinstance(v, list) and all(isinstance(i, int) for i in v) for k, v in
        k.items()), "k is not Dict[int, List[int]]"

    # Strict type checks for Dict[int, Dict[int, int]]
    assert isinstance(l, dict) and all(
        isinstance(k, int) and isinstance(v, dict) and all(
            isinstance(sub_k, int) and isinstance(sub_v, int) for sub_k, sub_v in v.items())
        for k, v in l.items()), "l is not Dict[int, Dict[int, int]]"

    # Strict type check for a generic dict
    assert isinstance(m, dict), "m is not dict"

    # Strict type check for FlyteFile
    assert isinstance(n, FlyteFile), "n is not FlyteFile"

    # Strict type check for FlyteDirectory
    assert isinstance(o, FlyteDirectory), "o is not FlyteDirectory"

    # # Strict type check for Enum
    assert isinstance(enum_status, Status), "enum_status is not Status"

    assert isinstance(sd, StructuredDataset), "sd is not StructuredDataset"
    print("sd:", sd.open(pd.DataFrame).all())

    assert isinstance(fsc, FlyteSchema), "fsc is not FlyteSchema"
    print("fsc: ", fsc.open().all())

    print("All attributes passed strict type checks.")


@workflow
def wf(bm: BM):
    t_bm(bm=bm)
    t_inner(inner_bm=bm.inner_bm)
    t_test_all_attributes(a=bm.a, b=bm.b, c=bm.c,
                          d=bm.d, e=bm.e, f=bm.f,
                          g=bm.g, h=bm.h, i=bm.i,
                          j=bm.j, k=bm.k, l=bm.l,
                          m=bm.m, n=bm.n, o=bm.o,
                          enum_status=bm.enum_status,
                          sd=bm.sd,
                          fsc=bm.fsc,
                          )

    t_test_all_attributes(a=bm.inner_bm.a, b=bm.inner_bm.b, c=bm.inner_bm.c,
                          d=bm.inner_bm.d, e=bm.inner_bm.e, f=bm.inner_bm.f,
                          g=bm.inner_bm.g, h=bm.inner_bm.h, i=bm.inner_bm.i,
                          j=bm.inner_bm.j, k=bm.inner_bm.k, l=bm.inner_bm.l,
                          m=bm.inner_bm.m, n=bm.inner_bm.n, o=bm.inner_bm.o,
                          enum_status=bm.inner_bm.enum_status,
                          sd=bm.inner_bm.sd,
                          fsc=bm.inner_bm.fsc,
                          )

if __name__ == "__main__":
    from flytekit.clis.sdk_in_container import pyflyte
    from click.testing import CliRunner

    runner = CliRunner()
    path = os.path.realpath(__file__)
    input_val = BM().model_dump_json()
    print(input_val)
    result = runner.invoke(pyflyte.main,
                           ["run", path, "wf", "--bm", input_val])
    print("Local Execution: ", result.output)

    result = runner.invoke(pyflyte.main,
                           ["run", "--remote", path, "wf", "--bm", input_val])
    print("Remote Execution: ", result.output)

Setup process

local and remote execution.
ImageSpec for the docker image.

Screenshots

local execution

remote execution

remote execution from flyte console input

Check all the applicable boxes

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

Signed-off-by: Future-Outlier <[email protected]>

Future-Outlier · 2024-10-08T15:34:34Z

flytekit/core/type_engine.py

+    if lv.scalar.primitive.float_value is not None:
+        logger.info(f"Converting literal float {lv.scalar.primitive.float_value} to int, might have precision loss.")
+        return int(lv.scalar.primitive.float_value)


This is for cases when you input from the flyte console, and you use attribute access directly, you have to convert the float to int.
Since javascript has only number, it can't tell the difference between int and float, and when goland (propeller) doing attribute access, it doesn't have the expected python type

class TrainConfig(BaseModel): lr: float = 1e-3 batch_size: int = 32 @workflow def wf(cfg: TrainConfig) -> TrainConfig: return t_args(a=cfg.lr, batch_size=cfg.batch_size)

the javascript issue and the attribute access issue are orthogonal right?

this should only be a javascript problem. attribute access should work since msgpack preserves float/int even in attribute access correct?

YES, the attribute access works well, it's because javascript pass float to golang, and golang pass float to python.

this should only be a javascript problem. attribute access should work since msgpack preserves float/int even in attribute access correct?

Yes, but when you are accessing a simple type, you have to change the behavior of SimpleTransformer.

For Pydantic Transformer, we will use strict=False as argument to convert it to right type.

def from_binary_idl(self, binary_idl_object: Binary, expected_python_type: Type[BaseModel]) -> BaseModel: if binary_idl_object.tag == MESSAGEPACK: dict_obj = msgpack.loads(binary_idl_object.value) python_val = expected_python_type.model_validate(obj=dict_obj, strict=False) return python_val

So we can delete this part after console is updated right?

If we can guarantee the console can generate an integer but not float from the input, then we can delete it.

how is this going to work though? Do we also do a version check of the backend?

After console does the right thing, won't this value be coming in through the binary value instead? Instead of lv.scalar.primitive.integer/float.

Future-Outlier · 2024-10-08T16:03:08Z

@lukas503
Hi, I saw you add an emoji to this PR!
Do you want to help me test this out?
Search "How to test it by others?" will have a guide for you!

lukas503 · 2024-10-08T21:43:14Z

Hi @Future-Outlier,

Thanks for working on the Pydantic TypeTransformer! Which "How to test it by others?" guide are you referring to?

I've been testing the code locally and wondered about a behavior related to caching. Specifically, I’m curious if model_json_schema is considered in the hash used for caching.

Here’s an example:

from flytekit import task, workflow
from pydantic import BaseModel

class Config(BaseModel):
    x: int = 1
    # y: int = 4

@task(cache=True, cache_version="v1")
def task1(val: int) -> Config:
    return Config()

@task(cache=True, cache_version="v1")
def task2(cfg: Config) -> Config:
    print("CALLED!", cfg)
    return cfg

@workflow
def my_workflow():
    config = task1(val=5)
    task2(cfg=config)

if __name__ == "__main__":
    print(Config.model_json_schema())
    my_workflow()

When I run the workflow for the first time, nothing is cached. On the second run, the results are cached, as expected. However, if I uncomment y: int = 4, the tasks still remain cached. I would assume that this schema change would trigger a cache bust and re-execute the tasks. This causes failure if I update the attributes and the cache_version of task2.

Is this the expected behavior? Shouldn't schema changes like this invalidate the cache?

Future-Outlier · 2024-10-09T01:05:57Z

Hi @Future-Outlier,

Thanks for working on the Pydantic TypeTransformer! Which "How to test it by others?" guide are you referring to?

I've been testing the code locally and wondered about a behavior related to caching. Specifically, I’m curious if model_json_schema is considered in the hash used for caching.

Here’s an example:
from flytekit import task, workflow
from pydantic import BaseModel

class Config(BaseModel):
    x: int = 1
    # y: int = 4

@task(cache=True, cache_version="v1")
def task1(val: int) -> Config:
    return Config()

@task(cache=True, cache_version="v1")
def task2(cfg: Config) -> Config:
    print("CALLED!", cfg)
    return cfg

@workflow
def my_workflow():
    config = task1(val=5)
    task2(cfg=config)

if __name__ == "__main__":
    print(Config.model_json_schema())
    my_workflow()
When I run the workflow for the first time, nothing is cached. On the second run, the results are cached, as expected. However, if I uncomment y: int = 4, the tasks still remain cached. I would assume that this schema change would trigger a cache bust and re-execute the tasks. This causes failure if I update the attributes and the cache_version of task2.

Is this the expected behavior? Shouldn't schema changes like this invalidate the cache?

good question, will test this out and ask other maintainers if I don't know what happened, thank you <3

Future-Outlier · 2024-10-09T03:34:16Z

Hi @Future-Outlier,

Thanks for working on the Pydantic TypeTransformer! Which "How to test it by others?" guide are you referring to?

I've been testing the code locally and wondered about a behavior related to caching. Specifically, I’m curious if model_json_schema is considered in the hash used for caching.

Here’s an example:
from flytekit import task, workflow
from pydantic import BaseModel

class Config(BaseModel):
    x: int = 1
    # y: int = 4

@task(cache=True, cache_version="v1")
def task1(val: int) -> Config:
    return Config()

@task(cache=True, cache_version="v1")
def task2(cfg: Config) -> Config:
    print("CALLED!", cfg)
    return cfg

@workflow
def my_workflow():
    config = task1(val=5)
    task2(cfg=config)

if __name__ == "__main__":
    print(Config.model_json_schema())
    my_workflow()
When I run the workflow for the first time, nothing is cached. On the second run, the results are cached, as expected. However, if I uncomment y: int = 4, the tasks still remain cached. I would assume that this schema change would trigger a cache bust and re-execute the tasks. This causes failure if I update the attributes and the cache_version of task2.

Is this the expected behavior? Shouldn't schema changes like this invalidate the cache?

@lukas503
sorry can you try again?
I've updated the above description.

Signed-off-by: Future-Outlier <[email protected]>

codecov · 2024-10-09T04:15:00Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 91.21%. Comparing base (8017ae8) to head (6993f09).
Report is 8 commits behind head on master.

Additional details and impacted files

@@             Coverage Diff             @@
##           master    #2792       +/-   ##
===========================================
+ Coverage   45.55%   91.21%   +45.65%     
===========================================
  Files         196       93      -103     
  Lines       20434     4620    -15814     
  Branches     2633        0     -2633     
===========================================
- Hits         9309     4214     -5095     
+ Misses      10676      406    -10270     
+ Partials      449        0      -449

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

Signed-off-by: Future-Outlier <[email protected]>

Future-Outlier · 2024-10-09T07:29:19Z

flytekit/types/file/file.py

+        if lv.scalar:
+            if lv.scalar.binary:
+                return self.from_binary_idl(lv.scalar.binary, expected_python_type)
+            if lv.scalar.generic:
+                return self.from_generic_idl(lv.scalar.generic, expected_python_type)


class DC(BaseModel): ff: FlyteFile = Field(default_factory=lambda: FlyteFile("s3://my-s3-bucket/example.txt")) @task(container_image=image) def t_args(dc: DC) -> DC: with open(dc.ff, "r") as f: print(f.read()) return dc @task(container_image=image) def t_ff(ff: FlyteFile) -> FlyteFile: with open(ff, "r") as f: print(f.read()) return ff @workflow def wf(dc: DC) -> DC: t_ff(dc.ff) return t_args(dc=dc)

this is for this case input from flyteconsole.

lukas503 · 2024-10-09T09:18:39Z

sorry can you try again?
I've updated the above description.

Thanks for updating the PR. I now understand the underlying issue better. It appears the caching mechanism is ignoring the output types/schema. What’s unclear to me is why the output types/schema aren’t factored into the hash used for caching. In my opinion, any interface change could invalidate the cache even the outputs. I don’t see how the old cached outputs can remain valid after an interface change.

That said, this concern isn’t directly related to the current PR, so feel free to proceed as is.

Update: It works as expected if remote flyte is used. The faulty behavior I described is happening only locally.

Signed-off-by: Future-Outlier <[email protected]>

Future-Outlier

Just do a final test and it works

eapolinario

Just a few minor things, otherwise, it's looking pretty good.

eapolinario · 2024-10-24T13:40:00Z

flytekit/extras/pydantic/__init__.py

+    logger.info(f"Meet error when importing pydantic: `{e}`")
+    logger.info("Flytekit only support pydantic version > 2.")


nit: those should be a warning.

eapolinario · 2024-10-24T13:40:31Z

flytekit/extras/pydantic/decorator.py

+    from pydantic import model_serializer, model_validator
+
+except ImportError:
+    logger.info(


eapolinario · 2024-10-24T13:41:59Z

flytekit/extras/pydantic/decorator.py

+    FuncType = TypeVar("FuncType", bound=Callable[..., Any])
+
+    from typing_extensions import Literal as typing_literal
+
+    def model_serializer(
+        __f: Union[Callable[..., Any], None] = None,
+        *,
+        mode: typing_literal["plain", "wrap"] = "plain",
+        when_used: typing_literal["always", "unless-none", "json", "json-unless-none"] = "always",
+        return_type: Any = None,
+    ) -> Callable[[Any], Any]:
+        """Placeholder decorator for Pydantic model_serializer."""
+
+        def decorator(fn: Callable[..., Any]) -> Callable[..., Any]:
+            def wrapper(*args, **kwargs):
+                raise Exception(
+                    "Pydantic is not installed.\n" "Please install Pydantic version > 2 to use this feature."
+                )
+
+            return wrapper
+
+        # If no function (__f) is provided, return the decorator
+        if __f is None:
+            return decorator
+        # If __f is provided, directly decorate the function
+        return decorator(__f)
+
+    def model_validator(
+        *,
+        mode: typing_literal["wrap", "before", "after"],
+    ) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+        """Placeholder decorator for Pydantic model_validator."""
+
+        def decorator(fn: Callable[..., Any]) -> Callable[..., Any]:
+            def wrapper(*args, **kwargs):
+                raise Exception(
+                    "Pydantic is not installed.\n" "Please install Pydantic version > 2 to use this feature."
+                )
+
+            return wrapper
+
+        return decorator


Aren't we supporting only pydantic v2? Why do we have these fallback definitions?

we want to make here to work.
This syntax is more readable then setattr
https://github.com/flyteorg/flytekit/pull/2792/files#diff-22cf9c7153b54371b4a77331ddf276a082cf4b3c5e7bd1595dd67232288594fdR168-R176

it's to support the case where pydantic is not installed at all. cuz yeah it looks nicer in the real File/Directory class, but we also want it to not fail ofc.

eapolinario · 2024-10-25T00:10:40Z

.github/workflows/pythonbuild.yml

+          # TODO: remove pydantic v1 plugin, since v2 is in core already
+          # flytekit-pydantic


Can we remove it?

eapolinario · 2024-10-25T00:10:57Z

flytekit/core/type_engine.py

@@ -215,16 +215,32 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type:
        )

    def from_binary_idl(self, binary_idl_object: Binary, expected_python_type: Type[T]) -> Optional[T]:
+        """
+        TODO: Add more comments to explain the lifecycle of attribute access.


fill in TODO

eapolinario · 2024-10-25T00:25:53Z

flytekit/core/type_engine.py

+    if lv.scalar.primitive.float_value is not None:
+        logger.info(f"Converting literal float {lv.scalar.primitive.float_value} to int, might have precision loss.")
+        return int(lv.scalar.primitive.float_value)


how is this going to work though? Do we also do a version check of the backend?

eapolinario · 2024-10-25T00:33:57Z

plugins/flytekit-pydantic/setup.py

@@ -4,7 +4,7 @@

 microlib_name = f"flytekitplugins-{PLUGIN_NAME}"

-plugin_requires = ["flytekit>=1.7.0b0", "pydantic"]
+plugin_requires = ["flytekit>=1.7.0b0", "pydantic<2"]


Can we also leave a warning in the README.md explaining that we're deprecating this plugin?

@eapolinario

how is this going to work though? Do we also do a version check of the backend?
No this is just for supporting the case I've mentioned above.
we didn't support this before and I think we should do it.

Signed-off-by: Future-Outlier <[email protected]>

Future-Outlier · 2024-10-25T07:53:19Z

Let's merge it.
cc @eapolinario @wild-endeavor @pingsutw

wild-endeavor

🙇 thank you @Future-Outlier 🙏 this is going to be fantastic.

wild-endeavor · 2024-10-25T19:51:03Z

flytekit/core/type_engine.py

@@ -1124,6 +1194,8 @@ def lazy_import_transformers(cls):
                from flytekit.extras import pytorch  # noqa: F401
            if is_imported("sklearn"):
                from flytekit.extras import sklearn  # noqa: F401
+            if is_imported("pydantic"):
+                from flytekit.extras import pydantic  # noqa: F401


can we change the name of this folder? pydantic can get confusing because the real library is also called pydantic right?

wild-endeavor · 2024-10-25T20:01:45Z

flytekit/core/type_engine.py

@@ -2194,6 +2304,34 @@ def _check_and_covert_float(lv: Literal) -> float:
    raise TypeTransformerFailedError(f"Cannot convert literal {lv} to float")


+def _handle_flyte_console_float_input_to_int(lv: Literal) -> int:
+    """
+    Flyte Console is written by JavaScript and JavaScript has only one number type which is float.


technically javascript's number type is Number but yeah, sometimes it keeps track of trailing 0s and sometimes it doesn't.

thank you, it's important to make words more accurately

wild-endeavor · 2024-10-25T20:05:43Z

flytekit/core/type_engine.py

+    if lv.scalar.primitive.float_value is not None:
+        logger.info(f"Converting literal float {lv.scalar.primitive.float_value} to int, might have precision loss.")
+        return int(lv.scalar.primitive.float_value)


After console does the right thing, won't this value be coming in through the binary value instead? Instead of lv.scalar.primitive.integer/float.

wild-endeavor · 2024-10-25T20:08:47Z

flytekit/extras/pydantic/decorator.py

+    FuncType = TypeVar("FuncType", bound=Callable[..., Any])
+
+    from typing_extensions import Literal as typing_literal
+
+    def model_serializer(
+        __f: Union[Callable[..., Any], None] = None,
+        *,
+        mode: typing_literal["plain", "wrap"] = "plain",
+        when_used: typing_literal["always", "unless-none", "json", "json-unless-none"] = "always",
+        return_type: Any = None,
+    ) -> Callable[[Any], Any]:
+        """Placeholder decorator for Pydantic model_serializer."""
+
+        def decorator(fn: Callable[..., Any]) -> Callable[..., Any]:
+            def wrapper(*args, **kwargs):
+                raise Exception(
+                    "Pydantic is not installed.\n" "Please install Pydantic version > 2 to use this feature."
+                )
+
+            return wrapper
+
+        # If no function (__f) is provided, return the decorator
+        if __f is None:
+            return decorator
+        # If __f is provided, directly decorate the function
+        return decorator(__f)
+
+    def model_validator(
+        *,
+        mode: typing_literal["wrap", "before", "after"],
+    ) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+        """Placeholder decorator for Pydantic model_validator."""
+
+        def decorator(fn: Callable[..., Any]) -> Callable[..., Any]:
+            def wrapper(*args, **kwargs):
+                raise Exception(
+                    "Pydantic is not installed.\n" "Please install Pydantic version > 2 to use this feature."
+                )
+
+            return wrapper
+
+        return decorator


it's to support the case where pydantic is not installed at all. cuz yeah it looks nicer in the real File/Directory class, but we also want it to not fail ofc.

wild-endeavor · 2024-10-25T20:12:06Z

flytekit/extras/pydantic/transformer.py

+        super().__init__("Pydantic Transformer", BaseModel, enable_type_assertions=False)
+
+    def get_literal_type(self, t: Type[BaseModel]) -> LiteralType:
+        schema = t.model_json_schema()


In a future PR, can we add some unit tests to ensure that we're correctly extracting default values into the schema?

I'll ask more question about this unit tests PR if not totally understand, thank you

wild-endeavor · 2024-10-25T20:25:08Z

tests/flytekit/unit/extras/pydantic/test_pydantic_transformer.py

+
+    bm = BM()
+    wf(bm=bm)
+


can you add two new lines between tests? I know it doesn't matter, but pycharm complains.

wild-endeavor · 2024-10-25T20:26:07Z

tests/flytekit/unit/extras/pydantic/test_pydantic_transformer.py

+
+def test_flytetypes_in_pydantic_basemodel_wf(local_dummy_file, local_dummy_directory):
+    class InnerBM(BaseModel):
+        flytefile: FlyteFile = field(default_factory=lambda: FlyteFile(local_dummy_file))


little bit confused about pydantic here. Are you supposed to use dataclasses.field here instead of pydantic.Field?

both are supported, this is a good catch, should tell users field from dataclass and Field from pydantic are supported.

wild-endeavor · 2024-10-25T20:27:05Z

tests/flytekit/unit/types/structured_dataset/test_structured_dataset.py

 def test_protocol():
    assert get_protocol("s3://my-s3-bucket/file") == "s3"
    assert get_protocol("/file") == "file"


 def generate_pandas() -> pd.DataFrame:
    return pd.DataFrame({"name": ["Tom", "Joseph"], "age": [20, 22]})
-


keep spaces plz

wild-endeavor · 2024-10-25T20:28:27Z

tests/flytekit/unit/extras/pydantic/test_pydantic_transformer.py

+        flytefile: FlyteFile = field(default_factory=lambda: FlyteFile(local_dummy_file))
+        flytedir: FlyteDirectory = field(default_factory=lambda: FlyteDirectory(local_dummy_directory))
+
+    class BM(BaseModel):


maybe in a different file, but can we add pydantic models that contain dataclass and also dataclasses that contain Pydantic models? I know some people have been asking for that, be good to have some tests for it.

Thank you!

pydantic models that contain dataclass

this is supported

dataclasses that contain Pydantic models

this is never supported, if people really want it, we can do it in a seperate PR

I'll add unit tests today, thank you so much

Future-Outlier · 2024-10-28T03:43:09Z

cc @wild-endeavor

After console does the right thing, won't this value be coming in through the binary value instead? Instead of lv.scalar.primitive.integer/float.

if console can turn input to expected python type (int or float), then yes.

Signed-off-by: Future-Outlier <[email protected]> Co-authored-by: Yee Hing Tong <[email protected]>

Signed-off-by: Kevin Su <[email protected]>

Signed-off-by: Thomas J. Fan <[email protected]>

* set parent node metadata on arraynode subnode Signed-off-by: Paul Dittamo <[email protected]> * clean up setting passed in metadata for array node Signed-off-by: Paul Dittamo <[email protected]> * add comment Signed-off-by: Paul Dittamo <[email protected]> * test Signed-off-by: Paul Dittamo <[email protected]> --------- Signed-off-by: Paul Dittamo <[email protected]>

* Update versioning Signed-off-by: Mecoli1219 <[email protected]> * Fix unit test Signed-off-by: Mecoli1219 <[email protected]> * Update path if jupyter file actually exist Signed-off-by: Mecoli1219 <[email protected]> * Fix pickled twice problem Signed-off-by: Mecoli1219 <[email protected]> --------- Signed-off-by: Mecoli1219 <[email protected]>

Signed-off-by: Future-Outlier <[email protected]>

Future-Outlier · 2024-10-28T07:12:03Z

cc @wild-endeavor @eapolinario @pingsutw

Summary of Changes:

Linting: Applied linting to all test files in the final commit.
Folder Renaming: Changed the folder name from flytekit.extras.pydantic to flytekit.extras.pydantic_transformer.
https://github.com/flyteorg/flytekit/pull/2792/files#diff-68624e201f644896b015de1a5716eabef9b8b5a106caba1d9f21c51c0708684bR1197-R1198
JavaScript Number Handling: Updated comments related to JavaScript's number handling.
https://github.com/flyteorg/flytekit/pull/2792/files#diff-68624e201f644896b015de1a5716eabef9b8b5a106caba1d9f21c51c0708684bR2308-R2312
Tests Added:
- Added tests using both dataclass.field and pydantic.Field.
- Created unit tests for Pydantic models inside dataclasses (note: this will raise an error, and an issue will be created).
  https://github.com/flyteorg/flytekit/pull/2792/files#diff-779944519da7ce37ef29492b53edb53d653884263629341e4721c22d2223dedd
- Created unit tests for dataclasses inside Pydantic models (using field).
  https://github.com/flyteorg/flytekit/pull/2792/files#diff-23d316d571d3169735f378aa89e677ee2d30f0d2b6c9004d08341180aa0d5e3a
Decorator Comments: Added comments in the Pydantic transformer decorator for scenarios where Pydantic is not installed.
https://github.com/flyteorg/flytekit/pull/2792/files#diff-65db41859ac618afd17929975a6975810003ee243e2d88f5898ba49667524e27R12-R15

Notes:

I used autopep8 as the linter for some test_*.py files, which resulted in around 1000 lines of code changes (please ignore these changes).
I encountered an issue while rebasing with the master branch and failed to pull others' commits into this PR (my mistake). Consequently, I had to manually set the DCO to pass. Apologies for the inconvenience.

Signed-off-by: Future-Outlier <[email protected]>

wild-endeavor · 2024-10-28T18:32:52Z

if people really want it, we can do it in a seperate PR

yeah, people really want it. let's do it in a separate PR.

Pydantic Transformer V2

333a05c

Signed-off-by: Future-Outlier <[email protected]>

Future-Outlier requested review from wild-endeavor, kumare3, eapolinario, pingsutw, cosmicBboy, samhita-alla and thomasjpfan as code owners October 8, 2024 10:08

Future-Outlier changed the title ~~Pydantic Transformer V2~~ [wip] Pydantic Transformer V2 Oct 8, 2024

Future-Outlier added 3 commits October 8, 2024 18:09

add __init__.py

9251508

Signed-off-by: Future-Outlier <[email protected]>

add json schema

4c46dee

Signed-off-by: Future-Outlier <[email protected]>

convert float to int

63ab9fd

Signed-off-by: Future-Outlier <[email protected]>

Future-Outlier commented Oct 8, 2024

View reviewed changes

change gitsha in test script mode

357ca00

Signed-off-by: Future-Outlier <[email protected]>

Future-Outlier added 5 commits October 9, 2024 12:22

change gitsha

cdd1b25

Signed-off-by: Future-Outlier <[email protected]>

use strict map=false

c89b59a

Signed-off-by: Future-Outlier <[email protected]>

Test flytefile console input + attr access

6b480da

Signed-off-by: Future-Outlier <[email protected]>

add conditional branch

15ce9ad

Signed-off-by: Future-Outlier <[email protected]>

better rx

94ce092

Signed-off-by: Future-Outlier <[email protected]>

Future-Outlier commented Oct 9, 2024

View reviewed changes

Future-Outlier added 4 commits October 9, 2024 22:13

Add flytedir generic -> flytedir

d97bc2a

Signed-off-by: Future-Outlier <[email protected]>

merge async type engine

4b0f008

Signed-off-by: Future-Outlier <[email protected]>

support enum

afd4344

Signed-off-by: Future-Outlier <[email protected]>

update

773a3b6

Signed-off-by: Future-Outlier <[email protected]>

fix syntax to support python 3.9

76ae0ef

Signed-off-by: Future-Outlier <[email protected]>

Future-Outlier commented Oct 24, 2024

View reviewed changes

Future-Outlier enabled auto-merge (squash) October 24, 2024 05:09

Future-Outlier requested a review from wild-endeavor October 24, 2024 05:09

eapolinario reviewed Oct 25, 2024

View reviewed changes

Future-Outlier added 2 commits October 25, 2024 15:40

add Eduardo's advice

dfe8762

Signed-off-by: Future-Outlier <[email protected]>

warning

7735352

Signed-off-by: Future-Outlier <[email protected]>

Future-Outlier requested a review from eapolinario October 25, 2024 07:47

Merge branch 'master' into pydantic-plugin-v2-with-msgpack

959f02b

wild-endeavor reviewed Oct 25, 2024

View reviewed changes

Future-Outlier and others added 6 commits October 28, 2024 11:57

Update Yee's advice

85df643

Signed-off-by: Future-Outlier <[email protected]> Co-authored-by: Yee Hing Tong <[email protected]>

Show traceback by default (#2862)

f440dab

Signed-off-by: Kevin Su <[email protected]>

Support Identifier in generate_console_url (#2868)

4f5f74e

Signed-off-by: Thomas J. Fan <[email protected]>

improved output handling in notebooks (#2869)

2bccd1b

Future-Outlier disabled auto-merge October 28, 2024 05:37

Future-Outlier enabled auto-merge (squash) October 28, 2024 05:37

Future-Outlier added 2 commits October 28, 2024 14:46

tests by Yee's suggestion

bf283f6

Signed-off-by: Future-Outlier <[email protected]>

fix tests

6993f09

Signed-off-by: Future-Outlier <[email protected]>

Merge branch 'master' into pydantic-plugin-v2-with-msgpack

a3fef67

Future-Outlier mentioned this pull request Oct 28, 2024

[Core feature] Pydantic BaseModel V2 in Dataclass flyteorg/flyte#5925

Open

2 tasks

format tests

d983c07

Signed-off-by: Future-Outlier <[email protected]>

Future-Outlier requested a review from wild-endeavor October 28, 2024 15:11

pingsutw approved these changes Oct 28, 2024

View reviewed changes

Future-Outlier merged commit ff2d0da into master Oct 28, 2024
102 checks passed

		logger.info(f"Meet error when importing pydantic: `{e}`")
		logger.info("Flytekit only support pydantic version > 2.")

		# TODO: remove pydantic v1 plugin, since v2 is in core already
		# flytekit-pydantic

Pydantic Transformer V2 #2792

Pydantic Transformer V2 #2792

Conversation

Future-Outlier commented Oct 8, 2024 • edited Loading

Tracking issue

How to test it by others

Not Sure

This case will fail in the Flyte Console

file tree structure

Why didn't integrate with pydantic v1 BaseModel? (make you run v1 and v2 BaseModel at the same time together)

Why are the changes needed?

What changes were proposed in this pull request?

How was this patch tested?

Setup process

Screenshots

Check all the applicable boxes

Future-Outlier Oct 8, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Future-Outlier commented Oct 8, 2024

lukas503 commented Oct 8, 2024 • edited Loading

Future-Outlier commented Oct 9, 2024

Future-Outlier commented Oct 9, 2024

codecov bot commented Oct 9, 2024 • edited Loading

Codecov Report

Choose a reason for hiding this comment

lukas503 commented Oct 9, 2024 • edited Loading

Future-Outlier left a comment • edited Loading

Choose a reason for hiding this comment

eapolinario left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Future-Outlier commented Oct 25, 2024

wild-endeavor left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Future-Outlier Oct 28, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Future-Outlier commented Oct 28, 2024

Future-Outlier commented Oct 28, 2024

Summary of Changes:

Notes:

wild-endeavor commented Oct 28, 2024 • edited Loading

Future-Outlier commented Oct 8, 2024 •

edited

Loading

Future-Outlier Oct 8, 2024 •

edited

Loading

lukas503 commented Oct 8, 2024 •

edited

Loading

codecov bot commented Oct 9, 2024 •

edited

Loading

lukas503 commented Oct 9, 2024 •

edited

Loading

Future-Outlier left a comment •

edited

Loading

Future-Outlier Oct 28, 2024 •

edited

Loading

wild-endeavor commented Oct 28, 2024 •

edited

Loading