🤔

Kubeflow Pipelines v2 で Pipeline の書き方がかなり変わる件について

2023/12/25に公開

Google Cloud

MLOps

Vertex AI

tech

この記事は Google Cloud Champion Innovators Advent Calendar 2023 の 12/23 の記事です。

2023年7月、Kubeflow Pipelines は version 2.0 をリリースしました。このリリースがされたことにより、パイプラインの記述が大きく変わります。この記事ではどのような変化があったのか確認しましょう。

Migration from v1 to v2

KFP SDK v2 へのバージョンアップについては、Migrate from KFP SDK v1 というドキュメントが提供されています。まずはこれを確認しましょう。

まず、慌てる必要はないことがすぐに分かります。

We suggest you migrate your code to the “New usage”, even though the “Previous usage” will still work with warnings.

一方、Vertex Pipelines のドキュメントには KFP SDK 1.8 系が 2024 年 12 月にサポートされなくなると記述されています。

これは Kubeflow Pipelines のアーキテクチャが関係しています。Kubeflow Pipelines はパイプラインを Argo Workflow を用いて実行するバックエンドと、パイプラインを記述する SDK、ユーザーインターフェイスを提供する Web サーバーから構成されています。Vertex Pipelines はこのうち、バックエンドとユーザーインターフェイスを提供しています。このため、バックエンドと SDK のバージョンに気をつける必要があります。

もう少し詳しく見ておきましょう、Kubeflow Pipelines では SDK のバージョンとバックエンドのバージョンを厳密に揃える必要はありません。Kubeflow Pipelines ではパイプラインを Python で記述したあと、pipeline_spec という中間言語にコンパイルします。この中間言語は 1.8 系と 2 系で共通するので、バージョンを厳密に揃える必要はありません。一方、1.8 系は 2 系の SDK が公開されてから 1 年でメンテナンスされなくなる方針が打ち出されているため、いずれにせよ 2 系への対応を計画したほうが良いでしょう。

次に、2系への対応についてですが、少々混乱しやすいので確認しておきましょう。次の記述は 1.8 系の記述です。

from kfp.v2 import dsl
from kfp.v2 import compiler

@dsl.pipeline(name='my-pipeline')
def pipeline():
  ...

compiler.Compiler().compile(...)

あたかも 2 系の SDK を使っているかのように見えるかもしれませんが、1.8 系の記述です。2系の同様の記述は次のようになります。

from kfp import dsl
from kfp import compiler

@dsl.pipeline(name='my-pipeline')
def pipeline():
  ...

compiler.Compiler().compile(...)

kfp.v2 名前空間は v2 移行のために 1.8 系に暫定的に用意されたものだったので、これがなくなるのは自然でしょう。

他にも、Vertex Pipelines 固有の機能が KFP SDK から削除されるなど、暫定的に行われた対策について整理が進んでいます。以前からのユーザーは Migrate from KFP SDK v1 を確認し、今後の対応計画を立てると良いでしょう。

また、大きな変更として fan-in への対応や pパイプラインを他のパイプラインのコンポーネントとして利用できる機能があります。これまで自力で実装しなければいけなかった fan-in が SDK の機能として提供されるようになったため、パイプラインの実装をシンプルにできないかあわせて検討すると良いでしょう。

KFP v2 SDK におけるコンポーネント

Kubeflow Pipelines では、前処理や訓練など、まとまった処理をコンテナーとして実装したものを、コンポーネントと呼びます。KFP v2 SDK ではコンポーネントの記述方法として次の3つが用意されています。

Lightweight Python Components
Containerized Python Components
Container Components

Containerized Python Components と Container Components は新しく導入された方法です。これらは次のように要約できます。

特徴	Lightweight Python Components	Containerized Python Components	Container Components
実装の記述量	小	中	大
抽象度	高度に抽象化されている	中間程度	ほぼ `docker run`
コンポーネントとパイプラインの記述の分離	困難	可能	可能
コンテナー内に KFP SDK が必要	必要	必要	不要
イメージのビルドに KFP SDK が必要	不要	必要	不要
Artifact の Type	利用可能	利用可能	非現実的

