插件#

插件的作用#

  • 在 AiiDA 的 entry point groups 中加入一个新类,包括:计算、解析器、workflows、数据类型、verdi 命令、调度器、传输和来自外部数据库的 importers/ 导出器。这通常需要对 AiiDA 为此提供的基类进行子类化。

  • 安装新的命令行和/或图形用户界面可执行程序

  • 依赖于任何其他插件并在其上构建(只要它们的要求不冲突)

插件不应做的事#

AiiDA 插件不应如此:

  • 更改 schema AiiDA 使用的数据库

  • 使用 AiiDA 受保护的函数、方法或类(以下划线 _ 开头的函数、方法或类)。

  • 猴子修补 aiida 命名空间(或命名空间本身)内的任何内容

否则,您的插件可能无法列入官方 AiiDA plugin registry

如果您发现自己需要执行上述任何操作,请在 AiiDA repository 上提出问题,我们会尽力为您提供建议。

插件设计指南#

CalcJob 和解析器插件#

在包装外部代码时,应牢记以下指导原则:

  • 简单开始 使用现有的类,如 Dict , SinglefileData , …只编写与 AiiDA 之间传递信息所需的内容。

  • 不要破坏数据 provenance. **至少 存储完全可重复性所需的数据。

  • 展示全部功能。 标准化是好事,但不要人为限制代码的功能,否则用户会感到沮丧。如果代码可以做到这一点,就应该有**种方法可以用你的插件做到。

  • 不要依赖AiiDA内部。 更深嵌套层的功能不被视为公共API的一部分,可能会在AiiDA小版本之间发生变化,从而破坏你的插件。

  • 分析要查询的内容 列出要查询的信息清单:

    1. 解析到数据库以便查询( Dict ,…)

    2. 存入文件库保管( SinglefileData ,……)。

    3. 在运行计算的计算机上离开( RemoteData ,……)。

什么是 entry point?#

setuptools 软件包(由 pip 使用)有一个名为 entry points 的功能,允许将一个字符串(entry point 标识符 )与 python 软件包内定义的任何 python 对象关联。Entry points 定义在 pyproject.toml 文件中,例如::

...
[project.entry-points."aiida.data"]
# entry point = path.to.python.object
"mycode.mydata = aiida_mycode.data.mydata:MyData",
...

在这里,我们向 entry point aiida.data 中添加一个新的 entry point mycode.mydata 。entry point 标识指向文件 mydata.py 中的 MyData 类,它是 aiida_mycode 软件包的一部分。

安装定义了 entry point 的 python 软件包时,entry point 规范会被写入发行版 .egg-info 文件夹中的一个文件。 setuptools 提供了一个软件包 pkg_resources ,用于按发行版、entry point 组和/或 entry point 名称查询这些 entry point 规范,并加载其指向的数据结构。

为什么是 entry point?#

AiiDA 定义了一组 entry point 组(见下文 AiiDA entry point 小组 )。通过检查 AiiDA 插件添加到这些组的 entry point,AiiDA 可以提供统一的接口与它们交互。例如

  • verdi plugin list aiida.workflows 提供 AiiDA 插件安装的所有 workflow 的概览。用户可以使用同一命令检查每个 workflow 的输入/输出,而无需研究插件的文档。

  • DataFactoryCalculationFactoryWorkflowFactory 方法允许通过简单的短字符串(如 quantumespresso.pw )实例化新类。用户无需记住类在插件包中的确切位置,而且插件可以重构,用户无需重新学习插件的 API。

AiiDA entry point 小组#

下面,我们列出了 AiiDA 定义和搜索的 entry point 组。你可以得到与 verdi plugin list 输出相同的列表。

aiida.calculations#

该组中的 Entry point 预计是 aiida.orm.JobCalculation 的子类。这取代了之前将包含相关类的 python 模块放在 aiida/orm/calculation/job 子包内的方法。

entry point 规格示例::

