GCPリソースへのラベル付与によるVertex AI Pipelinesのコスト可視化

🎄 本記事は ZOZO Advent Calendar 2024 シリーズ 2 の 12 日目です。
ぜひ他の記事もご覧ください。

はじめに

こんにちは。推薦基盤の関口です。ZOZOTOWNのパーソナライズ機能や推薦システム開発を担当する私たちのチームでは、複数の機械学習パイプラインをVertex AI Pipelinesで構築・運用しています。しかし、コストモニタリングに課題を感じていました。そこで本記事では、パイプライン単位でのコスト可視化を実現するための取り組みをまとめます。

課題感

単一の GCP プロジェクト上で複数のチームやアプリケーションが稼働している環境では、特定のパイプラインやアプリケーションによるコストを明確に把握するのが難しい状況です。

やりたいこと

リソースごとにラベルを付与し、Billing Reports や BigQuery を活用して、パイプライン単位でのコストを可視化します。

Vertex AI Pipelines におけるラベル付与の活用

参考：Vertex AI Pipelines によるリソースのラベル付け

PipelineJobへのラベル付与

PipelineJobにラベルを設定することで、Vertex AI Pipelinesのリソース（ジョブ全体）にラベルが自動的に適用されます。これにより、Google Cloudのコスト管理ツールを使って、リソースをラベル単位で分類・集計可能です

from google.cloud import aiplatform

# PipelineJobの設定
pl = aiplatform.PipelineJob(
    display_name=pipeline_name,
    enable_caching=enable_cache,
    template_path=path,
    parameter_values=arg_params,
    pipeline_root=pipeline_config["pipeline_root"],
    labels={
        "cost_monitoring_app_name": pipeline_name,  # アプリケーション名に基づくラベル
        "cost_monitoring_env_name": env,           # 環境情報（例: prod, staging）
        "cost_monitoring_team_name": "team_name",  # チーム情報
    },
)
pl.run()

ラベルはVertex AI Pipelinesジョブには適用されますが、サブコンポーネント（例: BigQuery、Dataflowジョブ）は自動的には伝播されません

BigQueryジョブへのラベル手動設定

BigQueryクエリを実行する際、QueryJobConfigにラベルを設定して、クエリジョブに適用します。

import os
import json
from google.cloud import bigquery

# BigQueryクエリの実行
client = bigquery.Client()

# 環境変数からPipelineラベルを取得
pipeline_labels = json.loads(os.getenv("VERTEX_AI_PIPELINES_RUN_LABELS", "{}"))

query = """
    SELECT *
    FROM `your_dataset.your_table`
    WHERE some_condition = True
"""

# BigQueryジョブのラベル設定
job_config = bigquery.QueryJobConfig(
    query_parameters=[
        bigquery.ScalarQueryParameter("param_name", "STRING", "param_value")
    ],
    labels=pipeline_labels
)

query_job = client.query(query, job_config=job_config)
query_job.result()  # クエリ実行

Dataflowジョブへのラベル手動設定

Dataflowジョブに対してもラベルを付与して、同様のコストモニタリングが可能です。

import os
import json
import datetime
from dateutil import tz
from apache_beam.options.pipeline_options import PipelineOptions

# 環境変数からラベルを取得し、更新
pipeline_labels = json.loads(os.getenv("VERTEX_AI_PIPELINES_RUN_LABELS", "{}"))
_labels.update(pipeline_labels)

# 現在時刻を使用してジョブ名を一意に設定
exec_ts = datetime.datetime.now().replace(tzinfo=tz.gettz("Asia/Tokyo"))
job_name = f'{dataflow_job_name}-{exec_ts.strftime("%Y%m%d-%H%M%s")}'

# PipelineOptionsの作成
options = PipelineOptions(
    project=project_id,  # プロジェクトID
    region=bigtable_region,  # リージョン
    job_name=job_name,  # ジョブ名
    service_account_email=dataflow_service_account,  # サービスアカウント
    subnetwork=f"regions/{bigtable_region}/subnetworks/{subnetwork}",  # サブネット
    staging_location=f"{dataflow_gcs_path}/staging",  # ステージングの場所
    temp_location=f"{dataflow_gcs_path}/temp",  # 一時ファイルの場所
    runner=worker_type,  # 実行方法（例: DataflowRunner）
    max_num_workers=max_num_workers,  # 最大ワーカー数
    machine_type=machine_type,  # ワーカーのマシンタイプ
    use_public_ips=bool_use_public_ips,  # パブリックIPの使用有無
    labels=_labels,  # ラベル（コストモニタリング等に活用）
)

Vertex AI Pipelines は Cloud Storage バケットなどのリソースに請求ラベルを伝播しません。
そのため、インフラ側の実装でリソースコストにもラベルを付与する必要があります。

Terraform を使用した Cloud Storage バケットへのラベル付与

以下は、Google Cloud Storage バケットにラベルを付与する Terraform のコード例です。この設定により、環境（Environment）、アプリケーション名（Application Name）、チーム名（Team Name）を明示的に指定し、ストレージコストもモニタリングできるようになります。
Cloud Storage バケット

resource "google_storage_bucket" "pipeline_bucket" {
  name          = "pipeline-data-storage-${local.env}" # バケット名
  location      = "US"                                # バケットのリージョン
  storage_class = "STANDARD"                          # ストレージクラス

  # ラベル設定
  labels = {
    cost_monitoring_env_name  = local.env              # 環境（dev, staging, productionなど）
    cost_monitoring_app_name  = "pipeline_name" # アプリケーション名
    cost_monitoring_team_name = "team_name"         # チーム名
  }

他のリソースへのラベル適用(一部)

Cloud Storage バケット以外のリソースにも同様にラベルを付与することで、Vertex AI Pipelines のコストをより詳細に可視化できます。以下は、その一例です。

BigQuery データセット

resource "google_bigquery_dataset" "pipeline_dataset" {
  dataset_id  = "pipeline_dataset_${local.env}"
  location    = "US"

  labels = {
    cost_monitoring_env_name  = local.env
    cost_monitoring_app_name  = "pipeline_name"
    cost_monitoring_team_name = "team_name"
  }
}

Redis インスタンス

resource "google_redis_instance" "pipeline_cache" {
  name           = "pipeline-cache-${local.env}"
  region         = "us-central1"
  tier           = "STANDARD_HA"
  memory_size_gb = 4

  labels = {
    cost_monitoring_env_name  = local.env
    cost_monitoring_app_name  = "pipeline_name"
    cost_monitoring_team_name = "team_name"
  }
}

まとめ

GCP のリソースに適切なラベルを付与することで、パイプライン単位やアプリケーション単位でのコスト管理が可能になりました。ラベル付与は、Terraform やコードレベルでの設定を通じて行われ、Billing Reports や BigQuery でのデータ分析を活用することで、詳細なコスト可視化が実現します。

これにより、システム単位でのコスト確認が可能となり、異常値発生時の原因追及を迅速化し、対応の効率化を図ることができます。また、新しいシステム構築時のコスト見積もりが容易になり、計画精度の向上にも繋がります。