Offene Gemma-Modelle mit TPUs in GKE mit Ray LLM bereitstellen

In dieser Anleitung wird beschrieben, wie Sie einen TPU-Inferenzdienst mit mehreren Hosts mithilfe von Ray Serve LLM bereitstellen. Durch die Nutzung der nativen TPU-Unterstützung von Ray, um verteilte Engine-Worker atomar über komplexe Beschleunigertopologien hinweg zu planen, können Sie große Modelle für die Inferenz über einen TPU-Slice mit mehreren Hosts hinweg bereitstellen.

Dieses Tutorial richtet sich an ML-Entwickler, Plattformadministratoren und ‑operatoren sowie an Daten- und KI-Spezialisten, die daran interessiert sind, Kubernetes-Container-Orchestrierungsfunktionen für die Bereitstellung von KI-/ML-Arbeitslasten auf verteilten TPU-Slices mit mehreren Hosts zu nutzen. Weitere Informationen zu gängigen Rollen und Beispielaufgaben, auf die wir in Google Cloud Inhalten verweisen, finden Sie unter Häufig verwendete GKE-Nutzerrollen und -Aufgaben.

Machen Sie sich vor dem Lesen dieser Seite mit den folgenden Themen vertraut:

Hintergrund

In diesem Abschnitt werden die in diesem Leitfaden verwendeten Schlüsseltechnologien beschrieben.

TPUs

Mit Tensor Processing Units (TPUs) können Sie bestimmte Arbeitslasten wie maschinelles Lernen und Datenverarbeitung beschleunigen, die auf Ihren Knoten ausgeführt werden. Der Hauptvorteil von TPUs ist die Leistung bei großem Umfang. In dieser Anleitung wird TPU Trillium verwendet, die sechste Generation von Cloud TPU. TPU-Slices mit mehreren Hosts bestehen aus mehreren physischen Knoten, die über eine Hochgeschwindigkeitsverbindung (Inter-Chip Interconnect, ICI) kommunizieren. Das eignet sich gut für Serving mit hohem Durchsatz und niedriger Latenz.

vLLM auf Ray

vLLM ist eine LLM-Bereitstellungs-Engine mit hohem Durchsatz und geringem Speicherbedarf. Durch die Integration in Ray Serve kann vLLM über mehrere Hosts skaliert werden und nativ auf physische Hardwaretopologien zugreifen. In dieser Anleitung wird gezeigt, wie Sie die Bereitstellungen LLMConfig und LLMServer von Ray Serve verwenden, um die vLLM-Inferenz über mehrere Host-Slices hinweg zu orchestrieren. Das Framework übernimmt dabei automatisch die Verteilung der Topologie und die Verteilung der Placement-Gruppe.

Ziele

Diese Anleitung bietet eine Grundlage für das Verständnis und die praktische Bereitstellung von LLMs für die Inferenz in einer verwalteten Kubernetes-Umgebung, in der Multi-Host-TPUs verwendet werden.

  1. Bereiten Sie Ihre Umgebung mit einem GKE-Cluster im Autopilot- oder Standardmodus vor.
  2. Erstellen Sie ein benutzerdefiniertes Container-Image mit integrierten Abhängigkeiten.
  3. Stellen Sie ein Ray LLM-Python-Script in Ihrem Cluster bereit, um die vLLM-Inferenz über einen TPU-Slice zu orchestrieren.
  4. Ray LLM verwenden, um das Gemma 4-Modell über curl und eine optionale Webchat-Oberfläche bereitzustellen.

Hinweis

Führen Sie die folgenden Aufgaben aus, bevor Sie beginnen:

  • Aktivieren Sie die Google Kubernetes Engine API.
    • Google Kubernetes Engine API aktivieren
  • Wenn Sie die Google Cloud CLI für diesen Task verwenden möchten, müssen Sie die gcloud CLI installieren und dann initialisieren. Wenn Sie die gcloud CLI bereits installiert haben, rufen Sie die neueste Version mit dem Befehl gcloud components update ab. In früheren gcloud CLI-Versionen werden die Befehle in diesem Dokument möglicherweise nicht unterstützt.
  • Prüfen Sie, ob Ihr Projekt ein ausreichendes Kontingent für TPU Trillium-Kapazität (v6e) in der ausgewählten Region hat. Weitere Informationen finden Sie unter Cloud TPU-Kontingente.
  • Ihr GKE-Cluster muss GKE Dataplane V2 verwenden und die Versionsanforderungen für DRANET erfüllen: 1.35.2-gke.1842000 oder höher für Standard und Autopilot.
  • Prüfen Sie, ob Sie die folgenden IAM-Rollen haben:
    • roles/container.admin
    • roles/iam.serviceAccountAdmin