[project.entry-points."aiida.calculations"]
"mycode.mycode" = "aiida_mycode.calcs.mycode:MycodeCalculation"

aiida_mycode/calcs/mycode.py ::

from aiida.orm import JobCalculation
class MycodeCalculation(JobCalculation):
   ...

将导致使用::

from aiida.plugins import CalculationFactory
calc = CalculationFactory('mycode.mycode')

aiida.parsers#

AiiDA 期望是 Parser 的子类。取代之前将解析器模块置于 aiida/parsers/plugins 下的方法。

规格示例::

[project.entry-points."aiida.parsers"]
"mycode.myparser" = "aiida_mycode.parsers.mycode:MycodeParser"

aida_mycode/parsers/myparser.py ::

from aiida.parsers import Parser
class MycodeParser(Parser)
   ...

使用方法::

from aiida.plugins import ParserFactory
parser = ParserFactory('mycode.mycode')

aiida.data#

Group for Data subclasses. Previously located in a subpackage of aiida/orm/data.

规格::

[project.entry-points."aiida.data"]
"mycode.mydata" = "aiida_mycode.data.mydata:MyData"

aiida_mycode/data/mydat.py ::

from aiida.orm import Data
class MyData(Data):
   ...

使用方法::

from aiida.plugins import DataFactory
params = DataFactory('mycode.mydata')

aiida.workflows#

AiiDA workflows 软件包如下:

规格::

[project.entry-points."aiida.workflows"]
"mycode.mywf" = "aiida_mycode.workflows.mywf:MyWorkflow"

aiida_mycode/workflows/mywf.py ::

from aiida.engine.workchain import WorkChain
class MyWorkflow(WorkChain):
   ...

使用方法::

from aiida.plugins import WorkflowFactory
wf = WorkflowFactory('mycode.mywf')

备注

老式 workflow 不支持插件系统的 entry point 机制。因此,无法使用 WorkflowFactory 加载这些 workflow。运行这些程序的唯一方法是将源代码存储在 aiida/workflows/user 目录中,然后使用普通 python imports 加载类。

aiida.cmdline#

verdi 使用 click_ 框架,可以为现有的 verdi 命令(如 verdi data mydata )添加新的子命令。AiiDA 希望每个 entry point 都是 click.Commandclick.Group 。目前可以在以下级别注入额外命令:

verdi data 的规格::

[project.entry-points."aiida.cmdline.data"]
"mydata" = "aiida_mycode.commands.mydata:mydata"

aiida_mycode/commands/mydata.py ::

import click
@click.group()
mydata():
   """commandline help for mydata command"""

@mydata.command('animate')
@click.option('--format')
@click.argument('pk')
create_fancy_animation(format, pk):
   """help"""
   ...

使用方法

verdi data mydata animate --format=Format PK

verdi data core.structure import 的规格::

entry_points={
   "aiida.cmdline.data.structure.import": [
      "myformat = aiida_mycode.commands.myformat:myformat"
   ]
}
[project.entry-points."aiida.cmdline.data.structure.import"]
"myformat" = "aiida_mycode.commands.myformat:myformat"

aiida_mycode/commands/myformat.py ::

import click
@click.group()
@click.argument('filename', type=click.File('r'))
myformat(filename):
   """commandline help for myformat import command"""
   ...

使用方法

verdi data core.structure import myformat a_file.myfmt

aiida.tools.dbexporters#

如果您的插件包添加了向外部数据库导出的支持,请使用 entry point 让 aiida 查找定义必要函数的模块。

aiida.tools.dbimporters#

如果您的插件包增加了对外部数据库 importing 的支持,请使用此 entry point 让 aiida 找到您定义必要函数的模块。

aiida.schedulers#

我们建议以调度程序的名称来命名插件包(如 aiida-myscheduler ),这样 entry point 的名称就可以与调度程序的名称相等:

规格::

[project.entry-points."aiida.schedulers"]
"myscheduler" = "aiida_myscheduler.myscheduler:MyScheduler"

aiida_myscheduler/myscheduler.py

