本頁面由 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) 是專門的硬體加速器，可大幅加快大規模機器學習模型的訓練和推論速度。搭配 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。

準備環境

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

在 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 的檔案。

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