Wdrażanie wnioskowania vLLM TPU na wielu hostach za pomocą Ray w GKE

1. Wprowadzenie

Z tego ćwiczenia dowiesz się, jak wdrożyć w Google Kubernetes Engine (GKE) usługi wnioskowania vLLM (Virtual Large Language Model) o wysokiej wydajności i wielu hostach przy użyciu jednostek Cloud TPU od Google. Skonfigurujesz rozproszone wnioskowanie za pomocą Ray i będziesz zarządzać zadaniami natywnie w GKE za pomocą LeaderWorkerSets.

Ten przewodnik symuluje konfigurację produkcyjną do obsługi dużych modeli, takich jak Qwen 30B.

Jakie zadania wykonasz

  • Utwórz niestandardową sieć VPC na potrzeby ruchu akceleratora.
  • Aprowizuj klaster GKE za pomocą operatora Ray i sterownika CSI GCS Fuse.
  • Zainicjuj szybką pamięć podręczną GCS, aby przyspieszyć wczytywanie modelu.
  • Utwórz pulę węzłów TPU v6e z wieloma hostami i zarezerwowaną pojemnością.
  • Skonfiguruj Workload Identity, aby zapewnić bezpieczny dostęp do wag modelu.
  • Wdróż i przetestuj silnik vLLM udostępniający model z 30 miliardami parametrów.

Czego potrzebujesz

  • Projekt Google Cloud z włączonymi płatnościami.
  • Rezerwacja Google Cloud dla zasobów TPU v6e (32 chipy, ct6e-standard-4t).
  • Dostęp do kopiowania wag modelu z zasobnika źródłowego.
  • Cloud Shell lub terminal lokalny z zainstalowanymi narzędziami gcloud, kubectl i helm.
  • Szacowany czas trwania: 60 minut
  • Szacowany koszt: poniżej 60 USD (zakładając, że analiza zostanie przeprowadzona szybko).

2. Zanim zaczniesz

Tworzenie lub wybieranie projektu Google Cloud

  1. W konsoli Google Cloud wybierz lub utwórz projekt w chmurze Google Cloud.
  2. Sprawdź, czy w projekcie w chmurze włączone są płatności.

Uruchamianie Cloud Shell

  1. Kliknij Aktywuj Cloud Shell u góry konsoli Google Cloud.
  2. Potwierdź uwierzytelnianie:
gcloud auth list
  1. Potwierdź projekt:
gcloud config get project
  1. W razie potrzeby ustaw:
export PROJECT_ID=<YOUR_PROJECT_ID>
gcloud config set project $PROJECT_ID

Ustawianie zmiennych środowiskowych

Aby ułatwić wykonywanie poleceń, zdefiniuj w powłoce te zmienne. Zastąp <YOUR_ZONE> przydzieloną strefą TPU, a <YOUR_RESERVATION_NAME> identyfikatorem rezerwacji. Aby pobrać wagi modelu z ograniczonym dostępem, musisz utworzyć token dostępu użytkownika Hugging Face. Po utworzeniu tokena zastąp ciąg <YOUR_HUGGING_FACE_TOKEN> nowo utworzonym tokenem.

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

Włącz interfejsy API

Włącz wymagane usługi Google Cloud:

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

3. Tworzenie sieci niestandardowej

Zbiory zadań TPU na wielu hostach wymagają określonych konfiguracji sieci, w tym większych rozmiarów MTU, aby zapewnić wydajną komunikację akceleratora. Utwórz niestandardową sieć VPC dla klastra.

  1. Utwórz sieć VPC z dużą wartością MTU (8896):
    gcloud compute --project=${PROJECT_ID} \
    networks create ${GVNIC_NETWORK_PREFIX}-main \
    --subnet-mode=custom \
    --mtu=8896
  2. Utwórz podsieć dla klastra:
    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. Utwórz reguły zapory sieciowej zezwalające na ruch wewnętrzny, aby umożliwić komunikację między procesami roboczymi:
    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. Aprowizowanie klastra GKE

Utwórz konfigurację standardowego klastra GKE skonfigurowaną do obsługi montowania GCS Fuse i zbiorów zadań Ray Operator.

  1. Utwórz klaster:
    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. Pobierz dane logowania do klastra:
    gcloud container clusters get-credentials ${CLUSTER_NAME} --region=${REGION}
  3. Utwórz tajny klucz Hugging Face: bezpiecznie zapisz token, aby pobierać kontenery:
    kubectl create secret generic hf-secret \
    --from-literal=hf_api_token=${HF_TOKEN} \
    --dry-run=client -o yaml | kubectl apply -f -
  4. Zainstaluj LeaderWorkerSet (LWS) za pomocą Helm. LWS zarządza grupami podów, które muszą być zaplanowane razem:
    helm install lws oci://registry.k8s.io/lws/charts/lws \
    --version=0.7.0 \
    --namespace lws-system \
    --create-namespace \
    --wait

5. Włączanie szybkiej pamięci podręcznej GCS

Aby przyspieszyć odczytywanie dziesiątek gigabajtów wag z Cloud Storage podczas obsługi, utwórz zasobnik GCS i włącz w swojej strefie szybką pamięć podręczną GCS.

  1. Utwórz zasobnik:
    gcloud storage buckets create gs://$BUCKET_NAME \
    --location=$REGION \
    --uniform-bucket-level-access
  2. Zainicjuj szybką pamięć podręczną w strefie TPU:
    gcloud storage buckets anywhere-caches create gs://$BUCKET_NAME $ZONE \
    --ttl=1d \
    --admission-policy=ADMIT_ON_FIRST_MISS

