插件#

插件的作用#

  • 在 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.manage.tests.pytest_fixtures']

只需添加这一行,pytest_fixtures 模块提供的固定装置就会自动 imported。该模块提供以下固定装置:

aiida_manager#

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

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

aiida_profile#

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

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

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

  • 临时 PostgreSQL 集群

  • 带有存储后台的临时测试配置文件(在临时 PostgreSQL 集群中创建数据库)

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

在测试套件开始时,创建临时实例和配置文件需要几秒钟来设置。要避免这种情况,可以一次性创建一个专用测试配置文件,并告诉夹具使用该配置文件,而不是每次都生成一个:

  • 使用 verdi setupverdi quicksetup 创建配置文件,并指定 --test-profile 标志

  • AIIDA_TEST_PROFILE 环境变量设置为测试配置文件的名称: export AIIDA_TEST_PROFILE=<test-profile-name>

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

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.clear_profile()

   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_instance#

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

def test(aiida_instance):
    aiida_instance.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['storage']['config']['repository_uri'] = '/some/custom/path'
    yield aiida_profile_factory(configuration)

请注意,这只有在需要自定义存储配置时才有用。如果任何配置都能正常工作,只需直接使用 aiida_profile 灯具即可,它默认使用 PsqlDosStorage 存储后端。

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_local_code_factory#

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

def test(aiida_local_code_factory):
    code = aiida_local_code_factory(
        entry_point='core.arithmetic.add',
        executable='/usr/bin/bash'
    )

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

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

submit_and_await#

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

def test(aiida_local_code_factory, submit_and_await):
    code = aiida_local_code_factory('core.arithmetic.add', '/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, 'custom.parser')

    assert ParserFactory('custom.parser', CustomParser)

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