Vertex AI Pipelinesに入門する

Vertex AI Pipelinesとは

Vertex AI Pipelines では、ML パイプラインを使用して ML ワークフローをオーケストレートすることで、サーバーレス方式で機械学習（ML）システムの自動化、モニタリング、管理を行うことができます。
Vertex AI Pipelines の概要 参照

Vertex AI Pipelinesの構成要素

• コンポーネント:
  • コンポーネントは、パイプライン内で実行される個々のタスクを表します。これにはデータの前処理、モデルの訓練、モデルの評価、データの保存などが含まれます。各コンポーネントは、特定の入力を受け取り、処理を実行した後に出力を生成します。コンポーネントは再利用可能であり、異なるパイプライン間で共有することができます。
• パイプライン:
  • パイプラインは、複数のコンポーネントを順序付けて実行するプロセスを定義します。パイプラインを使用すると、機械学習モデルの訓練やデプロイメントなど、複雑なワークフローを自動化できます。パイプラインは、特定のタスクが完了した後に次のタスクが開始されるように、コンポーネント間の依存関係を管理します。

いつも通りにDockerコンテナ（①）とその中で実行するPythonファイル（②）を用意します。それらを一つのコンポーネント（③）として定義します。そして、複数のコンポーネントをつなげたパイプライン（④）を定義して、コンパイルしてから実行します。

Vertex AI Pipelinesのはじめ方

簡単な例として、数字をaddしてからmultipyするパイプラインを作成します。

コードはこちら

0. ディレクトリ構成例

.
├── README.md
├── pyproject.toml
├── Dockerfile # 管理用
├── src
│   └── simple
│       ├── components
│       │   ├── add
│       │   │   ├── Dockerfile
│       │   │   ├── add.py
│       │   │   └── add_component.py
│       │   └── multiply
│       │       ├── Dockerfile
│       │       ├── multiply.py
│       │       └── multiply_component.py
│       └── pipelines
│           └── simple_pipeline.py
└── tests

1. Dockerイメージ準備

• ファイル準備

Dockerfile

FROM python:3.11

# install python package
COPY pyproject.toml ./
RUN pip install poetry
RUN poetry config virtualenvs.create false \
    && poetry install --no-root --with test,lint,simple

# mount dir
RUN mkdir -p /opt/mnt
WORKDIR /opt/mnt

# copy all files into container
COPY . /opt/mnt

# expose port
EXPOSE 8888

pyproject.toml

[tool.poetry]
name = "vap-example"
version = "0.1.0"
description = ""
authors = ["nokoxxx1212"]
packages = [
    { include = "src" },
]

[tool.poetry.dependencies]
python = ">=3.9,<3.12"
jupyterlab = "^4.0.8"

[tool.poetry.group.test]
optional = true
[tool.poetry.group.test.dependencies]
pytest = "^7.4.3"

[tool.poetry.group.lint]
optional = true
[tool.poetry.group.lint.dependencies]
isort = "^5.12.0"
black = "^23.10.1"
pyproject-flake8 = "^6.1.0"

[tool.poetry.group.simple]
optional = true
[tool.poetry.group.simple.dependencies]
google-cloud-pipeline-components = "^2.9.0"
python-dotenv = "^1.0.1"

[tool.pytest.ini_options]
testpaths = ["tests",]

• コンテナビルド

# pyproject.tomlを移動（共通で利用）
$ cp ../../../../pyproject.toml .

# ビルド
$ docker build -t gcr.io/your-project-id/add_docker .

2. Pythonスクリプト準備

• ファイル準備

add.py

import argparse

def add(x, y):
    ans = x + y
    print(ans)
    return ans

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Add two numbers')
    parser.add_argument('--x', type=int, required=True, help='first number')
    parser.add_argument('--y', type=int, required=True, help='second number')
    args = parser.parse_args()
    add(args.x, args.y)

multiply.py

import argparse

def multiply(z, w):
    ans = z * w
    print(ans)
    return ans

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Add two numbers')
    parser.add_argument('--z', type=int, required=True, help='first number')
    parser.add_argument('--w', type=int, required=True, help='second number')
    args = parser.parse_args()
    multiply(args.z, args.w)

• 動作確認

$ docker run gcr.io/your-project-id/add_docker python add.py --x 1 --y 2

# 開発中は `COPY . /opt/mnt` をコメントアウトして、下記でマウントしてランする
$ docker run -itd -v $(pwd):/opt/mnt -p 18888:8888 gcr.io/your-project-id/add_docker
$ docker ps
$ docker exec -it XXX /bin/bash

• リポジトリにpush

$ docker push gcr.io/your-project-id/add_docker

3. Component準備

• ファイル準備

add_component.py

from kfp import dsl
from kfp.v2.dsl import ContainerSpec

# コンポーネントを定義
@dsl.container_component
def add_component() -> ContainerSpec:
    return ContainerSpec(
        image='gcr.io/your-project-id/add_docker',
        command=["python3", "add.py", "--x", "1", "--y", "2"]
    )

4. Pipeline準備

• ファイル準備

simple_pipeline.py

from google.cloud import aiplatform
from kfp import dsl
from kfp.v2 import compiler
from dotenv import load_dotenv
import os
from src.simple.components.add.add_component import add_component
from src.simple.components.multiply.multiply_component import multiply_component


# パイプラインを定義
@dsl.pipeline(
    name='simple-pipeline',
    description='A pipeline to add and multiply numbers'
)
def simple_pipeline():
    # コンポーネントのインスタンスを作成
    add_task = add_component()
    multiply_task = multiply_component()
    multiply_task.after(add_task)

# パイプラインを実行
def run_pipeline(project: str, location: str, pipeline_root: str):
    # Vertex AI clientを初期化
    aiplatform.init(project=project, location=location)

    # パイプラインを実行
    job = aiplatform.PipelineJob(
        display_name='simple-pipeline',
        template_path='simple_pipeline_job.json',
        pipeline_root=pipeline_root,
    )
    job.run()

if __name__ == '__main__':
    # パイプラインを実行するためのパラメータ
    env_path = f".env.{os.getenv('APP_ENV', 'dev')}"
    load_dotenv(dotenv_path=env_path)
    PROJECT_ID = os.getenv('PROJECT_ID')
    LOCATION = os.getenv('LOCATION')
    PIPELINE_ROOT = os.getenv('PIPELINE_ROOT')

    # パイプラインをコンパイル
    compiler.Compiler().compile(
        pipeline_func=simple_pipeline,
        package_path='simple_pipeline_job.json'
    )

    # パイプラインを実行
    run_pipeline(PROJECT_ID, LOCATION, PIPELINE_ROOT)

.env.dev

PROJECT_ID = 'your-project-id'
LOCATION = 'us-central1'
PIPELINE_ROOT = 'gs://your-bucket/kh/pipeline_root'

• GCSバケット構築

5. Pipeline実行

# ルート直下Dockerfileのコンテナ内などから実行
# pathを通す
$ poetry install

# GCP認証。サービスアカウントのクレデンシャルファイルを配置して下記を実行する
$ export GOOGLE_APPLICATION_CREDENTIALS="/path/to/your/service-account-file.json"

# パイプライン実行
$ python simple_pipeline.py

Tips

• サービスアカウントは何が使われている？
  • Vertex AI Pipelinesで実行されるコンテナの中で使われるサービスアカウントは、パイプライン実行の際に指定するサービスアカウントではない。GCP側が用意したサービスアカウント（サービスエージェント）である Vertex AI Custom Code Service Agent と Vertex AI Service Agent が使われる。詳細は Vertex AI Pipelines のサービスアカウントで少しつまずいたので整理した