Implementa la inferencia de vLLM con TPU de varios hosts con Ray en GKE

1. Introducción

En este codelab, aprenderás a implementar servicios de inferencia de vLLM (modelo de lenguaje grande virtual) de alto rendimiento y varios hosts en Google Kubernetes Engine (GKE) con las TPU de Google Cloud. Configurarás la inferencia distribuida con Ray y administrarás la carga de trabajo de forma nativa en GKE con LeaderWorkerSets.

Este recorrido simula una configuración de producción para entregar modelos grandes, como Qwen 30B.

Actividades

  • Crea una red de VPC personalizada para el tráfico del acelerador.
  • Aprovisiona un clúster de GKE con el operador de Ray y el controlador de CSI de GCS FUSE.
  • Inicializa una caché rápida de GCS para acelerar la carga del modelo.
  • Aprovisiona un grupo de nodos TPU v6e de varios hosts con capacidad reservada.
  • Configura Workload Identity para acceder de forma segura a los pesos del modelo.
  • Implementa y prueba el motor de vLLM que entrega un modelo de 30 mil millones de parámetros.

Requisitos

  • Un proyecto de Google Cloud con facturación habilitada.
  • Una reservación de Google Cloud para recursos de TPU v6e (32 chips, ct6e-standard-4t).
  • Acceso para copiar los pesos del modelo desde un bucket de origen
  • Cloud Shell o una terminal local con gcloud, kubectl y helm instalados
  • Duración estimada: 60 minutos
  • Costo estimado: Menos de USD 60 (suponiendo que el desmontaje se realice de inmediato).

2. Antes de comenzar

Crea o selecciona un proyecto de Google Cloud

  1. En la consola de Google Cloud, selecciona o crea un proyecto de Google Cloud.
  2. Confirma que la facturación está habilitada para tu proyecto de Cloud.

Inicie Cloud Shell

  1. Haz clic en Activar Cloud Shell en la parte superior de la consola de Google Cloud.
  2. Verifica la autenticación:
gcloud auth list
  1. Confirma tu proyecto:
gcloud config get project
  1. Establécela si es necesario:
export PROJECT_ID=<YOUR_PROJECT_ID>
gcloud config set project $PROJECT_ID

Configura variables de entorno

Para facilitar la ejecución de comandos, define las siguientes variables en tu shell. Reemplaza <YOUR_ZONE> por la zona de TPU asignada y <YOUR_RESERVATION_NAME> por el ID de tu reserva. Deberás crear un token de acceso de usuario de Hugging Face para descargar los pesos del modelo restringido. Una vez que lo hayas creado, reemplaza <YOUR_HUGGING_FACE_TOKEN> por el token que acabas de crear.

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

Habilita las APIs

Habilita los servicios de Google Cloud necesarios:

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

3. Crea redes personalizadas

Las cargas de trabajo de TPU de varios hosts requieren configuraciones de red específicas, incluidos tamaños de MTU más altos para una comunicación eficiente del acelerador. Crea una red de VPC personalizada para tu clúster.

  1. Crea la red de VPC con una MTU grande (8896):
    gcloud compute --project=${PROJECT_ID} \
    networks create ${GVNIC_NETWORK_PREFIX}-main \
    --subnet-mode=custom \
    --mtu=8896
  2. Crea la subred para el clúster:
    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. Crea reglas de firewall que permitan el tráfico interno para que los trabajadores puedan comunicarse:
    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. Aprovisiona el clúster de GKE

Crea una configuración de clúster de GKE Standard configurada para admitir Ray Operator y las activaciones de GCS Fuse.

  1. Crea el clúster:
    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. Retrieve Cluster Credentials:
    gcloud container clusters get-credentials ${CLUSTER_NAME} --region=${REGION}
  3. Crea un secreto de Hugging Face: Guarda tu token de forma segura para las descargas de acceso al contenedor:
    kubectl create secret generic hf-secret \
    --from-literal=hf_api_token=${HF_TOKEN} \
    --dry-run=client -o yaml | kubectl apply -f -
  4. Instala LeaderWorkerSet (LWS) a través de Helm. LWS administra grupos de pods que se deben programar juntos:
    helm install lws oci://registry.k8s.io/lws/charts/lws \
    --version=0.7.0 \
    --namespace lws-system \
    --create-namespace \
    --wait

5. Habilita Rapid Cache de GCS

Para acelerar la lectura de decenas de GB de pesos de Cloud Storage durante la entrega, crea un bucket de GCS y habilita la caché rápida de GCS en tu zona.

  1. Crea el bucket:
    gcloud storage buckets create gs://$BUCKET_NAME \
    --location=$REGION \
    --uniform-bucket-level-access
  2. Inicializa Rapid Cache en tu zona de TPU:
    gcloud storage buckets anywhere-caches create gs://$BUCKET_NAME $ZONE \
    --ttl=1d \
    --admission-policy=ADMIT_ON_FIRST_MISS

