Flyte Sdk Types

Flyte 2 SDK Types Skill

Guide correct type annotations, I/O patterns, and serialization for Flyte 2 workflows.

Grounding References

| Resource | URL | |---|---| | Official docs | https://www.union.ai/docs/v2/flyte | | Docs index (LLMs) | https://www.union.ai/docs/v2/flyte/llms.txt | | SDK API reference | https://www.union.ai/docs/v2/union/api-reference/flyte-sdk/ | | CLI API reference | https://www.union.ai/docs/v2/union/api-reference/flyte-cli/ | | flyte-sdk source | https://github.com/flyteorg/flyte-sdk | | Example code | https://github.com/unionai/unionai-examples | | Flyte MCP tools | Available via the flyte-cluster and flyte-docs MCP servers |

Ground unfamiliar APIs in real examples. When unsure of a current Flyte 2 API, or for a pattern not shown below, and the flyte-docs search tools are available, search them first — by exact symbol (TaskEnvironment, flyte.io.File, map_task), since matching is literal substring, not semantic — then adapt a real example rather than inventing one, and cite the file or section you pulled it from. (Flyte 2 is not flytekit; priors are often wrong.)

Type System Overview

Flyte 2 uses Python type hints for serialization. Every task input/output must have a type annotation. Flyte's type transformer system handles conversion between Python types and remote storage (S3, GCS, etc.).

Supported Native Types

| Python Type | Flyte Type | Remote Transport | |---|---|---| | int, float, bool, str | Literal | Inline (JSON) | | list[T], dict[K, V] | Collection / Map | Inline (JSON) for small, blob for large | | flyte.io.File | Blob | Uploaded to metadata bucket | | flyte.io.Dir | Blob (directory) | Uploaded to metadata bucket | | flyte.io.DataFrame | DataFrame | Parquet in metadata bucket | | dataclass | Struct | JSON in metadata bucket | | pydantic.BaseModel | Struct | JSON in metadata bucket | | datetime, timedelta | DateTime / Duration | Inline |

flyte.io.File — Single Files

Use flyte.io.File for any single file that flows between tasks. Files are automatically uploaded to the metadata bucket at runtime.

import flyte
import flyte.io

@env.task
async def download_url(url: str) -> flyte.io.File:
    """Download a file and return as flyte.io.File."""
    import urllib.request
    local_path = f"/tmp/{url.split('/')[-1]}"
    urllib.request.urlretrieve(url, local_path)
    return flyte.io.File(path=local_path)

@env.task
async def process(file: flyte.io.File) -> dict:
    """Read a file — flyte.io.File downloads it automatically."""
    # file.path gives the local path (already downloaded)
    with open(file.path) as f:
        content = f.read()
    return {"lines": len(content.splitlines())}

@env.task
async def main(url: str) -> dict:
    downloaded = await download_url(url)
    return await process(downloaded)  # type: flyte.io.File flows as remote reference

flyte.io.File best practices

• Always pass flyte.io.File between tasks — never pass file paths as strings. Flyte serializes the reference to the remote blob.
• file.path gives the local download path inside the task container.
• Don't hardcode paths — let Flyte manage the download/upload lifecycle.
• Compression: Flyte infers format from extension (.parquet, .csv, .json, .pt, .png, etc.).

flyte.io.Dir — Directories

Use flyte.io.Dir for a collection of files (e.g., model checkpoints, output artifacts).

import flyte
import flyte.io

@env.task
async def train(checkpoint_dir: flyte.io.Dir) -> flyte.io.Dir:
    """Train and save checkpoints to a directory."""
    # Write checkpoints
    for i in range(10):
        path = f"{checkpoint_dir.path}/checkpoint_{i}.pt"
        save_model(path)
    return checkpoint_dir

@env.task
async def evaluate(checkpoints: flyte.io.Dir) -> dict:
    """Load checkpoints from a directory."""
    # List all files in the directory
    files = list(checkpoints.path.glob("*.pt"))
    ...

flyte.io.DataFrame — Polars DataFrames

Flyte 2 has built-in support for Polars DataFrames. They are passed by reference (Parquet in the metadata bucket), not inline.

import flyte
import flyte.io

@env.task
async def load_csv(url: str) -> flyte.io.DataFrame:
    """Load a CSV and return as Polars DataFrame."""
    import polars as pl
    df = pl.read_csv(url)
    return flyte.io.DataFrame(df)

