将 Google Kubernetes Engine (GKE) 上稳定扩散的启动时间缩短 4 倍

Cloud Ace 是 Google Cloud 全球战略合作伙伴，在亚太地区、欧洲、南北美洲和非洲拥有二十多个办公室。Cloud Ace 在谷歌专业领域认证及专业知识目前排名全球第一位，并连续多次获得 Google Cloud 各类奖项。作为谷歌云托管服务商，我们提供谷歌云、谷歌地图、谷歌办公套件、谷歌云认证培训等服务。

Cloud Ace - 谷歌云 | 谷歌云全球战略合作伙伴 | 云服务器据点最多经销商​cloud-ace.cn/​编辑

随着人工智能生成内容（AIGC）的日益流行，基于文本到图像的人工智能模型（例如稳定扩散）的开源项目已经出现。稳定扩散是一种扩散模型，可根据给定的文本输入生成逼真的图像。在此 GitHub 存储库中，我们提供了三种不同的解决方案，分别用于在 Google Cloud Vertex AI、Google Kubernetes Engine (GKE) 和基于 Agones 的平台上快速部署 Stable Diffusion，以通过弹性基础设施确保稳定的服务交付。本文将重点介绍 GKE 上的稳定扩散模型，并将启动时间缩短多达 4 倍。

问题陈述

Stable Diffusion 的容器镜像较大，大约达到 10-20GB，这会减慢容器启动时的镜像拉取过程，从而影响启动时间。在需要快速伸缩的场景下，启动新的容器副本可能需要10分钟以上，严重影响用户体验。

在容器启动过程中，我们可以按时间顺序看到以下事件：

1. 触发Cluster Autoscaler进行伸缩+节点启动和Pod调度：225秒
2. 图片拉取启动：4秒
3. 图片拉取：5分23秒
4. Pod 启动：1 秒
5. SD-WebUI服务：超过2分钟

分析这个时间序列可以看出，运行在容器中的 Stable Diffusion WebUI 启动缓慢主要是由于整个运行时的依赖较重，导致容器镜像体积过大，拉取镜像和初始化pod的时间变长。

因此，我们从以下三个方面考虑优化启动时间：

1. 优化Dockerfile：选择合适的基础镜像，尽量减少运行时依赖的安装，以减小镜像大小。
2. 将基础环境与运行时依赖分离：通过PD磁盘镜像加速运行时环境的创建。
3. 利用 GKE Image Streaming：利用 GKE Image Streaming 优化图像加载时间，并利用 Cluster Autoscaler 来增强弹性缩放和调整大小的速度。

本文重点介绍一种通过将基础环境与运行时依赖项分离并利用高性能磁盘映像来优化 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运行时之外，还安装了大量的库、依赖项和扩展。

优化前，容器镜像大小为16.3GB。

在优化Dockerfile方面，经过对Dockerfile的分析，我们发现nvidia运行时占用了大约2GB，而PyTorch库是一个非常大的包，占用了5GB左右。此外，稳定扩散及其扩展也占据了一些空间。

因此，遵循最小可行环境的原则，我们可以从环境中删除不必要的依赖项。

我们可以使用 NVIDIA 运行时作为基础镜像，将 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://github.com/toshiaki1729/stable-diffusion-webui-dataset-tag-editor.git extensions/stable-diffusion-webui-dataset-tag-editor

移出 Pytorch 库和稳定扩散后，我们只在基础映像中保留了 NVIDIA 运行时，这是新的 Dockerfile。

使用PD磁盘镜像来存储库PD 磁盘映像是 Google Cloud 中实例部署的基石。这些虚拟映像通常称为模板或引导磁盘，包含基线操作系统以及实例首次启动时将拥有的所有应用程序软件和配置。这里的想法是将所有运行时库和扩展存储在磁盘映像中，在本例中磁盘映像的大小为 6.77GB。使用磁盘镜像的优点是最多可以支持1000个磁盘同时恢复，适合大规模扩缩容的场景。

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

我们使用 DaemonSet 在 GKE 节点启动时挂载磁盘。

具体步骤如下：

如前几节所述，为了加快初始启动以获得更好的性能，我们尝试将永久性磁盘挂载到 GKE 节点以放置运行时库以实现稳定扩散。

利用 GKE Image Streaming 和集群自动缩放器

此外，如前所述，我们还启用了 GKE Image Streaming 来加速图像拉取和加载过程。GKE Image Streaming 的工作原理是使用网络挂载将容器的数据层附加到 containerd，并通过网络、内存和磁盘上的多个缓存层来支持它。一旦我们准备好 Image Streaming 挂载，无论容器大小如何，您的容器都会在几秒钟内从 ImagePulling 状态转换为 Running 状态。这有效地将应用程序启动与容器映像中所需数据的数据传输并行化。因此，您可以体验更快的容器启动时间和更快的自动扩展。

我们启用了 Cluster Autoscaler (CS) 功能，该功能允许 GKE 节点在请求增加时自动扩展。Cluster Autoscaler 触发并确定处理额外请求所需的节点数量。当 Cluster Autoscaler 启动新的扩展浪潮并且新的 GKE 节点在集群中注册时，DaemonSet 开始工作以协助挂载包含运行时依赖项的磁盘映像。然后，稳定扩散部署通过 HostPath 访问该磁盘。此外，我们还利用了Cluster Autoscaler 的优化利用率配置文件，这是 GKE CA 上的一个配置文件，它优先考虑优化利用率而不是保留集群中的备用资源，以减少扩展时间、节省成本并提高机器利用率。

最终结果

最终启动结果如下：

按年代顺序：

1. 触发集群自动缩放器进行缩放：38秒
2. 节点启动和Pod调度：89秒
3. 安装PVC：4秒
4. 图片拉取启动：10秒
5. 图片拉取：1秒
6. Pod 启动：1 秒
7. 能够提供sd-webui服务（大约）：65秒

总体而言，启动新的 Stable Diffusion 容器实例并开始在新的 GKE 节点上提供服务大约需要 3 分钟。与之前的12分钟相比，显然启动速度的显著提升，提升了用户体验。

在这里查看完整代码：

https://github.com/nonokangwei/Stable-Diffusion-on-GCP/tree/main/Stable-Diffusion-UI-Agones/optimized-init

注意事项：虽然上述技术分割了依赖关系，使容器大小更小，并且您可以从 PD 磁盘映像加载库，但仍有一些缺点需要考虑。将所有内容打包在一个容器映像中具有其优点，您可以拥有单个不可变且版本化的工件。将基本环境与运行时依赖项分离意味着您需要维护和更新多个工件。您可以通过构建工具来管理 PD 磁盘映像的更新来缓解这种情况。