Implantar a inferência do vLLM de TPU multihost com o Ray no GKE

1. Introdução

Neste codelab, você vai aprender a implantar serviços de inferência de vLLM (modelo de linguagem grande virtual) de alto desempenho e vários hosts no Google Kubernetes Engine (GKE) usando as TPUs do Google Cloud. Você vai configurar a inferência distribuída usando o Ray e gerenciar a carga de trabalho nativamente no GKE usando o LeaderWorkerSets.

Este tutorial simula uma configuração de produção para disponibilizar modelos grandes, como o Qwen 30B.

Atividades deste laboratório

  • Criar uma rede VPC personalizada para tráfego de acelerador.
  • Provisionar um cluster do GKE com o operador Ray e o driver CSI do GCS Fuse.
  • Inicializar um Rapid Cache do GCS para carregamento acelerado de modelos.
  • Provisionar um pool de nós de TPU v6e de vários hosts com capacidade reservada.
  • Configurar a Identidade da carga de trabalho para acesso seguro aos pesos do modelo.
  • Implantar e testar o mecanismo vLLM que disponibiliza um modelo de parâmetro 30B.

O que é necessário

  • Ter um projeto do Google Cloud com o faturamento ativado.
  • Uma reserva do Google Cloud para recursos de TPU v6e (32 chips, ct6e-standard-4t).
  • Acesso para copiar pesos de modelo de um bucket de origem.
  • Cloud Shell ou um terminal local com gcloud, kubectl e helm instalados.
  • Duração estimada:60 minutos
  • Custo estimado:menos de US $60 (supondo que a desmontagem seja feita imediatamente).

2. Antes de começar

Criar ou selecionar um projeto do Google Cloud

  1. No Console do Google Cloud, selecione ou crie um projeto na nuvem do Google Cloud.
  2. Verifique se o faturamento está ativado no seu projeto na nuvem.

Iniciar o Cloud Shell

  1. Clique em Ativar o Cloud Shell na parte de cima do console do Google Cloud.
  2. Verifique a autenticação:
gcloud auth list
  1. Confirme seu projeto:
gcloud config get project
  1. Defina-o, se necessário:
export PROJECT_ID=<YOUR_PROJECT_ID>
gcloud config set project $PROJECT_ID

Definir variáveis de ambiente

Para facilitar a execução de comandos, defina as seguintes variáveis no shell. Substitua <YOUR_ZONE> pela zona de TPU alocada e <YOUR_RESERVATION_NAME> pelo ID da reserva. Você precisará criar um token de acesso de usuário do Hugging Face para baixar pesos de modelo restritos. Depois de criar, substitua <YOUR_HUGGING_FACE_TOKEN> pelo token recém-criado.

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

Ativar APIs

Ative os serviços necessários do Google Cloud:

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

3. Criar rede personalizada

As cargas de trabalho de TPU de vários hosts exigem configurações de rede específicas, incluindo tamanhos de MTU maiores para comunicação eficiente do acelerador. Crie uma rede VPC personalizada para o cluster.

  1. Crie a rede VPC com um MTU grande (8896):
    gcloud compute --project=${PROJECT_ID} \
    networks create ${GVNIC_NETWORK_PREFIX}-main \
    --subnet-mode=custom \
    --mtu=8896
  2. Crie a sub-rede para o cluster:
    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. Crie regras de firewall que permitam o tráfego interno para que os workers possam se comunicar:
    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. Provisionar cluster do GKE

Crie uma configuração de cluster do GKE padrão configurada para oferecer suporte a montagens do GCS Fuse e cargas de trabalho do operador Ray.

  1. Crie o cluster:
    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. Recupere as credenciais do cluster:
    gcloud container clusters get-credentials ${CLUSTER_NAME} --region=${REGION}
  3. Crie o secret do Hugging Face: salve seu token com segurança para downloads de acesso ao contêiner:
    kubectl create secret generic hf-secret \
    --from-literal=hf_api_token=${HF_TOKEN} \
    --dry-run=client -o yaml | kubectl apply -f -
  4. Instale o LeaderWorkerSet (LWS) pelo Helm. O LWS gerencia grupos de pods que precisam ser programados juntos:
    helm install lws oci://registry.k8s.io/lws/charts/lws \
    --version=0.7.0 \
    --namespace lws-system \
    --create-namespace \
    --wait

5. Ativar o Rapid Cache do GCS

Para acelerar a leitura de dezenas de GBs de pesos do Cloud Storage durante a disponibilização, crie um bucket do GCS e ative o Rapid Cache do GCS na sua zona.

  1. Crie o bucket:
    gcloud storage buckets create gs://$BUCKET_NAME \
    --location=$REGION \
    --uniform-bucket-level-access
  2. Inicialize o Rapid Cache na sua zona de TPU:
    gcloud storage buckets anywhere-caches create gs://$BUCKET_NAME $ZONE \
    --ttl=1d \
    --admission-policy=ADMIT_ON_FIRST_MISS

