Diffuser des modèles ouverts Gemma à l'aide de TPU sur GKE avec Ray LLM

Ce tutoriel vous explique comment déployer un service d'inférence TPU multi-hôtes à l'aide de Ray Serve LLM. En tirant parti de la compatibilité native de Ray avec les TPU pour co-planifier de manière atomique les workers du moteur distribué sur des topologies d'accélérateur complexes, vous pouvez déployer de grands modèles sur une tranche de TPU multi-hôtes pour l'inférence.

Ce tutoriel est destiné aux ingénieurs en machine learning (ML), aux administrateurs et opérateurs de plate-forme, ainsi qu'aux spécialistes des données et de l'IA qui souhaitent utiliser les fonctionnalités d'orchestration de conteneurs Kubernetes pour diffuser des charges de travail d'IA/ML sur des tranches TPU distribuées et multihôtes. Pour en savoir plus sur les rôles courants et les exemples de tâches que nous citons dans le contenu Google Cloud , consultez Rôles utilisateur et tâches courantes de GKE.

Avant de lire cette page, assurez-vous de connaître les éléments suivants :

Arrière-plan

Cette section décrit les principales technologies utilisées dans ce guide.

TPU

Les Tensor Processing Units (TPU) vous permettent d'accélérer des charges de travail spécifiques exécutées sur vos nœuds, telles que le machine learning et le traitement des données. L'avantage principal des TPU est leur capacité à fournir des performances à grande échelle. Ce tutoriel utilise TPU Trillium, la sixième génération de Cloud TPU. Les tranches de TPU multi-hôtes se composent de plusieurs nœuds physiques communiquant à l'aide d'une interconnexion inter-chip (ICI) à haut débit, qui fonctionne bien pour le service à haut débit et à faible latence.

vLLM sur Ray

vLLM est un moteur de diffusion de LLM à haut débit et à faible consommation de mémoire. En s'intégrant à Ray Serve, vLLM peut évoluer sur plusieurs hôtes et accéder aux topologies matérielles physiques de manière native. Ce tutoriel montre comment utiliser les déploiements LLMConfig et LLMServer de Ray Serve pour orchestrer l'inférence vLLM sur des tranches multi-hôtes, en laissant le framework gérer automatiquement la distribution de la topologie et la répartition des groupes d'emplacements.

Objectifs

Ce tutoriel fournit les bases pour comprendre et explorer le déploiement pratique de LLM pour l'inférence dans un environnement Kubernetes géré qui utilise des TPU multi-hôtes.

  1. Préparez votre environnement avec un cluster GKE en mode Autopilot ou Standard.
  2. Créez une image de conteneur personnalisée avec des dépendances intégrées.
  3. Déployez un script Python Ray LLM sur votre cluster pour orchestrer l'inférence vLLM sur une tranche de TPU.
  4. Utilisez Ray LLM pour diffuser le modèle Gemma 4 via curl et une interface de chat Web facultative.

Avant de commencer

Avant de commencer, effectuez les tâches suivantes :

  • Activez l'API Google Kubernetes Engine.
    • Activer l'API Google Kubernetes Engine
  • Si vous souhaitez utiliser la Google Cloud CLI pour cette tâche, installez et initialisez la gcloud CLI. Si vous avez déjà installé la gcloud CLI, obtenez la dernière version en exécutant la commande gcloud components update. Il est possible que les versions antérieures de la gcloud CLI ne permettent pas d'exécuter les commandes de ce document.
  • Assurez-vous que votre projet dispose d'un quota suffisant pour la capacité TPU Trillium (v6e) dans la région sélectionnée. Pour en savoir plus, consultez Quotas Cloud TPU.
  • Assurez-vous que votre cluster GKE utilise GKE Dataplane V2 et qu'il répond aux exigences de version pour DRANET : 1.35.2-gke.1842000 ou version ultérieure pour les clusters Standard et Autopilot.
  • Assurez-vous de disposer des rôles IAM suivants :
    • roles/container.admin
    • roles/iam.serviceAccountAdmin

