Speed up Kubeflow Pipelines with Component Caching

Problem: redundant work slows iterative ML experiments

When you tweak a model or adjust hyper‑parameters, Kubeflow Pipelines re‑executes every step from scratch, even if the upstream data‑preparation or feature‑engineering logic hasn’t changed. This wastes compute time and delays feedback loops.

Thesis: component caching eliminates unnecessary recomputation

By storing the outputs of a step keyed to a hash of its inputs, Kubeflow Pipelines can skip execution when those inputs are identical, surfacing a “cache hit” in the UI and re‑using the persisted artifact.

How caching works in Kubeflow Pipelines

Each component’s inputs are serialized and hashed. The resulting key indexes a record in ML Metadata (MLMD) that points to the output artifact stored in the underlying artifact store (e.g., MinIO or S3). If a subsequent run finds a matching key, the engine marks the step as a cache hit, skips container execution, and returns the cached artifact.

Enabling caching per component

You can turn caching on or off with the enable_caching flag (or the @component decorator’s cache=True argument). The flag is evaluated at pipeline compile time, so changing it requires a re‑compile.

from kfp import dsl

@dsl.component(
    base_image='python:3.9',
    enable_caching=True  # <-- toggle caching for this step
)
def prepare_data(input_csv: str, output_parquet: dsl.Output[dsl.Parquet]):
    import pandas as pd
    df = pd.read_csv(input_csv)
    df.to_parquet(output_parquet.path)

@dsl.pipeline(name='training-pipeline')
def train_pipeline():
    prep = prepare_data(input_csv='gs://bucket/data/raw.csv')
    # downstream steps use prep.outputs['output_parquet']

Worked example: caching a preprocessing step

1. Define the pipeline as above with enable_caching=True on prepare_data.
2. Run the pipeline once; the step executes, taking e.g., 45 seconds, and writes a parquet file to the artifact store.
3. Modify only a downstream step (e.g., change the learning rate in the training component) and re‑run the pipeline.
4. In the Kubeflow UI, observe that prepare_data shows a “cache hit” badge and its duration drops to near 0 seconds.
5. Verify reuse: locate the artifact store bucket, compare the checksum of the parquet file from the first run with the file referenced in the second run – they should match.

Trade‑offs and limitations

• Non‑deterministic logic: If a component reads the current time, a random seed, or external state without those values appearing in its inputs, caching will return stale results, hiding bugs.
• Storage growth: Large outputs (e.g., trained models) accumulate in the artifact store until inputs change, potentially increasing costs.
• Image or code changes: Updating a component’s container image or code does not automatically invalidate the cache; you must manually bump a version or toggle enable_caching=False to force a re‑run.

Actionable closing

To start benefiting from caching:

1. Identify steps with stable, deterministic outputs (data ingestion, feature engineering, model validation).
2. Add enable_caching=True to their component definitions.
3. Re‑compile and run the pipeline; verify cache hits in the UI and artifact‑store reuse.
4. Monitor storage usage and set lifecycle policies on the artifact store to prune old cached artifacts when appropriate.

When used judiciously, component caching cuts redundant compute, shortens iteration cycles, and makes pipeline audits easier through MLMD‑recorded cache hits.