Umgebung vorbereiten

In dieser Anleitung verwenden Sie Cloud Shell zum Verwalten von Ressourcen, die in Google Cloudgehostet werden. Die Software, die Sie für diese Anleitung benötigen, ist in Cloud Shell vorinstalliert, einschließlich kubectl und der gcloud CLI.

So richten Sie Ihre Umgebung mit Cloud Shell ein:

  1. Starten Sie in der Google Cloud Console eine Cloud Shell-Sitzung, indem Sie auf Cloud Shell aktivieren Button zum Aktivieren von Cloud Shell klicken. Dadurch wird im unteren Bereich der Google Cloud Console eine Sitzung gestartet.

  2. Erstellen und aktivieren Sie eine virtuelle Python-Umgebung:

    python3 -m venv ray-env
source ray-env/bin/activate

  3. Installieren Sie die Ray-Befehlszeile:

    pip install "ray"

  4. Legen Sie die Standardumgebungsvariablen fest:

    export PROJECT_ID=$(gcloud config get project)
export CLUSTER_NAME=ray-llm-cluster
export REGION=REGION
export ZONE=ZONE
export NAMESPACE=default
export KSA_NAME=ray-ksa
export GSA_NAME=tpu-reader-sa
export NETWORK_NAME=${CLUSTER_NAME}-net
export GS_BUCKET=BUCKET_NAME
export REPO_NAME=ray-repo
export CUSTOM_IMAGE_URI=REGION-docker.pkg.dev/PROJECT_ID/REPOSITORY/vllm-tpu-ray:vllm-tpu

    Ersetzen Sie Folgendes:

    • PROJECT_ID: Projekt-ID in Google Cloud .
    • CLUSTER_NAME: Der Name Ihres Clusters.
    • REGION: Die Region, in der Ihre TPU Trillium-Kapazität verfügbar ist.
    • ZONE: Die Zone, in der Ihre TPU Trillium-Kapazität verfügbar ist. Weitere Informationen finden Sie unter TPU-Verfügbarkeit in GKE.
    • REPOSITORY: der Name Ihres Artifact Registry-Repositorys
    • BUCKET_NAME: Der Name Ihres Storage-Buckets.

Google Cloud -Ressourcen erstellen und konfigurieren

Folgen Sie dieser Anleitung, um die erforderlichen Ressourcen zu erstellen.

GKE-Cluster und -Knotenpool erstellen

Sie können Gemma auf TPUs in einem GKE Autopilot- oder Standard-Cluster bereitstellen. GKE-verwaltetes DRANET fordert dynamisch leistungsstarke Netzwerkressourcen für Ihre verteilten Pods an und verwaltet sie. So kann GKE automatisch sekundäre Hochgeschwindigkeitsnetzwerke für die Kommunikation zwischen Beschleunigern bereitstellen, ohne dass eine manuelle VPC-Einrichtung erforderlich ist.

Autopilot

  1. Erstellen Sie den Autopilot-Cluster in Cloud Shell:

    gcloud container clusters create-auto ${CLUSTER_NAME} \
    --project=${PROJECT_ID} \
    --enable-ray-operator \
    --location=${REGION}

  2. Konfigurieren Sie kubectl für die Kommunikation mit Ihrem Cluster:

    gcloud container clusters get-credentials ${CLUSTER_NAME} \
    --location=${REGION}

  3. Wenn Sie das von GKE verwaltete DRANET im Autopilot-Modus verwenden möchten, stellen Sie die benutzerdefinierte ComputeClass-Ressource bereit, die im Repository bereitgestellt wird, um dynamische Netzwerke zu aktivieren:

    apiVersion: cloud.google.com/v1
kind: ComputeClass
metadata:
  name: dranet-compute-class
spec:
  nodePoolAutoCreation:
    enabled: true
  nodePoolConfig:
    dra:
      networking:
        enabled: true
  priorities:
  - machineType: ct6e-standard-4t
    acceleratorNetworkProfile: auto

  4. Wenden Sie das Manifest auf Ihren Cluster an:

    kubectl apply -f ai-ml/gke-ray/rayserve/llm/tpu/networking/dranet-compute-class.yaml

