Esegui il deployment dell'inferenza vLLM TPU multihost con Ray su GKE

1. Introduzione

In questo codelab, imparerai a eseguire il deployment di servizi di inferenza vLLM (Virtual Large Language Model) multi-host ad alte prestazioni su Google Kubernetes Engine (GKE) utilizzando le TPU Google Cloud. Configurerai l'inferenza distribuita utilizzando Ray e gestirai il carico di lavoro in modo nativo su GKE utilizzando LeaderWorkerSets.

Questa procedura dettagliata simula una configurazione di produzione per la pubblicazione di modelli di grandi dimensioni come Qwen 30B.

In questo lab proverai a:

  • Creare una rete VPC personalizzata per il traffico degli acceleratori.
  • Eseguire il provisioning di un cluster GKE con l'operatore Ray e il driver CSI GCS Fuse.
  • Inizializzare una cache rapida GCS per il caricamento accelerato dei modelli.
  • Eseguire il provisioning di un node pool TPU v6e multi-host con capacità riservata.
  • Configurare Workload Identity per l'accesso sicuro alle ponderazioni dei modelli.
  • Eseguire il deployment e testare il motore vLLM che pubblica un modello con 30 miliardi di parametri.

Che cosa ti serve

  • Un progetto Google Cloud con la fatturazione abilitata.
  • Una prenotazione Google Cloud per le risorse TPU v6e (32 chip, ct6e-standard-4t).
  • Accesso per copiare le ponderazioni dei modelli da un bucket di origine.
  • Cloud Shell o un terminale locale con gcloud, kubectl e helm installati.
  • Durata stimata: 60 minuti
  • Costo stimato: meno di 60 $ (supponendo che la rimozione venga eseguita immediatamente).

2. Prima di iniziare

Creare o selezionare un progetto Google Cloud

  1. Nella console Google Cloud, seleziona o crea un progetto Google Cloud.
  2. Assicurati che la fatturazione sia abilitata per il tuo progetto Cloud.

Avvia Cloud Shell

  1. Fai clic su Attiva Cloud Shell nella parte superiore della console Google Cloud.
  2. Verifica l'autenticazione:
gcloud auth list
  1. Conferma il progetto:
gcloud config get project
  1. Impostalo, se necessario:
export PROJECT_ID=<YOUR_PROJECT_ID>
gcloud config set project $PROJECT_ID

Imposta le variabili di ambiente

Per semplificare l'esecuzione dei comandi, definisci le seguenti variabili nella shell. Sostituisci <YOUR_ZONE> con la zona TPU allocata e <YOUR_RESERVATION_NAME> con l'ID della prenotazione. Dovrai creare un token di accesso utente di Hugging Face per scaricare le ponderazioni dei modelli con accesso limitato. Una volta creato, sostituisci <YOUR_HUGGING_FACE_TOKEN> con il token appena creato.

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

Abilita API

Abilita i servizi Google Cloud richiesti:

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

3. Crea una rete personalizzata

I carichi di lavoro TPU multi-host richiedono configurazioni di rete specifiche, incluse dimensioni MTU maggiori per una comunicazione efficiente degli acceleratori. Crea una rete VPC personalizzata per il cluster.

  1. Crea la rete VPC con un MTU di grandi dimensioni (8896):
    gcloud compute --project=${PROJECT_ID} \
    networks create ${GVNIC_NETWORK_PREFIX}-main \
    --subnet-mode=custom \
    --mtu=8896
  2. Crea la subnet per il 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. Crea regole firewall che consentano il traffico interno per consentire ai worker di comunicare:
    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. Esegui il provisioning del cluster GKE

Crea una configurazione del cluster GKE standard configurata per supportare i montaggi GCS Fuse e i carichi di lavoro dell'operatore Ray.

  1. Crea il 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. Recupera le credenziali del cluster:
    gcloud container clusters get-credentials ${CLUSTER_NAME} --region=${REGION}
  3. Crea il secret di Hugging Face: salva il token in modo sicuro per i download di accesso ai container:
    kubectl create secret generic hf-secret \
    --from-literal=hf_api_token=${HF_TOKEN} \
    --dry-run=client -o yaml | kubectl apply -f -
  4. Installa LeaderWorkerSet (LWS) tramite Helm. LWS gestisce gruppi di pod che devono essere pianificati insieme:
    helm install lws oci://registry.k8s.io/lws/charts/lws \
    --version=0.7.0 \
    --namespace lws-system \
    --create-namespace \
    --wait

5. Abilita la cache rapida GCS

