How to set up a large scale project with Taskblaster using ASE (Atomic Simulation Environment)?

How to set up a large scale project with Taskblaster using ASE (Atomic Simulation Environment)?#

Read first How to set up a large scale project with Taskblaster?, which will explain the basic building blocks of a Python project. Here, we will focus on creating a new project, which contains those parts, but also adds ASE encoders and decoders, such that we store ASE objects (Atoms, BandStructure, etc.)

Run this script which will create the project files, and a virtual environement, and install the project to that virtual environment.

$ chmod +x setup_full_project_ase.sh && ./setup_full_project_ase.sh
Obtaining file:///home/me/repo/asedemo_venv/asedemo
  Installing build dependencies: started
  Installing build dependencies: finished with status 'done'
  Checking if build backend supports build_editable: started
  Checking if build backend supports build_editable: finished with status 'done'
  Getting requirements to build editable: started
  Getting requirements to build editable: finished with status 'done'
  Preparing editable metadata (pyproject.toml): started
  Preparing editable metadata (pyproject.toml): finished with status 'done'
Collecting taskblaster (from asedemo==0.1)
  Using cached taskblaster-0.2-py3-none-any.whl.metadata (1.4 kB)
Collecting ase (from asedemo==0.1)
  Downloading ase-3.28.0-py3-none-any.whl.metadata (4.3 kB)
Collecting numpy>=1.21.6 (from ase->asedemo==0.1)
  Downloading numpy-2.4.4-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.metadata (6.6 kB)
Collecting scipy>=1.8.1 (from ase->asedemo==0.1)
  Downloading scipy-1.17.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.metadata (62 kB)
Collecting matplotlib>=3.5.2 (from ase->asedemo==0.1)
  Downloading matplotlib-3.10.8-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (52 kB)
Collecting click>=8.0.0 (from taskblaster->asedemo==0.1)
  Using cached click-8.3.2-py3-none-any.whl.metadata (2.6 kB)
Collecting contourpy>=1.0.1 (from matplotlib>=3.5.2->ase->asedemo==0.1)
  Downloading contourpy-1.3.3-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.metadata (5.5 kB)
Collecting cycler>=0.10 (from matplotlib>=3.5.2->ase->asedemo==0.1)
  Downloading cycler-0.12.1-py3-none-any.whl.metadata (3.8 kB)
Collecting fonttools>=4.22.0 (from matplotlib>=3.5.2->ase->asedemo==0.1)
  Downloading fonttools-4.62.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (117 kB)
Collecting kiwisolver>=1.3.1 (from matplotlib>=3.5.2->ase->asedemo==0.1)
  Downloading kiwisolver-1.5.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (5.1 kB)
Collecting packaging>=20.0 (from matplotlib>=3.5.2->ase->asedemo==0.1)
  Downloading packaging-26.1-py3-none-any.whl.metadata (3.5 kB)
Collecting pillow>=8 (from matplotlib>=3.5.2->ase->asedemo==0.1)
  Downloading pillow-12.2.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.metadata (8.8 kB)
Collecting pyparsing>=3 (from matplotlib>=3.5.2->ase->asedemo==0.1)
  Downloading pyparsing-3.3.2-py3-none-any.whl.metadata (5.8 kB)
Collecting python-dateutil>=2.7 (from matplotlib>=3.5.2->ase->asedemo==0.1)
  Downloading python_dateutil-2.9.0.post0-py2.py3-none-any.whl.metadata (8.4 kB)
Collecting six>=1.5 (from python-dateutil>=2.7->matplotlib>=3.5.2->ase->asedemo==0.1)
  Downloading six-1.17.0-py2.py3-none-any.whl.metadata (1.7 kB)
Downloading ase-3.28.0-py3-none-any.whl (2.9 MB)
   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 2.9/2.9 MB 49.4 MB/s eta 0:00:00
Using cached taskblaster-0.2-py3-none-any.whl (122 kB)
Using cached click-8.3.2-py3-none-any.whl (108 kB)
Downloading matplotlib-3.10.8-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (8.7 MB)
   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 8.7/8.7 MB 141.0 MB/s eta 0:00:00
Downloading numpy-2.4.4-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (16.6 MB)
   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 16.6/16.6 MB 280.7 MB/s eta 0:00:00
Downloading scipy-1.17.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (35.2 MB)
   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 35.2/35.2 MB 279.1 MB/s eta 0:00:00
Downloading contourpy-1.3.3-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (362 kB)
Downloading cycler-0.12.1-py3-none-any.whl (8.3 kB)
Downloading fonttools-4.62.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (5.0 MB)
   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 5.0/5.0 MB 246.2 MB/s eta 0:00:00
