Creating Custom Type Transformers

Flyte uses a robust type system to ensure data consistency across tasks and workflows. While flyte-sdk supports many Python types out of the box (like int, str, list, dict, Pydantic models, and dataclasses), you may need to extend this system to support custom domain-specific types.

This tutorial walks you through creating a custom type transformer for a PositiveInt wrapper.

1. Define the Custom Python Type

First, define the Python class you want Flyte to recognize. In this example, we create a simple wrapper for an integer that validates the value is positive.

# custom_type.py
class PositiveInt:
    """A wrapper type that only accepts positive integers."""

    def __init__(self, value: int):
        if not isinstance(value, int):
            raise TypeError(f"Expected int, got {type(value).__name__}")
        if value <= 0:
            raise ValueError(f"Expected positive integer, got {value}")
        self._value = value

    @property
    def value(self) -> int:
        return self._value

    def __repr__(self) -> str:
        return f"PositiveInt({self._value})"

2. Implement the TypeTransformer

To teach flyte-sdk how to handle PositiveInt, you must implement a subclass of TypeTransformer. This class defines the mapping between your Python type and the Flyte IDL (Interface Definition Language).

Create a file named transformer.py and implement the following methods:

• get_literal_type: Defines the Flyte type representation.
• to_literal: Converts a Python object to a Flyte Literal.
• to_python_value: Converts a Flyte Literal back to a Python object.

# transformer.py
from typing import Type
from flyteidl2.core import literals_pb2, types_pb2
from flyte.types import TypeTransformer, TypeTransformerFailedError
from .custom_type import PositiveInt

class PositiveIntTransformer(TypeTransformer[PositiveInt]):
    def __init__(self):
        # Initialize with a name and the Python type it handles
        super().__init__(name="PositiveInt", t=PositiveInt)

    def get_literal_type(self, t: Type[PositiveInt]) -> types_pb2.LiteralType:
        """Maps PositiveInt to Flyte's INTEGER type."""
        return types_pb2.LiteralType(
            simple=types_pb2.SimpleType.INTEGER,
            structure=types_pb2.TypeStructure(tag="PositiveInt"),
        )

    async def to_literal(
        self,
        python_val: PositiveInt,
        python_type: Type[PositiveInt],
        expected: types_pb2.LiteralType,
    ) -> literals_pb2.Literal:
        """Converts PositiveInt instance to a Flyte Literal."""
        if not isinstance(python_val, PositiveInt):
            raise TypeTransformerFailedError(f"Expected PositiveInt, got {type(python_val).__name__}")

        return literals_pb2.Literal(
            scalar=literals_pb2.Scalar(
                primitive=literals_pb2.Primitive(integer=python_val.value)
            )
        )

    async def to_python_value(
        self, 
        lv: literals_pb2.Literal, 
        expected_python_type: Type[PositiveInt]
    ) -> PositiveInt:
        """Converts a Flyte Literal back to a PositiveInt instance."""
        if not lv.scalar or not lv.scalar.primitive:
            raise TypeTransformerFailedError("Missing scalar primitive in literal")

        value = lv.scalar.primitive.integer
        try:
            return PositiveInt(value)
        except (TypeError, ValueError) as e:
            raise TypeTransformerFailedError(f"Validation failed during deserialization: {e}")

    def guess_python_type(self, literal_type: types_pb2.LiteralType) -> Type[PositiveInt]:
        """Allows Flyte to infer the Python type from a stored LiteralType."""
        if (
            literal_type.simple == types_pb2.SimpleType.INTEGER
            and literal_type.structure
            and literal_type.structure.tag == "PositiveInt"
        ):
            return PositiveInt
        raise ValueError(f"Cannot guess PositiveInt from literal type {literal_type}")

3. Register the Transformer

For flyte-sdk to use your transformer, it must be registered with the TypeEngine.

Manual Registration

You can register the transformer manually in your application code:

from flyte.types import TypeEngine
from .transformer import PositiveIntTransformer

TypeEngine.register(PositiveIntTransformer())

Automatic Registration via Entry Points

For larger projects or plugins, use Python entry points in your pyproject.toml or setup.py. flyte-sdk automatically scans the flyte.plugins.types group and loads registered transformers.

In your pyproject.toml:

[project.entry-points."flyte.plugins.types"]
my_transformer = "my_package.transformer:register_transformer"

And in my_package/transformer.py:

def register_transformer():
    TypeEngine.register(PositiveIntTransformer())

4. Use the Custom Type in Tasks

Once registered, you can use PositiveInt as a type hint in any Flyte task. flyte-sdk will automatically invoke your transformer for serialization and deserialization.

import flyte
from .custom_type import PositiveInt

@flyte.task
def process_positive_int(x: PositiveInt) -> PositiveInt:
    # x is automatically converted from a Flyte Literal to a PositiveInt instance
    result = x.value + 10
    return PositiveInt(result) # Automatically converted back to a Literal

Using SimpleTransformer for Basic Types

If your custom type is a simple wrapper around a primitive and doesn't require complex logic, you can use SimpleTransformer to reduce boilerplate. It uses lambdas for conversion:

from flyte.types import SimpleTransformer, TypeEngine

# Define a transformer for a custom string-based ID type
MyIDTransformer = SimpleTransformer(
    name="MyID",
    t=str,
    lt=types_pb2.LiteralType(simple=types_pb2.SimpleType.STRING),
    to_literal_transformer=lambda x: literals_pb2.Literal(
        scalar=literals_pb2.Scalar(primitive=literals_pb2.Primitive(string_value=str(x)))
    ),
    from_literal_transformer=lambda lv: lv.scalar.primitive.string_value,
)

TypeEngine.register(MyIDTransformer)

Important Considerations

• Async Methods: Note that to_literal and to_python_value are async methods. This allows transformers to perform I/O (like uploading files or checking external registries) during conversion if necessary.
• Error Handling: Always raise TypeTransformerFailedError when conversion fails. This allows the Flyte engine to provide clear error messages to the user.
• Type Safety: In to_literal, prefer using the passed-in python_type instead of type(python_val) to ensure compatibility with generic types and inheritance.