Offene Gemma-Modelle mit mehreren Hosts in GKE mit Ray bereitstellen

Autopilot Standard

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 auf einem TPU-Slice mit mehreren Hosts für die Inferenz bereitstellen.

Tipp: Für Produktionsbereitstellungen in GKE empfehlen wir die Verwendung der Kurzanleitung zur Inferenz, um maßgeschneiderte Best Practices und Konfigurationen für die Modellinferenz zu erhalten.

Diese Anleitung richtet sich an ML-Entwickler, Plattformadministratoren und -operatoren sowie an Daten- und KI-Spezialisten, die Kubernetes-Container-Orchestrierungsfunktionen zum Bereitstellen von KI-/ML-Arbeitslasten auf verteilten TPU-Slices mit mehreren Hosts verwenden möchten. Weitere Informationen zu gängigen Rollen und Beispielaufgaben, auf die in Google Cloud Inhalten verwiesen wird, finden Sie unter Häufig verwendete GKE-Nutzerrollen und -Aufgaben.

Bevor Sie diese Seite lesen, sollten Sie mit den folgenden Themen vertraut sein:

• Autopilot-Modus und Standardmodus
• TPUs in GKE
• TPU-Slices mit mehreren Hosts
• GKE managed DRANET

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 im großen Maßstab. 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 Hochgeschwindigkeits-Inter-Chip-Interconnect-Verbindung (ICI) kommunizieren. Dies eignet sich gut für die Bereitstellung mit hohem Durchsatz und niedriger Latenz.

vLLM in Ray

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

Ziele

Diese Anleitung bietet eine Grundlage für das Verständnis und die Erkundung der praktischen LLM-Bereitstellung für die Inferenz in einer verwalteten Kubernetes-Umgebung, die TPUs mit mehreren Hosts verwendet.

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-Skript in Ihrem Cluster bereit, um die vLLM-Inferenz über einen TPU-Slice zu orchestrieren.
4. Verwenden Sie Ray LLM, 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 diese Aufgabe verwenden möchten, installieren und dann initialisieren Sie die gcloud CLI. Wenn Sie die gcloud CLI bereits installiert haben, rufen Sie die neueste Version mit dem gcloud components update Befehl ab. Ältere Versionen der gcloud CLI unterstützen möglicherweise nicht die Ausführung der Befehle in diesem Dokument.
  Hinweis: Legen Sie für vorhandene Installationen der gcloud CLI das Attribut compute/region fest. Wenn Sie hauptsächlich zonale Cluster verwenden, legen Sie stattdessen compute/zone fest. Wenn Sie einen standardmäßigen Speicherort festlegen, können Sie in der gcloud CLI Fehler wie diesen vermeiden: One of [--zone, --region] must be supplied: Please specify location. Bei bestimmten Befehlen müssen Sie möglicherweise den Speicherort angeben, wenn er sich von dem von Ihnen festgelegten Standard unterscheidet.

• 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.
• Prüfen Sie, ob Ihr GKE-Cluster GKE Dataplane V2 verwendet und die Versionsanforderungen für DRANET erfüllt: 1.35.2-gke.1842000 oder höher für Standard- und Autopilot-Cluster.
• 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 gehostet werden Google Cloud. 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. Klicken Sie dazu in der Console auf Cloud Shell aktivieren . Dadurch wird eine Sitzung im unteren Bereich der Google Cloud Console 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 CLI:

   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: Ihre Google Cloud Projekt-ID.
   • 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.

Ressourcen erstellen und konfigurieren Google Cloud

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

GKE-Cluster und -Knotenpool erstellen

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

Autopilot

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

   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 GKE managed DRANET im Autopilot-Modus verwenden möchten, stellen Sie die benutzerdefinierte ComputeClass-Ressource bereit, die im Repository bereitgestellt wird, um sich für die dynamische Vernetzung zu registrieren:

   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 Standard-Cluster, der den Ray-Operator aktiviert und GKE Dataplane V2 verwendet:

   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. Erstellen Sie einen TPU-Slice-Knotenpool mit mehreren Hosts, bei dem der DRANET-Treiber aktiviert ist:

   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 Storage-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ätslinks, um den Bucket mit den Gewichten sicher in Ihre GKE-Pods einzuhängen. Erstellen Sie zuerst ein dediziertes IAM-Dienstkonto und gewähren 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 Federation for 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 Gemma 4-Modellgewichte herunterladen möchten, müssen Sie die Lizenzvereinbarung von Google auf Hugging Face akzeptieren. 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 das Repository zugreifen) klicken.

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

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, das auf dem TPU-Image von vLLM basiert, 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 bei Ihrem Projekt:

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

