Dependencies: Update requirements for kiwipy and plumpy

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

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',

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

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

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

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

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'])

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.
     """

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
 

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