通过 Ray 使用 GKE 中的多主机 TPU 应用 Gemma 开放模型

本教程将引导您使用 Ray Serve LLM 部署多主机 TPU 推理服务。通过利用 Ray 的原生 TPU 支持，在复杂的加速器拓扑中以原子方式协同调度分布式引擎工作器，您可以在多主机 TPU 切片上部署大型模型以进行推理。

本教程适用于机器学习 (ML) 工程师、平台管理员和运维人员，以及对使用 Kubernetes 容器编排功能在分布式多主机 TPU 切片上部署 AI/ML 工作负载感兴趣的数据和 AI 专家。如需详细了解我们在内容中 Google Cloud 中提及的常见角色和示例任务，请参阅常见的 GKE 用户角色和任务。

在阅读本页面之前，请确保您熟悉以下内容：

背景

本部分介绍本指南中使用的关键技术。

TPU

借助张量处理单元 (TPU)，您可以加速在节点上运行的特定工作负载（例如机器学习和数据处理）。TPU 的主要优势在于大规模性能。本教程使用 TPU Trillium，这是第六代 Cloud TPU。多主机 TPU 切片由多个物理节点组成，这些节点使用高速芯片间互连 (ICI) 进行通信，非常适合高吞吐量和低延迟服务。

Ray 上的 vLLM

vLLM 是一种高吞吐量、高内存效率的 LLM 服务引擎。通过与 Ray Serve 集成，vLLM 可以跨多个主机进行扩缩，并以原生方式访问物理硬件拓扑。本教程演示了如何使用 Ray Serve 的 LLMConfig 和 LLMServer 部署来编排跨多主机切片的 vLLM 推理，让框架自动处理拓扑分布和展示位置组扩展。

目标

本教程为理解和探索在托管式 Kubernetes 环境中使用多主机 TPU 部署实际 LLM 以进行推理提供了基础。

  1. 使用 Autopilot 或 Standard 中的 GKE 集群准备环境。
  2. 构建包含内置依赖项的自定义容器映像。
  3. 将 Ray LLM Python 脚本部署到集群，以编排 TPU 切片上的 vLLM 推理。
  4. 通过 curl 和可选的网页聊天界面，使用 Ray LLM 提供 Gemma 4 模型。

准备工作

在开始之前，请确保您已执行以下任务：

  • 启用 Google Kubernetes Engine API。
    • 启用 Google Kubernetes Engine API
  • 如需使用 Google Cloud CLI 执行此任务， 请安装并 初始化 gcloud CLI。如果您之前安装了 gcloud CLI，请通过运行 gcloud components update 命令来获取最新 版本。较早版本的 gcloud CLI 可能不支持运行本文档中的命令。
  • 确保您的项目在所选区域中有足够的 TPU Trillium (v6e) 容量配额。如需了解详情，请参阅 Cloud TPU 配额。
  • 确保您的 GKE 集群使用 GKE Dataplane V2，并且满足 DRANET 的版本要求：对于 Standard 和 Autopilot，均为 1.35.2-gke.1842000 或更高版本 。
  • 确保您拥有以下 IAM 角色：
    • roles/container.admin
    • roles/iam.serviceAccountAdmin

准备环境

在本教程中，您将使用 Cloud Shell 来管理上托管的资源 Google Cloud。Cloud Shell 中预安装了本教程所需的软件，包括 kubectl 和 gcloud CLI。

如需使用 Cloud Shell 设置您的环境，请按照以下步骤操作：

  1. 在 Google Cloud 控制台中，点击 激活 Cloud Shell 激活 Shell 按钮 以启动 Cloud Shell 会话。这会在 Google Cloud 控制台的底部窗格中启动会话。

  2. 创建并激活 Python 虚拟环境：

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

  3. 安装 Ray CLI：

    pip install "ray"

  4. 设置默认环境变量：

    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

    替换以下内容：

    • PROJECT_ID：您的 Google Cloud 项目 ID。
    • CLUSTER_NAME：您的集群的名称。
    • REGION：您的 TPU Trillium 容量所在的区域。
    • ZONE：您的 TPU Trillium 容量所在的可用区。如需了解详情，请参阅 GKE 中的 TPU 可用性。
    • REPOSITORY：您的 Artifact Registry 代码库的名称。
    • BUCKET_NAME：存储桶的名称。

创建和配置 Google Cloud 资源

请按照以下说明创建所需的资源。

创建 GKE 集群和节点池

