GKE で Ray を使用してマルチホスト TPU vLLM 推論をデプロイする

1. はじめに

この Codelab では、Google Cloud TPU を使用して、Google Kubernetes Engine（GKE）に高性能のマルチホスト vLLM（Virtual Large Language Model）推論サービスをデプロイする方法について説明します。Ray を使用して分散推論を構成し、LeaderWorkerSet を使用して GKE でワークロードをネイティブに管理します。

このチュートリアルでは、Qwen 30B などの大規模モデルをサービングするための本番環境の設定をシミュレートします。

演習内容

  • アクセラレータ トラフィック用のカスタム VPC ネットワークを作成します。
  • Ray Operator と GCS Fuse CSI ドライバを使用して GKE クラスタをプロビジョニングします。
  • モデルの読み込みを高速化するために GCS Rapid Cache を初期化します。
  • 予約済みの容量を使用して、マルチホスト TPU v6e ノードプールをプロビジョニングします。
  • モデルの重みに安全にアクセスできるように Workload Identity を構成します。
  • 30B パラメータ モデルをサービングする vLLM エンジンをデプロイしてテストします。

必要なもの

  • 課金を有効にした Google Cloud プロジェクト
  • TPU v6e リソース（32 チップ、ct6e-standard-4t）のGoogle Cloud 予約
  • ソースバケットからモデルの重みをコピーする権限。
  • gcloud、kubectl、helm がインストールされた Cloud Shell またはローカル ターミナル。
  • 推定所要時間: 60 分
  • 推定費用: 60 ドル未満（すぐに削除する場合）。

2. 始める前に

Google Cloud プロジェクトを作成または選択する

  1. Google Cloud コンソールで、Google Cloud プロジェクトを作成または選択します。
  2. Cloud プロジェクトで課金が有効になっていることを確認します。

Cloud Shell の起動

  1. Google Cloud コンソールの上部にある [Cloud Shell をアクティブにする] をクリックします。
  2. 認証を検証します。
gcloud auth list
  1. プロジェクトを確認します。
gcloud config get project
  1. 必要に応じて設定します。
export PROJECT_ID=<YOUR_PROJECT_ID>
gcloud config set project $PROJECT_ID

環境変数の設定

コマンドの実行を簡単にするため、シェルで次の変数を定義します。<YOUR_ZONE> は割り当てられた TPU ゾーンに、<YOUR_RESERVATION_NAME> は予約 ID に置き換えてください。ゲートモデルの重みをダウンロードするには、Hugging Face ユーザー アクセス トークンを作成する必要があります。作成したら、<YOUR_HUGGING_FACE_TOKEN> を新しく作成したトークンに置き換えます。

export PROJECT_ID=$(gcloud config get-value project)
export PROJECT_NUMBER=$(gcloud projects describe ${PROJECT_ID} --format="value(projectNumber)")
export ZONE="<YOUR_ZONE>" # e.g., us-east5-a
export REGION=${ZONE%-*}
export CLUSTER_NAME="qwen-serving-cluster"
export GVNIC_NETWORK_PREFIX="qwen-serving"
export BUCKET_NAME="inf-demo-model-storage-${PROJECT_NUMBER}"
export RESERVATION_NAME="<YOUR_RESERVATION_NAME>"
export NODE_POOL_NAME="tpu-v6e-32-resvd-pool"
export MULTIHOST_COLLECTION_NAME="tpu-6-collection"
export HF_TOKEN="<YOUR_HUGGING_FACE_TOKEN>" # Token with access to Qwen model if restricted

API を有効にする

必要な Google Cloud サービスを有効にします。

gcloud services enable \
    container.googleapis.com \
    compute.googleapis.com \
    iam.googleapis.com \
    cloudresourcemanager.googleapis.com

3. カスタム ネットワーキングを作成する

