[WIP] Pydantic Plugin V2

.gitignore

@@ -35,7 +35,6 @@ docs/source/_tags/
 .hypothesis
 .npm
 /**/target
-coverage.xml
 
 # Version file is auto-generated by setuptools_scm
 flytekit/_version.py

Dockerfile.dev

@@ -36,7 +36,7 @@ COPY . /flytekit
 RUN SETUPTOOLS_SCM_PRETEND_VERSION_FOR_FLYTEKIT=$PSEUDO_VERSION \
     SETUPTOOLS_SCM_PRETEND_VERSION_FOR_FLYTEIDL=3.0.0dev0 \
         uv pip install --system --no-cache-dir -U \
-            "git+https://github.com/flyteorg/flyte.git@master#subdirectory=flyteidl" \
+            "git+https://github.com/flyteorg/flyte.git@32f929a6184d001cd7b471a9b19839649560cc06#subdirectory=flyteidl" \
             -e /flytekit \
             -e /flytekit/plugins/flytekit-k8s-pod \
             -e /flytekit/plugins/flytekit-deck-standard \

dev-requirements.in

@@ -1,5 +1,5 @@
 -e file:.
-flyteidl @ git+https://github.com/flyteorg/flyte.git@master#subdirectory=flyteidl
+flyteidl @ git+https://github.com/flyteorg/flyte.git@32f929a6184d001cd7b471a9b19839649560cc06#subdirectory=flyteidl
 
 coverage[toml]
 hypothesis

flytekit/core/promise.py

@@ -91,15 +91,15 @@ def my_wf(in1: int, in2: int) -> int:
         t = native_types[k]
         try:
             if type(v) is Promise:
-                v = resolve_attr_path_in_promise(v)
+                v = resolve_attr_path_in_promise(v, t)
             result[k] = TypeEngine.to_literal(ctx, v, t, var.type)
         except TypeTransformerFailedError as exc:
             raise TypeTransformerFailedError(f"Failed argument '{k}': {exc}") from None
 
     return result
 
 