Downloading kiwisolver-1.5.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (1.5 MB)
   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 1.5/1.5 MB 173.5 MB/s eta 0:00:00
Downloading packaging-26.1-py3-none-any.whl (95 kB)
Downloading pillow-12.2.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (7.1 MB)
   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 7.1/7.1 MB 243.4 MB/s eta 0:00:00
Downloading pyparsing-3.3.2-py3-none-any.whl (122 kB)
Downloading python_dateutil-2.9.0.post0-py2.py3-none-any.whl (229 kB)
Downloading six-1.17.0-py2.py3-none-any.whl (11 kB)
Building wheels for collected packages: asedemo
  Building editable for asedemo (pyproject.toml): started
  Building editable for asedemo (pyproject.toml): finished with status 'done'
  Created wheel for asedemo: filename=asedemo-0.1-0.editable-py3-none-any.whl size=2640 sha256=8403db6f9ad235b6befcaef8339191a99aca6775b194c05f00c5f418a6baafad
  Stored in directory: /tmp/pip-ephem-wheel-cache-zk4wmhv5/wheels/04/b7/b9/cd5573ffb870ddc60099d37e4f78dab4ea91b69797658f0d30
Successfully built asedemo
Installing collected packages: six, pyparsing, pillow, packaging, numpy, kiwisolver, fonttools, cycler, click, taskblaster, scipy, python-dateutil, contourpy, matplotlib, ase, asedemo
Successfully installed ase-3.28.0 asedemo-0.1 click-8.3.2 contourpy-1.3.3 cycler-0.12.1 fonttools-4.62.1 kiwisolver-1.5.0 matplotlib-3.10.8 numpy-2.4.4 packaging-26.1 pillow-12.2.0 pyparsing-3.3.2 python-dateutil-2.9.0.post0 scipy-1.17.1 six-1.17.0 taskblaster-0.2

You should have following files now in a package called asedemo in your virtual environment.

$ cd asedemo_venv/asedemo && find . | grep -v "egg-info"
.
./asedemo
./asedemo/main_workflow.py
./asedemo/mysubmodule
./asedemo/mysubmodule/tasks.py
./asedemo/mysubmodule/__init__.py
./asedemo/__init__.py
./pyproject.toml

There is a Python module called asedemo in the folder with the same name. We can now initialize a Taskblaster repo using tb init asedemo. We can do this command from anywhere, since we always can locate the installed module asedemo.

Warning

Before writing tb init asedemo make sure to go to a fresh folder (do not create a Taskblaster repository to your home folder, or to your virtual environment asedemo_venv-folder).

We can now activate the virtual environment.

source asedemo_venv/bin/activate
$ tb init asedemo
Created repository using module "asedemo" in "/home/me/repo".

We can now execute our workflow, which now executes a simple workflow in various folders.

$ tb workflow asedemo_venv/asedemo/asedemo/main_workflow.py
entry:                    add  new      0/0        tree/Cu/atoms 
                          add  new      0/1        tree/Cu/relaxed_atoms 
entry:                    add  new      0/0        tree/Ag/atoms 
                          add  new      0/1        tree/Ag/relaxed_atoms 
entry:                    add  new      0/0        tree/Fe/atoms 
                          add  new      0/1        tree/Fe/relaxed_atoms 
entry:                    add  new      0/0        tree/Ni/atoms 
                          add  new      0/1        tree/Ni/relaxed_atoms 

we can now run the tasks, and we observe the relaxation of atoms.

$ tb run .
Starting worker rank=000 size=001
[rank=000 2026-04-17 12:59:00 N/A-0/1] Worker class: —
[rank=000 2026-04-17 12:59:00 N/A-0/1] Required tags: —
[rank=000 2026-04-17 12:59:00 N/A-0/1] Supported tags: —
[rank=000 2026-04-17 12:59:00 N/A-0/1] name: None
    tags: —
    required_tags: —
    resources: None
    max_tasks: None
    subworker_size: None
    subworker_count: None
    wall_time: None
[rank=000 2026-04-17 12:59:00 N/A-0/1] Main loop
[rank=000 2026-04-17 12:59:01 N/A-0/1] Running Ag/atoms ...
[rank=000 2026-04-17 12:59:01 N/A-0/1] Task Ag/atoms finished in 0:00:00.016511
[rank=000 2026-04-17 12:59:01 N/A-0/1] Running Ag/relaxed_atoms ...

               Step     Time          Energy             fmax            smax          volume
