Eroga modelli aperti Gemma utilizzando le TPU su GKE con Ray LLM

Questo tutorial ti guida nel deployment di un servizio di inferenza TPU multi-host utilizzando Ray Serve LLM. Sfruttando il supporto TPU nativo di Ray per pianificare in modo atomico i worker del motore distribuito in topologie di acceleratori complesse, puoi eseguire il deployment di modelli di grandi dimensioni su una sezione TPU multi-host per l'inferenza.

Questo tutorial è destinato a ingegneri di machine learning (ML), amministratori e operatori di piattaforme e specialisti di dati e AI interessati a utilizzare le funzionalità di orchestrazione dei container Kubernetes per gestire i carichi di lavoro di AI/ML su sezioni TPU distribuite e multi-host. Per scoprire di più sui ruoli comuni e sulle attività di esempio a cui viene fatto riferimento nei contenuti di Google Cloud , consulta Ruoli utente e attività comuni di GKE.

Prima di leggere questa pagina, assicurati di avere familiarità con quanto segue:

Sfondo

Questa sezione descrive le tecnologie chiave utilizzate in questa guida.

TPU

Le Tensor Processing Unit (TPU) ti consentono di accelerare carichi di lavoro specifici in esecuzione sui tuoi nodi, come machine learning ed elaborazione di dati. Il vantaggio principale delle TPU è la scalabilità delle prestazioni. Questo tutorial utilizza TPU Trillium, la sesta generazione di Cloud TPU. Le slice TPU multi-host sono costituite da più nodi fisici che comunicano utilizzando un'interconnessione inter-chip (ICI) ad alta velocità, che funziona bene per il servizio a basso throughput e bassa latenza.

vLLM su Ray

vLLM è un motore di gestione degli LLM ad alta velocità e con un utilizzo efficiente della memoria. Grazie all'integrazione con Ray Serve, vLLM può scalare su più host e accedere in modo nativo alle topologie hardware fisiche. Questo tutorial mostra l'utilizzo dei deployment LLMConfig e LLMServer di Ray Serve per orchestrare l'inferenza vLLM su slice multi-host, consentendo al framework di gestire automaticamente la distribuzione della topologia e la diffusione del gruppo di posizionamento.

Obiettivi

Questo tutorial fornisce le basi per comprendere ed esplorare il deployment pratico di LLM per l'inferenza in un ambiente Kubernetes gestito che utilizza TPU multi-host.

  1. Prepara l'ambiente con un cluster GKE in modalità Autopilot o Standard.
  2. Crea un'immagine container personalizzata con dipendenze integrate.
  3. Esegui il deployment di uno script Python Ray LLM nel cluster per orchestrare l'inferenza vLLM su una sezione TPU.
  4. Utilizza Ray LLM per erogare il modello Gemma 4 tramite curl e un'interfaccia di chat web facoltativa.

Prima di iniziare

Prima di iniziare, assicurati di aver eseguito le seguenti operazioni:

  • Attiva l'API Google Kubernetes Engine.
    • Attiva l'API Google Kubernetes Engine
  • Se vuoi utilizzare Google Cloud CLI per questa attività, installala e poi inizializza gcloud CLI. Se hai già installato gcloud CLI, scarica l'ultima versione eseguendo il comando gcloud components update. Le versioni precedenti di gcloud CLI potrebbero non supportare l'esecuzione dei comandi in questo documento.
  • Assicurati che il tuo progetto disponga di una quota sufficiente per la capacità TPU Trillium (v6e) nella regione selezionata. Per maggiori informazioni, consulta la pagina Quote di Cloud TPU.
  • Assicurati che il cluster GKE utilizzi GKE Dataplane V2 e soddisfi i requisiti di versione per DRANET: 1.35.2-gke.1842000 o versioni successive sia per Standard che per Autopilot.
  • Assicurati di disporre dei seguenti ruoli IAM:
    • roles/container.admin
    • roles/iam.serviceAccountAdmin

