Flyte Sdk Eval

Flyte 2 SDK Eval Skill

Build evaluation harnesses, unit tests, and validation pipelines for Flyte 2 workflows.

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 |

Testing Patterns

Direct Task Invocation (unit testing)

Test task logic directly without remote execution:

import pytest
from pipeline import preprocess, train, evaluate

def test_preprocess():
    """Test preprocessing logic in isolation."""
    result = preprocess(["a", "b", "c"])
    assert result is not None
    assert len(result) == 3

def test_train():
    """Test training with a small dataset."""
    import flyte
    import flyte.io
    data = flyte.io.DataFrame(polars.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]}))
    model = train(data)
    assert model is not None

def test_evaluate():
    """Test evaluation metrics."""
    import flyte
    model = flyte.io.File(path="/tmp/mock_model.pt")
    metrics = evaluate(model)
    assert "accuracy" in metrics
    assert 0  flyte.io.DataFrame:
    """Load ground truth test data."""
    ...
    return flyte.io.DataFrame(df)

@env.task
async def load_model(model_uri: str) -> object:
    """Load a trained model."""
    ...

@env.task
async def predict(model: object, data: flyte.io.DataFrame) -> flyte.io.DataFrame:
    """Run model predictions on test data."""
    ...

@env.task
async def compute_metrics(
    predictions: flyte.io.DataFrame,
    ground_truth: flyte.io.DataFrame,
) -> dict:
    """Compute evaluation metrics."""
    from sklearn.metrics import (
        accuracy_score, f1_score, precision_score, recall_score,
        roc_auc_score, mean_squared_error,
    )
    y_true = ground_truth.to_polars()["label"].to_list()
    y_pred = predictions.to_polars()["prediction"].to_list()
    y_prob = predictions.to_polars()["probability"].to_list()

    return {
        "accuracy": accuracy_score(y_true, y_pred),
        "f1": f1_score(y_true, y_pred),
        "precision": precision_score(y_true, y_pred),
        "recall": recall_score(y_true, y_pred),
        "auc": roc_auc_score(y_true, y_prob),
    }

@env.task
async def generate_report(metrics: dict) -> flyte.io.File:
    """Generate an evaluation report."""
    import matplotlib.pyplot as plt
    fig, axes = plt.subplots(1, 2, figsize=(12, 4))
    # Confusion matrix, ROC curve, etc.
    path = "/tmp/eval_report.png"
    fig.savefig(path)
    return flyte.io.File(path=path)

@env.task
async def evaluate_pipeline(
    model_uri: str,
    test_data_uri: str,
) -> dict:
    """Full evaluation pipeline."""
    data = await load_test_data()
    model = await load_model(model_uri)
    preds = await predict(model, data)
    metrics = await compute_metrics(preds, data)
    report = await generate_report(metrics)
    return {"metrics": metrics, "report": report}

A/B Model Comparison

@env.task
async def compare_models(
    model_a_uri: str,
    model_b_uri: str,
    test_data: flyte.io.DataFrame,
) -> dict:
    """Compare two models on the same test data."""
    model_a = await load_model(model_a_uri)
    model_b = await load_model(model_b_uri)
    preds_a = await predict(model_a, test_data)
    preds_b = await predict(model_b, test_data)
    metrics_a = await compute_metrics(preds_a, test_data)
    metrics_b = await compute_metrics(preds_b, test_data)

    winner = "A" if metrics_a["accuracy"] > metrics_b["accuracy"] else "B"
    return {
        "model_a_metrics": metrics_a,
        "model_b_metrics": metrics_b,
        "winner": winner,
        "improvement": metrics_a["accuracy"] - metrics_b["accuracy"],
    }

Data Quality Validation

@env.task
async def validate_data(df: flyte.io.DataFrame) -> dict:
    """Run data quality checks."""
    inner = df.to_polars()
    checks = {}

    # Row count check
    row_count = len(inner)
    checks["row_count"] = row_count
    if row_count == 0:
        raise ValueError("Dataset is empty")

    # Null check
    null_counts = inner.null_count().to_dict()
    checks["null_counts"] = null_counts
    for col, count in null_counts.items():
        if count > 0 and count / row_count > 0.5:
            raise ValueError(f"Column {col} has >50% nulls")

    # Type check
    checks["dtypes"] = {str(k): str(v) for k, v in inner.schema.items()}

    # Value range check
    for col in inner.columns:
        if inner[col].dtype.is_float64():
            min_val = inner[col].min()
            max_val = inner[col].max()
            if min_val  1:
                checks[f"range_{col}"] = {"min": min_val, "max": max_val}

    return {"passed": True, "checks": checks}

