Using Dataclasses and Pydantic Models

Use structured type annotations

When a task input or output is a structured value, annotate it with a Python @dataclass or a Pydantic v2 BaseModel subclass instead of converting it to an untyped dictionary yourself. The runtime carries the declared type through interface publication and value conversion: _internal/runtime/types_serde.py calls TypeEngine.to_literal_type() for each interface variable, while _internal/runtime/convert.py supplies those same interface types when it serializes inputs and outputs.

A minimal dataclass shape used in the DataclassTransformer source is:

from dataclasses import dataclass


@dataclass
class Test:
    a: int
    b: str


t = Test(a=10, b="e")

For a Pydantic model, the SDK itself defines Dir as a Pydantic v2 model and a Mashumaro SerializableType:

from typing import Generic, Optional, TypeVar

from mashumaro.types import SerializableType
from pydantic import BaseModel

T = TypeVar("T")


class Dir(BaseModel, Generic[T], SerializableType):
    path: str
    name: Optional[str] = None
    format: str = ""
    hash: Optional[str] = None

You do not normally instantiate either transformer directly. TypeEngine discovers the appropriate transformer from the annotation. Dataclasses are selected by the final dataclasses.is_dataclass(python_type) fallback in TypeEngine.get_transformer(). Pydantic support is registered as the transformer for BaseModel, so subclasses are found through the type's method-resolution order.

What Flyte publishes for these types

Both transformers publish a Flyte STRUCT type, not a collection of separately named task variables. The schema is stored in the literal type's metadata, and the annotation records MessagePack as the serialization format:

from flyte.types import TypeEngine


dataclass_type = TypeEngine.to_literal_type(Test)
dir_type = TypeEngine.to_literal_type(Dir)

For Test, DataclassTransformer.get_literal_type() calls Mashumaro's build_json_schema() and places the resulting schema in metadata when generation succeeds. For Dir, PydanticTransformer.get_literal_type() calls Dir.model_json_schema() and places that Pydantic schema in metadata. Both return a LiteralType with SimpleType.STRUCT and annotation metadata equivalent to:

cache-key-metadata:
  serialization-format: msgpack

Interface publication is handled by _internal/runtime/types_serde.transform_type(), which wraps TypeEngine.to_literal_type(x) in an interface_pb2.Variable. Consequently, the declared dataclass or model type is available before any value is serialized.

Dataclass annotations wrapped in Annotated[...] are stripped for schema handling. DataclassTransformer.get_literal_type() logs that the annotations are skipped, and _get_origin_type_in_annotation() recursively handles annotated list and dictionary members. The annotation metadata above is added by the transformer itself; arbitrary Annotated metadata is not used to build the dataclass schema.

Runtime serialization and deserialization

The input path preserves the interface's type hints. convert_from_native_to_inputs() collects each declared input type, then calls TypeEngine.dict_to_literal_map(kwargs, type_hints). dict_to_literal_map() gives the supplied type hint precedence over type(v), then calls TypeEngine.to_literal() for each value. Outputs use the same declared type explicitly:

lit = await TypeEngine.to_literal(
    v,
    python_type,
    TypeEngine.to_literal_type(python_type),
)

That is the call made for each output in _internal/runtime/convert.py's convert_from_native_to_outputs(). On the way back into task code, convert_from_inputs_to_native() calls TypeEngine.literal_map_to_kwargs(), which calls TypeEngine.to_python_value() with the native input types. The task runner therefore surrounds user code with the same conversion in both directions:

native arguments
  -> convert_from_native_to_inputs
  -> dict_to_literal_map / to_literal
  -> Flyte Literal
  -> literal_map_to_kwargs / to_python_value
  -> native arguments

The normal representation for both supported structured types is a scalar whose binary value has the msgpack tag. TypeEngine.to_literal() also runs the transformer's type assertions when that transformer enables them and normalizes remote Flyte URIs after serialization.

Dataclasses: Mashumaro MessagePack conversion

DataclassTransformer.to_literal() accepts a dataclass instance and creates a cached MessagePackEncoder for the declared python_type. It encodes the instance and returns a MessagePack binary scalar. A dictionary is a special case: it is serialized directly with msgpack.dumps() rather than being passed to the dataclass encoder. A normal user-defined object that is neither a dataclass nor a dictionary raises TypeTransformerFailedError.

