Entrega modelos abiertos de Gemma con TPU de varios hosts en GKE con Ray

En este instructivo, se explica cómo implementar un servicio de inferencia de TPU de varios hosts con Ray Serve LLM. Si aprovechas la compatibilidad nativa de Ray con las TPU para programar de forma atómica los trabajadores del motor distribuido en topologías de aceleradores complejas, puedes implementar modelos grandes en una porción de TPU de varios hosts para la inferencia.

Este instructivo está dirigido a ingenieros de aprendizaje automático (AA), administradores y operadores de plataformas, y especialistas en datos y en IA que estén interesados en usar las capacidades de organización de contenedores de Kubernetes para entregar cargas de trabajo de IA/AA en segmentos de TPU distribuidos y con varios hosts. Para obtener más información sobre los roles comunes y las tareas de ejemplo a las que se hace referencia en el contenido de Google Cloud , consulta Roles y tareas comunes del usuario de GKE.

Antes de leer esta página, asegúrate de estar familiarizado con los siguientes temas:

Fondo

En esta sección, se describen las tecnologías clave que se usan en este instructivo.

TPU

Las unidades de procesamiento tensorial (TPU) te permiten acelerar cargas de trabajo específicas que se ejecutan en tus nodos, como el aprendizaje automático y el procesamiento de datos. La principal ventaja de las TPU es el rendimiento a gran escala. En este instructivo, se usa TPU Trillium, la sexta generación de Cloud TPU. Las porciones de TPU de varios hosts constan de varios nodos físicos que se comunican a través de una interconexión entre chips (ICI) de alta velocidad, que funciona bien para la publicación de alto rendimiento y baja latencia.

vLLM en Ray

vLLM es un motor de entrega de LLM de alta capacidad de procesamiento y eficiencia de memoria. Gracias a la integración con Ray Serve, vLLM puede escalar en varios hosts y acceder a topologías de hardware físicas de forma nativa. En este instructivo, se demuestra el uso de las implementaciones de LLMConfig y LLMServer de Ray Serve para coordinar la inferencia de vLLM en segmentos de varios hosts, lo que permite que el framework controle la distribución de la topología y la propagación del grupo de ubicación de forma automática.

Objetivos

En este instructivo, se proporcionan los conceptos básicos para comprender y explorar la implementación práctica de LLM para la inferencia en un entorno de Kubernetes administrado que usa TPU de varios hosts.

  1. Prepara tu entorno con un clúster de GKE en modo Autopilot o Standard.
  2. Compila una imagen de contenedor personalizada con dependencias integradas.
  3. Implementa una secuencia de comandos de Python de LLM de Ray en tu clúster para coordinar la inferencia de vLLM en una porción de TPU.
  4. Usa Ray LLM para entregar el modelo de Gemma 4 a través de curl y una interfaz de chat web opcional.

Antes de comenzar

Antes de comenzar, asegúrate de haber realizado las siguientes tareas:

  • Habilita la API de Google Kubernetes Engine.
    • Habilitar la API de Google Kubernetes Engine
  • Si deseas usar Google Cloud CLI para esta tarea, instala y, luego, inicializa gcloud CLI. Si ya instalaste gcloud CLI, ejecuta el comando gcloud components update para obtener la versión más reciente. Es posible que las versiones anteriores de gcloud CLI no admitan la ejecución de los comandos que se indican en este documento.
  • Asegúrate de que tu proyecto tenga suficiente cuota para la capacidad de TPU Trillium (v6e) en la región seleccionada. Para obtener más información, consulta Cuotas de Cloud TPU.
  • Asegúrate de que tu clúster de GKE use GKE Dataplane V2 y cumpla con los requisitos de versión para DRANET: 1.35.2-gke.1842000 o posterior para Standard y Autopilot.
  • Asegúrate de tener las siguientes funciones de IAM:
    • roles/container.admin
    • roles/iam.serviceAccountAdmin

