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 をアトミックにスケーリングできます。

環境を準備する

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

  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

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