如何为外部代码编写插件#

小技巧

在开始编写新插件之前,请检查 aiida plugin registry。如果已经有适用于您代码的插件,您可以直接跳到 如何运行外部代码

小技巧

本手册引导你了解 AiiDA 如何与外部代码交互的所有逻辑步骤。如果你已经了解基础知识,并希望快速开始使用新的插件包,请查看 如何打包插件

要使用 AiiDA 运行外部代码,你需要一个相应的 计算 插件,告诉 AiiDA 如何运行:

  1. 准备所需的输入文件。

  2. 使用正确的命令行参数运行代码。

最后,你可能需要一个 解析器 插件,它告诉 AiiDA 如何处理:

  1. 解析代码输出。

本教程将带您了解 creating a calculation plugin 的使用过程,将其用于 run the codewriting a parser 的输出。

在本例中,我们的 AbstractCode 将是 diff 可执行文件,它可以 ``computes`` 两个 ``input files`` 之间的差值,并将差值打印到标准输出:

$ cat file1.txt
file with content
content1

$ cat file2.txt
file with content
content2

$ diff file1.txt file2.txt
2c2
< content1
---
> content2

我们在这里使用 diff,因为几乎所有 UNIX 系统都默认使用它,而且它既接受命令行参数 (两个文件),也接受命令行选项 (例如 -i,用于不区分大小写的匹配)。这与许多科学模拟代码的可执行文件的工作方式类似,因此很容易将本示例调整为适合您的使用情况。

我们将以如下方式运行 diff

$ diff file1.txt file2.txt > diff.patch

因此将 file1.txtfile2.txt 之间的差值写入 diff.patch

连接外部代码#

首先创建文件 calculations.py,然后子类化 CalcJob 类:

from aiida.common import datastructures
from aiida.engine import CalcJob
from aiida.orm import SinglefileData

class DiffCalculation(CalcJob):
    """AiiDA calculation plugin wrapping the diff executable."""

下面,我们将通过实现两个关键方法来告诉 AiiDA 如何运行我们的代码:

  1. define()

  2. prepare_for_submission()

定义规格#

define 方法告诉 AiiDA CalcJob 期望的输入和输出(退出代码将是 discussed later)。这是通过 CalcJobProcessSpec 类的实例完成的,该实例作为 spec 参数传递给 define 方法。例如

    @classmethod
    def define(cls, spec):
        """Define inputs and outputs of the calculation."""
        super(DiffCalculation, cls).define(spec)

        # new ports
        spec.input('file1', valid_type=SinglefileData, help='First file to be compared.')
        spec.input('file2', valid_type=SinglefileData, help='Second file to be compared.')
        spec.output('diff', valid_type=SinglefileData, help='diff between file1 and file2.')

        spec.input('metadata.options.output_filename', valid_type=str, default='patch.diff')
        spec.inputs['metadata']['options']['resources'].default = {
            'num_machines': 1,
            'num_mpiprocs_per_machine': 1,
        }
        spec.inputs['metadata']['options']['parser_name'].default = 'diff-tutorial'

        spec.exit_code(
            300, 'ERROR_MISSING_OUTPUT_FILES', message='Calculation did not produce all expected output files.'
        )

该方法的第一行调用 CalcJob 父类的 define 方法。这一必要步骤定义了所有 CalcJob 共用的 inputsoutputs

接下来,我们使用 input() 方法定义两个 SinglefileData 类型的输入文件 file1file2

更多阅读

当使用 SinglefileData 时,AiiDA 会以 文件 的形式记录输入。这非常灵活,但缺点是很难查询这些文件中的信息,也很难确保输入是有效的。练习 - 支持命令行选项 演示了如何使用 Dict 类将 diff 命令行选项表示为 python 字典。aiida-diff 演示插件则更进一步,增加了自动验证功能。

然后我们使用 output() 来定义计算的唯一输出,标签为 diff。AiiDA 会使用提供的链接标签将这里定义的输出附加到(成功)完成的计算中。

