Source-linked AI summary
Workflows in AiiDA: Engineering a high-throughput, event-based engine for robust and modular computational workflows
Martin Uhrin, Sebastiaan P. Huber, Jusong Yu, Nicola Marzari, Giovanni Pizzi
TL;DR
High-throughput computational science needs workflow infrastructure that can manage large data volumes while preserving reproducibility and operating robustly across computing environments. The paper details AiiDA’s event-based, Python-driven workflow engine and its design for dynamic, modular, provenance-aware computation. Its extensible ecosystem had over 100 supported simulation code executables and many workflows at the time of writing.
Problem
High-throughput computational science requires infrastructure that manages large data volumes while preserving reproducible data and workflow provenance.
Method
The paper presents AiiDA’s scalable workflow engine, combining directly executed Python workflows, provenance storage, event-driven execution, and robust remote-job handling.
Results
Over 100 supported simulation code executables and many workflows were available with AiiDA 1.0 at the time of writing.
Takeaways & Limitations
AiiDA provides an extensible and modular foundation for interoperable, FAIR computational infrastructure supporting scientific discovery.
Abstract
from arXiv · showhide
Over the last two decades, the field of computational science has seen a dramatic shift towards incorporating high-throughput computation and big-data analysis as fundamental pillars of the scientific discovery process. This has necessitated the development of tools and techniques to deal with the generation, storage and processing of large amounts of data. In this work we present an in-depth look at the workflow engine powering AiiDA, a widely adopted, highly flexible and database-backed informatics infrastructure with an emphasis on data reproducibility. We detail many of the design choices that were made which were informed by several important goals: the ability to scale from running on individual laptops up to high-performance supercomputers, managing jobs with runtimes spanning from fractions of a second to weeks and scaling up to thousands of jobs concurrently, and all this while maximising robustness. In short, AiiDA aims to be a Swiss army knife for high-throughput computational science. As well as the architecture, we outline important API design choices made to give workflow writers a great deal of liberty whilst guiding them towards writing robust and modular workflows, ultimately enabling them to encode their scientific knowledge to the benefit of the wider scientific community.
I. INTRODUCTION
AiiDA addresses the reproducibility and scalability challenges of high-throughput computational science with an API-driven, provenance-aware workflow system. Its directly executed Python workflows support dynamic execution while integrating arbitrary codes through plugins.
- I. INTRODUCTION: Data provenance must be recorded automatically because provenance enables validation, verification, reuse, and reconstruction of data-producing workflows.The paper treats workflows themselves, not only their outputs, as part of tracked provenance.
- I. INTRODUCTION: AiiDA workflows execute directly as Python code, allowing runtime paths to evolve from completed-step results without an intermediate translation layer.This provides programming-language expressiveness and access to the stored provenance graph.
- I. INTRODUCTION: Static markup workflows require the exact flow to be specified before execution and naturally constrain workflows to DAG- or DCG-like structures.The paper contrasts this limitation with AiiDA’s directly programmed workflows.
- I. INTRODUCTION: AiiDA’s plugin system makes external codes compatible through registry-based plugins installable with pip.The paper presents this as part of AiiDA’s generic support for computational workflows.
- I. INTRODUCTION: AiiDA is designed to define dynamic workflows, support arbitrary external codes, automatically store queryable provenance, and scale without disproportionate storage overhead.The system separates a user API for workflow and provenance-graph interaction from an engine that executes workflows and stores results.
II. USER INTERFACE
AiiDA’s user interface combines unrestricted Python access with constructs that promote maintainable, modular, and robust workflows. This flexibility has a clear boundary: users can deliberately break full provenance when necessary.
- II. USER INTERFACE: AiiDA workflows are written and executed directly in Python, giving developers access to the full AiiDA API and Python libraries without a translation layer.This design also allows existing Python code to interface with the engine without additional custom development.
- II. USER INTERFACE: The engine adds workflow constructs that counteract incompatibility and non-modularity while preserving broad implementation freedom.These tools are intended to improve robustness and interoperability.
- II. USER INTERFACE: Full data provenance is not guaranteed under unrestricted workflow design, so AiiDA defines clear conditions for when provenance is preserved or broken.The design permits users to break provenance when they consider it necessary or justified.
- A. Process specification: A process is code implementing logical instructions that transforms inputs into outputs and may terminate prematurely through known failure modes.Its inputs, outputs, and failure modes are specified through ProcessSpec.
- 1. Ports and port namespaces: Port namespaces recursively validate nested ports, including accepted types and custom validators, and can be marked dynamic to accept unspecified ports.A namespace is valid only when its nested ports and the namespace itself pass validation.
2. Inputs and outputs
AiiDA’s ProcessSpec declaratively defines process inputs, outputs, and namespaces through typed, validated ports. Later declarations can override earlier ones, making specification construction composable but order-sensitive.
- 2. Inputs and outputs: ProcessSpec stores input and output ports in separate namespaces and provides methods to create individual ports or nested namespaces.The same interface supports both process boundaries and hierarchical port organization.
- 2. Inputs and outputs: A process specification can declare typed, defaulted, and validated ports with requiredness constraints.The examples show input and output declarations using valid_type, default, validator, and required attributes.
- 2. Inputs and outputs: A declaration such as an input or output call creates the corresponding port under the specified namespace key.Namespace methods interpret periods as separators and recursively create nested namespaces.
- 2. Inputs and outputs: The declarative ProcessSpec API allows later declarations to overwrite earlier declarations.In the example, a later input declaration changes port a from an integer specification to a float specification.
3. Exit codes
AiiDA represents process failure through declared exit codes, while decorators turn ordinary Python functions into provenance-tracked processes. Decorated calculations preserve outputs and execution relationships in the provenance graph.
- Exit codes: Exit codes are integer statuses returned by all processes: zero denotes successful execution, while non-zero values indicate mapped errors.Process specifications declare known failure modes using exit codes.
- Exit codes: A declared exit code combines an integer status, a reference label, and a human-readable message stored with the process.The example defines status 418, label ERROR_I_AM_A_TEAPOT, and an explanatory message.
- Calculation functions: The calcfunction decorator converts a regular Python function into an AiiDA process using its inspected signature to define input ports.This adds AiiDA process behavior while retaining the function's ordinary Python body.
- Calculation functions: Calling decorated functions with storable inputs returns ordinary values while automatically creating linked calculation and data nodes in the provenance graph.The nested example returns 35, with each calculation represented by a CalcFunctionNode and linked inputs and outputs.
2. Work functions
Work functions encode coordinated sequences of calculations as AiiDA workflows and record their call relationships in the provenance graph. They support nested workflows but block the interpreter during contiguous, computationally expensive function execution.
- Work functions: A work function distinguishes coordinated workflow logic from merely consecutive calculations by recording call links between the functions it invokes.The example adds two numbers, multiplies the result, and stores the resulting workflow process in the provenance graph.
- Work functions: Work functions can call both calculation functions and other work functions, enabling arbitrarily deeply nested workflows through CALL links.The two decorator-based process-function types provide the basic components for nested workflow construction.
- Limitations: Decorated functions execute as contiguous code that blocks the interpreter for the duration of computationally expensive operations.This makes process functions unsuitable for long-running calculations that need concurrent engine activity.
- Limitations: AiiDA therefore recommends using process functions sparingly and provides WorkChain as the construct addressing the workfunction's weak points.The passage frames WorkChain as the alternative for workflows requiring more robust execution behavior.
3. Work chains
WorkChain provides AiiDA's core workflow-development model by expressing logic as restartable outline steps with checkpointing, context-based data transfer, subprocess control, and output recording. Its API supports conditional and iterative workflows while preserving modular execution through child processes.
- Outline: WorkChain is a Process subclass whose outline encodes workflow logic as engine-executed steps and automatically saves progress between steps for restart after failure.The outline supports Python-like loops, conditionals, and returns.
- Outline and reporting: Outline steps are implemented as WorkChain methods, while reporting uses a custom REPORT log level integrated with AiiDA's database log handler.The outline syntax uses suffixed constructs such as while_, if_, elif_, and else_ to avoid Python keyword conflicts.
- Checkpoints and context: A WorkChain context, exposed through ctx, transfers data between outline steps that otherwise receive only self.The context is a Python dictionary maintained on each WorkChain instance.
- Calling subprocesses: WorkChains can submit CalcJobs or other WorkChains as child processes and use ToContext or to_context to pause until those subprocesses finish.Completed child process nodes are assigned to context keys and accessed in later outline steps.
- Calling subprocesses: Multiple independent subprocesses can run in parallel, with append_ collecting their completed process nodes into an ordered list in the context.Unique keys are unnecessary when results are collected under a shared list key.
- Recording outputs: The out method records output nodes in memory during a step and commits them to the database after the step, when output-port validation occurs.The method takes an outgoing link label and a node instance.
4. Calculation jobs
AiiDA automates remote calculation jobs across upload, submission, monitoring, and retrieval while using retries, pausing, and queued connections to remain robust under high-throughput conditions.
- AiiDA’s CalcJob automates remote calculation execution, including file upload, scheduler submission, job-status updates, and output retrieval.The engine creates a remote scratch folder, uploads inputs and the job script, submits it, monitors completion, and retrieves specified files.
- Error handling and robustness: Exponential-back-off retries reschedule failed transport tasks with increasing wait intervals and configurable retry limits.The mechanism catches transport exceptions, retries operations later, and configures the initial interval and maximum retries by task type.
- Error handling and robustness: When retries do not resolve a problem, AiiDA pauses the process for investigation and later resumption instead of allowing it to fail outright.Users can resume processes after fixing external causes or manually kill processes whose failures arise from workflow or parser code.
- Transport queue: A transport queue limits connection-opening rates for concurrent jobs, reducing the risk of exceeding remote-resource access limits.Each worker maintains its own queue, so the safe interval is guaranteed per worker and must be configured across workers using their known count.
- Bundling scheduler update requests: Bundling scheduler update requests reduces remote scheduler load, although high-throughput operation can still require additional request management.Shared transport connections alone do not prevent each active calculation from separately querying scheduler state.
III. ARCHITECTURE
AiiDA’s architecture separates persistence and communication through a database and message broker, supporting flexible deployments and robust workflow execution.
- The system is designed to scale from laptops to supercomputers, support runtimes from fractions of a second to weeks, and handle thousands of simultaneous processes.Deployments can be local, organisation-wide, or public-serving.
- AiiDA uses PostgreSQL to persist currently running process states and RabbitMQ to deliver messages between clients and workers.Clients and workers may run in the same Python instance or on separate computers.
- The database-backed state also acts as a proxy through which users can observe running processes.
- This decoupled architecture supports flexible deployment configurations and separates concerns, making correct and robust code easier to write.
A. The engine
AiiDA’s engine uses event-driven coroutines rather than polling to advance workflows responsively and manage multiple processes within one Python instance.
- Events trigger workflow actions such as resuming processes waiting for completion or orderly termination, avoiding periodic polling checks.
- Python’s asyncio event loop lets coroutines yield while waiting, allowing one Python instance to manage multiple AiiDA processes without multithreading.This approach avoids the complexity of writing correct multithreaded code.
- The Runner combines the event loop with persistence, communication, transport, and other components, and can run concurrent workflow processes within memory limits.The number of concurrently runnable processes on one runner is called the number of process slots.
1. The daemon
AiiDA’s daemon supervises runners and supports vertical or horizontal scaling according to workload characteristics.
- The daemon can automatically restart crashed runners while exposing process information and resource-usage monitoring.
- Scaling can be vertical through multiple slots per runner or horizontal through multiple Python instances, each with one runner.In-Python-heavy workloads favor more workers, whereas remote-calculation workloads can increase slots per worker to reduce daemon-host load.
B. The process
AiiDA models every executable entity as an event-driven Process with extended state and transition hooks. These mechanisms support controlled termination, state-dependent actions, and database updates during execution.
- B. The process: AiiDA’s Process class is an extended finite-state machine whose subclasses inherit shared workflow functionality.The state machine models both finite states and internal data members.
- B. The process: Event hooks run when processes enter or exit states, enabling actions to be tied consistently to state transitions.Hooks include on_entering, on_entered, and on_exiting callbacks.
- B. The process: Figure 6 represents the process state machine, with terminal states shown using double circles.The process commonly progresses through states while running functions or waiting on other processes.
- B. The process: State-transition hooks save process state to the database and broadcast updates to listeners.AiiDA uses them to reflect current state and save checkpoints.
- B. The process: Exceptions propagated to the Process level move it into the terminal EXCEPTED state and create a log containing the Python stack trace.Processes can also terminate prematurely through the kill method.
1. Persistence
AiiDA checkpoints process execution at state transitions so orderly or disorderly shutdowns can be followed by continuation from the last clean state.
- 1. Persistence: AiiDA persists each process’s context, outputs, and selected metadata by serializing them to the database at state transitions.The persister requests an out_state dictionary from the process before committing the checkpoint.
C. Communication
AiiDA uses RabbitMQ-based messaging and task queues to coordinate processes, preserve jobs across restarts, and support remote control and broadcast communication. Its architecture combines these mechanisms with scalable, fault-tolerant workflow execution and database-backed persistence.
- C. Communication: AiiDA’s messaging layer uses RabbitMQ to support external process control and high-throughput, fault-tolerant execution.The system uses durable and atomic message handling through a message broker.
- C. Communication: KiwiPy simplifies RabbitMQ interaction, offloads communication to a separate thread, and provides RPC and broadcast messaging.The separate thread helps processes respond to heartbeats during blocking workloads.
- C. Communication: Persistent RabbitMQ task-queue messages ensure that delivered jobs survive intentional or unintentional machine restarts.Heartbeats allow RabbitMQ to detect dead workers and trigger rescheduling.
- C. Communication: RPC messages pause, play, or kill active processes, while broadcasts control process groups and parent–child execution flow.Parent processes can listen for child broadcasts before continuing.
- C. Communication: AiiDA’s conclusions combine event-driven programming, futures, RabbitMQ, and PostgreSQL with a user-friendly API for modular, robust workflows and automatic provenance.The engine is designed for high-throughput workflows while remaining usable on personal desktops.
- C. Communication: At the time of writing, AiiDA 1.0 supported over 100 simulation-code executables and many workflows contributing to published scientific works.These resources supported the project’s aim of an interoperable FAIR computational infrastructure.