prepara l'ambiente

In questo tutorial utilizzerai Cloud Shell per gestire le risorse ospitate su Google Cloud. Cloud Shell include il software necessario per questo tutorial preinstallato, tra cui kubectl e l'interfaccia a riga di comando gcloud.

Per configurare l'ambiente con Cloud Shell:

  1. Nella console Google Cloud , avvia una sessione di Cloud Shell facendo clic su Attiva Cloud Shell Pulsante Attiva Cloud Shell. Viene avviata una sessione nel riquadro inferiore della console Google Cloud .

  2. Crea e attiva un ambiente virtuale Python:

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

  3. Installa la CLI Ray:

    pip install "ray"

  4. Imposta le variabili di ambiente predefinite:

    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

    Sostituisci quanto segue:

    • PROJECT_ID: il tuo ID progetto Google Cloud .
    • CLUSTER_NAME: il nome del tuo cluster.
    • REGION: la regione in cui è disponibile la capacità TPU Trillium.
    • ZONE: la zona in cui è disponibile la capacità TPU Trillium. Per maggiori informazioni, consulta Disponibilità delle TPU in GKE.
    • REPOSITORY: il nome del repository Artifact Registry.
    • BUCKET_NAME: il nome del bucket di archiviazione.

Creare e configurare risorse Google Cloud

Segui queste istruzioni per creare le risorse richieste.

Crea un cluster GKE e un pool di nodi

Puoi pubblicare Gemma su TPU in un cluster GKE Autopilot o Standard. GKE managed DRANET richiede e gestisce dinamicamente risorse di rete ad alte prestazioni per i tuoi pod distribuiti, consentendo a GKE di eseguire automaticamente il provisioning di reti secondarie ad alta velocità per l'intercomunicazione degli acceleratori senza richiedere la configurazione manuale del VPC.

Autopilot

  1. In Cloud Shell, crea il cluster Autopilot:

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

  2. Configura kubectl per comunicare con il cluster:

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

  3. Per utilizzare DRANET gestito da GKE in modalità Autopilot, esegui il deployment della risorsa ComputeClass personalizzata fornita nel repository per attivare il networking dinamico:

    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. Applica il manifest al cluster:

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

Standard

  1. In Cloud Shell, crea un cluster Standard che attiva l'operatore Ray e utilizza GKE Dataplane V2:

    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. Crea un pool di nodi TPU multi-host con il driver DRANET abilitato:

    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

Configura archiviazione e autenticazione

Crea un bucket Cloud Storage e inizializza un'istanza di Rapid Cache per accelerare il caricamento del modello, quindi configura l'autenticazione per Hugging Face:

  1. Nella zona TPU, crea un bucket di archiviazione e inizializza l'istanza Rapid Cache:

    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. Configura i collegamenti delle identità per montare in modo sicuro il bucket del peso nei tuoi pod GKE. Innanzitutto, crea un account di servizio IAM dedicato e concedigli le autorizzazioni di lettura del 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. Crea il binding di Workload Identity Federation for GKE e annota l'oggetto Kubernetes ServiceAccount:

    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. Per scaricare i pesi del modello Gemma 4, devi accettare il contratto di licenza di Google su Hugging Face. Vai alla pagina del modello Gemma 4 su Hugging Face.

  5. Accedi e accetta i termini della licenza facendo clic su Accetta e accedi al repository.

  6. Vai alle impostazioni del tuo account Hugging Face e genera un token di accesso con il ruolo Read.

  7. Esporta il token Hugging Face e crea un secret Kubernetes in modo che Ray possa estrarre i pesi del modello:

    export HF_TOKEN=YOUR_HUGGING_FACE_TOKEN

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

Crea l'immagine container personalizzata