Prepara el entorno

En este instructivo, usarás Cloud Shell para administrar recursos alojados en Google Cloud. Cloud Shell tiene preinstalado el software que necesitas para este instructivo, incluidos kubectl y la CLI de gcloud.

Para configurar tu entorno con Cloud Shell, sigue estos pasos:

  1. En la Google Cloud consola, haz clic en Activar Cloud Shell Botón de activar Shell para iniciar una sesión de Cloud Shell. Esto inicia una sesión en el panel inferior de la consola de Google Cloud .

  2. Crea y activa un entorno virtual de Python:

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

  3. Instala la CLI de Ray:

    pip install "ray"

  4. Configura las variables de entorno predeterminadas:

    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

    Reemplaza lo siguiente:

    • PROJECT_ID: Es el ID del proyecto de Google Cloud .
    • CLUSTER_NAME: El nombre de tu clúster.
    • REGION: Es la región en la que está disponible tu capacidad de TPU Trillium.
    • ZONE: Es la zona en la que está disponible tu capacidad de TPU Trillium. Para obtener más información, consulta la disponibilidad de TPU en GKE.
    • REPOSITORY: Es el nombre de tu repositorio de Artifact Registry.
    • BUCKET_NAME: Es el nombre de tu bucket de almacenamiento.

Crea y configura recursos de Google Cloud

Sigue estas instrucciones para crear los recursos necesarios.

Crea un clúster de GKE y un grupo de nodos

Puedes entregar Gemma en TPU en un clúster de GKE Autopilot o Standard. El DRANET administrado por GKE solicita y administra de forma dinámica recursos de redes de alto rendimiento para tus Pods distribuidos, lo que permite que GKE aprovisione automáticamente redes secundarias de alta velocidad para la intercomunicación del acelerador sin necesidad de configurar la VPC de forma manual.

Autopilot

  1. En Cloud Shell, crea el clúster de Autopilot:

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

  2. Configura kubectl para comunicarse con tu clúster:

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

  3. Para usar DRANET administrado por GKE en modo Autopilot, implementa el recurso ComputeClass personalizado que se proporciona en el repositorio para habilitar las redes dinámicas:

    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. Aplica el manifiesto al clúster:

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

Estándar

  1. En Cloud Shell, crea un clúster estándar que habilite el operador de Ray y use 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 grupo de nodos de porción de TPU de varios hosts con el controlador DRANET habilitado:

    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 el almacenamiento y la autenticación

Crea un bucket de Cloud Storage y, luego, inicializa una instancia de Rapid Cache para acelerar la carga del modelo. Luego, configura la autenticación para Hugging Face:

  1. En tu zona de TPU, crea un bucket de almacenamiento y, luego, inicializa la instancia de 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 vínculos de identidad para ayudar a montar de forma segura el bucket de peso en tus Pods de GKE. Primero, crea una cuenta de servicio de IAM dedicada y otórgale permisos de lectura 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 la vinculación de Workload Identity Federation for GKE y anota el objeto ServiceAccount de Kubernetes:

    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. Para descargar los pesos del modelo de Gemma 4, debes confirmar que aceptas el contrato de licencia de Google en Hugging Face. Ve a la página del modelo Gemma 4 en Hugging Face.

  5. Accede y acepta los términos de la licencia haciendo clic en Agree and access repository.

  6. Navega a la configuración de tu cuenta de Hugging Face y genera un token de acceso con el rol Read.

  7. Exporta tu token de Hugging Face y crea un Secret de Kubernetes para que Ray pueda extraer los pesos del modelo:

    export HF_TOKEN=YOUR_HUGGING_FACE_TOKEN

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

Compila la imagen del contenedor personalizado