@env.task
async def clean(df: flyte.io.DataFrame) -> flyte.io.DataFrame:
    """Clean the DataFrame."""
    inner = df.to_polars()  # Get the underlying Polars DataFrame
    cleaned = inner.drop_nulls()
    return flyte.io.DataFrame(cleaned)

@env.task
async def save_parquet(df: flyte.io.DataFrame, path: str) -> flyte.io.File:
    """Save DataFrame to Parquet."""
    inner = df.to_polars()
    inner.write_parquet(path)
    return flyte.io.File(path=path)

@env.task
async def main(url: str) -> flyte.io.DataFrame:
    raw = await load_csv(url)
    cleaned = await clean(raw)
    return cleaned  # Flows as Parquet reference

Polars DataFrame patterns

# Convert to Polars for manipulation
inner_df = df.to_polars()

# Convert from Polars
df = flyte.io.DataFrame(inner_df)

# Common operations
df = flyte.io.DataFrame(inner_df.filter(pl.col("age") > 18))
df = flyte.io.DataFrame(inner_df.group_by("category").agg(pl.col("value").mean()))

# Check shape
print(df.shape)  # (rows, cols)
print(df.schema)  # column names and types

Eager vs Lazy DataFrames

# Eager (loaded into memory)
df = flyte.io.DataFrame(inner_df)

# Lazy (streaming, for large datasets)
df = flyte.io.DataFrame(inner_df.lazy())

# Materialize lazy to eager
eager = df.to_polars()  # materializes

Dataclass and Pydantic Models

For structured data, use Python dataclasses or Pydantic models. Flyte serializes them to JSON.

from dataclasses import dataclass
from pydantic import BaseModel
import flyte
import flyte.io

@dataclass
class TrainingConfig:
    learning_rate: float
    batch_size: int
    epochs: int

class PredictionOutput(BaseModel):
    predictions: list[float]
    confidence: list[float]
    model_version: str

@env.task
async def train(config: TrainingConfig) -> flyte.io.File:
    # config.learning_rate, config.batch_size, etc.
    ...

@env.task
async def predict(model: flyte.io.File, data: flyte.io.DataFrame) -> PredictionOutput:
    return PredictionOutput(
        predictions=[0.5, 0.8, 0.3],
        confidence=[0.9, 0.7, 0.95],
        model_version="v1.0",
    )

Custom Type Transformers

Extend Flyte's type system to support custom types (e.g., PIL Images, HuggingFace datasets).

PIL Image transformer

from PIL import Image
import flyte
import flyte.io
from flyte.types import TypeTransformer

class PILImageTransformer(TypeTransformer[Image.Image]):
    _type = Image.Image

    def get_type(self, input: Image.Image) -> type:
        return Image.Image

    def save(self, img: Image.Image, path: str) -> None:
        img.save(path)

    def load(self, path: str) -> Image.Image:
        return Image.open(path)

# Register the transformer
flyte.types.TypeEngine.register(PILImageTransformer())

# Now use it in tasks
@env.task
async def process_image(img: Image.Image) -> flyte.io.File:
    # img is a PIL Image, already downloaded
    ...

HuggingFace Dataset transformer

from datasets import Dataset
import flyte

class HFDatasetTransformer(TypeTransformer[Dataset]):
    _type = Dataset

    def get_type(self, input: Dataset) -> type:
        return Dataset

    def save(self, ds: Dataset, path: str) -> None:
        ds.save_to_disk(path)

    def load(self, path: str) -> Dataset:
        return Dataset.load_from_disk(path)

flyte.types.TypeEngine.register(HFDatasetTransformer())

Data I/O Patterns by Domain

ETL / Data Engineering

# CSV → Parquet conversion
@env.task(cache="auto")
async def csv_to_parquet(csv_file: flyte.io.File, output_path: str) -> flyte.io.File:
    import polars as pl
    df = pl.read_csv(csv_file.path)
    df.write_parquet(output_path)
    return flyte.io.File(path=output_path)

# JsonlFile — batched JSONL reading
@env.task
async def process_jsonl(path: str) -> int:
    from flyte.extend import JsonlFile
    jf = JsonlFile(path)
    count = 0
    async for record in jf.stream():
        process(record)
        count += 1
    return count

# JsonlDir — batched JSONL directory
@env.task
async def process_jsonl_dir(dir_path: str) -> dict:
    from flyte.extend import JsonlDir
    jd = JsonlDir(dir_path)
    total = 0
    async for batch in jd.stream_batches():
        total += len(batch)
    return {"records": total}

