在 Google Kubernetes Engine (GKE) 上，使用 TPU 部署搭載 Stable Diffusion 模型的 Ray Serve 應用程式

標準

---

本指南說明如何使用 TPU、Ray Serve 和 Ray Operator 外掛程式，在 Google Kubernetes Engine (GKE) 上部署及提供 Stable Diffusion 模型。

本指南適用於 Generative AI 客戶、GKE 新使用者或現有使用者、機器學習工程師、MLOps (DevOps) 工程師，或是有意使用 Kubernetes 容器協調功能，透過 Ray 服務模型的平台管理員。

關於 Ray 和 Ray Serve

Ray 是開放原始碼的可擴充運算架構，適用於 AI/ML 應用程式。Ray Serve 是 Ray 的模型服務程式庫，用於在分散式環境中擴充及提供模型。詳情請參閱 Ray 說明文件中的「Ray Serve」。

關於 TPU

Tensor Processing Unit (TPU) 是專門的硬體加速器，可大幅加快大規模機器學習模型的訓練和推論速度。搭配 TPU 使用 Ray，可順暢擴展高效能機器學習應用程式。如要進一步瞭解 TPU，請參閱 Cloud TPU 文件中的「Cloud TPU 簡介」。

關於 KubeRay TPU 初始化 Webhook

做為 Ray Operator 外掛程式的一部分，GKE 提供驗證和變動網路掛鉤，可處理 TPU Pod 排程，以及容器初始化作業所需的特定 TPU 環境變數 (例如 JAX 架構)。KubeRay TPU Webhook 會使用下列屬性，變更要求 TPU 的 Pod (附有 app.kubernetes.io/name: kuberay 標籤)：

• TPU_WORKER_ID：TPU 配量中每個工作站 Pod 的專屬整數。
• TPU_WORKER_HOSTNAMES：所有需要彼此通訊的 TPU 工作站 DNS 主機名稱清單。這項變數只會注入多主機群組中的 TPU Pod。
• replicaIndex：Pod 標籤，內含 Pod 所屬工作者群組副本的專屬 ID。這對多主機工作站群組很有用，因為多個工作站 Pod 可能屬於同一個副本，而 Ray 會使用這項功能啟用多主機自動調度資源。
• TPU_NAME：代表這個 Pod 所屬 GKE TPU PodSlice 的字串，設定的值與 replicaIndex 標籤相同。
• podAffinity：確保 GKE 會在相同節點集區中，排定具有相符 replicaIndex 標籤的 TPU Pod。這樣一來，GKE 就能以節點集區為單位，而非單一節點，以不可分割的方式擴大多主機 TPU。

目標

• 建立含有 TPU 節點集區的 GKE 叢集。
• 部署含有 TPU 的 Ray 叢集。
• 部署 RayService 自訂資源。
• 與 Stable Diffusion 模型伺服器互動。

費用

在本文件中，您會使用下列 Google Cloud的計費元件：

• GKE

如要根據預測用量估算費用，請使用 Pricing Calculator。

初次使用 Google Cloud 的使用者可能符合免費試用資格。

完成本文所述工作後，您可以刪除已建立的資源，避免繼續計費。詳情請參閱清除所用資源一節。

事前準備

Cloud Shell 已預先安裝本教學課程所需的軟體，包括 kubectl 和 gcloud CLI。如果您未使用 Cloud Shell，請安裝 gcloud CLI。

Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.

Install the Google Cloud CLI.

注意：如果您先前已安裝 gcloud CLI，請執行 gcloud components update，確保您使用的是最新版本。

如果您使用外部識別資訊提供者 (IdP)，請先 使用聯合身分登入 gcloud CLI。

如要初始化 gcloud CLI，請執行下列指令：

gcloud init

Create or select a Google Cloud project.

Roles required to select or create a project

• Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
• Create a project: To create a project, you need the Project Creator (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

• Create a Google Cloud project:

  gcloud projects create PROJECT_ID

  Replace PROJECT_ID with a name for the Google Cloud project you are creating.

• Select the Google Cloud project that you created:

  gcloud config set project PROJECT_ID

  Replace PROJECT_ID with your Google Cloud project name.

Verify that billing is enabled for your Google Cloud project.

Enable the GKE API:

Roles required to enable APIs

To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

gcloud services enable container.googleapis.com

Install the Google Cloud CLI.

注意：如果您先前已安裝 gcloud CLI，請執行 gcloud components update，確保您使用的是最新版本。

如果您使用外部識別資訊提供者 (IdP)，請先 使用聯合身分登入 gcloud CLI。

如要初始化 gcloud CLI，請執行下列指令：

gcloud init

Create or select a Google Cloud project.

Roles required to select or create a project

• Select a project: Selecting a project doesn't require a specific IAM role—you can select any project that you've been granted a role on.
• Create a project: To create a project, you need the Project Creator (roles/resourcemanager.projectCreator), which contains the resourcemanager.projects.create permission. Learn how to grant roles.

Note: If you don't plan to keep the resources that you create in this procedure, create a project instead of selecting an existing project. After you finish these steps, you can delete the project, removing all resources associated with the project.

• Create a Google Cloud project:

  gcloud projects create PROJECT_ID

  Replace PROJECT_ID with a name for the Google Cloud project you are creating.

• Select the Google Cloud project that you created:

  gcloud config set project PROJECT_ID

  Replace PROJECT_ID with your Google Cloud project name.

Verify that billing is enabled for your Google Cloud project.

Enable the GKE API:

Roles required to enable APIs

To enable APIs, you need the Service Usage Admin IAM role (roles/serviceusage.serviceUsageAdmin), which contains the serviceusage.services.enable permission. Learn how to grant roles.

gcloud services enable container.googleapis.com

1. Grant roles to your user account. Run the following command once for each of the following IAM roles: roles/container.clusterAdmin, roles/container.admin

   gcloud projects add-iam-policy-binding PROJECT_ID --member="user:USER_IDENTIFIER" --role=ROLE

   Replace the following:

   • PROJECT_ID: your project ID.
   • USER_IDENTIFIER: the identifier for your user account—for example, [e-mail].
   • ROLE: the IAM role that you grant to your user account.

確保配額充足

確認 Google Cloud 專案在 Compute Engine 區域或地區中有足夠的 TPU 配額。詳情請參閱 Cloud TPU 說明文件中的「確保有足夠的 TPU 和 GKE 配額」。您可能也需要提高下列項目的配額：

• Persistent Disk SSD (GB)
• 使用中的 IP 位址

準備環境

如要準備環境，請按照下列步驟操作：

1. 在 Google Cloud 控制台中，按一下Google Cloud 控制台中的「啟用 Cloud Shell」，即可啟動 Cloud Shell 工作階段。系統會在 Google Cloud 控制台的底部窗格啟動工作階段。

2. 設定環境變數：

   export PROJECT_ID=PROJECT_ID
   export CLUSTER_NAME=ray-cluster
   export COMPUTE_REGION=us-central2-b
   export CLUSTER_VERSION=CLUSTER_VERSION
   

   更改下列內容：

   • PROJECT_ID：您的 Google Cloud 專案 ID。
   • CLUSTER_VERSION：要使用的 GKE 版本。必須為 1.30.1 或之後。

3. 複製 GitHub 存放區：

   git clone https://github.com/GoogleCloudPlatform/kubernetes-engine-samples
   

4. 變更為工作目錄：

   cd kubernetes-engine-samples/ai-ml/gke-ray/rayserve/stable-diffusion
   

建立含有 TPU 節點集區的叢集

建立具有 TPU 節點集區的標準 GKE 叢集：

1. 建立啟用 Ray Operator 的標準模式叢集：

   gcloud container clusters create ${CLUSTER_NAME} \
       --addons=RayOperator \
       --machine-type=n1-standard-8 \
       --cluster-version=${CLUSTER_VERSION} \
       --location=${COMPUTE_REGION}
   

2. 建立單一主機 TPU 節點集區：

   gcloud container node-pools create tpu-pool \
       --location=${COMPUTE_REGION} \
       --cluster=${CLUSTER_NAME} \
       --machine-type=ct4p-hightpu-4t \
       --num-nodes=1
   

如要在標準模式下使用 TPU，請選取：

• 有 TPU 加速器容量的 Compute Engine 位置
• TPU 適用的機器類型，以及
• TPU PodSlice 的實體拓撲

設定含 TPU 的 RayCluster 資源

設定 RayCluster 資訊清單，準備執行 TPU 工作負載：

設定 TPU nodeSelector

GKE 會使用 Kubernetes nodeSelectors，確保 TPU 工作負載排程在適當的 TPU 拓撲和加速器上。如要進一步瞭解如何選取 TPU nodeSelector，請參閱「在 GKE Standard 中部署 TPU 工作負載」。

更新 ray-cluster.yaml 資訊清單，在拓撲為 2x2x1 的 v4 TPU Podslice 上排定 Pod：

nodeSelector:
  cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice
  cloud.google.com/gke-tpu-topology: 2x2x1

設定 TPU 容器資源

如要使用 TPU 加速器，您必須設定 RayCluster 資訊清單 workerGroupSpecs 的 TPU 容器欄位 google.com/tpu資源 limits 和 requests，指定 GKE 應為每個 Pod 分配的 TPU 晶片數量。

使用資源限制和要求更新 ray-cluster.yaml 資訊清單：

resources:
  limits:
    cpu: "1"
    ephemeral-storage: 10Gi
    google.com/tpu: "4"
    memory: "2G"
   requests:
    cpu: "1"
    ephemeral-storage: 10Gi
    google.com/tpu: "4"
    memory: "2G"

設定工作站群組 numOfHosts

KubeRay v1.1.0 在 RayCluster 自訂資源中新增 numOfHosts 欄位，用於指定每個工作站群組副本要建立的 TPU 主機數量。對於多主機工作站群組，副本會視為 PodSlice，而非個別工作站，每個副本都會建立 numOfHosts 個工作站節點。

使用下列內容更新 ray-cluster.yaml 資訊清單：

workerGroupSpecs:
  # Several lines omitted
  numOfHosts: 1 # the number of "hosts" or workers per replica

建立 RayService 自訂資源

建立 RayService 自訂資源：

1. 請查看下列資訊清單：

   apiVersion: ray.io/v1
   kind: RayService
   metadata:
     name: stable-diffusion-tpu
   spec:
     serveConfigV2: |
       applications:
         - name: stable_diffusion
           import_path: ai-ml.gke-ray.rayserve.stable-diffusion.stable_diffusion_tpu:deployment
           runtime_env:
             working_dir: "https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/archive/refs/heads/main.zip"
             pip:
               - diffusers==0.7.2
               - flax
               - jax[tpu]==0.4.11
               - -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
               - fastapi
     rayClusterConfig:
       rayVersion: '2.9.0'
       headGroupSpec:
         rayStartParams: {}
         template:
           spec:
             containers:
             - name: ray-head
               image: rayproject/ray-ml:2.9.0-py310
               ports:
               - containerPort: 6379
                 name: gcs
               - containerPort: 8265
                 name: dashboard
               - containerPort: 10001
                 name: client
               - containerPort: 8000
                 name: serve
               resources:
                 limits:
                   cpu: "2"
                   memory: "8G"
                 requests:
                   cpu: "2"
                   memory: "8G"
       workerGroupSpecs:
       - replicas: 1
         minReplicas: 1
         maxReplicas: 10
         numOfHosts: 1
         groupName: tpu-group
         rayStartParams: {}
         template:
           spec:
             containers:
             - name: ray-worker
               image: rayproject/ray-ml:2.9.0-py310
               resources:
                 limits:
                   cpu: "100"
                   ephemeral-storage: 20Gi
                   google.com/tpu: "4"
                   memory: 200G
                 requests:
                   cpu: "100"
                   ephemeral-storage: 20Gi
                   google.com/tpu: "4"
                   memory: 200G
             nodeSelector:
               cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice
               cloud.google.com/gke-tpu-topology: 2x2x1

   這份資訊清單說明 RayService 自訂資源，該資源會建立 RayCluster 資源，其中包含 1 個頭部節點和 TPU 工作站群組 (拓撲為 2x2x1)，也就是每個工作站節點都會有 4 個 v4 TPU 晶片。

   TPU 節點屬於單一 v4 TPU Podslice，拓撲為 2x2x1。如要建立多主機工作站群組，請將 gke-tpu nodeSelector 值、google.com/tpu 容器限制和要求，以及 numOfHosts 值替換為多主機設定。如要進一步瞭解 TPU 多主機拓撲，請參閱 Cloud TPU 說明文件中的「系統架構」。

2. 將資訊清單套用至叢集：

   kubectl apply -f ray-service-tpu.yaml
   

3. 確認 RayService 資源是否正在執行：

   kubectl get rayservices
   

   輸出結果會與下列內容相似：

   NAME                   SERVICE STATUS   NUM SERVE ENDPOINTS
   stable-diffusion-tpu   Running          2
   

   在這個輸出中，SERVICE STATUS 欄中的 Running 表示 RayService 資源已準備就緒。

(選用) 查看 Ray 資訊主頁

您可以透過 Ray 資訊主頁查看 Ray Serve 部署作業和相關記錄。

1. 從 Ray 服務的 Ray 節點建立通訊埠轉送工作階段，連線至 Ray 資訊主頁：

   kubectl port-forward svc/stable-diffusion-tpu-head-svc 8265:8265
   

2. 在網路瀏覽器中前往 http://localhost:8265/。

3. 按一下「供應」分頁標籤。

將提示傳送至模型伺服器

1. 從 Ray head 服務建立 Serve 端點的通訊埠轉送工作階段：

   kubectl port-forward svc/stable-diffusion-tpu-serve-svc 8000
   

2. 開啟新的 Cloud Shell 工作階段。

3. 將文字轉圖像提示提交至 Stable Diffusion 模型伺服器：

   python stable_diffusion_tpu_req.py  --save_pictures
   

   穩定擴散推論結果會儲存至名為 diffusion_results.png 的檔案。

   

清除所用資源

刪除專案

注意：刪除專案會造成以下結果：

• 專案中的所有內容都會遭到刪除。 如果您使用現有專案來進行本文中的工作，刪除專案將一併刪除其中已完成的所有工作。
• 自訂的專案 ID 會消失。當您之前建立這個專案時，可能建立了想要在日後使用的自訂專案 ID。如要保留使用該專案 ID 的網址 (例如 appspot.com 網址)，請刪除在該專案中選取的資源，而不是刪除整個專案。

如果打算探索多種架構、教學課程或快速入門導覽課程，重複使用專案有助於避免超出專案配額限制。

Delete a Google Cloud project:

gcloud projects delete PROJECT_ID

刪除個別資源

如要刪除叢集，請輸入：

gcloud container clusters delete ${CLUSTER_NAME}

後續步驟

• 瞭解 Kubernetes 上的 Ray。
• 參閱 KubeRay 說明文件。
• 探索 Google Cloud 的參考架構、圖表和最佳做法。 歡迎瀏覽我們的雲端架構中心。