async def test_broadcast_filter_subject( communicator: kiwipy.rmq.RmqCommunicator): subjects = [] EXPECTED_SUBJECTS = ['purchase.car', 'purchase.piano'] done = asyncio.Future() def on_broadcast_1(_comm, _body, _sender=None, subject=None, _correlation_id=None): subjects.append(subject) if len(subjects) == len(EXPECTED_SUBJECTS): done.set_result(True) await communicator.add_broadcast_subscriber( kiwipy.BroadcastFilter(on_broadcast_1, subject='purchase.*')) for subj in [ 'purchase.car', 'purchase.piano', 'sell.guitar', 'sell.house' ]: await communicator.broadcast_send(None, subject=subj) await done assert len(subjects) == 2 assert EXPECTED_SUBJECTS == subjects
def test_broadcast_filter_sender_and_subject_regex(self): """As the standard broadcast test but using regular expressions instead of wildcards""" # pylint: disable=invalid-name senders_and_subjects = set() EXPECTED = { ('bob.jones', 'purchase.car'), ('bob.jones', 'purchase.piano'), ('alice.jones', 'purchase.car'), ('alice.jones', 'purchase.piano'), } done = kiwipy.Future() def on_broadcast_1(_communicator, _body, sender=None, subject=None, _correlation_id=None): senders_and_subjects.add((sender, subject)) if len(senders_and_subjects) == len(EXPECTED): done.set_result(True) filtered = kiwipy.BroadcastFilter(on_broadcast_1) filtered.add_sender_filter(re.compile('.*.jones')) filtered.add_subject_filter(re.compile('purchase.*')) self.communicator.add_broadcast_subscriber(filtered) for sender in ['bob.jones', 'bob.smith', 'martin.uhrin', 'alice.jones']: for subj in ['purchase.car', 'purchase.piano', 'sell.guitar', 'sell.house']: self.communicator.broadcast_send(None, sender=sender, subject=subj) done.result(timeout=self.WAIT_TIMEOUT) self.assertSetEqual(EXPECTED, senders_and_subjects)
async def test_broadcast_filter_sender( communicator: kiwipy.rmq.RmqCommunicator): EXPECTED_SENDERS = ['bob.jones', 'alice.jones'] senders = [] done = asyncio.Future() def on_broadcast_1(_comm, _body, sender=None, _subject=None, _correlation_id=None): senders.append(sender) if len(senders) == len(EXPECTED_SENDERS): done.set_result(True) await communicator.add_broadcast_subscriber( kiwipy.BroadcastFilter(on_broadcast_1, sender='*.jones')) for subj in ['bob.jones', 'bob.smith', 'martin.uhrin', 'alice.jones']: await communicator.broadcast_send(None, sender=subj) await done assert len(senders) == 2 assert senders == EXPECTED_SENDERS
def __init__(self, pk, loop=None, poll_interval=None, communicator=None): """Construct a future for a process node being finished. If a None poll_interval is supplied polling will not be used. If a communicator is supplied it will be used to listen for broadcast messages. :param pk: process pk :param loop: An event loop :param poll_interval: optional polling interval, if None, polling is not activated. :param communicator: optional communicator, if None, will not subscribe to broadcasts. """ from aiida.orm import load_node from .process import ProcessState super().__init__() assert not (poll_interval is None and communicator is None), 'Must poll or have a communicator to use' node = load_node(pk=pk) if node.is_terminated: self.set_result(node) else: self._communicator = communicator self.add_done_callback(lambda _: self.cleanup()) # Try setting up a filtered broadcast subscriber if self._communicator is not None: broadcast_filter = kiwipy.BroadcastFilter(lambda *args, **kwargs: self.set_result(node), sender=pk) for state in [ProcessState.FINISHED, ProcessState.KILLED, ProcessState.EXCEPTED]: broadcast_filter.add_subject_filter('state_changed.*.{}'.format(state.value)) self._broadcast_identifier = self._communicator.add_broadcast_subscriber(broadcast_filter) # Start polling if poll_interval is not None: loop.add_callback(self._poll_process, node, poll_interval)
def __init__(self, pk, loop=None, poll_interval=None, communicator=None): """ Get a future for a calculation node being finished. If a None poll_interval is supplied polling will not be used. If a communicator is supplied it will be used to listen for broadcast messages. :param pk: The calculation pk :param loop: An event loop :param poll_interval: The polling interval. Can be None in which case no polling. :param communicator: A communicator. Can be None in which case no broadcast listens. """ from aiida.orm import load_node from .processes import ProcessState super(CalculationFuture, self).__init__() assert not (poll_interval is None and communicator is None), 'Must poll or have a communicator to use' calc_node = load_node(pk=pk) if calc_node.is_terminated: self.set_result(calc_node) else: self._communicator = communicator self.add_done_callback(lambda _: self.cleanup()) # Try setting up a filtered broadcast subscriber if self._communicator is not None: self._filtered = kiwipy.BroadcastFilter(lambda *args, **kwargs: self.set_result(calc_node), sender=pk) for state in [ProcessState.FINISHED, ProcessState.KILLED, ProcessState.EXCEPTED]: self._filtered.add_subject_filter('state_changed.*.{}'.format(state.value)) self._communicator.add_broadcast_subscriber(self._filtered) # Start polling if poll_interval is not None: loop.add_callback(self._poll_calculation, calc_node, poll_interval)
def test_broadcast_filter_sender_and_subject(self): senders_and_subects = set() EXPECTED = { ('bob.jones', 'purchase.car'), ('bob.jones', 'purchase.piano'), ('alice.jones', 'purchase.car'), ('alice.jones', 'purchase.piano'), } done = kiwipy.Future() def on_broadcast_1(body, sender=None, subject=None, correlation_id=None): senders_and_subects.add((sender, subject)) if len(senders_and_subects) == len(EXPECTED): done.set_result(True) filtered = kiwipy.BroadcastFilter(on_broadcast_1) filtered.add_sender_filter("*.jones") filtered.add_subject_filter("purchase.*") self.communicator.add_broadcast_subscriber(filtered) for sender in ['bob.jones', 'bob.smith', 'martin.uhrin', 'alice.jones']: for subject in ['purchase.car', 'purchase.piano', 'sell.guitar', 'sell.house']: self.communicator.broadcast_send(None, sender=sender, subject=subject) self.communicator.await(done, timeout=self.WAIT_TIMEOUT) self.assertEqual(4, len(senders_and_subects)) self.assertSetEqual(EXPECTED, senders_and_subects)
def test_pause_play_kill(self): """ Test the pause/play/kill commands """ # pylint: disable=no-member from aiida.orm import load_node calc = self.runner.submit(test_processes.WaitProcess) start_time = time.time() while calc.process_state is not plumpy.ProcessState.WAITING: if time.time() - start_time >= self.TEST_TIMEOUT: self.fail('Timed out waiting for process to enter waiting state') # Make sure that calling any command on a non-existing process id will not except but print an error # To simulate a process without a corresponding task, we simply create a node and store it. This node will not # have an associated task at RabbitMQ, but it will be a valid `ProcessNode` so it will pass the initial # filtering of the `verdi process` commands orphaned_node = WorkFunctionNode().store() non_existing_process_id = str(orphaned_node.pk) for command in [cmd_process.process_pause, cmd_process.process_play, cmd_process.process_kill]: result = self.cli_runner.invoke(command, [non_existing_process_id]) self.assertClickResultNoException(result) self.assertIn('Error:', result.output) self.assertFalse(calc.paused) result = self.cli_runner.invoke(cmd_process.process_pause, [str(calc.pk)]) self.assertIsNone(result.exception, result.output) # We need to make sure that the process is picked up by the daemon and put in the Waiting state before we start # running the CLI commands, so we add a broadcast subscriber for the state change, which when hit will set the # future to True. This will be our signal that we can start testing waiting_future = Future() filters = kiwipy.BroadcastFilter( lambda *args, **kwargs: waiting_future.set_result(True), sender=calc.pk, subject='state_changed.*.waiting' ) self.runner.communicator.add_broadcast_subscriber(filters) # The process may already have been picked up by the daemon and put in the waiting state, before the subscriber # got the chance to attach itself, making it have missed the broadcast. That's why check if the state is already # waiting, and if not, we run the loop of the runner to start waiting for the broadcast message. To make sure # that we have the latest state of the node as it is in the database, we force refresh it by reloading it. calc = load_node(calc.pk) if calc.process_state != plumpy.ProcessState.WAITING: self.runner.loop.run_sync(lambda: with_timeout(waiting_future)) # Here we now that the process is with the daemon runner and in the waiting state so we can starting running # the `verdi process` commands that we want to test result = self.cli_runner.invoke(cmd_process.process_pause, ['--wait', str(calc.pk)]) self.assertIsNone(result.exception, result.output) self.assertTrue(calc.paused) result = self.cli_runner.invoke(cmd_process.process_play, ['--wait', str(calc.pk)]) self.assertIsNone(result.exception, result.output) self.assertFalse(calc.paused) result = self.cli_runner.invoke(cmd_process.process_kill, ['--wait', str(calc.pk)]) self.assertIsNone(result.exception, result.output) self.assertTrue(calc.is_terminated) self.assertTrue(calc.is_killed)
def __init__(self, pk: int, loop: Optional[asyncio.AbstractEventLoop] = None, poll_interval: Union[None, int, float] = None, communicator: Optional[kiwipy.Communicator] = None): """Construct a future for a process node being finished. If a None poll_interval is supplied polling will not be used. If a communicator is supplied it will be used to listen for broadcast messages. :param pk: process pk :param loop: An event loop :param poll_interval: optional polling interval, if None, polling is not activated. :param communicator: optional communicator, if None, will not subscribe to broadcasts. """ from .process import ProcessState # create future in specified event loop loop = loop if loop is not None else asyncio.get_event_loop() super().__init__(loop=loop) assert not (poll_interval is None and communicator is None ), 'Must poll or have a communicator to use' node = load_node(pk=pk) if node.is_terminated: self.set_result(node) else: self._communicator = communicator self.add_done_callback(lambda _: self.cleanup()) # Try setting up a filtered broadcast subscriber if self._communicator is not None: def _subscriber(*args, **kwargs): # pylint: disable=unused-argument if not self.done(): self.set_result(node) broadcast_filter = kiwipy.BroadcastFilter(_subscriber, sender=pk) for state in [ ProcessState.FINISHED, ProcessState.KILLED, ProcessState.EXCEPTED ]: broadcast_filter.add_subject_filter( f'state_changed.*.{state.value}') self._broadcast_identifier = self._communicator.add_broadcast_subscriber( broadcast_filter) # Start polling if poll_interval is not None: loop.create_task(self._poll_process(node, poll_interval))
def call_on_process_finish(self, pk: int, callback: Callable[[], Any]) -> None: """Schedule a callback when the process of the given pk is terminated. This method will add a broadcast subscriber that will listen for state changes of the target process to be terminated. As a fail-safe, a polling-mechanism is used to check the state of the process, should the broadcast message be missed by the subscriber, in order to prevent the caller to wait indefinitely. :param pk: pk of the process :param callback: function to be called upon process termination """ assert self.communicator is not None, 'communicator not set for runner' node = load_node(pk=pk) subscriber_identifier = str(uuid.uuid4()) event = threading.Event() def inline_callback(event, *args, **kwargs): # pylint: disable=unused-argument """Callback to wrap the actual callback, that will always remove the subscriber that will be registered. As soon as the callback is called successfully once, the `event` instance is toggled, such that if this inline callback is called a second time, the actual callback is not called again. """ if event.is_set(): return try: callback() finally: event.set() self.communicator.remove_broadcast_subscriber( subscriber_identifier) # type: ignore[union-attr] broadcast_filter = kiwipy.BroadcastFilter(functools.partial( inline_callback, event), sender=pk) for state in [ ProcessState.FINISHED, ProcessState.KILLED, ProcessState.EXCEPTED ]: broadcast_filter.add_subject_filter( f'state_changed.*.{state.value}') LOGGER.info('adding subscriber for broadcasts of %d', pk) self.communicator.add_broadcast_subscriber(broadcast_filter, subscriber_identifier) self._poll_process(node, functools.partial(inline_callback, event))
def test_broadcast_filter_match(self): """Test exact match broadcast filter""" EXPECTED_SENDERS = ['alice.jones'] # pylint: disable=invalid-name senders = [] done = kiwipy.Future() def on_broadcast_1(_comm, _body, sender=None, _subject=None, _correlation_id=None): senders.append(sender) if len(senders) == len(EXPECTED_SENDERS): done.set_result(True) self.communicator.add_broadcast_subscriber(kiwipy.BroadcastFilter(on_broadcast_1, sender='alice.jones')) for sender in ['bob.jones', 'bob.smith', 'martin.uhrin', 'alice.jones']: self.communicator.broadcast_send(None, sender=sender) done.result(timeout=self.WAIT_TIMEOUT) self.assertListEqual(EXPECTED_SENDERS, senders)
def test_broadcast_filter_subject(self): subjects = [] EXPECTED_SUBJECTS = ['purchase.car', 'purchase.piano'] # pylint: disable=invalid-name done = kiwipy.Future() def on_broadcast_1(_comm, _body, _sender=None, subject=None, _correlation_id=None): subjects.append(subject) if len(subjects) == len(EXPECTED_SUBJECTS): done.set_result(True) self.communicator.add_broadcast_subscriber(kiwipy.BroadcastFilter(on_broadcast_1, subject='purchase.*')) for subj in ['purchase.car', 'purchase.piano', 'sell.guitar', 'sell.house']: self.communicator.broadcast_send(None, subject=subj) done.result(timeout=self.WAIT_TIMEOUT) self.assertEqual(len(subjects), 2) self.assertListEqual(EXPECTED_SUBJECTS, subjects)
def test_broadcast_filter_subject(self): subjects = [] EXPECTED_SUBJECTS = ['purchase.car', 'purchase.piano'] done = kiwipy.Future() def on_broadcast_1(body, sender=None, subject=None, correlation_id=None): subjects.append(subject) if len(subjects) == len(EXPECTED_SUBJECTS): done.set_result(True) self.communicator.add_broadcast_subscriber( kiwipy.BroadcastFilter(on_broadcast_1, subject="purchase.*")) for subject in ['purchase.car', 'purchase.piano', 'sell.guitar', 'sell.house']: self.communicator.broadcast_send(None, subject=subject) self.communicator.await(done, timeout=self.WAIT_TIMEOUT) self.assertEqual(len(subjects), 2) self.assertListEqual(EXPECTED_SUBJECTS, subjects)
def test_broadcast_filter_sender(self): EXPECTED_SENDERS = ['bob.jones', 'alice.jones'] senders = [] done = kiwipy.Future() def on_broadcast_1(body, sender=None, subject=None, correlation_id=None): senders.append(sender) if len(senders) == len(EXPECTED_SENDERS): done.set_result(True) self.communicator.add_broadcast_subscriber( kiwipy.BroadcastFilter(on_broadcast_1, sender="*.jones")) for sender in ['bob.jones', 'bob.smith', 'martin.uhrin', 'alice.jones']: self.communicator.broadcast_send(None, sender=sender) self.communicator.await(done, timeout=self.WAIT_TIMEOUT) self.assertEqual(2, len(senders)) self.assertListEqual(EXPECTED_SENDERS, senders)
async def test_broadcast_filter_sender_and_subject( communicator: kiwipy.rmq.RmqCommunicator): senders_and_subects = set() EXPECTED = { ('bob.jones', 'purchase.car'), ('bob.jones', 'purchase.piano'), ('alice.jones', 'purchase.car'), ('alice.jones', 'purchase.piano'), } done = asyncio.Future() def on_broadcast_1(_comm, _body, sender=None, subject=None, _correlation_id=None): senders_and_subects.add((sender, subject)) if len(senders_and_subects) == len(EXPECTED): done.set_result(True) filtered = kiwipy.BroadcastFilter(on_broadcast_1) filtered.add_sender_filter('*.jones') filtered.add_subject_filter('purchase.*') await communicator.add_broadcast_subscriber(filtered) for sender in ['bob.jones', 'bob.smith', 'martin.uhrin', 'alice.jones']: for subj in [ 'purchase.car', 'purchase.piano', 'sell.guitar', 'sell.house' ]: await communicator.broadcast_send(None, sender=sender, subject=subj) await done assert len(senders_and_subects) == 4 assert senders_and_subects == EXPECTED
import kiwipy import sys import threading def callback(_comm, body, _sender, subject, _msg_id): print(f' [x] {subject!r}:{body!r}') with kiwipy.connect('amqp://localhost') as comm: binding_keys = sys.argv[1:] if not binding_keys: sys.stderr.write(f'Usage: {sys.argv[0]} [binding_key]...\n') sys.exit(1) for binding_key in binding_keys: comm.add_broadcast_subscriber( kiwipy.BroadcastFilter(callback, binding_key)) print(' [*] Waiting for logs. To exit press CTRL+C') threading.Event().wait()
def test_pause_play_kill(cmd_try_all, run_cli_command): """ Test the pause/play/kill commands """ # pylint: disable=no-member, too-many-locals from aiida.cmdline.commands.cmd_process import process_pause, process_play, process_kill from aiida.manage.manager import get_manager from aiida.engine import ProcessState from aiida.orm import load_node runner = get_manager().create_runner(rmq_submit=True) calc = runner.submit(test_processes.WaitProcess) test_daemon_timeout = 5. start_time = time.time() while calc.process_state is not plumpy.ProcessState.WAITING: if time.time() - start_time >= test_daemon_timeout: raise RuntimeError('Timed out waiting for process to enter waiting state') # Make sure that calling any command on a non-existing process id will not except but print an error # To simulate a process without a corresponding task, we simply create a node and store it. This node will not # have an associated task at RabbitMQ, but it will be a valid `ProcessNode` with and active state, so it will # pass the initial filtering of the `verdi process` commands orphaned_node = WorkFunctionNode() orphaned_node.set_process_state(ProcessState.RUNNING) orphaned_node.store() non_existing_process_id = str(orphaned_node.pk) for command in [process_pause, process_play, process_kill]: result = run_cli_command(command, [non_existing_process_id]) assert 'Error:' in result.output assert not calc.paused result = run_cli_command(process_pause, [str(calc.pk)]) # We need to make sure that the process is picked up by the daemon and put in the Waiting state before we start # running the CLI commands, so we add a broadcast subscriber for the state change, which when hit will set the # future to True. This will be our signal that we can start testing waiting_future = Future() filters = kiwipy.BroadcastFilter( lambda *args, **kwargs: waiting_future.set_result(True), sender=calc.pk, subject='state_changed.*.waiting' ) runner.communicator.add_broadcast_subscriber(filters) # The process may already have been picked up by the daemon and put in the waiting state, before the subscriber # got the chance to attach itself, making it have missed the broadcast. That's why check if the state is already # waiting, and if not, we run the loop of the runner to start waiting for the broadcast message. To make sure # that we have the latest state of the node as it is in the database, we force refresh it by reloading it. calc = load_node(calc.pk) if calc.process_state != plumpy.ProcessState.WAITING: runner.loop.run_until_complete(asyncio.wait_for(waiting_future, timeout=5.0)) # Here we now that the process is with the daemon runner and in the waiting state so we can starting running # the `verdi process` commands that we want to test result = run_cli_command(process_pause, ['--wait', str(calc.pk)]) assert calc.paused if cmd_try_all: cmd_option = '--all' else: cmd_option = str(calc.pk) result = run_cli_command(process_play, ['--wait', cmd_option]) assert not calc.paused result = run_cli_command(process_kill, ['--wait', str(calc.pk)]) assert calc.is_terminated assert calc.is_killed
# -*- coding: utf-8 -*- import threading import kiwipy def on_broadcast_send(_comm, body, sender, subject, __): print(' [x] listening on_broadcast_send:') print(f' body: {body}, sender {sender}, subject {subject}\n') def on_broadcast_filter(_comm, body, sender=None, subject=None, __=None): print(' [x] listening on_broadcast_filter:') print(f' body: {body}, sender {sender}, subject {subject}\n') if __name__ == '__main__': filtered = kiwipy.BroadcastFilter(on_broadcast_filter) # pylint: disable=invalid-name filtered.add_subject_filter('purchase.*') try: with kiwipy.connect('amqp://127.0.0.1') as comm: # Register a broadcast subscriber comm.add_broadcast_subscriber(on_broadcast_send) # Register a broadcast subscriber comm.add_broadcast_subscriber(filtered) # Now wait indefinitely for fibonacci calls threading.Event().wait() except KeyboardInterrupt: pass