Disponibilizar modelos abertos do Gemma usando TPUs de vários hosts no GKE com o Ray

Autopilot Standard

Neste tutorial, você vai aprender a implantar um serviço de inferência de TPU de vários hosts usando o Ray Serve LLM. Ao aproveitar o suporte nativo da TPU do Ray para agendar de maneira atômica trabalhadores de mecanismos distribuídos em topologias de aceleradores complexas, é possível implantar modelos grandes em uma fração de TPU de vários hosts para inferência.

Dica: para implantações de produção no GKE, recomendamos usar o Início rápido de inferência para receber práticas recomendadas e configurações personalizadas para a inferência do seu modelo.

Este tutorial é destinado a engenheiros de machine learning (ML), administradores e operadores de plataforma e especialistas em dados e IA interessados em usar os recursos de orquestração de contêineres do Kubernetes para veicular cargas de trabalho de IA/ML em frações de TPU distribuídas e com vários hosts. Para saber mais sobre papéis comuns e exemplos de tarefas referenciados no conteúdo do Google Cloud , consulte Funções e tarefas comuns do usuário do GKE.

Antes de ler esta página, confira se você conhece:

• Modo Autopilot e modo Standard
• TPUs no GKE
• Fatias de TPU de vários hosts
• DRANET gerenciado pelo GKE

Contexto

Esta seção descreve as principais tecnologias usadas neste guia.

TPUs

As Unidades de processamento de tensor (TPUs) permitem acelerar cargas de trabalho específicas em execução nos nós, como machine learning e processamento de dados. A principal vantagem das TPUs é o desempenho em escala. Este tutorial usa a TPU Trillium, a sexta geração da Cloud TPU. As frações de TPU de vários hosts consistem em vários nós físicos que se comunicam usando uma interconexão entre chips (ICI) de alta velocidade, o que funciona bem para serviços de alta capacidade e baixa latência.

vLLM no Ray

O vLLM é um mecanismo de disponibilização de LLM de alta capacidade de processamento e eficiência de memória. Ao se integrar ao Ray Serve, o vLLM pode escalonar em vários hosts e acessar topologias de hardware físico de forma nativa. Neste tutorial, mostramos como usar as implantações LLMConfig e LLMServer do Ray Serve para orquestrar a inferência do vLLM em fatias de vários hosts, permitindo que o framework processe a distribuição de topologia e a propagação do grupo de posicionamento automaticamente.

Objetivos

Este tutorial fornece uma base para entender e explorar a implantação prática de LLMs para inferência em um ambiente gerenciado do Kubernetes que usa TPUs de vários hosts.

1. Prepare seu ambiente com um cluster do GKE no modo Autopilot ou Standard.
2. Crie uma imagem de contêiner personalizada com dependências integradas.
3. Implante um script Python do Ray LLM no cluster para orquestrar a inferência do vLLM em uma fração de TPU.
4. Use o Ray LLM para disponibilizar o modelo Gemma 4 usando curl e uma interface de chat na Web opcional.

Antes de começar

Antes de começar, verifique se você realizou as tarefas a seguir:

• Ative a API Google Kubernetes Engine.
Ativar a API Google Kubernetes Engine
• Se você quiser usar a Google Cloud CLI para essa tarefa, instale e, em seguida, inicialize a CLI gcloud. Se você instalou a CLI gcloud anteriormente, instale a versão mais recente executando o comando gcloud components update. Talvez as versões anteriores da CLI gcloud não sejam compatíveis com a execução dos comandos neste documento.
  Observação:em instalações prévias da CLI gcloud, defina a propriedade compute/region. Se você usa principalmente clusters zonais, defina compute/zone. Ao definir um local padrão, é possível evitar erros na CLI gcloud como: One of [--zone, --region] must be supplied: Please specify location. Talvez seja necessário especificar o local em determinados comandos se o local do cluster for diferente do padrão definido.