from aiida.schedulers import Scheduler
class MyScheduler(Scheduler):
   ...

使用方法调度程序的使用方法很简单,在设置计算机时输入 ‘myscheduler’ 作为调度程序选项。

aiida.transports#

aiida-core 有两种将文件和文件夹传输到远程计算机的模式: core.sshcore.local (当远程计算机实际上相同时的存根)。我们建议以传输模式命名插件包(如 aiida-mytransport ),这样 entry point 的名称就可以简单地与传输模式的名称相等:

规格::

[project.entry-points."aiida.transports"]
"mytransport" = "aiida_mytransport.mytransport:MyTransport"

aiida_mytransport/mytransport.py ::

from aiida.transports import Transport
class MyTransport(Transport):
   ...

使用方法::

from aiida.plugins import TransportFactory
transport = TransportFactory('mytransport')

设置新计算机时,请将 mytransport 指定为传输模式。

插件测试夹具#

在开发 AiiDA 插件包时,建议使用 pytest 作为单元测试库,它是 Python 生态系统的事实标准。它提供了大量的 fixtures ,使设置和编写测试变得容易。 aiida-core 也提供了许多 AiiDA 专用的固定装置,可以轻松测试各种插件。

要使用这些固定装置,请在 tests 文件夹中创建一个 conftest.py 文件,并添加以下代码:

pytest_plugins = 'aiida.tools.pytest_fixtures

Just by adding this line, the fixtures that are provided by the pytest_fixtures module are automatically imported. The module provides the following fixtures:

备注

Before v2.6, test fixtures were located in aiida.manage.tests.pytest_fixtures. This module is now deprecated and will be removed in the future. Some fixtures have analogs in aiida.tools.pytest_fixtures that are drop-in replacements, but in general, there are differences in the interface and functionality.

aiida_manager#

返回 Manager 的全局实例。例如,可用于检索当前 Config 实例:

def test(aiida_manager):
   aiida_manager.get_config().get_option('logging.aiida_loglevel')

aiida_profile#

该夹具确保 AiiDA 配置文件与初始化的存储后端一起加载,以便存储数据。该夹具是会话作用域的,它设置了 autouse=True ,因此在测试会话中会自动启用。

默认情况下,夹具将生成一个完全临时独立的AiiDA实例和测试配置文件。这包括

  • 包含配置文件的 .aiida 临时配置文件夹

  • A temporary test profile configured with core.sqlite_dos storage backend

备注

The profile uses core.sqlite_dos instead of the standard core.psql_dos storage plugin as it doesn’t require PostgreSQL to be installed. Since the functionality of PostgreSQL is not needed for most common test cases, this choice makes it easier to start writing and running tests.

临时测试实例和配置文件会在测试会话结束时自动销毁。夹具保证 AiiDA 的实际实例及其配置和配置文件不会被更改。

备注

The profile does not configure RabbitMQ as a broker since it is not required for most test cases useful for plugins. This means, however, that any functionality that requires a broker is not available, such as running the daemon and submitting processes to the daemon. If that functionality is required, a profile should be created and loaded that configures a broker.

虽然该夹具会自动使用,因此不需要明确地将其传递到测试函数中,但它可能仍然有用,因为它可以用来清除存储后端的所有数据:

def test(aiida_profile):
   from aiida.orm import Data, QueryBuilder

   Data().store()
   assert QueryBuilder().append(Data).count() != 0

   # The following call clears the storage backend, deleting all data, except for the default user.
   aiida_profile.reset_storage()

   assert QueryBuilder().append(Data).count() == 0

aiida_profile_clean#

通过 aiida_profile 提供已加载的测试配置文件,但会在调用测试功能前清空存储空间。请注意,清空数据库后,将向数据库中插入一个默认用户。

def test(aiida_profile_clean):
    """The profile storage is guaranteed to be emptied at the start of this test."""

