> ## Documentation Index
> Fetch the complete documentation index at: https://notes.kodekloud.com/llms.txt
> Use this file to discover all available pages before exploring further.

# Demo Customizing Component Image

> Demonstrates creating and running a Kubeflow Pipeline component locally, installing missing Python packages in the component image to train a DecisionTree on the Iris dataset

Welcome — in this lesson you'll build a Kubeflow Pipeline (KFP) component that trains a simple classifier on the classic Iris dataset, run the component locally (each component executes in its own container), and learn how to customize the component image by installing additional Python packages.

This guide covers:

* A minimal KFP component that trains a DecisionTree on the Iris dataset.
* How to run the component locally using Docker.
* The common failure when required packages are missing and how to fix it with `packages_to_install`.
* A fully working script and expected runtime behavior.

Relevant links:

* [Kubeflow Pipelines Course](https://learn.kodekloud.com/user/courses/kubeflow)
* [scikit-learn documentation](https://scikit-learn.org/stable/)

## Dataset overview

* The Iris dataset contains feature vectors (`iris.data`) and class labels (`iris.target`).
* In this example we train a `DecisionTreeClassifier` and predict the class for a single sample.

## Minimal KFP component (initial example)

This minimal component uses the `python:3.10-slim` base image. Many minimal Python images do not include scientific libraries such as scikit-learn, so this example may fail unless scikit-learn is available in the image.

```python theme={null}
from kfp import dsl

@dsl.component(
    base_image="python:3.10-slim",
)
def classify_iris() -> str:
    from sklearn.datasets import load_iris
    from sklearn.tree import DecisionTreeClassifier

    # Load data
    iris = load_iris()

    # Train a simple model
    model = DecisionTreeClassifier()
    model.fit(iris.data, iris.target)

    # Predict one flower
    flower = [[5.1, 3.5, 1.4, 0.2]]
    prediction = model.predict(flower)[0]

    result = iris.target_names[prediction]

    print(f"Prediction: {result}")

    return result
```

## Run the component locally using Docker

During development it’s convenient to run KFP components locally. Use `kfp.local` with `DockerRunner()` so each component runs in its own Docker container.

```python theme={null}
from kfp import local

# Tell KFP to execute components using Docker
local.init(
    runner=local.DockerRunner()
)

# Invoke the component (this will execute it in a container)
task = classify_iris()

# Print the final result of the component execution
print(f"\nFinal result: {task.output}")
```

## What happens when a required package is missing

If the base image doesn’t include `scikit-learn`, the component container will fail with a `ModuleNotFoundError`. Typical error lines:

```text theme={null}
ModuleNotFoundError: No module named 'sklearn'
Traceback (most recent call last):
  File "/tmp/tmp.../ephemeral_component.py", line 8, in classify_iris
    from sklearn.datasets import load_iris
ModuleNotFoundError: No module named 'sklearn'
```

### Cause

The `python:3.10-slim` image is intentionally minimal and does not ship with many ML libraries.

### Fix

The `@dsl.component` decorator accepts a `packages_to_install` argument. KFP will run `pip install` for the listed packages inside the component container at startup. Use `packages_to_install` to install `scikit-learn` (or any other required packages) before your component code executes.

<Callout icon="lightbulb" color="#1CB2FE">
  Use `packages_to_install` when your component imports libraries that are not included in the base image. This runs `pip install` inside the component container at startup.
</Callout>

## Corrected full script (working version)

Below is the complete working script that initializes the local Docker runner and installs `scikit-learn` inside the component container before execution.

```python theme={null}
from kfp import dsl, local

# Tell KFP to execute components using Docker
local.init(
    runner=local.DockerRunner()
)

@dsl.component(
    base_image="python:3.10-slim",
    packages_to_install=["scikit-learn"]
)
def classify_iris() -> str:
    from sklearn.datasets import load_iris
    from sklearn.tree import DecisionTreeClassifier

    # Load data
    iris = load_iris()

    # Train a simple model
    model = DecisionTreeClassifier()
    model.fit(iris.data, iris.target)

    # Predict one flower
    flower = [[5.1, 3.5, 1.4, 0.2]]
    prediction = model.predict(flower)[0]

    result = iris.target_names[prediction]

    print(f"Prediction: {result}")

    return result

# Run the component locally
task = classify_iris()

# Print the final result (the output parameter)
print(f"\nFinal result: {task.output}")
```

## Expected runtime behavior

When you run the corrected script:

1. KFP pulls the `python:3.10-slim` base image (if not already present).
2. KFP installs `scikit-learn` inside the container (via pip).
3. The component code runs and prints the prediction.
4. The component returns the result as an output parameter.

Example summarized output:

```text theme={null}
Pulling image python:3.10-slim ...
Installing packages: scikit-learn ...
Prediction: setosa

Final result: setosa
```

Local output parameter (JSON):

```json theme={null}
{"parameterValues": {"Output": "setosa"}}
```

## Quick reference

| Topic | Description | Example |
| - | - | - |
| Base image | Minimal container image to run the component | `base_image="python:3.10-slim"` |
| Install packages | Packages installed at component startup via pip | `packages_to_install=["scikit-learn"]` |
| Local execution | Run components locally, each in its own Docker container | `local.init(runner=local.DockerRunner())` |
| Output parameter | The returned value from the component (accessible as `task.output`) | `{"parameterValues": {"Output": "setosa"}}` |

## Summary

* Define components with `@dsl.component(base_image=...)`.
* If your component imports packages not present in the base image, add them to `packages_to_install=[...]`.
* Use `kfp.local` plus `DockerRunner()` to run and test components locally before deploying to a cluster.

## Links and references

* [Kubeflow Pipelines — KodeKloud course](https://learn.kodekloud.com/user/courses/kubeflow)
* [scikit-learn](https://scikit-learn.org/stable/)
* [kfp Python SDK documentation](https://kubeflow-pipelines.readthedocs.io/en/latest/source/kfp.components.html)

<CardGroup>
  <Card title="Watch Video" icon="video" cta="Learn more" href="https://learn.kodekloud.com/user/courses/kubeflow/module/bece8da9-953e-480e-8774-b25b66c3830f/lesson/876c3017-9f08-4037-a5c8-03daeb7f8b4a" />
</CardGroup>


This documentation is built and hosted on [Mintlify](https://mintlify.com), a developer documentation platform.