Standard

  1. Erstellen Sie in Cloud Shell einen Standardcluster, in dem der Ray-Operator aktiviert ist und GKE Dataplane V2 verwendet wird:

    gcloud container clusters create ${CLUSTER_NAME} \
    --project=${PROJECT_ID} \
    --addons=RayOperator,GcsFuseCsiDriver \
    --machine-type=n2-standard-8 \
    --enable-dataplane-v2 \
    --workload-pool=${PROJECT_ID}.svc.id.goog \
    --location=${ZONE}

  2. So erstellen Sie einen TPU-Slice-Knotenpool mit mehreren Hosts und aktiviertem DRANET-Treiber:

    gcloud container node-pools create v6e-16 \
    --location=${ZONE} \
    --cluster=${CLUSTER_NAME} \
    --machine-type=ct6e-standard-4t \
    --tpu-topology=4x4 \
    --num-nodes=4 \
    --enable-gvnic \
    --scopes=https://www.googleapis.com/auth/cloud-platform \
    --accelerator-network-profile=auto \
    --node-labels=cloud.google.com/gke-networking-dra-driver=true

Speicher und Authentifizierung konfigurieren

Erstellen Sie einen Cloud Storage-Bucket und initialisieren Sie eine Rapid Cache-Instanz, um das Laden von Modellen zu beschleunigen. Konfigurieren Sie dann die Authentifizierung für Hugging Face:

  1. Erstellen Sie in Ihrer TPU-Zone einen Speicher-Bucket und initialisieren Sie die Rapid Cache-Instanz:

    gcloud storage buckets create gs://${GS_BUCKET} --project=${PROJECT_ID} --default-storage-class=STANDARD --location=${REGION}

gcloud storage buckets anywhere-caches create gs://${GS_BUCKET} ${ZONE} \
    --ttl=1d \
    --admission-policy=ADMIT_ON_FIRST_MISS

  2. Konfigurieren Sie Identitätsverknüpfungen, um den Gewichtungs-Bucket sicher in Ihre GKE-Pods einzubinden. Erstellen Sie zuerst ein dediziertes IAM-Dienstkonto und erteilen Sie ihm Leseberechtigungen für den Bucket:

    gcloud iam service-accounts create ${GSA_NAME}