如果没有预先存在的数据,设置和编写测试会更容易,那么该功能就会非常有用。不过,清理存储可能会耗费不可忽略的时间,因此只有在真正需要时才会使用,以保证测试尽可能快地运行。

aiida_profile_clean_class#

提供与 aiida_profile_clean 相同的功能,但带有 scope=class 。应用于测试类:

@pytest.mark.usefixtures('aiida_profile_clean_class')
class TestClass:

    def test():
        ...

在类初始化时,存储空间会被清理一次。

aiida_profile_factory#

创建临时配置文件,将其添加到已加载的 AiiDA 实例配置中,然后加载配置文件。可用于为自定义存储后端创建测试配置文件:

@pytest.fixture(scope='session')
def custom_storage_profile(aiida_profile_factory) -> Profile:
    """Return a test profile for a custom :class:`~aiida.orm.implementation.storage_backend.StorageBackend`"""
    from some_module import CustomStorage
    configuration = {
        'storage': {
            'backend': 'plugin_package.custom_storage',
            'config': {
                'username': 'joe'
                'api_key': 'super-secret-key'
            }
        }
    }
    yield aiida_profile_factory(configuration)

请注意,上述配置实际上并不实用,实际配置取决于所使用的存储实现。

aiida_config#

返回用于测试会话的 Config 实例。

def test(aiida_config):
    aiida_config.get_option('logging.aiida_loglevel')

config_psql_dos#

返回 PsqlDosBackend 的配置文件配置。该配置可与 aiida_profile_factory 夹具结合使用,以创建带有定制数据库参数的测试剖面:

@pytest.fixture(scope='session')
def psql_dos_profile(aiida_profile_factory, config_psql_dos) -> Profile:
    """Return a test profile configured for the :class:`~aiida.storage.psql_dos.PsqlDosStorage`."""
    configuration = config_psql_dos()
    configuration['repository_uri'] = '/some/custom/path'
    with aiida_profile_factory(storage_backend='core.psql_dos', storage_config=configuration) as profile:
        yield profile

Note that this is only useful if the storage configuration needs to be customized. If any configuration works, simply use the aiida_profile fixture straight away.

postgres_cluster#

使用 pgtest 创建一个临时和隔离的 PostgreSQL 群集,并在生成后进行清理。

@pytest.fixture()
def custom_postgres_cluster(postgres_cluster):
    yield postgres_cluster(
        database_name='some-database-name',
        database_username='guest',
        database_password='guest',
    )

aiida_localhost#

如果测试需要 Computer 实例,该测试将非常有用。该夹具将返回一个代表 localhostComputer 实例。

def test(aiida_localhost):
    aiida_localhost.get_minimum_job_poll_interval()

aiida_computer#

该固定装置用于创建和配置 Computer 实例。该固定装置提供了一个无需任何参数即可调用的工厂:

def test(aiida_computer):
    from aiida.orm import Computer
    computer = aiida_computer()
    assert isinstance(computer, Computer)

默认情况下,主机名使用 localhost,并随机生成一个标签。

def test(aiida_computer):
    custom_label = 'custom-label'
    computer = aiida_computer(label=custom_label)
    assert computer.label == custom_label

首先查询数据库,看是否已经存在带有给定标签的计算机。如果找到,则返回现有计算机,否则创建一个新实例。

返回的计算机也是为当前默认用户配置的。可通过 configuration_kwargs 字典对配置进行自定义:

def test(aiida_computer):
    configuration_kwargs = {'safe_interval': 0}
    computer = aiida_computer(configuration_kwargs=configuration_kwargs)
    assert computer.get_minimum_job_poll_interval() == 0

aiida_computer_local#

此灯具是 aiida_computer 使用本地传输设置 localhost 的快捷方式:

def test(aiida_computer_local):
    localhost = aiida_computer_local()
    assert localhost.hostname == 'localhost'
    assert localhost.transport_type == 'core.local'

要使新创建的计算机未配置,请输入 configure=False

def test(aiida_computer_local):
    localhost = aiida_computer_local(configure=False)
    assert not localhost.is_configured