Préparer votre environnement

Dans ce tutoriel, vous utilisez Cloud Shell pour gérer les ressources hébergées sur Google Cloud. Cloud Shell est préinstallé avec les logiciels dont vous avez besoin pour ce tutoriel, y compris kubectl et l'CLI gcloud.

Pour configurer votre environnement avec Cloud Shell, procédez comme suit :

  1. Dans la console Google Cloud , lancez une session Cloud Shell en cliquant sur Activer Cloud Shell Bouton d'activation de Cloud Shell. Une session s'ouvre dans le volet inférieur de la console Google Cloud .

  2. Créez et activez un environnement virtuel Python :

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

  3. Installez la CLI Ray :

    pip install "ray"

  4. Définissez les variables d'environnement par défaut :

    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

    Remplacez les éléments suivants :

    • PROJECT_ID : ID de votre projet Google Cloud .
    • CLUSTER_NAME : nom du cluster
    • REGION : région dans laquelle votre capacité de TPU Trillium est disponible.
    • ZONE : zone dans laquelle votre capacité de TPU Trillium est disponible. Pour en savoir plus, consultez la section Disponibilité des TPU dans GKE.
    • REPOSITORY : nom de votre dépôt Artifact Registry.
    • BUCKET_NAME : nom de votre bucket de stockage.

Créer et configurer des ressources Google Cloud

Suivez les instructions ci-dessous pour créer les ressources requises.

Créer un cluster GKE et un pool de nœuds

Vous pouvez diffuser les modèles Gemma sur des TPU dans un cluster GKE Autopilot ou GKE Standard. GKE DRANET géré demande et gère de manière dynamique des ressources réseau hautes performances pour vos pods distribués, ce qui permet à GKE de provisionner automatiquement des réseaux secondaires à haut débit pour l'intercommunication des accélérateurs sans nécessiter de configuration manuelle du VPC.

Autopilot

  1. Dans Cloud Shell, créez le cluster Autopilot :

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

  2. Configurez kubectl de manière à communiquer avec votre cluster :

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

  3. Pour utiliser le DRANET géré par GKE en mode Autopilot, déployez la ressource ComputeClass personnalisée fournie dans le dépôt pour activer la mise en réseau dynamique :

    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. Appliquez le fichier manifeste à votre cluster :

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

Standard

  1. Dans Cloud Shell, créez un cluster standard qui active l'opérateur Ray et utilise 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. Créez un pool de nœuds de tranche TPU multi-hôtes avec le pilote DRANET activé :

    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

Configurer le stockage et l'authentification

Créez un bucket Cloud Storage et initialisez une instance Rapid Cache pour accélérer le chargement du modèle, puis configurez l'authentification pour Hugging Face :

  1. Dans votre zone TPU, créez un bucket de stockage et initialisez l'instance 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. Configurez des liens d'identité pour monter le bucket de poids de manière sécurisée dans vos pods GKE. Commencez par créer un compte de service IAM dédié et accordez-lui des autorisations de lecture du 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. Créez l'association Workload Identity Federation for GKE et annotez l'objet 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. Pour télécharger les pondérations du modèle Gemma 4, vous devez accepter le contrat de licence de Google sur Hugging Face. Accédez à la page du modèle Gemma 4 sur Hugging Face.

  5. Connectez-vous et acceptez les conditions de licence en cliquant sur Accepter et accéder au dépôt.

  6. Accédez aux paramètres de votre compte Hugging Face et générez un jeton d'accès avec le rôle Read.

  7. Exportez votre jeton Hugging Face et créez un secret Kubernetes pour que Ray puisse extraire les pondérations du modèle :

    export HF_TOKEN=YOUR_HUGGING_FACE_TOKEN

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

Créer l'image de conteneur personnalisée

