在 GKE 上使用 Ray 部署多主机 TPU vLLM 推理

1. 简介

在此 Codelab 中，您将学习如何使用 Google Cloud TPU 在 Google Kubernetes Engine (GKE) 上部署高性能多主机 vLLM（虚拟大语言模型）推理服务。您将使用 Ray 配置分布式推理，并使用 LeaderWorkerSets 在 GKE 上原生管理工作负载。

此演练模拟了用于部署 Qwen 30B 等大型模型的生产设置。

您将执行的操作

  • 为加速器流量创建自定义 VPC 网络。
  • 预配具有 Ray Operator 和 GCS Fuse CSI 驱动程序的 GKE 集群。
  • 初始化 GCS 快速缓存，以加速模型加载。
  • 预配具有预留容量的多主机 TPU v6e 节点池。
  • 配置 Workload Identity 以安全访问模型权重。
  • 部署并测试 vLLM 引擎，该引擎可提供 300 亿形参模型。

所需条件

  • 启用了结算功能的 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 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

设置环境变量

为了更轻松地执行命令，请在 shell 中定义以下变量。将 <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 可用区中初始化快速缓存：
    gcloud storage buckets anywhere-caches create gs://$BUCKET_NAME $ZONE \
    --ttl=1d \
    --admission-policy=ADMIT_ON_FIRST_MISS

6. 设置 Workload Identity 和存储权限

配置身份链接，以安全地将权重桶装载到 GKE pod 中，而无需嵌入长期有效的密钥。

  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 命名空间 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. 模型权重设置

如需投放 300 亿参数模型，您需要将权重从 Hugging Face 下载到您的 GCS 存储分区中。如需绕过 Cloud Shell 磁盘配额限制 (5 GB)，请使用标准 Kubernetes 作业直接在集群内下载并安全地写入已装载的 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 个 slice 主机上动态启动 Ray head/worker 聚合。
    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 可能需要 5-10 分钟才能拉取容器映像、初始化 Ray 并完全变为 Ready。您可以通过观察 pod 初始化来跟踪状态：

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

等待所有 8 个 vllm-tpu-qwen- pod 的 Running 显示为 STATUS，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. 恭喜

恭喜！您已成功部署多主机 TPU 高推理速率 vLLM 技术栈，该技术栈通过 Google Kubernetes Engine 原生利用 Ray。

您学到的内容

  • 预配专为高速 TPU 流量量身定制的自定义路径。
  • 利用 GCS Fuse 和区域快速缓存装载权重。
  • 通过 LeaderWorkerSets 原生同步编排多主机工作负载切片。
  • 如需了解详情，请参阅 vLLM 用户指南和 llm-d 部署指南