利用 GPU 快照减少 GVisor 冷启动

# 利用内存快照减少 GPU 冷启动：数秒内恢复 CUDA 工作负载 来源：https://cerebrium.ai/blog/reducing-gpu-cold-starts-with-memory-snapshots-restoring-cuda-workloads-in-second 利用内存快照减少 GPU 冷启动：数秒内恢复 CUDA 工作负载 如果你在生产环境中运行 AI 模型，无论你愿不愿意，你都会与冷启动打交道。 三分钟的启动时间会改变你的扩容方式。你会让那些本可以释放的 GPU 保持温暖。你会过度配置，以避免用户等待。你会延长冷却期，因为过快缩容会在下一次流量高峰时带来痛苦。应用程序开始围绕一个问题积累复杂性：让模型准备好提供服务流量，并且要足够快。 在 Cerebrium，我们从第一天起就痴迷于冷启动问题。这种痴迷促使我们重新思考了基础设施的几乎每一层： - 定制虚拟机镜像，加快节点扩容速度（https://cerebrium.ai/blog/achieving-speed-improvements-in-custom-container-images） - 定制镜像运行时，实现亚秒级容器镜像冷启动（https://cerebrium.ai/blog/rethinking-container-image-distribution-to-eliminate-cold-starts） - 高可用、低延迟的编排器，用于跨区域和云路由工作负载（https://cerebrium.ai/blog/thalamus-our-highly-available-distributed-router-for-global-realtime-ai-workloads） - CPU 和 GPU 内存快照，用于在数秒内恢复完全预热状态的容器 随着越来越多的公司将大型自定义 AI 模型投入生产，他们遇到了同样的障碍。我们的客户运行大型语言模型、实时虚拟形象、转录模型、扩散模型以及其他 GPU 密集型工作负载，这些工作负载的启动时间从几秒到五分钟以上不等。 大部分时间都花在了让容器准备好处理请求的工作上：导入库、加载模型权重、初始化 CUDA、编译内核以及预热运行时。这正是检查点（checkpoint）要解决的核心问题。我们不是每次新容器启动时都从头重建相同的运行时，而是对完全初始化后的容器进行快照——包括 CPU 内存、GPU 内存、进程状态、模型权重和编译好的内核——然后直接将其恢复到新容器中，所需时间仅为原先的一小部分。 对于某些工作负载，这能将冷启动时间减少 80% 以上！ 这篇文章将解释我们如何在 Cerebrium 构建 CPU 和 GPU 内存检查点功能，它如何在我们高度定制的基于 gVisor 的运行时内部工作，以及要使像 vLLM 这样的真实 CUDA 工作负载能够可靠且快速地恢复，我们需要做哪些工作。 ### 时间到底花在了哪里？ 人们很容易将冷启动简单地理解为“拉取镜像”：将应用程序镜像下载到要运行容器的那台机器上。但对于 AI 工作负载来说，这只是让模型准备好提供服务流量的第一步，而且它不再是瓶颈。我们早已解决了容器下载问题（https://cerebrium.ai/blog/rethinking-container-image-distribution-to-eliminate-cold-starts）。在 CPU 或 GPU 容器中，真正的成本在于镜像下载到机器上之后，应用程序开始初始化时所发生的一切。 这个初始化路径包括导入 Python 模块、加载 PyTorch、组装模型权重、将其复制到 GPU 上，以及运行框架的预热路径——torch.compile、CUDA 图捕获、KV 缓存初始化，以及服务栈在能够接收流量之前所需的任何其他操作。 这些阶段中的每一个都是确定性的。 每次导入 PyTorch 都会产生相同加载的模块。构建模型并将权重复制到 GPU 上，每次都会在 GPU 内存中产生相同的字节。torch.compile 和 CUDA 图捕获每次都会产生相同的内核。 然而，在每次扩容时，我们都在为计算一个已知的结果而付出代价。 这正是检查点技术要改变的地方。 其想法很简单：执行一次昂贵的启动工作，冻结结果，然后按需恢复。 具体来说，创建检查点意味着： 1. **暂停执行：** 暂停所有应用程序进程、线程，并且至关重要的是，暂停 GPU 工作。 2. **转储内存：** 将 CPU 和 GPU 的内存状态序列化到文件中。 3. **上传：** 将这些文件推送到快速、持久的存储中。 恢复则逆向执行同样的过程。我们拉取检查点文件，重新填充 CPU 和 GPU 内存，修复那些无法在移动中幸存的状态片段，然后取消暂停工作负载。 恢复后的应用程序进程与我们之前冻结的、已经预热好的运行时一模一样：PyTorch 已经导入，模型权重已经驻留在 GPU 上，内核已经编译完毕，应用程序已经准备好处理流量。 这个思维模型很直接。但要使其在实际 GPU 工作负载上可靠运行，却并不简单。 ### 高层架构 从高层看，检查点功能需要位于能够控制容器完整生命周期的唯一位置：在容器运行时和运行工作负载的沙箱之间。 Cerebrium 在 gVisor（https://gvisor.dev/）沙箱内运行用户工作负载，以实现隔离。为了支持检查点，我们扩展了这个运行时路径，以便在容器启动时，我们可以在正常的启动序列完成之前做出一个决定： **这个容器应该从头启动，还是应该从检查点恢复？** 如果不存在检查点，容器将遵循正常路径。镜像启动，应用程序启动，模型加载，GPU 内存填充，工作负载变得就绪。一旦容器完全预热，用户可以触发检查点。此时，我们暂停工作负载，捕获其 CPU 和 GPU 状态，将检查点写入磁盘，然后上传到快速存储。 如果存在检查点，我们则跳过正常的启动路径。我们不是启动容器并等待 Python 导入、模型加载、GPU 传输、torch.compile 和 CUDA 图捕获，而是将保存的状态直接恢复到沙箱中。进程恢复运行，就像刚刚完成预热一样。 这听起来很简单，但需要运行时在恰当时机回答几个问题： - 正在启动的是哪个工作负载？ - 对于这个镜像、GPU 类型、机器类型和运行时版本，是否存在兼容的检查点？ - 检查点存储在哪里？ - 检查点是否已经缓存在本地主机上？ - 我们应该恢复，还是回退到干净的启动？ 为了实现这一点，我们在节点运行时中添加了两个组件。 第一个是一个小的检查点服务，运行在每个主机上。它处理检查点的操作方面：下载检查点、上传新检查点、本地缓存、驱逐过期或损坏的检查点，以及报告恢复状态。 第二个是一个修改过的 gVisor containerd shim。这是位于容器启动路径中的组件。它拦截容器创建，检查是否可以恢复检查点，然后要么继续正常的启动流程，要么用恢复操作替换该流程。 换句话说，检查点服务负责移动和管理快照文件。而 shim 则决定新容器是应该正常启动，还是从快照中唤醒。 最难的部分不是这两个组件之间的 API，而是时机。 *Containerd* 通过一个固定的序列来启动沙箱： 沙箱创建 → 沙箱启动 → 容器创建 → 容器启动 决定是否恢复的自然时机是沙箱启动时。但在那个时刻，我们还没有足够的信息关于容器镜像来判断是否存在检查点。镜像信息要等到稍后的容器创建阶段才能获得。 因此，我们不得不稍微调整启动顺序。 当 *containerd* 要求我们启动沙箱时，我们推迟了真正的启动。我们让 containerd 满意于预期的状态响应，但将实际的沙箱启动延迟到容器创建阶段，那时我们才知道正在启动的是哪个镜像，以及是否存在匹配的检查点。 这时，我们选择以下两条路径之一： 1. **正常启动：** 启动沙箱，启动容器，让应用程序初始化，并在预热后选择性地创建检查点。 2. **检查点恢复：** 下载或定位检查点，将 CPU 和 GPU 内存恢复到沙箱中，修复那些无法在移动中幸存的运行时状态，然后恢复进程。 这些工作大部分与运行时本已要做的相同。关键的变化是，我们将恢复决策从沙箱启动移到了容器创建阶段，此时镜像信息终于可用，我们可以判断是否存在匹配的检查点。 这个小小的重新排序使得检查点功能在用户看来是透明的。他们以同样的方式启动工作负载，但一旦存在检查点，未来的扩容将恢复已预热好的进程，而不是从头重建。 在我们测试和开发这个功能的过程中，我们遇到了几个从现有文档中并不明显的边缘情况。在可能的情况下，我们正在努力将这些修复推送到上游，以便下一个采用此技术的团队不必重新发现相同的问题。 我们发现的一些问题包括： - TCP 网络栈中的一个竞态条件，导致容器在检查点过程中接收到大量数据包时网络无法工作。 - 一个竞态条件，导致 gVisor 在 *containerd* 中运行时，如果检查点花费时间超过几秒就会崩溃。 - 支持用于 NVIDIA GPU 的容器设备接口（Container Device Interface）注入。 **检查点分发：为什么存储层比你想象的要重要得多** 一个预热好的 GPU 容器的检查点是很大的——我们的一个测试工作负载大约 9 GiB，然而使用 vLLM 恢复 Deepseek V4 FP8 将达到 640GB。只有当移动如此大量数据的速度快于容器自身冷启动所需时间时，恢复才有价值。这使得存储和网络路径成为整个系统中唯一最重要的设计决策。 这个数学计算很残酷： 对于我们的 9GB 容器，在 g5.12xlarge 上，完整的 vLLM 冷启动大约需要 50 秒。从 9 GiB 的检查点恢复，从 S3 恢复需要 2.25 秒，从本地 NVMe 恢复需要 9 秒。 我们使用 S3 作为默认恢复路径，因为它速度足够快，并且可跨 Cerebrium 支持的所有云和区域移植。当检查点已缓存在节点上时，本地 NVMe 速度很快，而对象存储仍然作为持久的真实数据源。 这些结果特定于 g5.12xlarge。在具有更高网络带宽或更快本地存储的节点上，恢复时间会进一步缩短。 **难点：真实的工作负载很混乱** 当工作负载的状态自包含于内存中时，检查点最容易。但真实的 GPU 工作负载很少如此干净。 快照可以保留预热好的运行时，但不能盲目地保留围绕它的每个外部依赖项。恢复后，应用程序可能仍然持有对文件系统路径、套接字、IP 地址、设备句柄或驱动程序状态的引用，这些在移动前有效，但移动后无效。这就是大多数棘手行为的来源。 网络状态是第一个明显的例子。打开的 TCP 连接与原始运行时环境绑定。恢复后，这些连接已终止，并且容器可能具有不同的外部 IP。这会破坏那些使用容器外部 IP 进行内部心跳、工作节点协调或控制平面通信的框架。例如，在 vLLM 中，这意味着进程可能成功恢复，但由于运行时的某些部分试图通过一个不再有效的地址进行通信而仍然在内部失败。解决方法是将内部框架通信绑定到回环地址，使用 `VLLM_HOST_IP=127.0.0.1`，这样工作节点协调就不再依赖于分配给容器的外部 IP。 多进程处理带来了另一类问题。许多 Python 服务框架使用工作进程，如果这些工作进程是通过 fork 创建的，它们可能会从父进程继承 NVIDIA 驱动程序的文件描述符。这一点很重要，因为检查点系统需要清楚地知道哪些进程实际上拥有 GPU 状态。泄露的驱动程序文件描述符可能使运行时认为 GPU 仍被某些进程占用，而这些进程不应该阻塞检查点，或者导致难以推理的恢复行为。对于 vLLM，解决方法是使用 spawn 而非 fork 来创建 GPU 工作进程，通过 `VLLM_WORKER_MULTIPROC_METHOD=spawn`，这样子进程会干净地启动，而不是从父进程继承 GPU 驱动程序状态。 本地运行时文件是另一个微妙的边缘情况。框架经常在本地磁盘上创建 Unix 套接字、临时文件、锁文件和协调状态。如果本地文件系统没有随检查点一起恢复，进程可能会在预期某些文件存在但实际已不存在的情况下启动。这是最烦人的故障模式之一，因为进程从外部看可能看起来健康，而工作进程却在静默地内部通信失败。在 vLLM 中，我们通过使用 `VLLM_RPC_BASE_PATH=/run/cuda-ckpt` 将恢复关键的 RPC 状态移动到一个小的保留路径来解决这个问题。 检查点的时机也很重要。检查点需要 CPU 和 GPU 内存的一致性视图。如果在拍摄快照时 CUDA 工作仍在运行，则检查点可能不一致或无法安全恢复。在实践中，这意味着检查点必须在工作负载完成预热并达到已知的空闲状态后进行。对于某些框架，这需要一个明确的就绪步骤：加载模型，运行预热过程，等待编译或 CUDA 图捕获完成，然后才触发检查点。 另一个优化是决定什么不应该被检查点化。vLLM 的睡眠模式在这里很有用，因为它可以在拍摄检查点之前丢弃瞬态状态，如 KV 缓存。KV 缓存可能很大，保留它会使检查点更大、上传更慢、恢复也更慢。对于许多工作负载，缓存不值得跨恢复携带，因为它与特定请求相关，并且可以在流量恢复后自然重建。在这些情况下，在检查点之前将 vLLM 置于睡眠模式可以显著减少快照大小并提高恢复性能。我们将其作为一种选择而非强制行为：用户可以决定是希望在恢复之间保留该状态，还是丢弃它以加快检查点速度。 最后一个限制是兼容性。GPU 内存检查点不像容器镜像那样是一个可移植的工件。它与创建它的环境紧密相关：GPU 类型、CPU 架构、机器类型、驱动程序/运行时兼容性以及 gVisor 版本。在一个硬件和运行时配置上创建的检查点不能安全地恢复到任意另一个配置上。因此，我们根据兼容性（而不仅仅是应用程序）来组织检查点。恢复路径仅在目标环境与原始检查点环境匹配时才使用检查点。 更大的模式是，GPU 内存检查点不仅仅是“转储内存并重新加载”。它关乎于分离那些可以冻结的状态和那些必须重新创建、重新连接或移动到检查点安全位置的状态。 这也是为什么该功能是选择加入并且是工作负载感知的。不同的服务栈依赖于不同的文件系统、套接字、设备句柄、网络假设和框架内部机制。要使检查点达到生产就绪状态，需要明确验证这些假设，而不是将每个 GPU 工作负载都视为可以完全相同的方式暂停、移动和恢复。 ### 结果：冷启动减少 71% 我们在六个具有不同功能的工作负载上，将 Cerebrium 与 Baseten 和 Modal 进行了基准测试。对于每个工作负载，