Image Processing

from PIL import Image
import flyte
import flyte.io

@env.task
async def resize_image(input_file: flyte.io.File, size: tuple[int, int]) -> flyte.io.File:
    img = Image.open(input_file.path)
    resized = img.resize(size)
    output_path = f"/tmp/resized_{size[0]}x{size[1]}.png"
    resized.save(output_path)
    return flyte.io.File(path=output_path)

@env.task
async def batch_resize(files: list[flyte.io.File], size: tuple[int, int]) -> list:
    import asyncio
    return await asyncio.gather(*(resize_image(f, size) for f in files))

Audio Processing

import librosa
import numpy as np
import flyte
import flyte.io

@env.task
async def extract_features(audio_file: flyte.io.File) -> flyte.io.File:
    """Extract MFCC features from audio file."""
    y, sr = librosa.load(audio_file.path, sr=None)
    mfccs = librosa.feature.mfcc(y=y, sr=sr)
    # Save as numpy array
    output_path = "/tmp/mfccs.npz"
    np.savez(output_path, mfccs=mfccs, sr=sr)
    return flyte.io.File(path=output_path)

HuggingFace Datasets

from datasets import load_dataset, Dataset
import flyte
import flyte.io

@env_task
async def load_dataset_from_hub(dataset_name: str, split: str = "train") -> flyte.io.File:
    """Load a HF dataset and save locally."""
    ds = load_dataset(dataset_name, split=split)
    path = f"/tmp/{dataset_name.replace('/', '_')}_{split}"
    ds.save_to_disk(path)
    return flyte.io.File(path=path)

@env.task
async def process_dataset(file: flyte.io.File) -> flyte.io.DataFrame:
    """Convert HF dataset to Flyte DataFrame."""
    ds = Dataset.load_from_disk(file.path)
    df = ds.to_pandas()
    return flyte.io.DataFrame(df)

Data Locality and Storage Best Practices

How data flows between tasks

1. By reference (default) — large data (DataFrames, Files, Directories) is uploaded to the metadata bucket. Tasks receive a remote reference and download on demand.
2. Inline (small data) — primitives (int, float, str, bool) and small collections are passed inline as JSON.

Choosing the right transport

| Data type | Transport | Max size | |---|---|---| | int, float, bool, str | Inline | None | | list, dict (small) | Inline | ~10 MB | | flyte.io.File | Reference (S3/GCS) | Unlimited | | flyte.io.Dir | Reference (S3/GCS) | Unlimited | | flyte.io.DataFrame (Polars) | Reference (Parquet) | Unlimited | | dataclass / BaseModel | Inline (JSON) | ~10 MB |

Storage best practices

1. Use flyte.io.File for files — don't pass strings. Flyte manages the upload/download.
2. Use flyte.io.DataFrame for tabular data — stored as Parquet, efficient for ETL pipelines.
3. Use flyte.io.Dir for collections — model checkpoints, output artifacts.
4. Set cache="auto" on idempotent tasks (ETL, transforms) to avoid re-processing.
5. Use raw_data_path for per-run customization: flyte.with_runcontext(raw_data_path="s3://my-bucket/{run_id}/")
6. Large data should never be inline — if your dataclass exceeds ~10 MB, switch to flyte.io.File or flyte.io.Dir.

Inline I/O Threshold

Control when data is passed inline vs by reference:

@env.task(inline_output_limit="10MB")  # data > 10MB goes by reference
async def process(data: dict) -> dict:
    ...

Default threshold is generous. For ML training outputs or large DataFrames, set it lower to avoid overhead.

Common Type Mistakes

1. Missing type hints — Flyte 2 requires type annotations on all task inputs/outputs. No type hint = serialization error.
2. Passing flyte.io.File as a string — always use flyte.io.File(path=...) objects. The path string alone won't serialize.
3. Using Pandas instead of Polars — Flyte 2's native DataFrame is Polars. Use df.to_polars() to get the underlying DataFrame.
4. Not registering custom transformers — if you register a custom type transformer, it must be registered before task execution.
5. Forgetting .path on flyte.io.File — inside a task, file is a flyte.io.File object, not a string. Use file.path for the local path.

Source & license

This open-source skill is cataloged on AgentStack and links to its original source — we do not rehost the code.

• Author: flyteorg
• Source: flyteorg/flyte-agent-plugins
• License: Apache-2.0

Install and usage instructions live in the source repository linked above.