Pass Dataclasses, Pydantic Models, and Complex Types Between Tasks with Serialization

AI engineering tip of the week: Pass dataclasses, Pydantic models, and complex types between tasks with Serialization

In most AI orchestrators, passing data between steps means converting everything to JSON strings or saving to files manually. Flyte handles serialization for you. Return a dataclass from one task, accept it in the next. Flyte serializes it automatically using MessagePack under the hood, and the types show up in the UI.

Dataclasses just work

from dataclasses import dataclass
import flyte

env = flyte.TaskEnvironment(name="types_demo")

@dataclass
class TrainingResult:
    accuracy: float
    loss: float
    epochs: int
    model_path: str

@env.task
async def train(epochs: int) -> TrainingResult:
    return TrainingResult(
        accuracy=0.95,
        loss=0.05,
        epochs=epochs,
        model_path="s3://models/latest.pt",
    )

@env.task
async def evaluate(result: TrainingResult) -> str:
    if result.accuracy > 0.9:
        return f"Model passed! {result.accuracy:.1%} accuracy after {result.epochs} epochs"
    return "Model needs more training"

@env.task
async def pipeline() -> str:
    result = await train(epochs=10)
    return await evaluate(result)

No JSON serialization code. No schema definitions. Just Python dataclasses flowing between tasks.

Pydantic models work too

If you prefer Pydantic for validation, those work the same way:

from pydantic import BaseModel
from typing import Optional

class ModelConfig(BaseModel):
    name: str
    learning_rate: float
    batch_size: int
    description: Optional[str] = None

@env.task
async def create_config() -> ModelConfig:
    return ModelConfig(
        name="bert-classifier",
        learning_rate=0.001,
        batch_size=32,
        description="Fine-tuned BERT for sentiment",
    )

@env.task
async def train_with_config(config: ModelConfig) -> str:
    return f"Training {config.name} with lr={config.learning_rate}"

You get Pydantic's validation on the way in, and Flyte handles the serialization between containers.

Nest complex types freely

Dataclasses can contain lists, dicts, other dataclasses, enums, datetimes, and more:

from dataclasses import dataclass
from datetime import datetime
from enum import Enum
from typing import List, Dict

class Status(str, Enum):
    PENDING = "pending"
    COMPLETE = "complete"
    FAILED = "failed"

@dataclass
class Metric:
    name: str
    value: float

@dataclass
class ExperimentResult:
    experiment_id: str
    status: Status
    started_at: datetime
    metrics: List[Metric]
    hyperparams: Dict[str, float]
    tags: List[str]

@env.task
async def run_experiment() -> ExperimentResult:
    return ExperimentResult(
        experiment_id="exp-001",
        status=Status.COMPLETE,
        started_at=datetime.now(),
        metrics=[
            Metric(name="accuracy", value=0.95),
            Metric(name="f1", value=0.93),
        ],
        hyperparams={"lr": 0.001, "dropout": 0.1},
        tags=["production", "bert"],
    )

Lists of dataclasses, enums inside dataclasses, nested dicts with typed values. It all serializes correctly.

Union types for flexible inputs

Need a task that accepts different types? Use Python's union syntax:

@dataclass
class TextInput:
    text: str

@dataclass
class FileInput:
    path: str
    format: str

@env.task
async def process(data: TextInput | FileInput) -> str:
    if isinstance(data, TextInput):
        return f"Processing text: {data.text[:50]}"
    return f"Processing file: {data.path}"

Flyte discriminates the union type at runtime and deserializes the correct variant.

TypedDict and NamedTuple

Python's structural types also work:

from typing import TypedDict, NamedTuple

class ModelMetrics(TypedDict):
    accuracy: float
    precision: float
    recall: float

class TrainOutput(NamedTuple):
    model_path: str
    score: float

@env.task
async def compute_metrics() -> ModelMetrics:
    return ModelMetrics(accuracy=0.95, precision=0.93, recall=0.91)

@env.task
async def train() -> TrainOutput:
    return TrainOutput(model_path="/models/v1", score=0.95)

Include files in your data models

Dataclasses and Pydantic models can contain `flyte.io.File` and `flyte.io.Dir` for mixing metadata with large artifacts:

from flyte.io import File

