Vertex AI Workbench の構成管理と運用

技術部データ基盤チーム データサイエンティストの zaimy です。

ペパボでは、機械学習関連のコンポーネントを運用するために Google の統合 ML プラットフォームである Vertex AI の導入を進めており、研究開発部門であるペパボ研究所の研究開発環境の一部として Vertex AI Workbench を採用しました。

この記事では、次のような要件を満たしながら、Vertex AI Workbench のインスタンス¹の構成および権限の管理を行う方法と、インスタンスの起動・終了を ChatOps で運用する方法についてご紹介します。

• 社内ユーザーが用意されたインスタンスにアクセスして Jupyter を使えること
• 社内ユーザーがインスタンスの起動と停止を自由に行えること
• データ基盤を運用しているチームがインスタンスなどの構成を管理すること
  • 即ち社内ユーザーが自由にインスタンスを作ったり構成を変更したりできる状態にはしないこと

Jupyter 環境や Vertex AI Workbench の導入を検討している方の参考になれば幸いです。

1. Vertex AI Workbench の概要
2. インスタンス構成の管理
3. 権限の管理
4. インスタンスの起動・終了を ChatOps で行う
5. まとめ
6. 脚注

Vertex AI Workbench の概要

Vertex AI Workbench は、Jupyter をベースとした開発環境を運用できるサービスです。任意のリソースサイズを指定して、主要な機械学習関連のパッケージを備えたフルマネージド環境を立ち上げることができます。

BigQuery や Vertex AI のその他のコンポーネントとの統合も簡単に行うことができるため、既に Google Cloud Platform を中心に構築しているペパボのデータ基盤「Bigfoot」のデータを活用するのにも適しています。

Jupyter を使える類似のサービスとしては Google Colab があり、ペパボでもデータサイエンティストやデータエンジニアが簡単な分析を行う際に重宝していますが、それなりにリソースを必要とする計算を行う場合や、本番環境への投入を見越したプロジェクトでの開発を行う場合などは、任意のリソースの指定、構成やコードの管理ができる Vertex AI Workbench を使うことになると思います。

インスタンス構成の管理

Vertex AI Workbench のインスタンスは terraform で管理できます。 Bigfoot でも、他の多くの Google Cloud Platform のリソースと併せて、terraform による Infrastructure as Code を行っています。

具体的には次のように google_notebooks_instance resource で定義が可能です。

resource "google_notebooks_instance" "research_and_development_foo" {
  name         = "research-and-development-foo"
  location     = "us-central1-a"
  machine_type = "n1-standard-8"

  install_gpu_driver = true
  accelerator_config {
    type       = "NVIDIA_TESLA_T4"
    core_count = 2
  }
  vm_image {
    project      = "deeplearning-platform-release"
    image_family = "pytorch-1-9-cu110-notebooks-ubuntu-1804"
  }

  service_account = google_service_account.vertex_ai_workbench_research_and_development.email
}

resource "google_service_account" "vertex_ai_workbench_research_and_development" {
  account_id   = "vertex-ai-workbench-rand"
  display_name = "vertex-ai-workbench-rand"
}

権限の管理

「社内ユーザーが、用意されたノートブックにアクセス・実行できるが、ノートブックインスタンスそのものは作成・変更できない」という状態にするためには、以下のように構成する必要があります。

• インスタンスに任意のサービスアカウントを指定する
  • 指定したサービスアカウントにインスタンスからアクセスするリソースの権限を付けます
  • 指定しないと Compute Engine のデフォルトサービスアカウント（roles/editor）で動作します
• 社内ユーザーのユーザーアカウントに roles/notebooks.viewer を付ける
• 社内ユーザーのユーザーアカウントにインスタンスに設定したサービスアカウントの roles/iam.serviceAccountUser を付ける

terraform では以下のようになります。

resource "google_project_iam_binding" "vertex_ai_workbench_user" {
  project = "foo"
  role    = "roles/notebooks.viewer"

  members = ["user:foo@example.com"]
}

resource "google_service_account_iam_binding" "workbench_service_account_user_research_and_development" {
  service_account_id = google_service_account.vertex_ai_workbench_research_and_development.name
  role               = "roles/iam.serviceAccountUser"

  members = ["user:foo@example.com"]
}

インスタンスの起動・終了を ChatOps で行う

ここまでで、「社内ユーザーが用意されたインスタンスにアクセスして Jupyter を使えること」と「データ基盤を運用しているチームがインスタンスなどの構成を管理すること」を実現できましたが、「社内ユーザーがインスタンスの起動と停止を自由に行えること」は実現していません。

インスタンスの起動と停止を行うためには、roles/notebooks.admin を持ったアカウントで操作する必要がありますが、この role は新規にインスタンスを作成したり、インスタンスの構成を変更したりできるため、社内ユーザーに付与すると「データ基盤を運用しているチームがインスタンスなどの構成を管理すること」が困難になってしまいます。そこで、社内ユーザーのアカウントには権限を付けずに、ChatOps でインスタンスの起動・停止を行えるようにしました。

ワークフローを呼び出した際の Slack のログ

具体的には、Slack から GitHub Actions をトリガーする pyama86/github-actions-trigger-bot （ペパボの Slack では actions-bot と呼んでいます）を使って、次のような流れで実行します。

1. ユーザーが Slack Workflow のダイアログから対象インスタンスと操作（起動 or 停止）を選択する
2. Slack Workflow が対象インスタンスと操作をパラメータに含んで actions-bot にメンションする
3. actions-bot が GitHub Action のワークフローを起動する
4. GitHub Actions のワークフローで gcloud コマンドを実行してインスタンスを起動・停止する

ユーザーが直接 @actions-bot org/repo control_vertex_ai_workbench_notebook notebook:foo command:start のようにメンションしても操作できますが、コマンドを覚えるのが大変なので Slack Workflow で UI を作った形です。

actions-bot にパラメータとして渡した対象インスタンスと操作は、ワークフロー内で gcloud コマンドの引数として埋め込みます。

# control_vertex_ai_workbench_notebook.yml
# 認証や通知関連の処理は省略
name: Control Vertex AI Workbench Notebook
on:
  repository_dispatch:
    types: [control_vertex_ai_workbench_notebook]

jobs:
  control_vertex_ai_workbench_notebook:
    runs-on: ...
    container: ...
    steps:
    - uses: actions/checkout@v2
    - name: Control
      run: |
        gcloud notebooks instances ${{ github.event.client_payload.command }} ${{ github.event.client_payload.notebook }} --location=us-central1-c

なお、ワークフローで使用するサービスアカウントには roles/notebooks.admin を付けることになりますが、リポジトリの設定により、ワークフローのデプロイにはデータ基盤を運用しているチームのレビューを必須としているため、社内ユーザーがこのサービスアカウントを使って任意の操作を行うことはできません。

まとめ

この記事では、terraform を使って Vertex AI Workbench のインスタンスの構成および権限の管理を行う方法と、インスタンスの起動・終了を ChatOps で運用する方法についてご紹介しました。

Vertex AI Workbench のようなサービスを上手に活用することで、データエンジニア・データサイエンティスト・研究者が開発環境の用意に腐心することがなく、データの活用そのものに集中することができると考えています。ペパボではこれらのサービスを使いながら、機械学習関連の研究開発も加速させていきます。

脚注