GPUスナップショットでgVisorのコールドスタートを削減
本記事では、CerebriumがGPUメモリのチェックポイント技術を用いて、gVisorサンドボックス内のGPUワークロードのコールドスタート時間を50秒から2.25秒まで短縮した方法を解説。初期化処理は決定的であり、一度実行した結果を凍結し、必要なときに復元するアイデア。実装にはgVisor containerd shimを修正し、コンテナ作成時に通常起動かチェックポイント復元かを判断。タイミング、ネットワーク状態、マルチプロセッシング、ファイルシステム、ストレージパフォーマンスなど多くのエッジケースに対処。
GPU加速AI推論において、コンテナのコールドスタートは顕著なボトルネックです。初期化プロセスには通常、Pythonモジュールのインポート、PyTorchのロード、モデルウェイトのGPUへのコピー、フレームワークのウォームアップパス（torch.compile、CUDAグラフキャプチャ、KVキャッシュ初期化など）が含まれます。これらの各段階は決定的です。同じ入力は毎回同じ出力を生成します。しかし、スケールアップのたびに既知の結果を再計算するコストを支払っています。
Cerebriumが提案する解決策はGPUメモリチェックポイントです。アイデアは単純です。高価な起動作業を一度実行し、結果を凍結し、必要に応じて復元します。具体的には、すべてのアプリケーションプロセスとGPU作業を一時停止し、CPUとGPUのメモリ状態をファイルにシリアライズし、高速な永続ストレージにアップロードします。復元は逆のプロセスです。チェックポイントファイルをダウンロードし、CPUとGPUメモリを再構築し、移動に耐えられない状態を修復し、ワークロードを再開します。復元されたアプリケーションプロセスは、以前に凍結したものと同じランタイムです。PyTorchは既にインポートされ、モデルウェイトはGPU上に存在し、カーネルはコンパイル済みで、トラフィックを処理する準備ができています。
この機能を実現するために、Cerebriumはコンテナランタイムとワークロードを実行するサンドボックスの間を拡張しました。2つのコンポーネントを追加しました。各ホストで動作する小さなチェックポイントサービス（ダウンロード、アップロード、キャッシュ、削除を担当）と、修正されたgVisor containerd shimです。shimはコンテナ作成をインターセプトし、チェックポイントが復元可能かどうかを確認し、通常起動か復元かを選択します。
重要な技術的課題はタイミングでした。containerdのサンドボックス起動シーケンスは固定されています。Sandbox Create → Sandbox Start → Container Create → Container Startです。復元を判断する自然な場所はSandbox Startですが、その時点ではコンテナイメージ情報が得られません。Cerebriumは起動シーケンスを若干並べ替えました。Sandbox Start時に実際の起動を遅延させ、Container Create段階でイメージ情報が利用可能になった時点で、通常起動か復元かを決定します。この小さな調整により、チェックポイント機能がユーザー透過的になりました。
開発とテスト中に、チームは多くのエッジケースに遭遇しました。チェックポイント中のTCPネットワークスタックの競合状態によるネットワーク障害、チェックポイントに数秒以上かかるとcontainerd内のgVisorがクラッシュする問題、NVIDIA GPUのコンテナデバイスインターフェース注入のサポート、ストレージレイヤーのパフォーマンス依存性などです。
ストレージとネットワークパスはシステム全体で最も重要な設計上の決定です。9GiBのチェックポイントの場合、g5.12xlargeインスタンスでは、完全なvLLMコールドスタートに約50秒かかりますが、チェックポイントからの復元はS3から2.25秒、ローカルNVMeから9秒に短縮されました。S3は十分な速度があり、Cerebriumがサポートするクラウドやリージョン間で移植可能なため、デフォルトの復元パスとして使用されます。
実際のワークロードは常にスナップショットに完璧に適合するとは限りません。ネットワーク状態は明白な例です。復元後、TCP接続は終了し、コンテナに異なる外部IPが割り当てられる可能性があり、外部IPに依存するフレームワーク（vLLMなど）の内部通信が壊れます。修正として、VLLM_HOST_IP=127.0.0.1を使用して内部通信をループバックにバインドします。
マルチプロセッシングも問題を引き起こします。Pythonフレームワークがforkを使用してGPUワーカーを作成すると、親プロセスからNVIDIAドライバのファイルディスクリプタを継承し、チェックポイントシステムがどのプロセスが実際にGPU状態を所有しているかを判断しにくくなります。解決策は、spawnでワーカーを作成することです（VLLM_WORKER_MULTIPROC_METHOD=spawn）。
ローカルランタイムファイル（Unixソケット、一時ファイル、ロックファイル）がチェックポイントに含まれていない場合、復元後にプロセスがファイルを見つけられずに失敗します。vLLMでは、VLLM_RPC_BASE_PATH=/run/cuda-ckptを使用して重要なRPC状態を保存されたパスに移動します。
チェックポイントのタイミングも重要です。CUDA作業が完了し、アプリケーションが既知のアイドル状態にあるときに取得する必要があります。vLLMのスリープモードは、チェックポイント前にKVキャッシュなどの一時的な状態を破棄し、スナップショットサイズを削減し、復元パフォーマンスを向上させるのに役立ちます。ユーザーはこれらの状態を保持するか、破棄するかを選択できます。
最後に、GPUメモリチェックポイントは作成時の環境に強く結びついています。GPUタイプ、CPUアーキテクチャ、マシンタイプ、ドライバ/ランタイムの互換性、gVisorバージョンが一致する必要があります。そのため、チェックポイントはアプリケーションだけでなく互換性によってキー付けされます。
要約すると、GPUメモリチェックポイントは単に「メモリをダンプしてリロードする」だけではありません。凍結できる状態と、再作成、再接続、または移動する必要がある状態を注意深く分離する必要があります。Cerebriumの実践は、慎重な設計により、この技術がGPUコンテナのコールドスタート時間を大幅に削減し、AI推論の迅速な弾性スケーリングを可能にすることを示しています。