Dask

Dask

Overview

 Dask  can be used to parallelize compute tasks. Dask can also be very powerful for modern data science tasks like those based on pandas and several machine learning libraries. It can also be used to parallelize any kind of Python function across the different cores on a compute node, which we will demonstrate below. Note that Dask is not currently suitable for distributing work over multiple nodes in a single Slurm allocation. If that is a feature you need, use  🌀⁠⁠Parsl⁠  instead.
To install Dask with its relevant add-ons, run the following:
pip install dask[distributed] dask-jobqueue

Simple Example

Motivation

Let's imagine I want to calculate the D4 dispersion correction for many materials. Each calculation takes only a few seconds, but if I'm doing this for thousands of materials, I don't want to have to wait all day. Each D4 calculation runs on a single-core, so it would be a waste of computing resources to request a full node without additional care. Ideally, what we would like to do is request a single node (e.g. 112 cores on a single Tiger node) so that we can run 112 concurrent D4 calculations across the full node. Dask lets us do exactly this.

Problem Setup

First, let's start with how to do the calculation without Dask. We will calculate the D4 correction for 100 bulk Cu structures as a toy example. On the tiger-arrk login node, it should take about 1.5 minutes — not very long, but this is merely a toy example.
import numpy as np
from ase import Atoms
from ase.build import bulk
from dftd4.ase import DFTD4

def run_D4(atoms: Atoms) -> tuple[float, np.ndarray, np.ndarray]:
    atoms.calc = DFTD4(method="r2SCAN")
    E = atoms.get_potential_energy()
    F = atoms.get_forces()
    S = atoms.get_stress()
    return E, F, S

atoms_list: list[Atoms] = [bulk("Cu")*(4,4,4) for _ in range(100)]
results = [run_D4(atoms) for atoms in atoms_list]

Scaling Up Locally

Let's go ahead and use Dask to scale this up. To start, we will simply parallelize over all available CPU cores on whatever machine we are using. Since we have our own dedicated login node on tiger-arrk, we can use that and run up to 48 calculations concurrently since it is a 48-core node and these calculations are not going to bother other users sharing the login node. You can also use your laptop, which has multiple cores.
The full workflow looks like the following and should take approximately 13 seconds. That's not a 48x speedup because the tasks themselves are very short and there is some overhead for distributing the tasks, but it is still a significant speedup.
import numpy as np
from ase import Atoms
from ase.build import bulk
from dask.distributed import Client, LocalCluster
from dftd4.ase import DFTD4

def run_D4(atoms: Atoms) -> tuple[float, np.ndarray, np.ndarray]:
    atoms = atoms.copy() # Atoms objects are mutable, so we be careful
    atoms.calc = DFTD4(method="r2SCAN")
    E = atoms.get_potential_energy()
    F = atoms.get_forces()
    S = atoms.get_stress()
    return E, F, S

cluster = LocalCluster()
client = Client(cluster)

atoms_list: list[Atoms] = [bulk("Cu")*(4,4,4) for _ in range(100)]
futures = [client.submit(run_D4, atoms) for atoms in atoms_list]
results = client.gather(futures)

client.close()
cluster.close()
LLMs like Claude know all about Dask, so this can be useful for troubleshooting and setting up Dask execution tasks.
You will see that the code is essentially the same as before except that we spin up a local Dask cluster, submit tasks to the cluster, and gather the results. The results will be returned much faster now.

Scaling Up On Slurm

In the previous example, we parallelized the single-core tasks across our local computing cores. However, some tasks may require larger amounts of resources and are better to run on Slurm. This can also be done. The only thing that changes is how we instantiate the client.
The example below will request a single Slurm job on Tiger and parallelize the single-core compute tasks across all 112 cores of the node. It is helpful to run this via a Jupyter Notebook so that you can iterate on the Slurm submission script as needed, hence why I have broken it up into two code blocks.
The Slurm job launched with the call to SLURMCluster() will remain active and waiting for work. To stop the Slurm job entirely, call client.shutdown() or manually scancel the job ID.
from dask.distributed import Client
from dask_jobqueue import SLURMCluster

slurm_jobs = 1
cores_per_node = 112
job_mem = "512G"
cluster_kwargs = {
    "cores": cores_per_node,
    "memory": job_mem,
    "shebang": "#!/bin/bash",
    "account": "rosengroup",
    "walltime": "00:00:30",
    "interface": "ib0",
    "job_cpu": cores_per_node,
    "job_mem": job_mem,
    "job_script_prologue": [
        "source ~/.bashrc",
        "module load anaconda3/2025.12",
        "conda activate cms",
    ]
}
cluster = SLURMCluster(**cluster_kwargs)
print(cluster.job_script())

