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
