Skip to main content

Core Type System Architecture

The core type system in this SDK is managed by the TypeEngine, which acts as a central registry and translation layer between native Python types and Flyte's internal Interface Definition Language (IDL) representation. This architecture allows the SDK to handle a wide range of types, from simple primitives to complex nested structures like Pydantic models and dataclasses.

The TypeEngine Hub​

The TypeEngine class in src/flyte/types/_type_engine.py is the primary entry point for all type-related operations. It maintains a registry of TypeTransformer objects that know how to convert specific Python types to and from Flyte Literal objects.

The engine provides three core capabilities:

  1. Type Resolution: Finding the correct transformer for a given Python type using get_transformer(python_type).
  2. Serialization: Converting Python values to Flyte literals via to_literal(python_val, python_type, expected).
  3. Deserialization: Converting Flyte literals back to Python values via to_python_value(lv, expected_python_type).

Transformer Resolution Logic​

When get_transformer is called, the engine follows a specific search order:

  1. Direct Registry Lookup: Checks if the exact type is in _REGISTRY.
  2. Special Handlers: Checks for specific structures like Annotated types, Enum, NamedTuple, TypedDict, and UnionType.
  3. MRO Search: If the type is a class, it inspects the Method Resolution Order (MRO) to find a transformer for any parent class.
  4. Dataclass/Pydantic Fallback: If no specific transformer is found, it checks if the type is a dataclass or a Pydantic model.
  5. Pickle Fallback: As a last resort, it uses the FlytePickleTransformer to serialize the object using Python's pickle module.

The TypeTransformer Interface​

Every supported type must have a corresponding implementation of the TypeTransformer abstract base class. A transformer defines how a Python type maps to a Flyte LiteralType and how values are moved between the two domains.

Implementing a Custom Transformer​

To extend the type system, you implement the following methods:

  • get_literal_type: Returns the Flyte IDL representation of the type.
  • to_literal: An async method that converts a Python value to a Literal.
  • to_python_value: An async method that converts a Literal back to a Python value.

Example of a custom transformer for a PositiveInt wrapper:

from flyte.types._type_engine import TypeTransformer, TypeEngine
from flyteidl2.core import types_pb2, literals_pb2

class PositiveIntTransformer(TypeTransformer[PositiveInt]):
    def __init__(self):
        super().__init__(name="PositiveInt", t=PositiveInt)

    def get_literal_type(self, t: Type[PositiveInt]) -> types_pb2.LiteralType:
        return types_pb2.LiteralType(simple=types_pb2.SimpleType.INTEGER)

    async def to_literal(self, python_val, python_type, expected):
        return literals_pb2.Literal(scalar=literals_pb2.Scalar(
            primitive=literals_pb2.Primitive(integer=python_val.value)
        ))

    async def to_python_value(self, lv, expected_python_type):
        return PositiveInt(lv.scalar.primitive.integer)

# Registering the transformer with the engine
TypeEngine.register(PositiveIntTransformer())

Handling Complex and Nested Types​

The TypeEngine natively supports complex Python structures by recursively applying transformers.

Dataclasses and Pydantic Models​

The engine includes a DataclassTransformer and logic for Pydantic BaseModel. These transformers inspect the fields of the class and use the TypeEngine to resolve transformers for each field. For serialization, these types often use MessagePack via the from_binary_idl method in the transformer.

Generics and Collections​

Generics like list[int] or dict[str, float] are handled by extracting the origin type (e.g., list) and the argument types (e.g., int). The engine recursively calls to_literal_type to build the nested Flyte LiteralType.

# Example of nested type resolution in tests/flyte/type_engine/test_type_engine.py
t = typing.Dict[str, typing.List[typing.Dict[str, timedelta]]]
lt = TypeEngine.to_literal_type(t)
# The engine builds a Map -> Collection -> Map -> Duration structure

LiteralsResolver and Remote Execution​

When working with FlyteRemote or raw LiteralMap objects, the LiteralsResolver class provides a high-level interface for lazy deserialization. It allows you to access elements of a map as native Python objects by providing type hints.

# Found in src/flyte/types/_type_engine.py
resolver = LiteralsResolver(literals={'x': some_literal}, variable_map=var_map)
# Lazy conversion with type hint
val = await resolver.get('x', as_type=int)

The LiteralsResolver caches converted values in _native_values to ensure that subsequent accesses do not repeat the transformation logic.

Architecture Constraints and Limitations​

The implementation in src/flyte/types/_type_engine.py enforces several strict rules to ensure compatibility with the Flyte platform:

  • Untyped Tuples: The to_literal_checks method explicitly forbids untyped tuples. Users must use typing.Tuple[int, str] or NamedTuple to provide explicit schemas.
  • Enum Restrictions: Flyte primarily supports string-valued Enums. The _ENUM_TRANSFORMER handles these, but integer-valued Enums may fail if they do not adhere to the expected IDL structure.
  • Lazy Loading: To minimize startup overhead, the engine uses lazy_import_transformers() to load heavyweight plugins (like DataFrames or Pytorch) only when a transformation is actually requested. This is protected by a lazy_import_lock to prevent race conditions during concurrent type resolution.
  • Type Assertions: If transformer.type_assertions_enabled is true, the engine calls transformer.assert_type before serialization to catch type mismatches early in the Python layer.