それぞれについて見ていきましょう。

Lightweight Python Components

Lightweight Python Components は以前から存在するコンポーネントの記述方法です。次のような記述で Python の関数としてコンポーネントを定義します。

from kfp import dsl

@dsl.component(base_image='python:3.8')
def add(a: int, b: int) -> int:
    return a + b

この方法は取り組み始めるのにはいいのですが、実運用しようとするとパイプラインの記述が複雑になりがちです。次は訓練と評価を行っているコンポーネントの例ですが、このようにモデルの定義を含めてすべてをコンポーネント内に記述するのは実際は難しいでしょう。

# https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/pipelines/metrics_viz_run_compare_kfp.ipynb より引用
@component(packages_to_install=["scikit-learn==1.2.2"], base_image="python:3.9")
def wine_classification(wmetrics: Output[ClassificationMetrics]):
    from sklearn.datasets import load_wine
    from sklearn.ensemble import RandomForestClassifier
    from sklearn.metrics import roc_curve
    from sklearn.model_selection import cross_val_predict, train_test_split

    X, y = load_wine(return_X_y=True)
    # Binary classification problem for label 1.
    y = y == 1

    X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=42)
    rfc = RandomForestClassifier(n_estimators=10, random_state=42)
    rfc.fit(X_train, y_train)
    y_scores = cross_val_predict(rfc, X_train, y_train, cv=3, method="predict_proba")
    fpr, tpr, thresholds = roc_curve(
        y_true=y_train, y_score=y_scores[:, 1], pos_label=True
    )
    wmetrics.log_roc_curve(fpr, tpr, thresholds)

これを避けるためにはあらかじめイメージを自分でビルドしておき base_image で利用する方法があります。個人的にはこの手法がもっとも柔軟性を保ちながら Kubeflow Pipelines の機能を活用できると考えていますが、実際に設計するためには Kubeflow Pipelines の実装に関する深い知識が必要となります。

Container Components

Container Components はもっともシンプルな方法です。基本的には次のような方法で、イメージとコマンド、引数を指定してコンテナーを起動します。