3. Prüfen Sie das 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, optimieren Sie die Leistung beim Laden von Modellen und sorgen Sie für hohe Verfügbarkeit auf Ihrem verteilten TPU-Slice, indem Sie die Modellgewichte mit einem eigenständigen Kubernetes-Job direkt in Ihrem Cloud Storage-Bucket bereitstellen. Dieser entkoppelte Ansatz ermöglicht ein koordiniertes paralleles Streaming und verkürzt die Cluster-Startzeiten.

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

   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. Behalten Sie den Job im Blick, bis der Downloadstream als erfolgreich gemeldet wird:

   kubectl logs -f job/model-downloader
   

Inferenzskript erstellen

Das folgende Python-Skript definiert eine Ray Serve-Anwendung, die vom LLMConfig-Wrapper auf hoher Ebene von Ray Serve unterstützt wird.

1. Prüfen Sie das Skript serve_tpu_multihost.py im Beispiel-Repository:

   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 das Wrapping der vLLM-Engine bietet Ray Serve LLM ein leistungsstarkes, skalierbares Framework, das speziell für hochgradig verteilte Inferenzarbeitslasten 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 massive Modelle bereitstellen, die sich über mehrere verteilte Hosts erstrecken (z. B. ein TPU-Slice mit mehreren Hosts), mit automatischem Placement, Koordination und Topologieverteilung.
• 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 über Ihren Kubernetes-Cluster hinweg skalieren.
• Produktionsreife Funktionen:Ray Serve LLM umfasst Funktionen der Enterprise-Klasse wie integriertes Autoscaling, benutzerdefiniertes Anfragenrouting 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. Es 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-Gruppe bereit, die genau Ihrem physischen 16-Chip-TPU-v6e-Slice entspricht.
• build_openai_app:Diese API wrappt die konfigurierte vLLM-Engine automatisch in einen OpenAI-kompatiblen FastAPI-Server. So erhalten Sie ohne benutzerdefinierten Servercode eine branchenübliche REST API (z. B. /v1/chat/completions).

RayService bereitstellen

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

1. Fordern Sie alle verfügbaren NetDevice-Schnittstellen auf jedem Knoten an, indem Sie die 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 Manifest der Vorlage auf Ihren Cluster an:

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

3. Das Bereitstellungsmanifest RayService ist im Repository verfügbar. Prüfen Sie die Manifestkonfiguration:

   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 den Opt-in-ComputeClass-nodeSelector hinzuzufügen, der für die DRANET-Vernetzung im Autopilot-Modus 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 Standard-Cluster 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

Hinweis: Bei TPU-Topologien mit mehreren Hosts führt nur der TPU-Koordinatorknoten die Bereitstellungsreplikat aus. Es ist zu erwarten, dass die :8000/healthz-Tests auf den anderen Worker-Pods im Slice fehlschlagen. Verwenden Sie die Bereitschaft einzelner Pods nicht, um den Zustand der Serve-Bereitstellung zu bestimmen. Warten Sie stattdessen, bis der Gesamtstatus von RayService ready ist, bevor Sie Prompt-Anfragen senden.

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

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

2. Prüfen Sie die Logs des Ray-Head-Pods, um zu bestätigen, dass das Modell erfolgreich geladen wurde:

   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 die Portweiterleitung zum Modell ein, indem Sie den folgenden Befehl ausführen:

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"
    }
  ]
}

Erfolg: Sie haben Gemma mithilfe von TPUs in GKE mit Ray LLM bereitgestellt. Sie können jetzt mit dem Modell interagieren.

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. Prüfen Sie die Manifestkonfiguration:

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 . 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 Beobachtbarkeitsmesswerte eines Modells aufrufen möchten, das in KubeRay ausgeführt wird, können Sie die entsprechenden Ray on GKE-Dashboards verwenden.

Eine detaillierte Anleitung zum Konfigurieren Ihres Clusters und zum Zugriff auf die Dashboards für die Beobachtbarkeit finden Sie unter Logs und Messwerte für RayCluster in Google Kubernetes Engine (GKE) erfassen und ansehen.

Auf das Ray-Dashboard zugreifen

Wenn Sie den Status Ihrer Ray-Akteure prüfen, detaillierte Anwendungsprotokolle aufrufen und die Nutzung auf Knotenebene nativ in Ray überwachen möchten, können Sie auf das Ray-Dashboard zugreifen.

1. Führen Sie eine Portweiterleitung des Ray-Head-Knotendienstes zu Ihrem lokalen Computer durch:

   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. Klicken Sie auf den Tab Serve (Bereitstellen), um Ihre vLLM-Bereitstellungen, den Zustand der Modellreplikate und die Abfragelatenzen aufzurufen.

Bereinigen

Wenn Sie vermeiden möchten, dass Ihrem Google Cloud Konto die in dieser Anleitung verwendeten Ressourcen in Rechnung gestellt werden, löschen Sie die Ressourcen:

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

• Mehr über Ray in Kubernetes erfahren.
• Informationen zum Bereitstellen von vLLM in GKE mit TPUs.
• Mehr über TPUs in GKE erfahren.