Evaluation

Evaluation answers whether a trained model is good enough to use, and whether it is better than the one you have. It has two parts with different costs: an inference pass that runs the model over a held-out set, which needs the model’s hardware, and the metrics computed from its predictions, which run on a CPU in seconds. Flyte runs them as separate tasks, so you can change a metric without running the model again.

Separate inference from metrics

Give each part its own environment. The prediction tasks get the hardware the model needs, and the metrics task gets a CPU:

evaluation.py
                
                    
                
            
                
            

                
                    
                
            
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

image = flyte.Image.from_uv_script(__file__, name="evaluation")

# In a real pipeline, give this environment the GPU the model needs.
predict_env = flyte.TaskEnvironment(
    name="eval-predict",
    image=image,
    resources=flyte.Resources(cpu=2, memory="4Gi"),
)

metrics_env = flyte.TaskEnvironment(
    name="eval-metrics",
    image=image,
    resources=flyte.Resources(cpu=1, memory="2Gi"),
)

driver_env = flyte.TaskEnvironment(
    name="eval-driver",
    image=image,
    resources=flyte.Resources(cpu=1, memory="2Gi"),
    depends_on=[predict_env, metrics_env],
)

Run the inference pass

The prediction task runs the model over one shard of the held-out set and returns every prediction, not only a score:

evaluation.py
                
@predict_env.task(cache="auto", retries=2)
async def predict(model: File, shard: pd.DataFrame) -> pd.DataFrame:
    """Run the model over one shard of the held-out set and keep every prediction."""
    clf = joblib.load(await model.download())
    features = shard.drop(columns="target")
    return pd.DataFrame(
        {
            "label": shard["target"],
            "prediction": clf.predict(features),
            "score": clf.predict_proba(features)[:, 1],
        }
    )

Keep the predictions. The task returns the labels, the predicted classes, and the scores as a DataFrame. You can compute new metrics from it, slice it by subgroup, or inspect the examples the model got wrong, without running the model again.
cache="auto" stores the predictions under the model and the shard. Evaluating the same checkpoint again, for example from a later pipeline, reuses them.
Shards in parallel. Each shard is its own task, so a large test set finishes in the time of one shard.

For an LLM or other large model, the inference pass is a batch inference job: load the model once per task, and send it as many examples at a time as the engine can batch. See Batch inference and Stream model weights.

Compute metrics from the predictions

The metrics task reads the stored predictions and returns typed metrics:

evaluation.py
                
@dataclass
class Metrics:
    accuracy: float
    f1: float
    roc_auc: float

@metrics_env.task
async def compute_metrics(predictions: pd.DataFrame) -> Metrics:
    """Compute metrics from stored predictions. Changing a metric reruns only this task."""
    from sklearn.metrics import accuracy_score, f1_score, roc_auc_score

    y, y_hat = predictions["label"], predictions["prediction"]
    return Metrics(
        accuracy=float(accuracy_score(y, y_hat)),
        f1=float(f1_score(y, y_hat)),
        roc_auc=float(roc_auc_score(y, predictions["score"])),
    )

It has no cache, because it is cheap. To add a metric, edit this task and rerun: the predictions come from the cache, and only the metrics run again.

Compare against a baseline

A metric means little on its own. Evaluate the candidate and the model it would replace on the same held-out set, and decide from the difference:

evaluation.py
                
                    
                
            
                
            

                
                    
                
            
async def run_eval(model: File, test: pd.DataFrame, n_shards: int) -> Metrics:
    shards = [test.iloc[i::n_shards].reset_index(drop=True) for i in range(n_shards)]
    predictions = await asyncio.gather(*(predict(model, s) for s in shards))
    return await compute_metrics(pd.concat(predictions, ignore_index=True))

@driver_env.task(report=True)
async def evaluate_model(max_depth: int = 8, baseline_depth: int = 2, n_shards: int = 4) -> bool:
    train, test = await make_split()
    candidate, baseline = await asyncio.gather(train_model(train, max_depth), train_model(train, baseline_depth))

    with flyte.group("candidate"):
        new = await run_eval(candidate, test, n_shards)
    with flyte.group("baseline"):
        old = await run_eval(baseline, test, n_shards)

    rows = "".join(
        f"<tr><td>{name}</td><td>{getattr(old, name):.4f}</td><td>{getattr(new, name):.4f}</td></tr>"
        for name in ("accuracy", "f1", "roc_auc")
    )
    await flyte.report.replace.aio(
        f"<h2>Evaluation</h2><table><tr><th>Metric</th><th>Baseline</th><th>Candidate</th></tr>{rows}</table>",
        do_flush=True,
    )

    # Promote the candidate only if it beats the baseline on the metric that matters.
    return new.roc_auc >= old.roc_auc

The driver evaluates both models with the same code, publishes a side-by-side table to the run’s report, and returns whether the candidate should be promoted. With a small test set or a noisy metric, a difference can be chance. To compare with repeated samples and a margin, see Evaluation gates.

Choose an evaluation framework

The framework runs inside the metrics task, or, for LLM harnesses that call the model themselves, inside the prediction task:

Model	Framework	What it computes
Classical ML	scikit-learn metrics	Accuracy, F1, ROC AUC, calibration, regression errors
Deep learning	TorchMetrics	Metrics that accumulate over batches, including across ranks
Hugging Face models	Hugging Face Evaluate	Task metrics such as BLEU, ROUGE, and seqeval
Object detection	pycocotools or TorchMetrics `MeanAveragePrecision`	COCO mAP
LLMs on standard benchmarks	lm-evaluation-harness or lighteval	Benchmark accuracy, with a vLLM backend for throughput
LLM agents and open-ended output	Inspect or an LLM judge	Multi-turn and tool-use tasks, and rubric scores

Run each benchmark as its own task, so benchmarks run in parallel and are cached separately. When the judge is a paid model API, wrap the call in @flyte.trace, so a retried task replays the recorded result instead of paying again. See Call a judge without paying for it twice.

Evaluate checkpoints while training runs

To see how quality changes during a long run, evaluate checkpoints as they are saved. Have the training task return checkpoints at intervals, or split training into stages, and launch an evaluation for each checkpoint with asyncio.create_task while the next stage trains. The evaluations run on their own hardware and never hold up the trainer. The last checkpoint isn’t always the best, so keep each checkpoint’s metrics in the run’s report and promote the one that scores highest.

Show results in the UI

report=True on the driver lets it publish tables and charts, such as a confusion matrix or a precision-recall curve, to the run’s report. To keep results across runs, log them from the metrics task to Weights & Biases or MLflow.

On Union.ai, publish the metrics alongside the model as an artifact, so the model registry shows the evaluation that each version passed.

BERT emotion classification: a confusion matrix and attention visualizations in the report.
RT-DETR object detection: COCO mAP evaluation and side-by-side bounding boxes.
Multimodal retrieval evaluation: benchmarking three retrieval models on GPUs.
RL evaluation: how evaluating an RL policy differs.
Deploy trained models: serve the model that passed.