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}

scikit-learn Training

import flyte
import flyte.io

env = flyte.TaskEnvironment(
    name="sklearn-training",
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "scikit-learn", "pandas", "polars", "joblib",
    ),
)

@env.task
async def train_sklearn(
    train_data: flyte.io.DataFrame,
    val_data: flyte.io.DataFrame,
    model_type: str = "random_forest",
) -> flyte.io.File:
    """Train a scikit-learn model."""
    from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
    from sklearn.linear_model import LogisticRegression
    import joblib

    X_train = train_data.to_polars().drop("label").to_numpy()
    y_train = train_data.to_polars()["label"].to_numpy()
    X_val = val_data.to_polars().drop("label").to_numpy()
    y_val = val_data.to_polars()["label"].to_numpy()

    if model_type == "random_forest":
        model = RandomForestClassifier(n_estimators=100)
    elif model_type == "gbm":
        model = GradientBoostingClassifier(n_estimators=100)
    else:
        model = LogisticRegression()

    model.fit(X_train, y_train)
    accuracy = model.score(X_val, y_val)

    path = f"/tmp/{model_type}_model.joblib"
    joblib.dump(model, path)
    return flyte.io.File(path=path)

HuggingFace Trainer

import flyte
import flyte.io

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

@env.task(
    requests=flyte.Resources(
        cpu="4", memory="16Gi", gpu="1", gpu_model="nvidia-a10g",
    ),
)
async def train_hf(
    dataset_name: str,
    model_name: str,
    hyperparams: dict,
) -> flyte.io.File:
    """Train with HuggingFace Trainer."""
    from datasets import load_dataset
    from transformers import (
        AutoModelForSequenceClassification,
        AutoTokenizer,
        Trainer,
        TrainingArguments,
    )

    train_dataset = load_dataset(dataset_name, split="train")
    val_dataset = load_dataset(dataset_name, split="validation")

    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSequenceClassification.from_pretrained(
        model_name, num_labels=2
    )

    def tokenize(examples):
        return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=512)

    train_dataset = train_dataset.map(tokenize)
    val_dataset = val_dataset.map(tokenize)

    training_args = TrainingArguments(
        output_dir="/tmp/training_output",
        learning_rate=hyperparams.get("lr", 2e-5),
        per_device_train_batch_size=hyperparams.get("batch_size", 16),
        num_train_epochs=hyperparams.get("epochs", 3),
        evaluation_strategy="epoch",
        save_strategy="epoch",
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=val_dataset,
    )

    trainer.train()
    trainer.save_model("/tmp/final_model")
    tokenizer.save_pretrained("/tmp/final_model")

    return flyte.io.File(path="/tmp/final_model")

Hyperparameter Optimization

Manual HPO with fan-out

import flyte

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

@env.task(
    requests=flyte.Resources(
        cpu="4", memory="16Gi", gpu="1", gpu_model="nvidia-a10g",
    ),
)
async def train_trial(hyperparams: dict) -> dict:
    """Run a single hyperparameter trial."""
    # hyperparams = {"model": "bert-base", "lr": 2e-5, "batch_size": 16, "epochs": 3}
    checkpoint = await train_hf(
        dataset_name="glue/mnli",
        model_name=hyperparams["model"],
        hyperparams=hyperparams,
    )
    # Evaluate
    metrics = await evaluate(checkpoint, "glue/mnli", split="validation")
    return {
        "hyperparams": hyperparams,
        "accuracy": metrics["accuracy"],
        "checkpoint": checkpoint,
    }

@env.task
async def hpo_search(
    param_grid: list[dict],
) -> dict:
    """Run hyperparameter search with parallel trials."""
    # Fan out all trials in parallel
    results = await flyte.map(train_trial, param_grid)
    best = max(results, key=lambda r: r["accuracy"])
    return best

@env.task
async def main() -> dict:
    param_grid = [
        {"model": "bert-base", "lr": 1e-5, "batch_size": 16, "epochs": 3},
        {"model": "bert-base", "lr": 2e-5, "batch_size": 16, "epochs": 3},
        {"model": "bert-base", "lr": 5e-5, "batch_size": 16, "epochs": 3},
        {"model": "bert-base", "lr": 2e-5, "batch_size": 32, "epochs": 3},
    ]
    return await hpo_search(param_grid)

Grid search pattern

from itertools import product