Pour vous assurer que l'environnement multi-hôtes dispose de toutes les dépendances requises, créez une image personnalisée basée sur l'image TPU de vLLM et copiez-y votre script de diffusion.

  1. Créer un dépôt Artifact Registry :

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

  2. Authentifiez Docker dans votre projet :

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

  3. Inspectez les Dockerfile dans l'exemple de dépôt :

    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. Créez et transférez l'image vers Artifact Registry :

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

Préparer les pondérations du modèle dans Cloud Storage

Avant de déployer le RayCluster, optimisez les performances de chargement du modèle et assurez une haute disponibilité dans votre tranche de TPU distribuée en pré-organisant les pondérations du modèle directement dans votre bucket Cloud Storage à l'aide d'un job Kubernetes autonome. Cette approche découplée permet un streaming parallèle coordonné, ce qui accélère les temps de démarrage du cluster.

  1. Le fichier manifeste pour le job de téléchargement est disponible dans le dépôt. Examinez la configuration du fichier manifeste :

    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. Créez le job de téléchargement en appliquant le fichier dans le dépôt :

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

  3. Surveillez le job jusqu'à ce que le flux de téléchargement indique qu'il a réussi :

    kubectl logs -f job/model-downloader

Créer le script d'inférence

Le script Python suivant définit une application Ray Serve optimisée par le wrapper LLMConfig de haut niveau de Ray Serve.

  1. Inspectez le script serve_tpu_multihost.py dans l'exemple de dépôt :

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

Comprendre l'API Ray LLM

Le script utilise la bibliothèque ray.serve.llm native de Ray Serve pour masquer la complexité de l'orchestration TPU multi-hôte. En encapsulant le moteur vLLM, Ray Serve LLM fournit un framework évolutif et hautes performances spécialement conçu pour les charges de travail d'inférence hautement distribuées en production.

L'utilisation de l'API Ray LLM présente plusieurs avantages clés :

  • Déploiements multinœuds : Ray Serve LLM permet aux utilisateurs de diffuser des modèles massifs qui s'étendent sur plusieurs hôtes distribués (comme une tranche TPU multi-hôte) avec placement, coordination et distribution de la topologie automatiques en mode natif.
  • Compatibilité avec vLLM : Ray Serve LLM fournit une API compatible avec OpenAI qui s'aligne sur le serveur vLLM. Vous pouvez également accéder à l'ensemble de fonctionnalités avancées de vLLM (telles que les sorties structurées, les capacités multimodales et les modèles de raisonnement) tout en mettant à l'échelle la charge de travail sur votre cluster Kubernetes.
  • Fonctionnalités prêtes pour la production : Ray Serve LLM inclut des fonctionnalités de niveau entreprise telles que l'autoscaling intégré, le routage des requêtes personnalisées pour maximiser les accès au cache et les intégrations intégrées pour les métriques et l'observabilité.

Dans le script d'inférence fourni, le déploiement est défini par deux composants principaux :

  • LLMConfig : cet objet définit la configuration de diffusion. Il spécifie la source du modèle, les paramètres du moteur pour vLLM et accelerator_config. En définissant {"kind": "tpu", "topology": "4x4"}, Ray Serve LLM provisionne automatiquement un groupe d'emplacements distribué qui correspond exactement à votre tranche TPU v6e physique à 16 puces.
  • build_openai_app : cette API encapsule automatiquement le moteur vLLM configuré dans un serveur FastAPI compatible avec OpenAI, ce qui vous donne une API REST standard (comme /v1/chat/completions) prête à l'emploi sans avoir à écrire de code serveur personnalisé.

Déployer le RayService

