Ensemble parallelism

Simulations are increasingly being performed not singly but as a collection of related tasks, which we will term an ensemble. This collection could consist of a set of replicas sampling a thermodynamic ensemble, but it could also represent a parameter sweep across experimental conditions or a set of simulations with different starting conformations as part of an advanced sampling strategy. Ensembles are treated as first-class objects in gmxapi, and the python interface is designed to facilitate parallelism across such ensembles, which is in turn implemented by the gmxapi backend. This is illustrated in Fig 1B. The ability to flexibly and simply articulate high-level ensemble logic is thus a key feature of gmxapi.

Composability

Another key design and usability feature of gmxapi is the notion of composability. Similar to Tinkertoys, different calls to GROMACS commands through the API are designed to be plugged together. This is illustrated in Fig 1C and has two advantages: a natural way to conceptualize a sequence of GROMACS calls (or calls to external programs that can use the gmxapi wrapper facility) and a way to parallelize using the ensemble logic and data-flow management by gmxapi.

Plugins

Gmxapi includes the ability to add user-defined plugins that can interact with GROMACS at runtime without modifying the GROMACS source. Plugins are currently used to implement custom force routines; these run using the native GROMACS parallel decomposition and thus can benefit from parallelism as well as acceleration. More details are given in the technical design section below.

Technical design

We used pybind11[23] to implement Python bindings to a C++ API for the GROMACS[21,22] molecular simulation library. The design concept of gmxapi, as with many modern high-level interfaces, is for the user to construct a computational graph where execution is then managed by lower layers of the software stack. Similar to Keras[24], we recognize that explicit data-flow programming is not natural for all users and thus provide an interface where data-flow can be specified explicitly or implicitly. Fig 2 illustrates the gmxapi paradigm in which control signals and data are treated in the same framework. Branching and other control events require that some work is "dynamic", such that the complete workflow graph may not be determined until runtime. However, the workflow can be fully specified in terms of workflow commands and data primitives (Fig 3).

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 2. Schematic of data flow and control flow for a segment of a complex simulation workflow.

Ensembles of simulations can be run (denoted by stacked rectangles) by gmxapi merely by passing an array of inputs instead of a single input. Custom plugins can interact with running MD simulations. Finally, conditional and looping logic can create high-level simulation algorithms.

https://doi.org/10.1371/journal.pcbi.1009835.g002

Download:

• PPT
  PowerPoint slide
• PNG
  larger image
• TIFF
  original image

Fig 3. Schematics for gmxapi operations.

As shown in panel (A), gmxapi operations have well defined inputs and outputs. These can operate on arrays of inputs and also include support for legacy GROMACS operations by wrapping the command-line toolset, as shown in panel (B). As shown in panel (C), the decorator @function_wrapper allows arbitrary user code to be transformed into a gmxapi operation.

https://doi.org/10.1371/journal.pcbi.1009835.g003

Data flow formalism

Data-flow formalisms have become increasingly popular in contemporary workflow engines due to their ability to separate work specification and task execution, permitting more straightforward resolution of dependencies and optimization of computation and data movement. For the same reasons, execution in gmxapi is deferred as much as possible to when and where it is required. The high-level gmxapi Python interface creates and executes a directed acyclic graph (DAG) specifying the computational work to be done. Nodes of the work graph represent discrete operations that produce and consume well defined inputs and outputs. Operation and data references in the work graph are proxy objects for the computational tasks and data until the graph is run. Execution management may be optimized by running only as much of the graph as necessary to satisfy explicit data dependencies. In addition, in order to minimize unnecessary data movement, most gmxapi operations do not transfer data from the C++ library to the Python interpreter unless explicitly requested by the user. When a gmxapi command is used to add work to the work graph, it returns a reference to the graph node representing the operation. The operation’s output attributes may be used as inputs to further commands, and may be used as “Future” results, as discussed below. Calling “result()” on a Future forces dependency resolution and data localization.

Python interface

The gmxapi Python interface consists of five categories of operations: 1) typing and logical operations, 2) bound GROMACS API calls, 3) legacy GROMACS command-line operations, 4) utilities for user creation of new gmxapi operations, and 5) looping and conditional operations. These are described in sequence below.

The simplest gmxapi operations manipulate signals, data topology, or typing. “join_arrays”, “logical_not”, and “make_constant” have strict definitions for inputs and outputs (Fig 3A). These operations are also composable: “concatenate_lists” is a helper function that reduces a number of inputs in terms of “join_arrays” (Fig 3B).

