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

このガイドでは、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 をインストールします。

  1. 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.
  2. Install the Google Cloud CLI.

  3. To initialize the gcloud CLI, run the following command: 

    gcloud init

  4. 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.

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

  6. Enable the GKE API: 

    gcloud services enable container.googleapis.com
  7. Install the Google Cloud CLI.

  8. To initialize the gcloud CLI, run the following command: 

    gcloud init

  9. 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.

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

  11. Enable the GKE API: 

    gcloud services enable container.googleapis.com

  12. 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: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 アドレス

環境を準備する

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

  1. Google Cloud コンソールで Cloud Shell 有効化アイコン「Cloud Shell をアクティブにする」をクリックして、Google Cloud コンソールから Cloud Shell セッションを起動します。Google Cloud コンソールの下部ペインでセッションが起動します。

  2. 環境変数を設定します。

    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 以降にする必要があります。

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

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

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

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

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

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

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

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

  2. 単一ホストの 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 カスタム リソースを作成します。

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

    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 ドキュメントでシステム アーキテクチャをご覧ください。

  2. マニフェストをクラスタに適用します。

    kubectl apply -f ray-service-tpu.yaml

  3. RayService リソースが実行されていることを確認します。

    kubectl get rayservices

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

    NAME                   SERVICE STATUS   NUM SERVE ENDPOINTS
stable-diffusion-tpu   Running          2

    この出力の SERVICE STATUS 列の Running は、RayService リソースの準備が完了したことを示します。

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

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

  1. Ray ヘッドサービスから Ray ダッシュボードへのポート転送セッションを確立します。

    kubectl port-forward svc/stable-diffusion-tpu-head-svc 8265:8265

  2. ウェブブラウザで http://localhost:8265/ にアクセスします。

  3. [Serve] タブをクリックします。

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

  1. Ray ヘッドサービスから Serve エンドポイントへのポート転送セッションを確立します。

    kubectl port-forward svc/stable-diffusion-tpu-serve-svc 8000

  2. 新しい Cloud Shell セッションを開きます。

  3. Stable Diffusion モデルサーバーにテキストから画像への変換を行うプロンプトを送信します。

    python stable_diffusion_tpu_req.py  --save_pictures

    Stable Diffusion 推論の結果は、diffusion_results.png という名前のファイルに保存されます。

    Stable Diffusion で生成された画像。8 つのセクション（緑色の椅子、家の外に立っている男性、通りを歩くロボット、テーブルに座っている家族、公園を歩く医者、飛んでいるドラゴン、日本風のクマの肖像画、滝）で構成されています。

クリーンアップ

プロジェクトの削除

    Delete a Google Cloud project:

    gcloud projects delete PROJECT_ID

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

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

gcloud container clusters delete ${CLUSTER_NAME}

次のステップ