[论文解读] Dflow, a Python framework for constructing cloud-native AI-for-Science workflows

Xinzijian Liu, Yanbo Han+16 more|arXiv (Cornell University)|Apr 29, 2024

Scientific Computing and Data ManagementDecision Sciences被引用 3

一句话总结

Dflow 是一个开源的 Python 框架，使科学家能够使用容器化、Kubernetes 管理的操作构建可扩展的、云原生的科学人工智能工作流。它支持跨异构基础设施的复杂、并发工作流，在并发中实现超过 1,200 个 GPU 节点，并通过最小的抽象开销，实现可重用、可移植且可观测的科学计算流水线。

ABSTRACT

In the AI-for-science era, scientific computing scenarios such as concurrent learning and high-throughput computing demand a new generation of infrastructure that supports scalable computing resources and automated workflow management on both cloud and high-performance supercomputers. Here we introduce Dflow, an open-source Python toolkit designed for scientists to construct workflows with simple programming interfaces. It enables complex process control and task scheduling across a distributed, heterogeneous infrastructure, leveraging containers and Kubernetes for flexibility. Dflow is highly observable and can scale to thousands of concurrent nodes per workflow, enhancing the efficiency of complex scientific computing tasks. The basic unit in Dflow, known as an Operation (OP), is reusable and independent of the underlying infrastructure or context. Dozens of workflow projects have been developed based on Dflow, spanning a wide range of projects. We anticipate that the reusability of Dflow and its components will encourage more scientists to publish their workflows and OP components. These components, in turn, can be adapted and reused in various contexts, fostering greater collaboration and innovation in the scientific community.

研究动机与目标

• 解决在超级计算机上使用传统手动或脚本化方式管理科学工作流时效率低下且缺乏可复现性的问题。
• 实现人工智能驱动的科学工作负载（如主动学习和并发学习）在云和高性能计算（HPC）环境之间的无缝集成。
• 提供一个用户友好、可扩展且可观测的工作流框架，抽象基础设施复杂性，同时支持复杂的流程控制和任务调度。
• 通过支持科学家发布和共享模块化、容器化的操作（OPs）和完整工作流，促进协作与复用。
• 通过将工作流编排与云原生和容器化执行统一，弥合人工智能科学中算法构思与实际部署之间的差距。

提出的方法

• Dflow 引入了操作（OP）作为一等公民，可重用且与基础设施无关的计算单元，将逻辑和依赖关系封装在容器化环境中。
• 它利用 Kubernetes 实现动态编排，支持在各种计算平台之间横向扩展至数千个并发节点。
• 该框架使用容器化技术确保环境可复现，并支持无容器本地调试，以提升开发便捷性。
• Dflow 集成云原生工具和平台（如 Bohrium），支持在托管云环境中部署工作流。
• 通过 Dflow-galaxy 扩展，支持复杂的工作流模式（如扇入/扇出和递归执行），该扩展提供高层级构建器 API，用于工作流组合。
• 该系统采用模块化架构，配备 Pythonic API，使科学家能够使用熟悉的编程原语组合工作流。

实验结果

研究问题

• RQ1 如何在异构的云和 HPC 环境中，使人工智能科学中的科学工作流具备可扩展性、可移植性和可复现性？
• RQ2 哪些架构模式能够实现对主动学习和高通量筛选等复杂科学工作负载的高效并发执行？
• RQ3 工作流框架如何在涉及机器学习和量子化学计算的科学工作负载中，同时支持高层抽象和底层控制？
• RQ4 工作流系统在迭代式科学流水线（如原子间势能训练）中，能在多大程度上减少开发和运维开销？
• RQ5 一个统一的开源框架能否通过支持组件复用和生态系统发展，加速人工智能驱动的科学发现的采用？

主要发现

• Dflow 支持最多约 1,500 个操作，最大并发节点数超过 1,200 个 GPU，证明了其在复杂科学工作负载中的高可扩展性。
• 该框架支持实现高级工作流，如 TESLA，一种用于机器学习势能主动学习的四步循环（包括训练、探索、筛选和 DFT 标注）。
• Dflow-galaxy 通过扩展 Dflow 的构建器模式和基于工件的数据流，支持灵活的输入处理和复杂并发学习（如氧化还原电势计算）。
• 该框架已在 Bohrium 等云平台成功部署生产级工作流，包括 Dynacat TESLA、CP2K Lightning 和 Dynacat MD，证明了其在现实场景中的可行性。
• Dflow 的模块化设计和开源特性已促成数十个跨不同科学领域的流程项目，推动了复用和生态系统发展。
• 该系统提供完整的可观测性，并支持无容器本地调试，提升了开发人员生产力并降低了部署摩擦。