@env.task
async def data_quality_gate(
    data: flyte.io.DataFrame,
    threshold: float = 0.9,
) -> bool:
    """Pass/fail gate based on data quality score."""
    result = await validate_data(data)
    score = result["checks"].get("quality_score", 1.0)
    if score  dict:
    """Validate that pipeline outputs meet quality thresholds."""
    validation = {
        "model_exists": model_path is not None,
        "metrics_valid": all(v >= 0 and v = min_accuracy,
    }

    if not validation["accuracy_threshold"]:
        raise ValueError(
            f"Model accuracy {metrics['accuracy']} below threshold {min_accuracy}"
        )

    return validation

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,
) -> flyte.io.File:
    """Track experiment results as a JSON file."""
    record = {
        "experiment": experiment_name,
        "timestamp": datetime.datetime.now().isoformat(),
        "hyperparameters": hyperparams,
        "metrics": metrics,
    }
    path = f"/tmp/experiments/{experiment_name}.json"
    with open(path, "w") as f:
        json.dump(record, f, indent=2)
    return flyte.io.File(path=path)

@env.task
async def run_experiment(
    config: dict,
    data: flyte.io.DataFrame,
) -> dict:
    """Run a single experiment and track results."""
    model = await train(data, config)
    metrics = await evaluate(model, data)
    await track_experiment(config["name"], config, metrics)
    return metrics

Hyperparameter Search with Tracking

@env.task
async def hpo_search(
    param_grid: list[dict],
    data: flyte.io.DataFrame,
) -> dict:
    """Run hyperparameter search with experiment tracking."""
    results = await flyte.map(
        lambda cfg: run_experiment(cfg, data),
        param_grid,
    )
    best = max(results, key=lambda r: r["accuracy"])
    return best

Performance Benchmarking

Task-level Benchmarking

import time
import flyte

@env.task
async def benchmark_task(
    task_fn,
    inputs: dict,
    num_runs: int = 5,
) -> dict:
    """Benchmark a task's performance."""
    durations = []
    for _ in range(num_runs):
        start = time.time()
        await task_fn(**inputs)
        durations.append(time.time() - start)

    return {
        "mean_ms": sum(durations) / len(durations) * 1000,
        "min_ms": min(durations) * 1000,
        "max_ms": max(durations) * 1000,
        "p95_ms": sorted(durations)[int(len(durations) * 0.95)] * 1000,
    }

Throughput Testing

@env.task
async def throughput_test(
    batch_sizes: list[int],
) -> dict:
    """Test throughput at different batch sizes."""
    results = {}
    for size in batch_sizes:
        data = create_batch(size)
        start = time.time()
        await process_batch(data)
        elapsed = time.time() - start
        results[size] = {
            "throughput": size / elapsed if elapsed > 0 else 0,
            "latency_ms": elapsed * 1000 / size,
        }
    return results

Testing with Flyte MCP

Inspecting runs with MCP

If Flyte MCP tools are available, use them to read a past run's inputs and outputs for validation, list recent runs for a task to compare against, and block until a run finishes before checking its status.

pytest Configuration

# pytest.ini
[pytest]
testpaths = tests
asyncio_mode = auto
markers =
    slow: marks tests as slow (deselect with '-m "not slow"')
    integration: marks tests as integration tests

# pyproject.toml
[tool.pytest.ini_options]
testpaths = ["tests"]
asyncio_mode = "auto"
markers = [
    "slow: marks tests as slow",
    "integration: marks tests as integration tests",
]

Test Structure

tests/
  __init__.py
  test_preprocess.py     # unit tests for preprocessing
  test_train.py           # unit tests for training
  test_evaluate.py        # unit tests for evaluation
  test_integration.py     # integration tests (flyte.run)
  test_data_quality.py    # data quality validation
  conftest.py             # shared fixtures

Anti-Patterns

1. Don't test against remote runs in unit tests — use direct function invocation for unit tests. Reserve flyte.run() for integration tests.
2. Don't hardcode test data paths — use flyte.io.File(path="/tmp/test_data") with temp directories.
3. Don't skip data quality gates — always validate data before and after transformations.
4. Don't use Union-only features — avoid ReusePolicy and other Union-specific APIs.
5. Don't test ML models with random data — use representative test datasets that match production distribution.

Source & license

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

• Author: flyteorg
• Source: flyteorg/flyte-agent-plugins
• License: Apache-2.0

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