Skip to main content

How to Add Support for a New Type

What you will build​

You will add a transformer that makes a Python type available to Flyte’s type system. The concrete walkthrough uses flyte-sdk’s FileTransformer: it maps File values to single-part blob literals, reconstructs File values from those literals, supports reverse type inference, and registers the adapter with TypeEngine. The same lifecycle applies when you add a transformer for another Python class: declare the Flyte IDL shape, implement both asynchronous conversion directions, and register the transformer before the type is used.

Prerequisites​

You need flyte-sdk, the flyteidl protobuf classes used by the transformer, and the Python type you are adapting. The public API for custom transformers is exported from flyte.types:

from flyte.types import TypeEngine, TypeTransformer, TypeTransformerFailedError

The protobuf values used by the file adapter come from flyteidl.core:

from flyteidl.core import literals_pb2, types_pb2

If your representation uses remote or offloaded data, the corresponding Flyte storage configuration must also be available at runtime. TypeEngine resolves offloaded literals through flyte.storage before invoking to_python_value.

1. Choose the Flyte literal shape​

Start by deciding which Flyte IDL representation describes your Python value. FileTransformer chooses a LiteralType containing a BlobType whose dimensionality is SINGLE:

def get_literal_type(self, t: Type[File]) -> types_pb2.LiteralType:
    """Get the Flyte literal type for a File type."""
    return types_pb2.LiteralType(
        blob=types_pb2.BlobType(
            format="",
            dimensionality=types_pb2.BlobType.BlobDimensionality.SINGLE,
        )
    )

The t argument is the declared Python type. The file implementation leaves the format empty in the interface type and supplies the actual format in the value’s blob metadata.

For a one-to-one mapping to an existing simple literal shape, flyte-sdk also provides SimpleTransformer. Its constructor takes the display name, Python type, LiteralType, and two callables:

DateTransformer = SimpleTransformer(
    "date",
    datetime.date,
    types_pb2.LiteralType(simple=types_pb2.SimpleType.DATETIME),
    lambda x: Literal(
        scalar=Scalar(primitive=Primitive(datetime=datetime.datetime.combine(x, datetime.time.min)))
    ),
    lambda x: x.scalar.primitive.datetime.date() if x.scalar.primitive.HasField("datetime") else None,
)

SimpleTransformer always returns its configured literal type. Its to_literal path requires type(python_val) is self._type, and its reverse path verifies the decoded value has that exact type. It is therefore a compact choice for a strict, single-shape conversion; use a TypeTransformer subclass when your type needs structural validation or custom reconstruction.

2. Subclass TypeTransformer​

Call the base constructor with a name and the Python class handled by the transformer. The name is used in registry diagnostics and the python_type property is the key used by TypeEngine.register:

class FileTransformer(TypeTransformer[File]):
    """
    Transformer for File objects. This type transformer does not handle any i/o. That is now the responsibility of the
    user.
    """

    def __init__(self):
        super().__init__(name="File", t=File)

TypeTransformer requires three methods: get_literal_type, asynchronous to_literal, and asynchronous to_python_value. Even when the conversion itself is synchronous, the two value-conversion methods must be declared with async def.

The base class can perform type assertions before serialization. With the default enable_type_assertions=True, TypeEngine.to_literal calls transformer.assert_type(python_type, python_val). The file adapter additionally validates the value inside to_literal, which lets it produce a domain-specific failure message.

3. Implement Python-to-Flyte serialization​

Construct a literals_pb2.Literal whose shape agrees with the LiteralType returned above. FileTransformer.to_literal checks the value, copies its path into the blob URI, writes its format into blob metadata, and preserves its optional hash:

async def to_literal(
    self,
    python_val: File,
    python_type: Type[File],
    expected: types_pb2.LiteralType,
) -> literals_pb2.Literal:
    """Convert a File object to a Flyte literal."""
    if not isinstance(python_val, File):
        raise TypeTransformerFailedError(f"Expected File object, received {type(python_val)}")

    return literals_pb2.Literal(
        scalar=literals_pb2.Scalar(
            blob=literals_pb2.Blob(
                metadata=literals_pb2.BlobMetadata(
                    type=types_pb2.BlobType(
                        format=python_val.format,
                        dimensionality=types_pb2.BlobType.BlobDimensionality.SINGLE,
                    )
                ),
                uri=python_val.path,
            )
        ),
        hash=python_val.hash if python_val.hash else None,
    )

The expected argument is the literal type returned by TypeEngine.to_literal_type(python_type). The base contract asks implementations to use the declared python_type rather than relying only on type(python_val); the stricter SimpleTransformer is the exception shown above. Raise TypeTransformerFailedError when the value cannot be represented. This exception is also a TypeError, AssertionError, and ValueError, so broad catches of those base exception types can catch transformer failures too.

4. Implement Flyte-to-Python deserialization​

Validate the incoming literal’s shape before constructing the Python object. The file adapter accepts only a blob with single-part dimensionality and derives the file name from the URI:

async def to_python_value(
    self,
    lv: literals_pb2.Literal,
    expected_python_type: Type[File],
) -> File:
    """Convert a Flyte literal to a File object."""
    if not lv.scalar.HasField("blob"):
        raise TypeTransformerFailedError(f"Expected blob literal, received {lv}")
    if not lv.scalar.blob.metadata.type.dimensionality == types_pb2.BlobType.BlobDimensionality.SINGLE:
        raise TypeTransformerFailedError(
            f"Expected single part blob, received {lv.scalar.blob.metadata.type.dimensionality}"
        )

    uri = lv.scalar.blob.uri
    filename = Path(uri).name
    hash_value = lv.hash if lv.hash else None
    f: File = File(path=uri, name=filename, format=lv.scalar.blob.metadata.type.format, hash=hash_value)
    return f

