Support pydantic plugin in 2.xx version (#2217)

* support pydantic when 2.xx version Signed-off-by: Future-Outlier <[email protected]> * update import Signed-off-by: Future-Outlier <[email protected]> * update Signed-off-by: Future-Outlier <[email protected]> * Revert Dev Dependencies Signed-off-by: Future-Outlier <[email protected]> * lint Signed-off-by: Future-Outlier <[email protected]> * pydantic Signed-off-by: Future-Outlier <[email protected]> * fix-plugin-ci Signed-off-by: Future-Outlier <[email protected]> * as pydantic, remove lazy load Signed-off-by: Future-Outlier <[email protected]> * lint Signed-off-by: Future-Outlier <[email protected]> * warning messages Signed-off-by: Future-Outlier <[email protected]> * lint Signed-off-by: Future-Outlier <[email protected]> * update pydantic version Signed-off-by: Future-Outlier <[email protected]> * fix build plugins CI Signed-off-by: Future-Outlier <[email protected]> * fix version problem Signed-off-by: Future-Outlier <[email protected]> * Update pingsu's advice Co-authored-by: Kevin Su <[email protected]> Signed-off-by: Future-Outlier <[email protected]> * update supported types Signed-off-by: Future-Outlier <[email protected]> * update pingsu's advice Signed-off-by: Future-Outlier <[email protected]> * remove tests Signed-off-by: Future-Outlier <[email protected]> --------- Signed-off-by: Future-Outlier <[email protected]> Co-authored-by: Kevin Su <[email protected]> Signed-off-by: Jan Fiedler <[email protected]>
flyteorg · Jul 25, 2024 · d1bacf3 · d1bacf3
1 parent a526636
commit d1bacf3
Show file tree

Hide file tree

Showing 7 changed files with 45 additions and 16 deletions.
diff --git a/.github/workflows/pythonbuild.yml b/.github/workflows/pythonbuild.yml
@@ -253,14 +253,17 @@ jobs:
           - flytekit-dolt
           - flytekit-duckdb
           - flytekit-envd
+          - flytekit-flyteinteractive
           - flytekit-greatexpectations
           - flytekit-hive
           - flytekit-huggingface
+          - flytekit-identity-aware-proxy
           - flytekit-k8s-pod
           - flytekit-kf-mpi
           - flytekit-kf-pytorch
           - flytekit-kf-tensorflow
           - flytekit-mlflow
+          - flytekit-mmcloud
           - flytekit-modin
           - flytekit-onnx-pytorch
           - flytekit-onnx-scikitlearn
@@ -271,12 +274,12 @@ jobs:
           - flytekit-pandera
           - flytekit-papermill
           - flytekit-polars
+          - flytekit-pydantic
           - flytekit-ray
           - flytekit-snowflake
           - flytekit-spark
           - flytekit-sqlalchemy
           - flytekit-vaex
-          - flytekit-flyteinteractive
           - flytekit-whylogs
         exclude:
           # flytekit-modin depends on ray which does not have a 3.11 wheel yet.

diff --git a/flytekit/interaction/click_types.py b/flytekit/interaction/click_types.py
@@ -28,11 +28,16 @@ def is_pydantic_basemodel(python_type: typing.Type) -> bool:
     Checks if the python type is a pydantic BaseModel
     """
     try:
-        import pydantic
+        import pydantic  # noqa: F401
     except ImportError:
         return False
     else:
-        return issubclass(python_type, pydantic.BaseModel)
+        try:
+            from pydantic.v1 import BaseModel
+        except ImportError:
+            from pydantic import BaseModel
+
+        return issubclass(python_type, BaseModel)
 
 
 def key_value_callback(_: typing.Any, param: str, values: typing.List[str]) -> typing.Optional[typing.Dict[str, str]]:

diff --git a/plugins/flytekit-pydantic/README.md b/plugins/flytekit-pydantic/README.md
@@ -13,7 +13,7 @@ pip install flytekitplugins-pydantic
 
 ## Type Example
 ```python
-from pydantic import BaseModel
+from pydantic.v1 import BaseModel
 
 
 class TrainConfig(BaseModel):

diff --git a/plugins/flytekit-pydantic/flytekitplugins/pydantic/basemodel_transformer.py b/plugins/flytekit-pydantic/flytekitplugins/pydantic/basemodel_transformer.py
@@ -5,13 +5,18 @@
 from google.protobuf import json_format
 from typing_extensions import Annotated
 
-from flytekit import FlyteContext, lazy_module
+from flytekit import FlyteContext
 from flytekit.core import type_engine
 from flytekit.models import literals, types
 
 from . import deserialization, serialization
 
-pydantic = lazy_module("pydantic")
+try:
+    # TODO: Use pydantic v2 to serialize/deserialize data
+    # https://github.com/flyteorg/flyte/issues/5033
+    import pydantic.v1 as pydantic
+except ImportError:
+    import pydantic
 
 BaseModelLiterals = Annotated[
     Dict[str, literals.Literal],
@@ -41,6 +46,12 @@ def to_literal(
         expected: types.LiteralType,
     ) -> literals.Literal:
         """Convert a given ``pydantic.BaseModel`` to the Literal representation."""
+        import warnings
+
+        warnings.warn(
+            "If you are using Pydantic version 2.0 or later, please import BaseModel using `from pydantic.v1 import BaseModel`.",
+            FutureWarning,
+        )
         return serialization.serialize_basemodel(python_val)
 
     def to_python_value(

diff --git a/plugins/flytekit-pydantic/flytekitplugins/pydantic/deserialization.py b/plugins/flytekit-pydantic/flytekitplugins/pydantic/deserialization.py
@@ -3,15 +3,20 @@
 
 from flytekitplugins.pydantic import commons, serialization
 
-from flytekit import lazy_module
 from flytekit.core import context_manager, type_engine
 from flytekit.models import literals
 from flytekit.types import directory, file
 
-pydantic = lazy_module("pydantic")
+try:
+    # TODO: Use pydantic v2 to serialize/deserialize data
+    # https://github.com/flyteorg/flyte/issues/5033
+    import pydantic.v1 as pydantic
+except ImportError:
+    import pydantic
 
 # this field is used by pydantic to get the validator method
 PYDANTIC_VALIDATOR_METHOD_NAME = pydantic.BaseModel.__get_validators__.__name__
+
 PythonType = TypeVar("PythonType")  # target type of the deserialization
 
 

diff --git a/plugins/flytekit-pydantic/flytekitplugins/pydantic/serialization.py b/plugins/flytekit-pydantic/flytekitplugins/pydantic/serialization.py
@@ -14,13 +14,17 @@
 from google.protobuf import json_format, struct_pb2
 from typing_extensions import Annotated
 
-from flytekit import lazy_module
 from flytekit.core import context_manager, type_engine
 from flytekit.models import literals
 
 from . import commons
 
-pydantic = lazy_module("pydantic")
+try:
+    # TODO: Use pydantic v2 to serialize/deserialize data
+    # https://github.com/flyteorg/flyte/issues/5033
+    import pydantic.v1 as pydantic
+except ImportError:
+    import pydantic
 
 BASEMODEL_JSON_KEY = "BaseModel JSON"
 OBJECTS_KEY = "Serialized Flyte Objects"

diff --git a/plugins/flytekit-pydantic/tests/test_type_transformer.py b/plugins/flytekit-pydantic/tests/test_type_transformer.py
@@ -7,15 +7,20 @@
 import pytest
 from flyteidl.core.types_pb2 import SimpleType
 from flytekitplugins.pydantic import BaseModelTransformer
-from flytekitplugins.pydantic.commons import PYDANTIC_SUPPORTED_FLYTE_TYPES
-from pydantic import BaseModel, Extra
 
 import flytekit
 from flytekit.core import context_manager
 from flytekit.core.type_engine import TypeEngine
 from flytekit.types import directory
 from flytekit.types.file import file
 
+try:
+    # TODO: Use pydantic v2 to serialize/deserialize data
+    # https://github.com/flyteorg/flyte/issues/5033
+    from pydantic.v1 import BaseModel, Extra
+except ImportError:
+    from pydantic import BaseModel, Extra
+
 
 class TrainConfig(BaseModel):
     """Config BaseModel for testing purposes."""
@@ -276,10 +281,6 @@ def wf():
     wf()
 
 
-def test_supported():
-    assert len(PYDANTIC_SUPPORTED_FLYTE_TYPES) == 9
-
-
 def test_single_df():
     ctx = context_manager.FlyteContextManager.current_context()
     lt = TypeEngine.to_literal_type(ConfigWithPandasDataFrame)