6. Configura Workload Identity y los permisos de almacenamiento

Configura vínculos de identidad para activar de forma segura el bucket de peso en tus Pods de GKE sin incorporar claves de larga duración.

  1. Crea una cuenta de servicio de IAM dedicada:
    gcloud iam service-accounts create tpu-reader-sa
  2. Otorga permisos de lectura del 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. Crea una vinculación de Workload Identity para la cuenta de servicio de Kubernetes del espacio de nombres 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. Anota la SA de Kubernetes:
    kubectl annotate serviceaccount default iam.gke.io/gcp-service-account=tpu-reader-sa@${PROJECT_ID}.iam.gserviceaccount.com

7. Configuración de los pesos del modelo

Para entregar un modelo de 30 mil millones de parámetros, debes descargar los pesos de Hugging Face en tu bucket de GCS. Para omitir el límite de cuota de disco de Cloud Shell (5 GB), usa un trabajo estándar de Kubernetes para descargar directamente dentro del clúster y escribir de forma segura en el volumen de GCS FUSE que se haya activado.

  1. Implementa el trabajo de Model Downloader: Crea y aplica el siguiente manifiesto para iniciar la descarga:
    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. Supervisa la descarga: Verifica los registros del pod del descargador para seguir el progreso:
    kubectl logs -f job/model-downloader
    Espera hasta que el trabajo se complete con el estado de éxito.

8. Crea un grupo de nodos TPU reservado

Aprovisiona la porción de TPU de varios hosts real con tu reserva de capacidad existente.

  1. Ejecuta el comando de creación:
    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. Espera a que se unan los nodos: Puedes observar el ajuste de escala de la agregación de nodos directamente. Espera hasta que 8 nodos que contengan ct6e se unan a kubectl get nodes.

9. Implementa el servicio de vLLM

  1. Crear reclamos de red: Debes solicitar el entorno de red:
    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. Implementa el extremo de la API del balanceador de cargas:
    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. Deploy LeaderWorkerSet workload: Este manifiesto inicia la agregación de líder y trabajador de Ray de forma dinámica en los 8 hosts de segmentación.
    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. Respuesta de prueba de implementación

Todos los Pods del LeaderWorkerSet pueden tardar entre 5 y 10 minutos en extraer imágenes de contenedores, inicializar Ray y estar completamente Ready. Puedes hacer un seguimiento del estado observando la inicialización del pod:

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

Espera hasta que los 8 Pods vllm-tpu-qwen- muestren STATUS como Running y READY como 2/2, y asegúrate de que el balanceador de cargas haya recibido una IP externa antes de continuar. Este proceso puede tardar entre 7 y 10 minutos.

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

Precaución: En un servicio de producción, este extremo debe protegerse con algo como Identity-Aware Proxy (IAP).

  1. Envía la solicitud de inferencia con 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 ""
    Deberías ver un resultado similar a una respuesta JSON que contiene el texto que representa tu inferencia generada.

11. Limpia

Para evitar que se apliquen cargos a tu cuenta de Google Cloud, borra los recursos que creaste durante este codelab.

  1. Borrar grupo de nodos:
    gcloud container node-pools delete "${NODE_POOL_NAME}" \
    --cluster="${CLUSTER_NAME}" \
    --region="${REGION}" \
    --project="${PROJECT_ID}" --quiet
  2. Borrar clúster:
    gcloud container clusters delete "${CLUSTER_NAME}" \
    --region="${REGION}" \
    --project="${PROJECT_ID}" --quiet
  3. Borra la configuración de red y 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. Desvincula y borra la cuenta de servicio:
        # 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. Borra el bucket de GCS. Navega a la consola de Cloud, selecciona Cloud Storage -> Buckets, selecciona inf-demo-model-storage y, luego, elige "Borrar".

12. ¡Felicitaciones!

¡Felicitaciones! Implementaste correctamente una pila de vLLM con una alta tasa de inferencia de TPU de varios hosts que utiliza Ray de forma nativa en Google Kubernetes Engine.

Qué aprendiste

  • Aprovisionamiento de rutas personalizadas diseñadas para el tráfico de TPU de alta velocidad
  • Se ajustan los pesos con GCS Fuse y cachés rápidas regionales.
  • Organizar porciones de cargas de trabajo de varios hosts sincronizadas de forma nativa a través de LeaderWorkerSets
  • Para obtener más información, consulta la Guía del usuario de vLLM y las Guías de implementación de llm-d.