Classical machine learning
Classical models such as gradient-boosted trees, linear models, and random forests train on CPUs in seconds to minutes. The cost is in how many times you train them: one fit per hyperparameter setting, per fold, per feature set. Flyte turns each fit into a task, runs them all at once, and caches every result, so a sweep that took an afternoon on a laptop finishes, given enough cluster capacity, in the time of its slowest trial.
Define the environments
Trials are small and CPU-only, so each gets two CPUs. The driver gets less, because it only launches trials and compares scores:
import asyncio
import tempfile
from dataclasses import dataclass
from pathlib import Path
import flyte
import joblib
import pandas as pd
from flyte.io import File
# The image installs the dependencies listed in this file's script header.
image = flyte.Image.from_uv_script(__file__, name="classical-ml")
env = flyte.TaskEnvironment(
name="classical-ml",
image=image,
resources=flyte.Resources(cpu=2, memory="4Gi"),
)
driver_env = flyte.TaskEnvironment(
name="classical-ml-driver",
image=image,
resources=flyte.Resources(cpu=1, memory="2Gi"),
depends_on=[env],
)
Image.from_uv_script builds the image from the dependencies in the script’s header. depends_on=[env] deploys the trial environment with the driver.
Cache the data
Load and clean the data in its own task, with cache="auto":
@env.task(cache="auto")
async def load_data() -> pd.DataFrame:
"""Load and clean the training data once. Later runs reuse the cached output."""
from sklearn.datasets import load_breast_cancer
data = load_breast_cancer(as_frame=True)
return data.frame.dropna()
The task returns a pandas.DataFrame, which Flyte stores as Parquet in your object store. Every trial receives a reference to the same stored data, and later runs skip the load entirely. See DataFrames.
Score one setting per task
Each hyperparameter setting is one call to an evaluation task:
@dataclass
class Params:
n_estimators: int
max_depth: int
@dataclass
class TrialResult:
params: Params
score: float
@env.task(cache="auto", retries=2)
async def evaluate(df: pd.DataFrame, params: Params) -> TrialResult:
"""Score one hyperparameter setting with cross-validation."""
from sklearn.ensemble import GradientBoostingClassifier
from sklearn.model_selection import cross_val_score
X, y = df.drop(columns="target"), df["target"]
model = GradientBoostingClassifier(
n_estimators=params.n_estimators, max_depth=params.max_depth, random_state=0
)
score = cross_val_score(model, X, y, cv=5).mean()
return TrialResult(params=params, score=float(score))
- Typed inputs and outputs.
ParamsandTrialResultare dataclasses, so the UI shows each trial’s settings and score, and the driver gets real objects back. See Data classes and structures. - Caching per trial. With
cache="auto", rerunning a sweep with one new value only runs the new trials. - Retries per trial. One trial that fails on a bad node retries on its own instead of failing the sweep.
Run the sweep in parallel
The final fit is one more task, which trains on all of the data and saves the model as a file:
@env.task
async def fit(df: pd.DataFrame, params: Params) -> File:
"""Fit the final model on all of the data and save it as a file."""
from sklearn.ensemble import GradientBoostingClassifier
X, y = df.drop(columns="target"), df["target"]
model = GradientBoostingClassifier(
n_estimators=params.n_estimators, max_depth=params.max_depth, random_state=0
).fit(X, y)
path = Path(tempfile.mkdtemp()) / "model.joblib"
joblib.dump(model, path)
return await File.from_local(str(path))
The driver builds the grid in plain Python and launches every trial at once with asyncio.gather:
@driver_env.task(report=True)
async def training_pipeline(
n_estimators: list[int] = [50, 100, 200],
max_depth: list[int] = [2, 3, 4],
) -> File:
df = await load_data()
# One task per setting, all running at once.
grid = [Params(n, d) for n in n_estimators for d in max_depth]
with flyte.group("sweep"):
results = await asyncio.gather(*(evaluate(df, p) for p in grid))
best = max(results, key=lambda r: r.score)
rows = "".join(
f"<tr><td>{r.params.n_estimators}</td><td>{r.params.max_depth}</td><td>{r.score:.4f}</td></tr>"
for r in sorted(results, key=lambda r: -r.score)
)
await flyte.report.replace.aio(
f"<h2>Sweep results</h2><table><tr><th>n_estimators</th><th>max_depth</th><th>CV accuracy</th></tr>"
f"{rows}</table>",
do_flush=True,
)
return await fit(df, best.params)
flyte.group("sweep") collects the trials under one node in the UI. The driver picks the winner with ordinary Python, publishes a results table to the run’s report, and calls fit with the best settings. The model comes back as a flyte.io.File, ready to pass to an evaluation task or a serving app. To limit how many trials run at once, see Controlling parallel execution.
Search more efficiently
A grid tries every combination. To sample settings at random, stop weak trials early, or let an Optuna study choose the next setting from earlier results, see Hyperparameter optimization.
Keep containers warm for short trials
When each trial takes seconds, starting a container for each one costs more than the trial. Reusable containers keep a pool of replicas running and send each trial to a warm one.
Related
- Fraud detection with Feast: train an XGBoost classifier on features from a feature store.
- MLE bot: an agent that designs and runs XGBoost experiments in parallel.
- Mapping over inputs:
flyte.mapfor sweeps over a list. - MLflow and Weights & Biases: log each trial to an experiment tracker.
- Deploy trained models: serve the fitted model behind a FastAPI app.