Timee Product Team Blog

タイミー開発者ブログ

Vertex AI Pipelinesテンプレートを管理するArtifact Registryの導入

Timee Product Advent Calendar 2024 13日目の記事です。

MLOpsエンジニアとして10月にタイミーにジョインした、ともっぴです。

データエンジニアリング部 データサイエンス（以下DS）グループに所属し、ML基盤の構築・改善に取り組んでいます。

概要

本記事では、Vertex AI Pipelinesを効率的に開発するために行った 「Vertex AI Pipelinesテンプレートを管理するArtifact Registryの導入」 の取り組みを紹介します。

過去、DSグループが取り組んできたVertex AI Pipelinesの開発効率化は、以下の記事を参照ください。

tech.timee.co.jp tech.timee.co.jp

背景と課題

背景

前提としてタイミーのML基盤では、サービスレベルに応じた複数のGoogle Projectが存在し、その中で複数のMLパイプラインが動いています。(参考記事)

Artifact Repositoryを導入する以前は、GCSにMLパイプラインテンプレートを格納していました。

下図は簡易的に表現したその時のアーキテクチャです。

Artifact Registry導入前のアーキテクチャ

課題

この状態での課題は2点ありました。

  1. 純粋な管理の煩雑さ

    複数の環境のGCSに、以下のような構成でMLパイプラインテンプレートが存在していました。

     ./project-internal
 ├── ML_Project_A_Bucket
 │   └── pipeline
 │       ├── v0.1.0
 │       │   └── pipeline-config
 │       │       └── pipeline.json
 │       └── v0.2.0
 │           └── pipeline-config
 │               └── pipeline.json
 └── ML_Project_B_Bucket
     └── pipeline
         ├── v1.0.0
         │   └── pipeline-config
         │       └── pipeline.json
         └── v1.1.0
             └── pipeline-config
                 └── pipeline.json

    こうした状況下で、パイプラインやバージョンが増えていくにつれ、見通しが悪くなっていました。

  2. ライフサイクル管理の煩雑さ

    MLパイプラインテンプレートすべてのファイルを保存し続ける必要はなく、開発用は一定の期間で削除、逆に本番用はバージョン管理して一定の世代分は必ず保持する、という管理が求められます。

    GCSでもオブジェクトライフサイクルが設定できるので、n世代前は削除する、n日前のパイプラインは削除する、等のルールは設定可能です。

    しかし従来の構成では、一つのMLプロジェクトにつき一つのバケットを作成していたので、MLプロジェクトが増えるたびに、ライフサイクル設定が必要になってしまい、手間が増える状況でした。

    GCSに集約するという方法も考えられますが、Artifact Registryの方が各種ポリシーを柔軟に設定できることから、Artifact Registryの採用に至りました。また、Artifact RegistryのKFP対応自体が最近*1 *2のことなので、これまで採用を見送ってきた背景もあります。

ソリューション

改善後のアーキテクチャ

Artifact RegistryにKFP formatのリポジトリを作成し、そのリポジトリに対してCDでMLパイプラインテンプレートをアップロードする構成となりました。

Artifact Registry導入後のアーキテクチャ

その構成にしたことで、前項で書いたような複数プロジェクト・複数のバケットで管理してきたMLパイプラインテンプレートを、一つのリポジトリの中で管理できるようになりました。 （以下のディレクトリ構成のイメージ）

./project-central-repository
└── kfp-repository
    ├── pipeline-1
    │   ├── version1.yaml
    │   └── version2.yaml
    └── pipeline-2
        └── version1.yaml

実装サマリー

具体的にどのような実装を行なったのか、概要を紹介します。

  1. CDスクリプトの変更

    CDで行うパイプラインテンプレートのアップロードは、kfpライブラリのRegistryClientを利用しました。KFP CLIから利用できたら便利なのですが、こちらのissueを読む限り現在サポートはされていないようでした。

    そのため、簡易なPythonスクリプトを用意して、 GitHub Actions上で呼び出すことにしました。

  2. cleanup_policyを設定しパイプラインのライフサイクルを管理

    Artifact Registryではcleanup_policyを設定することで、ライフサイクルを設定できます。Artifact Registryで一元管理することで、統一的なルールでライフサイクルを、各MLパイプラインテンプレートに適用でき、管理の煩雑さが解消できました。

実装途中で詰まったところ

実装をする上で、いくつか直面した問題があったので、その概要と解決方法を紹介します。

1) GitHub Actions上でkfp.RegistryClient.upload_pipelineが通らない

CDで呼び出しているアップロード用のPythonスクリプトですが、ローカルでは正常終了するにもかかわらず、GitHub Actions上では認証エラーになる現象に遭遇しました。

原因を探っていると、LayerXさんの記事に行き当たり、RegistryClientがWorkload Identityに対応していないことが原因と判明しました。

LayerXさんの記事の通り、kfp.registry.ApiAuthを利用することで問題は解決しました。

2) jsonではアップロードできず、yamlが必要になる

改善前の構成では、パイプラインはjsonでGCSに格納していました。

しかし、kfp.RegistryClient.upload_pipeline(…)を利用すると、jsonではinvalid extension type pipeline.json というエラーでBad Requestとなりました。

ドキュメントを確認すると、Artifact Registry REST APIを利用する場合、yamlである必要があると記載がありました。

KFP SDK v2ではyamlフォーマットが推奨されており、かつKFP SDK v1のサポートは2024/12/20に切れるとされているため、徐々にyamlに移行していくこととなりそうです。

