TPU を使用する Google Kubernetes Engine（GKE）に、Stable Diffusion モデルを使用する Ray Serve アプリケーションをデプロイする

Standard

このガイドでは、TPU、Ray Serve、Ray Operator アドオンを使用して、Google Kubernetes Engine（GKE）に Stable Diffusion モデルをデプロイして提供する方法について説明します。

このガイドは、生成 AI をご利用のお客様、GKE の新規または既存のユーザー、ML エンジニア、MLOps（DevOps）エンジニア、プラットフォーム管理者で、Ray を使用してモデルを提供するために Kubernetes コンテナオーケストレーション機能を使用することに関心のある方を対象としています。

Ray と Ray Serve について

Ray は、AI / ML アプリケーション向けのオープンソースのスケーラブルなコンピューティングフレームワークです。Ray Serve は、分散環境でのモデルのスケーリングとサービングに使用される Ray のモデルサービングライブラリです。詳細については、Ray ドキュメントの Ray Serve をご覧ください。

TPU について

Tensor Processing Unit（TPU）は、大規模な ML モデルのトレーニングと推論を大幅に高速化するように設計された特殊なハードウェアアクセラレータです。TPU で Ray を使用すると、高性能な ML アプリケーションをシームレスにスケーリングできます。TPU の詳細については、Cloud TPU ドキュメントで Cloud TPU の概要をご覧ください。

KubeRay TPU 初期化 Webhook について

GKE は、Ray Operator アドオンの一部として、TPU Pod のスケジューリングと、JAX などのフレームワークでコンテナの初期化に必要になる TPU 環境変数を処理する、検証と変更の Webhook を提供します。KubeRay TPU Webhook は、次のプロパティを持つ TPU をリクエストする app.kubernetes.io/name: kuberay ラベルを使用して Pod を変更します。

TPU_WORKER_ID: TPU スライス内のワーカー Pod ごとに一意の整数。
TPU_WORKER_HOSTNAMES: スライス内で相互に通信する必要のあるすべての TPU ワーカーの DNS ホスト名のリスト。この変数は、マルチホストグループ内の TPU Pod にのみ挿入されます。
replicaIndex: Pod が属するワーカーグループレプリカの一意の識別子を含む Pod ラベル。これは、複数のワーカー Pod が同じレプリカに属する可能性のあるマルチホストワーカーグループに役立ちます。また、マルチホストの自動スケーリングを有効にするために Ray で使用されます。
TPU_NAME: この Pod が属する GKE TPU PodSlice を表す文字列。replicaIndex ラベルと同じ値に設定します。
podAffinity: GKE が、同じノードプールに一致する replicaIndex ラベルを持つ TPU Pod をスケジュールするようにします。これにより、GKE は単一のノードではなく、ノードプールごとにマルチホスト TPU をアトミックにスケーリングできます。

目標

TPU ノードプールを含む GKE クラスタを作成します。
TPU を使用して Ray クラスタをデプロイします。
RayService カスタムリソースをデプロイします。
Stable Diffusion モデルサーバーを操作します。

費用

このドキュメントでは、Google Cloud の次の課金対象のコンポーネントを使用します。

料金計算ツールを使うと、予想使用量に基づいて費用の見積もりを生成できます。新しい Google Cloud ユーザーは無料トライアルをご利用いただける場合があります。

このドキュメントに記載されているタスクの完了後、作成したリソースを削除すると、それ以上の請求は発生しません。詳細については、クリーンアップをご覧ください。

始める前に

Cloud Shell には、kubectl、gcloud CLI など、このチュートリアルに必要なソフトウェアがプリインストールされています。Cloud Shell を使用しない場合は、gcloud CLI をインストールします。

Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.

Install the Google Cloud CLI.

To initialize the gcloud CLI, run the following command:

gcloud init

Note: If you installed the gcloud CLI previously, make sure you have the latest version by running

gcloud components
      update

Create or select a Google Cloud project.

Create a Google Cloud project:
```
gcloud projects create PROJECT_ID
```
Replace PROJECT_ID with a name for the Google Cloud project you are creating.
Select the Google Cloud project that you created:
```
gcloud config set project PROJECT_ID
```
Replace PROJECT_ID with your Google Cloud project name.

Make sure that billing is enabled for your Google Cloud project.

Enable the GKE API:

gcloud services enable container.googleapis.com

Install the Google Cloud CLI.

To initialize the gcloud CLI, run the following command:

gcloud init

Note: If you installed the gcloud CLI previously, make sure you have the latest version by running

gcloud components
      update

Create or select a Google Cloud project.

Create a Google Cloud project:
```
gcloud projects create PROJECT_ID
```
Replace PROJECT_ID with a name for the Google Cloud project you are creating.
Select the Google Cloud project that you created:
```
gcloud config set project PROJECT_ID
```
Replace PROJECT_ID with your Google Cloud project name.

Make sure that billing is enabled for your Google Cloud project.

Enable the GKE API:

