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

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

步骤 1：使用 TPU 创建 Kubernetes 集群#

按照使用 TPU 创建 GKE 集群用于 KubeRay 来创建一个包含 1 个 CPU 节点和 1 个 TPU 节点的 GKE 集群。

步骤 2：安装 KubeRay 操作员#

如果您的 GKE 集群中启用了 Ray Operator Addon，请跳过此步骤。按照部署 KubeRay 操作员说明从 Helm 仓库安装最新稳定的 KubeRay 操作员。多主机 TPU 支持在 KubeRay v1.1.0+ 中可用。请注意，此示例中的 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 操作符 v1.1.0 为 RayCluster CR 添加了一个新的 NumOfHosts 字段，支持多主机工作组。此字段指定每个副本要创建的工作者数量，每个副本代表一个多主机的 Pod 切片。NumOfHosts 的值应与给定的 cloud.google.com/gke-tpu-topology 节点选择器预期的 TPU VM 主机数量匹配。在这个例子中，Stable Diffusion 模型足够小，可以在单个 TPU 主机上运行，因此在 RayService 清单中将 numOfHosts 设置为 1。

步骤4：在Ray仪表板中查看Serve部署#

验证您已部署 RayService CR 并且它正在运行：

kubectl get rayservice

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

从 Ray 头服务端口转发 Ray 仪表板。要查看仪表板，请在本地机器上打开 http://localhost: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 的文件中。

扩散结果