ML workloads

The core patterns compose into the workloads data scientists and ML engineers run every day. Each of these is a complete v1→v2 pair in the examples repo. See Migration for the overall approach.

Small model training (scikit-learn / XGBoost)

Train a model, persist it as a File, and evaluate it. Image, resources, and caching move to the TaskEnvironment; FlyteFile becomes flyte.io.File.

Flyte 1Flyte 2
train_xgboost_v1.py 
import os

import joblib
from flytekit import task, workflow, ImageSpec, Resources, current_context
from flytekit.types.file import FlyteFile
from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split
from xgboost import XGBClassifier

image = ImageSpec(
    name="xgb-image",
    packages=["xgboost", "scikit-learn", "joblib"],
)

@task(container_image=image, requests=Resources(cpu="2", mem="4Gi"))
def train_model(n_estimators: int, max_depth: int) -> FlyteFile:
    data = load_breast_cancer()
    X_train, _, y_train, _ = train_test_split(data.data, data.target, random_state=42)
    model = XGBClassifier(n_estimators=n_estimators, max_depth=max_depth)
    model.fit(X_train, y_train)

    model_path = os.path.join(current_context().working_directory, "model.json")
    joblib.dump(model, model_path)
    return FlyteFile(path=model_path)

@task(container_image=image)
def evaluate(model_file: FlyteFile) -> float:
    model = joblib.load(model_file.download())
    data = load_breast_cancer()
    _, X_test, _, y_test = train_test_split(data.data, data.target, random_state=42)
    return float(model.score(X_test, y_test))

@workflow
def main(n_estimators: int, max_depth: int) -> float:
    model = train_model(n_estimators=n_estimators, max_depth=max_depth)
    return evaluate(model_file=model)
train_xgboost_v2.py 
import os

import joblib
import flyte
from flyte.io import File
from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split
from xgboost import XGBClassifier

env = flyte.TaskEnvironment(
    name="train_xgboost",
    image=flyte.Image.from_debian_base().with_pip_packages(
        "xgboost", "scikit-learn", "joblib"
    ),
    resources=flyte.Resources(cpu="2", memory="4Gi"),
)

@env.task
async def train_model(n_estimators: int, max_depth: int) -> File:
    data = load_breast_cancer()
    X_train, _, y_train, _ = train_test_split(data.data, data.target, random_state=42)
    model = XGBClassifier(n_estimators=n_estimators, max_depth=max_depth)
    model.fit(X_train, y_train)

    model_path = os.path.join(os.getcwd(), "model.json")
    joblib.dump(model, model_path)
    return await File.from_local(model_path)

@env.task
async def evaluate(model_file: File) -> float:
    local_path = await model_file.download()
    model = joblib.load(local_path)
    data = load_breast_cancer()
    _, X_test, _, y_test = train_test_split(data.data, data.target, random_state=42)
    return float(model.score(X_test, y_test))

@env.task
async def main(n_estimators: int, max_depth: int) -> float:
    model = await train_model(n_estimators, max_depth)
    return await evaluate(model)

Hyperparameter optimization

Fan out one training run per hyperparameter, then pick the best. In Flyte 1 the grid search runs through map_task and the “pick the best” step must itself be a task. In Flyte 2 you gather the runs and select the winner in plain Python.

Flyte 1Flyte 2
hpo_v1.py 
from flytekit import task, workflow, map_task
from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import cross_val_score

@task
def get_grid() -> list[int]:
    return [2, 4, 8, 16]

@task
def train_eval(max_depth: int) -> float:
    data = load_iris()
    model = RandomForestClassifier(max_depth=max_depth, random_state=42)
    scores = cross_val_score(model, data.data, data.target, cv=3)
    return float(scores.mean())

@task
def best_score(scores: list[float]) -> float:
    return max(scores)

@workflow
def main() -> float:
    grid = get_grid()
    # Fan out one training run per hyperparameter value.
    scores = map_task(train_eval)(max_depth=grid)
    return best_score(scores=scores)
hpo_v2.py 
import asyncio

import flyte
from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import cross_val_score

env = flyte.TaskEnvironment(
    name="hpo",
    image=flyte.Image.from_debian_base().with_pip_packages("scikit-learn"),
)

