Skip to content

Commit

Permalink
Automatically try and detect a RabbitMQ server
Browse files Browse the repository at this point in the history
The `verdi presto` command now checks whether it can connect to a
RabbitMQ instance on the localhost with default connection parameters.
If it is found, the profile is configured with RabbitMQ as a broker,
otehrwise the profile will simply not configure a broker.
  • Loading branch information
sphuber committed Apr 17, 2024
1 parent 3f079c4 commit d79ffc7
Show file tree
Hide file tree
Showing 2 changed files with 33 additions and 1 deletion.
21 changes: 21 additions & 0 deletions src/aiida/brokers/rabbitmq/defaults.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,9 @@
"""Defaults related to RabbitMQ."""

from __future__ import annotations

import typing as t

from aiida.common.extendeddicts import AttributeDict

__all__ = ('BROKER_DEFAULTS',)
Expand All @@ -19,3 +23,20 @@
'heartbeat': 600,
}
)


def detect_rabbitmq_config() -> dict[str, t.Any] | None:
"""Try to connect to a RabbitMQ server with the default connection parameters.
:returns: The connection parameters if the RabbitMQ server was successfully connected to, or ``None`` otherwise.
"""
from kiwipy.rmq.threadcomms import connect

connection_params = dict(BROKER_DEFAULTS)

try:
connect(connection_params=connection_params)
except ConnectionError:
return None

return connection_params
13 changes: 12 additions & 1 deletion src/aiida/cmdline/commands/cmd_presto.py
Original file line number Diff line number Diff line change
Expand Up @@ -93,12 +93,21 @@ def verdi_presto(ctx, profile_name, email, first_name, last_name, institution):
"""
import tempfile

from aiida.brokers.rabbitmq.defaults import detect_rabbitmq_config
from aiida.common import exceptions
from aiida.manage.configuration import create_profile, load_profile
from aiida.orm import Computer

storage_backend = 'core.sqlite_dos'
storage_config: dict[str, t.Any] = {}
storage_backend = 'core.sqlite_dos'

broker_config = detect_rabbitmq_config()
broker_backend = 'core.rabbitmq' if broker_config is not None else None

if broker_config is None:
echo.echo_report('RabbitMQ server not found: configuring the profile without a broker.')
else:
echo.echo_report('RabbitMQ server detected: configuring the profile with a broker.')

try:
profile = create_profile(
Expand All @@ -110,6 +119,8 @@ def verdi_presto(ctx, profile_name, email, first_name, last_name, institution):
institution=institution,
storage_backend=storage_backend,
storage_config=storage_config,
broker_backend=broker_backend,
broker_config=broker_config,
)
except (ValueError, TypeError, exceptions.EntryPointError, exceptions.StorageMigrationError) as exception:
echo.echo_critical(str(exception))
Expand Down

0 comments on commit d79ffc7

Please sign in to comment.