aiida.manage

Managing an AiiDA instance:

  • configuration file

  • profiles

  • databases

  • repositories

  • external components (such as Postgres, RabbitMQ)

Note

Modules in this sub package may require the database environment to be loaded

Package Contents

Classes

Config

Object that represents the configuration file of an AiiDA instance.

ConfigValidationError

Configuration error raised when the file contents fails validation.

ManagementApiConnectionError

Raised when no connection can be made to the management HTTP API.

Option

Represent a configuration option schema.

Postgres

Adds convenience functions to pgsu.PGSU.

ProcessLauncher

A sub class of plumpy.ProcessLauncher to launch a Process.

Profile

Class that models a profile as it is stored in the configuration file of an AiiDA instance.

RabbitmqManagementClient

Client for RabbitMQ Management HTTP API.

Functions

check_and_migrate_config

Checks if the config needs to be migrated, and performs the migration if needed.

config_needs_migrating

Checks if the config needs to be migrated.

config_schema

Return the configuration schema.

disable_caching

Context manager to disable caching, either for a specific node class, or globally.

downgrade_config

Run the registered configuration migrations down to the target version.

enable_caching

Context manager to enable caching, either for a specific node class, or globally.

get_current_version

Return the current version of the config.

get_launch_queue_name

Return the launch queue name with an optional prefix.

get_manager

Return the AiiDA global manager instance.

get_message_exchange_name

Return the message exchange name for a given prefix.

get_option

Return option.

get_option_names

Return a list of available option names.

get_rmq_url

Return the URL to connect to RabbitMQ.

get_task_exchange_name

Return the task exchange name for a given prefix.

get_use_cache

Return whether the caching mechanism should be used for the given process type according to the configuration.

parse_option

Parse and validate a value for a configuration option.

upgrade_config

Run the registered configuration migrations up to the target version.

Data

External

DEFAULT_DSN

PostgresConnectionMode

API

aiida.manage.BROKER_DEFAULTS = 'AttributeDict(...)'
aiida.manage.CURRENT_CONFIG_VERSION = 9
class aiida.manage.Config(filepath: str, config: dict, validate: bool = True)

Object that represents the configuration file of an AiiDA instance.

Initialization

Instantiate a configuration object from a configuration dictionary and its filepath.

If an empty dictionary is passed, the constructor will create the skeleton configuration dictionary.

Parameters:
  • filepath – the absolute filepath of the configuration file

  • config – the content of the configuration file in dictionary form

  • validate – validate the dictionary against the schema

KEY_VERSION = 'CONFIG_VERSION'
KEY_VERSION_CURRENT = 'CURRENT'
KEY_VERSION_OLDEST_COMPATIBLE = 'OLDEST_COMPATIBLE'
KEY_DEFAULT_PROFILE = 'default_profile'
KEY_PROFILES = 'profiles'
KEY_OPTIONS = 'options'
KEY_SCHEMA = '$schema'
classmethod from_file(filepath)

Instantiate a configuration object from the contents of a given file.

Note

if the filepath does not exist an empty file will be created with the current default configuration and will be written to disk. If the filepath does already exist but contains a configuration with an outdated schema, the content will be migrated and then written to disk.

Parameters:

filepath – the absolute path to the configuration file

Returns:

Config instance

classmethod _backup(filepath)

Create a backup of the configuration file with the given filepath.

Parameters:

filepath – absolute path to the configuration file to backup

Returns:

the absolute path of the created backup

static validate(config: dict, filepath: Optional[str] = None)

Validate a configuration dictionary.

__eq__(other)

Two configurations are considered equal, when their dictionaries are equal.

__ne__(other)

Two configurations are considered unequal, when their dictionaries are unequal.

handle_invalid(message)

Handle an incoming invalid configuration dictionary.

The current content of the configuration file will be written to a backup file.

Parameters:

message – a string message to echo with describing the infraction

property dictionary: dict

Return the dictionary representation of the config as it would be written to file.

Returns:

dictionary representation of config as it should be written to file

property version
property version_oldest_compatible
property version_settings
property filepath
property dirpath
property default_profile_name

Return the default profile name.

Returns:

the default profile name or None if not defined

property profile_names

Return the list of profile names.

Returns:

list of profile names

property profiles

Return the list of profiles.

Returns:

the profiles

Return type:

list of Profile instances

validate_profile(name)