@env.task
async def grid_search() -> dict:
    """Grid search over hyperparameter combinations."""
    lr_values = [1e-5, 2e-5, 5e-5]
    batch_sizes = [16, 32]
    epochs = [2, 3]

    param_grid = [
        {"model": "bert-base", "lr": lr, "batch_size": bs, "epochs": ep}
        for lr, bs, ep in product(lr_values, batch_sizes, epochs)
    ]

    results = await flyte.map(train_trial, param_grid)
    best = max(results, key=lambda r: r["accuracy"])
    return best

Experiment Tracking

Manual experiment tracking

import json
import datetime
import flyte
import flyte.io

@env.task
async def track_experiment(
    experiment_name: str,
    hyperparams: dict,
    metrics: dict,
    checkpoint: flyte.io.File,
) -> flyte.io.File:
    """Track experiment results as a JSON file in remote storage."""
    record = {
        "experiment": experiment_name,
        "timestamp": datetime.datetime.now().isoformat(),
        "hyperparameters": hyperparams,
        "metrics": metrics,
        "checkpoint_uri": checkpoint.path,
    }
    path = f"/tmp/experiments/{experiment_name}_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
    with open(path, "w") as f:
        json.dump(record, f, indent=2)
    return flyte.io.File(path=path)

@env.task
async def compare_experiments(
    experiment_names: list[str],
) -> dict:
    """Compare multiple experiments."""
    reports = []
    for name in experiment_names:
        report = await load_experiment(name)
        reports.append(report)

    # Find best by metric
    best = max(reports, key=lambda r: r["metrics"].get("accuracy", 0))
    return {"best_experiment": best, "all": reports}

Inference result tracking

@env.task
async def track_inference(
    model_uri: str,
    test_data: flyte.io.File,
    metrics: dict,
) -> flyte.io.File:
    """Track inference results."""
    record = {
        "model_uri": model_uri,
        "test_data": test_data.path,
        "metrics": metrics,
        "timestamp": datetime.datetime.now().isoformat(),
    }
    path = f"/tmp/inference/{model_uri.split('/')[-1]}_{datetime.datetime.now().strftime('%Y%m%d')}.json"
    with open(path, "w") as f:
        json.dump(record, f, indent=2)
    return flyte.io.File(path=path)

Model Evaluation and Selection

Evaluation pipeline

import flyte
import flyte.io

env = flyte.TaskEnvironment(
    name="evaluation",
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "scikit-learn", "scipy", "pandas", "matplotlib", "seaborn",
    ),
)

@env.task
async def evaluate_model(
    model_path: flyte.io.File,
    test_data: flyte.io.DataFrame,
) -> dict:
    """Evaluate a model and return metrics."""
    import joblib
    from sklearn.metrics import (
        accuracy_score, f1_score, precision_score, recall_score,
        roc_auc_score, confusion_matrix, classification_report,
    )

    model = joblib.load(model_path.path)
    X_test = test_data.to_polars().drop("label").to_numpy()
    y_test = test_data.to_polars()["label"].to_numpy()

    y_pred = model.predict(X_test)
    y_prob = model.predict_proba(X_test)[:, 1] if hasattr(model, "predict_proba") else y_pred

    return {
        "accuracy": accuracy_score(y_test, y_pred),
        "f1": f1_score(y_test, y_pred),
        "precision": precision_score(y_test, y_pred),
        "recall": recall_score(y_test, y_pred),
        "auc": roc_auc_score(y_test, y_prob),
        "confusion_matrix": confusion_matrix(y_test, y_pred).tolist(),
        "report": classification_report(y_test, y_pred, output_dict=True),
    }

@env.task
async def select_best_model(
    candidate_models: list[flyte.io.File],
    test_data: flyte.io.DataFrame,
) -> dict:
    """Evaluate all candidates and select the best."""
    evaluations = await flyte.map(
        lambda m: evaluate_model(m, test_data),
        candidate_models,
    )
    best = max(evaluations, key=lambda e: e["accuracy"])
    return {"best_metrics": best, "all_evaluations": evaluations}

Model comparison report

