Pydantic Transformer V2

flytekit/core/type_engine.py

@@ -2052,6 +2052,17 @@ def _check_and_covert_float(lv: Literal) -> float:
     raise TypeTransformerFailedError(f"Cannot convert literal {lv} to float")
 
 
+def _check_and_covert_int(lv: Literal) -> int:
+    if lv.scalar.primitive.integer is not None:
+        return lv.scalar.primitive.integer
+
+    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)
+
+    raise TypeTransformerFailedError(f"Cannot convert literal {lv} to int")
+
+
 def _check_and_convert_void(lv: Literal) -> None:
     if lv.scalar.none_type is None:
         raise TypeTransformerFailedError(f"Cannot convert literal {lv} to None")
@@ -2065,7 +2076,7 @@ def _register_default_type_transformers():
             int,
             _type_models.LiteralType(simple=_type_models.SimpleType.INTEGER),
             lambda x: Literal(scalar=Scalar(primitive=Primitive(integer=x))),
-            lambda x: x.scalar.primitive.integer,
+            _check_and_covert_int,
         )
     )
 

flytekit/interaction/click_types.py

@@ -36,7 +36,10 @@ def is_pydantic_basemodel(python_type: typing.Type) -> bool:
         return False
     else:
         try:
-            from pydantic.v1 import BaseModel
+            from pydantic import BaseModel as BaseModelV2
+            from pydantic.v1 import BaseModel as BaseModelV1
+
+            return issubclass(python_type, BaseModelV1) or issubclass(python_type, BaseModelV2)
         except ImportError:
             from pydantic import BaseModel
 

flytekit/types/directory/__init__.py

@@ -10,13 +10,14 @@
    :template: file_types.rst
 
    FlyteDirectory
+   FlyteDirToMultipartBlobTransformer
    TensorboardLogs
    TFRecordsDirectory
 """
 
 import typing
 
-from .types import FlyteDirectory
+from .types import FlyteDirectory, FlyteDirToMultipartBlobTransformer
 
 # The following section provides some predefined aliases for commonly used FlyteDirectory formats.
 

flytekit/types/file/__init__.py

@@ -10,6 +10,7 @@
    :template: file_types.rst
 
    FlyteFile
+   FlyteFilePathTransformer
    HDF5EncodedFile
    HTMLPage
    JoblibSerializedFile
@@ -25,7 +26,7 @@
 
 from typing_extensions import Annotated, get_args, get_origin
 
-from .file import FlyteFile
+from .file import FlyteFile, FlyteFilePathTransformer
 
 
 class FileExt:

flytekit/types/file/file.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import json
 import mimetypes
 import os
 import pathlib
@@ -11,6 +12,8 @@
 
 import msgpack
 from dataclasses_json import config
+from google.protobuf import json_format as _json_format
+from google.protobuf.struct_pb2 import Struct
 from marshmallow import fields
 from mashumaro.mixins.json import DataClassJSONMixin
 from mashumaro.types import SerializableType
@@ -549,12 +552,42 @@ def from_binary_idl(
         else:
             raise TypeTransformerFailedError(f"Unsupported binary format: `{binary_idl_object.tag}`")
 
+    def from_generic_idl(
+        self, generic: Struct, expected_python_type: typing.Union[typing.Type[FlyteFile], os.PathLike]
+    ):
+        json_str = _json_format.MessageToJson(generic)
+        python_val = json.loads(json_str)
+        path = python_val.get("path", None)
+
+        if path is None:
+            raise ValueError("FlyteFile's path should not be None")
+
+        return FlyteFilePathTransformer().to_python_value(
+            FlyteContextManager.current_context(),
+            Literal(
+                scalar=Scalar(
+                    blob=Blob(
+                        metadata=BlobMetadata(
+                            type=_core_types.BlobType(
+                                format="", dimensionality=_core_types.BlobType.BlobDimensionality.SINGLE
+                            )
+                        ),
+                        uri=path,
+                    )
+                )
+            ),
+            expected_python_type,
+        )
+
     def to_python_value(
         self, ctx: FlyteContext, lv: Literal, expected_python_type: typing.Union[typing.Type[FlyteFile], os.PathLike]
     ) -> FlyteFile:
         # Handle dataclass attribute access
-        if lv.scalar and lv.scalar.binary:
-            return self.from_binary_idl(lv.scalar.binary, expected_python_type)
+        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)
 
         try:
             uri = lv.scalar.blob.uri

flytekit/types/schema/__init__.py

@@ -1,5 +1,6 @@
 from .types import (
     FlyteSchema,
+    FlyteSchemaTransformer,
     LocalIOSchemaReader,
     LocalIOSchemaWriter,
     SchemaEngine,

flytekit/types/structured/__init__.py

@@ -7,9 +7,12 @@
    :template: custom.rst
    :toctree: generated/
 
-   StructuredDataset
-   StructuredDatasetEncoder
-   StructuredDatasetDecoder
+    StructuredDataset
+    StructuredDatasetDecoder
+    StructuredDatasetEncoder
+    StructuredDatasetMetadata
+    StructuredDatasetTransformerEngine
+    StructuredDatasetType
 """
 
 from flytekit.deck.renderer import ArrowRenderer, TopFrameRenderer