@env.task
async def train_eval(max_depth: int) -> float:
    data = load_iris()
    model = RandomForestClassifier(max_depth=max_depth, random_state=42)
    scores = cross_val_score(model, data.data, data.target, cv=3)
    return float(scores.mean())

@env.task
async def main() -> dict:
    grid = [2, 4, 8, 16]
    # Fan out one training run per hyperparameter value...
    scores = await asyncio.gather(*[train_eval(d) for d in grid])
    # ...then pick the best in plain Python (impossible in a Flyte 1 workflow).
    best_idx = max(range(len(scores)), key=lambda i: scores[i])
    return {"best_max_depth": grid[best_idx], "best_score": scores[best_idx]}

Large model training (deep learning)

GPU configuration moves to the TaskEnvironment: the Flyte 1 Resources(gpu="1") plus a separate accelerator=T4 become a single gpu="T4:1" string on flyte.Resources.

Flyte 1Flyte 2
train_deep_learning_v1.py 
from flytekit import task, workflow, ImageSpec, Resources
from flytekit.extras.accelerators import T4
import torch
import torch.nn as nn

image = ImageSpec(
    name="dl-image",
    packages=["torch"],
)

@task(
    container_image=image,
    requests=Resources(cpu="4", mem="16Gi", gpu="1"),
    accelerator=T4,
)
def train(epochs: int) -> float:
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = nn.Linear(10, 1).to(device)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
    loss_fn = nn.MSELoss()

    X = torch.randn(128, 10, device=device)
    y = torch.randn(128, 1, device=device)

    loss = torch.tensor(0.0)
    for _ in range(epochs):
        optimizer.zero_grad()
        loss = loss_fn(model(X), y)
        loss.backward()
        optimizer.step()
    return float(loss.item())

@workflow
def main(epochs: int) -> float:
    return train(epochs=epochs)
train_deep_learning_v2.py 
import flyte
import torch
import torch.nn as nn

# GPU type and count go in a single "T4:1"-style string. For multi-node
# distributed training, wrap the training task with the torch elastic plugin.
env = flyte.TaskEnvironment(
    name="train_deep_learning",
    image=flyte.Image.from_debian_base().with_pip_packages("torch"),
    resources=flyte.Resources(cpu="4", memory="16Gi", gpu="T4:1"),
)

@env.task
async def train(epochs: int) -> float:
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = nn.Linear(10, 1).to(device)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
    loss_fn = nn.MSELoss()

    X = torch.randn(128, 10, device=device)
    y = torch.randn(128, 1, device=device)

    loss = torch.tensor(0.0)
    for _ in range(epochs):
        optimizer.zero_grad()
        loss = loss_fn(model(X), y)
        loss.backward()
        optimizer.step()
    return float(loss.item())

@env.task
async def main(epochs: int) -> float:
    return await train(epochs)

For multi-node distributed training (PyTorch elastic, etc.), see Resources and the plugin integrations.

Batch inference

Load a trained model once and score many batches in parallel. map_task with a partial-bound model becomes asyncio.gather over the batches, reusing the same model reference.

Flyte 1Flyte 2
batch_inference_v1.py 
import os
from functools import partial

import joblib
from flytekit import task, workflow, map_task, ImageSpec, current_context
from flytekit.types.file import FlyteFile
from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier

image = ImageSpec(name="inference-image", packages=["scikit-learn", "joblib"])

@task(container_image=image)
def train_model() -> FlyteFile:
    data = load_iris()
    model = RandomForestClassifier().fit(data.data, data.target)
    model_path = os.path.join(current_context().working_directory, "model.joblib")
    joblib.dump(model, model_path)
    return FlyteFile(path=model_path)

@task(container_image=image)
def get_batches() -> list[list[list[float]]]:
    data = load_iris()
    rows = data.data.tolist()
    # Split the rows into batches of 30.
    return [rows[i : i + 30] for i in range(0, len(rows), 30)]

@task(container_image=image)
def score_batch(model_file: FlyteFile, batch: list[list[float]]) -> list[int]:
    model = joblib.load(model_file.download())
    return [int(p) for p in model.predict(batch)]