最后,我们设置了一些默认的 options,例如解析器的名称(我们稍后将实现)、输入和输出文件的名称,以及用于此类计算的计算资源。这些 options 已经通过 super().define(spec) 调用在 spec 上定义,可以通过 inputs 属性访问它们,其行为类似于字典。

define 中没有 return 语句:define 方法直接修改它接收到的 spec 对象。

备注

CalcJob 所需的另一项输入是使用哪个外部可执行文件。

外部可执行文件由 AbstractCode 实例表示,其中包含有关其所在计算机、文件系统路径等信息。外部可执行文件通过 code 输入传递给 CalcJob,该输入在 CalcJob 基类中定义,因此您不必这样做:

spec.input('code', valid_type=orm.AbstractCode, help='The `Code` to use for this job.')

更多阅读

有关设置 inputsoutputs 的详细信息(包括验证、动态输入数等),请参阅 Defining Processes 主题。

准备提交#

The prepare_for_submission() method has two jobs: Creating the input files in the format the external code expects and returning a CalcInfo object that contains instructions for the AiiDA engine on how the code should be run. For example:

    def prepare_for_submission(self, folder):
        """Create input files.

        :param folder: an `aiida.common.folders.Folder` where the plugin should temporarily place all files needed by
            the calculation.
        :return: `aiida.common.datastructures.CalcInfo` instance
        """
        codeinfo = datastructures.CodeInfo()
        codeinfo.cmdline_params = [self.inputs.file1.filename, self.inputs.file2.filename]
        codeinfo.code_uuid = self.inputs.code.uuid
        codeinfo.stdout_name = self.metadata.options.output_filename

        # Prepare a `CalcInfo` to be returned to the engine
        calcinfo = datastructures.CalcInfo()
        calcinfo.codes_info = [codeinfo]
        calcinfo.local_copy_list = [
            (self.inputs.file1.uuid, self.inputs.file1.filename, self.inputs.file1.filename),
            (self.inputs.file2.uuid, self.inputs.file2.filename, self.inputs.file2.filename),
        ]
        calcinfo.retrieve_list = [self.metadata.options.output_filename]

        return calcinfo

在调用 prepare_for_submission() 之前,所有提供给计算的输入都会根据 spec 进行验证。因此,在访问 inputs 属性时,可以放心地假设所有必需的输入都已设置,并且所有输入的类型都有效。

我们首先创建一个 CodeInfo 对象,让 AiiDA 知道如何运行代码,即这里:

$ diff file1.txt file2.txt > diff.patch

这包括命令行参数(此处:我们希望 diff 运行的文件名)和要运行的 AbstractCode 的 UUID。由于 diff 会直接写入标准输出,因此我们会将标准输出重定向到指定的输出文件名。

接下来,我们创建一个 CalcInfo 对象,让 AiiDA 知道哪些文件需要来回拷贝。在我们的例子中,两个输入文件已经存储在 AiiDA 文件库中,我们可以使用 local_copy_list 来传递它们。

备注

在其他使用情况下,您可能需要临时 创建 新文件。这就是 prepare_for_submission()folder 参数的作用:

with folder.open("filename", 'w') as handle:
    handle.write("file content")

在该沙盒文件夹中创建的任何文件和目录都将自动转移到进行实际计算的计算资源中。

另一方面,retrieve_list 则告诉 engine 在作业完成后从作业运行的目录中检索哪些文件。此处列出的所有文件都将存储在 FolderData node 中,作为标号为 retrieved 的输出 node 附加到计算中。

最后,我们将 CodeInfo 传递给一个 CalcInfo 对象。一个计算任务可能涉及多个可执行文件,因此 codes_info 是一个列表。如果 codes_info 中有多个可执行文件,可以设置 codes_run_mode 来指定执行这些文件的模式(默认为 CodeRunMode.SERIAL)。我们定义了 retrieve_list 文件名,engine 应在作业完成后从作业运行的目录中检索这些文件名。engine 会将这些文件存储在 FolderData node 中,并作为输出 node 附加到计算中,标签为 retrieved

更多阅读