请注意,如果计算机已经存在并在之前配置过,则不会取消配置。如果需要保证计算机未配置,请确保在测试前清理数据库或使用唯一标签:

def test(aiida_computer_local):
    import uuid
    localhost = aiida_computer_local(label=str(uuid.uuid4()), configure=False)
    assert not localhost.is_configured

aiida_computer_ssh#

此夹具是 aiida_computer 使用 SSH 传输设置 localhost 的快捷方式:

def test(aiida_computer_ssh):
    localhost = aiida_computer_ssh()
    assert localhost.hostname == 'localhost'
    assert localhost.transport_type == 'core.ssh'

如果需要测试的功能涉及测试 SSH 传输,这可能会很有用,但在 aiida-core 之外,这种用例应该很少见。要让新创建的计算机未配置,请通过 configure=False

def test(aiida_computer_ssh):
    localhost = aiida_computer_ssh(configure=False)
    assert not localhost.is_configured

请注意,如果计算机已经存在并在之前配置过,则不会取消配置。如果需要保证计算机未配置,请确保在测试前清理数据库或使用唯一标签:

def test(aiida_computer_ssh):
    import uuid
    localhost = aiida_computer_ssh(label=str(uuid.uuid4()), configure=False)
    assert not localhost.is_configured

aiida_code#

This fixture is useful if a test requires an AbstractCode instance. For example:

def test(aiida_localhost, aiida_code):
    from aiida.orm import InstalledCode
    code = aiida_code(
        'core.code.installed',
        label='test-code',
        computer=aiida_localhost,
        filepath_executable='/bin/bash'
    )
    assert isinstance(code, InstalledCode)

aiida_code_installed#

如果测试需要 InstalledCode 实例,则该测试非常有用。例如

def test(aiida_code_installed):
    from aiida.orm import InstalledCode
    code = aiida_code_installed()
    assert isinstance(code, InstalledCode)

默认情况下,它将使用 aiida_localhost 灯具返回的 localhost 计算机。

submit_and_await#

该夹具在测试向守护进程提交进程时非常有用。它将进程提交给守护进程,并等待进程达到特定状态。默认情况下,它会等待进程到达 ProcessState.FINISHED

def test(aiida_code_installed, submit_and_await):
    code = aiida_code_installed(filepath_executable='core.arithmetic.add', filepath_executable='/usr/bin/bash')
    builder = code.get_builder()
    builder.x = orm.Int(1)
    builder.y = orm.Int(1)
    node = submit_and_await(builder)
    assert node.is_finished_ok

请注意,灯具自动依赖于 started_daemon_client 灯具,以确保守护进程正在运行。

started_daemon_client#

该夹具确保测试配置文件的守护进程正在运行,并返回一个可用于控制守护进程的 DaemonClient 实例。

def test(started_daemon_client):
    assert started_daemon_client.is_daemon_running

stopped_daemon_client#

该夹具可确保停止测试配置文件的守护进程,并返回一个可用于控制守护进程的 DaemonClient 实例。

def test(stopped_daemon_client):
    assert not stopped_daemon_client.is_daemon_running

daemon_client#

返回一个可用于控制守护进程的 DaemonClient 实例:

def test(daemon_client):
    daemon_client.start_daemon()
    assert daemon_client.is_daemon_running
    daemon_client.stop_daemon(wait=True)

该夹具具有会话作用域。测试会话结束时,如果守护进程仍在运行,该夹具会自动关闭它。

entry_points#

返回一个 EntryPointManager 实例,用于添加和删除 entry point。

def test_parser(entry_points):
    """Test a custom ``Parser`` implementation."""
    from aiida.parsers import Parser
    from aiida.plugins import ParserFactory

    class CustomParser(Parser):
        """Parser implementation."""

    entry_points.add(CustomParser, 'aiida.parsers:custom.parser')

    assert ParserFactory('custom.parser', CustomParser)

任何 entry point 的添加和删除都会在测试结束时自动撤销。