gcloud storage buckets add-iam-policy-binding gs://${GS_BUCKET} \
    --member="serviceAccount:${GSA_NAME}@${PROJECT_ID}.iam.gserviceaccount.com" \
    --role="roles/storage.objectAdmin"

  3. Erstellen Sie die Workload Identity-Föderation für GKE-Bindung und annotieren Sie das Kubernetes-ServiceAccount-Objekt:

    gcloud iam service-accounts add-iam-policy-binding ${GSA_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \
    --role="roles/iam.workloadIdentityUser" \
    --member="serviceAccount:${PROJECT_ID}.svc.id.goog[${NAMESPACE}/${KSA_NAME}]"

kubectl create serviceaccount ${KSA_NAME} --namespace ${NAMESPACE}
kubectl annotate serviceaccount ${KSA_NAME} --namespace ${NAMESPACE} iam.gke.io/gcp-service-account=${GSA_NAME}@${PROJECT_ID}.iam.gserviceaccount.com

  4. Wenn Sie die Modellgewichte für Gemma 4 herunterladen möchten, müssen Sie die Lizenzvereinbarung von Google auf Hugging Face bestätigen. Rufen Sie die Seite zum Gemma 4-Modell auf Hugging Face auf.

  5. Melden Sie sich an und akzeptieren Sie die Lizenzbedingungen, indem Sie auf Agree and access repository (Zustimmen und auf Repository zugreifen) klicken.

  6. Rufen Sie die Einstellungen Ihres Hugging Face-Kontos auf und generieren Sie ein Zugriffstoken mit der Rolle Read.

  7. Exportieren Sie Ihr Hugging Face-Token und erstellen Sie ein Kubernetes-Secret, damit Ray die Modellgewichte abrufen kann:

    export HF_TOKEN=YOUR_HUGGING_FACE_TOKEN

kubectl create secret generic hf-secret \
  --from-literal=hf_api_token=${HF_TOKEN}

Benutzerdefiniertes Container-Image erstellen

Damit die Umgebung mit mehreren Hosts alle erforderlichen Abhängigkeiten enthält, erstellen Sie ein benutzerdefiniertes Image basierend auf dem TPU-Image von vLLM und kopieren Sie Ihr Bereitstellungsskript hinein.

  1. Erstellen Sie ein Artifact Registry-Repository:

    gcloud artifacts repositories create ${REPO_NAME} \
    --repository-format=docker \
    --location=${REGION}

  2. Authentifizieren Sie Docker für Ihr Projekt:

    gcloud auth configure-docker ${REGION}-docker.pkg.dev

  3. Prüfen Sie die Dockerfile im Beispiel-Repository:

    FROM vllm/vllm-tpu:v0.21.0

ENV VLLM_TARGET_DEVICE=tpu
ENV VLLM_XLA_CACHE_PATH=/data

USER root

RUN pip install --no-cache-dir -U \
    "https://s3-us-west-2.amazonaws.com/ray-wheels/master/75b85027a859439fae5634e49aa6443f6fbecfeb/ray-3.0.0.dev0-cp312-cp312-manylinux2014_x86_64.whl" && \
    pip install --no-cache-dir --no-deps "ray[llm]"

COPY serve_tpu_multihost.py /home/ray/serve_tpu_multihost.py

  4. Erstellen Sie das Image und übertragen Sie es per Push in die Artifact Registry:

    docker build -t ${CUSTOM_IMAGE_URI} .
docker push ${CUSTOM_IMAGE_URI}

Modellgewichte in Cloud Storage bereitstellen

Bevor Sie den RayCluster bereitstellen, sollten Sie die Leistung beim Laden von Modellen optimieren und für eine Hochverfügbarkeit in Ihrem verteilten TPU-Slice sorgen. Dazu können Sie die Modellgewichte mit einem eigenständigen Kubernetes-Job direkt in Ihrem Cloud Storage-Bucket vorab bereitstellen. Dieser entkoppelte Ansatz ermöglicht ein koordiniertes paralleles Streaming und beschleunigt die Cluster-Startzeiten.

  1. Das Manifest für den Downloader-Job ist im Repository verfügbar. Manifestkonfiguration prüfen:

    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: ${KSA_NAME}
      restartPolicy: OnFailure
      containers:
      - name: downloader
        image: python:3.10-slim
        command: ["/bin/sh", "-c"]
        args:
        - |
          pip install -U huggingface_hub filelock

          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="google/gemma-4-31B-it", 
              local_dir="/data/google/gemma-4-31B-it"
          )
          '
        env:
        - name: HF_TOKEN
          valueFrom:
            secretKeyRef:
              name: hf-secret
              key: hf_api_token
        volumeMounts:
        - name: gcs-fuse-csi-ephemeral
          mountPath: /data
      volumes:
      - name: gcs-fuse-csi-ephemeral
        csi:
          driver: gcsfuse.csi.storage.gke.io
          volumeAttributes:
            bucketName: ${GS_BUCKET}
            mountOptions: "implicit-dirs"

  2. Erstellen Sie den Downloader-Job, indem Sie die Datei im Repository anwenden:

    envsubst < ai-ml/gke-ray/rayserve/llm/tpu/components/model-downloader-job.yaml | kubectl apply -f -

  3. Beobachten Sie den Job, bis der Downloadstream einen Erfolg meldet:

    kubectl logs -f job/model-downloader

Inferenzskript erstellen

Im folgenden Python-Skript wird eine Ray Serve-Anwendung definiert, die auf dem LLMConfig-Wrapper auf hoher Ebene von Ray Serve basiert.

  1. Sehen Sie sich das serve_tpu_multihost.py-Skript im Beispiel-Repository an:

    import os
import ray
from ray import serve
from ray.serve.llm import LLMConfig, ModelLoadingConfig, LLMServingArgs, build_openai_app

# Read configurations from environment variables
MODEL_ID = os.environ.get("MODEL_ID", "google/gemma-4-31B-it")
MODEL_SOURCE = os.environ.get("MODEL_SOURCE", "/data/google/gemma-4-31B-it")

# TPU hardware options (i.e. TPU-V6E, TPU-V7X etc.)
ACCELERATOR_TYPE = os.environ.get("ACCELERATOR_TYPE", "TPU-V6E")
TPU_TOPOLOGY = os.environ.get("TPU_TOPOLOGY", "4x4")

