Multihost-TPU-vLLM-Inferenz mit Ray in GKE bereitstellen

1. Einführung

In diesem Codelab erfahren Sie, wie Sie leistungsstarke vLLM-Inferenzdienste (Virtual Large Language Model) mit mehreren Hosts in Google Kubernetes Engine (GKE) mithilfe von Google Cloud TPUs bereitstellen. Sie konfigurieren die verteilte Inferenz mit Ray und verwalten die Arbeitslast nativ in GKE mit LeaderWorkerSets.

In dieser Anleitung wird eine Produktionskonfiguration für die Bereitstellung großer Modelle wie Qwen 30B simuliert.

Aufgaben

  • Benutzerdefiniertes VPC-Netzwerk für Beschleunigertraffic erstellen
  • GKE-Cluster mit Ray Operator und GCS Fuse CSI-Treiber bereitstellen
  • GCS Rapid Cache für beschleunigtes Laden von Modellen initialisieren
  • TPU v6e-Knotenpool mit mehreren Hosts und reservierter Kapazität bereitstellen
  • Workload Identity für sicheren Zugriff auf Modellgewichtungen konfigurieren
  • vLLM-Engine bereitstellen und testen, die ein Modell mit 30 Milliarden Parametern bereitstellt

Voraussetzungen

  • Google Cloud-Projekt mit aktivierter Abrechnungsfunktion.
  • Eine Google Cloud-Reservierung für TPU v6e-Ressourcen (32 Chips, ct6e-standard-4t).
  • Zugriff zum Kopieren von Modellgewichtungen aus einem Quell-Bucket.
  • Cloud Shell oder ein lokales Terminal mit installierten gcloud, kubectl und helm.
  • Geschätzte Dauer:60 Minuten
  • Geschätzte Kosten:Weniger als 60 $ (bei umgehender Bereinigung)

2. Hinweis

Google Cloud-Projekt erstellen oder auswählen

  1. Wählen Sie in der Google Cloud Console ein Google Cloud-Projekt aus oder erstellen Sie eines.
  2. Prüfen Sie, ob für Ihr Cloud-Projekt die Abrechnung aktiviert ist.

Cloud Shell starten

  1. Klicken Sie oben in der Google Cloud Console auf Cloud Shell aktivieren.
  2. Authentifizierung überprüfen:
gcloud auth list
  1. Bestätigen Sie Ihr Projekt:
gcloud config get project
  1. Legen Sie es bei Bedarf fest:
export PROJECT_ID=<YOUR_PROJECT_ID>
gcloud config set project $PROJECT_ID

Umgebungsvariablen festlegen

Um die Ausführung von Befehlen zu erleichtern, definieren Sie die folgenden Variablen in Ihrer Shell. Ersetzen Sie <YOUR_ZONE> durch Ihre zugewiesene TPU-Zone und <YOUR_RESERVATION_NAME> durch Ihre Reservierungs-ID. Sie müssen ein Hugging Face-Nutzerzugriffstoken erstellen, um geschützte Modellgewichtungen herunterzuladen. Ersetzen Sie nach dem Erstellen <YOUR_HUGGING_FACE_TOKEN> durch Ihr neu erstelltes 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

APIs aktivieren

Aktivieren Sie die erforderlichen Google Cloud-Dienste:

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

3. Benutzerdefiniertes Netzwerk erstellen

Für TPU-Arbeitslasten mit mehreren Hosts sind bestimmte Netzwerkkonfigurationen erforderlich, einschließlich höherer MTU-Größen für eine effiziente Beschleunigerkommunikation. Erstellen Sie ein benutzerdefiniertes VPC-Netzwerk für Ihren Cluster.

  1. VPC-Netzwerk mit einer großen MTU (8896) erstellen :
    gcloud compute --project=${PROJECT_ID} \
    networks create ${GVNIC_NETWORK_PREFIX}-main \
    --subnet-mode=custom \
    --mtu=8896
  2. Subnetz für den Cluster erstellen:
    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. Firewallregeln erstellen , die internen Traffic zulassen, damit Worker kommunizieren können:
    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-Cluster bereitstellen

Erstellen Sie eine Standard-GKE-Clusterkonfiguration, die für GCS Fuse-Mounts und Ray Operator-Arbeitslasten konfiguriert ist.

  1. Cluster erstellen:
    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. Clusteranmeldedaten abrufen:
    gcloud container clusters get-credentials ${CLUSTER_NAME} --region=${REGION}
  3. Hugging Face-Secret erstellen: Speichern Sie Ihr Token sicher für Containerzugriffs-Downloads:
    kubectl create secret generic hf-secret \
    --from-literal=hf_api_token=${HF_TOKEN} \
    --dry-run=client -o yaml | kubectl apply -f -
  4. LeaderWorkerSet (LWS) über Helm installieren. LWS verwaltet Gruppen von Pods, die zusammen geplant werden müssen:
    helm install lws oci://registry.k8s.io/lws/charts/lws \
    --version=0.7.0 \
    --namespace lws-system \
    --create-namespace \
    --wait

5. GCS Rapid Cache aktivieren

Um das Lesen von Dutzenden von GB an Gewichtungen aus Cloud Storage während der Bereitstellung zu beschleunigen, erstellen Sie einen GCS-Bucket und aktivieren Sie GCS Rapid Cache in Ihrer Zone.

  1. Bucket erstellen:
    gcloud storage buckets create gs://$BUCKET_NAME \
    --location=$REGION \
    --uniform-bucket-level-access
  2. Rapid Cache in Ihrer TPU-Zone initialisieren:
    gcloud storage buckets anywhere-caches create gs://$BUCKET_NAME $ZONE \
    --ttl=1d \
    --admission-policy=ADMIT_ON_FIRST_MISS

