Flyte Migrate Data Io

Flyte 1 to 2 Migration: Data Types and I/O

Flyte 2 renames the offloaded-data types and makes their I/O async, but the mental model is the same: pass lightweight references to large data between tasks, not the materialized bytes. FlyteFile, FlyteDirectory, and StructuredDataset become flyte.io.File, flyte.io.Dir, and flyte.io.DataFrame. Plain dataclasses and Pydantic BaseModels work directly as task I/O with no JSON mixin.

Grounding References

| Resource | URL | |---|---| | Migration guide | https://www.union.ai/docs/v2/flyte/user-guide/migration/flyte-2/data-io/ | | 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/ | | Example code | https://github.com/unionai/unionai-examples | | Flyte MCP tools | Available via flyte-mcp server |

Type Mapping

| Flyte 1 | Flyte 2 | Notes | |---|---|---| | flytekit.types.file.FlyteFile | flyte.io.File | I/O is async | | flytekit.types.directory.FlyteDirectory | flyte.io.Dir | I/O is async | | flytekit.types.structured.StructuredDataset | flyte.io.DataFrame | build with from_df, read with open(...).all() | | @dataclass_json + @dataclass | plain @dataclass | no mixin needed | | Pydantic BaseModel (+ config) | plain Pydantic BaseModel | works directly as task I/O |

Offloaded Data: The Mental Model

File, Dir, and DataFrame are lightweight references (pointers) to data offloaded in blob storage — not the materialized bytes. In Flyte 2 the read/write operations are async: upload with await File.from_local(local_path), read with async with f.open("rb") as fh: await fh.read(), build a frame with flyte.io.DataFrame.from_df(df) (sync constructor), and read it with await fdf.open(pandas.DataFrame).all().

Files and Directories

FlyteFile and FlyteDirectory become flyte.io.File and flyte.io.Dir — the way you pass model artifacts and datasets between tasks. Use await File.from_local(...) to upload and async with file.open(...) to read.

Flyte 1

import os

from flytekit import task, workflow, current_context
from flytekit.types.file import FlyteFile

@task
def write_file(content: str) -> FlyteFile:
    path = os.path.join(current_context().working_directory, "out.txt")
    with open(path, "w") as f:
        f.write(content)
    return FlyteFile(path=path)

@task
def read_file(f: FlyteFile) -> str:
    with open(f.download()) as fh:
        return fh.read()

@workflow
def main(content: str) -> str:
    f = write_file(content=content)
    return read_file(f=f)

Flyte 2

import flyte
from flyte.io import File

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

@env.task
async def write_file(content: str) -> File:
    with open("out.txt", "w") as f:
        f.write(content)
    # File.from_local uploads the file to blob storage and returns a reference
    # (a lightweight pointer, not the materialized bytes).
    return await File.from_local("out.txt")

@env.task
async def read_file(f: File) -> str:
    async with f.open("rb") as fh:
        return (await fh.read()).decode("utf-8")

@env.task
async def main(content: str) -> str:
    f = await write_file(content)
    return await read_file(f)

Directories follow the same pattern: import Dir from flyte.io and use its async upload/read methods in place of FlyteDirectory. See Files and directories for more.

DataFrames

StructuredDataset becomes flyte.io.DataFrame. Construct one with flyte.io.DataFrame.from_df(df) and read it back with await df.open(pandas.DataFrame).all().

Flyte 1

import pandas as pd
from flytekit import task, workflow
from flytekit.types.structured import StructuredDataset

@task
def make_df() -> StructuredDataset:
    df = pd.DataFrame({"employee_id": [1, 2, 3], "salary": [50000, 60000, 70000]})
    return StructuredDataset(dataframe=df)

@task
def total_payroll(sd: StructuredDataset) -> float:
    df = sd.open(pd.DataFrame).all()
    return float(df["salary"].sum())

@workflow
def main() -> float:
    return total_payroll(sd=make_df())

Flyte 2

import pandas as pd
import flyte
import flyte.io

env = flyte.TaskEnvironment(
    name="dataframe",
    image=flyte.Image.from_debian_base().with_pip_packages("pandas", "pyarrow"),
)

@env.task
async def make_df() -> flyte.io.DataFrame:
    df = pd.DataFrame({"employee_id": [1, 2, 3], "salary": [50000, 60000, 70000]})
    # StructuredDataset becomes flyte.io.DataFrame.
    return flyte.io.DataFrame.from_df(df)

@env.task
async def total_payroll(fdf: flyte.io.DataFrame) -> float:
    df = await fdf.open(pd.DataFrame).all()
    return float(df["salary"].sum())

@env.task
async def main() -> float:
    return await total_payroll(await make_df())

Add the dataframe dependencies (for example pandas and pyarrow) to the TaskEnvironment image. See DataFrames for more.

Dataclasses and Structured Types

