利用GPU快照減少gVisor冷啟動時間

在GPU加速的AI推理場景中，容器冷啟動是一個顯著的瓶頸。初始化過程通常包括匯入Python模組、載入PyTorch、組裝模型權重並複製到GPU、執行框架預熱路徑（如torch.compile、CUDA圖捕獲、KV快取初始化）等步驟。每一個階段都是確定性的：相同的輸入每次產生相同的輸出。然而，每次擴容時，系統都在重複計算已知的結果。

Cerebrium提出的解決方案是GPU記憶體快照（checkpointing）。其核心思想很簡單：執行一次昂貴的啟動工作，將結果凍結，然後在需要時按需恢復。具體操作包括：暫停所有應用程序和GPU工作，將CPU和GPU記憶體狀態序列化到檔案，並將檔案上傳到快速持久化儲存。恢復則是逆向過程：下載檢查點檔案，重新填充CPU和GPU記憶體，修復無法在遷移中存活的狀態，然後解凍工作負載。恢復後的應用程序與之前凍結的執行時完全相同：PyTorch已匯入、模型權重已在GPU上、核心已編譯，準備直接處理流量。

為了實現這一功能，Cerebrium在容器執行時和執行工作負載的沙箱之間進行了擴充套件。他們新增了兩個元件：一個執行在每個主機上的小型檢查點服務，負責下載、上傳、快取和驅逐檢查點；一個修改後的gVisor containerd shim，它攔截容器建立過程，檢查是否可以恢復檢查點，並選擇正常啟動或恢復路徑。

關鍵的技術挑戰在於時序。containerd啟動沙箱的固定序列是：Sandbox Create → Sandbox Start → Container Create → Container Start。自然的選擇是在Sandbox Start時決定是否恢復，但此時還無法獲得容器映象資訊。Cerebrium重新排序了啟動序列：在Sandbox Start時延遲實際啟動，直到Container Create階段獲取映象資訊後，再決定是正常啟動還是恢復檢查點。這個小的調整使得檢查點功能對使用者透明。

在開發和測試中，團隊遇到了許多邊緣情況。包括：檢查點期間TCP網路棧的競態條件導致網路中斷；如果檢查點耗時超過幾秒，containerd內的gVisor會崩潰；需要支援NVIDIA GPU的容器裝置介面注入；以及檢查點分發依賴儲存層效能。

儲存和網路路徑是整個系統最重要的設計決策。對於一個9 GiB的檢查點，在g5.12xlarge例項上，完整的vLLM冷啟動約需50秒，而恢復檢查點只需2.25秒（從S3）或9秒（本地NVMe）。S3作為預設恢復路徑，因為速度足夠且跨雲和區域可移植。

實際工作負載並非總能完美適應快照。網路狀態是一個明顯的例子：恢復後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推理帶來更快的彈性伸縮能力。