GPUスナップショットでGVisorのコールドスタートを削減

メモリ・スナップショットによるGPUコールドスタートの削減: CUDAワークロードを数秒で復元 プロダクション環境でAIモデルを実行している場合、望むと望まざるとにかかわらず、コールドスタートとは切っても切れない関係にあります。3分間の起動時間は、スケーリング方法を変えます。解放できたはずのGPUをウォームアップし続けます。ユーザーを待たせないために過剰にプロビジョニングします。スケールダウンが早すぎると次のスパイク時に問題が生じるため、クールダウン期間を長くします。アプリケーションは、1つの問題、すなわちトラフィックを処理するためにモデルを十分に迅速に準備するという問題を中心に複雑さを増していきます。 Cerebriumでは、設立当初からコールドスタートの問題にこだわってきました。そのこだわりが、インフラのほぼすべてのレイヤーを再考させることにつながりました。 ノードのスケールアップを高速化するためのカスタムVMイメージ コンテナイメージのコールドスタートを数秒で実現するためのカスタムイメージランタイム リージョン間およびクラウド間でワークロードをルーティングするための高可用性、低遅延オーケストレーター 完全にウォームアップされたコンテナを数秒で復元するためのCPUおよびGPUメモリ・スナップショット より多くの企業が大規模なカスタムAIモデルをプロダクション環境に移行するにつれて、同じ壁にぶつかります。当社の顧客は、大規模言語モデル、リアルタイムアバター、文字起こしモデル、拡散モデル、その他のGPU負荷の高いワークロードを実行しており、起動時間は数秒から5分以上かかることがあります。この時間のほとんどは、コンテナをリクエスト処理可能な状態にするための作業に費やされます。ライブラリのインポート、モデルウェイトのロード、CUDAの初期化、カーネルのコンパイル、ランタイムのウォームアップなどです。これがチェックポイントが解決する核心的な問題です。新しいコンテナが起動するたびに同じランタイムをゼロから再構築する代わりに、完全に初期化されたコンテナ（CPUメモリ、GPUメモリ、プロセス状態、モデルウェイト、コンパイル済みカーネルを含む）をスナップショットし、それを新しいコンテナに短時間で直接復元します。一部のワークロードでは、これによりコールドスタート時間が80%以上短縮されます！ この投稿では、CerebriumでCPUおよびGPUメモリのチェックポイントを構築した方法、高度にカスタマイズされたgVisorベースのランタイム内でどのように機能するか、およびvLLMのような実際のCUDAワークロードを確実に迅速に復元するために何が必要だったかを説明します。 実際、時間はどこに消えるのか？ コールドスタートを単に「イメージのプル」、つまりコンテナを実行するマシンにアプリケーションイメージをダウンロードすることだと考えがちです。しかし、AIワークロードの場合、それはモデルをトラフィック処理可能な状態にするための最初の部分に過ぎず、もはやボトルネックではありません。コンテナのダウンロード問題はすでに解決済みです。CPUまたはGPUコンテナの実際のコストは、イメージがマシンに配置され、アプリケーションが初期化を開始した後に起こるすべてのことです。その初期化パスには、Pythonモジュールのインポート、PyTorchのロード、モデルウェイトのアセンブリ、それらをGPUへのコピー、およびフレームワークのウォームアップパス（torch.compile、CUDAグラフキャプチャ、KVキャッシュ初期化、およびサービングスタックがトラフィックを受け入れる前に必要なその他のすべて）の実行が含まれます。 これらのすべてのステージは決定論的です。PyTorchのインポートは毎回同じロードされたモジュールを生成します。モデルを構築し、ウェイトをGPUにコピーすると、毎回GPUメモリに同じバイトが生成されます。torch.compileとCUDAグラフキャプチャは毎回同じカーネルを生成します。しかし、スケールアップのたびに、既知の結果を再計算するためにコストを支払っています。それがチェックポイントが変えることです。アイデアは単純です。高価なスタートアップ作業を一度行い、その結果をフリーズし、オンデマンドで復元します。具体的には、チェックポイントを取得するとは、次のことを意味します。 実行を一時停止：すべてのアプリケーションプロセス、スレッド、そして重要なGPU作業を一時停止します。 メモリをダンプ：CPUとGPUの両方からインメモリ状態をファイルにシリアル化します。 アップロード：それらのファイルを高速で耐久性のあるストレージにプッシュします。 復元は同じプロセスを逆に行います。チェックポイントファイルをプルダウンし、CPUとGPUのメモリを再構築し、移動に耐えられない状態の断片を修復し、ワークロードを一時停止解除します。復元されたアプリケーションプロセスは、以前にフリーズしたのと同じウォームアップされたランタイムです。PyTorchはすでにインポートされており、モデルウェイトはすでにGPUに常駐しており、カーネルはすでにコンパイルされており、アプリケーションはトラフィックを処理する準備ができています。精神的なモデルは単純です。しかし、実際のGPUワークロードでそれを確実に機能させることは簡単ではありません。 高レベルアーキテクチャ 高レベルでは、チェックポイントはコンテナの完全なライフサイクルを制御できる唯一の場所、すなわちコンテナランタイムとワークロードを実行するサンドボックスの間に存在する必要があります。Cerebriumは、分離のためにユーザーワークロードをgVisorサンドボックス内で実行します。チェックポイントをサポートするために、そのランタイムパスを拡張し、コンテナが起動する際に通常のブートシーケンスが完了する前に決定を下せるようにしました。 このコンテナはゼロから開始すべきか、それともチェックポイントから復元すべきか？ チェックポイントが存在しない場合、コンテナは通常のパスに従います。イメージが起動し、アプリケーションがブートし、モデルがロードされ、GPUメモリが投入され、ワークロードが準備完了になります。コンテナが完全にウォームアップされると、ユーザーはチェックポイントをトリガーできます。その時点で、ワークロードを一時停止し、そのCPUおよびGPU状態をキャプチャし、チェックポイントをディスクに書き込み、高速ストレージにアップロードします。 チェックポイントが存在する場合、通常の起動パスはスキップされます。コンテナを起動し、Pythonのインポート、モデルのロード、GPU転送、torch.compile、CUDAグラフキャプチャを待つ代わりに、保存された状態をサンドボックスに直接復元します。プロセスは、ウォームアップを終えたばかりのように再開されます。これは単純に聞こえますが、ランタイムが適切なタイミングでいくつかの質問に答える必要があります。 どのワークロードが開始されているか？ このイメージ、GPUタイプ、マシンタイプ、およびランタイムバージョンに対応するチェックポイントが存在するか？ チェックポイントはどこに保存されているか？ チェックポイントはすでにホストにローカルでキャッシュされているか？ 復元すべきか、それともクリーンなブートにフォールバックすべきか？ これを機能させるために、ノードランタイムに2つのコンポーネントを追加しました。1つ目は、すべてのホストで実行される小さなチェックポイントサービスです。これはチェックポイントの運用面を処理します。チェックポイントのダウンロード、新しいチェックポイントのアップロード、ローカルでのキャッシュ、古いまたは破損したチェックポイントの削除、および復元ステータスの報告などです。2つ目は、変更されたgVisor containerd shimです。これはコンテナの起動パスに位置する部分です。コンテナの作成を傍受し、チェックポイントが復元できるかどうかをチェックし、通常のブートフローを続行するか、そのフローを復元で置き換えます。言い換えれば、チェックポイントサービスはスナップショットファイルを移動および管理します。シムは、新しいコンテナが通常にブートすべきか、それともスナップショットからウェイクアップすべきかを決定します。 最も困難だったのは、これら2つのコンポーネント間のAPIではありませんでした。タイミングでした。Containerdは、固定されたシーケンスでサンドボックスを起動します。 サンドボックス作成 → サンドボックス開始 → コンテナ作成 → コンテナ開始 復元するかどうかを決定する自然な場所は、サンドボックスが開始されるときです。しかし、その時点では、コンテナイメージに関する十分な情報がまだなく、チェックポイントが存在するかどうかを知ることができません。イメージ情報は、後でコンテナ作成時にのみ利用可能になります。そこで、起動シーケンスをわずかに変更する必要がありました。containerdがサンドボックスを開始するように要求したとき、実際の開始を延期します。想定されるステータス応答でcontainerdを満足させますが、コンテナ作成まで実際のサンドボックス起動を遅らせます。一度、どのイメージが起動されているか、および一致するチェックポイントが存在するかどうかがわかります。その時点で、次の2つのパスのいずれかを選択します。 通常のブート：サンドボックスを起動し、コンテナを起動し、アプリケーションを初期化させ、必要に応じてウォームアップ後にチェックポイントを作成します。 チェックポイント復元：チェックポイントをダウンロードまたは特定し、CPUおよびGPUメモリをサンドボックスに復元し、移動に耐えられないランタイム状態を修復し、プロセスを再開します。 作業は、ランタイムがすでに実行するであろう作業とほとんど同じです。主な変更点は、復元決定をサンドボックス開始からコンテナ作成に移動したことです。コンテナ作成時には、イメージ情報が最終的に利用可能になり、一致するチェックポイントが存在するかどうかを判断できます。このわずかな順序変更が、チェックポイントがユーザーの視点から透過的に感じられるようにするものです。ユーザーは同じ方法でワークロードを開始し、b