Per assicurarti che l'ambiente multi-host disponga di tutte le dipendenze richieste, crea un'immagine personalizzata basata sull'immagine TPU di vLLM e copia lo script di servizio.

  1. Crea un repository Artifact Registry:

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

  2. Autentica Docker nel tuo progetto:

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

  3. Esamina Dockerfile nel repository di esempio:

    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. Crea l'immagine ed eseguine il push in Artifact Registry:

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

Prepara i pesi del modello in Cloud Storage

Prima di eseguire il deployment di RayCluster, ottimizza le prestazioni di caricamento del modello e contribuisci a garantire l'alta affidabilità nella tua slice TPU distribuita pre-staging dei pesi del modello direttamente nel bucket Cloud Storage utilizzando un job Kubernetes autonomo. Questo approccio disaccoppiato consente lo streaming parallelo coordinato, accelerando i tempi di avvio del cluster.

  1. Il manifest per il job di download è disponibile nel repository. Rivedi la configurazione del manifest:

    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. Crea il job di download applicando il file nel repository:

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

  3. Monitora il job finché il flusso di download non segnala l'esito positivo:

    kubectl logs -f job/model-downloader

Crea lo script di inferenza

Lo script Python seguente definisce un'applicazione Ray Serve basata sul wrapper LLMConfig di alto livello di Ray Serve.

  1. Ispeziona lo script serve_tpu_multihost.py nel repository di esempio:

    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]
    )
)

Informazioni sull'API Ray LLM

Lo script utilizza la libreria ray.serve.llm nativa di Ray Serve per astrarre la complessità dell'orchestrazione delle TPU multi-host. Eseguendo il wrapping del motore vLLM, Ray Serve LLM fornisce un framework scalabile e ad alte prestazioni progettato specificamente per i carichi di lavoro di inferenza altamente distribuiti in produzione.

L'utilizzo dell'API Ray LLM offre diversi vantaggi principali:

  • Deployment multi-nodo:Ray Serve LLM consente agli utenti di erogare modelli di grandi dimensioni che si estendono su più host distribuiti (come una sezione multi-host TPU) con posizionamento automatico, coordinamento e distribuzione della topologia in modo nativo.
  • Compatibilità con vLLM:Ray Serve LLM fornisce un'API compatibile con OpenAI che si allinea al server di vLLM. Puoi anche accedere al set di funzionalità avanzate di vLLM (come output strutturato, funzionalità multimodali e modelli di ragionamento) durante lo scaling del workload nel cluster Kubernetes.
  • Funzionalità pronte per la produzione:Ray Serve LLM include funzionalità di livello aziendale come scalabilità automatica integrata, routing delle richieste personalizzato per massimizzare gli hit della cache e integrazioni integrate per metriche e osservabilità.

Nello script di inferenza fornito, il deployment è definito da due componenti principali:

  • LLMConfig: questo oggetto definisce la configurazione di pubblicazione. Specifica l'origine del modello, i parametri del motore per vLLM e accelerator_config. Se imposti {"kind": "tpu", "topology": "4x4"}, Ray Serve LLM esegue automaticamente il provisioning di un gruppo di posizionamento distribuito che corrisponde esattamente alla tua slice TPU v6e fisica a 16 chip.
  • build_openai_app: questa API esegue automaticamente il wrapping del motore vLLM configurato in un server FastAPI compatibile con OpenAI, offrendoti un'API REST standard del settore (come /v1/chat/completions) pronta all'uso senza scrivere codice server personalizzato.

Esegui il deployment di RayService