6. Workload Identity und Speicherberechtigungen einrichten

Konfigurieren Sie Identitätslinks, um den Bucket mit den Gewichtungen sicher in Ihre GKE-Pods einzubinden, ohne langlebige Schlüssel einzubetten.

  1. Dediziertes IAM-Dienstkonto erstellen:
    gcloud iam service-accounts create tpu-reader-sa
  2. Leseberechtigungen für den Bucket erteilen:
    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. Workload Identity-Bindung für das Kubernetes-Dienstkonto default erstellen:
    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-Dienstkonto annotieren:
    kubectl annotate serviceaccount default iam.gke.io/gcp-service-account=tpu-reader-sa@${PROJECT_ID}.iam.gserviceaccount.com

7. Modellgewichtungen einrichten

Um ein Modell mit 30 Milliarden Parametern bereitzustellen, müssen Sie Gewichtungen von Hugging Face in Ihren GCS-Bucket herunterladen. Um das Festplattenkontingent von Cloud Shell (5 GB) zu umgehen, verwenden Sie einen Standard-Kubernetes-Job , um die Gewichtungen direkt im Cluster herunterzuladen und sicher in das eingebundene GCS Fuse-Volume zu schreiben.

  1. Modell-Downloader-Job bereitstellen: Erstellen und wenden Sie das folgende Manifest an, um den Download zu starten:
    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. Download beobachten: Prüfen Sie die Logs des Downloader-Pods, um den Fortschritt zu verfolgen:
    kubectl logs -f job/model-downloader
    Warten Sie, bis der Job mit dem Status „Erfolg“ abgeschlossen ist.

8. Reservierten TPU-Knotenpool erstellen

Stellen Sie den eigentlichen TPU-Slice mit mehreren Hosts mithilfe Ihrer vorhandenen Kapazitätsreservierung bereit.

  1. Erstellungsbefehl ausführen:
    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. Warten, bis Knoten hinzugefügt werden: Sie können die Knotenzusammenfassung direkt beobachten. Warten Sie, bis 8 Knoten mit ct6e zu kubectl get nodes hinzugefügt wurden.

9. vLLM-Dienst bereitstellen

  1. Netzwerkansprüche erstellen: Sie müssen die Netzwerkumgebung anfordern:
    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. Load-Balancer-API-Endpunkt bereitstellen:
    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-Arbeitslast bereitstellen: Dieses Manifest startet die Ray-Head-/Worker-Aggregation dynamisch auf den 8 Slice-Hosts.
    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. Antwort der Bereitstellung testen

Es kann 5 bis 10 Minuten dauern, bis alle Pods im LeaderWorkerSet Container-Images abrufen, Ray initialisieren und vollständig Ready sind. Sie können den Status verfolgen, indem Sie die Pod-Initialisierung beobachten:

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

Warten Sie, bis für alle 8 vllm-tpu-qwen--Pods STATUS als Running und READY als 2/2 angezeigt wird. Prüfen Sie außerdem, ob der Load-Balancer eine externe IP-Adresse erhalten hat, bevor Sie fortfahren. Das kann 7 bis 10 Minuten dauern.

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

Achtung: In einem Produktionsdienst sollte dieser Endpunkt mit etwas wie Identity-Aware Proxy (IAP) gesichert werden.

  1. Inferenzanfrage mit curl senden:
        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 ""
    Sie sollten eine Ausgabe sehen, die einer JSON-Antwort ähnelt und den Text der generierten Inferenz enthält.

11. Bereinigen

Um laufende Kosten für Ihr Google Cloud-Konto zu vermeiden, löschen Sie die Ressourcen, die während dieses Codelabs erstellt wurden.

  1. Knotenpool löschen:
    gcloud container node-pools delete "${NODE_POOL_NAME}" \
    --cluster="${CLUSTER_NAME}" \
    --region="${REGION}" \
    --project="${PROJECT_ID}" --quiet
  2. Cluster löschen:
    gcloud container clusters delete "${CLUSTER_NAME}" \
    --region="${REGION}" \
    --project="${PROJECT_ID}" --quiet
  3. Netzwerk- und Firewalleinrichtungen löschen:
    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. Bindung des Dienstkontos aufheben und löschen:
        # 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-Bucket löschen : Rufen Sie die Cloud Console auf, wählen Sie „Cloud Storage“ > „Buckets“ aus, wählen Sie „inf-demo-model-storage“ aus und klicken Sie dann auf „Löschen“.

12. Glückwunsch

Glückwunsch! Sie haben erfolgreich einen vLLM-Stack mit mehreren Hosts und hoher Inferenzrate bereitgestellt, der Ray nativ in Google Kubernetes Engine nutzt.

Lerninhalte

  • Benutzerdefinierte Pfade für schnellen TPU-Traffic bereitstellen
  • Gewichtungen mit GCS Fuse und regionalen Rapid Caches einbinden
  • Arbeitslast-Slices mit mehreren Hosts nativ über LeaderWorkerSets synchronisieren
  • Weitere Informationen finden Sie im vLLM-Nutzerhandbuch und in den Bereitstellungsleitfäden für llm-d.