Zenn
DMM Data BlogPublicationへの投稿
🧪

Vertex AI Pipelinesにおけるローカル実行

Komekami
2024/04/22に公開
Docker
GCP
NVIDIA
Vertex AI
tech

要約

  • Vertex AI Pipelinesでローカル実行するには以下の2通りの方法がある
      1. gcloud ai custom-jobs local-runコマンドを使う方法（CustomJob単位）
      1. Kubeflow PipelinesのLocalRunnerを使う方法（CustomJob単位 or Pipeline単位）

背景

  • GCPで提供されているVertex AI Pipelinesを用いている
  • Vertex AI PipelinesでVertex AI CustomJobをつなげて動かしている
  • 毎回の実行をGCP上で行うのは、試行回数の観点において好ましくない
    • ローカル上でDockerイメージをpushして、インスタンスのマシンを割り当て、Dockerイメージをpullするのに数分はかかる
  • ローカル上のリソースを用いて、実行するようにしたい

ローカル実行

以下の2通りの方法がある

  1. gcloud ai custom-jobs local-runコマンドを使う方法（CustomJob単位）
  2. Kubeflow PipelinesのLocalRunnerを使う方法（CustomJob単位 or Pipeline単位）

1. gcloud ai custom-jobs local-runコマンドを使う方法

  • CustomJobで動かしたいDockerイメージを用意する
  • gcloud ai custom-jobs local-runコマンドを実行
コマンド例
gcloud ai custom-jobs local-run \
  --executor-image-uri=sample-custom-job:latest \ # Dockerイメージ
  --script=main.py \ # entrypoint
  --gpu \ # ローカル環境のGPUを使用
  -- \ # 以下、引数
  --model_path "gs://bucket_name/dir_name/2024-04-19/.../model_path" \
  --dt "2024-04-19" \
  --env "dev" \
  --epochs 10
出力例
Package is set to /home/jupyter/dir_name/component_name.
/usr/lib/google-cloud-sdk/platform/bundledpythonunix/lib/python3.11/subprocess.py:1010: RuntimeWarning: line buffering (buffering=1) isn't supported in binary mode, the default buffer size will be used
  self.stdin = io.open(p2cwrite, 'wb', bufsize)
/usr/lib/google-cloud-sdk/platform/bundledpythonunix/lib/python3.11/subprocess.py:1016: RuntimeWarning: line buffering (buffering=1) isn't supported in binary mode, the default buffer size will be used
  self.stdout = io.open(c2pread, 'rb', bufsize)
#0 building with "default" instance using docker driver
#1 [internal] load build definition from Dockerfile
#1 transferring dockerfile: 513B done
#1 DONE 0.0s
#2 [internal] load metadata for sample-custom-job:latest
#2 DONE 0.0s
#3 [internal] load .dockerignore
#3 transferring context: 2B done
#3 DONE 0.0s
#4 [1/4] FROM sample-custom-job:latest
#4 CACHED
#5 [internal] load build context
#5 transferring context: 637B done
#5 DONE 0.0s
#6 [2/4] RUN mkdir -m 777 -p /usr/app /home
#6 DONE 0.3s
#7 [3/4] WORKDIR /usr/app
#7 DONE 0.0s
#8 [4/4] COPY [., .]
#8 DONE 0.0s
#9 exporting to image
#9 exporting layers 0.0s done
#9 writing image sha256:5e4e2acc2618b112374cda70d7eef3e7bbbea323c5631abb68b3bc7015709ce0 done
#9 naming to docker.io/cloudai-autogenerated/main.py:20240403.04.17.51.683823 done
#9 DONE 0.0s
A training image is built.
Starting to run ...

注意点

  • ローカル実行では、CloudStorageが/gcsとしてmountされないので、直接gs://~のようにURIを指定する

    • --model_path /gcs/~ -> --model_path gs://~

  • コマンド実行時にDockerイメージ自動生成されるので、定期的にrmする必要がある

    • 公式Docによると、Dockerの知識がなくてもコンテナイメージを作成できる、という点が売りのようで、毎回の実行でイメージが生成される
$ docker images
REPOSITORY                      TAG                        IMAGE ID       CREATED         SIZE
cloudai-autogenerated/main.py   20240419.01.04.40.926317   5001268ebd42   5 minutes ago   5.61GB
cloudai-autogenerated/main.py   20240418.09.39.55.583863   da9333f40137   16 hours ago    5.61GB
cloudai-autogenerated/main.py   20240418.09.37.53.060252   aca894c902aa   16 hours ago    5.61GB
sample-custom-job               latest                     e03cdcae38a4   20 hours ago    5.61GB

