Vertex AI 上の Ray クラスタでアプリケーションを開発する

Vertex AI 上の Ray クラスタに接続し、アプリケーションを開発する方法は次のとおりです。

• Ray Client の機能を含む Vertex AI SDK for Python のバージョンを使用して、Ray Client 経由で Vertex AI の Ray クラスタに接続します。インタラクティブな Python 開発環境が必要な場合は、このオプションを使用します。

  • Google Cloud コンソールの Colab Enterprise ノートブック内で Vertex AI SDK for Python を使用します。

  • Python セッション、シェル、Jupyter ノートブック内で Vertex AI SDK for Python を使用します。

• Python スクリプトを作成し、Ray Jobs API を使用して Vertex AI の Ray クラスタにそのスクリプトを送信します。プログラマティックにジョブを送信する場合は、このオプションを使用します。

始める前に、Ray on Vertex AI の概要を読み、前提条件となるすべてのツールを設定してください。

Ray Client を介して Ray クラスタに接続する

インタラクティブな Ray Client を使用するには、Vertex AI の Ray クラスタに接続します。接続環境のネットワークは、クラスタのネットワーク構成によって異なります。クラスタが公共のインターネットにアクセスできる限り、接続環境に制限はありません。つまり、クラスタの作成時に VPC ネットワークが指定されていません。ただし、クラスタが Vertex AI とピアリングされているプライベート VPC ネットワーク上にある場合は、接続環境がクラスタと同じ VPC ネットワーク上にある必要があります。

クライアント側の Ray バージョンは、クラスタの Ray バージョンと一致する必要があります。pip install "google-cloud-aiplatform[ray]" は、デフォルトでクライアントサイドに Ray バージョン 2.33 をインストールします。クラスタの Ray バージョンが 2.9 の場合は、pip install ray==2.9.3 を使用して、クライアントサイドの Ray バージョンをクラスタの Ray バージョンと一致させる必要があります。

Console

OSS Ray のベスト プラクティスの推奨事項に従い、ヘッドノードでワークロードを実行しないように、Ray ヘッドノードで論理 CPU 数を 0 に設定します。

1. Google Cloud コンソールで、[Vertex AI での Ray] ページに移動します。

   [Vertex AI での Ray] ページに移動

2. 作成したクラスタの行で、[Colab Enterprise で開く] をクリックします。

3. Colab Enterprise ノートブックが開きます。Vertex AI SDK for Python を使用して Vertex AI の Ray クラスタに接続する手順を行います。

   注: Colab Enterprise では、ユーザー認証情報を使用して認証を行います。

   • API を有効にするよう求めるダイアログ画面が表示されたら、[有効にする] をクリックします。

   • クラスタに初めて接続する場合は [接続] をクリックします。クラスタに再接続する場合は [再接続] をクリックします。ノートブックがランタイムに接続するまでに数分かかります。

   • [+作成] をクリックして、新しいノートブックを作成します。

   • をクリックして、[Ray on Vertex AI] パネルを開きます。
     既存のクラスタが表示されます。

   • クラスタを選択し、[接続] をクリックします。
     選択したクラスタに接続するコードが、開いているノートブックに表示されます。

   • その他のアクション（省略可）: Ray on Vertex AI クラスタのリストページを開くには、[Ray on Vertex AI] パネルで [Manage clusters] をクリックします。

     • クラスタを選択して、（その他の操作）メニューをクリックします。
       その他の操作が表示されます。
       

   • Getting started コードセルを実行して Vertex AI SDK for Python をインポートし、Vertex AI の Ray クラスタに接続します。

Python

OSS Ray のベスト プラクティスの推奨事項に従い、ヘッドノードでワークロードを実行しないように、Ray ヘッドノードで論理 CPU 数を 0 に設定します。

インタラクティブな Python 環境から次を実行します。

import ray

# Necessary even if aiplatform.* symbol is not directly used in your program.
from google.cloud import aiplatform
import vertex_ray

import vertexai
vertexai.init()
# The CLUSTER_RESOURCE_NAME is the one returned from vertex_ray.create_ray_cluster.
CLUSTER_RESOURCE_NAME='projects/{}/locations/{}/persistentResources/{}'.format(PROJECT_ID, LOCATION, CLUSTER_NAME)

ray.init('vertex_ray://{}'.format(CLUSTER_RESOURCE_NAME))

ここで

• LOCATION: Vertex AI の Ray クラスタに指定したロケーション。

• PROJECT_ID: Google Cloud プロジェクト ID。プロジェクト ID は、Google Cloud コンソールの [ようこそ] ページで確認できます。

• CLUSTER_NAME: クラスタの作成時に指定した Vertex AI の Ray クラスタの名前。Google Cloud コンソールに移動して、プロジェクトのクラスタ名のリストを表示します。

出力は次のようになります。

Python version:  3.10.12
Ray version: 2.33
Vertex SDK version: 1.46.0
Dashboard: xxxx-dot-us-central1.aiplatform-training.googleusercontent.com

Dashboard URL を使用して、ブラウザから Ray ダッシュボードにアクセスできます。URI の形式は https://xxxx-dot-us-central1.aiplatform-training.googleusercontent.com/ です。ダッシュボードには、クラスタ内の各マシンについて送信されたジョブ、GPU または CPU の数、ディスク容量が表示されます。

Vertex AI の Ray クラスタに接続すると、通常の OSS Ray バックエンドでの開発と同じ方法で Ray プログラムを開発できます。

@ray.remote
def square(x):
  print(x)
  return x * x

# Launch four parallel square tasks.
futures = [square.remote(i) for i in range(4)]