TypeEngine.to_python_value selects the transformer from expected_python_type and then calls this method. If the literal contains offloaded_metadata, TypeEngine first downloads and loads the underlying Literal; the transformer still receives the resulting literal shape.

5. Add reverse type inference when needed​

Implement guess_python_type if Flyte interfaces or literal maps must be converted back into Python annotations. TypeEngine.guess_python_type tries registered transformers until one accepts the literal type. FileTransformer recognizes single-part blobs other than the pickle format:

def guess_python_type(self, literal_type: types_pb2.LiteralType) -> Type[File]:
    """Guess the Python type from a Flyte literal type."""
    if (
        literal_type.HasField("blob")
        and literal_type.blob.dimensionality == types_pb2.BlobType.BlobDimensionality.SINGLE
        and literal_type.blob.format != "PythonPickle"
    ):
        return File
    raise ValueError(f"Cannot guess python type from {literal_type}")

If you omit this method, the base implementation raises ValueError, so TypeEngine.guess_python_type cannot infer your Python type from a remote LiteralType. That does not prevent direct conversion when the declared Python type is already available.

6. Register the transformer​

Register an instance after the transformer and its Python type are defined:

TypeEngine.register(FileTransformer())

TypeEngine.register stores the transformer under transformer.python_type. It rejects an existing mapping with ValueError rather than replacing it. You can register aliases at the same time with additional_types:

TypeEngine.register(transformer, additional_types=[additional_type])

Use register_additional_type when one already-created transformer intentionally serves another concrete type. The dataframe subsystem uses this pattern for handler types and passes override=True:

engine = DataFrameTransformerEngine()
TypeEngine.register_additional_type(engine, h.python_type, override=True)

Register code is normally executed when the module is imported. Import timing matters: lookup can fall through to the dataclass transformer and, finally, a FlytePickleTransformer when no registered transformer matches. A custom transformer must be registered before Flyte needs the custom type’s interface or value conversion.

7. Verify the runtime path​

Task output serialization uses the registry rather than calling your transformer directly:

for (output_name, python_type), v in zip(interface.outputs.items(), o):
    try:
        lit = await TypeEngine.to_literal(v, python_type, TypeEngine.to_literal_type(python_type))
        named.append(run_definition_pb2.NamedLiteral(name=output_name, value=lit))
    except TypeTransformerFailedError as e:
        raise flyte.errors.RuntimeDataValidationError(output_name, e, task_name)

This gives you a concrete verification point: TypeEngine.to_literal_type must return the shape your to_literal emits, and an invalid value should produce a useful TypeTransformerFailedError that is then attached to the named task output as a runtime validation error.

Input reconstruction follows the reverse direction through literal_map_to_kwargs:

kwargs = await TypeEngine.literal_map_to_kwargs(lm, interface.outputs)

That method ultimately invokes TypeEngine.to_python_value with the declared Python types. dict_to_literal_map performs the corresponding batch conversion for Python dictionaries and gives explicit type hints precedence over type(value), which is important for generic collections whose element types are erased at runtime.

Common failure points​

  • Duplicate registration: TypeEngine.register raises ValueError if the Python type is already registered. Do not use override=True accidentally; that option belongs to register_additional_type.
  • Wrong method shape: to_literal and to_python_value are asynchronous abstract methods. Implement them with async def and return the protobuf Literal or reconstructed Python value.
  • Mismatched protobuf shape: Validate fields such as scalar.blob and dimensionality before reading them. Raise TypeTransformerFailedError for unsupported literal shapes.
  • Exact versus subclass validation: FileTransformer uses isinstance, while SimpleTransformer requires exact runtime type equality. Choose the behavior deliberately in your implementation.
  • Optional values and tuples: TypeEngine.to_literal_checks rejects None for a non-optional declared type and rejects tuples as individual Flyte values. Tuple types are registered as restricted types during default initialization.
  • Fallback behavior: An unregistered class can resolve to a pickle transformer after MRO and dataclass lookup. That is different from having a first-class custom literal representation and does not make guess_python_type recognize your custom literal.
  • Lazy imports: TypeEngine loads dataframe handlers on demand under a lock. Extra transformers such as PyTorch are not automatically imported according to the source TODO, so registration modules must be imported when required.

Complete implementation checklist​

Before using the new type in a workflow, confirm that you have:

  1. Called super().__init__(name=..., t=...) in a TypeTransformer subclass.
  2. Returned a LiteralType that matches the literal shape emitted by to_literal.
  3. Implemented both asynchronous conversion methods.
  4. Validated incoming values and literals with informative TypeTransformerFailedError messages.
  5. Implemented guess_python_type when reverse interface inference is required.
  6. Registered the transformer at import time with TypeEngine.register.
  7. Checked that no existing registry entry conflicts with the new type.
  8. Verified a Python-value-to-Literal-to-Python-value round trip, including invalid-value and invalid-literal cases.

The FileTransformer in io/_file.py is the complete blob-oriented reference: it defines the literal contract, performs both conversions, supports reverse inference, and ends with TypeEngine.register(FileTransformer()).