マルチホスト TPU ワークロードでは、アクセラレータの効率的な通信のために MTU サイズを大きくするなど、特定のネットワーク構成が必要です。クラスタのカスタム VPC ネットワークを作成します。

  1. 大きな MTU（8896）でVPC ネットワークを作成 します。
    gcloud compute --project=${PROJECT_ID} \
    networks create ${GVNIC_NETWORK_PREFIX}-main \
    --subnet-mode=custom \
    --mtu=8896
  2. クラスタのサブネットを作成 します。
    gcloud compute --project=${PROJECT_ID} \
    networks subnets create ${GVNIC_NETWORK_PREFIX}-tpu \
    --network=${GVNIC_NETWORK_PREFIX}-main \
    --region=${REGION} \
    --range=192.168.100.0/24
  3. ワーカーが通信できるように、内部トラフィックを許可するファイアウォール ルールを作成 します。
    gcloud compute --project=${PROJECT_ID} firewall-rules create ${GVNIC_NETWORK_PREFIX}-allow-internal \
    --network=${GVNIC_NETWORK_PREFIX}-main \
    --allow=all \
    --source-ranges=172.16.0.0/12,192.168.0.0/16,10.0.0.0/8 \
    --description="Allow all internal traffic within the network."

4. GKE クラスタをプロビジョニングする

GCS Fuse マウントと Ray Operator ワークロードをサポートするように構成された Standard GKE クラスタの設定を作成します。

  1. クラスタを作成 します。
    gcloud container clusters create ${CLUSTER_NAME} \
    --project=${PROJECT_ID} \
    --location=${REGION} \
    --release-channel=rapid \
    --machine-type=e2-standard-4 \
    --network=${GVNIC_NETWORK_PREFIX}-main \
    --subnetwork=${GVNIC_NETWORK_PREFIX}-tpu \
    --num-nodes=1 \
    --gateway-api=standard \
    --enable-managed-prometheus \
    --enable-dataplane-v2 \
    --enable-dataplane-v2-metrics \
    --workload-pool=${PROJECT_ID}.svc.id.goog \
    --addons=GcsFuseCsiDriver,RayOperator \
    --enable-ip-alias
  2. クラスタの認証情報を取得 します。
    gcloud container clusters get-credentials ${CLUSTER_NAME} --region=${REGION}
  3. Hugging Face Secret を作成 します。コンテナ アクセス ダウンロード用にトークンを安全に保存します。
    kubectl create secret generic hf-secret \
    --from-literal=hf_api_token=${HF_TOKEN} \
    --dry-run=client -o yaml | kubectl apply -f -
  4. Helm を使用してLeaderWorkerSet（LWS）をインストール します。LWS は、一緒にスケジュールする必要がある Pod のグループを管理します。
    helm install lws oci://registry.k8s.io/lws/charts/lws \
    --version=0.7.0 \
    --namespace lws-system \
    --create-namespace \
    --wait

5. GCS Rapid Cache を有効にする

サービング中に Cloud Storage から数十 GB の重みを読み取る速度を上げるには、GCS バケットを作成し、ゾーンで GCS Rapid Cache を有効にします。

  1. バケットを作成 します。
    gcloud storage buckets create gs://$BUCKET_NAME \
    --location=$REGION \
    --uniform-bucket-level-access
  2. TPU ゾーンで Rapid Cache を初期化 します。
    gcloud storage buckets anywhere-caches create gs://$BUCKET_NAME $ZONE \
    --ttl=1d \
    --admission-policy=ADMIT_ON_FIRST_MISS

6. Workload Identity とストレージ権限を設定する

有効期間の長い鍵を埋め込むことなく、重みバケットを GKE Pod に安全にマウントするように ID リンクを構成します。

  1. 専用の IAM サービス アカウントを作成 します。
    gcloud iam service-accounts create tpu-reader-sa
  2. バケットの読み取り権限を付与 します。
    gcloud storage buckets add-iam-policy-binding gs://${BUCKET_NAME} \
    --member="serviceAccount:tpu-reader-sa@${PROJECT_ID}.iam.gserviceaccount.com" \
    --role="roles/storage.objectAdmin"
  3. default Namespace Kubernetes サービス アカウントのWorkload Identity バインディングを作成 します。
    gcloud iam service-accounts add-iam-policy-binding tpu-reader-sa@${PROJECT_ID}.iam.gserviceaccount.com \
    --role="roles/iam.workloadIdentityUser" \
    --member="serviceAccount:${PROJECT_ID}.svc.id.goog[default/default]"
  4. Kubernetes SA にアノテーションを付けます：
    kubectl annotate serviceaccount default iam.gke.io/gcp-service-account=tpu-reader-sa@${PROJECT_ID}.iam.gserviceaccount.com

7. モデルの重みの設定