Déployez la configuration réseau de l'allocation dynamique des ressources (DRA) et le fichier manifeste de diffusion RayService :

  1. Demandez toutes les interfaces NetDevice disponibles sur chaque nœud en déployant le ResourceClaimTemplate fourni dans le dépôt :

    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. Appliquez le fichier manifeste du modèle à votre cluster :

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

  3. Le fichier manifeste de diffusion RayService est disponible dans le dépôt. Vérifiez la configuration du fichier manifeste :

    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. Déployez le service à l'aide du fichier manifeste :

    Autopilot

    1. Pour déployer le service dans un cluster Autopilot, vous devez d'abord télécharger le fichier manifeste et le modifier localement pour ajouter l'option ComputeClass nodeSelector, qui est requise pour la mise en réseau DRANET sur 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. Ajoutez le libellé sous le champ nodeSelector de sorte que le code se présente comme suit :

      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. Déployez ensuite le service à l'aide du fichier manifeste local modifié :

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

    Standard

    Pour déployer le service dans un cluster standard, déployez le fichier manifeste directement depuis le dépôt :

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

Validation

  1. Attendez que RayService soit disponible :

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

  2. Pour vérifier que le modèle a bien été chargé, consultez les journaux du pod principal Ray :

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

Diffuser le modèle

Dans cette section, vous allez interagir avec le modèle. Assurez-vous que le modèle est entièrement téléchargé avant de continuer.

Configurer le transfert de port

Configurez le transfert de port vers le modèle en exécutant la commande suivante :

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

Interagir avec le modèle à l'aide de curl

Cette section explique comment effectuer un test de fumée de base pour vérifier le modèle Gemma 4 déployé.

Dans une nouvelle session de terminal, utilisez curl pour discuter avec votre modèle :

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

La sortie ressemble à ceci :

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

(Facultatif) Interagir avec le modèle via une interface de chat Gradio

Dans cette section, vous allez créer une application de chat Web qui vous permet d'interagir avec votre modèle adapté aux instructions.

Gradio est une bibliothèque Python dotée d'un wrapper ChatInterface qui crée des interfaces utilisateur pour les chatbots.

Déployer l'interface de chat

Le fichier manifeste de l'interface de chat est disponible dans le dépôt. Vérifiez la configuration du fichier manifeste :

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

Appliquez le fichier manifeste :

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

Attendez que le déploiement soit disponible :

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

Utiliser l'interface de chat

Dans Cloud Shell, exécutez la commande suivante :

kubectl port-forward service/gradio 8080:8080

Cela crée un transfert de port de Cloud Shell vers le service Gradio.

Cliquez sur l'icône Aperçu sur le Web Bouton &quot;Aperçu sur le Web&quot; qui se trouve en haut à droite de la barre des tâches Cloud Shell. Cliquez sur Preview on Port 8080 (Aperçu sur le port 8080). Un nouvel onglet s'ouvre dans le navigateur.

Interagissez avec Gemma via l'interface de chat Gradio. Ajoutez une requête et cliquez sur Envoyer.

Observer les performances du modèle

Pour afficher les tableaux de bord des métriques d'observabilité d'un modèle exécuté sur KubeRay, vous pouvez utiliser les tableaux de bord Ray on GKE dédiés.

Pour obtenir des instructions détaillées sur la configuration de votre cluster et l'accès aux tableaux de bord d'observabilité, consultez Collecter et afficher les journaux et les métriques pour les RayClusters sur Google Kubernetes Engine (GKE).

Accéder au tableau de bord Ray

Pour inspecter l'état de vos acteurs Ray, afficher les journaux d'application détaillés et surveiller l'utilisation au niveau des nœuds de manière native dans Ray, vous pouvez accéder au tableau de bord Ray.

  1. Transférez le service de nœud principal Ray vers votre ordinateur local :

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

  2. Ouvrez votre navigateur et accédez à http://localhost:8265. Si vous utilisez Cloud Shell, cliquez sur le bouton Aperçu sur le Web et sélectionnez Prévisualiser sur le port 8265.

  3. Pour afficher vos déploiements vLLM, l'état des répliques de modèle et les latences des requêtes, cliquez sur l'onglet Serve (Servir).

Effectuer un nettoyage

Pour éviter que les ressources utilisées dans ce tutoriel ne soient facturées sur votre compte Google Cloud , supprimez-les :

  1. Supprimez le RayService :

    kubectl delete rayservice vllm-tpu-multihost

  2. Supprimez le cluster GKE :

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

Étapes suivantes