@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 $RESPONSE > $1
                            '''
        ],
        args=[name, greeting])

制約事項として、Kubeflow Pipelines のメタデータ管理の機能の利用が著しく困難になる点には注意が必要です。

Kubeflow Pipelines には Artifact という、コンポーネントの入出力の型を指定する機能があり、適切に型を指定することでリネージの管理や可視化に用いることができます。Vertex Pipelines におけるリネージや可視化については Output HTML and Markdown が良い例になっています。これらの機能が運用上必須でない場合は Container Components が良い選択肢になるでしょう。

Container Components については LayerX Engineering Blog に良い解説があります。

Containerized Python Components

Contaienrized Python Components は Lightweight Python Components と Containerized Python Components の中間となる方法で、Lightweight Python Components の成約を緩めた方法となっています。以降では Containerized Python Components にしたがって解説します。

まず、次のようなディレクトリ構造を用意します。

pipeline.py
src/
├── __init__.py
├── my_component.py
└── math_utils.py

math_utils.py は次のような内容で、コンポーネント内で行いたい処理を純粋に KFP SDK なしに記述します。

# src/math_utils.py
def add_numbers(num1, num2):
    return num1 + num2

my_component.py はコンポーネントのインターフェイスを Python で記述しているファイルで、KFP SDK を用いてコンポーネントを定義します。

# src/my_component.py
from kfp import dsl
from math_utils import add_numbers

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

次に、kfp component build を用いて、my_component.py をコンパイルします。

kfp component build src/ --component-filepattern my_component.py --no-push-image

コンパイルすると、次のようなファイルができあがります。

pipeline.py
src/
├──component_metadata/
|  └── my_component.yaml
├── .dockerignore
├── Dockerfile
├── kfp_config.ini
├── runtime-requirements.txt
├── __init__.py
├── my_component.py
└── math_utils.py

my_component.yaml はこのコンポーネントのみで構成されるパイプライン定義が pipeline_spec にしたがって記述されています。重要な部分だけ引用します。

deploymentSpec:
  executors:
    exec-my-component:
      container:
        args:
        - --executor_input
        - '{{$}}'
        - --function_to_execute
        - add
        command:
        - python3
        - -m
        - kfp.dsl.executor_main
        image: my-component

これはこのコンテナーの起動方法を示しており、コンテナーが python3 -m kfp.dsl.executor_main --executor_input '{{$}}' --function_to_execute add というコマンドを実行することを示しています。'{{$}}' はExecutor Input のプレースホルダーとなっており、実行時にはそのコンポーネントへの入出力のためのすべての情報を保持した JSON フォーマットの文字列が渡されます。

kfp.dsl.executor_main は Executor Input をパースし、Artifact を作成して --function-to-execute で指定した関数に渡します。

runtime-requirements.txt は dsl.component の packages_to_install で指定した、そのコンポーネントに追加で必要になるライブラリが列挙されます。仮に packages_to_install=['scikit-learn'] と指定すると次のようになります。

# 次のようにコンポーネントを定義した場合にできあがるファイル、指定しなければ空のファイルになる
# @dsl.component(
#   packages_to_install=['scikit-learn'],
# )
# def add(a: int, b: int) -> int:
#     return add_numbers(a, b)
scikit-learn

kfp_config.ini はコンポーネントの定義を行っている Python ファイルの情報を保持してます。

[Components]
my_component = my_component.py

Dockerfile では、上記の情報を元に次のようにしてイメージファイルを作成します。

# Generated by KFP.

FROM python:3.7 

WORKDIR /src/
COPY runtime-requirements.txt runtime-requirements.txt
RUN pip install --no-cache-dir -r runtime-requirements.txt

RUN pip install --no-cache-dir kfp==2.4.0
COPY . .

このようにして kfp component build を用いたコンポーネントは次のようにしてパイプライン中で用います。

# pipeline.py
from src.my_component import add
from kfp import dsl

@dsl.pipeline
def addition_pipeline(x: int, y: int) -> int:
    task1 = add(a=x, b=y)
    task2 = add(a=task1.output, b=x)
    return task2.output

compiler.Compiler().compile(addition_pipeline, 'pipeline.yaml')

注意点をいくつか述べておきます。

kfp component build を実行する際の Python 環境では、ビルド対象のライブラリすべてを利用可能にしておく必要があります
上記は v2.4.0 の実装について述べていますが、将来的に変更があるかもしれません、たとえば、コンポーネントごとにできあがる YAML ファイルは将来的に生成されなくなるかもしれません
結局、パイプラインでは Lightweight Python Components を使っています

Containerized Python Components は Kubeflow Pipelines のすべての機能を利用可能な方法で、柔軟にコンポーネントを記述可能です。一方、上記の方法が必ずしもそれぞれのパイプラインの設計に適しているとは言いがたいのも事実です。kfp component build で作成したファイルを初期ファイルとして、それぞれの現場で使いやすいように生成されたファイルを変更し、パイプラインでは使うことになるでしょう。

総じて、2023 年は Kubeflow Pipelines や Vertex Pipelines にとって大きな節目を迎えた年だったと言えます。

注意点として、これまでの Component Spec を手書きする方法は今後推奨されなくなっていくようです。もし機械学習パイプラインをすでに大規模にデプロイしている場合、v2 への移行に当たっては大規模な改修を覚悟しましょう。デッドラインは来年の今頃です。

一方で、fan-in などの機能追加も行われた年でもありました。よりシンプルな実装とできないか、いったん棚卸しするにはよい機会となるでしょう。