GKE Agent Sandbox:为智能体强化学习加速

Agentic RL 规模化之后,瓶颈往往不在 GPU,而在 GPU 前面的沙箱。模型在 GPU 上生成代码,然后要放到 CPU 沙箱里执行、拿奖励信号。如果沙箱冷启动要几十秒,几千个任务同时拉起时镜像拉取和调度排队会把集群拖死,加速器只能空转。Google Cloud 今天发布的 GKE Agent Sandbox 就是冲着这个问题去的,配套的还有 Agent Sandbox RL orchestration SDK,以及针对 Gymnasium、NVIDIA NeMo Gym、OpenHands 的原生集成。

原文把 agentic RL 的基础设施压力拆成三个具体问题。第一是加速器闲置成本:RL 是同步批处理,一个 batch 里最慢的沙箱决定整个训练步的进度,冷启动和尾延迟直接变成 GPU 空等时间。第二是镜像基数爆炸:SWE-bench 或 R2E-Gym 这类基准里,每个任务往往对应一个完全不同的 OCI 镜像,一次 run 要管理几千个几 GB 的大镜像,远超本地磁盘容量。第三是控制面饱和:agentic RL 是突发性批量负载,比如一个公开数据集有 4,578 个 R2E 镜像,每个镜像 4 个 rollout,就是 18,312 个任务同时要 18,312 个沙箱,标准 Kubernetes 控制面在这种压力下会出现 API-server 队列瓶颈、pod 状态错误,甚至误触发节点驱逐。

GKE Agent Sandbox 的核心机制是内置的 SandboxWarmPool:预先初始化并维持一批健康沙箱,消除冷启动开销。它和 GKE Image Streaming 配合,能支撑上千个超过 1.2GB 的大镜像,同时把 TTFC(time-to-first-command)压到极低。快照、挂起、恢复能力支持错误恢复检查点,沙箱分叉则支持并行 agent 分支探索。SDK 提供异步 Python API,研究者不用写 Kubernetes YAML,也不用往镜像里塞自定义 daemon。

基准数据来自一个 10 节点的 gVisor 沙箱池,跑了两个负载:500 镜像的 SWE-bench 环境,以及更苛刻的 4,578 镜像 R2E 语料库(本地磁盘放不下)。并发规模从 500 个任务一直到 18,312 个。结果:TTFC 平均值从原生 Kubernetes 的 44–85 秒降到 1.1–8.8 秒,约 10x 提升;最坏情况尾延迟从 7.5 分钟降到 10 秒以内,约 45x 提升。控制面压力方面,在 4,578 镜像 × 4 rollouts 的测试中,pod 创建数从 18,312 降到 5,869,控制面 churn 减少 3.1x。

性能提升主要来自两个设计。一是把镜像水合移出关键路径:SDK 提前规划镜像在节点上的放置,Image Streaming 和 warm-pooling 在 rollout 开始前就把镜像准备好,而不是训练中途现场拉取造成 I/O 争用。二是给控制面限速:Agent Sandbox Controller v1.0.0 增加了速率控制,限制 warm pool 的补充速度,让 etcd 和 API server 在突发流量下保持稳定。这些调优参数、基准测试工具和负载测试都包含在 agent-sandbox-rl 示例仓库里,可以自己在集群上复现对比。

代价也有:提前预热环境会占用便宜的 CPU 和后台集群时间,但换来的是昂贵的加速器不再空转,这笔账对 RL 训练来说是划算的。另一个细节是 SDK 会把每个失败的沙箱计为可重试,而不是静默丢弃——因为丢掉的 rollout 不只是损失一次计算,还会引入奖励偏差。Mistral AI 的研究工程师 Jean-Malo Delignon 在原文中引用说,他们用 GKE Agent Sandbox 在单集群上处理过超过 30,000 个沙箱的峰值,并跨集群编排数十万个安全环境。

Accelerate agentic RL with GKE Agent Sandbox

查看原文