30B パラメータ モデルをサービングするには、Hugging Face から GCS バケットに重みをダウンロードする必要があります。Cloud Shell ディスクの割り当て上限（5 GB）を回避するには、Standard Kubernetes Job を使用して、クラスタ内で直接ダウンロードし、マウントされた GCS Fuse ボリュームに安全に書き込みます。

  1. モデル ダウンローダー ジョブをデプロイ します。次のマニフェストを作成して適用し、ダウンロードを開始します。
    cat <<EOF | kubectl apply -f -
apiVersion: batch/v1
kind: Job
metadata:
  name: model-downloader
spec:
  ttlSecondsAfterFinished: 60
  template:
    metadata:
      annotations:
        gke-gcsfuse/volumes: "true"
        gke-gcsfuse/memory-limit: "0"
    spec:
      serviceAccountName: default
      restartPolicy: OnFailure
      containers:
      - name: downloader
        image: python:3.10-slim
        command: ["/bin/sh", "-c"]
        args:
        - |
          pip install -U "huggingface_hub[hf_transfer]" filelock
          export HF_HUB_ENABLE_HF_TRANSFER=1

          python -c '
          import filelock

          class DummyLock:
              def __init__(self, *args, **kwargs): pass
              def __enter__(self): return self
              def __exit__(self, *args): pass
              def acquire(self, *args, **kwargs): pass
              def release(self, *args, **kwargs): pass

          filelock.FileLock = DummyLock

          from huggingface_hub import snapshot_download
          snapshot_download(
              repo_id="Qwen/Qwen3-30B-A3B", 
              local_dir="/models/qwen3-30b-weights",
              local_dir_use_symlinks=False
          )
          '
        env:
        - name: HF_TOKEN
          valueFrom:
            secretKeyRef:
              name: hf-secret
              key: hf_api_token
        volumeMounts:
        - name: model-weights
          mountPath: /models
      volumes:
      - name: model-weights
        csi:
          driver: gcsfuse.csi.storage.gke.io
          volumeAttributes:
            bucketName: ${BUCKET_NAME}
            mountOptions: "implicit-dirs"
EOF
  2. ダウンロードをモニタリングします。ダウンローダー Pod のログを確認して、進行状況を確認します。
    kubectl logs -f job/model-downloader
    ジョブが完了し、ステータスが成功になるまで待ちます。

8. 予約済みの TPU ノードプールを作成する

既存の容量予約を使用して、実際のマルチホスト TPU スライスをプロビジョニングします。

  1. 作成コマンドを実行 します。
    gcloud beta container node-pools create ${NODE_POOL_NAME} \
    --project=${PROJECT_ID} \
    --cluster=${CLUSTER_NAME} \
    --region=${REGION} \
    --node-locations=${ZONE} \
    --machine-type=ct6e-standard-4t \
    --tpu-topology=4x8 \
    --num-nodes=8 \
    --scopes=https://www.googleapis.com/auth/cloud-platform \
    --reservation-affinity=specific \
    --reservation=${RESERVATION_NAME} \
    --accelerator-network-profile=auto \
    --node-labels=cloud.google.com/gke-nodepool-group-name=${MULTIHOST_COLLECTION_NAME} \
    --node-labels=cloud.google.com/gke-workload-type=HIGH_AVAILABILITY \
    --node-labels=cloud.google.com/gke-networking-dra-driver=true
  2. ノードが参加するまで待ちます。ノード集約のスケーリングを直接確認できます。ct6e を含む 8 つのノードが kubectl get nodes に参加するまで待ちます。

9. vLLM サービスをデプロイする

  1. ネットワーク クレームを作成 します。ネットワーク環境をリクエストする必要があります。
    cat <<EOF | kubectl apply -f -
apiVersion: resource.k8s.io/v1
kind: ResourceClaimTemplate
metadata:
  name: all-netdev
spec:
  spec:
    devices:
      requests:
      - name: req-netdev
        exactly:
          deviceClassName: netdev.google.com
          allocationMode: All
EOF
  2. ロードバランサ API エンドポイントをデプロイ します。
    cat <<EOF | kubectl apply -f -
apiVersion: v1
kind: Service
metadata:
  name: vllm-tpu-service
spec:
  type: LoadBalancer
  selector:
    leaderworkerset.sigs.k8s.io/name: vllm-tpu-qwen
    leaderworkerset.sigs.k8s.io/worker-index: "0"
  ports:
  - protocol: TCP
    port: 8000
    targetPort: 8000
