Google Kubernetes Engine（GKE）上の Stable Diffusion の起動時間を 4 分の 1 に短縮する
Google Cloud Japan Team
※この投稿は米国時間 2023 年 9 月 12 日に、Google Cloud blog に投稿されたものの抄訳です。
背景
AI 生成コンテンツ（AIGC）の人気が高まるなか、Stable Diffusion のような、テキストから画像を生成する AI モデルをベースとしたオープンソース プロジェクトが登場しました。Stable Diffusion は、与えられたテキスト入力に基づいて写実的な画像を生成する拡散モデルです。この GitHub リポジトリで提供されている 3 種類のソリューションは、Stable Diffusion をそれぞれ Google Cloud Vertex AI、Google Kubernetes Engine（GKE）、Agones ベースのプラットフォームにすばやくデプロイするためのもので、弾力性のあるインフラストラクチャによって安定したサービス提供を実現します。この記事では、GKE 上の Stable Diffusion モデルに焦点を当て、起動時間を最大 4 分の 1 に短縮する方法を示します。
問題提起
Stable Diffusion のコンテナ イメージはおよそ 10～20 GB のサイズに達するほど巨大です。そのため、コンテナ起動中のイメージ pull プロセスに時間がかかり、それが Stable Diffusion の起動時間に悪影響を及ぼしています。迅速なスケーリングが必要なシナリオでは、新しいコンテナ レプリカの起動に 10 分以上かかる場合があり、ユーザー エクスペリエンスが大幅に損なわれます。
コンテナの起動を観察したところ、時系列順に以下のイベントが発生しました。
- スケーリング用のクラスタ オートスケーラーのトリガー + Node の起動と Pod のスケジューリング: 225 秒
- イメージ pull の開始: 4 秒
- イメージの pull: 5 分 23 秒
- Pod の起動: 1 秒
- sd-webui のサービング: 2 分以上
この時系列を分析した結果、コンテナで実行される Stable Diffusion WebUI の起動が遅い主な原因は、ランタイム全体に膨大な依存関係があり、その結果コンテナ イメージのサイズが肥大してイメージの pull と Pod の初期化に長時間を要していることにあると判明しました。
したがって、以下の 3 つの側面から起動時間の最適化を検討します。
- Dockerfile の最適化: 適切なベースイメージを選択し、イメージサイズを縮小するためにランタイム依存関係のインストールを最小限に抑えます。
- ベース環境のランタイム依存関係からの分離: PD のディスク イメージを通じてランタイム環境の作成を高速化します。
- GKE イメージ ストリーミングの利用: GKE イメージ ストリーミングを利用してイメージ読み込み時間を最適化し、クラスタ オートスケーラーを使用して柔軟なスケーリングとサイズ変更をスピードアップします。
この記事では、ベース環境のランタイム依存関係からの分離とパフォーマンスの高いディスク イメージの利用によって Stable Diffusion WebUI コンテナの起動時間を最適化するソリューションを導入することに焦点を当てます。
Dockerfile の最適化
まず、Stable Diffusion WebUI の公式のインストール手順に基づくリファレンス Dockerfile を以下に示します。
この Stable Diffusion の初期構築版のコンテナ イメージには、ベースイメージの NVIDIA ランタイムに加えて、インストール対象となる数多くのライブラリ、依存関係、拡張機能も含まれています。
最適化前のコンテナ イメージのサイズは 16.3 GB でした。
Dockerfile の最適化という観点でこの Dockerfile を分析してみると、NVIDIA ランタイムが占めるサイズは約 2 GB でしたが、PyTorch ライブラリが巨大で、そのサイズは約 5 GB でした。さらに、Stable Diffusion とその拡張機能もかなりのスペースを占めていました。
したがって、実用最小限の環境の原則に従い、不必要な依存関係を環境から除去します。
NVIDIA ランタイムをベースイメージとして使用し、PyTorch ライブラリと Stable Diffusion のライブラリおよび拡張機能を元のイメージから分離して別のファイル システムに保存することにしました。以下は元の Dockerfile から抜粋したスニペットです。
Pytorch ライブラリと Stable Diffusion を除去し、NVIDIA ランタイムのみをベースイメージに保持した新しい Dockerfile は次のとおりです。
PD のディスク イメージを使用してライブラリを保存
PD のディスク イメージは、Google Cloud におけるインスタンス デプロイメントの基礎です。テンプレートまたはブートストラップ ディスクとも呼ばれるこの仮想イメージには、ベースライン オペレーティング システムとすべてのアプリケーション ソフトウェア、およびインスタンスの初回起動時に使用される構成が含まれます。ここに示すアイデアは、すべてのランタイム ライブラリと拡張機能を 1 つのディスク イメージに格納するというものです。この例では、ディスク イメージのサイズは 6.77 GB になります。ディスク イメージを使用することの利点は、同時に最大 1,000 のディスク復元をサポートできることです。そのため、ディスク イメージは大規模なスケーリングやサイズ変更を伴うシナリオに適しています。
GKE ノードの起動時に DaemonSet を使用してこのディスクをマウントします。
具体的な手順は次のとおりです。
前のセクションで説明したように、初期起動をスピードアップしてパフォーマンスを向上させるため、永続ディスクを GKE ノードにマウントして Stable Diffusion のランタイム ライブラリを配置します。
GKE イメージ ストリーミングとクラスタ オートスケーラーの利用
さらに、前述したように、GKE イメージ ストリーミングを有効にしてイメージの pull と読み込みのプロセスを高速化します。GKE イメージ ストリーミングの仕組みはこうです。ネットワーク マウントを使用してコンテナのデータレイヤーを containerd にアタッチし、ネットワーク上、メモリ内、およびディスク上の複数のキャッシュ レイヤでそれをサポートします。イメージ ストリーミング マウントの準備が完了すると、コンテナサイズを問わずほんの数秒でコンテナの状態が ImagePulling から Running に遷移します。このようにして、コンテナ イメージからの必要なデータの転送とアプリケーションの起動が実質的に並列化されます。その結果、コンテナの起動時間が短縮され、自動スケーリングも高速になります。
また、クラスタ オートスケーラー（CS）機能も有効にします。これにより、リクエストが増加したときに GKE ノードを自動的にスケールアップできます。クラスタ オートスケーラーは自動スケーリングをトリガーし、追加リクエストの処理に必要なノードの数を決定します。クラスタ オートスケーラーが新しいスケーリング ウェーブを開始し、新しい GKE ノードがクラスタに登録されると、DaemonSet が起動してランタイム依存関係を含むディスク イメージのマウントを支援します。その後、Stable Diffusion デプロイメントは、HostPath を通じてこのディスクにアクセスします。さらに、クラスタ オートスケーラーの optimize-utilization プロファイルを使用しました。これはクラスタに余剰リソースを保持するよりも使用率の最適化を優先させる GKE CA のプロファイルで、スケーリング時間の短縮、費用の節約、マシン使用率の向上といった効果をもたらします。
最終結果
最終的な起動結果は次のとおりです。
時系列順では、以下のようになりました。
- スケーリング用のクラスタ オートスケーラーのトリガー: 38 秒
- Node の起動と Pod のスケジューリング: 89 秒
- PVC のマウント: 4 秒
- イメージ pull の開始: 10 秒
- イメージの pull: 1 秒
- Pod の起動: 1 秒
- sd-webui サービスが提供可能になる（おおよその時間）: 65 秒
全体的に見ると、新しい Stable Diffusion コンテナ インスタンスを起動して新しい GKE ノードでのサービングを開始するまでに約 3 分かかりました。以前の 12 分と比較すると、起動速度の大幅な改善によってユーザー エクスペリエンスが向上したのは明らかです。
コード全体はこちらをご覧ください。https://github.com/nonokangwei/Stable-Diffusion-on-GCP/tree/main/Stable-Diffusion-UI-Agones/optimizated-init
考慮事項
上述の手法では、依存関係を分離することでコンテナサイズを縮小し、ライブラリを PD のディスク イメージから読み込めるようにしましたが、考慮すべきデメリットもあります。すべてを 1 つのコンテナに収めることには、アーティファクトが単一の変更不可なものとなり、そのバージョン管理ができるという利点があります。ベース環境をランタイム依存関係から分離すると、メンテナンスおよび更新するアーティファクトが複数存在することになります。この問題は、PD のディスク イメージの更新を管理するツールを作成することで軽減できます。
- Google Cloud、アプリケーション モダナイゼーション スペシャリスト Fan Liu