@env.task
async def generate_comparison_report(
    evaluations: list[dict],
    model_names: list[str],
) -> flyte.io.File:
    """Generate a model comparison report."""
    import matplotlib.pyplot as plt
    import pandas as pd

    df = pd.DataFrame({
        "model": model_names,
        "accuracy": [e["accuracy"] for e in evaluations],
        "f1": [e["f1"] for e in evaluations],
        "precision": [e["precision"] for e in evaluations],
        "recall": [e["recall"] for e in evaluations],
        "auc": [e["auc"] for e in evaluations],
    })

    fig, axes = plt.subplots(1, 3, figsize=(18, 5))
    metrics = ["accuracy", "f1", "precision", "recall", "auc"]
    for i, metric in enumerate(metrics[:3]):
        axes[i].bar(df["model"], df[metric])
        axes[i].set_title(metric)
        axes[i].tick_params(axis="x", rotation=45)

    path = "/tmp/model_comparison.png"
    fig.savefig(path, bbox_inches="tight")
    return flyte.io.File(path=path)

Batch Inference

Large-scale batch inference

import flyte
import flyte.io

env = flyte.TaskEnvironment(
    name="batch-inference",
    image=flyte.Image.from_debian_base(python_version=(3, 12)).with_pip_packages(
        "torch", "transformers", "pandas", "polars", "boto3",
    ),
)

@env.task(
    requests=flyte.Resources(
        cpu="4", memory="16Gi", gpu="1", gpu_model="nvidia-a10g",
    ),
)
async def load_model(model_uri: str) -> object:
    """Load model into memory."""
    from transformers import AutoModelForSequenceClassification, AutoTokenizer
    tokenizer = AutoTokenizer.from_pretrained(model_uri)
    model = AutoModelForSequenceClassification.from_pretrained(model_uri)
    model.eval()
    return {"model": model, "tokenizer": tokenizer}

@env.task(
    requests=flyte.Resources(
        cpu="2", memory="8Gi", gpu="1", gpu_model="nvidia-a10g",
    ),
)
async def batch_predict(
    model_ctx: object,
    data_file: flyte.io.File,
    batch_size: int = 32,
) -> flyte.io.File:
    """Run inference on a batch of data."""
    import torch
    import polars as pl

    model = model_ctx["model"]
    tokenizer = model_ctx["tokenizer"]

    df = pl.read_parquet(data_file.path)
    texts = df["text"].to_list()

    all_preds = []
    all_probs = []
    for i in range(0, len(texts), batch_size):
        batch = texts[i:i + batch_size]
        inputs = tokenizer(batch, padding=True, truncation=True, return_tensors="pt")
        with torch.no_grad():
            outputs = model(**inputs)
        probs = torch.softmax(outputs.logits, dim=1)
        preds = torch.argmax(probs, dim=1)
        all_preds.extend(preds.tolist())
        all_probs.extend(probs.tolist())

    results = pl.DataFrame({"prediction": all_preds, "probability": all_probs})
    path = f"/tmp/predictions_{data_file.path.split('/')[-1]}"
    results.write_parquet(path)
    return flyte.io.File(path=path)

@env.task
async def batch_inference(
    model_uri: str,
    data_files: list[str],
) -> list:
    """Run batch inference on multiple data files."""
    model_ctx = await load_model(model_uri)
    # Fan out inference across files
    results = await flyte.map(
        lambda f: batch_predict(model_ctx, flyte.io.File(path=f)),
        data_files,
    )
    return results

GPU batch inference optimization

@env.task
async def optimized_batch_inference(
    model_uri: str,
    data_files: list[str],
) -> list:
    """Optimized batch inference with dynamic batching."""
    # Use dynamic batcher for better GPU utilization
    # Combine small batches and shard large ones
    ...

Real-time Model Serving

FastAPI model serving (covered in flyte-sdk-app)

from fastapi import FastAPI
import flyte
from flyte.app.extras import FastAPIAppEnvironment

app = FastAPI()
model = None

@app.on_event("startup")
async def load_model():
    global model
    from transformers import AutoModelForSequenceClassification, AutoTokenizer
    model = AutoModelForSequenceClassification.from_pretrained("model-checkpoint")
    model.tokenizer = AutoTokenizer.from_pretrained("model-checkpoint")

@app.get("/predict")
async def predict(text: str) -> dict:
    inputs = model.tokenizer(text, return_tensors="pt", padding=True, truncation=True)
    with torch.no_grad():
        out

…

## Source & license

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

- **Author:** [flyteorg](https://github.com/flyteorg)
- **Source:** [flyteorg/flyte-agent-plugins](https://github.com/flyteorg/flyte-agent-plugins)
- **License:** Apache-2.0

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