Google 會運用 AI 技術將內容翻譯成你偏好的語言，但可能會出錯。

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

標準版

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

本指南適用於生成式 AI 客戶、GKE 新手或現有使用者、機器學習工程師、MLOps (DevOps) 工程師，或是有意使用 Kubernetes 容器自動化調度管理功能，透過 Ray 提供模型的平台管理員。

關於 Ray 和 Ray Serve

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

關於 TPU

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

關於 KubeRay TPU 初始化 Webhook

做為 Ray Operator 外掛程式的一部分，GKE 提供驗證和變動網路掛鉤，可處理 TPU Pod 排程，以及 JAX 等架構用於容器初始化作業的特定 TPU 環境變數。JAXKubeRay TPU 網路鉤子會使用下列屬性，變更要求 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的計費元件：

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

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

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

事前準備

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

登入 Google Cloud 帳戶。如果您是 Google Cloud新手，歡迎建立帳戶，親自評估產品在實際工作環境中的成效。新客戶還能獲得價值 $300 美元的免費抵免額，可用於執行、測試及部署工作負載。

安裝 Google Cloud CLI。

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

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

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

gcloud init

建立或選取 Google Cloud 專案。

選取或建立專案所需的角色

選取專案：選取專案時，不需要具備特定 IAM 角色，只要您已獲授角色，即可選取任何專案。
建立專案：如要建立專案，您需要「專案建立者」角色 (roles/resourcemanager.projectCreator)，其中包含 resourcemanager.projects.create 權限。瞭解如何授予角色。

建立 Google Cloud 專案：
```
gcloud projects create PROJECT_ID
```
將 PROJECT_ID 替換為您要建立的 Google Cloud 專案名稱。
選取您建立的 Google Cloud 專案：
```
gcloud config set project PROJECT_ID
```
將 PROJECT_ID 替換為 Google Cloud 專案名稱。

確認專案已啟用計費功能 Google Cloud 。

啟用 GKE API：

啟用 API 時所需的角色

如要啟用 API，您需要具備服務使用情形管理員 IAM 角色 (roles/serviceusage.serviceUsageAdmin)，其中包含 serviceusage.services.enable 權限。瞭解如何授予角色。

gcloud services enable container.googleapis.com

安裝 Google Cloud CLI。

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

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

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

gcloud init

建立或選取 Google Cloud 專案。

選取或建立專案所需的角色

選取專案：選取專案時，不需要具備特定 IAM 角色，只要您已獲授角色，即可選取任何專案。
建立專案：如要建立專案，您需要「專案建立者」角色 (roles/resourcemanager.projectCreator)，其中包含 resourcemanager.projects.create 權限。瞭解如何授予角色。

建立 Google Cloud 專案：
```
gcloud projects create PROJECT_ID
```
將 PROJECT_ID 替換為您要建立的 Google Cloud 專案名稱。
選取您建立的 Google Cloud 專案：
```
gcloud config set project PROJECT_ID
```
將 PROJECT_ID 替換為 Google Cloud 專案名稱。

確認專案已啟用計費功能 Google Cloud 。

啟用 GKE API：

啟用 API 時所需的角色

如要啟用 API，您需要具備服務使用情形管理員 IAM 角色 (roles/serviceusage.serviceUsageAdmin)，其中包含 serviceusage.services.enable 權限。瞭解如何授予角色。

gcloud services enable container.googleapis.com

將角色授予使用者帳戶。針對下列每個 IAM 角色，執行一次下列指令： roles/container.clusterAdmin, roles/container.admin
```
gcloud projects add-iam-policy-binding PROJECT_ID --member="user:USER_IDENTIFIER" --role=ROLE
```
更改下列內容：
- PROJECT_ID：專案 ID。
- USER_IDENTIFIER：使用者帳戶的 ID。例如：myemail@example.com。
- ROLE：授予使用者帳戶的 IAM 角色。

確保配額充足

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

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

準備環境

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

從 Google Cloud 控制台啟動 Cloud Shell 工作階段，方法是點選Google Cloud 控制台中的「啟用 Cloud Shell」。系統會在 Google Cloud 控制台的底部窗格啟動工作階段。
設定環境變數：
```
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 或之後。

複製 GitHub 存放區：

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

變更為工作目錄：

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

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

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

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

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

建立單一主機 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 自訂資源：

請查看下列資訊清單：

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 自訂資源，該資源會建立具有 1 個頭部節點的 RayCluster 資源，以及具有 2x2x1 拓撲的 TPU 工作站群組，也就是每個 worker 節點都會有 4 個 v4 TPU 晶片。

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

將資訊清單套用至叢集：
```
kubectl apply -f ray-service-tpu.yaml
```
確認 RayService 資源是否正在執行：
```
kubectl get rayservices
```
輸出結果會與下列內容相似：
```
NAME                   SERVICE STATUS   NUM SERVE ENDPOINTS
stable-diffusion-tpu   Running          2
```
在這個輸出內容中，SERVICE STATUS 欄中的 Running 表示 RayService 資源已準備就緒。

(選用) 查看 Ray 資訊主頁

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

從 Ray 首節點服務建立通訊埠轉送工作階段，連線至 Ray 資訊主頁：
```
kubectl port-forward svc/stable-diffusion-tpu-head-svc 8265:8265
```
使用網路瀏覽器前往 http://localhost:8265/。
按一下「供應」分頁標籤。

將提示傳送至模型伺服器

從 Ray head 服務建立通訊埠轉送工作階段至 Serve 端點：
```
kubectl port-forward svc/stable-diffusion-tpu-serve-svc 8000
```
開啟新的 Cloud Shell 工作階段。
將文字轉圖像提示提交至 Stable Diffusion 擴散模型伺服器：
```
python stable_diffusion_tpu_req.py  --save_pictures
```
穩定擴散推論結果會儲存至名為 diffusion_results.png 的檔案。

觀察 Ray 工作負載

如要查看 RayJob 的詳細資料，請前往 Google Cloud 控制台的「Kubernetes Engine > AI/ML > Jobs」部分。

在 Google Cloud 控制台中查看 RayJobs

清除所用資源

刪除專案

刪除 Google Cloud 專案：

gcloud projects delete PROJECT_ID

刪除個別資源

如要刪除叢集，請輸入：

gcloud container clusters delete ${CLUSTER_NAME}

後續步驟

瞭解 Kubernetes 上的 Ray。
參閱 KubeRay 說明文件。
查看 Google Cloud 的參考架構、圖表和最佳做法。歡迎瀏覽我們的 Cloud Architecture Center。

在 Google Kubernetes Engine (GKE) 上，使用 TPU 部署搭載 Stable Diffusion 模型的 Ray Serve 應用程式 透過集合功能整理內容 你可以依據偏好儲存及分類內容。

關於 Ray 和 Ray Serve

關於 TPU

關於 KubeRay TPU 初始化 Webhook

目標

費用

事前準備

確保配額充足

準備環境

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

設定含 TPU 的 RayCluster 資源

設定 TPU nodeSelector

設定 TPU 容器資源

設定工作人員群組 numOfHosts

建立 RayService 自訂資源

(選用) 查看 Ray 資訊主頁

將提示傳送至模型伺服器

觀察 Ray 工作負載

清除所用資源

刪除專案

刪除個別資源

後續步驟

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

設定 TPU `nodeSelector`

設定工作人員群組 `numOfHosts`