微软开源 TauGrid,一个云原生平台,用于在搭载 GPU 的 Kubernetes 集群上对 AI 工作负载进行管理、调度与监控。
在 Kubernetes 上运行 AI 工作负载通常需要平台团队组装和维护多个开源项目、自定义脚本和运维工具。
这其中还包括这些组件之间的“粘合剂”:提交脚本、队列封装器、健康检查和结果检索。
微软表示,TauGrid 面向工程团队和研究团队设计,提供了一套统一的技术栈。
平台团队可以使用工作区、队列、计算配置文件、存储、身份认证和可观测性等高级功能,而研究人员无需学习 Kubernetes 即可提交工作负载。
微软称,TauGrid 为 AI 工作负载提供端到端管理,涵盖从初始数据准备到分布式训练、微调和推理的所有环节。
它基于 Kubernetes 构建,采用专门的队列和拓扑感知调度技术,高效管理高强度的 GPU 工作负载。
除了 tau 命令行工具之外,TauGrid 还集成了 Kueue(用于工作负载排队和资源管理)、KubeRay(用于编排)、GPU 节点健康监控以及可观测性能力。
TauGrid 无需分别构建和维护这些组件及组件之间的集成,通过一次 Helm 安装即可提供所有功能,并具有清晰的权责边界。
TauGrid 使用 yaml 配置文件来定义工作负载,可通过 tau run 提交。
该命令会验证配置并创建 Kubernetes Job 或 KubeRay RayJob,然后由 Kueue 根据剩余配额和优先级进行排队。
执行时,TauGrid 会跟踪工作负载状态、日志和检查点。
它还会收集和存储实验证据,以便日后复现实验和诊断故障。
以下是在单个 A100 GPU 上运行 PyTorch 训练作业的 tau.yaml 配置示例: schema_version: 1name: aks-gpu-quickstartrun: entrypoint: train.py workload_kind: rayjobcompute: gpus: 1 workers: 1 cpus: 16 memory: 64Giruntime: image: mcr.microsoft.com/aks/ai-runtime/ray:py3.12-ray2.56.0-cuda13.0 pip: - torch>=2.4.0 当作业失败时,TauGrid 可以从检查点恢复作业。
TauGrid 仍在开发当中,路线图中列出了一系列规划中的功能,包括多租户工作区、RBAC 和配额、对 PyTorch DDP/FSDP、DeepSpeed 和 LoRA/QLoRA 工作流的支持、数据集生命周期管理、多集群/多云执行等。
🔗 原始来源
如果你要核对细节,可以再看原文: InfoQ AI原文链接
AI热榜