# vLLM engine parameters
TENSOR_PARALLEL_SIZE = int(os.environ.get("TENSOR_PARALLEL_SIZE", "16"))
MAX_MODEL_LEN = int(os.environ.get("MAX_MODEL_LEN", "8192"))
MAX_NUM_BATCHED_TOKENS = int(os.environ.get("MAX_NUM_BATCHED_TOKENS", "4096"))

# Define the multi-host TPU LLM config
llm_config = LLMConfig(
    model_loading_config=dict(
        model_id=MODEL_ID,
        model_source=MODEL_SOURCE
    ),
    accelerator_type=ACCELERATOR_TYPE,
    accelerator_config={"kind": "tpu", "topology": TPU_TOPOLOGY},
    engine_kwargs={
        "tensor_parallel_size": TENSOR_PARALLEL_SIZE,
        "max_model_len": MAX_MODEL_LEN,
        "max_num_batched_tokens": MAX_NUM_BATCHED_TOKENS,
        "distributed_executor_backend": "ray",
    }
)

deployment = build_openai_app(
    LLMServingArgs(
        llm_configs=[llm_config]
    )
)

Ray LLM API verstehen

Das Skript nutzt die native ray.serve.llm-Bibliothek von Ray Serve, um die Komplexität der TPU-Orchestrierung mit mehreren Hosts zu abstrahieren. Durch die Einbindung der vLLM-Engine bietet Ray Serve LLM ein leistungsstarkes, skalierbares Framework, das speziell für stark verteilte Inferenz-Arbeitslasten in der Produktion entwickelt wurde.

Die Verwendung der Ray LLM API bietet mehrere wichtige Vorteile:

  • Bereitstellungen mit mehreren Knoten:Mit Ray Serve LLM können Nutzer umfangreiche Modelle bereitstellen, die sich über mehrere verteilte Hosts erstrecken (z. B. ein TPU-Slice mit mehreren Hosts). Dabei werden automatisches Placement, Koordination und Topologieverteilung nativ automatisiert.
  • vLLM-Kompatibilität:Ray Serve LLM bietet eine OpenAI-kompatible API, die mit dem Server von vLLM übereinstimmt. Sie können auch auf das erweiterte Feature-Set von vLLM zugreifen (z. B. strukturierte Ausgabe, multimodale Funktionen und Reasoning-Modelle), während Sie die Arbeitslast in Ihrem Kubernetes-Cluster skalieren.
  • Funktionen für die Produktion:Ray Serve LLM bietet Funktionen auf Unternehmensniveau wie integriertes Autoscaling, benutzerdefiniertes Anforderungsrouting für maximale Cache-Treffer und integrierte Integrationen für Messwerte und Beobachtbarkeit.

Im bereitgestellten Inferenzskript wird die Bereitstellung durch zwei Hauptkomponenten definiert:

  • LLMConfig:Dieses Objekt definiert die Bereitstellungskonfiguration. Sie gibt die Modellquelle, die Engine-Parameter für vLLM und die accelerator_config an. Wenn Sie {"kind": "tpu", "topology": "4x4"} festlegen, stellt Ray Serve LLM automatisch eine verteilte Placement Group bereit, die genau Ihrer physischen TPU v6e-Slice mit 16 Chips entspricht.
  • build_openai_app:Bei dieser API wird die konfigurierte vLLM-Engine automatisch in einen OpenAI-kompatiblen FastAPI-Server eingebunden. So erhalten Sie ohne benutzerdefinierten Servercode eine branchenübliche REST API (wie /v1/chat/completions).

RayService bereitstellen

Stellen Sie die Netzwerkkonfiguration für die dynamische Ressourcenzuweisung (Dynamic Resource Allocation, DRA) und das RayService-Bereitstellungsmanifest bereit:

  1. Fordern Sie alle verfügbaren NetDevice-Schnittstellen auf jedem Knoten an, indem Sie das im Repository bereitgestellte ResourceClaimTemplate bereitstellen:

    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

  2. Wenden Sie das Vorlagenmanifest auf Ihren Cluster an:

    kubectl apply -f ai-ml/gke-ray/rayserve/llm/tpu/networking/all-netdev-template.yaml

  3. Das RayService-Manifest ist im Repository verfügbar. Manifestkonfiguration prüfen:

    apiVersion: ray.io/v1
kind: RayService
metadata:
  name: vllm-tpu-multihost
  labels:
    ai.gke.io/model: "gemma-4-31B-it"
    ai.gke.io/inference-server: "vllm"