-def resolve_attr_path_in_promise(p: Promise) -> Promise:
+def resolve_attr_path_in_promise(p: Promise, t: typing.Type) -> Promise:
     """
     resolve_attr_path_in_promise resolves the attribute path in a promise and returns a new promise with the resolved value
     This is for local execution only. The remote execution will be resolved in flytepropeller.
@@ -134,16 +134,26 @@ def resolve_attr_path_in_promise(p: Promise) -> Promise:
             break
 
     # If the current value is a dataclass, resolve the dataclass with the remaining path
-    if (
-        len(p.attr_path) > 0
-        and type(curr_val.value) is _literals_models.Scalar
-        and type(curr_val.value.value) is _struct.Struct
-    ):
-        st = curr_val.value.value
-        new_st = resolve_attr_path_in_pb_struct(st, attr_path=p.attr_path[used:])
-        literal_type = TypeEngine.to_literal_type(type(new_st))
-        # Reconstruct the resolved result to flyte literal (because the resolved result might not be struct)
-        curr_val = TypeEngine.to_literal(FlyteContextManager.current_context(), new_st, type(new_st), literal_type)
+    if len(p.attr_path) > 0 and type(curr_val.value) is _literals_models.Scalar:
+        import json
+
+        import msgpack
+
+        from flytekit.models.literals import Json
+
+        if type(curr_val.value.value) is _struct.Struct:
+            st = curr_val.value.value
+            new_st = resolve_attr_path_in_pb_struct(st, attr_path=p.attr_path[used:])
+            literal_type = TypeEngine.to_literal_type(type(new_st))
+            # Reconstruct the resolved result to flyte literal (because the resolved result might not be struct)
+            curr_val = TypeEngine.to_literal(FlyteContextManager.current_context(), new_st, type(new_st), literal_type)
+        elif type(curr_val.value.value) is Json:
+            json_bytes = curr_val.value.json.value
+            json_str = msgpack.loads(json_bytes)
+            dict_obj = json.loads(json_str)
+            v = resolve_attr_path_in_dict(dict_obj, attr_path=p.attr_path[used:])
+            literal_type = TypeEngine.to_literal_type(t)
+            curr_val = TypeEngine.to_literal(FlyteContextManager.current_context(), v, t, literal_type)
 
     p._val = curr_val
     return p
@@ -160,6 +170,20 @@ def resolve_attr_path_in_pb_struct(st: _struct.Struct, attr_path: List[Union[str
     return curr_val
 
 
+def resolve_attr_path_in_dict(d: dict, attr_path: List[Union[str, int]]) -> Any:
+    curr_val = d
+    for attr in attr_path:
+        try:
+            curr_val = curr_val[attr]
+        except (KeyError, IndexError, TypeError) as e:
+            raise FlytePromiseAttributeResolveException(
+                f"Failed to resolve attribute path {attr_path} in dict {curr_val}, attribute {attr} not found.\n"
+                f"Error Message: {e}"
+            )
+
+    return curr_val
+
+
 def get_primitive_val(prim: Primitive) -> Any:
     for value in [
         prim.integer,

flytekit/core/type_engine.py

@@ -487,12 +487,17 @@ def get_literal_type(self, t: Type[T]) -> LiteralType:
 
         ts = TypeStructure(tag="", dataclass_type=literal_type)
 
-        return _type_models.LiteralType(simple=_type_models.SimpleType.STRUCT, metadata=schema, structure=ts)
+        return _type_models.LiteralType(simple=_type_models.SimpleType.JSON, metadata=schema, structure=ts)
 
     def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], expected: LiteralType) -> Literal:
+        import msgpack
+
+        from flytekit.models.literals import Json
+
         if isinstance(python_val, dict):
             json_str = json.dumps(python_val)
-            return Literal(scalar=Scalar(generic=_json_format.Parse(json_str, _struct.Struct())))
+            json_bytes = msgpack.dumps(json_str)
+            return Literal(scalar=Scalar(json=Json(value=json_bytes)))
 
         if not dataclasses.is_dataclass(python_val):
             raise TypeTransformerFailedError(
@@ -519,7 +524,8 @@ def to_literal(self, ctx: FlyteContext, python_val: T, python_type: Type[T], exp
                 f" and implement _serialize and _deserialize methods."
             )
 
-        return Literal(scalar=Scalar(generic=_json_format.Parse(json_str, _struct.Struct())))  # type: ignore
+        json_bytes = msgpack.dumps(json_str)
+        return Literal(scalar=Scalar(json=Json(value=json_bytes)))
 
     def _get_origin_type_in_annotation(self, python_type: Type[T]) -> Type[T]:
         # dataclass will try to hash python type when calling dataclass.schema(), but some types in the annotation is
@@ -653,13 +659,21 @@ def _fix_dataclass_int(self, dc_type: Type[dataclasses.dataclass], dc: typing.An
         return dc
 
     def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[T]) -> T:
+        import msgpack
+
         if not dataclasses.is_dataclass(expected_python_type):
             raise TypeTransformerFailedError(
                 f"{expected_python_type} is not of type @dataclass, only Dataclasses are supported for "
                 "user defined datatypes in Flytekit"
             )
 
-        json_str = _json_format.MessageToJson(lv.scalar.generic)
+        scalar = lv.scalar
+        json_str = ""
+        if scalar.json:
+            json_bytes = lv.scalar.json.value
+            json_str = msgpack.loads(json_bytes)
+        elif scalar.generic:
+            json_str = _json_format.MessageToJson(scalar.generic)
 
         # The function looks up or creates a JSONDecoder specifically designed for the object's type.
         # This decoder is then used to convert a JSON string into a data class.
@@ -670,9 +684,12 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type:
             self._decoder[expected_python_type] = decoder
 
         dc = decoder.decode(json_str)
-
         dc = self._fix_structured_dataset_type(expected_python_type, dc)
-        return self._fix_dataclass_int(expected_python_type, dc)
+
+        if scalar.generic:
+            self._fix_dataclass_int(expected_python_type, dc)
+
+        return dc
 
     # This ensures that calls with the same literal type returns the same dataclass. For example, `pyflyte run``
     # command needs to call guess_python_type to get the TypeEngine-derived dataclass. Without caching here, separate
@@ -1636,6 +1653,28 @@ def dict_to_generic_literal(ctx: FlyteContext, v: dict, allow_pickle: bool) -> L
                 )
             raise e
 
+    @staticmethod
+    def dict_to_json_literal(ctx: FlyteContext, v: dict, allow_pickle: bool) -> Literal:
+        """
+        Creates a flyte-specific ``Literal`` value from a native python dictionary.
+        """
+        import msgpack
+
+        from flytekit.models.literals import Json
+        from flytekit.types.pickle import FlytePickle
+
+        try:
+            json_str = json.dumps(v)
+            json_bytes = msgpack.dumps(json_str)
+            return Literal(scalar=Scalar(json=Json(value=json_bytes)), metadata={"format": "json"})
+        except TypeError as e:
+            if allow_pickle:
+                remote_path = FlytePickle.to_pickle(ctx, v)
+                json_str = json.dumps({"pickle_file": remote_path})
+                json_bytes = msgpack.dumps(json_str)
+                return Literal(scalar=Scalar(json=Json(value=json_bytes)), metadata={"format": "pickle"})
+            raise e
+
     @staticmethod
     def is_pickle(python_type: Type[dict]) -> typing.Tuple[bool, Type]:
         base_type, *metadata = DictTransformer.extract_types_or_metadata(python_type)
@@ -1670,7 +1709,7 @@ def get_literal_type(self, t: Type[dict]) -> LiteralType:
                     return _type_models.LiteralType(map_value_type=sub_type)
                 except Exception as e:
                     raise ValueError(f"Type of Generic List type is not supported, {e}")
-        return _type_models.LiteralType(simple=_type_models.SimpleType.STRUCT)
+        return _type_models.LiteralType(simple=_type_models.SimpleType.JSON)
 
     def to_literal(
         self, ctx: FlyteContext, python_val: typing.Any, python_type: Type[dict], expected: LiteralType
@@ -1684,8 +1723,11 @@ def to_literal(
         if get_origin(python_type) is Annotated:
             allow_pickle, base_type = DictTransformer.is_pickle(python_type)
 
-        if expected and expected.simple and expected.simple == SimpleType.STRUCT:
-            return self.dict_to_generic_literal(ctx, python_val, allow_pickle)
+        if expected and expected.simple:
+            if expected.simple == SimpleType.STRUCT:
+                return self.dict_to_generic_literal(ctx, python_val, allow_pickle)
+            elif expected.simple == SimpleType.JSON:
+                return self.dict_to_json_literal(ctx, python_val, allow_pickle)
 
         lit_map = {}
         for k, v in python_val.items():
@@ -1720,17 +1762,36 @@ def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type:
 
         # for empty generic we have to explicitly test for lv.scalar.generic is not None as empty dict
         # evaluates to false
-        if lv and lv.scalar and lv.scalar.generic is not None:
-            if lv.metadata and lv.metadata.get("format", None) == "pickle":
-                from flytekit.types.pickle import FlytePickle
+        if lv and lv.scalar:
+            if lv.scalar.generic is not None:
+                if lv.metadata and lv.metadata.get("format", None) == "pickle":
+                    from flytekit.types.pickle import FlytePickle
 
-                uri = json.loads(_json_format.MessageToJson(lv.scalar.generic)).get("pickle_file")
-                return FlytePickle.from_pickle(uri)
+                    uri = json.loads(_json_format.MessageToJson(lv.scalar.generic)).get("pickle_file")
+                    return FlytePickle.from_pickle(uri)
 
-            try:
-                return json.loads(_json_format.MessageToJson(lv.scalar.generic))
-            except TypeError:
-                raise TypeTransformerFailedError(f"Cannot convert from {lv} to {expected_python_type}")
+                try:
+                    return json.loads(_json_format.MessageToJson(lv.scalar.generic))
+                except TypeError:
+                    raise TypeTransformerFailedError(f"Cannot convert from {lv} to {expected_python_type}")
+            elif lv.scalar.json is not None:
+                import msgpack
+
+                if lv.metadata and lv.metadata.get("format", None) == "pickle":
+                    from flytekit.types.pickle import FlytePickle
+
+                    json_bytes = lv.scalar.json.value
+                    json_str = msgpack.loads(json_bytes)
+                    uri = json.loads(json_str).get("pickle_file")
+                    return FlytePickle.from_pickle(uri)
+
+                try:
+                    json_bytes = lv.scalar.json.value
+                    json_str = msgpack.loads(json_bytes)
+                    return json.loads(json_str)
+
+                except TypeError:
+                    raise TypeTransformerFailedError(f"Cannot convert from {lv} to {expected_python_type}")
 
         raise TypeTransformerFailedError(f"Cannot convert from {lv} to {expected_python_type}")
 

flytekit/interaction/click_types.py

@@ -362,7 +362,7 @@ def literal_type_to_click_type(lt: LiteralType, python_type: typing.Type) -> cli
     Converts a Flyte LiteralType given a python_type to a click.ParamType
     """
     if lt.simple:
