פריסת הסקת מסקנות של vLLM ב-TPU מרובה מארחים באמצעות Ray ב-GKE

1. מבוא

בשיעור Codelab הזה תלמדו איך לפרוס שירותי הסקה (inferencing) של מודל שפה גדול וירטואלי (vLLM) עם ביצועים גבוהים וריבוי מארחים ב-Google Kubernetes Engine‏ (GKE) באמצעות Google Cloud TPU. תגדירו הסקת מסקנות מבוזרת באמצעות Ray ותנהלו את עומס העבודה באופן מקורי ב-GKE באמצעות LeaderWorkerSets.

המדריך הזה מדמה הגדרת ייצור להפעלת מודלים גדולים כמו Qwen 30B.

הפעולות שתבצעו:

  • יוצרים רשת VPC מותאמת אישית לתעבורת הנתונים של המאיץ.
  • הקצאת אשכול GKE עם Ray Operator ומנהל התקן GCS Fuse CSI.
  • מפעילים GCS Rapid Cache כדי לטעון מודלים מהר יותר.
  • הקצאת מאגר צמתים של TPU v6e עם כמה מארחים וקיבולת מוזמנת.
  • הגדרת Workload Identity לגישה מאובטחת למשקלים של המודל.
  • פריסה ובדיקה של מנוע vLLM שמשרת מודל עם 30 מיליארד פרמטרים.

הדרישות

  • פרויקט ב-Google Cloud שהחיוב בו מופעל.
  • הזמנה ב-Google Cloud למשאבי TPU v6e‏ (32 שבבים, ct6e-standard-4t).
  • גישה להעתקת משקלי המודל מקטגוריית מקור.
  • ‫Cloud Shell או טרמינל מקומי עם gcloud,‏ kubectl ו-helm מותקנים.
  • משך זמן משוער: 60 דקות
  • עלות משוערת: פחות מ-60$ (בהנחה שהפירוק יתבצע מיד).

‫2. לפני שמתחילים

יצירה או בחירה של פרויקט ב-Google Cloud

  1. ב-מסוף Google Cloud, בוחרים פרויקט או יוצרים פרויקט חדש ב-Google Cloud.
  2. מוודאים שהחיוב מופעל בפרויקט בענן שלכם.

הפעלת Cloud Shell

  1. לוחצים על Activate Cloud Shell בחלק העליון של מסוף Google Cloud.
  2. אימות האימות:
gcloud auth list
  1. מאשרים את הפרויקט:
gcloud config get project
  1. מגדירים אותו לפי הצורך:
export PROJECT_ID=<YOUR_PROJECT_ID>
gcloud config set project $PROJECT_ID

הגדרה של משתני סביבה

כדי להקל על הרצת הפקודות, מגדירים את המשתנים הבאים במעטפת. מחליפים את <YOUR_ZONE> באזור ה-TPU שהוקצה לכם ואת <YOUR_RESERVATION_NAME> במזהה ההזמנה. כדי להוריד משקלים של מודלים עם גישה מוגבלת, צריך ליצור טוקן גישה למשתמש ב-Hugging Face. אחרי שיוצרים את האסימון, מחליפים את <YOUR_HUGGING_FACE_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

הפעלת ממשקי ה-API

מפעילים את שירותי Google Cloud הנדרשים:

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

3. יצירת רשת בהתאמה אישית

עומסי עבודה של TPU עם כמה מארחים דורשים הגדרות רשת ספציפיות, כולל גדלים גדולים יותר של MTU לתקשורת יעילה בין המאיצים. יוצרים רשת VPC בהתאמה אישית לאשכול.

  1. יוצרים את רשת ה-VPC עם MTU גדול (8896):
    gcloud compute --project=${PROJECT_ID} \
    networks create ${GVNIC_NETWORK_PREFIX}-main \
    --subnet-mode=custom \
    --mtu=8896
  2. יוצרים את רשת המשנה לאשכול:
    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. יוצרים כללים של חומת אש שמאפשרים תעבורה פנימית כדי לאפשר לעובדים לתקשר:
    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

יוצרים הגדרת אשכול GKE רגיל שמוגדר לתמיכה בחיבורי GCS Fuse ובעומסי עבודה של Ray Operator.

  1. יוצרים את האשכול:
    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. אחזור פרטי הכניסה של האשכול:
    gcloud container clusters get-credentials ${CLUSTER_NAME} --region=${REGION}
  3. יצירת סוד ב-Hugging Face: שומרים את האסימון בצורה מאובטחת כדי להוריד קבצים מהמאגר:
    kubectl create secret generic hf-secret \
    --from-literal=hf_api_token=${HF_TOKEN} \
    --dry-run=client -o yaml | kubectl apply -f -
  4. מתקינים את LeaderWorkerSet ‏ (LWS) באמצעות Helm. שירות LWS מנהל קבוצות של פודים שצריך לתזמן יחד:
    helm install lws oci://registry.k8s.io/lws/charts/lws \
    --version=0.7.0 \
    --namespace lws-system \
    --create-namespace \
    --wait

5. הפעלת מטמון מהיר ב-GCS

כדי להאיץ את קריאת המשקלים של עשרות גיגה-בייט מ-Cloud Storage במהלך ההצגה, צריך ליצור קטגוריה ב-GCS ולהפעיל את GCS Rapid Cache באזור.

  1. יצירת הקטגוריה:
    gcloud storage buckets create gs://$BUCKET_NAME \
    --location=$REGION \
    --uniform-bucket-level-access
  2. מפעילים את Rapid Cache באזור ה-TPU:
    gcloud storage buckets anywhere-caches create gs://$BUCKET_NAME $ZONE \
    --ttl=1d \
    --admission-policy=ADMIT_ON_FIRST_MISS