spec:
  serveConfigV2: |
    http_options:
      host: 0.0.0.0
      port: 8000
    applications:
      - name: llm
        import_path: ai-ml.gke-ray.rayserve.llm.tpu.serve_tpu_multihost:deployment
        runtime_env:
          working_dir: "https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/archive/main.zip"
          env_vars:
            # Use local disk to prevent multi-host GCSFuse race conditions
            VLLM_XLA_CACHE_PATH: "/tmp/vllm_xla_cache"
  rayClusterConfig:
    headGroupSpec:
      rayStartParams: {}
      template:
        metadata:
          annotations:
            gke-gcsfuse/volumes: "true"
            gke-gcsfuse/cpu-limit: "0"
            gke-gcsfuse/memory-limit: "0"
            gke-gcsfuse/ephemeral-storage-limit: "0"
        spec:
          serviceAccountName: $KSA_NAME
          containers:
          - name: ray-head
            image: $CUSTOM_IMAGE_URI
            imagePullPolicy: Always
            ports:
            - containerPort: 6379
              name: gcs
            - containerPort: 8265
              name: dashboard
            - containerPort: 10001
              name: client
            - containerPort: 8000
              name: serve
            resources:
              limits:
                cpu: "2"
                memory: 16Gi
              requests:
                cpu: "2"
                memory: 16Gi
            volumeMounts:
            - name: dshm
              mountPath: /dev/shm
            - name: gcs-fuse-csi-ephemeral
              mountPath: /data
          volumes:
          - name: dshm
            emptyDir:
              medium: Memory
          - name: gke-gcsfuse-cache
            emptyDir:
              medium: Memory
          - name: gcs-fuse-csi-ephemeral
            csi:
              driver: gcsfuse.csi.storage.gke.io
              volumeAttributes:
                bucketName: $GS_BUCKET
                mountOptions: "implicit-dirs"
    workerGroupSpecs:
    - groupName: tpu-group
      replicas: 1
      minReplicas: 1
      maxReplicas: 1
      numOfHosts: 4
      rayStartParams: {}
      template:
        metadata:
          annotations:
            gke-gcsfuse/volumes: "true"
            gke-gcsfuse/cpu-limit: "0"
            gke-gcsfuse/memory-limit: "0"
            gke-gcsfuse/ephemeral-storage-limit: "0"
        spec:
          serviceAccountName: $KSA_NAME
          containers:
            - name: ray-worker
              image: $CUSTOM_IMAGE_URI
              imagePullPolicy: Always
              resources:
                limits:
                  cpu: "20"
                  google.com/tpu: "4"
                  memory: 200Gi
                requests:
                  cpu: "20"
                  google.com/tpu: "4"
                  memory: 200Gi
                claims:
                - name: netdev
              env:
                - name: HF_HOME
                  value: "/data/huggingface"
                - name: HF_TOKEN
                  valueFrom:
                    secretKeyRef:
                      name: hf-secret
                      key: hf_api_token
                - name: JAX_PLATFORMS
                  value: "tpu,cpu"
                - name: NODE_IP
                  valueFrom:
                    fieldRef:
                      fieldPath: status.hostIP
                - name: VBAR_CONTROL_SERVICE_URL
                  value: $(NODE_IP):8353
                - name: TPU_MULTIHOST_BACKEND
                  value: "ray"
                - name: TPU_BACKEND_TYPE
                  value: "jax"
                - name: ENABLE_PJRT_COMPATIBILITY
                  value: "true"
              volumeMounts:
              - name: dshm
                mountPath: /dev/shm
              - name: gcs-fuse-csi-ephemeral
                mountPath: /data
          volumes:
          - name: dshm
            emptyDir:
              medium: Memory
          - name: gke-gcsfuse-cache
            emptyDir:
              medium: Memory
          - name: gcs-fuse-csi-ephemeral
            csi:
              driver: gcsfuse.csi.storage.gke.io
              volumeAttributes:
                bucketName: $GS_BUCKET
                mountOptions: "implicit-dirs"
          resourceClaims:
            - name: netdev
              resourceClaimTemplateName: all-netdev
          nodeSelector:
            cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
            cloud.google.com/gke-tpu-topology: 4x4

  4. Stellen Sie den Dienst mit dem Manifest bereit:

    Autopilot

    1. Wenn Sie den Dienst in einem Autopilot-Cluster bereitstellen möchten, müssen Sie zuerst das Manifest herunterladen und lokal bearbeiten, um die Opt-in-Anmerkung ComputeClass nodeSelector hinzuzufügen, die für die DRANET-Netzwerkfunktionen in Autopilot erforderlich ist:

      curl -O https://raw.githubusercontent.com/GoogleCloudPlatform/kubernetes-engine-samples/main/ai-ml/gke-ray/rayserve/llm/tpu/ray-service.tpu-v6e-multihost.yaml

    2. Fügen Sie das Label unter dem Feld nodeSelector hinzu, sodass es so aussieht:

      nodeSelector:
  cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
  cloud.google.com/gke-tpu-topology: 4x4
  cloud.google.com/compute-class: dranet-compute-class

    3. Stellen Sie den Dienst dann mit dem geänderten lokalen Manifest bereit:

      envsubst < ray-service.tpu-v6e-multihost.yaml | kubectl apply -f -

    Standard

    Wenn Sie den Dienst in einem Standardcluster bereitstellen möchten, stellen Sie das Manifest direkt aus dem Repository bereit:

    envsubst < ai-ml/gke-ray/rayserve/llm/tpu/ray-service.tpu-v6e-multihost.yaml | kubectl apply -f -