Para garantizar que el entorno de varios hosts tenga todas las dependencias necesarias, crea una imagen personalizada basada en la imagen de TPU de vLLM y copia tu secuencia de comandos de servicio en ella.

  1. Crea un repositorio de Artifact Registry:

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

  2. Autentica Docker en tu proyecto:

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

  3. Inspecciona el Dockerfile en el repositorio de ejemplo:

    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. Compila y envía la imagen a Artifact Registry:

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

Transferencia previa de las ponderaciones del modelo a Cloud Storage

Antes de implementar el RayCluster, optimiza el rendimiento de la carga del modelo y ayuda a garantizar la alta disponibilidad en toda tu porción de TPU distribuida. Para ello, realiza una preetapa de los pesos del modelo directamente en tu bucket de Cloud Storage con un trabajo de Kubernetes independiente. Este enfoque desacoplado permite la transmisión paralela coordinada, lo que acelera los tiempos de inicio del clúster.

  1. El manifiesto del trabajo de descarga está disponible en el repositorio. Revisa la configuración del manifiesto:

    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. Aplica el archivo en el repositorio para crear el trabajo de descarga:

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

  3. Supervisa el trabajo hasta que la transmisión de descarga informe que se completó correctamente:

    kubectl logs -f job/model-downloader

Crea la secuencia de comandos de inferencia

La siguiente secuencia de comandos de Python define una aplicación de Ray Serve potenciada por el wrapper LLMConfig de alto nivel de Ray Serve.

  1. Inspecciona la secuencia de comandos serve_tpu_multihost.py en el repositorio de ejemplo:

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

Información sobre la API de LLM de Ray

La secuencia de comandos aprovecha la biblioteca ray.serve.llm nativa de Ray Serve para abstraer la complejidad de la organización de TPU de varios hosts. Al encapsular el motor de vLLM, Ray Serve LLM proporciona un framework escalable y de alto rendimiento diseñado específicamente para cargas de trabajo de inferencia altamente distribuidas en producción.

El uso de la API de LLM de Ray proporciona varios beneficios clave:

  • Implementaciones de varios nodos: Ray Serve LLM permite a los usuarios entregar modelos masivos que abarcan varios hosts distribuidos (como una porción de TPU de varios hosts) con ubicación, coordinación y distribución de topología automáticas de forma nativa.
  • Compatibilidad con vLLM: Ray Serve LLM proporciona una API compatible con OpenAI que se alinea con el servidor de vLLM. También puedes acceder al conjunto de funciones avanzadas de vLLM (como la salida estructurada, las capacidades multimodales y los modelos de razonamiento) mientras escalas la carga de trabajo en tu clúster de Kubernetes.
  • Funciones listas para la producción: Ray Serve LLM incluye capacidades de nivel empresarial, como el ajuste de escala automático integrado, el enrutamiento de solicitudes personalizado para maximizar los aciertos de caché y las integraciones integradas para las métricas y la observabilidad.

En la secuencia de comandos de inferencia proporcionada, la implementación se define mediante dos componentes principales:

  • LLMConfig: Este objeto define la configuración de la publicación. Especifica la fuente del modelo, los parámetros del motor para vLLM y el accelerator_config. Si configuras {"kind": "tpu", "topology": "4x4"}, Ray Serve LLM aprovisiona automáticamente un grupo de colocación distribuido que se asigna exactamente a tu segmento físico de TPU v6e de 16 chips.
  • build_openai_app: Esta API encapsula automáticamente el motor de vLLM configurado en un servidor FastAPI compatible con OpenAI, lo que te brinda una API de REST estándar de la industria (como /v1/chat/completions) lista para usar sin escribir ningún código de servidor personalizado.

Implementa el RayService

