Skip to main content

Handling Restricted Types

Restricted types in flyte-sdk are Python types that are explicitly forbidden from being used as inputs or outputs in tasks and workflows. This mechanism ensures that only types with well-defined Flyte IDL mappings are used, preventing runtime errors during serialization or deserialization that might occur with unsupported complex types.

Restricting a Custom Type​

If you have a custom class that should not be passed between tasks—perhaps because it contains non-serializable state like network sockets or file handles—you can register it as a restricted type using the TypeEngine.

To restrict a type, call TypeEngine.register_restricted_type with a descriptive name and the Python class:

from flyte.types import TypeEngine

class LocalOnlyState:
    def __init__(self):
        self.connection = "open"

# Prevent LocalOnlyState from being used in task signatures
TypeEngine.register_restricted_type("local_state", LocalOnlyState)

When a type is registered this way, flyte-sdk internally creates a RestrictedTypeTransformer for it. Any attempt to use this type in a @task or @workflow signature will trigger a RestrictedTypeError when the TypeEngine attempts to resolve the literal type or perform a conversion.

Default Restricted Types​

By default, flyte-sdk restricts several built-in Python types that do not have a direct mapping to the Flyte IDL or whose behavior in a distributed environment would be ambiguous. These are registered at the end of types/_type_engine.py:

TypeEngine.register_restricted_type("non typed tuple", tuple)
TypeEngine.register_restricted_type("non typed tuple", typing.Tuple)
TypeEngine.register_restricted_type("named tuple", NamedTuple)

While NamedTuple is restricted from being used as a nested type within other structures, flyte-sdk allows it as a top-level return type for tasks to support multiple output values. However, using a raw tuple or a generic typing.Tuple as an input or output will result in a RestrictedTypeError.

Internal Implementation​

The restriction mechanism is implemented through two primary components in types/_type_engine.py:

  1. RestrictedTypeTransformer: A specialized TypeTransformer that overrides to_literal, to_python_value, and get_literal_type to always raise a RestrictedTypeError.
  2. TypeEngine.register_restricted_type: A convenience method that appends the type to the internal _RESTRICTED_TYPES list and registers the RestrictedTypeTransformer in the TypeEngine registry.

When the TypeEngine processes a task signature, it calls get_transformer(python_type). If the type is restricted, it retrieves the RestrictedTypeTransformer. Any subsequent call to to_literal_type (to generate the Flyte IDL) or to_literal (to serialize data) will fail:

# Internal behavior of RestrictedTypeTransformer
def get_literal_type(self, t: Optional[Type[T]] = None) -> LiteralType:
    raise RestrictedTypeError(f"Transformer for type {self.python_type} is restricted currently")

Handling Restricted Type Errors​

If you encounter a RestrictedTypeError, it typically means you are trying to use a type that flyte-sdk cannot safely serialize.

from flyte.core.task import task
from typing import Tuple

@task
def my_task(x: int) -> Tuple[int, int]:  # This will raise RestrictedTypeError
    return x, x

To resolve this for tuples, you should use a NamedTuple for the return type or specific Flyte types like FlyteFile or FlyteDirectory for complex IO, rather than raw Python objects that lack a defined schema in the Flyte ecosystem.