Feat: Add type support for pydantic BaseModels

flytekit/clis/sdk_in_container/run.py

@@ -111,15 +111,13 @@ class PickleParamType(click.ParamType):
     def convert(
         self, value: typing.Any, param: typing.Optional[click.Parameter], ctx: typing.Optional[click.Context]
     ) -> typing.Any:
-
         uri = FlyteContextManager.current_context().file_access.get_random_local_path()
         with open(uri, "w+b") as outfile:
             cloudpickle.dump(value, outfile)
         return FileParam(filepath=str(pathlib.Path(uri).resolve()))
 
 
 class DateTimeType(click.DateTime):
-
     _NOW_FMT = "now"
     _ADDITONAL_FORMATS = [_NOW_FMT]
 
@@ -276,7 +274,6 @@ def get_uri_for_dir(
     def convert_to_structured_dataset(
         self, ctx: typing.Optional[click.Context], param: typing.Optional[click.Parameter], value: Directory
     ) -> Literal:
-
         uri = self.get_uri_for_dir(ctx, value, "00000.parquet")
 
         lit = Literal(
@@ -338,7 +335,7 @@ def convert_to_union(
                 python_val = converter._click_type.convert(value, param, ctx)
                 literal = converter.convert_to_literal(ctx, param, python_val)
                 return Literal(scalar=Scalar(union=Union(literal, variant)))
-            except (Exception or AttributeError) as e:
+            except Exception or AttributeError as e:
                 logging.debug(f"Failed to convert python type {python_type} to literal type {variant}", e)
         raise ValueError(f"Failed to convert python type {self._python_type} to literal type {lt}")
 
@@ -399,7 +396,10 @@ def convert_to_struct(
         Convert the loaded json object to a Flyte Literal struct type.
         """
         if type(value) != self._python_type:
-            o = cast(DataClassJsonMixin, self._python_type).from_json(json.dumps(value))
+            if is_pydantic_basemodel(self._python_type):
+                o = self._python_type.parse_raw(json.dumps(value))  # type: ignore
+            else:
+                o = cast(DataClassJsonMixin, self._python_type).from_json(json.dumps(value))
         else:
             o = value
         return TypeEngine.to_literal(self._flyte_ctx, o, self._python_type, self._literal_type)
@@ -446,6 +446,15 @@ def convert(self, ctx, param, value) -> typing.Union[Literal, typing.Any]:
             raise click.BadParameter(f"Failed to convert param {param}, {value} to {self._python_type}") from e
 
 
+def is_pydantic_basemodel(python_type: typing.Type) -> bool:
+    try:
+        import pydantic
+    except ImportError:
+        return False
+    else:
+        return issubclass(python_type, pydantic.BaseModel)
+
+
 def to_click_option(
     ctx: click.Context,
     flyte_ctx: FlyteContext,

flytekit/core/type_engine.py

@@ -710,7 +710,7 @@ def register_additional_type(cls, transformer: TypeTransformer, additional_type:
             cls._REGISTRY[additional_type] = transformer
 
     @classmethod
-    def get_transformer(cls, python_type: Type) -> TypeTransformer[T]:
+    def get_transformer(cls, python_type: T) -> TypeTransformer[T]:
         """
         The TypeEngine hierarchy for flyteKit. This method looksup and selects the type transformer. The algorithm is
         as follows

plugins/flytekit-pydantic/README.md

@@ -0,0 +1,28 @@
+# Flytekit Pydantic Plugin
+
+Pydantic is a data validation and settings management library that uses Python type annotations to enforce type hints at runtime and provide user-friendly errors when data is invalid. Pydantic models are classes that inherit from `pydantic.BaseModel` and are used to define the structure and validation of data using Python type annotations.
+
+The plugin adds type support for pydantic models.
+
+To install the plugin, run the following command:
+
+```bash
+pip install flytekitplugins-pydantic
+```
+
+
+## Type Example
+```python
+from pydantic import BaseModel
+
+
+class TrainConfig(BaseModel):
+    lr: float = 1e-3
+    batch_size: int = 32
+    files: List[FlyteFile]
+    directories: List[FlyteDirectory]
+
+@task
+def train(cfg: TrainConfig):
+    ...
+```

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

@@ -0,0 +1,4 @@
+from .basemodel_transformer import BaseModelTransformer
+from .deserialization import set_validators_on_supported_flyte_types as _set_validators_on_supported_flyte_types
+
+_set_validators_on_supported_flyte_types()  # enables you to use flytekit.types in pydantic model

plugins/flytekit-pydantic/flytekitplugins/pydantic/basemodel_transformer.py

@@ -0,0 +1,65 @@
+"""Serializes & deserializes the pydantic basemodels """
+
+from typing import Type
+
+import pydantic
+from google.protobuf import json_format
+from typing_extensions import Annotated
+
+from flytekit import FlyteContext
+from flytekit.core import type_engine
+from flytekit.models import literals, types
+
+from . import deserialization, serialization
+
+BaseModelLiteralValue = Annotated[
+    literals.LiteralMap,
+    """
+    BaseModel serialized to a LiteralMap consisting of: 
+        1) the basemodel json with placeholders for flyte types 
+        2) mapping from placeholders to serialized flyte type values in the object store
+    """,
+]
+
+
+class BaseModelTransformer(type_engine.TypeTransformer[pydantic.BaseModel]):
+    _TYPE_INFO = types.LiteralType(simple=types.SimpleType.STRUCT)
+
+    def __init__(self):
+        """Construct pydantic.BaseModelTransformer."""
+        super().__init__(name="basemodel-transform", t=pydantic.BaseModel)
+
+    def get_literal_type(self, t: Type[pydantic.BaseModel]) -> types.LiteralType:
+        return types.LiteralType(simple=types.SimpleType.STRUCT)
+
+    def to_literal(
+        self,
+        ctx: FlyteContext,
+        python_val: pydantic.BaseModel,
+        python_type: Type[pydantic.BaseModel],
+        expected: types.LiteralType,
+    ) -> BaseModelLiteralValue:
+        """Convert a given ``pydantic.BaseModel`` to the Literal representation."""
+        return serialization.serialize_basemodel(python_val)
+
+    def to_python_value(
+        self,
+        ctx: FlyteContext,
+        lv: BaseModelLiteralValue,
+        expected_python_type: Type[pydantic.BaseModel],
+    ) -> pydantic.BaseModel:
+        """Re-hydrate the pydantic BaseModel object from Flyte Literal value."""
+        basemodel_json_w_placeholders = read_basemodel_json_from_literalmap(lv)
+        flyte_obj_literalmap = lv.literals[serialization.FLYTETYPE_OBJSTORE_KEY]
+        with deserialization.PydanticDeserializationLiteralStore.attach(flyte_obj_literalmap):
+            return expected_python_type.parse_raw(basemodel_json_w_placeholders)
+
+
+def read_basemodel_json_from_literalmap(lv: BaseModelLiteralValue) -> serialization.SerializedBaseModel:
+    basemodel_literal: literals.Literal = lv.literals[serialization.BASEMODEL_JSON_KEY]
+    basemodel_json_w_placeholders = json_format.MessageToJson(basemodel_literal.scalar.generic)
+    assert isinstance(basemodel_json_w_placeholders, str)
+    return basemodel_json_w_placeholders
+
+
+type_engine.TypeEngine.register(BaseModelTransformer())

plugins/flytekit-pydantic/flytekitplugins/pydantic/commons.py

@@ -0,0 +1,29 @@
+import builtins
+import datetime
+import typing
+
+from typing_extensions import Annotated
+
+import pandas as pd
+
+
+from flytekit.core import type_engine
+
+MODULES_TO_EXLCLUDE_FROM_FLYTE_TYPES = {m.__name__ for m in [builtins, typing, datetime]}
+
+
+def include_in_flyte_types(t: type) -> bool:
+    if t is None:
+        return False
+    if t.__module__ in MODULES_TO_EXLCLUDE_FROM_FLYTE_TYPES:
+        return False
+    return True
+
+
+PYDANTIC_SUPPORTED_FLYTE_TYPES = tuple(
+    filter(include_in_flyte_types, type_engine.TypeEngine.get_available_transformers())
+) + (pd.DataFrame,)
+
+# this is the UUID placeholder that is set in the serialized basemodel JSON, connecting that field to
+# the literal map that holds the actual object that needs to be deserialized (w/ protobuf)
+LiteralObjID = Annotated[str, "Key for unique object in literal map."]

plugins/flytekit-pydantic/flytekitplugins/pydantic/deserialization.py

@@ -0,0 +1,148 @@
+import contextlib
+from typing import Any, Callable, Dict, Generator, Iterator, List, Optional, Type, TypeVar, Union
+
+import pydantic
+from flytekit.core import context_manager, type_engine
+from flytekit.models import literals
+
+from flytekit.types import directory, file
+
+from flytekitplugins.pydantic import commons, serialization
+
+
+# this field is used by pydantic to get the validator method
+PYDANTIC_VALIDATOR_METHOD_NAME = (
+    pydantic.BaseModel.__get_validators__.__name__
+    if pydantic.__version__ < "2.0.0"
+    else pydantic.BaseModel.__get_pydantic_core_schema__.__name___  # type: ignore
+)
+PythonType = TypeVar("PythonType")  # target type of the deserialization
+
+
+class PydanticDeserializationLiteralStore:
+    """
+    The purpose of this class is to provide a context manager that can be used to deserialize a basemodel from a
+    literal map.
+
+    Because pydantic validators are fixed when subclassing a BaseModel, this object is a singleton that
+    serves as a namespace that can be set with the attach_to_literalmap context manager for the time that
+    a basemodel is being deserialized. The validators are then accessing this namespace for the flyteobj
+    placeholders that it is trying to deserialize.
+    """
+
+    literal_store: Optional[serialization.LiteralStore] = None  # attachement point for the literal map
+
+    def __init__(self) -> None:
+        raise Exception("This class should not be instantiated")
+
+    def __init_subclass__(cls) -> None:
+        raise Exception("This class should not be subclassed")
+
+    @classmethod
+    @contextlib.contextmanager
+    def attach(cls, literal_map: literals.LiteralMap) -> Generator[None, None, None]:
+        """
+        Read a literal map and populate the object store from it.
+
+        This can be used as a context manager to attach to a literal map for the duration of a deserialization
+        Note that this is not threadsafe, and designed to manage a single deserialization at a time.
+        """
+        assert not cls.is_attached(), "can only be attached to one literal map at a time."
+        try:
+            cls.literal_store = literal_map.literals
+            yield
+        finally:
+            cls.literal_store = None
+
+    @classmethod
+    def contains(cls, item: commons.LiteralObjID) -> bool:
+        assert cls.is_attached(), "can only check for existence of a literal when attached to a literal map"
+        assert cls.literal_store is not None
+        return item in cls.literal_store
+
+    @classmethod
+    def is_attached(cls) -> bool:
+        return cls.literal_store is not None
+
+    @classmethod
+    def get_python_object(
+        cls, identifier: commons.LiteralObjID, expected_type: Type[PythonType]
+    ) -> Optional[PythonType]:
+        """Deserialize a flyte literal and return the python object."""
+        if not cls.is_attached():
+            raise Exception("Must attach to a literal map before deserializing")
+        literal = cls.literal_store[identifier]  # type: ignore
+        python_object = deserialize_flyte_literal(literal, expected_type)
+        return python_object
+
+
+def set_validators_on_supported_flyte_types() -> None:
+    """
+    Set pydantic validator for the flyte types supported by this plugin.
+    """
+    for flyte_type in commons.PYDANTIC_SUPPORTED_FLYTE_TYPES:
+        setattr(flyte_type, PYDANTIC_VALIDATOR_METHOD_NAME, add_flyte_validators_for_type(flyte_type))
+
+
+def add_flyte_validators_for_type(
+    flyte_obj_type: Type[type_engine.T],
+) -> Callable[[Any], Iterator[Callable[[Any], type_engine.T]]]:
+    """
+    Add flyte deserialisation validators to a type.
+    """
+
+    previous_validators = getattr(flyte_obj_type, PYDANTIC_VALIDATOR_METHOD_NAME, lambda *_: [])()
+
+    def validator(object_uid_maybe: Union[commons.LiteralObjID, Any]) -> Union[type_engine.T, Any]:
+        """Partial of deserialize_flyte_literal with the object_type fixed"""
+        if not PydanticDeserializationLiteralStore.is_attached():
+            return object_uid_maybe  # this validator should only trigger when we are deserializeing
+        if not isinstance(object_uid_maybe, str):
+            return object_uid_maybe  # object uids are strings and we dont want to trigger on other types
+        if not PydanticDeserializationLiteralStore.contains(object_uid_maybe):
+            return object_uid_maybe  # final safety check to make sure that the object uid is in the literal map
+        return PydanticDeserializationLiteralStore.get_python_object(object_uid_maybe, flyte_obj_type)
+
+    def validator_generator(*args, **kwags) -> Iterator[Callable[[Any], type_engine.T]]:
+        """Generator that returns validators."""
+        yield validator
+        yield from previous_validators
+        yield from ADDITIONAL_FLYTETYPE_VALIDATORS.get(flyte_obj_type, [])
+
+    return validator_generator
+
+
+def validate_flytefile(flytefile: Union[str, file.FlyteFile]) -> file.FlyteFile:
+    """Validate a flytefile (i.e. deserialize)."""
+    if isinstance(flytefile, file.FlyteFile):
+        return flytefile
+    if isinstance(flytefile, str):  # when e.g. initializing from config
+        return file.FlyteFile(flytefile)
+    else:
+        raise ValueError(f"Invalid type for flytefile: {type(flytefile)}")
+
+
+def validate_flytedir(flytedir: Union[str, directory.FlyteDirectory]) -> directory.FlyteDirectory:
+    """Validate a flytedir (i.e. deserialize)."""
+    if isinstance(flytedir, directory.FlyteDirectory):
+        return flytedir
+    if isinstance(flytedir, str):  # when e.g. initializing from config
+        return directory.FlyteDirectory(flytedir)
+    else:
+        raise ValueError(f"Invalid type for flytedir: {type(flytedir)}")
+
+
+ADDITIONAL_FLYTETYPE_VALIDATORS: Dict[Type, List[Callable[[Any], Any]]] = {
+    file.FlyteFile: [validate_flytefile],
+    directory.FlyteDirectory: [validate_flytedir],
+}
+
+
+def deserialize_flyte_literal(
+    flyteobj_literal: literals.Literal, python_type: Type[PythonType]
+) -> Optional[PythonType]:
+    """Deserialize a Flyte Literal into the python object instance."""
+    ctx = context_manager.FlyteContext.current_context()
+    transformer = type_engine.TypeEngine.get_transformer(python_type)
+    python_obj = transformer.to_python_value(ctx, flyteobj_literal, python_type)
+    return python_obj