Déployer l'inférence vLLM TPU multihôte avec Ray sur GKE

1. Introduction

Dans cet atelier de programmation, vous apprendrez à déployer des services d'inférence vLLM (Virtual Large Language Model) multi-hôtes hautes performances sur Google Kubernetes Engine (GKE) à l'aide de Google Cloud TPU. Vous configurerez l'inférence distribuée à l'aide de Ray et gérerez la charge de travail de manière native sur GKE à l'aide de LeaderWorkerSets.

Ce tutoriel simule une configuration de production pour la diffusion de grands modèles tels que Qwen 30B.

Objectifs de l'atelier

  • Créer un réseau VPC personnalisé pour le trafic de l'accélérateur.
  • Provisionner un cluster GKE avec l'opérateur Ray et le pilote CSI GCS Fuse.
  • Initialiser un cache rapide GCS pour accélérer le chargement du modèle.
  • Provisionner un pool de nœuds TPU v6e multi-hôtes avec une capacité réservée.
  • Configurer Workload Identity pour un accès sécurisé aux pondérations de modèle.
  • Déployer et tester le moteur vLLM diffusant un modèle de paramètre 30B.

Ce dont vous avez besoin

  • Un projet Google Cloud avec facturation activée.
  • Une réservation Google Cloud pour les ressources TPU v6e (32 puces, ct6e-standard-4t).
  • Accès pour copier les pondérations de modèle à partir d'un bucket source.
  • Cloud Shell ou un terminal local avec gcloud, kubectl et helm installés.
  • Durée estimée : 60 minutes
  • Coût estimé : moins de 60 $ (en supposant que le démontage suit rapidement).

2. Avant de commencer

Créer ou sélectionner un projet Google Cloud

  1. Dans la Google Cloud Console, sélectionnez ou créez un projet Google Cloud.
  2. Vérifiez que la facturation est activée pour votre projet Cloud.

Démarrer Cloud Shell

  1. Cliquez sur Activer Cloud Shell en haut de la console Google Cloud.
  2. Vérifiez l'authentification :
gcloud auth list
  1. Confirmez votre projet :
gcloud config get project
  1. Définissez-le si nécessaire :
export PROJECT_ID=<YOUR_PROJECT_ID>
gcloud config set project $PROJECT_ID

Définir des variables d'environnement

Pour faciliter l'exécution des commandes, définissez les variables suivantes dans votre shell. Remplacez <YOUR_ZONE> par la zone TPU qui vous est attribuée et <YOUR_RESERVATION_NAME> par l'ID de votre réservation. Vous devrez créer un jeton d'accès utilisateur Hugging Face pour télécharger les pondérations de modèle limitées. Une fois que vous l'avez créé, remplacez <YOUR_HUGGING_FACE_TOKEN> par le jeton que vous venez de créer.

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

Activer les API

Activez les services Google Cloud requis :

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

3. Créer une mise en réseau personnalisée

Les charges de travail TPU multi-hôtes nécessitent des configurations réseau spécifiques, y compris des tailles MTU plus élevées pour une communication efficace de l'accélérateur. Créez un réseau VPC personnalisé pour votre cluster.

  1. Créez le réseau VPC avec une MTU élevée (8896) :
    gcloud compute --project=${PROJECT_ID} \
    networks create ${GVNIC_NETWORK_PREFIX}-main \
    --subnet-mode=custom \
    --mtu=8896
  2. Créez le sous-réseau pour le 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. Créez des règles de pare-feu autorisant le trafic interne pour permettre aux nœuds de calcul de communiquer :
    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. Provisionner un cluster GKE

Créez une configuration de cluster GKE Standard configurée pour prendre en charge les montages GCS Fuse et les charges de travail de l'opérateur Ray.

  1. Créez le 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. Récupérez les identifiants du cluster :
    gcloud container clusters get-credentials ${CLUSTER_NAME} --region=${REGION}
  3. Créez un secret Hugging Face : enregistrez votre jeton de manière sécurisée pour les téléchargements d'accès aux conteneurs :
    kubectl create secret generic hf-secret \
    --from-literal=hf_api_token=${HF_TOKEN} \
    --dry-run=client -o yaml | kubectl apply -f -
  4. Installez LeaderWorkerSet (LWS) via Helm. LWS gère les groupes de pods qui doivent être planifiés ensemble :
    helm install lws oci://registry.k8s.io/lws/charts/lws \
    --version=0.7.0 \
    --namespace lws-system \
    --create-namespace \
    --wait

5. Activer le cache rapide GCS

Pour accélérer la lecture de dizaines de Go de pondérations à partir de Cloud Storage lors de la diffusion, créez un bucket GCS et activez le cache rapide GCS dans votre zone.

  1. Créez le bucket:
    gcloud storage buckets create gs://$BUCKET_NAME \
    --location=$REGION \
    --uniform-bucket-level-access
  2. Initialisez le cache rapide dans votre zone TPU :
    gcloud storage buckets anywhere-caches create gs://$BUCKET_NAME $ZONE \
    --ttl=1d \
    --admission-policy=ADMIT_ON_FIRST_MISS