通过 other file lists available,您可以轻松自定义如何将文件移入或移出远程工作目录,以防止创建不必要的副本。有关 CalcJob 类的更多详情,请参阅 defining calculations 上的主题部分。

解析输出结果#

将代码生成的输出文件解析为 AiiDA nodes 是可选的,但它可以使你的数据可查询,从而更容易访问和分析。

要创建解析器插件,请在名为 parsers.py 的文件中子类化 Parser 类。

from aiida.engine import ExitCode
from aiida.orm import SinglefileData
from aiida.parsers.parser import Parser
from aiida.plugins import CalculationFactory

DiffCalculation = CalculationFactory('diff-tutorial')


class DiffParser(Parser):

在调用 parse() 方法之前,会在 Parser 实例上设置两个 important 属性:

  1. self.retrievedFolderData 的实例,该实例指向包含 CalcJob 指示检索的所有输出文件的文件夹,并提供了对其中包含的任何文件进行 open() 的方法。

  2. self.node:代表已完成计算的 CalcJobNode,除其他外,可访问其所有输入 (self.node.inputs)。

现在将其 parse() 方法执行为

    def parse(self, **kwargs):
        """Parse outputs, store results in database."""
        output_filename = self.node.get_option('output_filename')

        # add output file
        self.logger.info(f"Parsing '{output_filename}'")
        with self.retrieved.open(output_filename, 'rb') as handle:
            output_node = SinglefileData(file=handle)
        self.out('diff', output_node)

        return ExitCode(0)

The get_option() convenience method is used to get the filename of the output file.

最后,使用 out() 方法返回输出文件,作为计算的 diff 输出:第一个参数是连接计算和数据 node 的链接标签名称。第二个参数是应作为输出记录的 node。

备注

输出及其类型必须与相应 CalcJob 的流程规范相符(否则会出现异常)。

在这个简约的示例中,实际上并没有进行太多解析–我们只是将输出文件作为 SinglefileData node 传递。如果你的代码会产生结构化格式的输出,你可能不想只返回文件,而是想将其解析为 python 字典 (Dict node),以便于搜索结果。

运动

考虑一下您最喜欢的仿真代码产生的不同输出文件。您希望获得哪些信息?

  1. 解析到数据库中进行查询(如 DictStructureData……)?

  2. 存储在 AiiDA 文件库中保管(例如作为 SinglefileData,……)?

  3. 在运行计算的计算机上留下(例如,使用 RemoteData 记录其远程位置或直接忽略它们)?

知道这些问题的答案后,您就可以开始为代码编写解析器了。

为了要求自动解析 CalcJob (一旦完成),用户可以在启动作业时设置 metadata.options.parser_name 输入。如果默认使用特定的解析器,CalcJob define 方法可以为解析器名称设置默认值,如 previous section 所做的那样:

@classmethod
def define(cls, spec):
    ...
    spec.inputs['metadata']['options']['parser_name'].default = 'diff-tutorial'

请注意,默认值不是设置为 Parser 类本身,而是设置为*entry point 字符串*,解析器类是在该字符串下注册的。我们稍后将为解析器注册 entry point。

处理解析错误#

到目前为止,我们还没有花太多精力来处理运行外部代码时可能出现的错误。然而,有很多方法可以导致代码无法正常执行。Parser 可以在检测和通报此类错误方面发挥 important 的作用,然后 workflows 可以决定如何继续,例如修改输入参数并重新提交计算。

解析器通过 exit codes 传递错误,这些错误定义在其解析的 CalcJobspec 中。DiffCalculation 示例定义了以下退出代码:

spec.exit_code(300, 'ERROR_MISSING_OUTPUT_FILES', message='Calculation did not produce all expected output files.')

exit_code 定义:

  • 退出状态(一个正整数,沿用 退出代码惯例)、

  • 可用于引用 parse 方法中代码的标签(通过 self.exit_codes 属性,如下所示),以及

  • 信息,对问题进行更详细的描述。

