PythonJob

Preparing inputs for PythonJob

The prepare_pythonjob_inputs function is available for setting up the inputs for a PythonJob calculation. This function simplifies the process of preparing and serializing data, and configuring the execution environment.

Code: You can specify the computer where the job will run, which will create a python3@computer code if it doesn’t already exist. Alternatively, if the code has already been created, you can set the code directly.
Python Version: Ensure the Python version on the remote computer matches the local environment. This is important since pickle is used for data storage and retrieval. Use conda to create and activate a virtual environment with the same Python version. Pass metadata to the scheduler to activate the environment during the job execution:
```
metadata = {
    "options": {
        'custom_scheduler_commands': 'module load anaconda\nconda activate py3.11\n',
    }
}
```

Create a conda environment on the remote computer

One can use the create_conda_env function to create a conda environment on the remote computer. The function will create a conda environment with the specified packages and modules. The function will update the packages if the environment already exists.

from aiida_pythonjob.utils import create_conda_env
# create a conda environment on remote computer
create_conda_env(
       "merlin6",                # Remote computer
       "test_pythonjob",         # Name of the conda environment
       modules=["anaconda"],     # Modules to load (e.g., Anaconda)
       pip=["numpy", "matplotlib"],  # Python packages to install via pip
       conda={                   # Conda-specific settings
           "channels": ["conda-forge"],  # Channels to use
           "dependencies": ["qe"]       # Conda packages to install
       }
   )

Annotating outputs and handling data

PythonJob supports advanced features for structuring outputs, handling custom data types, and defining exit codes for error handling. These functionalities are shared with pyfunction and are detailed in a separate guide.

Using parent folder

The parent_folder parameter allows a task to access the output files of a parent task. This feature is particularly useful when you want to reuse data generated by a previous computation in subsequent computations. In the following example, the multiply task uses the result.txt file created by the add task.

def add(x, y):
    z = x + y
    with open("result.txt", "w") as f:
        f.write(str(z))
    return x + y


def multiply(x, y):
    with open("parent_folder/result.txt", "r") as f:
        z = int(f.read())
    return x * y + z


inputs1 = prepare_pythonjob_inputs(
    add,
    function_inputs={"x": 1, "y": 2},
)

result1, node1 = run_get_node(PythonJob, inputs=inputs1)

inputs2 = prepare_pythonjob_inputs(
    multiply,
    function_inputs={"x": 1, "y": 2},
    parent_folder=result1["remote_folder"],
)

result2, node2 = run_get_node(PythonJob, inputs=inputs2)

print("result: ", result2)

result:  {'remote_folder': <RemoteData: uuid: 00642d22-0073-4df5-be25-40b3995fb39c (pk: 49)>, 'retrieved': <FolderData: uuid: 6bc6cc32-27af-4f1c-abd4-72d47489de07 (pk: 50)>, 'result': <Int: uuid: 5160efdf-fc46-4092-862f-5663df06cd2e (pk: 51) value: 5>}

Upload files or folders to the remote computer

The upload_files parameter allows users to upload files or folders to the remote computer. The files will be uploaded to the working directory of the remote computer.

import os  # noqa: E402

# create a temporary file "input.txt" in the current directory
with open("/tmp/input.txt", "w") as f:
    f.write("2")

# create a temporary folder "inputs_folder" in the current directory
# and add a file "another_input.txt" in the folder
os.makedirs("/tmp/inputs_folder", exist_ok=True)
with open("/tmp/inputs_folder/another_input.txt", "w") as f:
    f.write("3")


def add():
    with open("input.txt", "r") as f:
        a = int(f.read())
    with open("inputs_folder/another_input.txt", "r") as f:
        b = int(f.read())
    return a + b


# ------------------------- Submit the calculation -------------------
# we need use full path to the file
input_file = os.path.abspath("/tmp/input.txt")
input_folder = os.path.abspath("/tmp/inputs_folder")
inputs = prepare_pythonjob_inputs(
    add,
    upload_files={
        "input.txt": input_file,
        "inputs_folder": input_folder,
    },
)
result, node = run_get_node(PythonJob, inputs=inputs)
print("result: ", result["result"])

result:  uuid: f2bcfaf9-b671-4798-aa00-bad482d53062 (pk: 58) value: 5

Retrieve additional files from the remote computer

Sometimes, one may want to retrieve additional files from the remote computer after the job has finished. For example, one may want to retrieve the output files generated by the pw.x calculation in Quantum ESPRESSO.

One can use the additional_retrieve_list parameter to specify which files should be retrieved from the working directory and stored in the local repository after the job has finished

def add(x, y):
    z = x + y
    with open("result.txt", "w") as f:
        f.write(str(z))
    return x + y


inputs = prepare_pythonjob_inputs(
    add,
    function_inputs={"x": 1, "y": 2},
    metadata={
        "options": {
            "additional_retrieve_list": ["result.txt"],
        }
    },
)

result, node = run_get_node(PythonJob, inputs=inputs)
print("retrieved files: ", result["retrieved"].list_object_names())

retrieved files:  ['_error.json', '_scheduler-stderr.txt', '_scheduler-stdout.txt', 'aiida.out', 'result.txt', 'results.pickle']

Using MPI for parallel execution

If your function uses MPI for parallel execution, you can set the withmpi option to True in the metadata.options dictionary. This will ensure that the job is executed using mpirun or an equivalent command on the (remote) computer.

metadata = {
    "options": {
        'withmpi': True,
    }
}

Note: In order to run MPI jobs, the remote computer must have installed the mpi4py package in the Python environment.

What if my calculation fails?

The PythonJobParser can return specialized exit codes when different kinds of errors occur during the calculation:

ERROR_READING_OUTPUT_FILE (310):
The retrieved output file (e.g., results.pickle) could not be opened or read.
ERROR_INVALID_OUTPUT (320):
The output file is corrupt or contains unexpected/invalid data structures.
ERROR_RESULT_OUTPUT_MISMATCH (321):
The number of actual results does not match the number/structure of expected outputs.
ERROR_IMPORT_CLOUDPICKLE_FAILED (322):
The script on the remote machine failed to import cloudpickle. The script writes error.json describing the ImportError.
ERROR_UNPICKLE_INPUTS_FAILED (323):
The script failed to unpickle the input data (e.g., inputs.pickle).
ERROR_UNPICKLE_FUNCTION_FAILED (324):
The script failed to load the pickled function (e.g., function.pkl).
ERROR_FUNCTION_EXECUTION_FAILED (325):
An exception was raised during the function call.
ERROR_PICKLE_RESULTS_FAILED (326):
The script failed to pickle (serialize) the final results.
ERROR_SCRIPT_FAILED (327):
A catch-all exit code if none of the above match. Indicates an unknown/unrecognized error.

Using register_pickle_by_value

If the function is defined inside an external module that is not installed on the remote computer, this can cause import errors during execution.

Solution: By enabling register_pickle_by_value=True, the function is serialized by value instead of being referenced by its module path. This embeds the function unpickled even if the original module is unavailable on the remote computer.

Example:

inputs = prepare_pythonjob_inputs(
    my_function,
    function_inputs={"x": 1, "y": 2},
    computer="localhost",
    register_pickle_by_value=True,  # Ensures function is embedded
)

Important Considerations:: If the function contains import statements, the imported modules must still be installed on the remote computer.

What’s Next

Tutorials

Real-world examples in computational materials science and more.

Total running time of the script: (0 minutes 10.189 seconds)

Gallery generated by Sphinx-Gallery