Überprüfung

  1. Warten Sie, bis der RayService verfügbar ist:

    kubectl wait --for=condition=Ready --timeout=1800s rayservice/vllm-tpu-multihost

  2. Wenn Sie prüfen möchten, ob das Modell erfolgreich geladen wurde, sehen Sie sich die Protokolle des Ray-Head-Pods an:

    kubectl logs -f -l ray.io/node-type=head -c ray-head

Modell bereitstellen

In diesem Abschnitt interagieren Sie mit dem Modell. Das Modell muss vollständig heruntergeladen sein, bevor Sie fortfahren.

Portweiterleitung einrichten

Richten Sie mit dem folgenden Befehl die Portweiterleitung zum Modell ein:

kubectl port-forward svc/vllm-tpu-multihost-head-svc 8000:8000 2>&1 >/dev/null &

Mithilfe von curl mit dem Modell interagieren

In diesem Abschnitt wird gezeigt, wie Sie einen einfachen Smoke Test machen, um das bereitgestellte Gemma 4-Modell zu prüfen.

Verwenden Sie in einer neuen Terminalsitzung curl, um mit Ihrem Modell zu chatten:

curl -X POST http://127.0.0.1:8000/v1/chat/completions \
    -H "Content-Type: application/json" \
    -d '{
        "model": "google/gemma-4-31B-it",
        "messages": [
            {
              "role": "user",
              "content": "Why is GKE managed DRANET preferred for multi-host TPU networking?"
            }
        ],
        "max_tokens": 256
    }'

Die Ausgabe sieht dann ungefähr so aus:

{
  "id": "chatcmpl-392692d3-5325-4832-a3a3-0b084c1045b0",
  "object": "chat.completion",
  "created": 1779883255,
  "model": "google/gemma-4-31B-it",
  "choices": [
    {
      "index": 0,
      "message": {
        "role": "assistant",
        "content": "To understand why GKE-managed **DRANET** (Distributed RANET) is preferred for multi-host TPU networking, it is first necessary to understand the fundamental challenge of TPU pods: **the need for massive, low-latency, all-to-all communication.**\n\nWhen you scale a model across multiple TPU hosts (multi-host), the hosts must synchronize gradients and weights constantly. Standard TCP/IP networking introduces too much overhead (latency and CPU jitter) for these operations.\n\nHere is the detailed breakdown of why GKE-managed DRANET is the preferred architecture:\n\n### 1. Bypassing the Kernel (Zero-Copy Networking)\nStandard networking requires the operating system kernel to handle packets, moving data from the network card to kernel space and then to user space.\n*   **The DRANET Advantage:** DRANET implements a specialized networking stack that allows for **Kernel Bypass**. It enables the TPU hardware/drivers to write data directly into the memory of the destination host. This reduces latency and eliminates the CPU overhead associated with processing network interrupts.\n\n### 2. High-Bandwidth, Low-Latency Interconnect\nMulti-host TPU training relies on a specialized topology (like a 2D or 3D"
      },
      "finish_reason": "length"
    }
  ]
}

Optional: Über eine Gradio-Chat-Oberfläche mit dem Modell interagieren