Esegui il deployment della configurazione di rete Dynamic Resource Allocation (DRA) e del manifest di pubblicazione RayService:

  1. Richiedi tutte le interfacce NetDevice disponibili su ogni nodo eseguendo il deployment di ResourceClaimTemplate fornito nel repository:

    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. Applica il manifest del modello al tuo cluster:

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

  3. Il manifest di servizio RayService è disponibile nel repository. Esamina la configurazione del manifest:

    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. Esegui il deployment del servizio utilizzando il manifest:

    Autopilot

    1. Per eseguire il deployment del servizio in un cluster Autopilot, devi prima scaricare il manifest e modificarlo localmente per aggiungere l'opt-in ComputeClass nodeSelector, necessario per il networking DRANET su Autopilot:

      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. Aggiungi l'etichetta sotto il campo nodeSelector in modo che risulti come segue:

      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. Poi, esegui il deployment del servizio utilizzando il manifest locale modificato:

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

    Standard

    Per eseguire il deployment del servizio in un cluster Standard, esegui il deployment del manifest direttamente dal repository:

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

Verifica

  1. Attendi che RayService sia disponibile:

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

  2. Per verificare che il modello sia stato caricato correttamente, visualizza i log del pod Ray head:

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

Eroga il modello

In questa sezione, interagisci con il modello. Assicurati che il modello sia stato scaricato completamente prima di procedere.

Configurare il port forwarding

Configura il port forwarding al modello eseguendo questo comando:

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

Interagire con il modello utilizzando curl

Questa sezione mostra come eseguire un test di fumo di base per verificare il modello Gemma 4 di cui è stato eseguito il deployment.

In una nuova sessione del terminale, utilizza curl per chattare con il tuo modello:

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
    }'

L'output è simile al seguente:

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

(Facoltativo) Interagisci con il modello tramite un'interfaccia di chat Gradio

In questa sezione, crei un'applicazione di chat web che ti consente di interagire con il modello ottimizzato per le istruzioni.

Gradio è una libreria Python che include un wrapper ChatInterface che crea interfacce utente per i chatbot.

Implementare l'interfaccia di chat

Il manifest per l'interfaccia di chat è disponibile nel repository. Esamina la configurazione del manifest:

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

Applica il manifest:

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

Attendi che il deployment sia disponibile:

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

Utilizzare l'interfaccia di chat

In Cloud Shell, esegui questo comando:

kubectl port-forward service/gradio 8080:8080

In questo modo viene creato un port forwarding da Cloud Shell al servizio Gradio.

Fai clic sull'icona Anteprima web Pulsante Anteprima web, che si trova in alto a destra nella barra delle app di Cloud Shell. Fai clic su Anteprima sulla porta 8080. Nel browser si apre una nuova scheda.

Interagisci con Gemma utilizzando l'interfaccia di chat di Gradio. Aggiungi un prompt e fai clic su Invia.

Osserva le prestazioni del modello

Per visualizzare le dashboard per le metriche di osservabilità di un modello in esecuzione su KubeRay, puoi utilizzare le dashboard Ray on GKE dedicate.

Per istruzioni dettagliate sulla configurazione del cluster e sull'accesso ai dashboard di osservabilità, consulta Raccogliere e visualizzare log e metriche per RayCluster su Google Kubernetes Engine (GKE).

Accedere alla dashboard Ray

Per esaminare lo stato degli attori Ray, visualizzare i log dettagliati delle applicazioni e monitorare l'utilizzo a livello di nodo in modo nativo in Ray, puoi accedere alla dashboard Ray.

  1. Esegui il port forwarding del servizio del nodo head di Ray alla tua macchina locale:

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

  2. Apri il browser e vai alla pagina http://localhost:8265. Se utilizzi Cloud Shell, fai clic sul pulsante Anteprima web e seleziona Anteprima sulla porta 8265.

  3. Per visualizzare i deployment di vLLM, l'integrità delle repliche del modello e le latenze delle query, fai clic sulla scheda Servi.

Esegui la pulizia

Per evitare che al tuo account Google Cloud vengano addebitati costi relativi alle risorse utilizzate in questo tutorial, elimina le risorse:

  1. Elimina RayService:

    kubectl delete rayservice vllm-tpu-multihost

  2. Elimina il cluster GKE:

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

Passaggi successivi