Validate that a profile exists.

Parameters:

name – name of the profile:

Raises:

aiida.common.ProfileConfigurationError – if the name is not found in the configuration file

get_profile(name: Optional[str] = None) aiida.manage.configuration.profile.Profile

Return the profile for the given name or the default one if not specified.

Returns:

the profile instance or None if it does not exist

Raises:

aiida.common.ProfileConfigurationError – if the name is not found in the configuration file

add_profile(profile)

Add a profile to the configuration.

Parameters:

profile – the profile configuration dictionary

Returns:

self

update_profile(profile)

Update a profile in the configuration.

Parameters:

profile – the profile instance to update

Returns:

self

remove_profile(name)

Remove a profile from the configuration.

Parameters:

name – the name of the profile to remove

Raises:

aiida.common.ProfileConfigurationError – if the given profile does not exist

Returns:

self

delete_profile(name: str, include_database: bool = True, include_database_user: bool = False, include_repository: bool = True)

Delete a profile including its storage.

Parameters:
  • include_database – also delete the database configured for the profile.

  • include_database_user – also delete the database user configured for the profile.

  • include_repository – also delete the repository configured for the profile.

set_default_profile(name, overwrite=False)

Set the given profile as the new default.

Parameters:
  • name – name of the profile to set as new default

  • overwrite – when True, set the profile as the new default even if a default profile is already defined

Raises:

aiida.common.ProfileConfigurationError – if the given profile does not exist

Returns:

self

property options
set_option(option_name, option_value, scope=None, override=True)

Set a configuration option for a certain scope.

Parameters:
  • option_name – the name of the configuration option

  • option_value – the option value

  • scope – set the option for this profile or globally if not specified

  • override – boolean, if False, will not override the option if it already exists

Returns:

the parsed value (potentially cast to a valid type)

unset_option(option_name: str, scope=None)

Unset a configuration option for a certain scope.

Parameters:
  • option_name – the name of the configuration option

  • scope – unset the option for this profile or globally if not specified

get_option(option_name, scope=None, default=True)

Get a configuration option for a certain scope.

Parameters:
  • option_name – the name of the configuration option

  • scope – get the option for this profile or globally if not specified

  • default – boolean, If True will return the option default, even if not defined within the given scope

Returns:

the option value or None if not set for the given scope

get_options(scope: Optional[str] = None) Dict[str, Tuple[aiida.manage.configuration.options.Option, str, Any]]

Return a dictionary of all option values and their source (‘profile’, ‘global’, or ‘default’).

Parameters:

scope – the profile name or globally if not specified

Returns:

(option, source, value)

store()

Write the current config to file.

Note

if the configuration file already exists on disk and its contents differ from those in memory, a backup of the original file on disk will be created before overwriting it.

Returns:

self

_atomic_write(filepath=None)

Write the config as it is in memory, i.e. the contents of self.dictionary, to disk.

Note

this command will write the config from memory to a temporary file in the same directory as the target file filepath. It will then use os.rename to move the temporary file to filepath which will be overwritten if it already exists. The os.rename is the operation that gives the best guarantee of being atomic within the limitations of the application.

Parameters:

filepath – optional filepath to write the contents to, if not specified, the default filename is used.

class aiida.manage.ConfigValidationError(message: str, keypath: Sequence[Any] = (), schema: Optional[dict] = None, filepath: Optional[str] = None)

Bases: aiida.common.exceptions.ConfigurationError

Configuration error raised when the file contents fails validation.

Initialization

__str__() str
aiida.manage.MIGRATIONS = ()
class aiida.manage.ManagementApiConnectionError

Bases: aiida.common.exceptions.AiidaException

Raised when no connection can be made to the management HTTP API.

aiida.manage.OLDEST_COMPATIBLE_CONFIG_VERSION = 9
class aiida.manage.Option(name: str, schema: Dict[str, Any])

Represent a configuration option schema.

Initialization

__str__() str
property name: str
property schema: Dict[str, Any]
property valid_type: Any
property default: Any
property description: str
property global_only: bool
validate(value: Any, cast: bool = True) Any

Validate a value

Parameters:
  • value – The input value

  • cast – Attempt to cast the value to the required type

Returns:

The output value

Raise:

ConfigValidationError

class aiida.manage.Postgres(dbinfo=None, **kwargs)

Bases: pgsu.PGSU

Adds convenience functions to pgsu.PGSU.

