AI-Assisted Workflow in Scientific Computing

• AI-assisted workflows are automated pipelines that integrate machine learning, simulation engines, and high-performance computing for complex scientific tasks.
• They use modular task design, dynamic scheduling, and decoupled resource configurations to achieve scalability, portability, and fault tolerance.
• The exa-AMD system exemplifies these workflows by combining database integration, AI/ML screening, and high-fidelity simulations across diverse computing infrastructures.

AI-assisted workflow refers to a computational orchestration that integrates artificial intelligence or machine learning components—such as predictors, optimizers, or surrogates—within a robust, automated architecture to perform complex, multi-stage scientific or industrial tasks. These workflows enable the execution of data-driven pipelines at scale, often bridging disparate elements such as domain databases, simulation engines, and high-performance computing, all while supporting fine-grained task parallelism and portability across resources. The paradigm is exemplified by exa-AMD, a Python-based, Parsl-driven system for accelerated materials discovery and design, in which modular workflow stages are dynamically scheduled and scaled from laptops to supercomputers, providing a template for general AI-augmented scientific pipelines (Moraru et al., 26 Jun 2025).

1. Pipeline Architecture and Core Components

A typical AI-assisted workflow consists of a sequence of logically distinct stages, each implemented as an autonomous or semi-autonomous task node. In exa-AMD, these stages are:

1. Input Acquisition: Initial entities, such as crystal structure templates, are sourced from standardized materials databases (e.g., Materials Project, AFLOW, OQMD, GNoME).

1. Candidate Generation: Algorithmic procedures generate new candidates, e.g., via element substitution or combinatorial transformations.
2. AI/ML Screening: ML surrogates (e.g., Crystal Graph Convolutional Neural Networks) rapidly predict quantities of interest (e.g., formation energy) at low computational cost.
3. Deduplication: Near-duplicate candidates are removed using structural fingerprints or checksums.
4. High-Fidelity Simulation: Expensive, often quantum-mechanical, simulations are submitted for the most promising candidates (e.g., DFT with VASP).
5. Postprocessing and Feedback: Key properties (e.g., energy above hull) are extracted and synthesis feedback loops (e.g., convex hull updates) are triggered.

Each stage is coded as a Python function with Parsl task decorators. Data dependencies are specified via function signatures, yielding an implicit, dynamically constructed dataflow graph at runtime.

2. Workflow Logic, Dataflow, and Execution Model

The workflow logic is decoupled from physical execution and resource configuration. In exa-AMD:

• Task Abstraction: Each step is a fine-grained, restartable task returning Parsl DataFutures. Downstream dependencies are tracked via futures rather than explicit DAG files.
• Decoupled Configuration: Resource parameters (scheduler, node count, wall time, thread count) are externalized in YAML/JSON and passed to Parsl’s Executors (HighThroughputExecutor, ThreadPoolExecutor, etc.).
• Dynamic Scheduling and Robustness: Parsl’s dataflow scheduler launches tasks as soon as dependencies resolve, with execution details invisible to the workflow code. Persistent metadata (e.g., SQLite) enables automatic resumption after failures and prevents redundant computations.

This decoupling allows a single workflow logic to execute consistently from a local ThreadPoolExecutor (e.g., on a laptop, max_threads=4) up to a 4,096-core supercomputer run (e.g., on ALCF’s ThetaGPU, 128 GPUs), with resource changes requiring only configuration edits, not code rewrites.

3. Performance and Scaling Properties

Strong- and weak-scaling performance is central to AI-assisted workflow design. In exa-AMD:

• Performance Model: The end-to-end runtime is modeled as

$$\text{Runtime} \approx \frac{T \cdot \tau_{\text{avg}}}{C \cdot \epsilon(C)}$$

where $T$ is the number of candidates, $C$ is the number of compute units (cores or GPUs), $\tau_{\text{avg}}$ is the average per-task wall-time at the slowest stage, and $\epsilon(C)$ is the parallel efficiency.

• Empirical Scaling: For large $T \gg C$, parallel efficiency is sustained, with $\epsilon(C) \approx 1 - \alpha \log C$, $\alpha \ll 1$. Scaling to 128 GPUs, exa-AMD demonstrated speedup $S(C) \approx C^{0.98}$.
• Overhead: Workflow orchestration overheads (dispatch, I/O, queuing) scale sublinearly, $O(\log T + \log C)$, due to batching and allocation reuse.

These properties are generalizable to other scientific workflows leveraging concurrent ML screening and high-throughput simulation.

4. Patterns for Scalability, Portability, and Fault Tolerance

exa-AMD instantiates a series of workflow principles critical for robust, scalable AI-assisted scientific computation:

1. Task Modularity: The pipeline is built from atomic, self-contained modules (e.g., separate tasks for candidate generation, ML prediction, DFT calculation), each swappable or extendable.
2. Execution/Logic Decoupling: All resource-specific details (scheduler, node count, environmental variables) are isolated in configuration files, allowing rapid retargeting to new infrastructure.
3. Fine-Grained Tasking and Resumability: Tasks are defined at a granularity (e.g., per-candidate) that permits granular retry on failure and efficient checkpointing. Task success/failure, input/output and state are persisted in lightweight local stores.
4. Elastic Resource Allocation: Executors can dynamically grow or shrink worker pools based on task load and resource requirements, releasing (e.g.) unused GPUs when only CPU-bound stages remain.
5. Hierarchical Data Management: Outputs are organized by pipeline stage and candidate ID; checksums/fingerprints enforce deduplication, with intermediate data staged locally and final results stored on parallel filesystems.
6. Performance Instrumentation: Each task records timing and resource usage, enabling adaptive workload balancing (e.g., increasing concurrency for candidates with heavier DFT requirements).

5. Generalization and Cross-Domain Applications

The exa-AMD blueprint is abstractable to any AI-assisted computational pipeline that requires:

• Integration of large static databases, fast AI/ML inference engines, and high-fidelity simulators.
• Fault-tolerant, elastic orchestration of thousands of loosely coupled tasks.
• Portability between workstation, cloud, and HPC environments.
• Persistent, resumable execution with robust task failure handling.

The workflow idioms—Pythonic dataflow graphs, declarative config, fine-grained tasking, elastic execution—are equally applicable to materials design, computational chemistry, computational biology, and applied physics campaigns that exploit AI surrogates to triage candidates for costly simulations (Moraru et al., 26 Jun 2025).

6. Best Practices and Design Recommendations

From implementation experience, essential recommendations for AI-assisted workflow design include:

• Emphasize atomic task granularity to maximize fault isolation and minimize re-execution cost.
• Always separate workflow logic from resource definitions for maximal portability and ease of scaling.
• Persist minimal but sufficient task state and metadata (e.g., via SQLite or equivalent) to enable resumability and reproducibility.
• Instrument all stages for wall-time and resource tracking, enabling responsive load balancing and identification of pipeline bottlenecks.
• Establish clear data hierarchies and deduplication strategies early, especially when generating or triaging large candidate sets.
• Employ workflow engines (Parsl, Dask, Swift/T) that support dynamic, implicit DAG construction and elastic resource management, minimizing user intervention as workload scales.

By formally adhering to these patterns, AI-assisted workflows can be composable, robust, and performant across a diverse landscape of scientific and engineering domains (Moraru et al., 26 Jun 2025).