为了通知 AiiDA 计算失败,只需从 parse 方法返回与检测到的问题相对应的退出代码。下面是上一节 Parser 例子的更完整版本:

    def parse(self, **kwargs):
        """Parse outputs, store results in database.

        :returns: non-zero exit code, if parsing fails
        """
        output_filename = self.node.get_option('output_filename')

        # Check that folder content is as expected
        files_retrieved = self.retrieved.list_object_names()
        files_expected = [output_filename]
        # Note: set(A) <= set(B) checks whether A is a subset of B
        if not set(files_expected) <= set(files_retrieved):
            self.logger.error(f"Found files '{files_retrieved}', expected to find '{files_expected}'")
            return self.exit_codes.ERROR_MISSING_OUTPUT_FILES

        # add output file
        self.logger.info(f"Parsing '{output_filename}'")
        with self.retrieved.open(output_filename, 'rb') as handle:
            output_node = SinglefileData(file=handle)
        self.out('diff', output_node)

        return ExitCode(0)

这个简单的检查可以确保从运行计算的计算机获取的文件中包含预期输出文件 diff.patch。生产插件通常会扫描输出的其他方面(如标准错误、输出文件等),查找可能表明计算有问题的任何问题,并返回相应的退出代码。

AiiDA 将 parse 方法返回的退出代码存储在正在解析的计算 node,在那里可以进一步检查(详见 defining processes 主题)。请注意,某些调度程序插件可以在调度程序级别(通过解析作业调度程序输出)检测问题并设置退出代码。有关 scheduler exit codes 的主题部分解释了如何在解析器内部检测这些问题,以及如何选择性地覆盖这些问题。

定制#

工艺标签#

每次运行 Process 时,都会在数据库中存储 ProcessNode 以记录执行情况。process_label 属性中会存储一个人类可读的标签。默认情况下,进程类的名称被用作此标签。如果默认值不够翔实,可通过重载 _build_process_label(): 方法进行自定义:

class SomeProcess(Process):

    def _build_process_label(self):
        return 'custom_process_label'

通过执行此进程类创建的 Node 将显示 node.process_label == 'custom_process_label'

注册 entry points#

Entry points are the preferred method of registering new calculation, parser and other plugins with AiiDA.

有了 calculations.pyparsers.py 文件,让我们为它们包含的插件注册 entry point:

  • 将两个脚本移至 aiida_diff_tutorial 子文件夹:

$ mkdir aiida_diff_tutorial
$ mv calculations.py parsers.py aiida_diff_tutorial/
$ touch aiida_diff_tutorial/__init__.py

您刚刚创建了一个 aiida_diff_tutorial Python package

  • 通过编写 pyproject.toml 文件,为软件包添加一组最基本的元数据:

[build-system]
# build the package with [flit](https://flit.readthedocs.io)
requires = ["flit_core >=3.4,<4"]
build-backend = "flit_core.buildapi"

[project]
# See https://www.python.org/dev/peps/pep-0621/
name = "aiida-diff-tutorial"
version = "0.1.0"
description = "AiiDA demo plugin"
dependencies = [
    "aiida-core>=2.0,<3",
]

[project.entry-points."aiida.calculations"]
"diff-tutorial" = "aiida_diff_tutorial.calculations:DiffCalculation"

[project.entry-points."aiida.parsers"]
"diff-tutorial" = "aiida_diff_tutorial.parsers:DiffParser"

[tool.flit.module]
name = "aiida_diff_tutorial"

备注

这样就可以在 pyproject.toml 文件中使用 PEP 621 格式完全指定项目元数据。

  • 安装新的 aiida-diff-tutorial 插件包。

$ pip install -e .  # install package in "editable mode"

详见 如何安装插件 部分。

之后,您应该会看到您的插件列表:

$ verdi plugin list aiida.calculations
$ verdi plugin list aiida.calculations diff-tutorial
$ verdi plugin list aiida.parsers

运行计算#

设置好 entry point 后,您就可以使用新插件启动第一次计算了:

  • 如果您还没有这样做,set up your computer。在下文中,我们假定它是 localhost:

$ verdi computer setup -L localhost -H localhost -T core.local -S core.direct -w `echo $PWD/work` -n
$ verdi computer configure core.local localhost --safe-interval 5 -n
  • 为我们的计算创建输入文件

