Google AX 开源:像调度 Pod 一样调度 AI Agent

发布时间:2026/9/28 23:49:09
Google AX 开源:像调度 Pod 一样调度 AI Agent 1. 从 Pod 到 Agent一次调度范式的迁移Google 这次开源的 AX 项目核心命题其实一句话就能说清楚把 AI Agent 当成 Kubernetes 里的 Pod 来调度。这个思路第一次看到的时候我愣了一下因为过去两年我接触的 Agent 项目绝大多数都是一个进程跑一个 Agent自己管自己的生命周期调度这件事基本靠人肉或者简单的队列。而 AX 做的事情是把 Agent 的创建、分配、执行、回收这一整套流程塞进一个类似 K8s 调度器的框架里。先说清楚这个项目解决的是什么问题。假设你手上有几十上百个 Agent 任务有的要调大模型推理有的要跑工具调用有的只是做简单的数据搬运。如果每个 Agent 都自己起一个进程、自己抢资源很快就会遇到几个典型问题GPU 被某个长任务占死、短任务排队等到超时、某个 Agent 崩了没人知道、扩缩容全靠手动。这些问题的本质和当年容器编排出现之前运维面对的问题一模一样——资源没有统一视图调度没有统一入口。AX 的价值就在于它把这套已经被验证过的调度思路搬到了 Agent 场景。Pod 是 K8s 里最小的调度单元Agent 在 AX 里就是最小的调度单元。你可以给 Agent 声明资源需求需要多少显存、多少并发槽位、声明依赖关系这个 Agent 必须在那个 Agent 完成之后才能启动、声明重启策略失败了重试几次、超时多久杀掉。这些声明式的配置让 Agent 的管理从写代码控制变成了写配置描述。适合谁来参考这个项目我的判断是三类人。第一类是正在做 Agent 平台化的团队手上有多个 Agent 需要统一管理正在纠结要不要自己造调度轮子第二类是做 LLM 应用基础设施的工程师想理解 Agent 调度和传统任务调度到底差在哪第三类是对 K8s 调度机制熟悉、想把这套经验迁移到 AI 场景的运维同学。如果你只是写单个 Agent 做 demoAX 可能有点重但如果你已经在考虑Agent 多了怎么管这个问题那这个项目的设计思路值得花时间啃一啃。2. AX 的整体设计思路拆解2.1 为什么是像 Pod 一样而不是像 Job 一样这是我在读 AX 设计文档时第一个想搞清楚的问题。K8s 里既有 Pod 也有 JobJob 更适合跑一次性任务为什么 AX 选择对标 Pod 而不是 Job原因在于 Agent 的生命周期特征。一个 Agent 往往不是跑完就死的一次性任务它可能有多个交互轮次可能需要在执行过程中保持上下文可能被暂停后恢复。Job 的语义是确保完成 N 次而 Pod 的语义是维持这个实例的运行状态。Agent 更像后者——你需要的是一个持续存在、可以被调度、可以被观测、可以被替换的运行单元。另一个关键点是调度粒度。Job 的调度粒度是任务级别的一个 Job 对应一批 Pod。而 Agent 的调度粒度更细每个 Agent 实例可能对应不同的资源需求、不同的优先级、不同的亲和性规则。用 Pod 模型每个 Agent 就是一个独立的调度对象调度器可以针对每个 Agent 单独决策。这种细粒度控制在 Agent 场景下非常重要因为不同 Agent 的资源画像差异极大——一个做代码生成的 Agent 可能需要大显存一个做意图识别的 Agent 可能只需要少量 CPU。2.2 调度层的核心抽象AX 的调度层抽象了几个关键概念我按自己的理解梳理一下。Agent Spec这是描述一个 Agent 该怎么跑的核心配置。类比 Pod Spec它包含镜像Agent 运行时环境、资源请求CPU、内存、GPU、环境变量、启动命令、健康检查等。不同之处在于Agent Spec 还会包含模型相关的配置比如用哪个模型、推理参数是什么、工具集怎么挂载。Agent Scheduler调度器本身负责把 Agent 分配到合适的节点上。它的决策依据包括节点资源余量、Agent 的资源请求、亲和性与反亲和性规则、优先级等。这里有个细节值得注意——Agent 调度器需要考虑的不只是节点有没有资源还要考虑节点上的模型服务能不能支撑这个 Agent。比如一个节点上已经加载了某个模型把需要同款模型的 Agent 调度过去就能复用避免重复加载。Agent RuntimeAgent 实际运行的载体。它负责拉起 Agent 进程、注入配置、上报状态、处理生命周期事件。Runtime 和 Scheduler 之间通过一个类似 kubelet 的组件通信Scheduler 决定放哪Runtime 负责跑起来。Agent Controller控制器负责维持期望状态。你声明我要 3 个某类 Agent 在跑Controller 就确保实际运行数量匹配。Agent 崩了它负责重启节点挂了它负责重新调度。这套控制循环和 K8s 的 ReplicaSet 逻辑几乎一致。2.3 与传统任务调度的本质差异很多人第一反应是这不就是个任务调度器吗和 Airflow、DolphinScheduler 有什么区别。我一开始也这么想但仔细对比后发现差异挺大。传统任务调度器比如 DAG 调度的核心是依赖关系和执行顺序。它关心的是任务 A 完成后触发任务 B调度单位是任务生命周期是触发-执行-完成。而 AX 关心的是资源分配和实例维持调度单位是 Agent 实例生命周期是创建-运行-可能重启-销毁。举个具体例子。在 DolphinScheduler 里你定义一个工作流里面有 5 个任务调度器按依赖顺序触发它们每个任务跑完就结束。在 AX 里你声明我要 3 个 Agent 实例持续运行调度器负责把这 3 个实例分配到节点上某个实例挂了就补一个节点资源不够就排队等。前者是流程编排后者是资源编排。这个差异决定了 AX 更适合的场景是长期运行的 Agent 集群管理而不是一次性数据处理流水线。当然两者可以结合——你可以用传统调度器触发一个工作流工作流里的某个步骤是向 AX 提交一批 Agent 任务。3. 核心机制与实操要点解析3.1 Agent 的资源声明与调度决策Agent 的资源声明是调度的基础。在 AX 里你需要为每个 Agent 声明它需要多少资源。这里有个容易踩的坑Agent 的资源需求往往不是静态的。一个 Agent 在处理简单请求时可能只用 1GB 显存处理复杂请求时可能飙到 8GB。如果你按峰值声明资源利用率会很低如果按均值声明高峰期会 OOM。我的建议是采用分级声明的方式。把 Agent 按资源画像分成几档比如 small2GB 显存、medium8GB、large24GB每档对应不同的调度策略。small 档可以高密度部署large 档独占节点。这样既保证了资源利用率又避免了资源争抢。调度决策的流程大致是这样的Agent 提交后进入待调度队列调度器遍历可用节点对每个节点打分选最高分的节点绑定。打分维度包括打分维度权重建议说明资源余量高余量越充足得分越高避免碎片化模型亲和性高节点已加载所需模型则加分负载均衡中避免所有 Agent 挤在少数节点数据本地性中数据所在节点优先反亲和性视场景同组 Agent 分散到不同节点这套打分逻辑和 K8s 默认调度器非常像如果你熟悉 K8s 的调度框架上手 AX 会很快。3.2 Agent 的生命周期管理Agent 的生命周期比 Pod 复杂因为它多了推理执行这个阶段。一个 Agent 实例的完整生命周期包括Pending等待调度、Starting拉取运行时、加载模型、Ready就绪可接收任务、Running执行中、Paused暂停、Terminating终止中、Failed失败。这里的关键难点是就绪判断。Pod 的就绪探针通常是 HTTP 探针或命令探针判断进程是否活着。但 Agent 的就绪不只是进程活着还要确认模型加载完成、工具集初始化完成、上下文就绪。AX 的做法是让 Agent Runtime 暴露一个就绪接口Agent 自己上报就绪状态。这个设计很合理因为只有 Agent 自己知道什么时候真正准备好了。另一个难点是优雅终止。Agent 执行到一半被 kill 掉可能导致上下文丢失、任务状态不一致。AX 支持配置 terminationGracePeriod给 Agent 一段时间做清理。我的实操经验是这个时间要设得足够长尤其是涉及外部工具调用的 Agent因为工具调用可能正在等待外部响应强行中断会留下脏状态。3.3 配置示例与参数说明下面给一个 Agent Spec 的配置示例这是我根据 AX 的设计思路整理的实际使用时需要对照官方文档调整字段名。apiVersion: ax.io/v1 kind: Agent metadata: name: code-review-agent labels: tier: medium team: devtools spec: runtime: python-agent-runtime:v1.2 model: name: code-llm-7b endpoint: http://model-service:8080 maxTokens: 4096 resources: requests: cpu: 2 memory: 8Gi nvidia.com/gpu: 1 limits: cpu: 4 memory: 16Gi nvidia.com/gpu: 1 tools: - name: git-reader config: repoPath: /workspace/repo - name: linter config: rules: /etc/lint/rules.yaml healthCheck: readiness: path: /healthz/ready initialDelaySeconds: 30 periodSeconds: 10 liveness: path: /healthz/live initialDelaySeconds: 60 periodSeconds: 30 restartPolicy: OnFailure terminationGracePeriodSeconds: 120 affinity: nodeAffinity: preferredDuringScheduling: - weight: 80 preference: matchExpressions: - key: model-cached operator: In values: [code-llm-7b]几个参数值得单独说。initialDelaySeconds在 Agent 场景下要比普通 Pod 设得大因为模型加载慢30 秒起步是常态大模型可能要 120 秒。terminationGracePeriodSeconds建议设 120 秒以上给 Agent 足够时间保存上下文。restartPolicy推荐用 OnFailure 而不是 Always因为 Agent 正常完成后不应该被重启。3.4 调度器的扩展点AX 的调度器设计留了扩展点这点很关键。因为不同团队的 Agent 调度需求差异很大硬编码的调度逻辑很难覆盖所有场景。AX 支持自定义调度插件你可以在调度决策的各个阶段插入自己的逻辑。常见的扩展场景包括按成本调度优先用便宜的节点、按延迟调度优先用离用户近的节点、按合规调度数据不能出某个区域。这些需求在通用调度器里很难内置但通过插件机制就能灵活支持。我个人的经验是不要一上来就写调度插件。先用默认调度器跑一段时间观察实际的调度瓶颈在哪再针对性地写插件。很多团队一开始就想着我要自定义调度策略结果写出来的策略还不如默认的好用。4. 实操过程与核心环节实现4.1 环境准备与依赖安装假设你已经有一个 K8s 集群部署 AX 的步骤大致如下。首先确认集群版本AX 对 K8s 版本有要求建议 1.24 以上因为一些调度相关的 API 在低版本里不稳定。# 确认集群版本 kubectl version --short # 确认节点资源 kubectl describe nodes | grep -A 5 Allocatable # 确认 GPU 插件已安装如果要用 GPU kubectl get pods -n kube-system | grep nvidia然后安装 AX 的 CRD 和控制器。AX 的安装包通常包含 CRD 定义、控制器 Deployment、调度器 Deployment 三部分。# 安装 CRD kubectl apply -f https://github.com/google/ax/releases/latest/download/crds.yaml # 安装控制器和调度器 kubectl apply -f https://github.com/google/ax/releases/latest/download/ax-system.yaml # 确认组件运行 kubectl get pods -n ax-system这里有个坑要注意调度器需要和默认调度器共存。AX 的调度器通常以 secondary scheduler 的形式部署通过 schedulerName 字段指定哪些 Pod 用 AX 调度器。如果你直接替换默认调度器集群里其他工作负载会受影响。4.2 第一个 Agent 的部署与验证环境准备好后部署第一个 Agent 验证链路。建议从最简单的开始不要一上来就上大模型。apiVersion: ax.io/v1 kind: Agent metadata: name: hello-agent spec: runtime: python-agent-runtime:v1.2 model: name: tiny-model endpoint: http://model-service:8080 resources: requests: cpu: 500m memory: 1Gi command: [python, -m, agent.main] restartPolicy: OnFailure应用后观察状态kubectl apply -f hello-agent.yaml kubectl get agents kubectl describe agent hello-agent正常的话你会看到 Agent 从 Pending 变成 Running。如果卡在 Pending用 describe 看事件通常是资源不足或调度失败。如果卡在 Starting看 Runtime 的日志通常是模型加载失败或依赖缺失。4.3 多 Agent 编排与依赖管理单个 Agent 跑通后下一步是多 Agent 编排。AX 支持 Agent 之间的依赖声明比如 Agent B 依赖 Agent A 的输出。apiVersion: ax.io/v1 kind: Agent metadata: name: downstream-agent spec: dependsOn: - name: upstream-agent condition: Completed runtime: python-agent-runtime:v1.2 # ... 其他配置依赖管理的实现方式是控制器监听上游 Agent 的状态状态满足条件后触发下游 Agent 的创建。这里有个细节依赖条件要设计得足够细。不只是上游完成还可能是上游成功、上游产出特定数据、上游运行超过 N 秒。AX 的 condition 字段支持这些细粒度条件。实操中我遇到的一个问题是循环依赖检测。如果 A 依赖 B、B 依赖 A控制器会陷入死锁。AX 在提交时会做依赖图检测发现环就拒绝创建。但如果你通过动态方式创建依赖可能绕过检测。建议在 CI 阶段就做依赖图校验别等到运行时才发现。4.4 监控与可观测性配置Agent 跑起来之后可观测性是刚需。AX 暴露了 Prometheus 格式的指标包括 Agent 数量、调度延迟、执行时长、失败率等。apiVersion: v1 kind: ServiceMonitor metadata: name: ax-metrics spec: selector: matchLabels: app: ax-controller endpoints: - port: metrics interval: 30s关键指标我建议重点盯这几个ax_agent_scheduling_duration_seconds调度延迟反映调度器压力、ax_agent_pending_total待调度数量反映资源缺口、ax_agent_failure_total失败数反映稳定性。这三个指标能覆盖大部分运维场景。日志方面AX 的组件都输出结构化日志建议接入统一的日志平台。Agent 自身的日志通过 Runtime 收集可以配置输出到 stdout 或文件。我的经验是 Agent 日志一定要带 trace id否则多 Agent 协作时根本追不清调用链。5. 常见问题与排查技巧实录5.1 调度失败类问题调度失败是最常见的问题表现是 Agent 一直 Pending。排查思路是看 describe 输出的事件。现象可能原因排查方法解决方式一直 Pending资源不足看节点 Allocatable 和 Requests扩容节点或降低请求一直 Pending亲和性不满足看 nodeAffinity 配置调整亲和性规则调度后立即失败镜像拉取失败看 Runtime 日志检查镜像地址和凭证调度后卡 Starting模型加载慢看 Agent 日志增大 initialDelaySeconds频繁重新调度节点不稳定看节点状态排查节点健康有个隐蔽的坑是资源碎片化。集群总资源够但分散在各个节点上单个节点都放不下一个 Agent。这种情况 describe 会显示insufficient resources但你看总资源是够的。解决办法是配置资源整理策略或者用更大的节点。5.2 运行时异常类问题Agent 跑起来之后崩了排查起来更麻烦因为涉及 Agent 自身逻辑。常见的一类问题是上下文丢失。Agent 执行到一半重启之前的上下文没了导致任务失败。这个问题的根源是 Agent 没有做状态持久化。解决办法是在 Agent 里实现 checkpoint 机制定期把上下文写到外部存储重启后恢复。另一类问题是工具调用超时。Agent 调用外部工具工具响应慢Agent 一直等。AX 支持配置超时但超时后怎么处理需要 Agent 自己决定——是重试、是降级、还是直接失败。我的建议是给每个工具调用配独立的超时和重试策略不要用全局配置。5.3 性能调优经验跑了一段时间后性能调优是绕不开的。我总结了几个调优方向。调度吞吐如果 Agent 创建频率很高调度器可能成为瓶颈。AX 的调度器支持并发调度可以调大并发数。但并发太高会导致调度决策质量下降需要权衡。我的经验是并发数设为节点数的 2-3 倍比较合适。资源利用率默认调度器倾向于分散部署资源利用率不高。可以通过配置 bin-packing 策略提高利用率但代价是故障域变大。这个取舍要看业务对可用性的要求。冷启动优化Agent 冷启动慢是普遍问题主要慢在模型加载。优化手段包括模型预热提前加载到节点、模型缓存节点本地缓存模型文件、镜像优化减小镜像体积。这几个手段组合使用冷启动时间能从几分钟降到几十秒。5.4 独家避坑清单最后分享几个我踩过的坑都是文档里不会写的。注意Agent 的 restartPolicy 不要设 Always。Agent 正常完成后如果被重启会重复执行任务产生副作用。用 OnFailure 更安全。注意terminationGracePeriodSeconds 要设得比 Agent 最长单次执行时间还长。否则 Agent 正在执行长任务时被强杀状态不一致。注意多 Agent 共享模型服务时要评估模型服务的并发能力。Agent 数量上去了模型服务可能先扛不住。注意Agent 的日志量可能很大尤其是调试阶段。提前配置日志轮转否则磁盘很快满。注意调度器的自定义插件要幂等。调度器可能对同一个 Agent 多次调用插件插件如果有副作用会出问题。6. 这套调度思路还能怎么用AX 把 Agent 当 Pod 调度这个思路其实可以延展到很多场景。我最近在想的几个方向一是把 Agent 调度和 CI/CD 结合让 Agent 作为流水线的一个环节被调度二是把 Agent 调度和边缘计算结合让 Agent 在边缘节点上按需拉起三是把 Agent 调度和成本优化结合按实时价格选择最便宜的节点。我个人在实际操作中的体会是Agent 调度这件事难点不在调度算法本身而在 Agent 的状态管理和生命周期语义。Pod 的状态是相对简单的运行/终止Agent 的状态复杂得多思考中、调用工具中、等待输入中。AX 在这方面的抽象还在演进但方向是对的。如果你正在做 Agent 平台建议先把 Agent 的状态机定义清楚再考虑调度否则调度器再强也管不好状态混乱的 Agent。最后再分享一个小技巧调试调度问题时把调度器的日志级别调到 debug能看到每个 Agent 的完整调度决策过程包括每个节点的打分。这个信息比 describe 输出详细得多排查疑难调度问题非常有用。