@dataclass
class ModelArtifact:
    name: str
    version: str
    weights: File      # large file handled by Flyte's object store
    accuracy: float

@env.task
async def save_model() -> ModelArtifact:
    weights = File.new_remote()
    async with weights.open("wb") as f:
        await f.write(b"model weights here")
    return ModelArtifact(
        name="bert-v2",
        version="2.1.0",
        weights=weights,
        accuracy=0.96,
    )

The metadata (name, version, accuracy) is serialized inline. The weights file is stored in the object store and referenced by URI. Best of both worlds.

What types are supported

• Primitives: str, int, float, bool
• Date/time: datetime, timedelta
• Collections: list, dict, tuple (typed)
• Structures: dataclass, BaseModel, TypedDict, NamedTuple
• Enums: str enums, int enums
• Unions: `X | Y` or `Union[X, Y]`
• Optional: `Optional[X]`
• Files: `flyte.io.File`, `flyte.io.Dir`
• Custom types: Register your own TypeTransformer for anything else

How do you pass agent state between tasks?

Return a Pydantic model or dataclass from the task and accept it in the next one. Flyte serializes it for you, so agent state moves between containers without you writing JSON encoders or parsing strings back into objects.

This could matter more for agents than for classic pipelines. An agent step produces messy, nested state: a message history, a list of tool calls, token counts, a confidence score. If your orchestrator only passes strings, every hop becomes `json.dumps` on the way out and a fragile `dict["key"]["maybe_here"]` on the way in. One schema change and you are debugging a `KeyError` in a container you cannot see.

Because LLM structured output is already a Pydantic model, you hand Flyte the same model you gave the LLM:

from pydantic import BaseModel

class AgentStep(BaseModel):
    messages: list[dict]
    tool_calls: list[str]
    tokens_used: int
    confidence: float

@env.task
async def run_agent(prompt: str) -> AgentStep:
    ...  # call your LLM, return the parsed model

@env.task
async def review(step: AgentStep) -> str:
    if step.confidence < 0.7:
        return "escalate to human"
    return "approved"

What you get from typed serialization in an agent loop:

• Validation at every hop, so a malformed LLM response fails at the task boundary instead of three steps later
• Readable inputs and outputs in the UI, so you can see the exact state an agent step received when it went off the rails
• Safe schema evolution, since adding a field to the model updates every task that passes it
• Caching and retries that work on real objects, not on strings you have to re-parse

Full type system docs: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-programming/dataclasses-and-structures/

See what's happening in the Flyte Community:

Latest from the blog

• Flyte 2 Is Generally Available: The Durable, Open-Source AI Runtime - Read on Union.ai
• Why Untrusted Kernel Evaluation Needs Process Isolation (and How We Built It) - Read on Union
• A Memory Store Built on Flyte and Cognee - Read on Union.ai
• Building Grounded Agents on Fresh Web Data - Read on Union.ai
• From DNA to 3D Fold: Compare a Gene Across Six Species with Carbon and ESMFold - Read on Union.ai
• Run Models, Agents and Apps on Infrastructure You Own - Read on union.ai
• Agents That Survive Production: Rebuilding 21 Design Patterns on Flyte - Read on union.ai
• Introducing Queues and Cluster Controls: Durable Workloads Under Contention - Read on union.ai

Recent talks & recordings

• When the Pipeline Breaks: Building ML Infrastructure for Biotech R&D | Session 1 - Watch on YouTube
• Building Code Mode Agents - Watch on YouTube
• LLM fine-tuning with GRPO - Watch on YouTube
• LLM fine-tuning with LoRA & QLoRA - Watch on YouTube

Upcoming events

• Aug 18th: Flyte 2: The Durable Runtime Built for AI - RSVP on Luma
• Aug 20th: Seattle RAG & Agent Context with Vector Stores | AI Hacknight - RSVP on Luma

Releases & updates

• Flyte 2 Is Generally Available: The Durable, Open-Source AI Runtime - Read on Union.ai

<div class="button-group is-center"><a class="button" target="_blank" rel="noopener noreferrer" href="https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox/">Download Devbox</a></div>

From the community

• Open-Source Music Generation: Text-to-Music & Lyrics-to-Song - RSVP on Luma
• AI Book Club: Build a Reasoning Model (From Scratch) - RSVP on Luma

That's all for this week! - Sage Elliott