Provides convenience functions for
  • creating/dropping users

  • creating/dropping databases

etc.

Example:

postgres = Postgres()
postgres.create_dbuser('username', 'password')
if not postgres.db_exists('dbname'):
    postgres.create_db('username', 'dbname')

Initialization

See documentation of pgsu.PGSU.__init__().

classmethod from_profile(profile: aiida.manage.configuration.Profile, **kwargs)

Create Postgres instance with dbinfo from AiiDA profile data.

Note: This only uses host and port from the profile, since the others are not going to be relevant for the

database superuser.

Parameters:
  • profile – AiiDA profile instance

  • kwargs – keyword arguments forwarded to PGSU constructor

Returns:

Postgres instance pre-populated with data from AiiDA profile

dbuser_exists(dbuser)

Find out if postgres user with name dbuser exists

Parameters:

dbuser (str) – database user to check for

Returns:

(bool) True if user exists, False otherwise

create_dbuser(dbuser, dbpass, privileges='')

Create a database user in postgres

Parameters:
  • dbuser (str) – Name of the user to be created.

  • dbpass (str) – Password the user should be given.

Raises:

psycopg2.errors.DuplicateObject if user already exists and self.connection_mode == PostgresConnectionMode.PSYCOPG

drop_dbuser(dbuser)

Drop a database user in postgres

Parameters:

dbuser (str) – Name of the user to be dropped.

check_dbuser(dbuser)

Looks up if a given user already exists, prompts for using or creating a differently named one.

Parameters:

dbuser (str) – Name of the user to be created or reused.

Returns:

tuple (dbuser, created)

db_exists(dbname)

Check wether a postgres database with dbname exists

Parameters:

dbname (str) – Name of the database to check for

Returns:

(bool), True if database exists, False otherwise

create_db(dbuser, dbname)

Create a database in postgres

Parameters:
  • dbuser (str) – Name of the user which should own the db.

  • dbname (str) – Name of the database.

drop_db(dbname)

Drop a database in postgres

Parameters:

dbname (str) – Name of the database.

copy_db(src_db, dest_db, dbuser)
check_db(dbname)

Looks up if a database with the name exists, prompts for using or creating a differently named one.

Parameters:

dbname (str) – Name of the database to be created or reused.

Returns:

tuple (dbname, created)

create_dbuser_db_safe(dbname, dbuser, dbpass)

Create DB and user + grant privileges.

Prompts when reusing existing users / databases.

property host_for_psycopg2

Return correct host for psycopg2 connection (as required by regular AiiDA operation).

property port_for_psycopg2

Return port for psycopg2 connection (as required by regular AiiDA operation).

property dbinfo

Alias for Postgres.dsn.

class aiida.manage.ProcessLauncher

Bases: plumpy.ProcessLauncher

A sub class of plumpy.ProcessLauncher to launch a Process.

It overrides the _continue method to make sure the node corresponding to the task can be loaded and that if it is already marked as terminated, it is not continued but the future is reconstructed and returned

static handle_continue_exception(node, exception, message)

Handle exception raised in _continue call.

If the process state of the node has not yet been put to excepted, the exception was raised before the process instance could be reconstructed, for example when the process class could not be loaded, thereby circumventing the exception handling of the state machine. Raising this exception will then acknowledge the process task with RabbitMQ leaving an uncleaned node in the CREATED state for ever. Therefore we have to perform the node cleaning manually.

Parameters:
  • exception – the exception object

  • message – string message to use for the log message

async _continue(communicator, pid, nowait, tag=None)

Continue the task.

Note that the task may already have been completed, as indicated from the corresponding the node, in which case it is not continued, but the corresponding future is reconstructed and returned. This scenario may occur when the Process was already completed by another worker that however failed to send the acknowledgment.

Parameters:
  • communicator – the communicator that called this method

  • pid – the pid of the process to continue

  • nowait – if True don’t wait for the process to finish, just return the pid, otherwise wait and return the results

  • tag – the tag of the checkpoint to continue from

class aiida.manage.Profile(name: str, config: Mapping[str, Any], validate=True)

Class that models a profile as it is stored in the configuration file of an AiiDA instance.

Initialization

Load a profile with the profile configuration.

