Men-deploy Inferensi vLLM TPU Multi-host dengan Ray di GKE

1. Pengantar

Dalam codelab ini, Anda akan mempelajari cara men-deploy layanan inferensi vLLM (Virtual Large Language Model) multi-host berperforma tinggi di Google Kubernetes Engine (GKE) menggunakan TPU Google Cloud. Anda akan mengonfigurasi inferensi terdistribusi menggunakan Ray dan mengelola workload secara native di GKE menggunakan LeaderWorkerSets.

Panduan ini menyimulasikan penyiapan produksi untuk menyajikan model besar seperti Qwen 30B.

Yang akan Anda lakukan

  • Buat jaringan VPC kustom untuk traffic akselerator.
  • Sediakan cluster GKE dengan driver CSI GCS Fuse dan Ray Operator.
  • Lakukan inisialisasi GCS Rapid Cache untuk pemuatan model yang dipercepat.
  • Sediakan node pool TPU v6e multi-host dengan kapasitas yang dipesan.
  • Konfigurasi Workload Identity untuk akses yang aman ke bobot model.
  • Deploy dan uji mesin vLLM yang melayani model parameter 30B.

Yang Anda butuhkan

  • Project Google Cloud yang mengaktifkan penagihan.
  • Reservasi Google Cloud untuk resource TPU v6e (32 chip, ct6e-standard-4t).
  • Akses untuk menyalin bobot model dari bucket sumber.
  • Cloud Shell atau terminal lokal dengan gcloud, kubectl, dan helm yang terinstal.
  • Perkiraan Durasi: 60 menit
  • Perkiraan Biaya: Di bawah $60 (dengan asumsi pembongkaran dilakukan segera).

2. Sebelum memulai

Membuat atau Memilih Project Google Cloud

  1. Di Konsol Google Cloud, pilih atau buat project Google Cloud.
  2. Pastikan penagihan diaktifkan untuk project Cloud Anda.

Mulai Cloud Shell

  1. Klik Activate Cloud Shell di bagian atas konsol Google Cloud.
  2. Verifikasi autentikasi:
gcloud auth list
  1. Konfirmasi project Anda:
gcloud config get project
  1. Tetapkan jika perlu:
export PROJECT_ID=<YOUR_PROJECT_ID>
gcloud config set project $PROJECT_ID

Menetapkan Variabel Lingkungan

Untuk mempermudah eksekusi perintah, tentukan variabel berikut di shell Anda. Ganti <YOUR_ZONE> dengan zona TPU yang dialokasikan dan <YOUR_RESERVATION_NAME> dengan ID reservasi Anda. Anda harus membuat Token Akses Pengguna Hugging Face untuk mendownload bobot model yang dibatasi. Setelah Anda membuatnya, ganti <YOUR_HUGGING_FACE_TOKEN> dengan token yang baru dibuat.

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

Mengaktifkan API

Aktifkan layanan Google Cloud yang diperlukan:

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

3. Membuat Jaringan Kustom

Workload TPU multi-host memerlukan konfigurasi jaringan tertentu, termasuk ukuran MTU yang lebih tinggi untuk komunikasi akselerator yang efisien. Buat jaringan VPC kustom untuk cluster Anda.

  1. Buat jaringan VPC dengan MTU besar (8896):
    gcloud compute --project=${PROJECT_ID} \
    networks create ${GVNIC_NETWORK_PREFIX}-main \
    --subnet-mode=custom \
    --mtu=8896
  2. Buat subnet untuk 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. Buat aturan firewall yang mengizinkan traffic internal agar pekerja dapat berkomunikasi:
    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. Menyediakan Cluster GKE

Buat penyiapan cluster GKE Standar yang dikonfigurasi untuk mendukung pemasangan GCS Fuse dan beban kerja Ray Operator.

  1. Buat 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. Ambil Kredensial Cluster:
    gcloud container clusters get-credentials ${CLUSTER_NAME} --region=${REGION}
  3. Buat Secret Hugging Face: Simpan token Anda dengan aman untuk download akses penampung:
    kubectl create secret generic hf-secret \
    --from-literal=hf_api_token=${HF_TOKEN} \
    --dry-run=client -o yaml | kubectl apply -f -
  4. Instal LeaderWorkerSet (LWS) melalui Helm. LWS mengelola grup pod yang harus dijadwalkan bersama:
    helm install lws oci://registry.k8s.io/lws/charts/lws \
    --version=0.7.0 \
    --namespace lws-system \
    --create-namespace \
    --wait

5. Mengaktifkan GCS Rapid Cache