2. Kubeflow PipelinesのLocalRunnerを使う方法

$ pip install kfp=="2.7.0"
$ pip install google-cloud-pipeline-components=="2.13.0"
  • CustomJobで動かしたいDockerイメージを用意する
  • Pipeline定義ファイルを用意し、実行
pipeline.py
from kfp import dsl, local
from kfp.dsl import Model, Output

local.init(runner=local.DockerRunner())

@dsl.container_component
def task(
    model_path: Output[Model],
    dt: str,
    env: str,
    epochs: int = 5,
):
    return dsl.ContainerSpec(
        image="sample-custom-job:latest", # Dockerイメージ
        command=["python3", "main.py"], # entrypoint
        # 引数
        args=[
            "--model_path",
            model_path.path,
            "--dt",
            dt,
            "--env",
            env,
            "--epochs",
            epochs,
        ],
    )

task(dt="2024-04-19", env="dev")
出力例
$ python pipeline.py
04:10:36.568 - INFO - Executing task 'task'
04:10:36.568 - INFO - Streamed logs:

    Found image 'sample-custom-job:latest'

...

Pipeline単位での実行も可能

補足

GPUを用いた実行について

  • Kubeflow PipelinesのLocalRunnerは、ローカルのGPUソフトウェア(NVIDIA CUDA Toolkit, cuDNN)に対して依存がありそう
    • GPUを用いた処理をローカル実行した際に、error code CUDA driver version is insufficient for CUDA runtime versionというエラーが出たため
from kfp import dsl, local

local.init(runner=local.DockerRunner())

@dsl.container_component
def gpu_processing():
    return dsl.ContainerSpec(
        image="gcr.io/google_containers/cuda-vector-add:v0.1",
    )

gpu_processing()
出力
$ python pipeline.py
10:39:44.956 - INFO - Executing task 'gpu-processing'
10:39:44.957 - INFO - Streamed logs:

    Found image 'gcr.io/google_containers/cuda-vector-add:v0.1'

    Failed to allocate device vector A (error code CUDA driver version is insufficient for CUDA runtime version)!
    [Vector addition of 50000 elements]
Traceback (most recent call last):
  File "/home/jupyter/gcp-batch-workflow/vertex_ai_pipelines/owl_top_adult_u2i/pipeline.py", line 39, in <module>
    gpu_processing()
  File "/opt/conda/lib/python3.10/site-packages/kfp/dsl/base_component.py", line 101, in __call__
    return pipeline_task.PipelineTask(
  File "/opt/conda/lib/python3.10/site-packages/kfp/dsl/pipeline_task.py", line 184, in __init__
    self._execute_locally(args=args)
  File "/opt/conda/lib/python3.10/site-packages/kfp/dsl/pipeline_task.py", line 199, in _execute_locally
    self._outputs = task_dispatcher.run_single_task(
  File "/opt/conda/lib/python3.10/site-packages/kfp/local/task_dispatcher.py", line 57, in run_single_task
    outputs, _ = run_single_task_implementation(
  File "/opt/conda/lib/python3.10/site-packages/kfp/local/task_dispatcher.py", line 179, in run_single_task_implementation
    raise RuntimeError(msg)
RuntimeError: Task 'gpu-processing' finished with status FAILURE
  • 一方で、gcloud ai custom-jobs local-runでの実行はDocker内で閉じているのでGPUソフトウェアが不要だったり、ローカル環境を汚さずに済みそう
    • NVIDIA Container Toolkitを用いて実行しているため(?)

おわりに

  • Vertex AI Pipelinesには、楽にローカル実行できる機能が揃っていた
  • GPUソフトウェアに対する依存を持ちたくないので、今後はgcloud ai custom-jobs local-runを使う方向でいきたい
  • ここら辺のローカル実行まわりの文献がネット上にあまりなかったので、記事にすることができてよかった

参考

https://cloud.google.com/vertex-ai/docs/training/containerize-run-code-local?hl=ja

https://www.kubeflow.org/docs/components/pipelines/v2/local-execution/

https://caddi.tech/archives/4850

GitHubで編集を提案
DMM Data Blog
DMM Data BlogPublication

DMM.comに所属するデータ分析関連組織の有志によるテックブログです。データ分析に関する知識・技術を発信します。

Discussion