您可以在 GKE Autopilot 或 Standard 集群中的 TPU 上应用 Gemma。GKE 托管 DRANET 会为分布式 Pod 动态请求和管理高性能网络资源，从而让 GKE 自动为加速器互通信预配辅助高速网络，而无需手动设置 VPC。

Autopilot

  1. 在 Cloud Shell 中，创建 Autopilot 集群：

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

  2. 配置 kubectl 以与您的集群通信：

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

  3. 如需在 Autopilot 模式下使用 GKE 托管 DRANET，请部署代码库中提供的自定义 ComputeClass 资源，以选择启用动态网络：

    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. 将清单应用到您的集群：

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

标准版

  1. 在 Cloud Shell 中，创建一个启用 Ray 运算符并使用 GKE Dataplane V2 的 Standard 集群：

    gcloud container clusters create ${CLUSTER_NAME} \
    --project=${PROJECT_ID} \
    --addons=RayOperator,GcsFuseCsiDriver \
    --machine-type=n2-standard-8 \
    --enable-dataplane-v2 \
    --workload-pool=${PROJECT_ID}. \
    --location=${ZONE}

  2. 创建启用了 DRANET 驱动程序的多主机 TPU 切片节点池：

    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

配置存储空间和身份验证

创建 Cloud Storage 存储桶并初始化 Rapid Cache 实例以加快模型加载速度，然后为 Hugging Face 配置身份验证：

  1. 在 TPU 可用区中，创建一个存储桶并初始化 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. 配置身份链接，以帮助将权重存储桶安全地装载到 GKE Pod 中。首先，创建一个专用 IAM 服务帐号并授予其存储桶读取权限：

    gcloud iam service-accounts create ${GSA_NAME}

gcloud storage buckets add-iam-policy-binding gs://${GS_BUCKET} \
    --member="serviceAccount:${GSA_NAME}@${PROJECT_ID}." \
    --role="roles/storage.objectAdmin"

  3. 创建 Workload Identity Federation for GKE 绑定并为 Kubernetes ServiceAccount 对象添加注解：

    gcloud iam service-accounts add-iam-policy-binding ${GSA_NAME}@${PROJECT_ID}. \
    --role="roles/iam.workloadIdentityUser" \
    --member="serviceAccount:${PROJECT_ID}.[${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}.

  4. 如需下载 Gemma 4 模型权重，您必须在 Hugging Face 上确认 Google 的许可协议。前往 Hugging Face 上的 Gemma 4 模型页面。

  5. 点击同意并访问代码库 ，登录并接受许可条款。

  6. 前往 Hugging Face 账号设置，并生成具有 访问令牌 角色的 Read。

  7. 导出 Hugging Face 令牌并创建一个 Kubernetes Secret，以便 Ray 可以拉取模型权重：

    export HF_TOKEN=YOUR_HUGGING_FACE_TOKEN

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

构建自定义容器映像

为确保多主机环境具有所有必需的依赖项，请基于 vLLM 的 TPU 映像构建自定义映像，并将服务脚本复制到其中。

  1. 创建 Artifact Registry 代码库：

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

  2. 向您的项目进行 Docker 身份验证：

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

  3. 检查示例代码库中的 Dockerfile：

    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. 构建映像并将其推送到 Artifact Registry：

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

将模型权重预先暂存到 Cloud Storage

在部署 RayCluster 之前，请使用独立的 Kubernetes 作业将模型权重直接预先暂存到 Cloud Storage 存储桶中，以优化模型加载性能并帮助确保分布式 TPU 切片的高可用性。这种分离式方法允许协调的并行流式传输，从而缩短集群启动时间。

  1. 下载器作业的清单位于代码库中。检查清单配置：

    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. 通过应用代码库中的文件来创建下载器作业：

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

  3. 监控作业，直到下载流报告成功：

    kubectl logs -f job/model-downloader

创建推理脚本

以下 Python 脚本定义了一个由 Ray Serve 的高级 LLMConfig 封装容器提供支持的 Ray Serve 应用。

  1. 检查示例代码库中的 serve_tpu_multihost.py 脚本：

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

了解 Ray LLM API

该脚本利用 Ray Serve 的原生 ray.serve.llm 库来抽象化多主机 TPU 编排的复杂性。通过封装 vLLM 引擎，Ray Serve LLM 提供了一个高性能、可伸缩的框架，专门为生产环境中的高度分布式推理工作负载而设计。