6. הגדרה של Workload Identity והרשאות אחסון

אפשר להגדיר קישורי זהויות כדי לטעון בצורה מאובטחת את מאגר המשקלים בתרמילי GKE בלי להטמיע מפתחות לטווח ארוך.

  1. יוצרים חשבון שירות ייעודי ב-IAM:
    gcloud iam service-accounts create tpu-reader-sa
  2. מתן הרשאות קריאה לקטגוריה:
    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 Binding לחשבון השירות של Kubernetes במרחב השמות 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. הוספת הערות ל-SA של Kubernetes:
    kubectl annotate serviceaccount default iam.gke.io/gcp-service-account=tpu-reader-sa@${PROJECT_ID}.iam.gserviceaccount.com

7. הגדרת משקלים של מודלים

כדי להכניס לשימוש בסביבת הייצור מודל עם 30 מיליארד פרמטרים, צריך להוריד משקלים מ-Hugging Face לקטגוריה שלכם ב-GCS. כדי לעקוף את מגבלת מכסת האחסון בדיסק של Cloud Shell (5GB), אפשר להשתמש ב-Job רגיל של Kubernetes כדי להוריד ישירות לתוך האשכול ולכתוב לנפח GCS Fuse המצורף בצורה מאובטחת.

  1. פריסת המשימה Model Downloader: יוצרים את המניפסט הבא ומחילים אותו כדי להתחיל את ההורדה:
    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. מעקב אחר ההורדה: בודקים את היומנים של ה-pod של כלי ההורדה כדי לעקוב אחר ההתקדמות:
    kubectl logs -f job/model-downloader
    מחכים עד שהעבודה מסתיימת עם סטטוס הצלחה.

8. יצירת מאגר צמתים שמורים של TPU

הקצאת חלקה של TPU עם כמה מארחים באמצעות הזמנת הקיבולת הקיימת.

  1. מריצים את פקודת היצירה:
    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. המתנה להצטרפות של צמתים: אפשר לראות את שינוי הגודל של צבירת הצמתים ישירות. צריך לחכות עד ש-8 צמתים שמכילים ct6e יצטרפו ל-kubectl get nodes.

9. פריסת שירות vLLM

  1. יצירת תביעות בעלות על רשתות: צריך לבקש את סביבת הרשת:
    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. פריסת נקודת קצה של API למאזן עומסים:
    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: המניפסט הזה מתחיל צבירה של Ray head/worker באופן דינמי ב-8 מארחי הפרוסות.
    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. תגובה לבדיקת פריסה

יכול להיות שיעברו 5-10 דקות עד שכל הפודים ב-LeaderWorkerSet ימשכו תמונות של קונטיינרים, יאתחלו את Ray ויהפכו לזמינים באופן מלא Ready. אפשר לעקוב אחרי הסטטוס באמצעות צפייה באתחול של ה-Pod:

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

מחכים עד שכל 8 ה-vllm-tpu-qwen- pods יציגו STATUS כ-Running ו-READY כ-2/2, ומוודאים שמאזן העומסים קיבל כתובת IP חיצונית לפני שממשיכים. התהליך יכול להימשך 7 עד 10 דקות.

  1. אחזור כתובת IP חיצונית: 
    export EXTERNAL_IP=$(kubectl get svc vllm-tpu-service -o jsonpath='{.status.loadBalancer.ingress[0].ip}')
echo $EXTERNAL_IP

זהירות: בשירות הפקה, נקודת הקצה הזו צריכה להיות מאובטחת באמצעות משהו כמו שרת proxy לאימות זהויות (IAP)

  1. שליחת בקשת הסקה באמצעות 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 ""
    אמור להופיע פלט שדומה לתגובת JSON שמכילה את הטקסט שמייצג את ההסקת המסקנות שנוצרה!

11. הסרת המשאבים

כדי להימנע מחיובים שוטפים בחשבון Google Cloud, מוחקים את המשאבים שנוצרו במהלך ה-codelab הזה.

  1. מחיקת מאגר צמתים:
    gcloud container node-pools delete "${NODE_POOL_NAME}" \
    --cluster="${CLUSTER_NAME}" \
    --region="${REGION}" \
    --project="${PROJECT_ID}" --quiet
  2. מחיקת האשכול:
    gcloud container clusters delete "${CLUSTER_NAME}" \
    --region="${REGION}" \
    --project="${PROJECT_ID}" --quiet
  3. מחיקת הגדרות של רשת וחומת אש:
    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. ביטול הקישור ומחיקה של חשבון שירות:
        # 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 עוברים למסוף הענן, בוחרים באפשרות Cloud Storage -> Buckets, בוחרים באפשרות inf-demo-model-storage ואז בוחרים באפשרות Delete.

12. מזל טוב

מעולה! הצלחתם לפרוס מחסנית vLLM עם קצב גבוה של הסקת מסקנות TPU מרובת-מארחים באמצעות Ray באופן מקורי ב-Google Kubernetes Engine.

מה למדתם

  • הקצאת נתיבים מותאמים אישית שמותאמים לתעבורת נתונים מהירה של TPU.
  • העלאת משקלים באמצעות GCS Fuse ומטמונים אזוריים מהירים.
  • תזמור של פרוסות עומס עבודה מרובות מארחים שמסונכרנות באופן טבעי באמצעות LeaderWorkerSets.
  • מידע נוסף זמין במדריך למשתמש של vLLM ובמדריכי הפריסה של llm-d.