diff --git a/environment.yml b/environment.yml index 98dd997ba1..99d7748c64 100644 --- a/environment.yml +++ b/environment.yml @@ -8,7 +8,6 @@ dependencies: - python~=3.9 - alembic~=1.2 - archive-path~=0.4.2 -- aio-pika~=6.6 - circus~=0.18.0 - click-spinner~=0.1.8 - click~=8.1 @@ -19,11 +18,11 @@ dependencies: - ipython>=7 - jedi<0.19 - jinja2~=3.0 -- kiwipy[rmq]~=0.7.7 +- kiwipy[rmq]~=0.8.4 - importlib-metadata~=6.0 - numpy~=1.21 - paramiko>=2.7.2,~=2.7 -- plumpy~=0.21.6 +- plumpy~=0.22.3 - pgsu~=0.2.1 - psutil~=5.6 - psycopg2-binary~=2.8 diff --git a/pyproject.toml b/pyproject.toml index c70c7a96de..5f31cef2a0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,7 +20,6 @@ classifiers = [ dependencies = [ 'alembic~=1.2', 'archive-path~=0.4.2', - 'aio-pika~=6.6', 'circus~=0.18.0', 'click-spinner~=0.1.8', 'click~=8.1', @@ -31,11 +30,11 @@ dependencies = [ 'ipython>=7', 'jedi<0.19', 'jinja2~=3.0', - 'kiwipy[rmq]~=0.7.7', + 'kiwipy[rmq]~=0.8.4', 'importlib-metadata~=6.0', 'numpy~=1.21', 'paramiko~=2.7,>=2.7.2', - 'plumpy~=0.21.6', + 'plumpy~=0.22.3', 'pgsu~=0.2.1', 'psutil~=5.6', 'psycopg2-binary~=2.8', diff --git a/requirements/requirements-py-3.10.txt b/requirements/requirements-py-3.10.txt index 7bd1c23ce4..d6ca92f6c3 100644 --- a/requirements/requirements-py-3.10.txt +++ b/requirements/requirements-py-3.10.txt @@ -6,8 +6,8 @@ # accessible-pygments==0.0.5 aiida-export-migration-tests==0.9.0 -aio-pika==6.8.1 -aiormq==3.3.1 +aio-pika==9.4.0 +aiormq==6.8.0 alabaster==0.7.13 alembic==1.11.1 aniso8601==9.0.1 @@ -78,7 +78,7 @@ jupyter-server==2.6.0 jupyter-server-terminals==0.4.4 jupyterlab-pygments==0.2.2 jupyterlab-widgets==3.0.7 -kiwipy[rmq]==0.7.7 +kiwipy[rmq]==0.8.4 kiwisolver==1.4.4 latexcodec==2.0.1 mako==1.2.4 @@ -106,7 +106,7 @@ numpy==1.25.0 overrides==7.3.1 packaging==23.1 palettable==3.3.3 -pamqp==2.3.0 +pamqp==3.3.0 pandas==2.0.2 pandocfilters==1.5.0 paramiko==2.12.0 @@ -120,7 +120,7 @@ pillow==9.5.0 platformdirs==3.6.0 plotly==5.15.0 pluggy==1.0.0 -plumpy==0.21.8 +plumpy==0.22.3 prometheus-client==0.17.0 prompt-toolkit==3.0.38 psutil==5.9.5 diff --git a/requirements/requirements-py-3.11.txt b/requirements/requirements-py-3.11.txt index db6593c6ba..95347a0980 100644 --- a/requirements/requirements-py-3.11.txt +++ b/requirements/requirements-py-3.11.txt @@ -6,8 +6,8 @@ # accessible-pygments==0.0.5 aiida-export-migration-tests==0.9.0 -aio-pika==6.8.1 -aiormq==3.3.1 +aio-pika==9.4.0 +aiormq==6.8.0 alabaster==0.7.13 alembic==1.11.1 aniso8601==9.0.1 @@ -77,7 +77,7 @@ jupyter-server==2.6.0 jupyter-server-terminals==0.4.4 jupyterlab-pygments==0.2.2 jupyterlab-widgets==3.0.7 -kiwipy[rmq]==0.7.7 +kiwipy[rmq]==0.8.4 kiwisolver==1.4.4 latexcodec==2.0.1 mako==1.2.4 @@ -105,7 +105,7 @@ numpy==1.25.0 overrides==7.3.1 packaging==23.1 palettable==3.3.3 -pamqp==2.3.0 +pamqp==3.3.0 pandas==2.0.2 pandocfilters==1.5.0 paramiko==2.12.0 @@ -119,7 +119,7 @@ pillow==9.5.0 platformdirs==3.6.0 plotly==5.15.0 pluggy==1.0.0 -plumpy==0.21.8 +plumpy==0.22.3 prometheus-client==0.17.0 prompt-toolkit==3.0.38 psutil==5.9.5 diff --git a/requirements/requirements-py-3.12.txt b/requirements/requirements-py-3.12.txt index 78f4e3a8f5..15d59944df 100644 --- a/requirements/requirements-py-3.12.txt +++ b/requirements/requirements-py-3.12.txt @@ -6,8 +6,8 @@ # accessible-pygments==0.0.4 aiida-export-migration-tests==0.9.0 -aio-pika==6.8.1 -aiormq==3.3.1 +aio-pika==9.4.0 +aiormq==6.8.0 alabaster==0.7.13 alembic==1.12.0 aniso8601==9.0.1 @@ -77,7 +77,7 @@ jupyter-server==2.8.0 jupyter-server-terminals==0.4.4 jupyterlab-pygments==0.2.2 jupyterlab-widgets==3.0.9 -kiwipy[rmq]==0.7.8 +kiwipy[rmq]==0.8.4 kiwisolver==1.4.5 latexcodec==2.0.1 mako==1.2.4 @@ -105,7 +105,7 @@ numpy==1.26.1 overrides==7.4.0 packaging==23.2 palettable==3.3.3 -pamqp==2.3.0 +pamqp==3.3.0 pandas==2.1.1 pandocfilters==1.5.0 paramiko==2.12.0 @@ -119,7 +119,7 @@ pillow==10.1.0 platformdirs==3.11.0 plotly==5.17.0 pluggy==1.3.0 -plumpy==0.21.10 +plumpy==0.22.3 prometheus-client==0.17.1 prompt-toolkit==3.0.39 psutil==5.9.6 diff --git a/requirements/requirements-py-3.9.txt b/requirements/requirements-py-3.9.txt index a576ca238d..1a7d1b2704 100644 --- a/requirements/requirements-py-3.9.txt +++ b/requirements/requirements-py-3.9.txt @@ -6,8 +6,8 @@ # accessible-pygments==0.0.5 aiida-export-migration-tests==0.9.0 -aio-pika==6.8.1 -aiormq==3.3.1 +aio-pika==9.4.0 +aiormq==6.8.0 alabaster==0.7.13 alembic==1.11.1 aniso8601==9.0.1 @@ -80,7 +80,7 @@ jupyter-server==2.6.0 jupyter-server-terminals==0.4.4 jupyterlab-pygments==0.2.2 jupyterlab-widgets==3.0.7 -kiwipy[rmq]==0.7.7 +kiwipy[rmq]==0.8.4 kiwisolver==1.4.4 latexcodec==2.0.1 mako==1.2.4 @@ -108,7 +108,7 @@ numpy==1.25.0 overrides==7.3.1 packaging==23.1 palettable==3.3.3 -pamqp==2.3.0 +pamqp==3.3.0 pandas==2.0.2 pandocfilters==1.5.0 paramiko==2.12.0 @@ -122,7 +122,7 @@ pillow==9.5.0 platformdirs==3.6.0 plotly==5.15.0 pluggy==1.0.0 -plumpy==0.21.8 +plumpy==0.22.3 prometheus-client==0.17.0 prompt-toolkit==3.0.38 psutil==5.9.5 diff --git a/src/aiida/brokers/rabbitmq/broker.py b/src/aiida/brokers/rabbitmq/broker.py index 5321f6d400..c4ecfa2400 100644 --- a/src/aiida/brokers/rabbitmq/broker.py +++ b/src/aiida/brokers/rabbitmq/broker.py @@ -122,4 +122,4 @@ def get_rabbitmq_version(self): """ from packaging.version import parse - return parse(self.get_communicator().server_properties['version'].decode('utf-8')) + return parse(self.get_communicator().server_properties['version']) diff --git a/src/aiida/engine/processes/process.py b/src/aiida/engine/processes/process.py index f4dc9f9d69..5eabfd56f7 100644 --- a/src/aiida/engine/processes/process.py +++ b/src/aiida/engine/processes/process.py @@ -39,9 +39,10 @@ import plumpy.futures import plumpy.persistence import plumpy.processes -from aio_pika.exceptions import ConnectionClosed from kiwipy.communications import UnroutableError from plumpy.process_states import Finished, ProcessState +from plumpy.processes import ConnectionClosed # type: ignore[attr-defined] +from plumpy.processes import Process as PlumpyProcess from plumpy.utils import AttributesFrozendict from aiida import orm @@ -66,7 +67,7 @@ @plumpy.persistence.auto_persist('_parent_pid', '_enable_persistence') -class Process(plumpy.processes.Process): +class Process(PlumpyProcess): """This class represents an AiiDA process which can be executed and will have full provenance saved in the database. """ diff --git a/src/aiida/manage/manager.py b/src/aiida/manage/manager.py index c84b39f903..8621b324f4 100644 --- a/src/aiida/manage/manager.py +++ b/src/aiida/manage/manager.py @@ -68,6 +68,8 @@ class Manager: def __init__(self) -> None: """Construct a new instance.""" + from aiida.common.log import AIIDA_LOGGER + # note: the config currently references the global variables self._broker: Optional['Broker'] = None self._profile: Optional['Profile'] = None @@ -76,6 +78,7 @@ def __init__(self) -> None: self._process_controller: Optional['RemoteProcessThreadController'] = None self._persister: Optional['AiiDAPersister'] = None self._runner: Optional['Runner'] = None + self.logger = AIIDA_LOGGER.getChild(__name__) @staticmethod def get_config(create=False) -> 'Config': @@ -165,8 +168,15 @@ def reset_profile_storage(self) -> None: def reset_broker(self) -> None: """Reset the communicator.""" + from concurrent import futures + if self._broker is not None: + try: + self._broker.close() + except futures.TimeoutError as exception: + self.logger.warning(f'Call to close the broker timed out: {exception}') self._broker.close() + self._broker = None self._process_controller = None diff --git a/tests/manage/test_manager.py b/tests/manage/test_manager.py new file mode 100644 index 0000000000..4359a2ab48 --- /dev/null +++ b/tests/manage/test_manager.py @@ -0,0 +1,33 @@ +"""Tests for :mod:`aiida.manage.manager`.""" + +import pytest +from aiida import engine, orm + + +@engine.calcfunction +def add_calcfunction(data): + return orm.Int(data.value + 1) + + +@pytest.mark.requires_rmq +def test_disconnect(): + """Test the communicator disconnect. + + When the dependency ``kiwipy`` was updated to v0.8, it introduced a problem with shutting down the communicator. + After at least one process would have been run, trying to disconnect the communcitor would time out. The problem + is related to the update of the lower lying libraries ``aio-pika`` and ``aiormq`` to v9.4 and v6.8, respectively. + After much painstaking debugging the cause could not be determined, nor a solution. This test is added to + demonstrate the problematic behavior. Getting the communicator and then disconnecting it (through calling + :meth:`aiida.manage.manager.Manager.reset_profile`) works fine. However, if a process is a run before closing it, + for example running a calcfunction, the closing of the communicator will raise a ``TimeoutError``. + """ + from aiida.manage import get_manager + + manager = get_manager() + manager.get_communicator() + manager.reset_profile() # This returns just fine + + result, node = add_calcfunction.run_get_node(1) + assert node.is_finished_ok + assert result == 2 + manager.reset_profile() # This hangs before timing out