本教學課程將逐步說明如何使用 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 的 LLMConfig 和 LLMServer 部署作業,在多主機切片中協調 vLLM 推論,讓架構自動處理拓撲分配和放置群組擴散。
目標
本教學課程提供基礎知識,協助您瞭解及探索如何使用多主機 TPU,在 Kubernetes 代管環境中實際部署 LLM,以進行推論。
- 在 Autopilot 或 Standard 模式中,使用 GKE 叢集準備環境。
- 建構內建依附元件的自訂容器映像檔。
- 將 Ray LLM Python 指令碼部署至叢集,以便在 TPU 節點上協調 vLLM 推論。
- 透過
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.adminroles/iam.serviceAccountAdmin
準備環境
在本教學課程中,您將使用 Cloud Shell 管理 Google Cloud上託管的資源。Cloud Shell 已預先安裝本教學課程所需的軟體,包括 kubectl 和 gcloud CLI。
如要使用 Cloud Shell 設定環境,請按照下列步驟操作:
在 Google Cloud 控制台中,按一下「啟用 Cloud Shell」圖示
,啟動 Cloud Shell 工作階段。系統會在 Google Cloud 控制台的底部窗格啟動工作階段。建立並啟用 Python 虛擬環境:
python3 -m venv ray-env source ray-env/bin/activate安裝 Ray CLI:
pip install "ray"設定預設環境變數:
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
在 Cloud Shell 中建立 Autopilot 叢集:
gcloud container clusters create-auto ${CLUSTER_NAME} \ --project=${PROJECT_ID} \ --enable-ray-operator \ --location=${REGION}設定
kubectl與叢集通訊:gcloud container clusters get-credentials ${CLUSTER_NAME} \ --location=${REGION}如要在 Autopilot 模式中使用 GKE 代管 DRANET,請部署存放區中提供的自訂 ComputeClass 資源,選擇啟用動態網路:
將資訊清單套用至叢集:
kubectl apply -f ai-ml/gke-ray/rayserve/llm/tpu/networking/dranet-compute-class.yaml
標準
在 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}建立啟用 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 的驗證:
在 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設定身分連結,協助將權重值安全地掛接到 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"建立 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如要下載 Gemma 4 模型權重,您必須在 Hugging Face 上確認 Google 的授權協議。前往 Hugging Face 上的 Gemma 4 模型頁面。
按一下「同意並存取存放區」,登入並接受授權條款。
匯出 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 映像檔建構自訂映像檔,並將服務指令碼複製到該映像檔中。
建立 Artifact Registry 存放區:
gcloud artifacts repositories create ${REPO_NAME} \ --repository-format=docker \ --location=${REGION}向專案驗證 Docker:
gcloud auth configure-docker ${REGION}-docker.pkg.dev檢查範例存放區中的
Dockerfile:建構映像檔並推送至 Artifact Registry:
docker build -t ${CUSTOM_IMAGE_URI} . docker push ${CUSTOM_IMAGE_URI}
將模型權重預先暫存至 Cloud Storage
部署 RayCluster 前,請先使用獨立的 Kubernetes Job,直接在 Cloud Storage 值區中預先暫存模型權重,以提升模型載入效能,並確保分散式 TPU 節點的高可用性。這種分離式方法可協調並行串流,縮短叢集啟動時間。
下載器工作的資訊清單位於存放區中。檢查資訊清單設定:
套用存放區中的檔案,建立下載器工作:
envsubst < ai-ml/gke-ray/rayserve/llm/tpu/components/model-downloader-job.yaml | kubectl apply -f -監控工作,直到下載串流回報成功為止:
kubectl logs -f job/model-downloader
建立推論指令碼
下列 Python 指令碼定義的 Ray Serve 應用程式,是由 Ray Serve 的高階 LLMConfig 包裝函式所驅動。
檢查範例存放區中的
serve_tpu_multihost.py指令碼:
瞭解 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 放送資訊清單:
在每個節點上部署存放區提供的
ResourceClaimTemplate,要求所有可用的 NetDevice 介面:將範本資訊清單套用至叢集:
kubectl apply -f ai-ml/gke-ray/rayserve/llm/tpu/networking/all-netdev-template.yaml存放區提供
RayService服務資訊清單。檢查資訊清單設定:使用資訊清單部署服務:
Autopilot
如要在 Autopilot 叢集中部署服務,您必須先下載資訊清單,並在本地編輯該資訊清單,加入選擇加入的
ComputeClassnodeSelector,這是 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在
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然後使用修改後的本機資訊清單部署服務:
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 -
驗證
等待 RayService 可用:
kubectl wait --for=condition=Ready --timeout=1800s rayservice/vllm-tpu-multihost如要確認模型是否已成功載入,請查看 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 包裝函式,可為聊天機器人建立使用者介面。
部署對話介面
對話介面的資訊清單位於存放區。檢查資訊清單設定:
套用資訊清單:
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 資訊主頁。
將 Ray 頭部節點服務轉送至本機電腦:
kubectl port-forward svc/vllm-tpu-multihost-head-svc 8265:8265開啟瀏覽器並前往
http://localhost:8265。 如果您使用 Cloud Shell,請按一下「網頁預覽」按鈕,然後選取「透過以下通訊埠預覽:8265」。如要查看 vLLM 部署作業、模型副本健康狀態和查詢延遲時間,請按一下「提供服務」分頁標籤。
清除所用資源
如要避免系統向您的 Google Cloud 帳戶收取本教學課程所用資源的費用,請刪除這些資源:
刪除 RayService:
kubectl delete rayservice vllm-tpu-multihost刪除 GKE 叢集:
gcloud container clusters delete ${CLUSTER_NAME} --zone=${ZONE}
後續步驟
- 瞭解 Kubernetes 上的 Ray。
- 瞭解如何在 GKE 上使用 TPU 提供 vLLM。
- 進一步瞭解 GKE 中的 TPU。