Usability: Make it easy to run Python-based codes

[sphuber]

Motivation

AiiDA is mostly used in the chemistry and materials science domains, where Python is very popular. Many users, therefore, will have Python-based codes that they want to use through or integrate with AiiDA.

If the code can be used as a library, the calcfunction concept makes it easy to use the code in AiiDA while keeping the provenance. However, the calcfunction has limitations. It can only be run on the machine where AiiDA runs, and so for heavy operations, this solution can overload the daemon workers. In addition, the calcfunctions execution does not support checkpointing, meaning the execution has to be finished in one go and cannot be interrupted to be restarted later. This poses a problem for code that has a long execution time.

The solution to both these problems is to run the code as a standalone script and run it on another machine as a calculation job. However, this requires writing a dedicated CalcJob plugin, and potentially a Parser plugin to parse the results, which is complicated and takes a lot of time. In addition, these Python scripts are often "moving targets" where the interface changes fast. This means the plugins have to change with it, essentially doubling the development cost and severely slowing down the speed of development.

This limitation often make AiiDA too restrictive to use for many use-cases where code is still in development and developing CalcJob and Parser plugins are simply not worth time.

Desired Outcome

There should be a way to allow users to easily run Python based codes without the requirements of writing dedicated plugins.

Impact

Since Python is very popular in the field where most AiiDA users are active, but also in general in computational science, making it easier to run these codes with AiiDA would greatly increase adoption and improve the efficiency of users.

Complexity

This section should comment on the complexity of resolving the item, and make a rough estimate as to the amount of work/time it would take to reach the desired outcome.

Background

This issue is very closely related to roadmap item #5 which discusses making it easy to run "codes" (or executables) without requiring dedicated plugins. In this case, the code is Python itself and the relevant script would be a command line argument.

Progress

As mentioned in the previous section, this use-case is strongly related to #5 which already has a concrete solution in the form of aiida-shell. This plugin package also makes it easy to run arbitrary Python scripts on any computer configured in AiiDA and it is already being used in production.

As an example, consider the following script:

import numpy as np
from scipy import linalg

matrix = np.load('matrix.npy', allow_pickle=True)
inverse = linalg.inv(matrix)
np.save('inverse.npy', inverse)

It loads a numpy matrix from the file matrix.npy, computes the inverse and writes the result to the file inverse.npy. It is currently not possible to have AiiDA run this on another computer, except if a dedicated CalcJob plugin is written first.

Now consider how aiida-shell actually makes this possible. Imagine that the numpy script is written in the current working directory as numpy_script.py. The following is a complete script that works with aiida-shell v0.4.0:

#!/usr/bin/env runaiida
import io
import numpy as np
from aiida.orm import SinglefileData, load_computer
from aiida_shell import launch_shell_job

# Write the contents of the input matrix to be inverted to stream
stream = io.BytesIO()
np.save(stream, np.array([[1, 3, 5], [2, 5, 1], [2, 3, 8]]))
stream.seek(0)

# Define a custom parser on-the-fly to convert the output file into an `ArrayData` node
def parser(self, dirpath):
    """Parse the content of the ``inverse.npy`` file and return as an ``ArrayData`` node."""
    import numpy as np
    from aiida.orm import ArrayData
    array = ArrayData()
    array.set_array('inverse', np.load(dirpath / 'inverse.npy'))
    return {'inverse': array}

# Launch a "shell" job to the computer `localhost`
results, node = launch_shell_job(
    'python',  # The executable is the Python interpreter
    arguments=['{script}'],  # Only command line argument is the input script
    nodes={
        'script': 'numpy_script.py',  # Automatically convert the `numpy_script.py` file to a `SinglefileData`
        'matrix': SinglefileData(stream, filename='matrix.npy')  # Write the matrix dump file to the working directory of the job
    },
    parser=parser,  # Attach the custom parser
    metadata={
        'computer': load_computer('localhost'),  # Submit to `localhost`, could be replaced with any other computer
        'options': {
            'additional_retrieve': ['inverse.npy']  # Retrieve the file that `numpy_script.py` will write to the working directory of the job
        }
    }
)
print(results['inverse'].get_array('inverse'))

With a few lines of Python, the script can be run (on any computer configured in AiiDA that has Python installed) without having to write a CalcJob plugin. The custom parser is even fully optional. Without it, the inverse.npy would automatically be attached as a SinglefileData. But since it is useful to have it as an ArrayData, it makes sense to use the on-the-fly parser to change the data type of the output.

Note that in this example, the job was "run", however, it can just as easily be submitted to the daemon by passing submit=True to the launch_shell_job function.

[ltalirz]

I think https://github.com/microsoft/aiida-dynamic-workflows is another great solution to this problem.

Its approach is simply to treat Python Code as Data, and execute it via the Python interpreter on the remote machine.
Makes it much easier to play around and explore, directly from Jupyter notebooks.

[sphuber]

Good point, I really need to find some time to give aiida-dynamic-workflows a try some day. What I couldn't find from the documentation or examples is whether the compute environment needs access to the AiiDA profile? That is to say, can you submit a Python snippet to a remote cluster that in the code access data from AiiDA's storage through its API? None of the examples actually use AiiDA's ORM or accept Data nodes as arguments.

If that is the case, then this would of course restrict the compute environments that are compatible. Most computing environments have Python installed, but don't necessarily have access to the same AiiDA environment.

Of course one could then only submit code that doesn't use AiiDA's API, but that means you also have to pass non Data instances as arguments to the function steps. This would cause provenance to be lost.

[ltalirz]

What I couldn't find from the documentation or examples is whether the compute environment needs access to the AiiDA profile? That is to say, can you submit a Python snippet to a remote cluster that in the code access data from AiiDA's storage through its API?

The remote compute environment does not have access to the AiiDA profile.

Of course one could then only submit code that doesn't use AiiDA's API, but that means you also have to pass non Data instances as arguments to the function steps. This would cause provenance to be lost.

It includes automatic type conversions for calculation/workflow inputs (where this is possible)
https://github.com/microsoft/aiida-dynamic-workflows/blob/4d452ed3be4192dc5b2c8f40690f82c3afcaa7a8/aiida_dynamic_workflows/data.py#L355-L356

(inside the decorated Python functions, AiiDA objects are of course not allowed).

[JPchico]

This seems quite interesting, I can see use cases for some software packages, e.g. thermocalc where one can use a python interface to submit calculations. Writing a full plugin for this might be not desirable as the python interface already exists.

AiiDA making it more flexible to make use of these codes would make it much easier to use for codes which are not DFT/MD which has been the focus until now.