print(ray.get(futures))
# Returns [0, 1, 4, 9]

Ray Jobs API を使用してアプリケーションを開発する

このセクションでは、Ray Jobs API を使用して、Python プログラムを Vertex AI の Ray クラスタに送信する方法について説明します。

Python スクリプトを作成する

任意のテキスト エディタで、アプリケーションを Python スクリプトとして開発します。たとえば、次のスクリプトを my_script.py ファイルに配置します。

注: プログラムのエントリポイントのコードはヘッドノードで実行されます。関数を適切に並列化するには、@ray.remote で関数をラップする必要があります。

import ray
import time

@ray.remote
def hello_world():
    return "hello world"

@ray.remote
def square(x):
    print(x)
    time.sleep(100)
    return x * x

ray.init()  # No need to specify address="vertex_ray://...."
print(ray.get(hello_world.remote()))
print(ray.get([square.remote(i) for i in range(4)]))

Ray Jobs API を使用して Ray ジョブを送信する

Python、Ray Jobs CLI、または Ray ダッシュボードの公開アドレスを使用して、Ray ジョブを送信できます。

注: working_dir ディレクトリ内のすべてのコンテンツは Ray クラスタに送信されますが、ディレクトリが大きすぎるとタイムアウトになることがあります。送信するファイルのみを含む working_dir ディレクトリを指定することをおすすめします。

Python - クラスタ リソース名

Python 環境を使用して Ray ジョブを送信します。

import ray
import vertex_ray
from ray.job_submission import JobSubmissionClient
from google.cloud import aiplatform  # Necessary even if aiplatform.* symbol is not directly used in your program.

CLUSTER_RESOURCE_NAME='projects/{}/locations/REGION/persistentResources/{}'.format(PROJECT_ID, CLUSTER_NAME)

client = JobSubmissionClient("vertex_ray://{}".format(CLUSTER_RESOURCE_NAME))

job_id = client.submit_job(
  # Entrypoint shell command to execute
  entrypoint="python my_script.py",
  # Path to the local directory that contains the my_script.py file.
  runtime_env={
    "working_dir": "./directory-containing-my-script",
    "pip": ["numpy",
            "setuptools<70.0.0",
            "xgboost",
            "ray==CLUSTER_RAY_VERSION", # pin the Ray version to the same version as the cluster
           ]
  }
)

# Ensure that the Ray job has been created.
print(job_id)

ここで

• REGION: Vertex AI の Ray クラスタに指定したリージョン。

• PROJECT_ID: Google Cloud プロジェクト番号。 プロジェクト ID は、Google Cloud コンソールの [ようこそ] ページで確認できます。

• CLUSTER_NAME: クラスタの作成時に指定した Vertex AI の Ray クラスタの名前。Google Cloud コンソールに移動して、プロジェクトのクラスタ名のリストを表示します。

• CLUSTER_RAY_VERSION: Ray バージョンをクラスタと同じバージョンに固定します。例: 2.33.0。

Python - Ray ダッシュボード

Ray ダッシュボードのアドレスには VPC の外部からアクセスできます（公共のインターネットからもアクセスできます）。自動的に認証を取得するには、vertex_ray が必要です。

from ray.job_submission import JobSubmissionClient
import vertex_ray

DASHBOARD_ADDRESS=DASHBOARD_ADDRESS

client = JobSubmissionClient(
  "vertex_ray://{}".format(DASHBOARD_ADDRESS),
)

job_id = client.submit_job(
  # Entrypoint shell command to execute
  entrypoint="python my_script.py",
  # Path to the local directory that contains the my_script.py file
  runtime_env={
    "working_dir": "./directory-containing-my-script",
    "pip": ["numpy",
            "setuptools<70.0.0",
            "xgboost",
            "ray==CLUSTER_RAY_VERSION", # pin the Ray version to the same version as the cluster
           ]
  }
)
print(job_id)

ここで

DASHBOARD_ADDRESS: クラスタの Ray ダッシュボード アドレス。ダッシュボード アドレスは、Vertex AI SDK for Python を使用して確認できます。

Ray Jobs CLI

Ray Jobs CLI コマンドは、ピアリングされた VPC ネットワーク内でのみ使用できます。

$ ray job submit --working-dir ./ --address vertex_ray://{CLUSTER_RESOURCE_NAME} -- python my_script.py

長時間実行される Ray Job を送信した後、client.get_job_status(job_id) を使用してジョブのステータスをモニタリングする場合は、JobSubmissionClient(client = JobSubmissionClient("vertex_ray://{}".format(CLUSTER_RESOURCE_NAME)) ) を再インスタンス化して認証トークンを更新する必要があります。

VPC ピアリングとカスタム サービス アカウントのサポート

Ray on Vertex AI は、デフォルトのサービス エージェントとカスタム サービス アカウントのパブリック ネットワークで Ray Client と Ray Jobs API（JobSubmissionClient）をサポートしています。

Ray on Vertex AI の VPC ピアリングのサポート（Ray クラスタが VPC ネットワークで作成されている場合）は、次の表のとおりです。

VPC ピアリング	デフォルトのサービス エージェント	カスタム サービス アカウント
Ray Client（インタラクティブ モード）	○	×
Ray JobSubmissionClient	○	○

VPC Service Controls（VPC-SC）では追加の構成が必要です。詳細については、プライベート接続とパブリック接続をご覧ください。

次のステップ

• BigQuery と Ray on Vertex AI を使用する

• Vertex AI にモデルをデプロイして予測を取得する

• Vertex AI 上の Ray クラスタのログを表示する

• Ray クラスタを削除する