优化 Stable Diffusion 在 GKE 上的启动体验
以下文章来源于谷歌云服务，作者 Google Cloud
背景
现如今随着 AIGC 这个话题越来越热，越来越多优秀的开源项目基于文生图的 AI 模型如 MidJourney，Stable Diffusion 等应运而生。Stable Diffusion 是一个文字生成图像的 Diffusion 模型，它能够根据给定任何文本输入生成逼真的图像。我们在 GitHub Repo 中提供了三种不同的解决方案 (可参考https://github.com/nonokangwei/Stable-Diffusion-on-GCP)，可以快速地分别在 GCP Vertex AI，GKE，和基于 Agones 的平台上部署 Stable Diffusion，以提供弹性的基础设施保证 Stable Diffusion 提供稳定的服务。本文将重点讨论 Stable Diffusion 模型在 GKE 上的实践。
提出问题
在实践中，我们也遇到了一些问题，例如 Stable Diffusion 的容器镜像较大，大约达到 10-20GB，导致容器在启动过程中拉取镜像的速度变慢，从而影响了启动时间。在需要快速扩容的场景下，启动新的容器副本需要超过 10 分钟的时间，严重影响了用户体验。
我们看到容器的启动过程，按时序排列：
●触发 Cluster Autoscaler 扩容 + Node 启动并调度 Pod: 225s
●启动 Pull Image: 4s
●拉取镜像: 5m 23s
●启动 Pod: 1s
●能够提供 sd-webui 的服务 (大约): > 2m
在这段时序分析中，我们可以看到，在 Stable Diffusion WebUI 运行在容器上启动慢主要面临的问题是由于整个 runtime 依赖较多，导致容器镜像太大从而花费了很长时间拉取下载、也造成了 pod 启动初始化加载时间过长。于是，我们考虑优化启动时间从以下三个方面入手：
●优化 Dockerfile，选择正确的 base image，精简 runtime 的依赖安装，减小镜像大小。
●借助基础环境与 runtime 依赖分离方式，通过磁盘复制方式加速运行环境的创建。
●通过 GKE Image Streaming 优化镜像加载时间，利用 Cluster Autoscaler 提升弹性扩缩容速度。
本文着重为大家介绍通过基础环境与 runtime 依赖分离方式，借助磁盘复制的高性能来优化 Stable Diffusion WebUI 容器启动时间的方案。
优化 Dockerfile
首先，我们可以参考官方 Stable Diffusion WebUI 安装说明，生成其 Dockerfile。在这里给大家一个参考: https://github.com/nonokangwei/Stable-Diffusion-on-GCP/blob/main/Stable-Diffusion-UI-Agones/sd-webui/Dockerfile
在初始构建的 Stable Diffusion 的容器镜像中，我们发现除了基础镜像 nvidia runtime 之外，还安装了大量的库和扩展等。
▲ 调优之前容器镜像大小为 16.3GB
在 Dockerfile 优化方面，我们对 Dockerfile 进行分析后，发现 nvidia runtime 约占 2GB，而 PyTorch 库是一个非常大的包，约占 5GB。另外 Stable Diffusion 及其扩展等也占据了一定的空间。因此，我们按照最小可用环境为原则，去除环境中不必要的依赖。将 nvidia runtime 作为基础镜像，然后把 PyTorch、Stable Diffusion 的库和扩展等从原始镜像中分离出来，单独存放在文件系统中。
以下是初始的 Dockerfile 的片段。
# Base image
FROM nvidia/cuda:11.8.0-runtime-ubuntu22.04
RUN set -ex && \
apt update && \
apt install -y wget git python3 python3-venv python3-pip libglib2.0-0 pkg-config libcairo2-dev && \
rm -rf /var/lib/apt/lists/*
# Pytorch
RUN python3 -m pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 --extra-index-url https://download.pytorch.org/whl/cu117
…
# Stable Diffusion
RUN git clone https://github.com/AUTOMATIC1111/stable-diffusion-webui.git
RUN git clone https://github.com/Stability-AI/stablediffusion.git /stable-diffusion-webui/repositories/stable-diffusion-stability-ai
RUN git -C /stable-diffusion-webui/repositories/stable-diffusion-stability-ai checkout cf1d67a6fd5ea1aa600c4df58e5b47da45f6bdbf
…
# Stable Diffusion extensions
RUN set -ex && cd stable-diffusion-webui \
&& git clone https://gitcode.net/ranting8323/sd-webui-additional-networks.git extensions/sd-webui-additional-networks \
&& git clone https://gitcode.net/ranting8323/sd-webui-cutoff extensions/sd-webui-cutoff \
&& git clone https://ghproxy.com/https://github.com/toshiaki1729/stable-diffusion-webui-dataset-tag-editor.git extensions/stable-diffusion-webui-dataset-tag-editor
我们在移除 Pytorch 的库和 Stable Diffusion 之后，我们只保留了基础镜像 nvidia runtime 在新的 Dockerfile 中。
FROM nvidia/cuda:11.8.0-runtime-ubuntu22.04
RUN set -ex && \
apt update && \
apt install -y wget git python3 python3-venv python3-pip libglib2.0-0 && \
rm -rf /var/lib/apt/lists/*
▲ 基础镜像变成了 2GB
其余的运行时类库和 extension 等存放在磁盘镜像中，磁盘镜像的大小为 6.77GB。采用磁盘镜像的好处是，它可以最多支持同时恢复 1,000 块磁盘，完全能满足大规模扩缩容的使用场景。
挂载磁盘到 GKE 节点
然而，问题来了，如何将这个单独的文件系统挂载到容器运行时中呢？一种想法是使用 Persistent VolumeClaim (PVC) 进行挂载，但由于 Stable Diffusion WebUI 在运行时既需要读取又需要写入磁盘，而 GKE 的 PD CSI 驱动程序目前不支持多写入 ReadWriteMany，只有像 Filestore 这样的 NFS 文件系统才能支持，但是通过网络挂载的 Filestore 就延迟来说仍然无法达到快速启动的效果。同时，由于 GKE 目前不支持在创建或更新 Nodepool 时挂载磁盘，所以我们考虑使用 DaemonSet 在 GKE 节点启动时挂载磁盘。具体做法如下：
那么如何将磁盘挂载到 GKE 的节点上呢？可以直接调用 Cloud SDK，创建基于磁盘镜像的磁盘。
gcloud compute disks create sd-lib-disk-$NOW --type=pd-balanced --size=30GB --zone=$ZONE --image=$IMAGE_NAME
gcloud compute instances attach-disk ${MY_NODE_NAME} --disk=projects/$PROJECT_ID/zones/$ZONE/disks/sd-lib-disk-$NOW --zone=$ZONE
利用 GKE Image Streaming
和 Cluster Autoscaler
另外，正如我们前面提到的那样，在优化镜像下载和加载时间方面，我们还启用了 GKE Image Streaming 来加速镜像的拉取速度。它的工作原理是使用网络挂载将容器数据层挂载到 containerd 中，并在网络、内存和磁盘上使用多个缓存层对其进行支持。一旦我们准备好 Image Streaming 挂载，您的容器就会在几秒钟内从 ImagePulling 状态转换为 Running (无论容器大小)；这有效地将应用程序启动与容器映像中所需数据的数据传输并行化。因此，您可以看到更快的容器启动时间和更快速的自动缩放。
我们开启了 Cluster Autoscaler 功能，让有更多的请求到来时，GKE 节点自动进行弹性扩展。通过 Cluster Autoscaler 触发并决定扩展到多少个节点来接收新增的请求。当 CA 触发了新的一轮扩容，新的 GKE 节点注册到集群以后，Daemonset 就会开始工作，帮助挂载存储了 runtime 依赖的磁盘镜像，而 Stable Diffusion Deployment 则会通过 HostPath 来访问这个挂载在节点上的磁盘。
我们还使用了 Cluster Autoscaler 的 Optimization Utilization Profile 来缩短扩缩容时间、节省成本并提高机器利用率。
最后的启动效果如下：
按时序排列
●触发 Cluster Autoscaler 扩容: 38s
●Node 启动并调度 Pod: 89s
●挂载 PVC: 4s
●启动 Pull Image: 10s
●拉取镜像: 1s
●启动 Pod: 1s
●能够提供 sd-webui 的服务 (大约): 65s
总共经历了大约 3 分钟的时间，就完成了启动一个新的 Stale Diffusion 容器实例，并在一个新的 GKE 节点上进行正常服务的过程。相比于之前的 12 分钟，可以看见，明显的提升启动速度改善了用户体验。
完整代码: https://github.com/nonokangwei/Stable-Diffusion-on-GCP/tree/main/Stable-Diffusion-UI-Agones/optimizated-init
通过 VolumeSnapshot
除了挂载 Disk 到 GKE 节点上，还有一种尝试，我们可以使用 StatefulSet 来挂载 PVC。
具体做法如下: 先定义一个 storageclass，注意我们使用 DiskImageType: images 来指定 PVC从 Disk Image 来恢复，而不是 Snapshot。
●Snapshot 每 10 分钟只能恢复一次，一小时以内恢复 6 次 Disk 的限制。
●而 Image 可以支持每 30 秒恢复一次，最多 1,000 个 Disk。
apiVersion: snapshot.storage.k8s.io/v1
kind: VolumeSnapshotClass
metadata:
name: image-class
driver: pd.csi.storage.gke.io
deletionPolicy: Delete
parameters:
DiskImageType: images
再定义一个 VoluemSnapShotContent，它指定了 source 为一个 Disk Image sd-image。
apiVersion: snapshot.storage.k8s.io/v1
kind: VolumeSnapshotContent
metadata:
name: test-snapshotcontent-from-image
spec:
deletionPolicy: Retain
driver: pd.csi.storage.gke.io
volumeSnapshotClassName: image-class
source:
snapshotHandle:projects/flius-vpc-2/global/images/sd-image
volumeSnapshotRef:
name: test-snapshot
namespace: default
接下来，我们再创建一个 VolumeSnapShot，指定它的 source 是刚刚定义的VoluemSnapShotContent。
apiVersion: snapshot.storage.k8s.io/v1
kind: VolumeSnapshot
metadata:
name: test-snapshot
namespace: default
spec:
source:
volumeSnapshotContentName:test-snapshotcontent-from-image
最后，我们创建一个 StatefulSet 来挂载这个 VolumeSnapShot。
apiVersion: apps/v1
kind: StatefulSet
metadata:
name: stable-diffusion-statefulset-image
labels:
app: stable-diffusion
spec:
podManagementPolicy: "Parallel"
replicas: 1
selector:
matchLabels:
app: stable-diffusion
template:
metadata:
labels:
app: stable-diffusion
spec:
containers:
- name: stable-diffusion-webui
image: us-central1-docker.pkg.dev/flius-vpc-2/stable-diffusion-repo/sd-webui-final:0.1
command: ["/bin/bash"]
args: ["-c", "source /runtime-lib/bin/activate; cp/user-watch.py /runtime-lib/stable-diffusion-webui/user-watch.py;cp/start.sh /runtime-lib/stable-diffusion-webui/start.sh; cd /runtime-lib/stable-diffusion-webui; python3 launch.py --listen --xformers --enable-insecure-extension-access--no-gradio-queue" ]
volumeMounts:
- mountPath: "/runtime-lib"
name: runtime-lib
resources:
limits:
nvidia.com/gpu: 1
ports:
- containerPort: 7860
volumeClaimTemplates:
- metadata:
name: runtime-lib
spec:
dataSource:
name: test-snapshot
kind: VolumeSnapshot
apiGroup: snapshot.storage.k8s.io
accessModes: [ "ReadWriteOnce" ]
storageClassName: "standard-rwo"
resources:
requests:
storage: 30Gi
我们尝试扩容更多的副本。
kubectl scale statefulset stable-diffusion-statefulset-image --replicas=15
可见 GKE 可以支持并行的启动这些 Pod，并且分别挂载相应的磁盘。
PersistentVolumeClaims
完整代码:
https://github.com/Leisureroad/volumesnapshot-from-diskimage
最终，我们可以看见如下的 Stable Diffusion 的 WebUI。