6. Konfigurowanie Workload Identity i uprawnień do przechowywania

Skonfiguruj połączenia tożsamości, aby bezpiecznie zamontować zasobnik wagi w podach GKE bez osadzania długotrwałych kluczy.

  1. Utwórz dedykowane konto usługi Uprawnienia:
    gcloud iam service-accounts create tpu-reader-sa
  2. Przyznaj uprawnienia do odczytu zasobnika:
    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. Utwórz wiązanie Workload Identity dla konta usługi Kubernetes w przestrzeni nazw 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. Dodaj adnotację do konta usługi Kubernetes:
    kubectl annotate serviceaccount default iam.gke.io/gcp-service-account=tpu-reader-sa@${PROJECT_ID}.iam.gserviceaccount.com

7. Konfiguracja wag modelu

Aby udostępnić model z 30 miliardami parametrów, musisz pobrać wagi z Hugging Face do zasobnika GCS. Aby ominąć limit miejsca na dysku Cloud Shell (5 GB), użyj standardowego zadania Kubernetes, aby pobrać dane bezpośrednio do klastra i bezpiecznie zapisać je w zamontowanym woluminie GCS Fuse.

  1. Wdróż zadanie pobierania modelu: utwórz i zastosuj ten plik manifestu, aby rozpocząć pobieranie:
    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. Monitorowanie pobierania: sprawdź logi poda pobierającego, aby śledzić postępy:
    kubectl logs -f job/model-downloader
     Poczekaj, aż zadanie się zakończy i zostanie wyświetlony stan „Ukończono”.

8. Tworzenie zarezerwowanej puli węzłów TPU

Zainicjuj rzeczywisty wycinek TPU z wieloma hostami, korzystając z istniejącej rezerwacji mocy obliczeniowej.

  1. Uruchom polecenie tworzenia:
    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. Czekaj na dołączenie węzłów: możesz bezpośrednio obserwować skalowanie agregacji węzłów. Poczekaj, aż 8 węzłów zawierających ct6e dołączy do kubectl get nodes.

9. Wdrażanie usługi vLLM

  1. Tworzenie roszczeń sieciowych: musisz poprosić o środowisko sieciowe:
    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. Wdróż punkt końcowy interfejsu API usługi Load Balancer:
    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. Wdróż zadanie LeaderWorkerSet: ten plik manifestu dynamicznie uruchamia agregację węzłów głównych i roboczych Ray na 8 hostach z podziałem.
    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. Odpowiedź na wdrożenie testowe

Pobranie obrazów kontenerów, zainicjowanie Ray i osiągnięcie stanu pełnej Ready przez wszystkie pody w LeaderWorkerSet może potrwać 5–10 minut. Stan możesz śledzić, obserwując inicjowanie poda:

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

Zanim przejdziesz dalej, poczekaj, aż wszystkie 8 podów vllm-tpu-qwen- będzie miało stan STATUS jako Running i READY jako 2/2, a system równoważenia obciążenia otrzyma zewnętrzny adres IP. Może to potrwać 7–10 minut.

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

Uwaga: w usłudze produkcyjnej ten punkt końcowy powinien być zabezpieczony za pomocą funkcji Identity Aware Proxy (IAP).

  1. Prześlij prośbę o wnioskowanie za pomocą 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 ""
    Powinny się wyświetlić dane wyjściowe przypominające odpowiedź w formacie JSON zawierającą tekst wygenerowanego wnioskowania.

11. Czyszczenie danych

Aby uniknąć obciążenia konta Google Cloud bieżącymi opłatami, usuń zasoby utworzone podczas tego ćwiczenia.

  1. Usuń pulę węzłów:
    gcloud container node-pools delete "${NODE_POOL_NAME}" \
    --cluster="${CLUSTER_NAME}" \
    --region="${REGION}" \
    --project="${PROJECT_ID}" --quiet
  2. Usuń klaster:
    gcloud container clusters delete "${CLUSTER_NAME}" \
    --region="${REGION}" \
    --project="${PROJECT_ID}" --quiet
  3. Usuwanie konfiguracji sieci i zapory sieciowej:
    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. Odłączanie i usuwanie konta usługi:
        # 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. Usuń zasobnik GCS: otwórz konsolę w chmurze, wybierz Cloud Storage –> Zasobniki, wybierz inf-demo-model-storage, a następnie kliknij „Usuń”.

12. Gratulacje

Gratulacje! Udało Ci się wdrożyć stos vLLM z wieloma hostami TPU i wysoką częstotliwością wnioskowania, który natywnie wykorzystuje Ray w Google Kubernetes Engine.

Czego się dowiedziałeś(-aś)

  • udostępnianie niestandardowych ścieżek dostosowanych do szybkiego ruchu TPU.
  • Wagi montażowe wykorzystujące GCS Fuse i regionalne szybkie pamięci podręczne.
  • Orkiestracja fragmentów zadań obejmujących wiele hostów, które są natywnie synchronizowane za pomocą LeaderWorkerSets.
  • Więcej informacji znajdziesz w przewodniku użytkownika vLLM i przewodnikach wdrażania llm-d.