• Verifique se o projeto tem cota suficiente para a capacidade da TPU Trillium (v6e) na região selecionada. Para mais informações, consulte Cotas do Cloud TPU.
• Verifique se o cluster do GKE usa o GKE Dataplane V2 e atende aos requisitos de versão do DRANET: 1.35.2-gke.1842000 ou mais recente para Standard e Autopilot.
• Verifique se você tem os seguintes papéis do IAM:
  • roles/container.admin
  • roles/iam.serviceAccountAdmin

Preparar o ambiente

Neste tutorial, você vai usar o Cloud Shell para gerenciar recursos hospedados no Google Cloud. O Cloud Shell vem pré-instalado com o software necessário para este tutorial, incluindo kubectl e a CLI gcloud.

Para configurar o ambiente com o Cloud Shell, siga estas etapas:

1. No console do Google Cloud , inicie uma sessão do Cloud Shell clicando em Ativar o Cloud Shell . Isso inicia uma sessão no painel inferior do console Google Cloud .

2. Crie e ative um ambiente virtual Python:

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

3. Instale a CLI do Ray:

   pip install "ray"
   

4. Defina as variáveis de ambiente padrão:

   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
   

   Substitua:

   • PROJECT_ID: o ID do projeto Google Cloud .
   • CLUSTER_NAME: o nome do cluster.
   • REGION: a região em que sua capacidade do TPU Trillium está disponível.
   • ZONE: a zona em que sua capacidade de TPU Trillium está disponível. Para mais informações, consulte Disponibilidade da TPU no GKE.
   • REPOSITORY: o nome do repositório do Artifact Registry.
   • BUCKET_NAME: o nome do bucket de armazenamento.

Criar e configurar recursos Google Cloud

Siga estas instruções para criar os recursos necessários.

Criar um cluster do GKE e um pool de nós

É possível disponibilizar o Gemma em TPUs em um cluster do GKE Autopilot ou Standard. O DRANET gerenciado pelo GKE solicita e gerencia dinamicamente recursos de rede de alta performance para seus pods distribuídos, permitindo que o GKE provisione automaticamente redes secundárias de alta velocidade para a intercomunicação do acelerador sem exigir configuração manual da VPC.

Piloto automático

1. No Cloud Shell, crie o cluster do Autopilot:

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

2. Configure kubectl para se comunicar com o cluster:

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

3. Para usar o DRANET gerenciado pelo GKE no modo Autopilot, implante o recurso ComputeClass personalizado fornecido no repositório para ativar a rede dinâmica:

   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. Aplique o manifesto ao cluster:

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

Padrão

1. No Cloud Shell, crie um cluster padrão que ative o operador do Ray e use o 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. Crie um pool de nós de fração de TPU com vários hosts e o driver DRANET ativado:

   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
   

Configurar armazenamento e autenticação

Crie um bucket do Cloud Storage e inicialize uma instância do Rapid Cache para acelerar o carregamento do modelo. Em seguida, configure a autenticação para o Hugging Face:

1. Na sua zona de TPU, crie um bucket de armazenamento e inicialize a instância do 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. Configure links de identidade para ajudar a ativar com segurança o bucket de peso nos seus pods do GKE. Primeiro, crie uma conta de serviço dedicada do IAM e conceda a ela permissões de leitura do 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. Crie a vinculação da Federação de Identidade da Carga de Trabalho para GKE e anote o objeto ServiceAccount do 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 fazer o download dos pesos do modelo Gemma 4, você precisa confirmar o contrato de licença do Google no Hugging Face. Acesse a página do modelo Gemma 4 no Hugging Face.

5. Faça login e aceite os termos de licença clicando em Concordar e acessar o repositório.

6. Acesse as configurações da sua conta do Hugging Face e gere um token de acesso com a função Read.

7. Exporte seu token do Hugging Face e crie um secret do Kubernetes para que o Ray possa extrair os pesos do modelo:

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

Criar a imagem de contêiner personalizada

