Zenn
🐺

VertexAIのCustomContainerTrainingJobでTrainする

2024/11/15に公開
Google Cloud
Vertex AI
LLM
tech

Vertex AI の CustomContainerTrainingJob を使ってみる

Vertex AI では自前で作成したコンテナをArtifact Registryから指定して使えるので触ってみる

https://cloud.google.com/python/docs/reference/aiplatform/1.0.1/google.cloud.aiplatform.CustomContainerTrainingJob

Vertex AI とは

Vertex AI は、MLモデルとAIアプリケーションのトレーニングとデプロイを行う機械学習プラットフォームで
データエンジニアリング、データサイエンス、MLエンジニアリングのワークフローを統合している。

https://cloud.google.com/vertex-ai/docs/start/introduction-unified-platform?hl=ja

ステップバイステップで実施

公式リポジトリにステップバイステップで実施できそうなREADMEがあるので
そちらを参考にしながら進めてみる

※今回は、あくまでモデルのtrainを実施したいので対象となるステップを参考にして進めています。
主に「Create a custom training job with your container」の部分

TENSORBOARD_NAME = "[your-tensorboard-name]"

if (
    TENSORBOARD_NAME == ""
    or TENSORBOARD_NAME is None
    or TENSORBOARD_NAME == "[your-tensorboard-name]"
):
    TENSORBOARD_NAME = PROJECT_ID + "-tb-" + UUID

tensorboard = aiplatform.Tensorboard.create(
    display_name=TENSORBOARD_NAME, project=PROJECT_ID, location=LOCATION
)
TENSORBOARD_RESOURCE_NAME = tensorboard.gca_resource.name
print("TensorBoard resource name:", TENSORBOARD_RESOURCE_NAME)

https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/tensorboard/tensorboard_custom_training_with_custom_container.ipynb

流れ

Artifact Registryに登録する

まずは、作成したdockerイメージをpushする

https://cloud.google.com/artifact-registry/docs/overview?hl=ja

cliを実行するpythonを実装する

google.cloud.aiplatformを使ってCLIを作成していきます。
ここではローカルから実行する想定です。

https://pypi.org/project/google-cloud-aiplatform/

import google.cloud.aiplatform as aiplatform

aiplatform.init(
    project=project_id,
    location=location,
    experiment=experiment,
    experiment_tensorboard=False,
    staging_bucket="gs://my-bucket",
)

job = aiplatform.CustomContainerTrainingJob(
    display_name="my-training-job",
    container_uri="us-docker.pkg.dev/vertex-ai/training/tf-cpu.2-2:latest",
)

job.run(
    machine_type=machine_type,
    accelerator_type=accelerator_type,
    accelerator_count=accelerator_count,
    sync=False,
    environment_variables={"HOGE": os.environ["HOGE"]},
)

python-aiplatformリポジトリにもサンプル実装が用意されています。

https://github.com/googleapis/python-aiplatform/blob/11e53784398b4db210c04b778a4ff39e0b9ee41a/samples/model-builder/create_training_pipeline_custom_container_job_sample.py#L43

aiplatform.init

いわゆる初期化設定

https://github.com/googleapis/python-aiplatform/tree/11e53784398b4db210c04b778a4ff39e0b9ee41a?tab=readme-ov-file#initialization

aiplatform.CustomContainerTrainingJob

Artifact Registoryにpushしたイメージを使用したJob設定ができる

Requiredなparameter

  • display_name: TrainingPipelineの名前
  • container_uri: GCR内にあるTrainingContainerのイメージuri

https://cloud.google.com/python/docs/reference/aiplatform/latest/google.cloud.aiplatform.CustomContainerTrainingJob#google_cloud_aiplatform_CustomContainerTrainingJob

aiplatform.CustomContainerTrainingJob.run

  • machine_type: トレーニングで使用するマシンを指定
  • accelerator_type: ハードウェアアクセラレータのタイプを指定
  • accelerator_count: ワーカレプリカにアタッチするアクセラレータの数を指定
  • sync: このメソッドを同期的に実行するかどうか
  • environment_variables: コンテナに渡す環境変数 ※最大10個の環境変数を指定できるらしい。

https://cloud.google.com/python/docs/reference/aiplatform/latest/google.cloud.aiplatform.CustomContainerTrainingJob#google_cloud_aiplatform_CustomContainerTrainingJob_run

scriptを実行する

最後に作成したpythonを実行します。
argparseなど導入してCLI化しておくと使い回せて便利です。

実行後はGCPのVertexAIのコンソール画面を確認すると実行ログが確認できます。

Discussion