使用 Ray LLM API 具有多项关键优势：

  • 多节点部署： Ray Serve LLM 使用户能够以原生方式提供跨多个分布式主机（例如 TPU 多主机切片）的大型模型，并自动进行放置、协调和拓扑分布。
  • vLLM 兼容性： Ray Serve LLM 提供与 vLLM 服务器一致的 OpenAI 兼容 API。您还可以在 Kubernetes 集群中伸缩工作负载的同时，访问 vLLM 的高级功能集（例如结构化输出、多模态功能和推理模型）。
  • 可用于生产用途的功能： Ray Serve LLM 包含企业级功能，例如内置自动扩缩、用于最大限度提高缓存命中率的自定义请求路由，以及用于指标和可观测性的内置集成。

在提供的推理脚本中，部署由两个主要组件定义：

  • LLMConfig： 此对象定义了服务配置。它指定了模型来源、vLLM 的引擎参数和 accelerator_config。通过设置 {"kind": "tpu", "topology": "4x4"}，Ray Serve LLM 会自动预配一个分布式展示位置组，该组与您的物理 16 芯片 TPU v6e 切片完全匹配。
  • build_openai_app： 此 API 会自动将配置的 vLLM 引擎封装在 OpenAI 兼容的 FastAPI 服务器中，让您无需编写任何自定义服务器代码即可获得行业标准的 REST API（例如 /v1/chat/completions）。

部署 RayService

部署动态资源分配 (DRA) 网络配置和 RayService 服务清单：

  1. 通过部署代码库中提供的 ResourceClaimTemplate，请求每个节点上的所有可用 NetDevice 接口：

    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. 将模板清单应用于您的集群：

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

  3. RayService 服务清单位于代码库中。检查清单配置：

    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. 使用清单部署服务：

    Autopilot

    1. 如需在 Autopilot 集群中部署服务，您必须先下载清单并在本地对其进行修改，以添加选择启用 ComputeClass nodeSelector，这是在 Autopilot 上进行 DRANET 网络连接所必需的：

      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. 在 nodeSelector 字段下添加标签，使其如下所示：

      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. 然后，使用修改后的本地清单部署服务：

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

    标准版

    如需在 Standard 集群中部署服务，请直接从代码库部署清单：

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

验证

  1. 等待 RayService 可用：

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

  2. 如需确认模型已成功加载，请查看 Ray head Pod 中的日志：

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

应用模型

在本部分中，您将与模型互动。确保模型已完全下载，然后再继续。

设置端口转发

运行以下命令，设置到模型的端口转发：

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

使用 curl 与模型互动

本部分介绍如何执行基本的冒烟测试来验证所部署的 Gemma 4 模型。

在新的终端会话中，使用 curl 与模型聊天：

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

输出类似于以下内容：

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

（可选）通过 Gradio 聊天界面与模型互动

在本部分中，您将构建一个网页聊天应用，可让您与指令调优模型互动。

Gradio 是一个 Python 库，它具有一个可为聊天机器人创建界面的 ChatInterface 封装容器。

部署聊天界面

聊天界面的清单位于代码库中。检查清单配置：

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

应用清单：

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

等待部署成为可用状态：

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

使用聊天界面

在 Cloud Shell 中，运行以下命令：

kubectl port-forward service/gradio 8080:8080

这会创建从 Cloud Shell 到 Gradio 服务的端口转发。

点击 Cloud Shell 任务栏右上角的网页预览 图标 “网页预览”按钮。点击在端口 8080 上预览 。浏览器中会打开一个新的标签页。

使用 Gradio 聊天界面与 Gemma 互动。添加提示，然后点击提交 。

观察模型性能

如需查看在 KubeRay 上运行的模型的观测指标的信息中心，您可以使用专用的 Ray on GKE 信息中心。

如需详细了解如何配置集群和访问可观测性信息中心，请参阅收集和查看 Google Kubernetes Engine (GKE) 上的 RayCluster 的日志和指标。

访问 Ray 信息中心

如需检查 Ray 执行器的状态、查看详细的应用日志以及在 Ray 中以原生方式监控节点级利用率，您可以访问 Ray 信息中心。

  1. 将 Ray head 节点服务通过端口转发到您的本地机器：

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

  2. 打开浏览器并访问 http://localhost:8265。 如果您使用的是 Cloud Shell，请点击网页预览 按钮，然后选择在端口 8265 上预览 。

  3. 如需查看 vLLM 部署、模型副本运行状况和查询延迟时间，请点击 Serve 选项卡。

清理

为避免因本教程中使用的资源导致您的 Google Cloud 账号产生费用，请删除这些资源：

  1. 删除 RayService：

    kubectl delete rayservice vllm-tpu-multihost

  2. 删除 GKE 集群：

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

后续步骤