Implementa la configuración de red de la asignación dinámica de recursos (DRA) y el manifiesto de entrega RayService:

  1. Implementa el ResourceClaimTemplate proporcionado en el repositorio para solicitar todas las interfaces de NetDevice disponibles en cada nodo:

    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. Aplica el manifiesto de la plantilla a tu clúster:

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

  3. El manifiesto de publicación de RayService está disponible en el repositorio. Revisa la configuración del manifiesto:

    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. Implementa el servicio con el manifiesto:

    Autopilot

    1. Para implementar el servicio en un clúster de Autopilot, primero debes descargar el manifiesto y editarlo de forma local para agregar el parámetro de activación ComputeClass nodeSelector, que se requiere para la conexión en red de DRANET en 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. Agrega la etiqueta en el campo nodeSelector para que se vea de la siguiente manera:

      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. Luego, implementa el servicio con el manifiesto local modificado:

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

    Estándar

    Para implementar el servicio en un clúster Estándar, implementa el manifiesto directamente desde el repositorio:

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

Verificación

  1. Espera a que el RayService esté disponible:

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

  2. Para confirmar que el modelo se cargó correctamente, consulta los registros del Pod principal de Ray:

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

Entrega el modelo

En esta sección, interactuarás con el modelo. Asegúrate de que el modelo se haya descargado por completo antes de continuar.

Configura la redirección de puertos

Ejecuta el siguiente comando para configurar la redirección de puertos al modelo:

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

Interactúa con el modelo con curl

En esta sección, se muestra cómo puedes realizar una prueba de humo básica para verificar el modelo de Gemma 4 implementado.

En una sesión de terminal nueva, usa curl para chatear con tu modelo:

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

El resultado es similar al siguiente:

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

Interactúa con el modelo a través de una interfaz de chat de Gradio (opcional)

En esta sección, compilarás una aplicación de chat web que te permita interactuar con el modelo ajustado a instrucciones.

Gradio es una biblioteca de Python que tiene un wrapper ChatInterface que crea interfaces de usuario para chatbots.

Implementa la interfaz de chat

El manifiesto de la interfaz de chat está disponible en el repositorio. Revisa la configuración del manifiesto:

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

Aplica el manifiesto

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

Espera a que la implementación esté disponible:

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

Usa la interfaz de chat

En Cloud Shell, ejecuta el siguiente comando:

kubectl port-forward service/gradio 8080:8080

Esto crea una redirección de puertos desde Cloud Shell al servicio de Gradio.

Haz clic en el ícono de Vista previa en la Web Botón Vista previa en la Web, que se encuentra en la parte superior derecha de la barra de tareas de Cloud Shell. Haga clic en Vista previa en el puerto 8080. Se abrirá una pestaña nueva en el navegador.

Interactúa con Gemma con la interfaz de chat de Gradio. Agrega un mensaje y haz clic en Enviar.

Observa el rendimiento del modelo

Para ver los paneles de las métricas de observabilidad de un modelo que se ejecuta en KubeRay, puedes usar los paneles dedicados de Ray en GKE.

Para obtener instrucciones detalladas sobre cómo configurar tu clúster y acceder a los paneles de observabilidad, consulta Recopila y consulta registros y métricas de los clústeres de Ray en Google Kubernetes Engine (GKE).

Accede al panel de Ray

Para inspeccionar el estado de tus actores de Ray, ver registros detallados de la aplicación y supervisar la utilización a nivel del nodo de forma nativa en Ray, puedes acceder al panel de Ray.

  1. Redirecciona el servicio del nodo principal de Ray a tu máquina local:

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

  2. Abre el navegador y navega a http://localhost:8265. Si usas Cloud Shell, haz clic en el botón Vista previa en la Web y selecciona Vista previa en el puerto 8265.

  3. Para ver tus implementaciones de vLLM, el estado de las réplicas del modelo y las latencias de las consultas, haz clic en la pestaña Serve.

Realiza una limpieza

Para evitar que se generen cargos en tu cuenta de Google Cloud por los recursos que usaste en este instructivo, borra los recursos:

  1. Borra el RayService:

    kubectl delete rayservice vllm-tpu-multihost

  2. Borra el clúster de GKE:

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

¿Qué sigue?