Vertex AI Pipelinesを作って動かす

はじめに

私は業務でVertex AI Pipelinesを使ったパイプラインを作っています．日ごろ作っているパイプラインを，作って実行するまでのプロセスについて記載します．

お題

今回は簡単なパイプラインを作ります．
行なう処理は，Big Queryにあるデータセットからデータを抽出し保存．保存したものを読み取って，printするものです．

全体図

Vertex AI Pipelinesにてパイプラインを実行するにあたり，以下のような構成を作っています．

• Cloud Buildにて各コンポーネントのビルドとパイプラインのコンパイル
• Artifact Registryに，ビルドしたコンポーネントのDockerイメージとパイプライン定義ファイルを保存
• ローカルPCから，パイプラインの実行を指示

開発環境

Pythonを使ってパイプラインを構築します．環境はpyenvとPoetryを用いており，パイプライン作成にあたって，以下メインとなる2つのライブラリをインストールします．

• kfp (^2.0.0) ドキュメント
• google-cloud-aiplatform ドキュメント

フォルダ構成

詳細はこちらのGithubに公開しております．
フォルダ構成は以下のようにしています．

フォルダ構成

<simple-pipeline>
 ├── .cloudbuild
 │   └── deployment.yaml
 ├── components
 │   ├── bq-component
 │   │   ├── src
 │   │   │   ├── __init__.py
 │   │   │   └── main.py
 │   │   ├── component.yaml
 │   │   ├── dependencies.txt
 │   │   ├── Dockerfile
 │   │   ├── Makefile
 │   │   ├── poetry.lock
 │   │   └── pyproject.toml #(root下pyproject.tomlのシンボリックリンク)
 │   └── display-component
 │       ├── ...
 ├── deployment
 │   ├── build.sh
 │   ├── deployment_component.sh
 │   ├── deployment_pipeline.yaml
 │   └── exec.sh
 ├── pipeline
 │   ├── src
 │   │   ├── __init__.py
 │   │   ├── compile_pipeline.py
 │   │   ├── exec_pipeline.py
 │   │   └── simple_pipeline.py
 │   ├── dependencies.txt
 │   ├── Dockerfile
 │   ├── Makefile
 │   ├── poetry.lock
 │   └── pyproject.toml #(root下pyproject.tomlのシンボリックリンク)
 ├── (.env)
 ├── .gcloudignore
 ├── .gitignore
 ├── poetry.lock
 └── pyproject.toml

ポイントは以下です

• コンポーネントは，閉じた環境で管理できるようにするため，基本的には各ディレクトリを設けて，そこでコードを作っています．
• パイプラインの入出力を定義するため，各コンポーネントのディレクトリにはcomponent.yamlを設けています．(※ただし，kfp^2.0.0が正式リリースされたことで，今後は使わない方向で考えています(参考))．
• pipeline/src/simple_pipeline.pyにて，コンポーネント同士のつなぎこみを行なっています．このとき，@dsl.pipeline()のように，kfpのメソッドをオーバーライドすることで簡単にパイプラインを定義することができます．
• kfp^2.0.0より，パイプラインの定義ファイルをArtifact Registryに格納する機能が搭載されました．これまではGCSに保存していたのですが，Artifact Registryに保存できるようになったことで，バージョン管理ができるようになりました．

ビルド

ここまで構成できれば，あとはコマンドを実行するだけです．以下コマンドを使って，Cloud Buildにてビルドを行ないます．パイプラインのコンポーネントとなるDockerイメージのビルドと，パイプラインのコンパイルを行なって，成果物をArtifact Registryに格納します．

#!/usr/bin/env bash
export $(cat .env | xargs)

gcloud builds submit \
    --config .cloudbuild/deployment.yaml\
    --region=asia-northeast1\
	--substitutions=TAG_NAME="${TAG_NAME}",\
_COMPILE_FILE="${COMPILE_FILE}",\
_PIPELINE_NAME="${PIPELINE_NAME}",\
_REGISTRY_PATH="${REGISTRY_PATH}"

deployment.yaml

# @format

steps:
  - name: gcr.io/cloud-builders/docker
    entrypoint: bash
    args: ["deployment/deployment_component.sh"]
    env:
      - GCP_PROJECT_ID=${PROJECT_ID}
      - DOCKER_BUILDKIT=1
      - TAG=${TAG_NAME}
    id: Build & push components docker images

  - name: "gcr.io/cloud-builders/gcloud"
    args:
      [
        "builds",
        "submit",
        "--config=deployment/deployment_pipeline.yaml",
        "--region=${LOCATION}",
        "--substitutions=_TAG=${TAG_NAME},\
        _COMPILE_FILE=${_COMPILE_FILE},\
        _PIPELINE_NAME=${_PIPELINE_NAME},\
        _REGISTRY_PATH=${_REGISTRY_PATH}"
      ]
    dir: .
    id: Compile pipeline
options:
  logging: CLOUD_LOGGING_ONLY

実行

コンパイルして格納したパイプライン定義ファイルを指定(registry_template_path)して，パイプラインを実行します．

#!/usr/bin/env bash
export $(cat .env | xargs)

poetry run python pipeline/src/exec_pipeline.py \
    --project="${PROJECT_ID}" \
    --location="${LOCATION}" \
    --pipeline_root="${PIPELINE_ROOT}" \
    --registry_template_path="${REGISTRY_PATH}/${PIPELINE_NAME}/${TAG}" \
    --table_id="${TABLE_ID}"

exec_pipeline.py内でgoogle-cloud-aiplatformのPipelineJobを実行しています．

exec_pipeline.py

from google.cloud import aiplatform

...
aiplatform.init(
    project=project,
    location=location,
)

job = aiplatform.PipelineJob(
    display_name="simple_pipeline",
    template_path=template_path,
    pipeline_root=pipeline_root,
    parameter_values={
        "table_id": table_id,
    },
)
job.submit(
    service_account=service_account,
)
...

実行結果

以下のようにパイプラインが動き，最終的にデータフレームをプリントしてくれています．

まとめ

本記事では，簡単なパイプラインをVertex AI Pipelinesに構築しました．何番煎じかわかりませんが，普段行なっている業務のアウトプットとして記載しました．業務では，パイプラインの実行は定期的に行なうことができるように，Cloud FunctionsとCloud Schedulerを使っていたり，githubでのタグ付けをトリガにしてビルドするようにしたりと，よりCI/CDを充実させています．今後その旨を記載できればと思います．

参考

• kfp (^2.0.0)
• google-cloud-aiplatform
• Vertex Pipelinesによる機械学習パイプラインの実行