flyte-sdk-ml

Flyte 2 SDK ML Skill

Build ML training, HPO, evaluation, and inference pipelines with Flyte 2.

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.)

Model Training

PyTorch Training

import flyte
import flyte.io

env = flyte.TaskEnvironment(
    name="training",
    image=flyte.Image.from_base("pytorch/pytorch:2.1-cuda12.1-cudnn8-devel").with_pip_packages(
        "transformers", "datasets", "accelerate",
    ),
)

@env.task(
    requests=flyte.Resources(
        cpu="4", memory="16Gi", gpu="1", gpu_model="nvidia-a10g",
    ),
)
async def train(
    train_data: flyte.io.File,
    val_data: flyte.io.File,
    hyperparams: dict,
) -> flyte.io.File:
    """Train a model and save checkpoint."""
    import torch
    from transformers import AutoModelForSequenceClassification, AutoTokenizer

    # Load data
    tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
    model = AutoModelForSequenceClassification.from_pretrained(
        "bert-base-uncased", num_labels=2
    )

    # Train
    for epoch in range(hyperparams["epochs"]):
        # ... training loop ...
        pass

    # Save checkpoint
    output_path = "/tmp/model_checkpoint"
    model.save_pretrained(output_path)
    tokenizer.save_pretrained(output_path)
    return flyte.io.File(path=output_path)

@env.task
async def main(
    train_uri: str,
    val_uri: str,
    lr: float = 0.001,
    batch_size: int = 32,
    epochs: int = 3,
) -> dict:
    hyperparams = {"lr": lr, "batch_size": batch_size, "epochs": epochs}
    checkpoint = await train(
        train_data=flyte.io.File(path=train_uri),
        val_data=flyte.io.File(path=val_uri),
        hyperparams=hyperparams,
    )
    return {"checkpoint": checkpoint, "hyperparams": hyperparams}