gcloud services enable container.googleapis.com

Grant roles to your user account. Run the following command once for each of the following IAM roles: roles/container.clusterAdmin, roles/container.admin
```
gcloud projects add-iam-policy-binding PROJECT_ID --member="USER_IDENTIFIER" --role=ROLE
```
- Replace PROJECT_ID with your project ID.
- Replace USER_IDENTIFIER with the identifier for your user account. For example, user:myemail@example.com.
- Replace ROLE with each individual role.

十分な割り当てを確保する

Google Cloud プロジェクトの Compute Engine リージョンまたはゾーンに十分な TPU 割り当てがあることを確認します。詳細については、Cloud TPU のドキュメントで十分な TPU と GKE の割り当てを確保するをご覧ください。また、次の割り当ての増加が必要になる場合があります。

Persistent Disk SSD（GB）
使用中の IP アドレス

環境を準備する

環境を準備する手順は次のとおりです。

Google Cloud コンソールで「Cloud Shell をアクティブにする」をクリックして、Google Cloud コンソールから Cloud Shell セッションを起動します。Google Cloud コンソールの下部ペインでセッションが起動します。
環境変数を設定します。
```
export PROJECT_ID=PROJECT_ID
export CLUSTER_NAME=ray-cluster
export COMPUTE_REGION=us-central2-b
export CLUSTER_VERSION=CLUSTER_VERSION
```
次のように置き換えます。
- PROJECT_ID: Google Cloud プロジェクト ID。
- CLUSTER_VERSION: 使用する GKE のバージョン。1.30.1 以降にする必要があります。

GitHub リポジトリのクローンを作成します。

git clone https://github.com/GoogleCloudPlatform/kubernetes-engine-samples

作業ディレクトリを変更します。

cd kubernetes-engine-samples/ai-ml/gke-ray/rayserve/stable-diffusion

TPU ノードプールを使用してクラスタを作成する

TPU ノードプールを含む GKE Standard クラスタを作成します。

Ray Operator を有効にして Standard モードのクラスタを作成します。

gcloud container clusters create ${CLUSTER_NAME} \
    --addons=RayOperator \
    --machine-type=n1-standard-8 \
    --cluster-version=${CLUSTER_VERSION} \
    --location=${COMPUTE_REGION}

単一ホストの TPU ノードプールを作成します。

gcloud container node-pools create tpu-pool \
    --location=${COMPUTE_REGION} \
    --cluster=${CLUSTER_NAME} \
    --machine-type=ct4p-hightpu-4t \
    --num-nodes=1 \
    --tpu-topology=2x2x1

Standard モードで TPU を使用するには、次のものを選択する必要があります。

TPU アクセラレータの容量がある Compute Engine ロケーション
TPU と互換性のあるマシンタイプ
TPU PodSlice の物理トポロジ

TPU を使用して RayCluster リソースを構成する

TPU ワークロードを準備するように RayCluster マニフェストを構成します。

TPU `nodeSelector` を構成する

GKE は Kubernetes nodeSelectors を使用して、TPU ワークロードが適切な TPU トポロジとアクセラレータでスケジュールされるようにします。TPU nodeSelector の選択の詳細については、GKE Standard に TPU ワークロードをデプロイするをご覧ください。

ray-cluster.yaml マニフェストを更新して、2 x 2 x 1 トポロジの v4 TPU PodSlice で Pod をスケジュールします。

nodeSelector:
  cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice
  cloud.google.com/gke-tpu-topology: 2x2x1

TPU コンテナリソースを構成する

TPU アクセラレータを使用するには、RayCluster マニフェスト workerGroupSpecs の TPU コンテナフィールドで google.com/tpu リソース limits と requests を構成して、GKE が各 Pod に割り当てる TPU チップの数を指定する必要があります。

リソースの上限とリクエストを指定して ray-cluster.yaml マニフェストを更新します。

resources:
  limits:
    cpu: "1"
    ephemeral-storage: 10Gi
    google.com/tpu: "4"
    memory: "2G"
   requests:
    cpu: "1"
    ephemeral-storage: 10Gi
    google.com/tpu: "4"
    memory: "2G"

ワーカーグループ `numOfHosts` を構成する

KubeRay v1.1.0 では、RayCluster カスタムリソースに numOfHosts フィールドが追加され、ワーカーグループのレプリカごとに作成する TPU ホストの数を指定します。マルチホストワーカーグループの場合、レプリカは個々のワーカーではなく PodSlice として扱われ、レプリカごとに numOfHosts ワーカーノードが作成されます。

ray-cluster.yaml マニフェストを次のように更新します。

workerGroupSpecs:
  # Several lines omitted
  numOfHosts: 1 # the number of "hosts" or workers per replica

RayService カスタムリソースを作成する

RayService カスタムリソースを作成します。

次のマニフェストを確認します。

apiVersion: ray.io/v1
kind: RayService
metadata:
  name: stable-diffusion-tpu