3) Cloud Composer(Airflow)のRunPipelineJobOperatorからArtifact Registryのパイプラインを利用する際にタグ指定だと権限エラーとなる

最初、以下のようなコードで動作すると思っていたところ、権限エラーとなりました。

from airflow.decorators import dag
from airflow.providers.google.cloud.operators.vertex_ai.pipeline_job import RunPipelineJobOperator

@dag(...)
def main():
    TEMPLATE_REPOSITORY_URL="https://asia-northeast1-kfp.pkg.dev/project-central-repository/pipelines/ml-project/v1.0.0"
    
    vertexai_pipeline = RunPipelineJobOperator(
        task_id=task_id,
        project_id=gcp_project_id,
        region=gcp_location,
        display_name=display_name,
        template_path=TEMPLATE_REPOSITORY_URL,  # KFP Repository上のテンプレートのURL
        service_account=pipeline_service_account
    )

発生したエラーは以下です。

Service account XXX does not have permission to get ArtifactRegistry tag projects/project-central-repository/locations/asia-northeast1/repositories/pipelines/packages/ml-project/tags/v1.0.0 in region asia-northeast1.

タグを取得する権限がないと言われています。

ここで、Artifact Registryにおけるタグとバージョンの違いについて確認します。

タグとは、MLパイプラインテンプレートをアップロードする際に、任意で付与できるものです。一方バージョンは、アップロード時に自動で付与される、sha256: で始まるハッシュ値です。

MLテンプレートのタグとバージョン

関係するサービスアカウントと、Artifact RegistryへのREAD権限は以下の通りです。

サービスアカウント Artifact RegistryへのREAD権限
Cloud Composer用のサービスアカウント あり
MLパイプライン用のサービスアカウント (pipeline_service_account) なし

一時的なエラー解消方法は、2通りあります。

  1. タグではなくバージョンハッシュを指定する
  2. MLパイプライン用のサービスアカウントにArtifact RegistryへのREAD権限を付与する。

しかし、それぞれデメリットがあり、

1.の方法では、MLパイプライン更新のたびにバージョンハッシュを指定し直す運用となり、意図せぬバージョン指定のミスや、手間がかかることがデメリットです。

2.の方法でも、既存のMLパイプライン用サービスアカウントのIAMを変更し、新規のパイプラインを作る際も権限を毎度付与する手間が発生します。タイミーのML基盤では、 MLパイプラインごとに柔軟に権限付与するために、MLパイプラインごとにサービスアカウントを作成する運用になっているため、その工数は無視できないものになります。

Airflowのソースコードを確認したところ、RunPipelineJobOperatorの内部では、PipelineJob初期化時にArtifact Registryからyamlを取得し、Vertex AI Pipelinesのジョブを作成するリクエストを送っている流れになっていることがわかりました。

上述のエラーは、ジョブを作成するリクエストをまさに送っている部分で発生しており、この部分でMLパイプライン用のサービスアカウントから共有リポジトリへのアクセスが発生していると判断できました。

原因を踏まえると、RunPipelineJobOperatorにバージョン付きのMLパイプラインテンプレートを渡せるようにする方法が良いと考えました。

タグからバージョンを取得するには、kfp.RegistryClient.get_tag(…)を利用すれば良いので、最終的に以下のような実装となりました。

from airflow.decorators import dag
from airflow.providers.google.cloud.operators.vertex_ai.pipeline_job import RunPipelineJobOperator

@dag(...)
def main():
    TEMPLATE_REPOSITORY_URL="https://asia-northeast1-kfp.pkg.dev/project-central-repository/pipelines/ml-project/v1.0.0"
    
    # kfp.RegistryClient.get_tag(...)をもとに、タグ付きのURLからバージョン付きのURLを取得するメソッド
    template_url_with_version = resoleve_template_version(TEMPLATE_REPOSITORY_URL)
    
    vertexai_pipeline = RunPipelineJobOperator(
        task_id=task_id,
        project_id=gcp_project_id,
        region=gcp_location,
        display_name=display_name,
        template_path=template_url_with_version,  # KFP Repository上のテンプレートのURL w/ version
        service_account=pipeline_service_account
    )

まとめ

取り組みの内容自体はちょっとしたものですが、ML基盤の拡張性や利便性を高められるものかと思っています。一方、当初想定したよりも特に権限周りで詰まりどころが多く、この記事の内容が誰かの参考になれば嬉しいです。

入社2ヶ月ということもあり最初は入社エントリでも書こうと思っていたのですが、せっかくなら直近取り組んできたことを書きたいと思い、本記事の執筆に至りました。

最後に、入社エントリっぽいことを簡単に書いておきたいと思います。

タイミーではデータを活用した施策が活発に行われています。 MLOpsエンジニアとして、より広く、安全で使い勝手の良いML基盤にするため、日々ワクワクしながら取り組んでいます。

We’re Hiring!

タイミーのデータエンジニアリング部・データアナリティクス部では、ともに働くメンバーを募集しています！！

現在募集中のポジションはこちらです！

「話を聞きたい」と思われた方は、是非一度カジュアル面談でお話ししましょう！

*1:Vertex AI PipelinesテンプレートがArtifact Registryに対応したのは、2022年12月です。 Release Note: https://cloud.google.com/vertex-ai/docs/release-notes#December_05_2022

*2:kfp.RegistryClientのリリースはKFP SDK 2.0.0からで、2023年6月です。Release Note: https://github.com/kubeflow/pipelines/blob/master/sdk/RELEASE.md#200