-        if lt.simple == SimpleType.STRUCT:
+        if lt.simple == SimpleType.STRUCT or lt.simple == SimpleType.JSON:
             ct = JsonParamType(python_type)
             ct.name = f"JSON object {python_type.__name__}"
             return ct

flytekit/interaction/string_literals.py

@@ -47,6 +47,8 @@ def scalar_to_string(scalar: Scalar) -> typing.Any:
         return MessageToDict(scalar.generic)
     if scalar.union:
         return literal_string_repr(scalar.union.value)
+    if scalar.json:
+        return scalar.json.value
     raise ValueError(f"Unknown scalar type {scalar}")
 
 

flytekit/models/literals.py

@@ -298,6 +298,32 @@ def from_flyte_idl(cls, proto):
         return cls()
 
 
+class Json(_common.FlyteIdlEntity):
+    def __init__(self, value: bytes):
+        self._value = value
+
+    @property
+    def value(self):
+        """
+        :rtype: bytes
+        """
+        return self._value
+
+    def to_flyte_idl(self):
+        """
+        :rtype: flyteidl.core.literals_pb2.Json
+        """
+        return _literals_pb2.Json(value=self.value)
+
+    @classmethod
+    def from_flyte_idl(cls, proto):
+        """
+        :param flyteidl.core.literals_pb2.Json proto:
+        :rtype: Json
+        """
+        return cls(value=proto.value)
+
+
 class BindingDataMap(_common.FlyteIdlEntity):
     def __init__(self, bindings):
         """
@@ -712,6 +738,7 @@ def __init__(
         error: Error = None,
         generic: Struct = None,
         structured_dataset: StructuredDataset = None,
+        json: Json = None,
     ):
         """
         Scalar wrapper around Flyte types.  Only one can be specified.