@@ -19,7 +22,9 @@
     StructuredDataset,
     StructuredDatasetDecoder,
     StructuredDatasetEncoder,
+    StructuredDatasetMetadata,
     StructuredDatasetTransformerEngine,
+    StructuredDatasetType,
 )
 
 

plugins/flytekit-pydantic-v2/README.md

@@ -0,0 +1 @@
+# TMP

plugins/flytekit-pydantic-v2/flytekitplugins/pydantic/v2/__init__.py

@@ -0,0 +1,25 @@
+from flytekit.types.directory import FlyteDirectory
+from flytekit.types.file import FlyteFile
+from flytekit.types.schema import FlyteSchema
+from flytekit.types.structured import StructuredDataset
+
+from . import transformer
+from .custom import (
+    deserialize_flyte_dir,
+    deserialize_flyte_file,
+    deserialize_flyte_schema,
+    deserialize_structured_dataset,
+    serialize_flyte_dir,
+    serialize_flyte_file,
+    serialize_flyte_schema,
+    serialize_structured_dataset,
+)
+
+setattr(FlyteFile, "serialize_flyte_file", serialize_flyte_file)
+setattr(FlyteFile, "deserialize_flyte_file", deserialize_flyte_file)
+setattr(FlyteDirectory, "serialize_flyte_dir", serialize_flyte_dir)
+setattr(FlyteDirectory, "deserialize_flyte_dir", deserialize_flyte_dir)
+setattr(FlyteSchema, "serialize_flyte_schema", serialize_flyte_schema)
+setattr(FlyteSchema, "deserialize_flyte_schema", deserialize_flyte_schema)
+setattr(StructuredDataset, "serialize_structured_dataset", serialize_structured_dataset)
+setattr(StructuredDataset, "deserialize_structured_dataset", deserialize_structured_dataset)

plugins/flytekit-pydantic-v2/flytekitplugins/pydantic/v2/custom.py

@@ -0,0 +1,130 @@
+from typing import Dict
+
+from flytekit.core.context_manager import FlyteContextManager
+from flytekit.models.core import types as _core_types
+from flytekit.models.literals import Blob, BlobMetadata, Literal, Scalar, Schema
+from flytekit.types.directory import FlyteDirectory, FlyteDirToMultipartBlobTransformer
+from flytekit.types.file import FlyteFile, FlyteFilePathTransformer
+from flytekit.types.schema import FlyteSchema, FlyteSchemaTransformer
+from flytekit.types.structured import (
+    StructuredDataset,
+    StructuredDatasetMetadata,
+    StructuredDatasetTransformerEngine,
+    StructuredDatasetType,
+)
+from pydantic import model_serializer, model_validator
+
+
+@model_serializer
+def serialize_flyte_file(self) -> Dict[str, str]:
+    lv = FlyteFilePathTransformer().to_literal(FlyteContextManager.current_context(), self, type(self), None)
+    return {"path": lv.scalar.blob.uri}
+
+
+@model_validator(mode="after")
+def deserialize_flyte_file(self) -> FlyteFile:
+    pv = FlyteFilePathTransformer().to_python_value(
+        FlyteContextManager.current_context(),
+        Literal(
+            scalar=Scalar(
+                blob=Blob(
+                    metadata=BlobMetadata(
+                        type=_core_types.BlobType(
+                            format="", dimensionality=_core_types.BlobType.BlobDimensionality.SINGLE
+                        )
+                    ),
+                    uri=self.path,
+                )
+            )
+        ),
+        type(self),
+    )
+    pv._remote_path = None
+    return pv
+
+
+@model_serializer
+def serialize_flyte_dir(self) -> Dict[str, str]:
+    lv = FlyteDirToMultipartBlobTransformer().to_literal(FlyteContextManager.current_context(), self, type(self), None)
+    return {"path": lv.scalar.blob.uri}
+
+
+@model_validator(mode="after")
+def deserialize_flyte_dir(self) -> FlyteDirectory:
+    pv = FlyteDirToMultipartBlobTransformer().to_python_value(
+        FlyteContextManager.current_context(),
+        Literal(
+            scalar=Scalar(
+                blob=Blob(
+                    metadata=BlobMetadata(
+                        type=_core_types.BlobType(
+                            format="", dimensionality=_core_types.BlobType.BlobDimensionality.MULTIPART
+                        )
+                    ),
+                    uri=self.path,
+                )
+            )
+        ),
+        type(self),
+    )
+    pv._remote_directory = None
+    return pv
+
+
+@model_serializer
+def serialize_flyte_schema(self) -> Dict[str, str]:
+    FlyteSchemaTransformer().to_literal(FlyteContextManager.current_context(), self, type(self), None)
+    return {"remote_path": self.remote_path}
+
+
+@model_validator(mode="after")
+def deserialize_flyte_schema(self) -> FlyteSchema:
+    # If we call the method to_python_value, FlyteSchemaTransformer will overwrite the local_path,
+    # which will lose our data.
+    # If this data is from an existed FlyteSchema, local path will be None.
+
+    if hasattr(self, "_local_path"):
+        return self
+
+    t = FlyteSchemaTransformer()
+    return t.to_python_value(
+        FlyteContextManager.current_context(),
+        Literal(scalar=Scalar(schema=Schema(self.remote_path, t._get_schema_type(type(self))))),
+        type(self),
+    )
+
+
+@model_serializer
+def serialize_structured_dataset(self) -> Dict[str, str]:
+    lv = StructuredDatasetTransformerEngine().to_literal(FlyteContextManager.current_context(), self, type(self), None)
+    sd = StructuredDataset(uri=lv.scalar.structured_dataset.uri)
+    sd.file_format = lv.scalar.structured_dataset.metadata.structured_dataset_type.format
+    return {
+        "uri": sd.uri,
+        "file_format": sd.file_format,
+    }
+
+
+@model_validator(mode="after")
+def deserialize_structured_dataset(self) -> StructuredDataset:
+    # If we call the method to_python_value, StructuredDatasetTransformerEngine will overwrite the 'dataframe',
+    # which will lose our data.
+    # If this data is from an existed StructuredDataset, dataframe will be None.
+
+    if hasattr(self, "dataframe"):
+        return self
+
+    return StructuredDatasetTransformerEngine().to_python_value(
+        FlyteContextManager.current_context(),
+        Literal(
+            scalar=Scalar(
+                structured_dataset=StructuredDataset(
+                    metadata=StructuredDatasetMetadata(
+                        structured_dataset_type=StructuredDatasetType(format=self.file_format)
+                    ),
+                    uri=self.uri,
+                )
+            )
+        ),
+        type(self),
+    )

