Skip to content

Kubeflow Pipelines

Kubeflow Pipelines (KFP) lets you define ML workflows as Python functions that are compiled to a portable pipeline YAML and run on the cluster.

Core concepts

  • Component — a containerised Python function (a single step)
  • Pipeline — a DAG of components with defined inputs/outputs
  • Run — a single execution of a pipeline
  • Experiment — a named group of runs for comparison

Installation

pip install kfp kfp-kubernetes

Minimal example

from kfp import dsl
from kfp.dsl import Input, Output, Dataset, Model
import kfp

@dsl.component(
    base_image="python:3.11",
    packages_to_install=["scikit-learn", "pandas", "joblib"],
)
def train_model(
    train_data: Input[Dataset],
    output_model: Output[Model],
    n_estimators: int = 100,
):
    import pandas as pd
    from sklearn.ensemble import RandomForestClassifier
    import joblib, json

    df = pd.read_csv(train_data.path)
    X, y = df.drop("label", axis=1), df["label"]
    model = RandomForestClassifier(n_estimators=n_estimators)
    model.fit(X, y)
    joblib.dump(model, output_model.path)
    print(f"Trained with {n_estimators} estimators")


@dsl.component(
    base_image="python:3.11",
    packages_to_install=["scikit-learn", "pandas", "joblib"],
)
def evaluate_model(
    model: Input[Model],
    test_data: Input[Dataset],
    metrics: Output[dsl.Metrics],
):
    import pandas as pd
    from sklearn.metrics import accuracy_score
    import joblib

    df = pd.read_csv(test_data.path)
    X, y = df.drop("label", axis=1), df["label"]
    clf = joblib.load(model.path)
    acc = accuracy_score(y, clf.predict(X))
    metrics.log_metric("accuracy", acc)
    print(f"Accuracy: {acc:.4f}")


@dsl.pipeline(name="random-forest-pipeline")
def rf_pipeline(
    train_path: str,
    test_path: str,
    n_estimators: int = 100,
):
    train_ds = dsl.importer(artifact_uri=train_path, artifact_class=Dataset)
    test_ds = dsl.importer(artifact_uri=test_path, artifact_class=Dataset)

    train_step = train_model(
        train_data=train_ds.output,
        n_estimators=n_estimators,
    )
    evaluate_model(
        model=train_step.outputs["output_model"],
        test_data=test_ds.output,
    )


# Compile
kfp.compiler.Compiler().compile(rf_pipeline, "rf_pipeline.yaml")

Submitting a run

client = kfp.Client(host="https://kubeflow.pais.auckland.ac.nz")

run = client.create_run_from_pipeline_func(
    rf_pipeline,
    arguments={
        "train_path": "/vast/rg-compsci/data/train.csv",
        "test_path": "/vast/rg-compsci/data/test.csv",
        "n_estimators": 200,
    },
    experiment_name="random-forest-v1",
    run_name="n200-run",
)

GPU components

To run a component on GPU, add resource requests:

@dsl.component(base_image="nvcr.io/nvidia/pytorch:24.08-py3")
def gpu_train(...):
    ...

# In the pipeline:
step = gpu_train(...)
step.set_accelerator_type("GPU")
step.set_gpu_limit(1)

See also