在 GKE 上使用 TPU 部署 Stable Diffusion 模型#

注意： Ray Serve 应用及其客户端的 Python 文件位于 ray-project/serve_config_examples。本指南改编自 tensorflow/tpu 示例。

步骤 1：创建带 TPU 的 Kubernetes 集群#

按照为 KubeRay 创建带 TPU 的 GKE 集群创建一个包含 1 个 CPU 节点和 1 个 TPU 节点的 GKE 集群。

步骤 2：安装 KubeRay operator#

如果在 GKE 集群中启用了 Ray Operator Addon，请跳过此步骤。请按照部署 KubeRay Operator 的说明，从 Helm 仓库安装最新的稳定版 KubeRay Operator。KubeRay v1.1.0+ 支持多主机 TPU。请注意，此示例中的 YAML 文件使用了 serveConfigV2，KubeRay 从 v0.6.0 开始支持。

步骤 3：安装 RayService CR#

# Creates a RayCluster with a single-host v4 TPU worker group of 2x2x1 topology.
kubectl apply -f https://raw.githubusercontent.com/ray-project/kuberay/master/ray-operator/config/samples/ray-service.tpu-single-host.yaml

KubeRay Operator v1.1.0 在 RayCluster CR 中添加了一个新的 NumOfHosts 字段，支持多主机工作节点组。此字段指定每个副本要创建的工作节点数量，其中每个副本代表一个多主机 Pod 切片。 NumOfHosts 的值应与给定 cloud.google.com/gke-tpu-topology 节点选择器期望的 TPU VM 主机数量匹配。在此示例中，Stable Diffusion 模型足够小，可以在单个 TPU 主机上运行，因此 RayService manifest 中的 numOfHosts 设置为 1。

步骤 4：在 Ray Dashboard 中查看 Serve 部署#

验证 RayService CR 已部署并正在运行

kubectl get rayservice

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

通过 Ray Head 服务端口转发 Ray Dashboard。要在本地查看仪表板，请在浏览器中打开 https://:8265/。

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

在 Ray Dashboard 的“Serve”选项卡中监控 RayService CR 的状态。已安装的 RayService CR 应创建一个名为“stable_diffusion”的运行中应用。该应用应有两个部署：接收输入提示的 API 入口点，以及 Stable Diffusion 模型服务器。

serve_dashboard

步骤 5：向模型服务器发送文本到图像的提示#

端口转发 Ray Serve 服务

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

在单独的终端中，下载 Python 提示脚本

curl -LO https://raw.githubusercontent.com/ray-project/serve_config_examples/master/stable_diffusion/stable_diffusion_tpu_req.py

安装运行 Python 脚本所需的依赖项

# Create a Python virtual environment.
python3 -m venv myenv
source myenv/bin/activate

pip install numpy pillow requests tqdm

向 Stable Diffusion 模型服务器提交文本到图像的提示

python stable_diffusion_tpu_req.py  --save_pictures

Python 提示脚本将 Stable Diffusion 推理的结果保存到名为 diffusion_results.png 的文件中。

diffusion_results