In diesem Abschnitt erstellen Sie eine Webchat-Anwendung, mit der Sie mit Ihrem abgestimmten Modell für Anweisungen interagieren können.

Gradio ist eine Python-Bibliothek mit einem ChatInterface-Wrapper, der Benutzeroberflächen für Chatbots erstellt.

Chatoberfläche bereitstellen

Das Manifest für die Chatoberfläche ist im Repository verfügbar. Manifestkonfiguration prüfen:

apiVersion: apps/v1
kind: Deployment
metadata:
  name: gradio
  labels:
    app: gradio
spec:
  replicas: 1
  selector:
    matchLabels:
      app: gradio
  template:
    metadata:
      labels:
        app: gradio
    spec:
      containers:
      - name: gradio
        image: us-docker.pkg.dev/google-samples/containers/gke/gradio-app:v1.0.7
        resources:
          requests:
            cpu: "250m"
            memory: "512Mi"
          limits:
            cpu: "500m"
            memory: "512Mi"
        env:
        - name: CONTEXT_PATH
          value: "/v1/chat/completions"
        - name: HOST
          value: "http://vllm-tpu-multihost-serve-svc:8000"
        - name: LLM_ENGINE
          value: "openai-chat"
        - name: MODEL_ID
          value: "google/gemma-4-31B-it"
        - name: DISABLE_SYSTEM_MESSAGE
          value: "true"
        ports:
        - containerPort: 7860
---
apiVersion: v1
kind: Service
metadata:
  name: gradio
spec:
  selector:
    app: gradio
  ports:
  - protocol: TCP
    port: 8080
    targetPort: 7860
  type: ClusterIP

Wenden Sie das Manifest an:

kubectl apply -f ai-ml/gke-ray/rayserve/llm/tpu/components/gradio.yaml

Warten Sie, bis die Bereitstellung verfügbar ist:

kubectl wait --for=condition=Available --timeout=900s deployment/gradio

Chatoberfläche verwenden

Führen Sie in Cloud Shell den folgenden Befehl aus:

kubectl port-forward service/gradio 8080:8080

Dadurch wird eine Portweiterleitung von Cloud Shell zum Gradio-Dienst erstellt.

Klicken Sie oben rechts in der Cloud Shell-Taskleiste auf das Symbol Webvorschau Schaltfläche für Webvorschau. Klicken Sie auf Vorschau auf Port 8080. Im Browser wird ein neuer Tab geöffnet.

Interagieren Sie über die Gradio-Chat-Oberfläche mit Gemma. Fügen Sie einen Prompt hinzu und klicken Sie auf Senden.

Modellleistung beobachten

Wenn Sie die Dashboards für die Messwerte zur Beobachtbarkeit eines Modells ansehen möchten, das in KubeRay ausgeführt wird, können Sie die speziellen Ray on GKE-Dashboards verwenden.

Eine detaillierte Anleitung zum Konfigurieren Ihres Clusters und zum Aufrufen der Observability-Dashboards finden Sie unter Logs und Messwerte für Ray-Cluster in Google Kubernetes Engine (GKE) erfassen und ansehen.

Auf das Ray-Dashboard zugreifen

Um den Status Ihrer Ray-Akteure zu prüfen, detaillierte Anwendungslogs aufzurufen und die Nutzung auf Knotenebene nativ in Ray zu überwachen, können Sie auf das Ray-Dashboard zugreifen.

  1. Leiten Sie den Ray-Head-Knotendienst an Ihren lokalen Computer weiter:

    kubectl port-forward svc/vllm-tpu-multihost-head-svc 8265:8265

  2. Öffnen Sie Ihren Browser und rufen Sie http://localhost:8265 auf. Wenn Sie Cloud Shell verwenden, klicken Sie auf den Button Webvorschau und wählen Sie Vorschau auf Port 8265 aus.

  3. Wenn Sie Ihre vLLM-Deployments, den Systemzustand von Modellreplikaten und die Abfragelatenzen aufrufen möchten, klicken Sie auf den Tab Bereitstellen.

Bereinigen

So vermeiden Sie, dass Ihrem Google Cloud Konto die in dieser Anleitung verwendeten Ressourcen in Rechnung gestellt werden:

  1. Löschen Sie den RayService:

    kubectl delete rayservice vllm-tpu-multihost

  2. Löschen Sie den GKE-Cluster:

    gcloud container clusters delete ${CLUSTER_NAME} --zone=${ZONE}

Nächste Schritte