Deserialization first checks for a binary scalar. from_binary_idl() accepts only the exact MESSAGEPACK binary tag and uses a cached Mashumaro MessagePackDecoder. Dataclasses inheriting DataClassJSONMixin take a compatibility branch: the bytes are unpacked, converted to JSON, and passed to from_json(). to_python_value() also has a legacy generic protobuf-Struct path, using either DataClassJSONMixin.from_json() or a cached Mashumaro JSONDecoder; the source marks this path as one that should be revisited because v2 values should be binary.

Type assertions are more specific than checking only the outer Python class when the value and expected class differ. DataclassTransformer.assert_type() compares field names and declared field types, recursively checks nested dataclasses, allows optional fields to be absent, and rejects missing or extra dictionary keys. It also supports the case where a remote interface supplies a schema-derived dataclass rather than the original user-defined class.

Custom fields need Mashumaro support. If MessagePackEncoder.encode() raises NotImplementedError, the transformer raises a targeted error stating that the type should inherit mashumaro.types.SerializableType and implement _serialize() and _deserialize(). The SDK's DataFrame is an example of this convention:

from dataclasses import dataclass, field
from typing import Optional

from mashumaro.mixins.json import DataClassJSONMixin
from mashumaro.types import SerializableType


@dataclass
class DataFrame(SerializableType, DataClassJSONMixin):
    uri: Optional[str] = field(default=None)
    file_format: Optional[str] = field(default="")

In io/_dataframe/dataframe.py, DataFrame._serialize() obtains its literal type with TypeEngine.to_literal_type(type(self)) and delegates to DataFrameTransformerEngine. This specialized behavior takes precedence over the generic dataclass fallback.

Pydantic models: JSON normalization followed by MessagePack

PydanticTransformer is initialized with enable_type_assertions=False. Flyte therefore does not run the ordinary transformer-level assertion before serialization; Pydantic performs validation and coercion during model validation.

For serialization, PydanticTransformer.to_literal() follows this concrete sequence:

BaseModel instance
  -> model_dump_json()
  -> json.loads() to a dictionary
  -> msgpack.dumps()
  -> MESSAGEPACK binary scalar

The reverse binary path unpacks MessagePack with strict_map_key=False, converts the result to JSON, and calls the expected model's model_validate_json() with strict=False and context={"deserialize": True}. The expected model class matters: the transformer does not manufacture a generic dictionary when restoring a Pydantic value.

Pydantic values can also arrive as generic protobuf Struct scalars, notably through the UI compatibility path. If the literal is not a binary scalar, PydanticTransformer.to_python_value() converts lv.scalar.generic with MessageToJson() and validates that JSON through the expected model class using the same non-strict deserialization settings.

The SDK's Dir demonstrates that a model can also implement Mashumaro's SerializableType; its Pydantic inheritance is what makes the generic Pydantic transformer applicable, while specialized SDK types may register their own transformer. The same precedence rule applies to SDK File, DataFrame, and other types: an explicitly registered transformer or an MRO match is considered before the generic dataclass fallback.

Remote inspection and compatibility details

A remote interface can be converted back into Python type hints with TypeEngine.guess_python_types(). For a STRUCT containing schema metadata with a title, DataclassTransformer.guess_python_type() calls the schema conversion helper and returns a generated dataclass. This generated class is not necessarily the original class definition. The transformer comments specifically connect stable repeated reconstruction with commands such as pyflyte run and with TypeEngine.assert_type().

Pydantic models do not have an equivalent reverse-schema construction in PydanticTransformer; its deserialization requires the expected BaseModel subclass. When using a literal map directly, provide the type to LiteralsResolver.get(..., as_type=...) when possible rather than relying on type guessing.

Keep these failure modes in mind:

• Dataclass and Pydantic binary deserialization reject any binary tag other than MESSAGEPACK with TypeTransformerFailedError.
• Dataclass schema extraction catches schema-building exceptions, logs the error, and still returns a STRUCT literal type with no schema metadata. The logged guidance specifically mentions removing incompatible DataClassJsonMixin/dataclass_json combinations.
• Dataclass field validation rejects extra or missing non-optional fields and mismatched field types. This is especially visible when values cross a remote boundary and the expected type is schema-derived.
• DataclassTransformer.guess_python_type() requires STRUCT metadata containing a schema title; without it, it raises ValueError.
• Pydantic deserialization uses model_validate_json(strict=False, context={"deserialize": True}). Do not assume that Flyte performs a separate strict pre-validation pass.
• Dataclass support is a late fallback. Registering a transformer for the dataclass itself or a matching base class can change which transformer TypeEngine selects; the source notes that custom dataclass transformers have no compatibility guarantee with Flyte's transformer.