cluster.scale(jobs=slurm_jobs)
client = Client(cluster)
And the code otherwise remains the same as before, using the updated Dask cluster and client:
import numpy as np
from ase import Atoms
from ase.build import bulk
from dask.distributed import Client, LocalCluster
from dftd4.ase import DFTD4

def run_D4(atoms: Atoms) -> tuple[float, np.ndarray, np.ndarray]:
    atoms = atoms.copy() # Atoms objects are mutable, so we be careful
    atoms.calc = DFTD4(method="r2SCAN")
    E = atoms.get_potential_energy()
    F = atoms.get_forces()
    S = atoms.get_stress()
    return E, F, S

client = Client(cluster)

atoms_list: list[Atoms] = [bulk("Cu")*(4,4,4) for _ in range(100)]
futures = [client.submit(run_D4, atoms) for atoms in atoms_list]
results = client.gather(futures)

client.close()
cluster.close()

Simple Workflow Example

The above example involved parallelizing a single task that is called many times. If you have a more complex workflow with multiple steps and dependencies, it is better to use a feature called Dask delayed. In the example below, we will consider a toy example where we do call addition functions whose outputs are added together.

Local Parallelization

Consider the following toy example. Normally, you would expect this calculation to take a total of 10 seconds since there are two sequential 5-second add functions being called. We will pretend that the add function is a surrogate for some compute-heavy task (e.g. a DFT calculation).
import time

def add(a, b):
    time.sleep(5)
    return a + b

def workflow(a, b):
    output1 = add(a, b)
    output2 = add(a, b)
    return output1 + output2

result = workflow(1, 2)
When using Dask, the above code could be run in only 5 seconds by parallelizing the work over two CPU cores. This is achieved by decorating the compute task with the @delayed decorator.
import time
from dask import delayed
from dask.distributed import Client, LocalCluster

@delayed
def add(a, b):
    time.sleep(5)
    return a + b

def workflow(a, b):
    output1 = add(a, b)
    output2 = add(a, b)
    return output1 + output2

cluster = LocalCluster()
client = Client(cluster)

future = client.compute(workflow(1, 2))
result = future.result()

client.close()
cluster.close()
The code will now finish in 5 seconds instead of 10. Note that add(a, b) now returns a Delayed object rather than the actual result. To fetch the result, one must call .result() on the Delayed object — but note that this "blocks" the execution and Python will wait until the value is returned before continuing, so be careful where you call it. Dask will automatically handle task dependencies. For instance, if output1 is passed to add(output1, b), Dask will know not to call the second add() task until the first one is finished.
Note that if you have a list of Delayed objects being returned, you can easily convert them all to results via client.gather(futures).

Slurm Parallelization

Like before, we can also parallelize via a Slurm job.
The Slurm job launched with the call to SLURMCluster() will remain active and waiting for work. To stop the Slurm job entirely, call client.shutdown() or manually scancel the job ID.
from dask.distributed import Client
from dask_jobqueue import SLURMCluster

slurm_jobs = 1
cores_per_node = 112  # change as needed
job_mem = "512G"
cluster_kwargs = {
    "cores": cores_per_node,
    "memory": job_mem,
    "shebang": "#!/bin/bash",
    "account": "rosengroup",
    "walltime": "00:00:30",
    "interface": "ib0",
    "job_cpu": cores_per_node,
    "job_mem": job_mem,
    "job_script_prologue": [
        "source ~/.bashrc",
        "module load anaconda3/2025.12",
        "conda activate cms",
    ]
}
cluster = SLURMCluster(**cluster_kwargs)
print(cluster.job_script())

cluster.scale(jobs=slurm_jobs)
client = Client(cluster)
Then you can run the following code, which will dispatch the add() tasks across the compute node such that it finishes in 5 seconds instead of 10. To launch the calculation, the client.compute() function must be used. A Future is returned, which can be converted to its true value via a blocking .result() call. The @delayed decorator allows for asynchronous execution of the task.
import time
from dask import delayed

@delayed
def add(a, b):
    time.sleep(5)
    return a + b

def workflow(a, b):
    output1 = add(a, b)
    output2 = add(a, b)
    return output1 + output2

future = client.compute(workflow(1, 2))
result = future.result()