KEY_UUID = 'PROFILE_UUID'
KEY_DEFAULT_USER_EMAIL = 'default_user_email'
KEY_STORAGE = 'storage'
KEY_PROCESS = 'process_control'
KEY_STORAGE_BACKEND = 'backend'
KEY_STORAGE_CONFIG = 'config'
KEY_PROCESS_BACKEND = 'backend'
KEY_PROCESS_CONFIG = 'config'
KEY_OPTIONS = 'options'
KEY_TEST_PROFILE = 'test_profile'
REQUIRED_KEYS = ()
__repr__() str
copy()

Return a copy of the profile.

property uuid: str

Return the profile uuid.

Returns:

string UUID

property default_user_email: Optional[str]

Return the default user email.

property storage_backend: str

Return the type of the storage backend.

property storage_config: Dict[str, Any]

Return the configuration required by the storage backend.

set_storage(name: str, config: Dict[str, Any]) None

Set the storage backend and its configuration.

Parameters:
  • name – the name of the storage backend

  • config – the configuration of the storage backend

property storage_cls: Type[aiida.orm.implementation.StorageBackend]

Return the storage backend class for this profile.

property process_control_backend: str

Return the type of the process control backend.

property process_control_config: Dict[str, Any]

Return the configuration required by the process control backend.

set_process_controller(name: str, config: Dict[str, Any]) None

Set the process control backend and its configuration.

Parameters:
  • name – the name of the process backend

  • config – the configuration of the process backend

property options
get_option(option_key, default=None)
set_option(option_key, value, override=True)

Set a configuration option for a certain scope.

Parameters:
  • option_key – the key of the configuration option

  • option_value – the option value

  • override – boolean, if False, will not override the option if it already exists

unset_option(option_key)
property name

Return the profile name.

Returns:

the profile name

property dictionary: Dict[str, Any]

Return the profile attributes as a dictionary with keys as it is stored in the config

Returns:

the profile configuration dictionary

property is_test_profile: bool

Return whether the profile is a test profile

Returns:

boolean, True if test profile, False otherwise

property repository_path: pathlib.Path

Return the absolute path of the repository configured for this profile.

The URI should be in the format protocol://address

Note:

At the moment, only the file protocol is supported.

Returns:

absolute filepath of the profile’s file repository

property rmq_prefix: str

Return the prefix that should be used for RMQ resources

Returns:

the rmq prefix string

get_rmq_url() str

Return the RMQ url for this profile.

property filepaths

Return the filepaths used by this profile.

Returns:

a dictionary of filepaths

class aiida.manage.RabbitmqManagementClient(username: str, password: str, hostname: str, virtual_host: str)

Client for RabbitMQ Management HTTP API.

