Skip to content

Commit

Permalink
Keep on mess up with coordinator loop
Browse files Browse the repository at this point in the history
When calling add_rpc_subscriber and add_task_subscriber, the event loop
of caller may from random event loop. But the target event loop is the
runner one. Therefore it requires to pass the loop to the broker when
creating the runner and runner's broker.
  • Loading branch information
unkcpz committed Dec 27, 2024
1 parent 02a939e commit 329c51c
Show file tree
Hide file tree
Showing 4 changed files with 17 additions and 5 deletions.
5 changes: 5 additions & 0 deletions src/aiida/brokers/rabbitmq/coordinator.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
# -*- coding: utf-8 -*-
from asyncio import AbstractEventLoop
import asyncio
from typing import Generic, TypeVar, final
import kiwipy
import concurrent.futures
Expand Down Expand Up @@ -84,3 +85,7 @@ def task_send(self, task, no_reply=False):

def close(self):
self._comm.close()

def is_closed(self) -> bool:
"""Return `True` if the communicator was closed"""
return self._comm.is_closed()
3 changes: 1 addition & 2 deletions src/aiida/cmdline/utils/decorators.py
Original file line number Diff line number Diff line change
Expand Up @@ -326,8 +326,7 @@ def start_daemon():

assert profile is not None

loop = asyncio.get_event_loop()
if manager.create_broker(loop) is None:
if manager.get_broker() is None:
echo.echo_critical(
f'profile `{profile.name}` does not define a broker and so cannot use this functionality.'
f'See {URL_NO_BROKER} for more details.'
Expand Down
4 changes: 4 additions & 0 deletions src/aiida/engine/runners.py
Original file line number Diff line number Diff line change
Expand Up @@ -92,9 +92,13 @@ def __init__(
self._plugin_version_provider = PluginVersionProvider()

# FIXME: broker and coordinator overlap the concept there for over-abstraction, remove the abstraction
# Broker should always create inside runner? since they should share the loop.
if broker is not None:
self._coordinator = broker.get_coordinator()
self._controller = broker.get_controller()

# FIXME: why with wrapper, the pending task not exist??
# self._coordinator = wrap_communicator(broker.get_coordinator().communicator, self._loop)
elif self._broker_submit:
# FIXME: if broker then broker_submit else False
LOGGER.warning('Disabling broker submission, no coordinator provided')
Expand Down
10 changes: 7 additions & 3 deletions src/aiida/manage/manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -286,9 +286,13 @@ def get_profile_storage(self) -> 'StorageBackend':
return self._profile_storage

def get_broker(self) -> 'Broker | None':
return self._broker
if self._broker is not None:
return self._broker

_default_loop = asyncio.get_event_loop()
return self._create_broker(_default_loop)

def create_broker(self, loop) -> 'Broker | None':
def _create_broker(self, loop) -> 'Broker | None':
"""Return an instance of :class:`aiida.brokers.broker.Broker` if the profile defines a broker.
:returns: The broker of the profile, or ``None`` if the profile doesn't define one.
Expand Down Expand Up @@ -427,7 +431,7 @@ def create_runner(
_default_loop = asyncio.get_event_loop()

loop = loop or _default_loop
_default_broker = self.create_broker(loop)
_default_broker = self._create_broker(loop)

runner = runners.Runner(
poll_interval=poll_interval or _default_poll_interval,
Expand Down

0 comments on commit 329c51c

Please sign in to comment.