Skip to main content

An adaptive and highly asynchronous ensemble simulation workflow manager MatEnsemble (https://github.com/Q-CAD/MatEnsemble) built jointly on top of the hierarchical graph based scheduler FLUX and concurrent-futures infrastructure of python

Project description

PyPI version Documentation Python License

MatEnsemble

MatEnsemble

MatEnsemble is a framework to build, orchestrate, and asynchronously manage extremely scalable adaptive-learning workflows, especially targeted for compute-intensive AI-driven high-throughput and ensemble-driven materials modeling simulations (e.g., atomistic modeling, Phase-Field, etc.) as efficiently as possible.

Here is an example which defines a simple MPI chore, adds ten independent instances of it to a pipeline, and submits the workflow.

pipe = Pipeline()

# register a function the MatEnsemble
@pipe.chore(num_tasks=10, cores_per_task=1, gpus_per_task=0, mpi=True)
def mpi_hello_world():
    size = MPI.COMM_WORLD.Get_size()
    rank = MPI.COMM_WORLD.Get_rank()
    name = MPI.Get_processor_name()

    print(f"Hello World! I am process {rank} of {size} on {name}.")

# adding 10 mpi_hello_world chores to the pipeline
for _ in range(10):
    mpi_hello_world()

pipe.submit(log_delay=1)

Rather than launching the MPI jobs directly, you describe the work declaratively with the @pipe.chore decorator and enqueue each invocation. MatEnsemble then schedules the chores through Flux, allocating the requested resources (num_tasks, CPU cores, GPUs, and MPI support) for each job.

Installation

If you are on the OLCF Frontier or Pathfinder HPC systems or the NERSC Perlmutter system then you can use our installation script to install MatEnsemble.

curl -fsSL https://raw.githubusercontent.com/Q-CAD/MatEnsemble/refs/heads/main/install.sh | bash

For general installation see our documentation


While MatEnsemble can operate on a personal macOS or Linux workstation, orchestrating arbitrary Python callables and shell commands through explicit, resource- and dependency-aware execution graphs from a single Python workflow driver, it is primarily designed for the autonomous execution of large batches of user-defined, adaptively and hierarchically scheduled tasks on HPC systems, especially petascale and exascale platforms such as Perlmutter, Frontier, and Aurora.

Minimal Code Example

MatEnsemble workflows are ordinary Python scripts (and/or shell commands) which can be use to: 1. define resource-aware chores, 2. pass chore outputs into later chores to create a DAG, and 3. add a strategy when the workflow should decide what to launch next while the campaign is already running.

from matensemble.pipeline import Pipeline
from matensemble.model import Resources
from matensemble.chore import ChoreSpec

pipe = Pipeline()

md_resources = dict(num_tasks=128, cores_per_task=1, gpus_per_task=4, mpi=True)
analysis_resources = dict(num_tasks=1, cores_per_task=8)


@pipe.chore(name="simulate", **md_resources)
def simulate(candidate):
    # Run LAMMPS, DFT, phase-field, or another science application here.
    return {"trajectory": "traj.dump", "candidate": candidate}


@pipe.chore(name="score", **analysis_resources)
def score(simulation):
    # Analyze the completed simulation and propose the next high-value sample.
    return {
        "uncertainty": 0.18,
        "next_candidate": {"temperature": 1750, "composition": "SiO2"},
    }


@pipe.strategy(bolo_list=["score"], **analysis_resources)
def adapt(report):
    if report["uncertainty"] < 0.05:
        return None

    return ChoreSpec(
        args=(report["next_candidate"],),
        kwargs={},
        resources=Resources(**md_resources),
        qualname="simulate",
    )


seed = {"temperature": 1600, "composition": "SiO2"}
trajectory = simulate(seed)
score(trajectory)  # OutputReference creates the simulate -> score DAG edge.

future = pipe.submit(log_delay=10)
results = future.result()

Publications

  1. Bagchi, Soumendu, et al. "Towards “on-demand” van der Waals epitaxy with adaptive ensemble sampling atomistic workflows." Digital Discovery (2026) https://doi.org/10.1039/d6dd00049e.
  2. Morelock, Ryan, et al. "pyRMG: A framework for high-throughput, large-cell DFT calculations on supercomputers." The Journal of Chemical Physics 164.5 (2026).

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

matensemble-0.5.5.tar.gz (67.2 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

matensemble-0.5.5-py3-none-any.whl (83.8 kB view details)

Uploaded Python 3

File details

Details for the file matensemble-0.5.5.tar.gz.

File metadata

  • Download URL: matensemble-0.5.5.tar.gz
  • Upload date:
  • Size: 67.2 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: uv/0.11.28 {"installer":{"name":"uv","version":"0.11.28","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"Ubuntu","version":"24.04","id":"noble","libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":true}

File hashes

Hashes for matensemble-0.5.5.tar.gz
Algorithm Hash digest
SHA256 823f8accebc772ce7cee87001b86cb0c392d502e6f603fefd0c18c58b0389e5a
MD5 2931bb85a24ba37399c43aada761c9ed
BLAKE2b-256 733fdae5a29176d4afec632148a3077b4feb0805896b831f526ff0ba588cca41

See more details on using hashes here.

File details

Details for the file matensemble-0.5.5-py3-none-any.whl.

File metadata

  • Download URL: matensemble-0.5.5-py3-none-any.whl
  • Upload date:
  • Size: 83.8 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: uv/0.11.28 {"installer":{"name":"uv","version":"0.11.28","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"Ubuntu","version":"24.04","id":"noble","libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":true}

File hashes

Hashes for matensemble-0.5.5-py3-none-any.whl
Algorithm Hash digest
SHA256 638adb7e209ce3e1fcda83a8f01dd1835991b2b4745a6c2e3138413303082f1a
MD5 b537fff674de1fb7fd07561675dda57b
BLAKE2b-256 60fba9be3d68003d7bc15ad2789c5e95d34bbd2a59eeca73ef70ddb2514f8fbb

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page