@workflow
def main() -> list[list[int]]:
    model = train_model()
    batches = get_batches()
    return map_task(partial(score_batch, model_file=model))(batch=batches)
batch_inference_v2.py 
import asyncio
import os

import joblib
import flyte
from flyte.io import File
from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier

env = flyte.TaskEnvironment(
    name="batch_inference",
    image=flyte.Image.from_debian_base().with_pip_packages("scikit-learn", "joblib"),
)

@env.task
async def train_model() -> File:
    data = load_iris()
    model = RandomForestClassifier().fit(data.data, data.target)
    model_path = os.path.join(os.getcwd(), "model.joblib")
    joblib.dump(model, model_path)
    return await File.from_local(model_path)

@env.task
async def score_batch(model_file: File, batch: list[list[float]]) -> list[int]:
    local_path = await model_file.download()
    model = joblib.load(local_path)
    return [int(p) for p in model.predict(batch)]

@env.task
async def main() -> list[list[int]]:
    model = await train_model()
    rows = load_iris().data.tolist()
    batches = [rows[i : i + 30] for i in range(0, len(rows), 30)]
    # Score every batch in parallel, reusing the same model reference.
    coros = [score_batch(model, batch) for batch in batches]
    return list(await asyncio.gather(*coros))

A complete example: end-to-end ML pipeline

Putting it together — a load / train / evaluate pipeline shows the image, resources, caching, file I/O, and orchestration changes in one place.

Flyte 1Flyte 2
ml_pipeline_v1.py 
import os

import joblib
import pandas as pd
from flytekit import task, workflow, ImageSpec, Resources, current_context
from flytekit.types.file import FlyteFile
from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier

image = ImageSpec(
    name="ml-image",
    packages=["pandas", "scikit-learn", "joblib"],
)

@task(
    container_image=image,
    requests=Resources(cpu="2", mem="4Gi"),
    cache=True,
    cache_version="1.0",
)
def load_data() -> pd.DataFrame:
    data = load_iris(as_frame=True)
    df = data.frame
    df["species"] = data.target
    return df

@task(container_image=image)
def train_model(data: pd.DataFrame) -> FlyteFile:
    model = RandomForestClassifier()
    X = data.drop("species", axis=1)
    y = data["species"]
    model.fit(X, y)

    model_path = os.path.join(current_context().working_directory, "model.joblib")
    joblib.dump(model, model_path)
    return FlyteFile(path=model_path)

@task(container_image=image)
def evaluate(model_file: FlyteFile, data: pd.DataFrame) -> float:
    model = joblib.load(model_file.download())
    X = data.drop("species", axis=1)
    y = data["species"]
    return float(model.score(X, y))

@workflow
def main() -> float:
    data = load_data()
    model = train_model(data=data)
    return evaluate(model_file=model, data=data)
ml_pipeline_v2.py 
import os

import joblib
import pandas as pd
import flyte
from flyte.io import File
from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier

# Image, resources, and cache are set once on the TaskEnvironment.
env = flyte.TaskEnvironment(
    name="ml_pipeline",
    image=flyte.Image.from_debian_base().with_pip_packages(
        "pandas", "scikit-learn", "joblib"
    ),
    resources=flyte.Resources(cpu="2", memory="4Gi"),
    cache="auto",
)

@env.task
async def load_data() -> pd.DataFrame:
    data = load_iris(as_frame=True)
    df = data.frame
    df["species"] = data.target
    return df

@env.task
async def train_model(data: pd.DataFrame) -> File:
    model = RandomForestClassifier()
    X = data.drop("species", axis=1)
    y = data["species"]
    model.fit(X, y)

    model_path = os.path.join(os.getcwd(), "model.joblib")
    joblib.dump(model, model_path)
    return await File.from_local(model_path)

@env.task
async def evaluate(model_file: File, data: pd.DataFrame) -> float:
    local_path = await model_file.download()
    model = joblib.load(local_path)
    X = data.drop("species", axis=1)
    y = data["species"]
    return float(model.score(X, y))

# The "workflow" is just an orchestrating task.
@env.task
async def main() -> float:
    data = await load_data()
    model = await train_model(data)
    return await evaluate(model, data)

Next