@@ -724,6 +751,7 @@ def __init__(
         :param Error error:
         :param google.protobuf.struct_pb2.Struct generic:
         :param StructuredDataset structured_dataset:
+        :param Json json:
         """
 
         self._primitive = primitive
@@ -735,6 +763,7 @@ def __init__(
         self._error = error
         self._generic = generic
         self._structured_dataset = structured_dataset
+        self._json = json
 
     @property
     def primitive(self):
@@ -796,6 +825,10 @@ def generic(self):
     def structured_dataset(self) -> StructuredDataset:
         return self._structured_dataset
 
+    @property
+    def json(self) -> Json:
+        return self._json
+
     @property
     def value(self):
         """
@@ -812,6 +845,7 @@ def value(self):
             or self.error
             or self.generic
             or self.structured_dataset
+            or self.json
         )
 
     def to_flyte_idl(self):
@@ -828,6 +862,7 @@ def to_flyte_idl(self):
             error=self.error.to_flyte_idl() if self.error is not None else None,
             generic=self.generic,
             structured_dataset=self.structured_dataset.to_flyte_idl() if self.structured_dataset is not None else None,
+            json=self.json.to_flyte_idl() if self.json is not None else None,
         )
 
     @classmethod
@@ -849,6 +884,7 @@ def from_flyte_idl(cls, pb2_object):
             structured_dataset=StructuredDataset.from_flyte_idl(pb2_object.structured_dataset)
             if pb2_object.HasField("structured_dataset")
             else None,
+            json=Json.from_flyte_idl(pb2_object.json) if pb2_object.HasField("json") else None,
         )
 
 

flytekit/models/types.py

@@ -22,6 +22,7 @@ class SimpleType(object):
     BINARY = _types_pb2.BINARY
     ERROR = _types_pb2.ERROR
     STRUCT = _types_pb2.STRUCT
+    JSON = _types_pb2.JSON
 
 
 class SchemaType(_common.FlyteIdlEntity):

flytekit/types/directory/__init__.py

@@ -16,7 +16,7 @@
 
 import typing
 
-from .types import FlyteDirectory
+from .types import FlyteDirectory, FlyteDirToMultipartBlobTransformer
 
 # The following section provides some predefined aliases for commonly used FlyteDirectory formats.