利用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推理帶來更快的彈性伸縮能力。