kfpをlocalで動かす

# kubeflow pipeline の ローカル実行とコンポーネント定義

[@nishikoh](https://github.com/nishikoh)

## Contents

- kfp をローカルで動かす
- kfp のコンポーネント定義
- その他、気になったもの

## kfp とは

Kubeflow Pipelines（KFP）は、コンテナイメージを使ってポータブルでスケーラブルな機械学習（ML）ワークフローを構築し、デプロイするためのプラットフォーム。

Vertex AI Pipelines の裏で kfp が動いている。

## kfp を使う際の困りごと

- 毎回リモート環境で動作確認すると大変
- それぞれのコンポーネントはトレードオフの関係
  - Lightweight Python Components
  - Containerized Python Components
  - Container Components

発表の中でこれらをどう解消するか紹介

## kfp を手元の開発環境で実行する

### コンポーネントの動作確認どうする？

機械学習用のコンテナイメージは GB 単位になる。コンテナイメージを毎回 push してリモートのパイプラインで動作確認すると待機時間が長くなる。

↪︎ 手元の開発環境で動作確認できると生産性が上がる。

（余談：2024 年 1 月にローカル実行のドキュメントが公開された）

### ローカル環境でコンポーネント実行

シンプルな足し算の例

```python [|4-6|12-13]
from kfp import local
from kfp import dsl

# 関数定義の後に実行しても良い
# 実行にはdockerが必要
local.init(runner=local.DockerRunner())

@dsl.component
def add(a: int, b: int) -> int:
    return a + b

task = add(a=1, b=2)
assert task.output == 3
```

```bash
pants run sandbox/kfp_local/01_function.py
```

```text title="result"
...
  return component_factory.create_component_from_func(
18:45:51.384 - INFO - Executing task 'add'
18:45:51.384 - INFO - Streamed logs:

Found image 'python:3.7'

WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv
    [KFP Executor 2024-01-15 18:45:55,664 INFO]: Looking for component `add` in --component_module_path `/tmp/tmp.ti3t2K9DwH/ephemeral_component.py`
    [KFP Executor 2024-01-15 18:45:55,664 INFO]: Loading KFP component "add" from /tmp/tmp.ti3t2K9DwH/ephemeral_component.py (directory "/tmp/tmp.ti3t2K9DwH" and module name "ephemeral_component")
    [KFP Executor 2024-01-15 18:45:55,665 INFO]: Got executor_input:
    {
        "inputs": {
            "parameterValues": {
                "a": 1,
                "b": 2
            }
        },
        "outputs": {
            "parameters": {
                "Output": {
                    "outputFile": "<my_path>/local_outputs/add-2024-01-15-18-45-51-383673/add/Output"
                }
            },
            "outputFile": "<my_path>/local_outputs/add-2024-01-15-18-45-51-383673/add/executor_output.json"
        }
    }
    [KFP Executor 2024-01-15 18:45:55,665 INFO]: Wrote executor output file to <my_path>/local_outputs/add-2024-01-15-18-45-51-383673/add/executor_output.json.
18:45:55.877 - INFO - Task 'add' finished with status SUCCESS
18:45:55.878 - INFO - Task 'add' outputs:
    Output: 3
```

出力

```json title="local_outputs/add-2024-01-15-18-45-51-383673/add/executor_output.json"
{ "parameterValues": { "Output": 3 } }
```

### アーティファクトを出力

足し算の結果をファイルに書き込む例

```python [|14-17|23-25]
from kfp import local
from kfp import dsl
from kfp.dsl import Output, Artifact
import json

local.init(runner=local.DockerRunner())

@dsl.component
def add(a: int, b: int, out_artifact: Output[Artifact]):
    import json

result = json.dumps(a + b)

with open(out_artifact.path, 'w') as f:
        f.write(result)

out_artifact.metadata['operation'] = 'addition'

task = add(a=1, b=2)
# can read artifact contents
with open(task.outputs['out_artifact'].path) as f:
    contents = f.read()
assert json.loads(contents) == 3
assert task.outputs['out_artifact'].metadata['operation'] == 'addition'
```

```bash
pants run sandbox/kfp_local/02_output.py
```

```text title="result"
...
  return component_factory.create_component_from_func(
20:38:28.731 - INFO - Executing task 'add'
20:38:28.731 - INFO - Streamed logs:

Found image 'python:3.7'

WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv
    [KFP Executor 2024-01-15 20:38:32,768 INFO]: Looking for component `add` in --component_module_path `/tmp/tmp.PbVALy4cam/ephemeral_component.py`
    [KFP Executor 2024-01-15 20:38:32,769 INFO]: Loading KFP component "add" from /tmp/tmp.PbVALy4cam/ephemeral_component.py (directory "/tmp/tmp.PbVALy4cam" and module name "ephemeral_component")
    [KFP Executor 2024-01-15 20:38:32,769 INFO]: Got executor_input:
    {
        "inputs": {
            "parameterValues": {
                "a": 1,
                "b": 2
            }
        },
        "outputs": {
            "artifacts": {
                "out_artifact": {
                    "artifacts": [
                        {
                            "name": "out_artifact",
                            "type": {
                                "schemaTitle": "system.Artifact",
                                "schemaVersion": "0.0.1"
                            },
                            "uri": "<my_path>/local_outputs/add-2024-01-15-20-38-28-731045/add/out_artifact",
                            "metadata": {}
                        }
                    ]
                }
            },
            "outputFile": "<my_path>/local_outputs/add-2024-01-15-20-38-28-731045/add/executor_output.json"
        }
    }
    [KFP Executor 2024-01-15 20:38:32,771 INFO]: Wrote executor output file to <my_path>/local_outputs/add-2024-01-15-20-38-28-731045/add/executor_output.json.

__import__(pkg_name)
20:38:32.975 - INFO - Task 'add' finished with status SUCCESS
20:38:32.975 - INFO - Task 'add' outputs:
    out_artifact: Artifact( name='out_artifact',
                            uri='<my_path>/local_outputs/add-2024-01-15-20-38-28-731045/add/out_artifact',
                            metadata={'operation': 'addition'} )
```

出力

```json
{
  "artifacts": {
    "out_artifact": {
      "artifacts": [
        {
          "name": "out_artifact",
          "uri": "<my_path>/local_outputs/add-2024-01-15-20-38-28-731045/add/out_artifact",
          "metadata": { "operation": "addition" }
        }
      ]
    }
  }
}
```

### アーティファクト

アーティファクトはパイプラインと紐づくもの。

- Dataset
- Model
- Metrics
- HTML
- Markdown ...など

### 関数コンポーネントと 3rd party package への依存

事前に依存関係を全て含んだイメージを作り、base_image として設定

```python [|9|7]
from kfp import dsl, local
from kfp.dsl import Artifact, Output

local.init(runner=local.DockerRunner())

@dsl.component(base_image="kfp_base:latest")
def dump_pydantic_version(out_artifact: Output[Artifact]):
    from sandbox.kfp_local.util import get_pydantic_version

version = get_pydantic_version()

with open(out_artifact.path, "w") as f:
        f.write(version)

out_artifact.metadata["operation"] = "dump_pydantic_version"

task = dump_pydantic_version()
with open(task.outputs["out_artifact"].path) as f:
    contents = f.read()
print(contents)
```

```bash
pants run projects/sandbox/kfp_local/03_base_image.py
```

```text title="result"
18:42:12.76 [INFO] Completed: Building 3 requirements for kfp_local.pex from the 3rdparty/python/default.lock resolve: docker==7.0.0, kfp==2.6.0, pydantic==2.5.3
18:42:13.385 - INFO - Executing task 'dump-pydantic-version'
18:42:13.385 - INFO - Streamed logs:

Found image 'kfp_base:latest'

[KFP Executor 2024-01-16 18:42:18,760 INFO]: Looking for component `dump_pydantic_version` in --component_module_path `/tmp/tmp.R4Nb4FGxn6/ephemeral_component.py`
    [KFP Executor 2024-01-16 18:42:18,760 INFO]: Loading KFP component "dump_pydantic_version" from /tmp/tmp.R4Nb4FGxn6/ephemeral_component.py (directory "/tmp/tmp.R4Nb4FGxn6" and module name "ephemeral_component")
    [KFP Executor 2024-01-16 18:42:18,762 INFO]: Got executor_input:
    {
        "inputs": {},
        "outputs": {
            "artifacts": {
                "out_artifact": {
                    "artifacts": [
                        {
                            "name": "out_artifact",
                            "type": {
                                "schemaTitle": "system.Artifact",
                                "schemaVersion": "0.0.1"
                            },
                            "uri": "<my_path>/local_outputs/dump-pydantic-version-2024-01-16-18-42-13-384523/dump-pydantic-version/out_artifact",
                            "metadata": {}
                        }
                    ]
                }
            },
            "outputFile": "<my_path>/local_outputs/dump-pydantic-version-2024-01-16-18-42-13-384523/dump-pydantic-version/executor_output.json"
        }
    }
    2.5.3
    [KFP Executor 2024-01-16 18:42:18,767 INFO]: Wrote executor output file to <my_path>/local_outputs/dump-pydantic-version-2024-01-16-18-42-13-384523/dump-pydantic-version/executor_output.json.
18:42:18.990 - INFO - Task 'dump-pydantic-version' finished with status SUCCESS
18:42:18.991 - INFO - Task 'dump-pydantic-version' outputs:
    out_artifact: Artifact( name='out_artifact',
                            uri='<my_path>/local_outputs/dump-pydantic-version-2024-01-16-18-42-13-384523/dump-pydantic-version/out_artifact',
                            metadata={'operation': 'dump_pydantic_version'} )

2.5.3
```

### docker

`dsl.ContainerSpec`をつかうと任意のコンテナイメージで実行ができる。文字列をファイルに書き込む例。

```python [|5]
from kfp import dsl, local

local.init(runner=local.DockerRunner())

@dsl.container_component
def say_hello(name: str, greeting: dsl.OutputPath(str)):
    """Log a greeting and return it as an output."""

return dsl.ContainerSpec(
        image="alpine",
        command=[
            "sh",
            "-c",
            """RESPONSE="Hello, $0!"\
                            && echo $RESPONSE\
                            && mkdir -p $(dirname $1)\
                            && echo hi $RESPONSE > $1
                            """,
        ],
        args=[name, greeting],
    )

task = say_hello(name="World")
print(task.outputs)
```

```bash
pants run projects/sandbox/kfp_local/04_docker.py
```

```text title="result"
06:39:37.468 - INFO - Executing task 'say-hello'
06:39:37.468 - INFO - Streamed logs:

Found image 'alpine:latest'

Hello, World!
06:39:37.953 - INFO - Task 'say-hello' finished with status SUCCESS
06:39:37.953 - INFO - Task 'say-hello' outputs:
    greeting: 'hi Hello, World!
'

{'greeting': 'hi Hello, World!\n'}
```

### GPU を使ったコンポーネント

GPU を使ったコンポーネントをローカルで動かす

```python [|9]
from kfp import dsl, local

local.init(runner=local.DockerRunner())

@dsl.container_component
def gpu_processing():
    return dsl.ContainerSpec(
        image="gcr.io/google_containers/cuda-vector-add:v0.1",
    )

task = gpu_processing()
print(task.outputs)
```

```bash
pants run projects/sandbox/kfp_local/06_gpu_component.py
```

```text title="result"
16:21:16.599 - INFO - Executing task 'gpu-processing'
16:21:16.600 - INFO - Streamed logs:

Found image 'gcr.io/google_containers/cuda-vector-add:v0.1'

[Vector addition of 50000 elements]
    Copy input data from the host memory to the CUDA device
    CUDA kernel launch with 196 blocks of 256 threads
    Copy output data from the CUDA device to the host memory
    Test PASSED
    Done
16:21:18.040 - INFO - Task 'gpu-processing' finished with status SUCCESS
16:21:18.040 - INFO - Task 'gpu-processing' has no outputs
{}
```

GPU を使ったコンポーネントもローカルで動かせる

### pipeline 実行

パイプライン実行をやってみる

```python [|17-25]
from kfp import dsl, local

local.init(runner=local.DockerRunner())

@dsl.component
def square(x: float) -> float:
    return x ** 2

@dsl.component
def add(x: float, y: float) -> float:
    return x + y

@dsl.component
def square_root(x: float) -> float:
    return x ** .5

@dsl.pipeline
def pythagorean(a: float, b: float) -> float:
    a_sq_task = square(x=a)
    b_sq_task = square(x=b)
    sum_task = add(x=a_sq_task.output, y=b_sq_task.output)
    return square_root(x=sum_task.output).output

result = pythagorean(a=3.0, b=4.0)
print(result)
```

```bash
pants run projects/sandbox/kfp_local/05_pipeline.py
```

```text title="result"
...
    raise NotImplementedError(
NotImplementedError: Local pipeline execution is not currently supported.
zsh: exit 1     pants run projects/sandbox/kfp_local/05_pipeline.py
```

パイプライン実行はローカルではサポートされていない

パイプライン全ての機能は利用できないけど、簡単なオーケストレーションの確認くらいはできる

```python [|17]
from kfp import dsl, local

local.init(runner=local.DockerRunner())

@dsl.component
def square(x: float) -> float:
    return x ** 2

@dsl.component
def add(x: float, y: float) -> float:
    return x + y

@dsl.component
def square_root(x: float) -> float:
    return x ** .5

#@dsl.pipeline # ここをコメントアウト
def pythagorean(a: float, b: float) -> float:
    a_sq_task = square(x=a)
    b_sq_task = square(x=b)
    sum_task = add(x=a_sq_task.output, y=b_sq_task.output)
    return square_root(x=sum_task.output).output

result = pythagorean(a=3.0, b=4.0)
print(result)
```

3^2 + 4^2 の平方根は 5 なので正しく動いている

```text title="result"
...
06:59:08.912 - INFO - Task 'square-root' finished with status SUCCESS
06:59:08.912 - INFO - Task 'square-root' outputs:
    Output: 5.0
```

### ここまでのまとめ

- コンポーネント、パイプラインの簡単な確認は手元の環境でできる。毎回リモート環境で動作確認しなくて良い
- GPU を使ったコンポーネントも手元の環境で動かせる

## コンポーネントの依存管理

### コンポーネント定義の方法によって依存管理の方法が異なる

- Lightweight Python Components
- Containerized Python Components
- Container Components

### Lightweight Python Components

関数からコンポーネント作成

```python [|1]
@dsl.component(packages_to_install=['numpy==1.21.6'])
def sin(val: float = 3.14) -> float：
   return np.sin(val).item()
```

依存管理のメンテナンス性は正直微妙。アーティファクトは制限なく使える。

### Containerized Python Components

Lightweight Python Components を拡張して、コンテナイメージとして事前にビルドして依存管理

`kfp components build`でコンテナイメージをビルドする。アーティファクトは制限なく使える。

これをやるくらいなら`pants`使ってビルドした方が楽

### Container Components

任意のコンテナイメージをコンポーネントとして使う

```python
from kfp import dsl

@dsl.container_component
def say_hello():
    return dsl.ContainerSpec(image='alpine', command=['echo'], args=['Hello'])
```

依存管理は普通のコンテナビルド。一部のアーティファクトのみ使える。

### コンポーネントの比較

|  | func | 中間 | Container |
| --- | :---: | :---: | :---: |
| 依存管理 | pip | pip | docker |
| アーティファクト | ○ | ○ | △ |
| ポータビリティ | x | △ | ○ |
| ビルドの手間 | - | x | ○ |

### コンポーネント定義で実現したいこと

- 楽な依存管理と高いポータビリティ
- アーティファクトの制限なし

### アーティファクト利用と依存管理を両立させる

- 依存を含んだコンテナイメージを`pants`でビルド
- Lightweight Python Components の base image でビルドしたイメージを利用

関数にコアなロジックを書くと密結合になってテストしにくくなるので、コアなロジックは別のモジュールに切り出して記述したい。

## まとめ

- kfp をローカルで動かせる
- コンポーネントは pants と func で構築すると依存管理とアーティファクトの利用を両立できる

<!--

時間が余ったら

- fiftyone
- skypilot
- zenml
- cirun
- numaflow -->