優(yōu)化 Stable Diffusion 在 GKE 上的啟動(dòng)體驗(yàn)

現(xiàn)如今隨著 AIGC 這個(gè)話題越來越熱，越來越多優(yōu)秀的開源項(xiàng)目基于文生圖的 AI 模型如 MidJourney，Stable Diffusion 等應(yīng)運(yùn)而生。Stable Diffusion 是一個(gè)文字生成圖像的 Diffusion 模型，它能夠根據(jù)給定任何文本輸入生成逼真的圖像。我們在 GitHub Repo 中提供了三種不同的解決方案 (可參考https://github.com/nonokangwei/Stable-Diffusion-on-GCP)，可以快速地分別在 GCP Vertex AI，GKE，和基于 Agones 的平臺上部署 Stable Diffusion，以提供彈性的基礎(chǔ)設(shè)施保證 Stable Diffusion 提供穩(wěn)定的服務(wù)。本文將重點(diǎn)討論 Stable Diffusion 模型在 GKE 上的實(shí)踐。

在實(shí)踐中，我們也遇到了一些問題，例如 Stable Diffusion 的容器鏡像較大，大約達(dá)到 10-20GB，導(dǎo)致容器在啟動(dòng)過程中拉取鏡像的速度變慢，從而影響了啟動(dòng)時(shí)間。在需要快速擴(kuò)容的場景下，啟動(dòng)新的容器副本需要超過 10 分鐘的時(shí)間，嚴(yán)重影響了用戶體驗(yàn)。

我們看到容器的啟動(dòng)過程，按時(shí)序排列：

在這段時(shí)序分析中，我們可以看到，在 Stable Diffusion WebUI 運(yùn)行在容器上啟動(dòng)慢主要面臨的問題是由于整個(gè) runtime 依賴較多，導(dǎo)致容器鏡像太大從而花費(fèi)了很長時(shí)間拉取下載、也造成了 pod 啟動(dòng)初始化加載時(shí)間過長。于是，我們考慮優(yōu)化啟動(dòng)時(shí)間從以下三個(gè)方面入手：

本文著重為大家介紹通過基礎(chǔ)環(huán)境與 runtime 依賴分離方式，借助磁盤復(fù)制的高性能來優(yōu)化 Stable Diffusion WebUI 容器啟動(dòng)時(shí)間的方案。

首先，我們可以參考官方 Stable Diffusion WebUI 安裝說明，生成其 Dockerfile。在這里給大家一個(gè)參考: https://github.com/nonokangwei/Stable-Diffusion-on-GCP/blob/main/Stable-Diffusion-UI-Agones/sd-webui/Dockerfile

在初始構(gòu)建的 Stable Diffusion 的容器鏡像中，我們發(fā)現(xiàn)除了基礎(chǔ)鏡像 nvidia runtime 之外，還安裝了大量的庫和擴(kuò)展等。

▲調(diào)優(yōu)之前容器鏡像大小為 16.3GB

在 Dockerfile 優(yōu)化方面，我們對 Dockerfile 進(jìn)行分析后，發(fā)現(xiàn) nvidia runtime 約占 2GB，而 PyTorch 庫是一個(gè)非常大的包，約占 5GB。另外 Stable Diffusion 及其擴(kuò)展等也占據(jù)了一定的空間。因此，我們按照最小可用環(huán)境為原則，去除環(huán)境中不必要的依賴。將 nvidia runtime 作為基礎(chǔ)鏡像，然后把 PyTorch、Stable Diffusion 的庫和擴(kuò)展等從原始鏡像中分離出來，單獨(dú)存放在文件系統(tǒng)中。

以下是初始的 Dockerfile 的片段。

我們在移除 Pytorch 的庫和 Stable Diffusion 之后，我們只保留了基礎(chǔ)鏡像 nvidia runtime 在新的 Dockerfile 中。

▲基礎(chǔ)鏡像變成了 2GB

其余的運(yùn)行時(shí)類庫和 extension 等存放在磁盤鏡像中，磁盤鏡像的大小為 6.77GB。采用磁盤鏡像的好處是，它可以最多支持同時(shí)恢復(fù) 1,000 塊磁盤，完全能滿足大規(guī)模擴(kuò)縮容的使用場景。

然而，問題來了，如何將這個(gè)單獨(dú)的文件系統(tǒng)掛載到容器運(yùn)行時(shí)中呢？一種想法是使用 Persistent VolumeClaim (PVC) 進(jìn)行掛載，但由于 Stable Diffusion WebUI 在運(yùn)行時(shí)既需要讀取又需要寫入磁盤，而 GKE 的 PD CSI 驅(qū)動(dòng)程序目前不支持多寫入 ReadWriteMany，只有像 Filestore 這樣的 NFS 文件系統(tǒng)才能支持，但是通過網(wǎng)絡(luò)掛載的 Filestore 就延遲來說仍然無法達(dá)到快速啟動(dòng)的效果。同時(shí)，由于 GKE 目前不支持在創(chuàng)建或更新 Nodepool 時(shí)掛載磁盤，所以我們考慮使用 DaemonSet 在 GKE 節(jié)點(diǎn)啟動(dòng)時(shí)掛載磁盤。具體做法如下：

那么如何將磁盤掛載到 GKE 的節(jié)點(diǎn)上呢？可以直接調(diào)用 Cloud SDK，創(chuàng)建基于磁盤鏡像的磁盤。

另外，正如我們前面提到的那樣，在優(yōu)化鏡像下載和加載時(shí)間方面，我們還啟用了 GKE Image Streaming 來加速鏡像的拉取速度。它的工作原理是使用網(wǎng)絡(luò)掛載將容器數(shù)據(jù)層掛載到 containerd 中，并在網(wǎng)絡(luò)、內(nèi)存和磁盤上使用多個(gè)緩存層對其進(jìn)行支持。一旦我們準(zhǔn)備好 Image Streaming 掛載，您的容器就會在幾秒鐘內(nèi)從 ImagePulling 狀態(tài)轉(zhuǎn)換為 Running (無論容器大小)；這有效地將應(yīng)用程序啟動(dòng)與容器映像中所需數(shù)據(jù)的數(shù)據(jù)傳輸并行化。因此，您可以看到更快的容器啟動(dòng)時(shí)間和更快速的自動(dòng)縮放。

我們開啟了 Cluster Autoscaler 功能，讓有更多的請求到來時(shí)，GKE 節(jié)點(diǎn)自動(dòng)進(jìn)行彈性擴(kuò)展。通過 Cluster Autoscaler 觸發(fā)并決定擴(kuò)展到多少個(gè)節(jié)點(diǎn)來接收新增的請求。當(dāng) CA 觸發(fā)了新的一輪擴(kuò)容，新的 GKE 節(jié)點(diǎn)注冊到集群以后，Daemonset 就會開始工作，幫助掛載存儲了 runtime 依賴的磁盤鏡像，而 Stable Diffusion Deployment 則會通過 HostPath 來訪問這個(gè)掛載在節(jié)點(diǎn)上的磁盤。

我們還使用了 Cluster Autoscaler 的 Optimization Utilization Profile 來縮短擴(kuò)縮容時(shí)間、節(jié)省成本并提高機(jī)器利用率。

最后的啟動(dòng)效果如下：

按時(shí)序排列

總共經(jīng)歷了大約 3 分鐘的時(shí)間，就完成了啟動(dòng)一個(gè)新的 Stale Diffusion 容器實(shí)例，并在一個(gè)新的 GKE 節(jié)點(diǎn)上進(jìn)行正常服務(wù)的過程。相比于之前的 12 分鐘，可以看見，明顯的提升啟動(dòng)速度改善了用戶體驗(yàn)。

完整代碼: https://github.com/nonokangwei/Stable-Diffusion-on-GCP/tree/main/Stable-Diffusion-UI-Agones/optimizated-init

除了掛載 Disk 到 GKE 節(jié)點(diǎn)上，還有一種嘗試，我們可以使用 StatefulSet 來掛載 PVC。

具體做法如下:先定義一個(gè) storageclass，注意我們使用DiskImageType: images來指定 PVC從 Disk Image 來恢復(fù)，而不是 Snapshot。

再定義一個(gè) VoluemSnapShotContent，它指定了 source 為一個(gè) Disk Image sd-image。

接下來，我們再創(chuàng)建一個(gè) VolumeSnapShot，指定它的 source 是剛剛定義的VoluemSnapShotContent。

最后，我們創(chuàng)建一個(gè) StatefulSet 來掛載這個(gè) VolumeSnapShot。

我們嘗試擴(kuò)容更多的副本。

可見 GKE 可以支持并行的啟動(dòng)這些 Pod，并且分別掛載相應(yīng)的磁盤。

PersistentVolumeClaims

完整代碼:

https://github.com/Leisureroad/volumesnapshot-from-diskimage

最終，我們可以看見如下的 Stable Diffusion 的 WebUI。