6. Configurer Workload Identity et les autorisations de stockage

Configurez des liens d'identité pour monter de manière sécurisée le bucket de pondération dans vos pods GKE sans intégrer de clés à longue durée de vie.

  1. Créez un compte de service IAM dédié :
    gcloud iam service-accounts create tpu-reader-sa
  2. Accordez des autorisations de lecture de 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. Créez une liaison Workload Identity pour le compte de service Kubernetes de l'espace de noms 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. Annotez le compte de service Kubernetes :
    kubectl annotate serviceaccount default iam.gke.io/gcp-service-account=tpu-reader-sa@${PROJECT_ID}.iam.gserviceaccount.com

7. Configurer les pondérations de modèle

Pour diffuser un modèle de paramètre 30B, vous devez télécharger les pondérations de Hugging Face dans votre bucket GCS. Pour contourner la limite de quota de disque Cloud Shell (5 Go), utilisez un job Kubernetes Standard pour télécharger directement dans le cluster et écrire de manière sécurisée dans le volume GCS Fuse monté.

  1. Déployez le job de téléchargement de modèle : créez et appliquez le manifeste suivant pour lancer le téléchargement :
    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. Surveillez le téléchargement : consultez les journaux du pod de téléchargement pour suivre la progression :
    kubectl logs -f job/model-downloader
    Attendez que le job se termine avec l'état "Réussite".

8. Créer un pool de nœuds TPU réservé

Provisionnez la tranche TPU multi-hôtes réelle à l'aide de votre réservation de capacité existante.

  1. Exécutez la commande de création :
    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. Attendez que les nœuds rejoignent le cluster : vous pouvez observer directement la mise à l'échelle de l'agrégation de nœuds. Attendez que huit nœuds contenant ct6e rejoignent kubectl get nodes.

9. Déployer le service vLLM

  1. Créez des revendications réseau : vous devez demander l'environnement réseau :
    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. Déployez le point de terminaison de l'API de l'équilibreur de charge :
    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. Déployez la charge de travail LeaderWorkerSet : ce manifeste démarre l'agrégation dynamique des nœuds de calcul/de tête Ray sur les huit hôtes de tranche.
    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. Tester la réponse du déploiement

L'ensemble des pods du LeaderWorkerSet peut mettre entre 5 et 10 minutes pour extraire les images de conteneur, initialiser Ray et devenir entièrement Ready. Vous pouvez suivre l'état en observant l'initialisation des pods :

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

Attendez que les huit pods vllm-tpu-qwen- affichent STATUS comme Running et READY comme 2/2, et assurez-vous que l'équilibreur de charge a reçu une adresse IP externe avant de continuer. Cela peut prendre entre 7 et 10 minutes.

  1. Récupérez l'adresse IP externe : 
    export EXTERNAL_IP=$(kubectl get svc vllm-tpu-service -o jsonpath='{.status.loadBalancer.ingress[0].ip}')
echo $EXTERNAL_IP

Attention : Dans un service de production, ce point de terminaison doit être sécurisé avec un outil tel que Identity-Aware Proxy (IAP).

  1. Envoyez la requête d'inférence à l'aide de 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 ""
    Vous devriez voir une sortie semblable à une réponse JSON contenant le texte représentant votre inférence générée.

11. Effectuer un nettoyage

Pour éviter que votre compte Google Cloud ne soit facturé en permanence, supprimez les ressources créées lors de cet atelier de programmation.

  1. Supprimez le pool de nœuds :
    gcloud container node-pools delete "${NODE_POOL_NAME}" \
    --cluster="${CLUSTER_NAME}" \
    --region="${REGION}" \
    --project="${PROJECT_ID}" --quiet
  2. Supprimez le cluster :
    gcloud container clusters delete "${CLUSTER_NAME}" \
    --region="${REGION}" \
    --project="${PROJECT_ID}" --quiet
  3. Supprimez les configurations réseau et de pare-feu :
    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. Dissociez et supprimez le compte de service :
        # 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. Supprimez le bucket GCS : accédez à votre console Cloud, sélectionnez Cloud Storage > Buckets, sélectionnez inf-demo-model-storage, puis choisissez "Supprimer".

12. Félicitations

Félicitations ! Vous avez déployé une pile vLLM à taux d'inférence élevé multi-hôtes TPU à l'aide de Ray de manière native sur Google Kubernetes Engine.

Connaissances acquises

  • Provisionner des chemins d'accès personnalisés adaptés au trafic TPU à haut débit.
  • Monter des pondérations à l'aide de GCS Fuse et de caches rapides régionaux.
  • Orchestrer des tranches de charge de travail multi-hôtes synchronisées de manière native via LeaderWorkerSets.
  • Pour en savoir plus, consultez le guide de l'utilisateur vLLM et les guides de déploiement llm-d.