$ echo -e "File with content\ncontent1" > file1.txt
$ echo -e "File with content\ncontent2" > file2.txt
$ mkdir input_files
$ mv file1.txt file2.txt input_files
  • 编写 launch.py 脚本:

"""Launch a calculation using the 'diff-tutorial' plugin"""
from pathlib import Path

from aiida import engine, orm
from aiida.common.exceptions import NotExistent

INPUT_DIR = Path(__file__).resolve().parent / 'input_files'

# Create or load code
computer = orm.load_computer('localhost')
try:
    code = orm.load_code('diff@localhost')
except NotExistent:
    # Setting up code via python API (or use "verdi code setup")
    code = orm.InstalledCode(
        label='diff', computer=computer, filepath_executable='/usr/bin/diff', default_calc_job_plugin='diff-tutorial'
    )

# Set up inputs
builder = code.get_builder()
builder.file1 = orm.SinglefileData(file=INPUT_DIR / 'file1.txt')
builder.file2 = orm.SinglefileData(file=INPUT_DIR / 'file2.txt')
builder.metadata.description = 'Test job submission with the aiida_diff_tutorial plugin'

# Run the calculation & parse results
result = engine.run(builder)
computed_diff = result['diff'].get_content()
print(f'Computed diff between files:\n{computed_diff}')

备注

launch.py 脚本设置了一个 AiiDA AbstractCode 实例,将 /usr/bin/diff 可执行文件与 DiffCalculation 类关联(通过其 entry point diff 关联)。

该代码会自动设置在生成器的 code 输入端口上,并作为输入传递给计算插件。

  • 启动计算:

$ verdi run launch.py

如果一切顺利,就会打印出计算结果,类似于这样:

$ verdi run launch.py
Computed diff between files:
2c2
< content1
---
> content2

小技巧

如果遇到解析错误,制作一个 试运行 可能会有帮助,它允许你在启动任何计算之前检查 AiiDA 生成的输入文件夹。

最后,你可以将计算提交给 AiiDA 守护进程,而不是在当前 shell 中运行:

  • (重新)启动守护进程,更新其 Python 环境:

$ verdi daemon restart --reset
  • 更新启动脚本以使用

# Submit calculation to the aiida daemon
node = engine.submit(builder)
print("Submitted calculation {}".format(node))

备注

nodeCalcJobNode,代表底层计算过程的状态(可能尚未完成)。

  • 启动计算:

$ verdi run launch.py

这将打印已提交计算的 UUID 和 PK。

您可以使用 verdi 命令行界面来执行 monitor 进程:

$ verdi process list -a -p1

这将显示刚刚运行的两个计算的进程。使用 verdi calcjob outputcat <pk> 查看提交给守护进程的计算结果。

恭喜 - 您现在可以为外部仿真代码编写插件,并使用它们提交计算结果!

如果您还有时间,可以考虑做下面的选择性练习。

为现有计算编写 importers#

插件的新用户可能经常在没有使用AiiDA的情况下完成了许多以前的计算,他们希望将这些计算import到AiiDA中。在这种情况下,可以为他们的输入/输出写一个 importer ,为相应的 CalcJob 生成 provenance nodes 。

importer 必须写成 CalcJobImporter 的子类,示例见 aiida.calculations.importers.arithmetic.add.ArithmeticAddCalculationImporter

要将 importer 与 CalcJob 类联系起来,必须将 importer 与 aiida.calculations.importers 组中的 entry point 注册在一起。

[project.entry-points."aiida.calculations.importers"]
"core.arithmetic.add" = "aiida.calculations.importers.arithmetic.add:ArithmeticAddCalculationImporter"

备注

请注意,entry point 名称可以是任何有效的 entry point 名称。如果 importer 插件与相应的 CalcJob 插件由同一软件包提供,建议 importer 和 CalcJob 插件的 entry point 名称相同。这将允许 get_importer() 方法自动获取相关的 importer。如果 entry point 名称不同,则需要将所需 importer 实现的 entry point 名称作为参数传递给 get_importer()