Para garantir que o ambiente de vários hosts tenha todas as dependências necessárias, crie uma imagem personalizada com base na imagem de TPU do vLLM e copie o script de disponibilização nela.

1. Crie um repositório do Artifact Registry:

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

2. Autentique o Docker no projeto:

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

3. Inspecione o Dockerfile no repositório de exemplo:

   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. Crie e envie a imagem para o Artifact Registry:

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

Pré-armazenar pesos de modelo no Cloud Storage

Antes de implantar o RayCluster, otimize o desempenho de carregamento do modelo e ajude a garantir a alta disponibilidade em toda a fração de TPU distribuída pré-transferindo os pesos do modelo diretamente no bucket do Cloud Storage usando um job independente do Kubernetes. Essa abordagem desacoplada permite o streaming paralelo coordenado, acelerando os tempos de inicialização do cluster.

1. O manifesto do job de download está disponível no repositório. Analise a configuração do manifesto:

   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. Crie o job de download aplicando o arquivo no repositório:

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

3. Monitore o job até que o stream de download informe sucesso:

   kubectl logs -f job/model-downloader
   

Criar o script de inferência

O script Python a seguir define um aplicativo do Ray Serve com tecnologia do wrapper LLMConfig de alto nível do Ray Serve.

1. Inspecione o script serve_tpu_multihost.py no repositório de exemplo:

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

Entender a API Ray LLM

O script usa a biblioteca ray.serve.llm nativa do Ray Serve para abstrair a complexidade da orquestração de TPUs de vários hosts. Ao encapsular o mecanismo vLLM, o Ray Serve LLM oferece um framework escalonável e de alto desempenho projetado especificamente para cargas de trabalho de inferência altamente distribuídas em produção.

O uso da API Ray LLM oferece vários benefícios importantes:

• Implantações de vários nós:o Ray Serve LLM permite que os usuários disponibilizem modelos enormes que abrangem vários hosts distribuídos (como uma fração de TPU de vários hosts) com posicionamento, coordenação e distribuição de topologia automáticos de forma nativa.
• Compatibilidade com vLLM:o Ray Serve LLM oferece uma API compatível com a OpenAI que se alinha ao servidor do vLLM. Você também pode acessar o conjunto de recursos avançados do vLLM (como saída estruturada, recursos multimodais e modelos de raciocínio) ao escalonar a carga de trabalho no cluster do Kubernetes.
• Recursos prontos para produção:o Ray Serve LLM inclui recursos de nível empresarial, como escalonamento automático integrado, roteamento de solicitações personalizadas para maximizar os acertos de cache e integrações integradas para métricas e capacidade de observação.

No script de inferência fornecido, a implantação é definida por dois componentes principais:

• LLMConfig:esse objeto define a configuração de exibição. Ele especifica a origem do modelo, os parâmetros do mecanismo para vLLM e o accelerator_config. Ao definir {"kind": "tpu", "topology": "4x4"}, o Ray Serve LLM provisiona automaticamente um grupo de posicionamento distribuído que corresponde exatamente à sua fração física de TPU v6e de 16 chips.
• build_openai_app:essa API encapsula automaticamente o mecanismo vLLM configurado em um servidor FastAPI compatível com a OpenAI, oferecendo uma API REST padrão do setor (como /v1/chat/completions) sem escrever nenhum código de servidor personalizado.

Implantar o RayService

Implante a configuração de rede de alocação dinâmica de recursos (DRA) e o manifesto de veiculação RayService:

1. Implante o ResourceClaimTemplate fornecido no repositório para solicitar todas as interfaces NetDevice disponíveis em cada nó:

   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. Aplique o manifesto de modelo ao cluster:

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

