透過 Ray 在 GKE 上使用多主機 TPU 提供 Gemma 開放原始碼模型

本教學課程將逐步說明如何使用 Ray Serve LLM 部署多主機 TPU 推論服務。您可以運用 Ray 的原生 TPU 支援功能,在複雜的加速器拓撲中,以原子方式共同排程分散式引擎工作人員,在多主機 TPU 配量中部署大型模型以進行推論。

本教學課程的適用對象為機器學習 (ML) 工程師、平台管理員和營運人員,以及有興趣使用 Kubernetes 容器自動化調度管理功能,在分散式多主機 TPU 節點上提供 AI/ML 工作負載的資料和 AI 專家。如要進一步瞭解 Google Cloud 內容中提及的常見角色和範例工作,請參閱「常見的 GKE 使用者角色和工作」。

閱讀本頁面之前,請先熟悉下列項目:

背景

本節說明本指南中使用的重要技術。

TPU

Tensor Processing Unit (TPU) 可加速處理節點上執行的特定工作負載,例如機器學習和資料處理。TPU 的主要優勢是可大規模提升效能。本教學課程使用第六代 Cloud TPU TPU Trillium。多主機 TPU 配量由多個實體節點組成,這些節點會使用高速晶片間互連 (ICI) 通訊,非常適合高輸送量和低延遲的服務。

Ray 上的 vLLM

vLLM 是高處理量、記憶體效率高的 LLM 服務引擎。透過與 Ray Serve 整合,vLLM 可跨多部主機擴充,並以原生方式存取實體硬體拓撲。本教學課程示範如何使用 Ray Serve 的 LLMConfigLLMServer 部署作業,在多主機切片中協調 vLLM 推論,讓架構自動處理拓撲分配和放置群組擴散。

目標

本教學課程提供基礎知識,協助您瞭解及探索如何使用多主機 TPU,在 Kubernetes 代管環境中實際部署 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 的版本需求:標準和 Autopilot 叢集皆須使用 1.35.2-gke.1842000 以上版本
  • 確認您具備下列 IAM 角色
    • roles/container.admin
    • roles/iam.serviceAccountAdmin

準備環境

在本教學課程中,您將使用 Cloud Shell 管理 Google Cloud上託管的資源。Cloud Shell 已預先安裝本教學課程所需的軟體,包括 kubectlgcloud CLI。

如要使用 Cloud Shell 設定環境,請按照下列步驟操作:

  1. 在 Google Cloud 控制台中,按一下「啟用 Cloud Shell」圖示 「Activate 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:儲存空間 bucket 的名稱。

建立及設定 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}.svc.id.goog \
        --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 bucket 並初始化 Rapid Cache 執行個體,加快模型載入速度,然後設定 Hugging Face 的驗證:

  1. 在 TPU 可用區中建立儲存空間 bucket,並初始化 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}.iam.gserviceaccount.com" \
        --role="roles/storage.objectAdmin"
    
  3. 建立 Workload Identity Federation for GKE 繫結,並為 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. 如要下載 Gemma 4 模型權重,您必須在 Hugging Face 上確認 Google 的授權協議。前往 Hugging Face 上的 Gemma 4 模型頁面

  5. 按一下「同意並存取存放區」,登入並接受授權條款。

  6. 前往 Hugging Face 帳戶設定,然後產生具有 Read 角色的「存取權杖」

  7. 匯出 Hugging Face 權杖並建立 Kubernetes 密鑰,讓 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 Job,直接在 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 應用程式,是由 Ray Serve 的高階 LLMConfig 包裝函式所驅動。

  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 編排作業的複雜性抽象化。Ray Serve LLM 封裝 vLLM 引擎,提供高效能、可擴充的架構,專為正式環境中高度分散的推論工作負載設計。

使用 Ray LLM API 有幾項主要優點:

  • 多節點部署:Ray Serve LLM 可讓使用者服務跨多個分散式主機 (例如 TPU 多主機切片) 的大型模型,並自動放置、協調及分配拓撲。
  • vLLM 相容性:Ray Serve LLM 提供與 OpenAI 相容的 API,與 vLLM 的伺服器一致。您也可以在 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 -
      

    標準

    如要在標準叢集中部署服務,請直接從存放區部署資訊清單:

    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 頭部節點服務轉送至本機電腦:

    kubectl port-forward svc/vllm-tpu-multihost-head-svc 8265:8265
    
  2. 開啟瀏覽器並前往 http://localhost:8265。 如果您使用 Cloud Shell,請按一下「網頁預覽」按鈕,然後選取「透過以下通訊埠預覽:8265」

  3. 如要查看 vLLM 部署作業、模型副本健康狀態和查詢延遲時間,請按一下「提供服務」分頁標籤。

清除所用資源

如要避免系統向您的 Google Cloud 帳戶收取本教學課程所用資源的費用,請刪除這些資源:

  1. 刪除 RayService:

    kubectl delete rayservice vllm-tpu-multihost
    
  2. 刪除 GKE 叢集:

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

後續步驟