Run calculations remotely

Introduction

There are three ways to run calculations remotely using AiiDA:

• CalcJob: run a shell command (via dedicated AiiDA plugins to manage inputs, outputs) on a remote computer (see Interoperate with aiida-core components).

• ShellJob: run any shell command on a remote computer (see Run shell commands as a task).

• PythonJob: run any Python function on a remote computer.

In this tutorial, we introduce how to run PythonJob tasks inside a WorkGraph. For more details on the aiida-pythonjob package itself, please refer to the AiiDA-PythonJob documentation.

from aiida import load_profile

from aiida_workgraph import task

load_profile()

Profile<uuid='2c6af82ef2b241408a4c27c7c32d5490' name='presto'>

PythonJob

One can create a PythonJob task using the @task.pythonjob decorator. This allows you to define a Python function that can be executed as a job on a remote computer.

@task.pythonjob
def add(x: int, y: int) -> int:
    return x + y


@task.graph
def RemoteAdd(x: int, y: int, computer: str) -> int:
    return add(x=x, y=y, computer=computer).result


wg = RemoteAdd.build(x=1, y=2, computer='localhost')
wg.run()

print('\nResult: ', wg.outputs.result.value)

Result:  uuid: b1e415b3-83d8-4cd2-9c1b-8a81a963c2f1 (pk: 239) value: 3

Understanding inputs and outputs

When you apply the @task.pythonjob decorator, your Python function (like add) transforms into an AiiDA PythonJob task. This task requires and produces more than just your function’s direct inputs and outputs:

• Inputs: In addition to your function’s arguments (e.g., x, y), the task takes additional inputs (e.g., computer) to manage its execution on a remote machine.

• Outputs: In addition to your function’s return value (available as .result by default), the task provides comprehensive outputs, such as remote_folder (a reference to the job’s directory on the remote computer).

Prepare Python environment on remote computer

PythonJob requires that the Python version on the remote computer matches the one used on the localhost where AiiDA is running. You can use conda to create a virtual environment with the same Python version, then activate the environment in the metadata of the PythonJob task.

Here’s an example demonstrating how to submit a PythonJob and chain it with another Python function within a WorkGraph. We’ll also show how to pass custom scheduler commands via metadata, though for this example, they will be empty.

metadata = {
    'options': {
        # "custom_scheduler_commands" : "module load anaconda\nconda activate py3.11\n",
        'custom_scheduler_commands': '',  # Keeping it empty for this example
    }
}

Note

Uncomment the custom_scheduler_commands line to modify it for your environment.

from typing import Annotated


@task
def multiply(x, y):
    return x * y


@task.graph
def RemoteAddLocalMultiply(x: int, y: int, computer: str, add_metadata: Annotated[dict, add.inputs.metadata]) -> int:
    the_sum = add(x=x, y=y, computer=computer, metadata=add_metadata).result
    return multiply(x=the_sum, y=4)  # this will run locally


wg = RemoteAddLocalMultiply.build(
    x=2,
    y=3,
    computer='localhost',
    add_metadata=metadata,
)
wg.run()

print('\nResult:', wg.outputs.result.value)

Result: uuid: c3f5d21c-2e5a-4f8f-abfb-b8c8de41d86c (pk: 252) value: 20

Conclusion

In this tutorial, we briefly discussed the three ways to run calculations remotely using AiiDA: CalcJob, ShellJob, and PythonJob. We demonstrated how to use PythonJob to define and execute Python functions as remote jobs, including managing the remote Python environment via custom scheduler commands.