EOF
  3. LeaderWorkerSet ワークロードをデプロイ します。このマニフェストは、8 つのスライスホスト間で Ray ヘッド/ワーカーの集約を動的に開始します。
    cat <<EOF | kubectl apply -f -
apiVersion: leaderworkerset.x-k8s.io/v1
kind: LeaderWorkerSet
metadata:
  name: vllm-tpu-qwen
spec:
  replicas: 1
  leaderWorkerTemplate:
    size: 8
    restartPolicy: RecreateGroupOnPodRestart
    workerTemplate:
      metadata:
        annotations:
          gke-gcsfuse/volumes: "true"
          gke-gcsfuse/memory-limit: "0"
        labels:
          leaderworkerset.sigs.k8s.io/name: vllm-tpu-qwen
          gke-gcsfuse/volumes: "true"
      spec:
        hostname: vllm-tpu-qwen
        serviceAccountName: default
        containers:
        - name: vllm-tpu
          image: vllm/vllm-tpu:nightly
          command: ["sh", "-c"]
          args:
          - |
            MY_TPU_IP=\$(hostname -I | awk '{print \$1}')
            echo "My TPU Network IP is: \$MY_TPU_IP"

            LEADER_DNS="vllm-tpu-qwen-0.vllm-tpu-qwen"
            until getent hosts \$LEADER_DNS; do
              echo "DNS not ready. Sleeping 5s..."
              sleep 5
              done
            LEADER_IP=\$(getent hosts \$LEADER_DNS | awk '{print \$1}')

            export JAX_PLATFORMS=''
            export SCAN_TPU_CHIPS=True
            export TPU_MULTIHOST_BACKEND=ray
            export JAX_DISTRIBUTED_INITIALIZATION_TIMEOUT=300
            export LD_LIBRARY_PATH=\$LD_LIBRARY_PATH:/usr/local/lib
            export VLLM_HOST_IP=\$MY_TPU_IP

            if [ "\$LWS_WORKER_INDEX" = "0" ]; then
              echo "Starting Ray Head..."
              ray start --head --port=6379 --node-ip-address=\$MY_TPU_IP --resources='{"TPU": 4}' --block &
              sleep 20
              until ray status; do sleep 5; done

              echo "Starting vLLM API Server..."
              python3 -m vllm.entrypoints.openai.api_server \
                --model=/models/qwen3-30b-weights \
                --tensor-parallel-size=32 \
                --pipeline-parallel-size=1 \
                --distributed-executor-backend=ray \
                --host=0.0.0.0 --port=8000 \
                --enforce-eager \
                --gpu-memory-utilization=0.90
            else
              ray start --address=\${LEADER_IP}:6379 --node-ip-address=\$MY_TPU_IP --resources='{"TPU": 4}' --block
            fi
          ports:
          - containerPort: 8000
          - containerPort: 6379
          volumeMounts:
          - name: model-weights
            mountPath: /models
            readOnly: true
          - name: dshm
            mountPath: /dev/shm
          resources:
            claims:
            - name: net-resources
            limits:
              google.com/tpu: 4
              memory: "100Gi"
            requests:
              google.com/tpu: 4
              memory: "100Gi"
        nodeSelector:
          cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
          cloud.google.com/gke-tpu-topology: 4x8
          gke.networks.io/accelerator-network-profile: auto
        resourceClaims:
        - name: net-resources
          resourceClaimTemplateName: all-netdev
        volumes:
        - name: model-weights
          csi:
            driver: gcsfuse.csi.storage.gke.io
            volumeAttributes:
              bucketName: ${BUCKET_NAME}
              mountOptions: "implicit-dirs"
        - name: dshm
          emptyDir:
            medium: Memory
EOF

10. デプロイ レスポンスをテストする

LeaderWorkerSet 内のすべての Pod がコンテナ イメージを pull し、Ray を初期化して完全に Ready になるまでに 5 ～ 10 分かかることがあります。Pod の初期化を監視してステータスを追跡できます。

kubectl get pods -l leaderworkerset.sigs.k8s.io/name=vllm-tpu-qwen -w

8 つの vllm-tpu-qwen- Pod の STATUS が Running、READY が 2/2 になるまで待ちます。また、続行する前にロードバランサが外部 IP を受信していることを確認します。 これには 7 ～ 10 分かかることがあります。

  1. **外部 IP を取得** します。
    export EXTERNAL_IP=$(kubectl get svc vllm-tpu-service -o jsonpath='{.status.loadBalancer.ingress[0].ip}')