Untuk mempercepat pembacaan bobot berukuran puluhan GB dari Cloud Storage selama penayangan, buat bucket GCS dan aktifkan GCS Rapid Cache di zona Anda.

  1. Buat bucket:
    gcloud storage buckets create gs://$BUCKET_NAME \
    --location=$REGION \
    --uniform-bucket-level-access
  2. Lakukan inisialisasi Rapid Cache di zona TPU Anda:
    gcloud storage buckets anywhere-caches create gs://$BUCKET_NAME $ZONE \
    --ttl=1d \
    --admission-policy=ADMIT_ON_FIRST_MISS

6. Menyiapkan Izin Penyimpanan & Workload Identity

Konfigurasi penautan identitas untuk memasang bucket berat ke pod GKE Anda secara aman tanpa menyematkan kunci yang berlaku lama.

  1. Buat Akun Layanan IAM khusus:
    gcloud iam service-accounts create tpu-reader-sa
  2. Memberikan izin Baca 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. Buat Binding Workload Identity untuk Akun Layanan Kubernetes 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. Anotasikan SA Kubernetes:
    kubectl annotate serviceaccount default iam.gke.io/gcp-service-account=tpu-reader-sa@${PROJECT_ID}.iam.gserviceaccount.com

7. Penyiapan Bobot Model

Untuk menayangkan model parameter 30B, Anda perlu mendownload bobot dari Hugging Face ke bucket GCS Anda. Untuk melewati batas kuota disk Cloud Shell (5 GB), gunakan Job Kubernetes Standar untuk mendownload langsung di dalam cluster dan menulis ke volume GCS Fuse yang terpasang secara aman.

  1. Deploy Tugas Model Downloader: Buat dan terapkan manifes berikut untuk memulai 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. Memantau Download: Periksa log pod downloader untuk mengikuti progres:
    kubectl logs -f job/model-downloader
    Tunggu hingga tugas selesai dengan status berhasil.

8. Membuat Node Pool TPU yang dipesan

Sediakan slice TPU multi-host yang sebenarnya menggunakan reservasi kapasitas yang ada.

  1. Jalankan perintah pembuatan:
    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. Menunggu node bergabung: Anda dapat mengamati penskalaan agregasi node secara langsung. Tunggu hingga 8 node yang berisi ct6e bergabung ke kubectl get nodes.

9. Men-deploy Layanan vLLM

  1. Buat Klaim Jaringan: Anda perlu meminta lingkungan jaringan:
    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. Deploy Load Balancer API Endpoint:
    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: Manifes ini memulai agregasi head/worker Ray secara dinamis di 8 host 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. Respons Deployment Pengujian

Diperlukan waktu 5-10 menit agar semua pod di LeaderWorkerSet menarik image container, menginisialisasi Ray, dan menjadi Ready sepenuhnya. Anda dapat melacak status dengan memantau inisialisasi pod:

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

Tunggu hingga semua 8 pod vllm-tpu-qwen- menampilkan STATUS sebagai Running dan READY sebagai 2/2, dan pastikan load balancer telah menerima IP Eksternal sebelum melanjutkan. Proses ini dapat memakan waktu 7-10 menit.

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

Perhatian: Dalam layanan produksi, endpoint ini harus diamankan dengan sesuatu seperti Identity-Aware Proxy (IAP)

  1. Kirim permintaan inferensi menggunakan 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 ""
    Anda akan melihat output yang menyerupai respons JSON yang berisi teks yang merepresentasikan inferensi yang dihasilkan.

11. Pembersihan

Untuk menghindari biaya berkelanjutan pada akun Google Cloud Anda, hapus resource yang dibuat selama codelab ini.

  1. Hapus Node Pool:
    gcloud container node-pools delete "${NODE_POOL_NAME}" \
    --cluster="${CLUSTER_NAME}" \
    --region="${REGION}" \
    --project="${PROJECT_ID}" --quiet
  2. Hapus Cluster:
    gcloud container clusters delete "${CLUSTER_NAME}" \
    --region="${REGION}" \
    --project="${PROJECT_ID}" --quiet
  3. Menghapus penyiapan Jaringan dan 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. Batalkan Pengikatan dan Hapus Akun Layanan:
        # 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. Hapus Bucket GCS Buka konsol cloud Anda, pilih Cloud Storage -> Buckets, pilih inf-demo-model-storage, lalu pilih 'Delete'.

12. Selamat

Selamat! Anda telah berhasil men-deploy stack vLLM dengan kecepatan inferensi tinggi TPU multi-host yang memanfaatkan Ray secara native melalui Google Kubernetes Engine.

Yang telah Anda pelajari

  • Menyediakan jalur kustom yang disesuaikan untuk traffic TPU berkecepatan tinggi.
  • Memasang bobot menggunakan GCS Fuse dan cache cepat regional.
  • Mengorkestrasi slice workload multi-host yang disinkronkan secara native melalui LeaderWorkerSets.
  • Untuk mempelajari lebih lanjut, lihat Panduan Pengguna vLLM dan Panduan Deployment llm-d