[Paper Review] 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 Sciences3 citations

TL;DR

Dflow is an open-source Python framework that enables scientists to build scalable, cloud-native AI-for-Science workflows using containerized, Kubernetes-managed operations. It supports complex, concurrent workflows across heterogeneous infrastructure, achieving over 1,200 GPU nodes in concurrency, and enables reusable, portable, and observable scientific computing pipelines with minimal abstraction overhead.

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.

Motivation & Objective

• To address the inefficiency and lack of reproducibility in traditional manual or scripted scientific workflow management on supercomputers.
• To enable seamless integration of AI-driven scientific workloads—such as active learning and concurrent learning—across cloud and high-performance computing (HPC) environments.
• To provide a user-friendly, extensible, and observable workflow framework that abstracts infrastructure complexity while supporting complex process control and task scheduling.
• To foster collaboration and reuse by enabling scientists to publish and share modular, containerized operations (OPs) and complete workflows.
• To bridge the gap between algorithm conception and practical deployment in AI for Science by unifying workflow orchestration with cloud-native and containerized execution.

Proposed method

• Dflow introduces the Operation (OP) as a first-class, reusable, and infrastructure-agnostic computational unit, encapsulating logic and dependencies in a containerized environment.
• It leverages Kubernetes for dynamic orchestration, enabling horizontal scaling to thousands of concurrent nodes across diverse computing platforms.
• The framework uses containerization to ensure environment reproducibility and supports local debugging without containers for development ease.
• Dflow integrates with cloud-native tools and platforms such as Bohrium, enabling deployment of workflows on managed cloud environments.
• It supports complex workflow patterns like fan-in/fan-out and recursive execution through the Dflow-galaxy extension, which provides a high-level builder API for workflow composition.
• The system is built on a modular architecture with a Pythonic API, allowing scientists to compose workflows using familiar programming constructs.

Experimental results

Research questions

• RQ1 How can scientific workflows in AI for Science be made scalable, portable, and reproducible across heterogeneous cloud and HPC environments?
• RQ2 What architectural patterns enable efficient, concurrent execution of complex scientific workloads such as active learning and high-throughput screening?
• RQ3 How can workflow frameworks support both high-level abstraction and low-level control for scientific workloads involving machine learning and quantum chemistry calculations?
• RQ4 To what extent can a workflow system reduce the development and operational overhead in iterative scientific pipelines like those used in interatomic potential training?
• RQ5 Can a unified, open-source framework accelerate the adoption of AI-driven scientific discovery by enabling component reuse and ecosystem growth?

Key findings

• Dflow supports workflows with up to approximately 1,500 Operations, achieving a maximum concurrency of over 1,200 GPU nodes, demonstrating high scalability in complex scientific workloads.
• The framework enables the implementation of advanced workflows such as TESLA, a four-step loop for active learning of machine learning potentials, including training, exploration, screening, and DFT labeling.
• Dflow-galaxy enables flexible input handling and supports complex concurrent learning, such as redox potential calculations, by extending Dflow with a builder pattern and artifact-based data flow.
• The framework has been used to deploy production-grade workflows on cloud platforms like Bohrium, including Dynacat TESLA, CP2K Lightning, and Dynacat MD, demonstrating real-world deployment viability.
• Dflow’s modular design and open-source nature have enabled dozens of workflow projects across diverse scientific domains, promoting reuse and ecosystem growth.
• The system provides full observability and supports local debugging without containers, improving developer productivity and reducing deployment friction.