spec:
  serveConfigV2: |
    applications:
      - name: stable_diffusion
        import_path: ai-ml.gke-ray.rayserve.stable-diffusion.stable_diffusion_tpu:deployment
        runtime_env:
          working_dir: "https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/archive/refs/heads/main.zip"
          pip:
            - diffusers==0.7.2
            - flax
            - jax[tpu]==0.4.11
            - -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
            - fastapi
  rayClusterConfig:
    rayVersion: '2.9.0'
    headGroupSpec:
      rayStartParams: {}
      template:
        spec:
          containers:
          - name: ray-head
            image: rayproject/ray-ml:2.9.0-py310
            ports:
            - containerPort: 6379
              name: gcs
            - containerPort: 8265
              name: dashboard
            - containerPort: 10001
              name: client
            - containerPort: 8000
              name: serve
            resources:
              limits:
                cpu: "2"
                memory: "8G"
              requests:
                cpu: "2"
                memory: "8G"
    workerGroupSpecs:
    - replicas: 1
      minReplicas: 1
      maxReplicas: 10
      numOfHosts: 1
      groupName: tpu-group
      rayStartParams: {}
      template:
        spec:
          containers:
          - name: ray-worker
            image: rayproject/ray-ml:2.9.0-py310
            resources:
              limits:
                cpu: "100"
                ephemeral-storage: 20Gi
                google.com/tpu: "4"
                memory: 200G
              requests:
                cpu: "100"
                ephemeral-storage: 20Gi
                google.com/tpu: "4"
                memory: 200G
          nodeSelector:
            cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice
            cloud.google.com/gke-tpu-topology: 2x2x1

このマニフェストでは、1 つのヘッドノードと 2 x 2 x 1 トポロジの TPU ワーカーグループを含む RayCluster リソースを作成する RayService カスタムリソースを記述しています。つまり、各ワーカーノードには 4 つの v4 TPU チップが割り当てられます。

TPU ノードは、2 x 2 x 1 トポロジの単一の v4 TPU PodSlice に属しています。マルチホストワーカーグループを作成するには、gke-tpu nodeSelector 値、google.com/tpu コンテナの上限とリクエスト、numOfHosts 値をマルチホスト構成に置き換えます。TPU マルチホストトポロジの詳細については、Cloud TPU ドキュメントでシステムアーキテクチャをご覧ください。

マニフェストをクラスタに適用します。
```
kubectl apply -f ray-service-tpu.yaml
```
RayService リソースが実行されていることを確認します。
```
kubectl get rayservices
```
出力は次のようになります。
```
NAME                   SERVICE STATUS   NUM SERVE ENDPOINTS
stable-diffusion-tpu   Running          2
```
この出力の SERVICE STATUS 列の Running は、RayService リソースの準備が完了したことを示します。

（省略可）Ray ダッシュボードを表示する

Ray Serve デプロイとその関連ログは、Ray ダッシュボードで確認できます。

Ray ヘッドサービスから Ray ダッシュボードへのポート転送セッションを確立します。
```
kubectl port-forward svc/stable-diffusion-tpu-head-svc 8265:8265
```
ウェブブラウザで http://localhost:8265/ にアクセスします。
[Serve] タブをクリックします。

モデルサーバーにプロンプトを送信する

Ray ヘッドサービスから Serve エンドポイントへのポート転送セッションを確立します。
```
kubectl port-forward svc/stable-diffusion-tpu-serve-svc 8000
```
新しい Cloud Shell セッションを開きます。
Stable Diffusion モデルサーバーにテキストから画像への変換を行うプロンプトを送信します。
```
python stable_diffusion_tpu_req.py  --save_pictures
```
Stable Diffusion 推論の結果は、diffusion_results.png という名前のファイルに保存されます。

クリーンアップ

プロジェクトを削除する

注意: プロジェクトを削除すると、次のような影響があります。

プロジェクト内のすべてのものが削除されます。このドキュメントのタスクで既存のプロジェクトを使用した場合、それを削除すると、そのプロジェクトで行った他の作業もすべて削除されます。
カスタムプロジェクト ID が失われます。このプロジェクトを作成したときに、将来使用するカスタムプロジェクト ID を作成した可能性があります。そのプロジェクト ID を使用した URL（たとえば、appspot.com）を保持するには、プロジェクト全体ではなくプロジェクト内の選択したリソースだけを削除します。

複数のアーキテクチャ、チュートリアル、クイックスタートを実施する予定がある場合は、プロジェクトを再利用すると、プロジェクトの割り当て上限を超えないようにすることができます。

Delete a Google Cloud project:

gcloud projects delete PROJECT_ID

リソースを個別に削除する

クラスタを削除するには、次のように入力します。

gcloud container clusters delete ${CLUSTER_NAME}

次のステップ

Ray on Kubernetes について学習する。
KubeRay のドキュメントを確認する。
Google Cloud に関するリファレンスアーキテクチャ、図、ベストプラクティスを確認する。Cloud アーキテクチャセンターをご覧ください。