This requires the rabbitmq_management plugin (https://www.rabbitmq.com/management.html) to be enabled. Typically this is enabled by running rabbitmq-plugins enable rabbitmq_management.

Initialization

Construct a new instance.

Parameters:
  • username – The username to authenticate with.

  • password – The password to authenticate with.

  • hostname – The hostname of the RabbitMQ server.

  • virtual_host – The virtual host.

format_url(url: str, url_params: dict[str, str] | None = None) str

Format the complete URL from a partial resource path with placeholders.

The base URL will be automatically prepended.

Parameters:
  • url – The resource path with placeholders, e.g., queues/{virtual_host}/{queue}.

  • url_params – Dictionary with values for the placeholders in the url. The virtual_host value is automatically inserted and should not be specified.

Returns:

The complete URL.

request(url: str, url_params: dict[str, str] | None = None, method: str = 'GET', params: dict[str, Any] | None = None) requests.Response

Make a request.

Parameters:
  • url – The resource path with placeholders, e.g., queues/{virtual_host}/{queue}.

  • url_params – Dictionary with values for the placeholders in the url. The virtual_host value is automatically inserted and should not be specified.

  • method – The HTTP method.

  • params – Query parameters to add to the URL.

Returns:

The response of the request.

Raises:

ManagementApiConnectionError – If connection to the API cannot be made.

property is_connected: bool

Return whether the API server can be connected to.

Note

Tries to reach the server at the /api/cluster-name end-point.

Returns:

True if the server can be reached, False otherwise.

aiida.manage.check_and_migrate_config(config, filepath: Optional[str] = None)

Checks if the config needs to be migrated, and performs the migration if needed.

Parameters:
  • config – the configuration dictionary

  • filepath – the path to the configuration file (optional, for error reporting)

Returns:

the migrated configuration dictionary

aiida.manage.config_needs_migrating(config, filepath: Optional[str] = None)

Checks if the config needs to be migrated.

If the oldest compatible version of the configuration is higher than the current configuration version defined in the code, the config cannot be used and so the function will raise.

Parameters:

filepath – the path to the configuration file (optional, for error reporting)

Returns:

True if the configuration has an older version and needs to be migrated, False otherwise

Raises:

aiida.common.ConfigurationVersionError – if the config’s oldest compatible version is higher than the current

aiida.manage.config_schema() Dict[str, Any]

Return the configuration schema.

aiida.manage.disable_caching(*, identifier=None)

Context manager to disable caching, either for a specific node class, or globally.

Warning

this does not affect the behavior of the daemon, only the local Python interpreter.

Parameters:

identifier (str) – Process type string of the node, or a pattern with ‘*’ wildcard that matches it. If not provided, caching is disabled for all classes.

aiida.manage.downgrade_config(config: aiida.manage.configuration.migrations.migrations.ConfigType, target: int, migrations: Iterable[Type[aiida.manage.configuration.migrations.migrations.SingleMigration]] = MIGRATIONS) aiida.manage.configuration.migrations.migrations.ConfigType

Run the registered configuration migrations down to the target version.

Parameters:

config – the configuration dictionary

Returns:

the migrated configuration dictionary

aiida.manage.enable_caching(*, identifier=None)

Context manager to enable caching, either for a specific node class, or globally.

Warning

this does not affect the behavior of the daemon, only the local Python interpreter.

Parameters:

identifier (str) – Process type string of the node, or a pattern with ‘*’ wildcard that matches it. If not provided, caching is enabled for all classes.

aiida.manage.get_current_version(config)

Return the current version of the config.

Returns:

current config version or 0 if not defined

aiida.manage.get_launch_queue_name(prefix=None)

Return the launch queue name with an optional prefix.

Returns:

launch queue name

aiida.manage.get_manager() Manager

Return the AiiDA global manager instance.

aiida.manage.get_message_exchange_name(prefix)

Return the message exchange name for a given prefix.

Returns:

message exchange name

aiida.manage.get_option(name: str) aiida.manage.configuration.options.Option

Return option.

aiida.manage.get_option_names() List[str]

Return a list of available option names.

aiida.manage.get_rmq_url(protocol=None, username=None, password=None, host=None, port=None, virtual_host=None, **kwargs)

Return the URL to connect to RabbitMQ.

Note

The default of the host is set to 127.0.0.1 instead of localhost because on some computers localhost resolves first to IPv6 with address ::1 and if RMQ is not running on IPv6 one gets an annoying warning. For more info see: https://github.com/aiidateam/aiida-core/issues/1142

Parameters:
  • protocol – the protocol to use, amqp or amqps.

  • username – the username for authentication.

  • password – the password for authentication.

  • host – the hostname of the RabbitMQ server.

  • port – the port of the RabbitMQ server.

  • virtual_host – the virtual host to connect to.

  • kwargs – remaining keyword arguments that will be encoded as query parameters.

Returns:

the connection URL string.

aiida.manage.get_task_exchange_name(prefix)

Return the task exchange name for a given prefix.

Returns:

task exchange name

aiida.manage.get_use_cache(*, identifier=None)

Return whether the caching mechanism should be used for the given process type according to the configuration.

Parameters:

identifier (str) – Process type string of the node

Returns:

boolean, True if caching is enabled, False otherwise

Raises:

~aiida.common.exceptions.ConfigurationError if the configuration is invalid, either due to a general configuration error, or by defining the class both enabled and disabled

aiida.manage.parse_option(option_name: str, option_value: Any) Tuple[aiida.manage.configuration.options.Option, Any]

Parse and validate a value for a configuration option.

Parameters:
  • option_name – the name of the configuration option

  • option_value – the option value

Returns:

a tuple of the option and the parsed value

aiida.manage.upgrade_config(config: aiida.manage.configuration.migrations.migrations.ConfigType, target: int = CURRENT_CONFIG_VERSION, migrations: Iterable[Type[aiida.manage.configuration.migrations.migrations.SingleMigration]] = MIGRATIONS) aiida.manage.configuration.migrations.migrations.ConfigType

Run the registered configuration migrations up to the target version.

Parameters:

config – the configuration dictionary

Returns:

the migrated configuration dictionary