plugins/flytekit-pydantic-v2/flytekitplugins/pydantic/v2/transformer.py

@@ -0,0 +1,73 @@
+import json
+from typing import Type
+
+import msgpack
+from google.protobuf import json_format as _json_format
+
+from flytekit import FlyteContext
+from flytekit.core.constants import MESSAGEPACK
+from flytekit.core.type_engine import TypeEngine, TypeTransformer, TypeTransformerFailedError
+from flytekit.loggers import logger
+from flytekit.models import types
+from flytekit.models.literals import Binary, Literal, Scalar
+from flytekit.models.types import LiteralType, TypeStructure
+from pydantic import BaseModel
+
+
+class PydanticTransformer(TypeTransformer[BaseModel]):
+    def __init__(self):
+        super().__init__("Pydantic Transformer", BaseModel, enable_type_assertions=False)
+
+    def get_literal_type(self, t: Type[BaseModel]) -> LiteralType:
+        schema = t.model_json_schema()
+        literal_type = {}
+        fields = t.__annotations__.items()
+
+        for name, python_type in fields:
+            try:
+                literal_type[name] = TypeEngine.to_literal_type(python_type)
+            except Exception as e:
+                logger.warning(
+                    "Field {} of type {} cannot be converted to a literal type. Error: {}".format(name, python_type, e)
+                )
+
+        ts = TypeStructure(tag="", dataclass_type=literal_type)
+
+        return types.LiteralType(simple=types.SimpleType.STRUCT, metadata=schema, structure=ts)
+
+    def to_literal(
+        self,
+        ctx: FlyteContext,
+        python_val: BaseModel,
+        python_type: Type[BaseModel],
+        expected: types.LiteralType,
+    ) -> Literal:
+        dict_obj = python_val.model_dump()
+        msgpack_bytes = msgpack.dumps(dict_obj)
+        return Literal(scalar=Scalar(binary=Binary(value=msgpack_bytes, tag="msgpack")))
+
+    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, raw=False, strict_map_key=False)
+            python_val = expected_python_type.model_validate(obj=dict_obj, strict=False)
+            return python_val
+        else:
+            raise TypeTransformerFailedError(f"Unsupported binary format: `{binary_idl_object.tag}`")
+
+    def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[BaseModel]) -> BaseModel:
+        """
+        There will have 2 kinds of literal values:
+        1. protobuf Struct (From Flyte Console)
+        2. binary scalar (Others)
+        Hence we have to handle 2 kinds of cases.
+        """
+        if lv and lv.scalar and lv.scalar.binary is not None:
+            return self.from_binary_idl(lv.scalar.binary, expected_python_type)  # type: ignore
+
+        json_str = _json_format.MessageToJson(lv.scalar.generic)
+        dict_obj = json.loads(json_str)
+        python_val = expected_python_type.model_validate(obj=dict_obj, strict=False)
+        return python_val
+
+
+TypeEngine.register(PydanticTransformer())