3. O manifesto de serviço RayService está disponível no repositório. Analise a configuração do manifesto:

   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. Implante o serviço usando o manifesto:

   Piloto automático

   1. Para implantar o serviço em um cluster do Autopilot, primeiro faça o download do manifesto e edite-o localmente para adicionar a opção de ativação ComputeClass nodeSelector, que é necessária para a rede DRANET no 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. Adicione o rótulo no campo nodeSelector para que ele fique assim:

      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. Em seguida, implante o serviço usando o manifesto local modificado:

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

   Padrão

   Para implantar o serviço em um cluster Standard, implante o manifesto diretamente do repositório:

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

Verificação

Observação: com topologias de TPU de vários hosts, apenas o nó coordenador de TPU executa a réplica de disponibilização. É um comportamento esperado que as sondagens :8000/healthz nos outros pods de worker na fração falhem. Evite usar a prontidão individual do pod para determinar a integridade da implantação do serviço. Em vez disso, aguarde o status geral RayService se tornar ready antes de enviar solicitações de comandos.

1. Aguarde até que o RayService esteja disponível:

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

2. Para confirmar que o modelo foi carregado, confira os registros do pod principal do Ray:

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

Disponibilizar o modelo

Nesta seção, você vai interagir com o modelo. Verifique se o download do modelo foi concluído antes de prosseguir.

Configurar o encaminhamento de portas

Configure o encaminhamento de portas para o modelo executando o seguinte comando:

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

Interagir com o modelo usando curl

Nesta seção, mostramos como realizar um teste preliminar básico para verificar o modelo do Gemma 4 implantado.

Em uma nova sessão do terminal, use curl para conversar com seu 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
    }'

A saída será assim:

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

Sucesso :você veiculou o Gemma usando TPUs no GKE com o Ray LLM. Agora é possível interagir com o modelo.

(Opcional) Interagir com o modelo usando uma interface de chat do GRadio

Nesta seção, você vai criar um aplicativo de chat na Web que permite interagir com seu modelo ajustado por instruções.

O Gradio é uma biblioteca Python que tem um wrapper ChatInterface que cria interfaces de usuário para chatbots.

Implantar a interface de chat

O manifesto da interface de chat está disponível no repositório. Analise a configuração do manifesto:

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

Aplique o manifesto:

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

Aguarde até que a implantação esteja disponível:

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

Usar a interface de chat

No Cloud Shell, execute este comando:

kubectl port-forward service/gradio 8080:8080

Isso cria um encaminhamento de porta do Cloud Shell para o serviço GRadio.

Clique no ícone Visualização da Web no canto superior direito da barra de tarefas do Cloud Shell. Clique em Visualizar na porta 8080. Uma nova guia será aberta no navegador.

Interaja com Gemma usando a interface de chat do GRadio. Adicione uma solicitação e clique em Enviar.

Observar a performance do modelo

Para conferir os painéis das métricas de observabilidade de um modelo em execução no KubeRay, use os painéis dedicados do Ray no GKE.

Para instruções detalhadas sobre como configurar o cluster e acessar os painéis de observabilidade, consulte Coletar e visualizar registros e métricas dos RayClusters no Google Kubernetes Engine (GKE).

Acessar o painel do Ray

Para inspecionar o status dos atores do Ray, ver registros detalhados de aplicativos e monitorar a utilização no nível do nó de forma nativa no Ray, acesse o painel do Ray.

1. Encaminhe o serviço do nó principal do Ray para sua máquina local:

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

2. Abra o navegador e acesse http://localhost:8265. Se você estiver usando o Cloud Shell, clique no botão Visualização da Web e selecione Visualizar na porta 8265.

3. Para conferir as implantações do vLLM, a integridade da réplica do modelo e as latências de consulta, clique na guia Veicular.

Limpar

Para evitar cobranças na sua conta do Google Cloud pelos recursos usados neste tutorial, exclua os recursos:

1. Exclua o RayService:

   kubectl delete rayservice vllm-tpu-multihost
   

2. Exclua o cluster do GKE:

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

A seguir

• Saiba mais sobre o Ray no Kubernetes.
• Saiba como disponibilizar o vLLM no GKE com TPUs.
• Saiba mais sobre TPUs no GKE.