Integrate ZMQ broker into daemon, status, and CI (aiidateam#7284)

agoscinski · agoscinski · commit d2201e70119b · 2026-04-24T10:39:18.000+02:00
Wire the ZMQ broker into the daemon lifecycle:
- Add `verdi daemon broker` hidden command for circus to manage
- Add ZMQ broker as a circus watcher started before workers
- Show ZMQ broker status in `verdi daemon status` and `verdi status`

Fix SQLite stale-read issue by committing the session after RPC calls
in `control.py`.

Update CI to run test matrix with both `rmq` and `zmq` broker backends.
Remove RabbitMQ service dependency from minimum-requirements and presto
test jobs.
diff --git a/.github/workflows/ci-code.yml b/.github/workflows/ci-code.yml
@@ -28,6 +28,7 @@ jobs:
       matrix:
         python-version: ['3.10', '3.14']
         database-backend: [psql]
+        broker-backend: [rmq, zmq]
 
     services:
       postgres:
@@ -72,7 +73,7 @@ jobs:
         AIIDA_TEST_PROFILE: test_aiida
         AIIDA_WARN_v3: 1
       run: |
-        pytest -n auto --db-backend ${{ matrix.database-backend }} -m 'not nightly' tests/ ${{ matrix.python-version == '3.14' && '--cov aiida' || '' }}
+        pytest -n auto --db-backend ${{ matrix.database-backend }} --broker-backend ${{ matrix.broker-backend }} -m 'not nightly' tests/ ${{ matrix.python-version == '3.14' && '--cov aiida' || '' }}
 
     - name: Upload coverage report
       if: matrix.python-version == 3.14 && github.repository == 'aiidateam/aiida-core'
@@ -104,13 +105,6 @@ jobs:
         python-version: ['3.10']
         database-backend: [sqlite]
 
-    services:
-      rabbitmq:
-        image: rabbitmq:3.8.14-management
-        ports:
-        - 5672:5672
-        - 15672:15672
-
     steps:
     - uses: actions/checkout@v6
 
@@ -131,7 +125,7 @@ jobs:
     - name: Run test suite
       env:
         AIIDA_WARN_v3: 0
-      run: pytest --disable-warnings -n auto --db-backend ${{ matrix.database-backend }} -m 'not nightly' tests/
+      run: pytest --disable-warnings -n auto --db-backend ${{ matrix.database-backend }} --broker-backend zmq -m 'not nightly' tests/
 
 
   tests-presto:
@@ -156,7 +150,7 @@ jobs:
     - name: Run test suite
       env:
         AIIDA_WARN_v3: 0
-      run: pytest -n auto -m 'presto' tests/
+      run: pytest -n auto --broker-backend zmq -m 'presto' tests/
 
 
   verdi:
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
@@ -63,21 +63,14 @@ jobs:
     runs-on: ubuntu-24.04
     timeout-minutes: 30
 
-    services:
-      rabbitmq:
-        image: rabbitmq:3.8.14-management
-        ports:
-        - 5672:5672
-        - 15672:15672
-
     steps:
     - uses: actions/checkout@v6
 
     - name: Install aiida-core
       uses: ./.github/actions/install-aiida-core
 
     - name: Run sub-set of test suite
-      run: pytest -s -m requires_rmq --db-backend=sqlite tests/
+      run: pytest -s -m requires_broker --broker-backend zmq --db-backend=sqlite tests/
 
   publish-pypi:
     if: startsWith(github.ref, 'refs/tags/')
diff --git a/src/aiida/cmdline/commands/cmd_daemon.py b/src/aiida/cmdline/commands/cmd_daemon.py
@@ -109,6 +109,7 @@ def status(ctx, all_profiles, timeout):
 
     from aiida.cmdline.utils.common import format_local_time
     from aiida.engine.daemon.client import DaemonException, get_daemon_client
+    from aiida.manage.manager import get_manager
 
     if all_profiles is True:
         profiles = [profile for profile in ctx.obj.config.profiles if not profile.is_test_profile]
@@ -146,8 +147,29 @@ def status(ctx, all_profiles, timeout):
             workers_info = '--> No workers are running. Use `verdi daemon incr` to start some!\n'
 
         start_time = format_local_time(daemon_response['info']['create_time'])
+
+        # Build broker status line for managed brokers (e.g., ZMQ)
+        broker_line = ''
+        broker = get_manager().get_broker()
+        from aiida.brokers.zmq.broker import ZmqBroker
+
+        if isinstance(broker, ZmqBroker):
+            if broker.is_running():
+                status_info = broker.get_service_status()
+                if status_info:
+                    broker_pid = status_info.get('pid', '?')
+                    pending = status_info.get('pending_tasks', 0)
+                    processing = status_info.get('processing_tasks', 0)
+                    broker_line = (
+                        f'Broker is running as PID {broker_pid} [{pending} pending, {processing} processing]\n'
+                        f'Broker directory: {broker.base_path}\n'
+                    )
+            else:
+                broker_line = 'Broker is NOT running\n'
+
         echo.echo(
             f'Daemon is running as PID {daemon_response["info"]["pid"]} since {start_time}\n'
+            f'{broker_line}'
             f'Active workers [{len(workers)}]:\n{workers_info}\n'
             f'Log file: {client.daemon_log_file}\n'
             'Use `verdi daemon [incr | decr] [num]` to increase / decrease the number of workers'
@@ -279,3 +301,20 @@ def worker():
     from aiida.engine.daemon.worker import start_daemon_worker
 
     start_daemon_worker(foreground=True)
+
+
+@verdi_daemon.command('broker', hidden=True)
+@decorators.with_dbenv()
+@decorators.requires_broker
+def broker():
+    """Run the ZMQ broker server in the current process.
+
+    .. note:: this should not be called directly from the commandline!
+    """
+    from aiida.brokers.zmq.broker import get_broker_base_path
+    from aiida.brokers.zmq.service import run_broker_service
+    from aiida.manage.manager import get_manager
+
+    profile = get_manager().get_profile()
+    assert profile is not None
+    run_broker_service(base_path=get_broker_base_path(profile))
diff --git a/src/aiida/cmdline/commands/cmd_status.py b/src/aiida/cmdline/commands/cmd_status.py
@@ -130,46 +130,69 @@ def verdi_status(print_traceback: bool, no_rmq: bool) -> None:
             version=3,
         )
 
-    # Getting the broker status
+    # Getting the daemon and broker status
     broker = manager.get_broker()
 
     if broker:
-        try:
-            broker.get_communicator()
-        except Exception as exc:
-            message = f'Unable to connect to broker: {broker}'
-            print_status(ServiceStatus.ERROR, 'broker', message, exception=exc, print_traceback=print_traceback)
-            exit_code = ExitCode.CRITICAL
-        else:
-            print_status(ServiceStatus.UP, 'broker', str(broker))
-        finally:
-            broker.close()
-    else:
+        from aiida.brokers.zmq.broker import ZmqBroker
+
+        # For RabbitMQ: verify broker connectivity as a separate status line
+        # For ZMQ: broker info is shown alongside the daemon status below
+        if not isinstance(broker, ZmqBroker):
+            try:
+                broker.get_communicator()
+            except Exception as exc:
+                message = f'Unable to connect to broker: {broker}'
+                print_status(ServiceStatus.ERROR, 'broker', message, exception=exc, print_traceback=print_traceback)
+                exit_code = ExitCode.CRITICAL
+            else:
+                print_status(ServiceStatus.UP, 'broker', str(broker))
+            finally:
+                broker.close()
+
+    if not broker:
         print_status(
             ServiceStatus.WARNING,
             'broker',
             f'No broker defined for this profile: certain functionality not available.\nSee {URL_NO_BROKER}',
         )
-
-    # Getting the daemon status
-    try:
-        status = manager.get_daemon_client().get_status()
-    except ConfigurationError:
         print_status(
             ServiceStatus.WARNING,
             'daemon',
-            'No broker defined for this profile: daemon is not available.',
+            f'No broker defined for this profile: daemon is not available. See {URL_NO_BROKER}',
         )
-    except DaemonNotRunningException as exception:
-        print_status(ServiceStatus.WARNING, 'daemon', str(exception))
-    except DaemonException as exception:
-        print_status(ServiceStatus.ERROR, 'daemon', str(exception))
-    except Exception as exception:
-        message = 'Error getting daemon status'
-        print_status(ServiceStatus.ERROR, 'daemon', message, exception=exception, print_traceback=print_traceback)
-        exit_code = ExitCode.CRITICAL
     else:
-        print_status(ServiceStatus.UP, 'daemon', f'Daemon is running with PID {status["pid"]}')
+        try:
+            daemon_status = manager.get_daemon_client().get_status()
+        except ConfigurationError:
+            print_status(
+                ServiceStatus.WARNING,
+                'daemon',
+                f'No broker defined for this profile: daemon is not available. See {URL_NO_BROKER}',
+            )
+        except DaemonNotRunningException as exception:
+            print_status(ServiceStatus.WARNING, 'daemon', str(exception))
+        except DaemonException as exception:
+            print_status(ServiceStatus.ERROR, 'daemon', str(exception))
+        except Exception as exception:
+            message = 'Error getting daemon status'
+            print_status(ServiceStatus.ERROR, 'daemon', message, exception=exception, print_traceback=print_traceback)
+            exit_code = ExitCode.CRITICAL
+        else:
+            daemon_msg = f'Daemon is running with PID {daemon_status["pid"]}'
+            # Append broker info for managed brokers (e.g., ZMQ)
+
+            if isinstance(broker, ZmqBroker):
+                if broker.is_running():
+                    status_info = broker.get_service_status()
+                    if status_info:
+                        broker_pid = status_info.get('pid', '?')
+                        pending = status_info.get('pending_tasks', 0)
+                        processing = status_info.get('processing_tasks', 0)
+                        daemon_msg += f', Broker PID {broker_pid} [{pending} pending, {processing} processing]'
+                else:
+                    daemon_msg += ', Broker is NOT running'
+            print_status(ServiceStatus.UP, 'daemon', daemon_msg)
 
     # Note: click does not forward return values to the exit code, see https://github.com/pallets/click/issues/747
     if exit_code != ExitCode.SUCCESS:
diff --git a/src/aiida/engine/daemon/client.py b/src/aiida/engine/daemon/client.py
@@ -149,6 +149,11 @@ def cmd_start_daemon_worker(self) -> list[str]:
         """Return the command to start a daemon worker process."""
         return [self._verdi_bin, '-p', self.profile.name, 'daemon', 'worker']
 
+    @property
+    def cmd_start_broker(self) -> list[str]:
+        """Return the command to start the ZMQ broker server process."""
+        return [self._verdi_bin, '-p', self.profile.name, 'daemon', 'broker']
+
     @property
     def loglevel(self) -> str:
         return get_config_option('logging.circus_loglevel')
@@ -692,6 +697,61 @@ def _start_daemon(self, number_workers: int = 1, foreground: bool = False) -> No
         if not foreground:
             logoutput = self.circus_log_file
 
+        watchers = []
+
+        # Start ZMQ broker before workers so its sockets are ready when workers connect.
+        # Skip if a broker is already running (e.g. started by the test fixture).
+        if self.profile.process_control_backend == 'core.zmq':
+            from aiida.manage.manager import get_manager
+
+            broker_instance = get_manager().get_broker()
+            broker_already_running = (
+                broker_instance is not None and hasattr(broker_instance, 'is_running') and broker_instance.is_running()
+            )
+
+            if not broker_already_running:
+                watchers.append(
+                    {
+                        'cmd': ' '.join(self.cmd_start_broker),
+                        'name': f'{self.daemon_name}-broker',
+                        'numprocesses': 1,
+                        'virtualenv': self.virtualenv,
+                        'copy_env': True,
+                        'stdout_stream': {
+                            'class': 'FileStream',
+                            'filename': self.daemon_log_file,
+                            'time_format': '%Y-%m-%d %H:%M:%S',
+                        },
+                        'stderr_stream': {
+                            'class': 'FileStream',
+                            'filename': self.daemon_log_file,
+                            'time_format': '%Y-%m-%d %H:%M:%S',
+                        },
+                        'env': self.get_env(),
+                    }
+                )
+
+        watchers.append(
+            {
+                'cmd': ' '.join(self.cmd_start_daemon_worker),
+                'name': self.daemon_name,
+                'numprocesses': number_workers,
+                'virtualenv': self.virtualenv,
+                'copy_env': True,
+                'stdout_stream': {
+                    'class': 'FileStream',
+                    'filename': self.daemon_log_file,
+                    'time_format': '%Y-%m-%d %H:%M:%S',
+                },
+                'stderr_stream': {
+                    'class': 'FileStream',
+                    'filename': self.daemon_log_file,
+                    'time_format': '%Y-%m-%d %H:%M:%S',
+                },
+                'env': self.get_env(),
+            }
+        )
+
         arbiter_config = {
             'controller': self.get_controller_endpoint(),
             'pubsub_endpoint': self.get_pubsub_endpoint(),
@@ -701,26 +761,7 @@ def _start_daemon(self, number_workers: int = 1, foreground: bool = False) -> No
             'debug': False,
             'statsd': True,
             'pidfile': self.circus_pid_file,
-            'watchers': [
-                {
-                    'cmd': ' '.join(self.cmd_start_daemon_worker),
-                    'name': self.daemon_name,
-                    'numprocesses': number_workers,
-                    'virtualenv': self.virtualenv,
-                    'copy_env': True,
-                    'stdout_stream': {
-                        'class': 'FileStream',
-                        'filename': self.daemon_log_file,
-                        'time_format': '%Y-%m-%d %H:%M:%S',
-                    },
-                    'stderr_stream': {
-                        'class': 'FileStream',
-                        'filename': self.daemon_log_file,
-                        'time_format': '%Y-%m-%d %H:%M:%S',
-                    },
-                    'env': self.get_env(),
-                }
-            ],
+            'watchers': watchers,
         }
 
         if not foreground:
diff --git a/src/aiida/engine/processes/control.py b/src/aiida/engine/processes/control.py
@@ -237,6 +237,13 @@ def _perform_actions(
 
     _resolve_futures(futures, infinitive, present, timeout)
 
+    # End the current read transaction so that subsequent attribute reads see
+    # the state committed by the daemon.  Without this, SQLite's transaction
+    # snapshot isolation can cause stale reads after the RPC round-trip.
+    storage = get_manager().get_profile_storage()
+    if hasattr(storage, 'get_session'):
+        storage.get_session().commit()
+
 
 def _resolve_futures(
     futures: dict[concurrent.futures.Future, ProcessNode],
diff --git a/tests/cmdline/commands/test_daemon.py b/tests/cmdline/commands/test_daemon.py
diff --git a/tests/cmdline/commands/test_status.py b/tests/cmdline/commands/test_status.py
diff --git a/tests/engine/test_zmq.py b/tests/engine/test_zmq.py