CellAwareBFGS:    0 12:59:01        0.974356         0.000000        0.082643       29.556545
CellAwareBFGS:    1 12:59:01        0.752733         0.000000        0.088820       26.977125
CellAwareBFGS:    2 12:59:01        0.090994         0.000000        0.104051       14.805360
CellAwareBFGS:    3 12:59:01        0.056929         0.000000        0.052913       18.723637
CellAwareBFGS:    4 12:59:01        0.001306         0.000000        0.010788       17.079643
CellAwareBFGS:    5 12:59:01       -0.000320         0.000000        0.001882       16.724754
[rank=000 2026-04-17 12:59:01 N/A-0/1] Task Ag/relaxed_atoms finished in 0:00:00.068191
[rank=000 2026-04-17 12:59:01 N/A-0/1] Running Cu/atoms ...
[rank=000 2026-04-17 12:59:01 N/A-0/1] Task Cu/atoms finished in 0:00:00.000620
[rank=000 2026-04-17 12:59:01 N/A-0/1] Running Cu/relaxed_atoms ...

               Step     Time          Energy             fmax            smax          volume
CellAwareBFGS:    0 12:59:01        0.969364         0.000000        0.130132       20.323821
CellAwareBFGS:    1 12:59:01        0.605548         0.000000        0.136057       17.601865
CellAwareBFGS:    2 12:59:01        0.168729         0.000000        0.212532        9.660108
CellAwareBFGS:    3 12:59:01        0.035427         0.000000        0.066366       12.742629
CellAwareBFGS:    4 12:59:01       -0.005787         0.000000        0.013094       11.753569
CellAwareBFGS:    5 12:59:01       -0.007020         0.000000        0.001545       11.544196
[rank=000 2026-04-17 12:59:01 N/A-0/1] Task Cu/relaxed_atoms finished in 0:00:00.076775
[rank=000 2026-04-17 12:59:01 N/A-0/1] Running Fe/atoms ...
[rank=000 2026-04-17 12:59:01 N/A-0/1] Task Fe/atoms finished in 0:00:00.000563
[rank=000 2026-04-17 12:59:01 N/A-0/1] Running Fe/relaxed_atoms ...
[rank=000 2026-04-17 12:59:01 N/A-0/1] Traceback (most recent call last):
  File "/home/me/repo/asedemo_venv/lib/python3.12/site-packages/taskblaster/worker.py", line 497, in process_one_task
    loaded_task.run(self)
  File "/home/me/repo/asedemo_venv/lib/python3.12/site-packages/taskblaster/worker.py", line 218, in run
    output = function(**kwargs)
             ^^^^^^^^^^^^^^^^^^
  File "/home/me/repo/asedemo_venv/asedemo/asedemo/mysubmodule/tasks.py", line 16, in relax
    relax.run()
  File "/home/me/repo/asedemo_venv/lib/python3.12/site-packages/ase/optimize/cellawarebfgs.py", line 110, in run
    return Dynamics.run(self)
           ^^^^^^^^^^^^^^^^^^
  File "/home/me/repo/asedemo_venv/lib/python3.12/site-packages/ase/optimize/optimize.py", line 375, in run
    for converged in Dynamics.irun(self, steps=steps):
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/me/repo/asedemo_venv/lib/python3.12/site-packages/ase/optimize/optimize.py", line 322, in irun
    gradient = self.optimizable.get_gradient()
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/me/repo/asedemo_venv/lib/python3.12/site-packages/ase/filters.py", line 32, in get_gradient
    return -self.filterobj.get_forces().ravel()
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/me/repo/asedemo_venv/lib/python3.12/site-packages/ase/filters.py", line 618, in get_forces
    stress = self.atoms.get_stress(**kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/me/repo/asedemo_venv/lib/python3.12/site-packages/ase/atoms.py", line 2043, in get_stress
    stress = self._calc.get_stress(self)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/me/repo/asedemo_venv/lib/python3.12/site-packages/ase/calculators/abc.py", line 36, in get_stress
    return self.get_property('stress', atoms)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/me/repo/asedemo_venv/lib/python3.12/site-packages/ase/calculators/calculator.py", line 520, in get_property
    self.calculate(atoms, [name], system_changes)
  File "/home/me/repo/asedemo_venv/lib/python3.12/site-packages/ase/calculators/emt.py", line 179, in calculate
    self.initialize(self.atoms)
  File "/home/me/repo/asedemo_venv/lib/python3.12/site-packages/ase/calculators/emt.py", line 87, in initialize
    raise NotImplementedError(f'No EMT-potential for {sym}')
NotImplementedError: No EMT-potential for Fe

[rank=000 2026-04-17 12:59:01 N/A-0/1] Task Fe/relaxed_atoms failed: No EMT-potential for Fe
[rank=000 2026-04-17 12:59:01 N/A-0/1] WARDEN: Task Fe/relaxed_atoms has handlers [].
[rank=000 2026-04-17 12:59:01 N/A-0/1] WARDEN: No handlers found.
[rank=000 2026-04-17 12:59:01 N/A-0/1] Running Ni/atoms ...
[rank=000 2026-04-17 12:59:01 N/A-0/1] Task Ni/atoms finished in 0:00:00.000636
[rank=000 2026-04-17 12:59:01 N/A-0/1] Running Ni/relaxed_atoms ...

               Step     Time          Energy             fmax            smax          volume
CellAwareBFGS:    0 12:59:01        1.238133         0.000000        0.174489       18.841338
CellAwareBFGS:    1 12:59:01        0.636987         0.000000        0.184091       15.537429
CellAwareBFGS:    2 12:59:01        0.307421         0.000000        0.376803        8.527122
CellAwareBFGS:    3 12:59:01        0.041458         0.000000        0.090590       11.712732
CellAwareBFGS:    4 12:59:01       -0.010779         0.000000        0.022212       10.825183
CellAwareBFGS:    5 12:59:01       -0.013279         0.000000        0.002392       10.578098
[rank=000 2026-04-17 12:59:01 N/A-0/1] Task Ni/relaxed_atoms finished in 0:00:00.078761
[rank=000 2026-04-17 12:59:01 N/A-0/1] No available tasks, end worker main loop

We can check the outcome by listing the outputs. We on purpose included an error to task (missing EMT potential) to show how that would show up.

$ tb ls
state    info       tags        worker         time     folder
──────── ────────── ─────────── ─────────── ─────────── ─────────────────────────────
done     0/0                    N/A-0/1        00:00:01 tree/Ag/atoms
done     1/1                    N/A-0/1        00:00:00 tree/Ag/relaxed_atoms
done     0/0                    N/A-0/1        00:00:00 tree/Cu/atoms
done     1/1                    N/A-0/1        00:00:00 tree/Cu/relaxed_atoms
done     0/0                    N/A-0/1        00:00:00 tree/Fe/atoms
fail     1/1                    N/A-0/1        00:00:00 tree/Fe/relaxed_atoms
^^^^  NotImplementedError: No EMT-potential for Fe
done     0/0                    N/A-0/1        00:00:00 tree/Ni/atoms
done     1/1                    N/A-0/1        00:00:00 tree/Ni/relaxed_atoms

We can examine the output of the task returning the relaxed atoms using tb view command,

$ tb view tree/Ni/relaxed_atoms
name: Ni/relaxed_atoms
  location:        /home/me/repo/tree/Ni/relaxed_atoms
  state:           done
  target:          asedemo.mysubmodule.tasks.relax(…)
  wait for:        0 dependencies
  depth:           1
  source workflow: Ni
  frozen by: (not frozen)

  latest handled inputs:
     None

  handlers:
    []

  handler data:
    <None>

  parents:
    Ni/atoms

  input:
    ["asedemo.mysubmodule.tasks.relax", {"atoms": {"__tb_type__": "ref", "index": [], "name": "Ni/atoms"}}]

  output:
    {'relaxation_result': Atoms(symbols='Ni', pbc=True, cell=[[2.04825585111184e-15, 1.7423117537339736, 1.7423117537339718], [1.7423117537339685, 1.9669463159053335e-15, 1.7423117537339718], [1.7423117537339696, 1.742311753733974, 2.2549505358870145e-15]], initial_magmoms=...), 'total_steps': 5}

Run information:
    Worker name: N/A-0/1
    Start time: 2026-04-17 12:59:01
    End time: 2026-04-17 12:59:01
    Duration: 0:00:00
    Error: None

No custom actions defined for this task.

or we can see how the ase.Atoms objects are encoded into the output dictionary

$ cat tree/Ni/relaxed_atoms/output.json
{"relaxation_result": {"__ase_objtype__": "atoms", "cell": {"__ndarray__": [[3, 3], "float64", [2.04825585111184e-15, 1.7423117537339736, 1.7423117537339718, 1.7423117537339685, 1.9669463159053335e-15, 1.7423117537339718, 1.7423117537339696, 1.742311753733974, 2.2549505358870145e-15]]}, "initial_magmoms": {"__ndarray__": [[1], "float64", [0.6]]}, "numbers": {"__ndarray__": [[1], "int64", [28]]}, "pbc": {"__ndarray__": [[3], "bool", [true, true, true]]}, "positions": {"__ndarray__": [[1, 3], "float64", [4.4045923070222535e-18, -7.49624593758901e-17, 3.2911723159269653e-16]]}}, "total_steps": 5}