本頁面由 Cloud Translation API 翻譯而成。

在 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) 是專門的硬體加速器，可大幅加快大規模機器學習模型的訓練和推論速度。搭配使用 Ray 和 TPU，即可順暢擴展高效能機器學習應用程式。如要進一步瞭解 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的下列計費元件：

如要根據預測用量估算費用，請使用 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.

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.

Make sure that billing is enabled for your Google Cloud project.

Enable the GKE API:

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.

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.

Make sure that billing is enabled for your Google Cloud project.

Enable the GKE API:

gcloud services enable container.googleapis.com

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 PROJECT_ID with your project ID.
- Replace USER_IDENTIFIER with the identifier for your user account. For example, user:myemail@example.com.
- Replace ROLE with each individual role.

確保配額充足

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

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

準備環境

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

在 Google Cloud 控制台中，按一下Google Cloud 控制台中的「啟用 Cloud Shell」，即可啟動 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 自訂資源，該資源會建立 RayCluster 資源，其中包含 1 個頭部節點和 TPU 工作站群組 (拓撲為 2x2x1)，也就是每個工作站節點都會有 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 的檔案。

清除所用資源

刪除專案

Delete a Google Cloud project:

gcloud projects delete PROJECT_ID

刪除個別資源

如要刪除叢集，請輸入：

gcloud container clusters delete ${CLUSTER_NAME}

後續步驟

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

在 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 資訊主頁

將提示傳送至模型伺服器

清除所用資源

刪除專案

刪除個別資源

後續步驟

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

設定 TPU `nodeSelector`

設定工作站群組 `numOfHosts`