Diffuser des modèles ouverts Gemma à l'aide de TPU multi-hôtes sur GKE avec Ray

Ce tutoriel vous explique comment déployer un service d'inférence TPU multi-hôte à l'aide de Ray Serve LLM. En tirant parti de la prise en charge native des TPU par Ray pour co-planifier de manière atomique les nœuds de calcul du moteur distribué sur des topologies d'accélérateur complexes, vous pouvez déployer des modèles volumineux sur une tranche de TPU multi-hôte 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 de TPU multi-hôtes distribuées. Pour en savoir plus sur les rôles courants et les exemples de tâches que nous citons dans le Google Cloud contenu, consultez Rôles utilisateur et tâches courantes de GKE.

Avant de lire cette page, assurez-vous de bien 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 de données. Le principal avantage des TPU est leur performance à grande échelle. Ce tutoriel utilise TPU Trillium, la sixième génération de Cloud TPU. Les tranches de TPU multi-hôtes sont constituées de plusieurs nœuds physiques communiquant à l'aide d'une interconnexion entre puces (ICI) à haut débit, ce qui est idéal pour la diffusion à haut débit et à faible latence.

vLLM sur Ray

vLLM est un moteur de diffusion 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 explique 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 une base 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 Google Cloud CLI pour cette tâche, installez puis initialisez la gcloud CLI. Si vous avez déjà installé la gcloud CLI, obtenez la dernière version en exécutant la gcloud components update commande. 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 modes 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 la CLI gcloud.

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

  1. Dans la Google Cloud console, 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 Google Cloud console.

  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 Google Cloud projet.
    • 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 Google Cloud des ressources

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. Le DRANET géré par GKE demande et gère de manière dynamique les 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ôte 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 en toute sécurité le bucket de pondération dans vos pods GKE. Commencez par créer un compte de service IAM dédié et accordez-lui des autorisations de lecture de 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 la fédération d'identité de charge de travail pour la liaison 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 Agree and access repository (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 Read rôle.

  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ôte 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 le 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é-organiser 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é sur 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 une diffusion parallèle coordonnée, ce qui accélère les temps de démarrage du cluster.

  1. Le fichier manifeste du job de téléchargement est disponible dans le dépôt. Vérifiez 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 alimenté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 exploite la bibliothèque ray.serve.llm native de Ray Serve pour abstraire la complexité de l'orchestration des TPU multi-hôtes. 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 de TPU multi-hôte) avec un placement, une coordination et une distribution de la topologie automatiques de manière native.
  • Compatibilité vLLM : Ray Serve LLM fournit une API compatible avec OpenAI qui s'aligne sur le serveur de vLLM. Vous pouvez également accéder à l'ensemble de fonctionnalités avancées de vLLM (telles que la sortie structurée, les fonctionnalités multimodales et les modèles de raisonnement) tout en faisant évoluer 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é 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 le accelerator_config. En définissant {"kind": "tpu", "topology": "4x4"}, Ray Serve LLM provisionne automatiquement un groupe d'emplacements distribué qui correspond exactement à votre tranche de 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 écrire de code serveur personnalisé.

Déployer le RayService

Déployez la configuration réseau DRA (Dynamic Resource Allocation) 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 le nodeSelector ComputeClass d'activation, qui est requis 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 pour qu'il ressemble à ceci :

      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 à partir du dépôt :

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

Validation

  1. Attendez que le 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 sur 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 des 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 Preview on port 8265 (Aperçu sur le port 8265).

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

Libérer de l'espace

Pour éviter que les ressources utilisées dans ce tutoriel ne soient facturées sur votre Google Cloud compte, 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}

Étape suivante