MLOps基盤のフルマネージド化に向けたVertex AI Pipelinesへの移行

こんにちは。エンジニアの竹田です。
BUYMAの検索システムやMLOps基盤の開発・運用を担当しております。

こちらはEnigmo Advent Calendar 2023の21日目の記事です 🎄

弊社では2021年頃よりMLOps基盤をGoogle Cloud PlatformのAI Platform Pipelines上に構築して開発・運用を行っています。
この度、Vertex AI Pipelinesへの移行を全面的に進めることになりましたので、ご紹介も兼ねて記事にしたいと思います。

背景

2024/07にAI Platform Pipelinesが非推奨になるという通知を受けたことがきっかけです。
AI Platform Pipelines deprecations | Google Cloud

非推奨の通知が移行を開始するトリガーではあったものの、かねてからGKEクラスタの運用をどうにかできないかなと考えていました。
Vertex AI Pipelinesへの移行によりフルマネージド化できるのは大きなモチベーションとなっています。

移行における課題

この機会にKubeflow Pipelines SDK(以下、kfp) v1からv2へのアップグレードを進めています。
移行ドキュメントも提供されているため、それほど苦労はしないものと考えておりました。
が、結果としてこの選択が多くの苦労を抱えることになってしまいました 😢
実際に kfp v2 利用してみて、良かった点・苦労している点を交えてご紹介いたします。

※下位互換で動作させることも可能でしたが、kfp v1での機能拡張は行われない、そのうち書き換えが必要になる、という点を考慮してkfp v2への書き換えに踏み切りました。

kfp v1からv2へ

kfpは、kubernetes上で機械学習パイプラインを動作させるためのツールキットです。
コードベースはPythonです。
An introduction to Kubeflow

パイプラインを作成してVertex AI Pipelines上で動作させると、動作フローが視覚的に分かりやすく表現されます。

Pythonで記述したコードをコンパイルして利用する性質上、kfp v1の頃からかなりクセが強いなとは感じていました。
実際に利用してみた上での良い点、苦労している点を列挙し、所感を書いていこうと思います。

kfp v2の利用バージョン

# pip list | grep kfp
kfp                              2.3.0
kfp-pipeline-spec                0.2.2
kfp-server-api                   2.0.3

kfp v2の良い点

入出力に利用するInput[xxx] / Output[xxx] が便利
ParallelForによる並列処理の結果をCollectedで受け取れるようになった
@dsl.componentやset_accelerator_typeが直感的

入出力に利用する`Input[xxx]` / `Output[xxx]` が便利

kfp v2では入出力に利用するオブジェクトがこの形(基本的に Input[Artifact] / Output[Artifact] の利用)にほぼ統一されており、GCS上のパスを意識せず利用できるため非常に便利です。
https://www.kubeflow.org/docs/components/pipelines/v2/data-types/artifacts/

`ParallelFor`による並列処理の結果を`Collected`で受け取れるようになった

kfp v1でもParallelForは利用できましたが、fan-in(複数の入力を一つにまとめること)が厄介でした。
kfp v2では最近になってCollectedが利用可能となり、ParallelForの後に呼び出すことで結果をリスト形式でfan-inできるため、コードの可読性も飛躍的に向上します。

`@dsl.component`や`set_accelerator_type`が直感的

個人的な好みの部類かもしれませんが、kfp v2は定義周りが直感的になった印象があります。
https://www.kubeflow.org/docs/components/pipelines/v2/migration/#create_component_from_func-and-func_to_container_op-support

kfp v2で苦労している点

変数展開がおかしくなることがある
型指定の厳密化により、何を渡せばよいのか分からなくなることがある

変数展開がおかしくなることがある

kfp v2への移行で最も困っている点です。以下のようなissueも挙がっています。
https://github.com/kubeflow/pipelines/issues/10261
変数内に何らかの文字列や数値を入れているはずが、実際に利用する場合に以下のような展開がされてしまいます。
{{channel:task=;name=g;type=String;}}
コンポーネントの出力結果をうまく展開できない場合は以下のような内容です。 {{channel:task=term-calc;name=list_date;type=typing.List[str];}}
機械学習の初回実行プロセスの多くがBigQueryからのデータ取得を行っており、データ取得期間や特徴量を変数で管理しているため、既存のパイプラインコードではPythonのformatメソッドによる書式変換を多用しています。
この書式変換のほとんどが正常に動かなくなってしまい、試行錯誤を繰り返すことになってしまいました。
以下、期待した変数展開とならずにエラーとなるパターンの一部です。

パイプライン引数やコンポーネントの返却結果をdictに加えてコンポーネントに渡した場合

@dsl.component
def convert_str(tmpl: str, value: dict, output: Output[Artifact]):
    with open(output.path, "w") as f:
        f.write(tmpl.format(value))
  
@dsl.pipeline(name="test", description="test prediction")
def test_pipeline(table_name: str = "sample"):
    value = {
        "table_name": table_name
    }
    sql = "SELECT * FROM {0[table_name]}"
    convert_str_op = convert_str(tmpl=sql, value=value)

コンパイル時のエラー内容

ValueError: Value must be one of the following types: str, int, float, bool, dict, and list. Got: "{{channel:task=;name=table_name;type=String;}}" of type "<class 'kfp.dsl.pipeline_channel.PipelineParameterChannel'>".

パイプライン引数をパイプライン本体で利用しようとした場合

@dsl.pipeline(name="test", description="test prediction")
def test_pipeline(periods: str = "1m"):
    target_file = f"periods_{periods}.yaml"

    with open(target_file, mode="r") as f:
        periods_conf = yaml.safe_load(f)

コンパイル時のエラー内容

FileNotFoundError: [Errno 2] No such file or directory: 'periods_{{channel:task=;name=periods;type=String;}}.yaml'

型指定の厳密化により、何を渡せばよいのか分からなくなることがある

パイプライン引数を利用しようとして以下のようなエラーが出たり

TypeError: PipelineParameterChannel is not a parameter type.

特定コンポーネントにinputを渡した場合に以下のようなエラーが出たり

ValueError: Constant argument inputs must be one of type ['String', 'Integer', 'Float', 'Boolean', 'List', 'Dict'] Got: <kfp.dsl.pipeline_task.PipelineTask object at 0x7f8a03f89880> of type <class 'kfp.dsl.pipeline_task.PipelineTask'>.

といったことが割と発生します。
自分としては正しい型での引き渡し、および参照をしているつもりのため、どう対処してよいか分からなくなることが多いです。
事前にキャストすることで正常に動作することもあれば、そもそもデータ型の扱いを見直す必要があったりします。