Flyte 1 required a @dataclass_json mixin for dataclass I/O. In Flyte 2, plain dataclasses (and Pydantic BaseModels) work directly as task inputs and outputs — handy for passing around a training config.

Flyte 1

from dataclasses import dataclass

from dataclasses_json import dataclass_json
from flytekit import task, workflow

@dataclass_json
@dataclass
class TrainingConfig:
    learning_rate: float
    n_estimators: int
    max_depth: int = 6

@task
def make_config(learning_rate: float, n_estimators: int) -> TrainingConfig:
    return TrainingConfig(learning_rate=learning_rate, n_estimators=n_estimators)

@task
def train(config: TrainingConfig) -> str:
    return (
        f"trained with lr={config.learning_rate}, "
        f"n_estimators={config.n_estimators}, max_depth={config.max_depth}"
    )

@workflow
def main(learning_rate: float, n_estimators: int) -> str:
    config = make_config(learning_rate=learning_rate, n_estimators=n_estimators)
    return train(config=config)

Flyte 2

from dataclasses import dataclass

import flyte

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

# Plain dataclasses work directly as task I/O -- no @dataclass_json mixin needed.
# Pydantic BaseModels work the same way.
@dataclass
class TrainingConfig:
    learning_rate: float
    n_estimators: int
    max_depth: int = 6

@env.task
def make_config(learning_rate: float, n_estimators: int) -> TrainingConfig:
    return TrainingConfig(learning_rate=learning_rate, n_estimators=n_estimators)

@env.task
def train(config: TrainingConfig) -> str:
    return (
        f"trained with lr={config.learning_rate}, "
        f"n_estimators={config.n_estimators}, max_depth={config.max_depth}"
    )

@env.task
def main(learning_rate: float, n_estimators: int) -> str:
    config = make_config(learning_rate, n_estimators)
    return train(config)

Data ETL: Putting It Together

Extract, clean, aggregate, and write out a feature table. StructuredDataset becomes flyte.io.DataFrame, and the tasks become async.

Flyte 1

import pandas as pd
from flytekit import task, workflow
from flytekit.types.structured import StructuredDataset

@task
def extract() -> pd.DataFrame:
    # Read raw transaction records (stand-in for a real source).
    return pd.DataFrame(
        {
            "user_id": [1, 1, 2, 3, 3, 3],
            "amount": [10.0, 5.0, 20.0, 7.5, 2.5, 1.0],
        }
    )

@task
def transform(df: pd.DataFrame) -> StructuredDataset:
    # Clean and aggregate into a per-user feature table.
    df = df[df["amount"] > 0]
    agg = df.groupby("user_id", as_index=False)["amount"].sum()
    return StructuredDataset(dataframe=agg)

@task
def load(sd: StructuredDataset) -> int:
    df = sd.open(pd.DataFrame).all()
    return len(df)

@workflow
def main() -> int:
    raw = extract()
    features = transform(df=raw)
    return load(sd=features)

Flyte 2

import pandas as pd
import flyte
import flyte.io

env = flyte.TaskEnvironment(
    name="data_etl",
    image=flyte.Image.from_debian_base().with_pip_packages("pandas", "pyarrow"),
)

@env.task
async def extract() -> pd.DataFrame:
    # Read raw transaction records (stand-in for a real source).
    return pd.DataFrame(
        {
            "user_id": [1, 1, 2, 3, 3, 3],
            "amount": [10.0, 5.0, 20.0, 7.5, 2.5, 1.0],
        }
    )

@env.task
async def transform(df: pd.DataFrame) -> flyte.io.DataFrame:
    # Clean and aggregate into a per-user feature table.
    df = df[df["amount"] > 0]
    agg = df.groupby("user_id", as_index=False)["amount"].sum()
    # StructuredDataset becomes flyte.io.DataFrame.
    return flyte.io.DataFrame.from_df(agg)

@env.task
async def load(sd: flyte.io.DataFrame) -> int:
    df = await sd.open(pd.DataFrame).all()
    return len(df)

@env.task
async def main() -> int:
    raw = await extract()
    features = await transform(raw)
    return await load(features)

Anti-Patterns

1. Don't call the offloaded-data I/O synchronously — File.from_local, file.open(...).read(), and DataFrame.open(...).all() are async in Flyte 2; await them inside async tasks.
2. Don't keep the @dataclass_json mixin — plain @dataclass and Pydantic BaseModels serialize as task I/O directly; drop dataclasses_json.
3. Don't return StructuredDataset(dataframe=df) — use flyte.io.DataFrame.from_df(df) instead.
4. Don't materialize large data into task outputs — return File, Dir, or DataFrame references, not the raw bytes or full frames.
5. Don't forget the dataframe dependencies — add pandas and pyarrow (or your engine) to the TaskEnvironment image so DataFrame I/O works remotely.
6. Don't import from flytekit.types.* — import File and Dir from flyte.io, and use flyte.io.DataFrame.

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.