6. Configurar a Identidade da carga de trabalho e as permissões de armazenamento

Configure links de identidade para montar o bucket de peso com segurança nos pods do GKE sem incorporar chaves de longa duração.

  1. Crie uma conta de serviço do IAM dedicada:
    gcloud iam service-accounts create tpu-reader-sa
  2. Conceda permissões de leitura do 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. Crie a vinculação da Identidade da carga de trabalho para a conta de serviço do Kubernetes do namespace 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. Anote a SA do Kubernetes:
    kubectl annotate serviceaccount default iam.gke.io/gcp-service-account=tpu-reader-sa@${PROJECT_ID}.iam.gserviceaccount.com

7. Configuração de pesos do modelo

Para disponibilizar um modelo de parâmetro 30B, é necessário baixar pesos do Hugging Face para o bucket do GCS. Para ignorar o limite de cota de disco do Cloud Shell (5 GB), use um job do Kubernetes padrão para fazer o download diretamente no cluster e gravar no volume montado do GCS Fuse com segurança.

  1. Implante o job do downloader de modelos: crie e aplique o seguinte manifesto para iniciar o download:
    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. Monitore o download: verifique os registros do pod do downloader para acompanhar o progresso:
    kubectl logs -f job/model-downloader
    Aguarde até que o job seja concluído com o status de sucesso.

8. Criar pool de nós de TPU reservados

Provisione a fração de TPU de vários hosts usando a reserva de capacidade atual.

  1. Execute o comando de criação:
    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. Aguarde os nós entrarem: é possível observar o escalonamento da agregação de nós diretamente. Aguarde até que oito nós que contenham ct6e entrem em kubectl get nodes.

9. Implantar o serviço vLLM

  1. Crie declarações de rede: é necessário solicitar o ambiente de rede:
    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. Implante o endpoint da API do balanceador de carga:
    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. Implante a carga de trabalho do LeaderWorkerSet: esse manifesto inicia a agregação de head/worker do Ray de forma dinâmica nos oito hosts de fração.
    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. Resposta do teste de implantação

Pode levar de 5 a 10 minutos para que todos os pods no LeaderWorkerSet extraiam imagens de contêiner, inicializem o Ray e fiquem totalmente Ready. É possível acompanhar o status observando a inicialização do pod:

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

Aguarde até que todos os oito pods vllm-tpu-qwen- mostrem STATUS como Running e READY como 2/2 e verifique se o balanceador de carga recebeu um IP externo antes de continuar. Isso pode levar de 7 a 10 minutos.

  1. Recupere o IP externo: 
    export EXTERNAL_IP=$(kubectl get svc vllm-tpu-service -o jsonpath='{.status.loadBalancer.ingress[0].ip}')
echo $EXTERNAL_IP

Atenção: em um serviço de produção, esse endpoint precisa ser protegido com algo como o Identity-Aware Proxy (IAP).

  1. Envie a solicitação de inferência usando 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 ""
    Você verá uma saída semelhante a uma resposta JSON contendo o texto que representa sua inferência gerada.

11. Limpar

Para evitar cobranças contínuas na sua conta do Google Cloud, exclua os recursos criados durante este codelab.

  1. Excluir pool de nós:
    gcloud container node-pools delete "${NODE_POOL_NAME}" \
    --cluster="${CLUSTER_NAME}" \
    --region="${REGION}" \
    --project="${PROJECT_ID}" --quiet
  2. Excluir cluster:
    gcloud container clusters delete "${CLUSTER_NAME}" \
    --region="${REGION}" \
    --project="${PROJECT_ID}" --quiet
  3. Excluir configurações de rede e firewall:
    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. Desvincular e excluir conta de serviço:
        # 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. Excluir bucket do GCS : acesse o console do Cloud, selecione Cloud Storage > Buckets, selecione inf-demo-model-storage e escolha "Excluir".

12. Parabéns

Parabéns! Você implantou com sucesso uma pilha de vLLM de alta taxa de inferência de TPU de vários hosts usando o Ray nativamente no Google Kubernetes Engine.

O que você aprendeu

  • Provisionamento de caminhos personalizados para tráfego de TPU de alta velocidade.
  • Montagem de pesos usando o GCS Fuse e caches rápidos regionais.
  • Orquestração de frações de carga de trabalho de vários hosts sincronizadas nativamente pelo LeaderWorkerSets.
  • Para saber mais, consulte o guia do usuário do vLLM e os guias de implantação do llm-d.