然后,用户可以通过 get_importer() 方法进行 import 计算:

from aiida.plugins import CalculationFactory

ArithmeticAddCalculation = CalculationFactory('arithmetic.add')
importer = ArithmeticAddCalculation.get_importer()
remote_data = RemoteData('/some/absolute/path', computer=load_computer('computer'))
inputs = importer.parse_remote_data(remote_data)
results, node = run.get_node(ArithmeticAddCalculation, **inputs)
assert node.is_imported

参见

请参阅 AEP 004: Infrastructure to import completed calculation jobs ,了解有关此功能的设计考虑因素。

练习 - 支持命令行选项#

如前所述,diff 知道几个命令行选项:

$ diff --help
Usage: diff [OPTION]... FILES
Compare files line by line.
...
-i, --ignore-case               ignore case differences in file contents
-E, --ignore-tab-expansion      ignore changes due to tab expansion
-b, --ignore-space-change       ignore changes in the amount of white space
-w, --ignore-all-space          ignore all white space
-B, --ignore-blank-lines        ignore changes where lines are all blank
-I, --ignore-matching-lines=RE  ignore changes where all lines match RE
...

为简单起见,让我们把重点放在上面显示的选项摘录上,并允许我们的插件用户传递这些选项。

请注意,其中一个选项(--ignore-matching-lines)要求用户传递一个正则表达式字符串,而其他选项不需要任何值。

表示一组命令行选项的一种方法是

diff --ignore-case --ignore-matching-lines='.*ABC.*'

将使用 python 字典:

parameters = {
  'ignore-case': True,
  'ignore-space-change': False,
  'ignore-matching-lines': '.*ABC.*'
 }

下面是将字典翻译成命令行选项列表的简单代码片段:

def cli_options(parameters):
     """Return command line options for parameters dictionary.

     :param dict parameters: dictionary with command line parameters
     """
     options = []
     for key, value in parameters.items():
         # Could validate: is key a known command-line option?
         if isinstance(value, bool) and value:
             options.append(f'--{key}')
         elif isinstance(value, str):
             # Could validate: is value a valid regular expression?
             options.append(f'--{key}')
             options.append(value)

     return options

备注

在向仿真代码传递参数时,请尝试对参数进行 验证。这样可以在 提交*计算时直接检测出错误,从而防止畸形输入的计算进入高性能计算系统的队列。

为了简洁起见,我们在这里不执行验证,但有许多 Python 库,如 voluptuous (被 aiida-diff 使用,参见 example )、marshmallowpydantic ,可以帮助您定义一个 schema,以验证输入。

让我们打开之前的 calculations.py 文件,开始修改 DiffCalculation 类:

  1. define 方法中,在标号为 'parameters'spec 中添加新的 input 并键入 Dict (from aiida.orm import Dict)

  2. prepare_for_submission 方法中,在 self.inputs.parameters.get_dict() 上运行上面的 cli_options 函数,以获取命令行选项列表。将它们添加到 codeinfo.cmdline_params 中。

解决方案

For 1. add the following line to the define method:

spec.input('parameters', valid_type=Dict, help='diff command-line parameters')

复制 calculations.py 末尾的 cli_options 代码段,并将 cmdline_params 设置为:

codeinfo.cmdline_params = cli_options(self.inputs.parameters.get_dict()) + [ self.inputs.file1.filename, self.inputs.file2.filename]

就是这样。现在让我们打开 launch.py 脚本,并传递命令行参数:

...
builder.parameters = orm.Dict(dict={'ignore-case': True})
...

更改 file1.txt 第一行中一个字符的大小写。然后,重新启动守护进程并提交新的计算结果:

$ verdi daemon restart
$ verdi run launch.py

如果一切正常,第一行中的大小写差异应被忽略(因此不会显示在输出中)。

本教程到此结束。

CalcJobParser 插件仍然相当基本,而 aiida-diff-tutorial 插件包则缺少许多有用的功能,例如包元数据、文档、测试、CI 等。请继续阅读 如何打包插件,了解如何从零开始快速创建功能丰富的新插件包。