echo $EXTERNAL_IP

注意: 本番環境のサービスでは、このエンドポイントは Identity Aware Proxy（IAP）などで保護する必要があります。

  1. curl を使用して推論リクエストを送信 します。
        curl -N -s http://$EXTERNAL_IP:8000/v1/chat/completions \
        -H "Content-Type: application/json" \
        -d '{
            "model": "/models/qwen3-30b-weights",
            "messages": [{"role": "user", "content": "Write a haiku about high-performance computing on TPUs."}],
            "temperature": 0.7,
            "max_tokens": 100,
            "stream": true
        }' | sed 's/^data: //' | grep -v '\[DONE\]' | grep -v '^$' | jq -rj '.choices[0].delta.content // empty' ; echo ""
    生成された推論を表すテキストを含む JSON レスポンスのような出力が表示されます。

11. クリーンアップ

Google Cloud アカウントに継続的に課金されないようにするには、この Codelab で作成したリソースを削除します。

  1. ノードプールを削除 します。
    gcloud container node-pools delete "${NODE_POOL_NAME}" \
    --cluster="${CLUSTER_NAME}" \
    --region="${REGION}" \
    --project="${PROJECT_ID}" --quiet
  2. クラスタを削除 します。
    gcloud container clusters delete "${CLUSTER_NAME}" \
    --region="${REGION}" \
    --project="${PROJECT_ID}" --quiet
  3. ネットワークとファイアウォールの設定を削除 します。
    gcloud compute firewall-rules delete \
    "${GVNIC_NETWORK_PREFIX}-allow-internal" \
    --project="${PROJECT_ID}" --quiet

gcloud compute networks subnets delete "${GVNIC_NETWORK_PREFIX}-tpu" \
    --region="${REGION}" --quiet

gcloud compute networks delete "${GVNIC_NETWORK_PREFIX}-main" --quiet
  4. サービス アカウントのバインドを解除して削除 します。
        # 1. Create the cleanup script
    cat << 'EOF' > clean_up_sa.sh
    #!/bin/bash

    # Validate that PROJECT_ID is available
    if [ -z "$PROJECT_ID" ]; then
      echo "Error: PROJECT_ID environment variable is not set."
      exit 1
    fi

    SA_EMAIL="tpu-reader-sa@${PROJECT_ID}.iam.gserviceaccount.com"
    SA_MEMBER="serviceAccount:${SA_EMAIL}"

    echo "Gathering IAM policy for ${SA_EMAIL}..."

    # Fetch roles assigned to this specific SA
    ROLES=$(gcloud projects get-iam-policy ${PROJECT_ID} \
        --flatten="bindings[].members" \
        --filter="bindings.members:${SA_MEMBER}" \
        --format="value(bindings.role)")

    if [ -z "$ROLES" ]; then
        echo "No IAM bindings found for this service account."
    else
        for ROLE in $ROLES; do
            echo "Removing binding for: ${ROLE}..."
            gcloud projects remove-iam-policy-binding ${PROJECT_ID} \
                --member="${SA_MEMBER}" \
                --role="${ROLE}" --quiet > /dev/null
        done
        echo "Successfully unbound all roles."
    fi

    # 2. Delete the service account itself
    echo "Deleting service account..."
    gcloud iam service-accounts delete ${SA_EMAIL} --project=${PROJECT_ID} --quiet

    echo "Cleanup complete."
    EOF

    # 2. Make the script executable and run it
    chmod +x clean_up_sa.sh
    ./clean_up_sa.sh
  5. GCS バケットを削除 します。Cloud コンソールに移動し、[Cloud Storage] -> [バケット] を選択して、[inf-demo-model-storage] を選択し、[削除] を選択します。

12. 完了

おめでとうございます！Google Kubernetes Engine 上で Ray をネイティブに使用して、マルチホスト TPU 高推論レート vLLM スタックを正常にデプロイできました。

学習した内容

  • 高速 TPU トラフィックに合わせてカスタマイズされたカスタム パスウェイをプロビジョニングする。
  • GCS Fuse とリージョン Rapid Cache を使用して重みをマウントする。
  • LeaderWorkerSet を介してネイティブに同期されたマルチホスト ワークロード スライスをオーケストレートする。
  • 詳細については、vLLM ユーザーガイドと llm-d デプロイガイドをご覧ください。