Per velocizzare la lettura di decine di GB di ponderazioni da Cloud Storage durante la pubblicazione, crea un bucket GCS e abilita la cache rapida GCS nella tua zona.

  1. Crea il bucket:
    gcloud storage buckets create gs://$BUCKET_NAME \
    --location=$REGION \
    --uniform-bucket-level-access
  2. Inizializza la cache rapida nella zona TPU:
    gcloud storage buckets anywhere-caches create gs://$BUCKET_NAME $ZONE \
    --ttl=1d \
    --admission-policy=ADMIT_ON_FIRST_MISS

6. Configura Workload Identity e le autorizzazioni di archiviazione

Configura i link di identità per montare in modo sicuro il bucket di ponderazione nei pod GKE senza incorporare chiavi a lunga durata.

  1. Crea un account di servizio IAM dedicato:
    gcloud iam service-accounts create tpu-reader-sa
  2. Concedi le autorizzazioni di lettura 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 un binding di Workload Identity per il service account Kubernetes default namespace:
    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. Aggiungi un'annotazione al service account Kubernetes:
    kubectl annotate serviceaccount default iam.gke.io/gcp-service-account=tpu-reader-sa@${PROJECT_ID}.iam.gserviceaccount.com

7. Configurazione delle ponderazioni dei modelli

Per pubblicare un modello con 30 miliardi di parametri, devi scaricare le ponderazioni da Hugging Face nel tuo bucket GCS. Per aggirare il limite della quota disco di Cloud Shell (5 GB), utilizza un job Kubernetes standard per eseguire il download direttamente all'interno del cluster e scrivere in modo sicuro nel volume GCS Fuse montato.

  1. Esegui il deployment del job di download del modello: crea e applica il seguente manifest per avviare il 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. Monitora il download: controlla i log del pod del downloader per seguire l'avanzamento:
    kubectl logs -f job/model-downloader
    Attendi il completamento del job con lo stato di successo.

8. Crea un node pool TPU riservato

Esegui il provisioning della slice TPU multi-host effettiva utilizzando la prenotazione della capacità esistente.

  1. Esegui il comando di creazione:
    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. Attendi che i nodi si uniscano: puoi osservare direttamente lo scaling dell'aggregazione dei nodi. Attendi che 8 nodi contenenti ct6e si uniscano a kubectl get nodes.

9. Esegui il deployment del servizio vLLM

  1. Crea richieste di rete: devi richiedere l'ambiente di rete:
    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. Esegui il deployment dell'endpoint API del bilanciatore del carico:
    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. Esegui il deployment del carico di lavoro LeaderWorkerSet: questo manifest avvia l'aggregazione dinamica di head/worker di Ray sugli 8 host della slice.
    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. Testa la risposta al deployment

Possono essere necessari 5-10 minuti prima che tutti i pod in LeaderWorkerSet eseguano il pull delle immagini container, inizializzino Ray e diventino completamente Ready. Puoi monitorare lo stato osservando l'inizializzazione dei pod:

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

Attendi che tutti gli 8 pod vllm-tpu-qwen- mostrino STATUS come Running e READY come 2/2 e assicurati che il bilanciatore del carico abbia ricevuto un IP esterno prima di procedere. Questa operazione può richiedere 7-10 minuti.

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

Attenzione: in un servizio di produzione, questo endpoint deve essere protetto con un servizio come Identity Aware Proxy (IAP)

  1. Invia una richiesta di inferenza utilizzando 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 ""
    Dovresti visualizzare un output simile a una risposta JSON contenente il testo che rappresenta l'inferenza generata.

11. Libera spazio

Per evitare addebiti continui sul tuo account Google Cloud, elimina le risorse create durante questo codelab.

  1. Elimina il node pool:
    gcloud container node-pools delete "${NODE_POOL_NAME}" \
    --cluster="${CLUSTER_NAME}" \
    --region="${REGION}" \
    --project="${PROJECT_ID}" --quiet
  2. Elimina il cluster:
    gcloud container clusters delete "${CLUSTER_NAME}" \
    --region="${REGION}" \
    --project="${PROJECT_ID}" --quiet
  3. Elimina le configurazioni di rete 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. Dissocia ed elimina il service account:
        # 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. Elimina il bucket GCS : vai alla console Google Cloud, seleziona Cloud Storage -> Bucket, seleziona inf-demo-model-storage e poi scegli "Elimina".

12. Complimenti

Complimenti! Hai eseguito correttamente il deployment di uno stack vLLM multi-host TPU ad alta velocità di inferenza utilizzando Ray in modo nativo su Google Kubernetes Engine.

Che cosa hai imparato

  • Eseguire il provisioning di percorsi personalizzati adatti al traffico TPU ad alta velocità.
  • Montare le ponderazioni utilizzando GCS Fuse e le cache rapide regionali.
  • Orchestrare le slice di carichi di lavoro multi-host sincronizzate in modo nativo tramite LeaderWorkerSets.
  • Per saperne di più, consulta la Guida per l'utente di vLLM e le guide al deployment di llm-d