The preferred mechanism for interfacing to GROMACS is via bound API calls. Commands like “read_tpr”, “modify_input”, and “mdrun” use a binary Python extension module written in C++ using pybind11 to interact with GROMACS operations and data through libgmxapi, which is a C++ library installed by default with recent versions of GROMACS.

In order to provide legacy support for the full range of GROMACS command-line operations, “gmxapi.commandline.cli()” is a simple pure-Python gmxapi operation that wraps command-line tools. It is a thin wrapper of the Python “subprocess” module, so command inputs and outputs are embedded in the argument list parameter. Instances of “gmxapi.commandline.cli” are not by themselves conducive to DAG representations of data flow, so a helper function creates additional gmxapi primitives to provide a consistent set of named inputs and outputs. “gmxapi.commandline_operation()” generates a graph of gmxapi primitives around “cli” to translate input and output arguments into operation inputs and outputs, as specified by the user (Fig 3B).

In order to provide extensibility, arbitrary user code can be transformed into a gmxapi operation by decorating a Python function definition with “@function_wrapper” (Fig 3C). As an example, “commandline_operation” is written in this manner.

Looping and conditional commands are critical to adaptive ensemble simulation[25] and thus form a key part of the gmxapi repertoire. Conditional iteration takes the form of conventional gmxapi command syntax but relies on some metaprogramming under the hood. The “while_loop” command allows a graph to be dynamically extended as a chain of repeated operations, including fused operations constructed with the “subgraph” tool. A subgraph allows inputs and outputs to be expressed in terms of other gmxapi operations. When used with “while_loop”, internal subgraph state can be propagated from one iteration to the next to provide a consistent environment similar to a standard loop.

The work graph

The work graph permits data flow topology to be represented independently from execution strategies or run time resource assignment details. An array of input sources provided to a command generates a corresponding set of tasks, such as for trajectory ensemble simulations. Because such ensembles may be coupled (and gmxapi currently does not have a mechanism to specify that an ensemble is uncoupled), arrays of simulations must currently be co-executed through the gmxapi mpi4py executor. Operations defined with “@function_wrapper”, including the other built-in operations, are assumed to be uncoupled, and are launched in a sequence determined by the DAG topology during recursive resolution of data dependencies, executed sequentially by the simple built-in gmxapi 0.1 executor.

The work graph enables straightforward dependency resolution: commands are specified to operate on abstracted “handles” or references to work inputs rather than requiring the fully instantiated objects, so chains of simulations can be expressed where one command depends on the outputs of a prior command. Each individual command is then ready to execute once its inputs have been fully resolved. Because operations and data flow are represented as a DAG, arbitrarily complex topologies can be expressed unambiguously without unexpected side-effects. Trajectories can be forked or extended without re-executing simulation segments or overwriting previous results.

Execution and control flow

The gmxapi interface enables several levels of control logic, described below. Native Python logic (if/else statements, for loops, etc.) can be used with gmxapi operations. However, since gmxapi operations construct a work graph that is then executed, expressing conditionals within this work graph requires special constructs rather than standard Python operators. Gmxapi therefore provides logic such as “while_loop” as an efficient way to express adaptivity within a sequence of gmxapi operations. Python-level logic can be utilized by forcing work graph execution using the result() method, which causes explicit execution of all code required to produce the requested result. In contrast, gmxapi-level logic operates on Futures and is thus deferred until graph execution.

Gmxapi also provides low-level logic for simulation control and ensemble operators that can be used by third party code. For instance, gmxapi registers with the GROMACS “StopSignal” facility, enabling plugin code to stop a simulation based on external criteria, such as when a statistical estimator has converged. Additional examples of low-level operations include a “ReduceAll” to collect data across an in-flight ensemble. This operation can be used for adaptive updates across an ensemble of simulations, such as modifying biasing forces in restrained-ensemble simulations[26,27] or to update estimators and determine when a simulation ensemble should be terminated.

Gmxapi as a means to extend MD code without source modification

Gmxapi establishes a facility[28] for providing MD extension code to GROMACS during launch. Binary objects are then loaded at runtime via the Python interpreter. This allows custom code to be executed during the MD integration loop with minimal overhead. Python facilitates plugin binding, but once the simulation launches, GROMACS and the extension code communicate directly via C++ interfaces, so there is no overhead added by Python. This extension code executes within the existing GROMACS parallelism framework, so it can take advantage of domain decomposition and MPI parallelism. Since this only requires a standard GROMACS installation with gmxapi, it permits flexible extension of the molecular dynamics code without modifying the main codebase or libraries.

Gmxapi does not prescribe a C/C++ bindings strategy. We used pybind11 for GROMACS bindings, but Python interaction with compiled code relies only on the Python C API (PyCapsule) and the public libgmxapi C++ interface.