Развертывание многохостового вывода TPU vLLM с использованием Ray на GKE

1. Введение

В этом практическом занятии вы узнаете, как развернуть высокопроизводительные многохостовые сервисы вывода виртуальных больших языковых моделей (vLLM) на Google Kubernetes Engine (GKE) с использованием Google Cloud TPU. Вы настроите распределенный вывод с помощью Ray и будете управлять рабочей нагрузкой непосредственно в GKE с помощью LeaderWorkerSets.

В этом пошаговом руководстве имитируется производственная конфигурация для обслуживания крупных моделей, таких как Qwen 30B .

Что вы будете делать

  • Создайте собственную сеть VPC для трафика ускорителя.
  • Разверните кластер GKE с Ray Operator и драйвером GCS Fuse CSI.
  • Инициализируйте кэш GCS Rapid Cache для ускоренной загрузки моделей.
  • Создайте многохостовый пул узлов TPU v6e с зарезервированной емкостью.
  • Настройте идентификацию рабочей нагрузки для безопасного доступа к весам модели.
  • Разверните и протестируйте движок vLLM, работающий с моделью, имеющей 30 миллиардов параметров.

Что вам понадобится

  • Проект Google Cloud с включенной функцией выставления счетов.
  • Резервирование ресурсов TPU v6e в Google Cloud (32 чипа, ct6e-standard-4t ).
  • Доступ к копированию весов модели из исходного хранилища.
  • Cloud Shell или локальный терминал с установленными gcloud , kubectl и helm .
  • Примерное время: 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

Установка переменных среды

Для упрощения выполнения команд определите следующие переменные в вашей оболочке. Замените <YOUR_ZONE> на выделенную вам зону TPU, а <YOUR_RESERVATION_NAME> на идентификатор вашего бронирования. Вам потребуется создать токен доступа пользователя 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. Создайте сеть VPC с большим значением MTU (8896): 
    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

Создайте стандартную конфигурацию кластера GKE, настроенную для поддержки монтирования GCS Fuse и рабочих нагрузок Ray Operator.

  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. Создайте секретное изображение обнимающегося лица : надежно сохраните свой токен для доступа к контейнерам и загрузки файлов: 
    kubectl create secret generic hf-secret \
    --from-literal=hf_api_token=${HF_TOKEN} \
    --dry-run=client -o yaml | kubectl apply -f -
  4. Установите LeaderWorkerSet (LWS) через Helm. LWS управляет группами подов, которые должны быть запланированы на выполнение одновременно: 
    helm install lws oci://registry.k8s.io/lws/charts/lws \
    --version=0.7.0 \
    --namespace lws-system \
    --create-namespace \
    --wait

5. Включите GCS Rapid Cache.

Чтобы ускорить чтение десятков гигабайт данных о весовых коэффициентах из облачного хранилища во время обслуживания, создайте сегмент GCS и включите GCS Rapid Cache в своей зоне.

  1. Создайте хранилище : 
    gcloud storage buckets create gs://$BUCKET_NAME \
    --location=$REGION \
    --uniform-bucket-level-access
  2. Инициализируйте Rapid Cache в вашей зоне TPU: 
    gcloud storage buckets anywhere-caches create gs://$BUCKET_NAME $ZONE \
    --ttl=1d \
    --admission-policy=ADMIT_ON_FIRST_MISS

6. Настройка прав доступа к учетной записи рабочей нагрузки и хранилищу.

Настройте идентификационные связи для безопасного крепления контейнера с грузами в ваши модули GKE без встраивания долговременных ключей.

  1. Создайте выделенную учетную запись службы IAM : 
    gcloud iam service-accounts create tpu-reader-sa
  2. Предоставить Bucket права на чтение : 
    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. Создайте привязку идентификатора рабочей нагрузки для учетной записи службы Kubernetes в пространстве имен default : 
    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 : 
    kubectl annotate serviceaccount default iam.gke.io/gcp-service-account=tpu-reader-sa@${PROJECT_ID}.iam.gserviceaccount.com

7. Настройка весов модели

Для обслуживания модели с 30 байтами параметров необходимо загрузить веса из Hugging Face в ваш GCS-хранилище. Чтобы обойти ограничение на дисковую квоту Cloud Shell (5 ГБ), используйте стандартное задание 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. Отслеживайте процесс загрузки : проверяйте журналы загрузки, чтобы следить за ходом загрузки:
    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. Дождитесь присоединения узлов : Вы можете наблюдать масштабирование агрегации узлов напрямую. Дождитесь, пока присоединятся 8 узлов, содержащих ct6e , с помощью 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 : Этот манифест динамически запускает агрегацию голов/рабочих узлов Ray на 8 хостах срезов. 
    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. Ответ на запрос о развертывании тестового оборудования

Для загрузки образов контейнеров, инициализации Ray и полной Ready всех подов в LeaderWorkerSet может потребоваться от 5 до 10 минут. Вы можете отслеживать статус, наблюдая за инициализацией подов: 

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

Дождитесь, пока STATUS всех 8 vllm-tpu-qwen- pods не изменится на 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, удалите ресурсы, созданные в ходе этого практического занятия.

  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. Перейдите в консоль облачного хранилища, выберите «Облачное хранилище» -> «Корзины», выберите inf-demo-model-storage, а затем нажмите «Удалить».

12. Поздравляем!

Поздравляем! Вы успешно развернули многохостовый стек vLLM с высокой скоростью вывода TPU, использующий Ray непосредственно в Google Kubernetes Engine.

Что вы узнали

  • Настройка пользовательских маршрутов, адаптированных для высокоскоростного трафика TPU.
  • Установка противовесов с использованием системы GCS Fuse и региональных станций быстрого доступа.
  • Организация работы с несколькими хостами, синхронизированными напрямую с помощью LeaderWorkerSets.
  • Для получения более подробной информации ознакомьтесь с Руководством пользователя vLLM и Руководствами по развертыванию llm-d.