def setUp(self):
super(TestStatusProvider, self).setUp()
self._response = None
self._corr_id = None
try:
self._connection = pika.BlockingConnection()
except pika.exceptions.ConnectionClosed:
self.fail(
"Couldn't open connection. Make sure _rmq server is running")
self.channel = self._connection.channel()
# Set up the request exchange
self.request_exchange = '{}.{}.task_control'.format(
self.__class__, uuid.uuid4())
self.channel.exchange_declare(exchange=self.request_exchange,
type='fanout')
# Set up the response queue
result = self.channel.queue_declare(exclusive=True)
self.response_queue = result.method.queue
self.channel.basic_consume(self._on_response,
no_ack=True,
queue=self.response_queue)
self.controller = ProcessController()
self.provider = ProcessStatusSubscriber(
self._connection,
exchange=self.request_exchange,
process_controller=self.controller)
def __init__(self,
connection,
queue=Defaults.TASK_QUEUE,
decoder=json.loads,
process_controller=None):
"""
:param connection: The pika RabbitMQ connection
:type connection: :class:`pika.Connection`
:param queue: The queue name to use
:param decoder: A function to deserialise incoming messages
:param process_controller: The process controller to use, one will be created if None is passed
:type process_controller: :class:`plum.process_controller.ProcessController`
"""
if process_controller is None:
self._controller = ProcessController()
else:
self._controller = process_controller
self._decode = decoder
self._running_processes = {}
self._stopping = False
self._num_processes = 0
# Set up communications
self._channel = connection.channel()
self._channel.basic_qos(prefetch_count=1)
self._channel.queue_declare(queue=queue, durable=True)
self._channel.basic_consume(self._on_launch, queue=queue)
class TestTaskControllerAndRunner(TestCase):
def setUp(self):
super(TestTaskControllerAndRunner, self).setUp()
try:
self._connection = pika.BlockingConnection()
except pika.exceptions.ConnectionClosed:
self.fail(
"Couldn't open connection. Make sure _rmq server is running")
self.controller = ProcessController()
queue = "{}.{}.tasks".format(self.__class__.__name__, uuid.uuid4())
self.publisher = ProcessLaunchPublisher(self._connection, queue=queue)
self.subscriber = ProcessLaunchSubscriber(
self._connection, queue=queue, process_controller=self.controller)
def tearDown(self):
self._connection.close()
num_procs = self.controller.get_num_processes()
if num_procs != 0:
warnings.warn(
"Process manager is still running '{}' processes".format(
num_procs))
# Kill any still running processes
self.controller.abort_all(timeout=10.)
self.assertEqual(self.controller.get_num_processes(), 0,
"Failed to abort all processes")
def test_launch(self):
class RanLogger(ProcessMonitorListener):
def __init__(self):
self.ran = []
def on_monitored_process_registered(self, process):
self.ran.append(process.__class__)
# Try launching some processes
for proc_class in TEST_PROCESSES:
self.publisher.launch(proc_class)
l = RanLogger()
with MONITOR.listen(l):
# Now make them run
num_ran = 0
for _ in range(0, 10):
num_ran += self.subscriber.poll(0.2)
if num_ran >= len(TEST_PROCESSES):
break
self.assertEqual(num_ran, len(TEST_PROCESSES))
self.assertListEqual(TEST_PROCESSES, l.ran)
def setUp(self):
super(TestControl, self).setUp()
self._connection = self._create_connection()
self.controller = ProcessController()
self.exchange = "{}.{}.control".format(self.__class__, uuid.uuid4())
self.publisher = ProcessControlPublisher(self._connection,
exchange=self.exchange)
self.subscriber_thread = SubscriberThread(
self._create_connection, self._create_control_subscriber)
self.subscriber_thread.set_poll_time(0.1)
self.subscriber_thread.start()
self.subscriber_thread.wait_till_started()
def setUp(self):
super(TestTaskControllerAndRunner, self).setUp()
try:
self._connection = pika.BlockingConnection()
except pika.exceptions.ConnectionClosed:
self.fail(
"Couldn't open connection. Make sure _rmq server is running")
self.controller = ProcessController()
queue = "{}.{}.tasks".format(self.__class__.__name__, uuid.uuid4())
self.publisher = ProcessLaunchPublisher(self._connection, queue=queue)
self.subscriber = ProcessLaunchSubscriber(
self._connection, queue=queue, process_controller=self.controller)
class TestStatusRequesterAndProvider(TestCase):
def setUp(self):
super(TestStatusRequesterAndProvider, self).setUp()
self.response = None
# Set up communications
try:
self._connection = pika.BlockingConnection()
except pika.exceptions.ConnectionClosed:
self.fail(
"Couldn't open connection. Make sure _rmq server is running")
exchange = "{}.{}.status_request".format(self.__class__.__name__,
uuid.uuid4())
self.requester = ProcessStatusRequester(self._connection,
exchange=exchange)
self.controller = ProcessController()
self.provider = ProcessStatusSubscriber(
self._connection,
exchange=exchange,
process_controller=self.controller)
def tearDown(self):
self.controller.remove_all(timeout=10.)
self.assertEqual(self.controller.get_num_processes(), 0,
"Failed to remove all processes")
super(TestStatusRequesterAndProvider, self).tearDown()
self._connection.close()
def test_request(self):
procs = []
for i in range(0, 10):
p = WaitForSignalProcess.new()
procs.append(p)
self.controller.insert_and_play(p)
responses = self._send_request_poll_response(0.2)
self.assertEqual(len(responses), 1)
procs_info = responses[0][status.PROCS_KEY]
self.assertEqual(len(procs_info), len(procs))
self.assertSetEqual(set(procs_info.keys()), {p.pid for p in procs})
self.controller.remove_all(timeout=10.)
self.assertEqual(self.controller.get_num_processes(), 0,
"Failed to abort all processes")
responses = self._send_request_poll_response(0.2)
self.assertEqual(len(responses), 1)
self.assertEqual(len(responses[0][status.PROCS_KEY]), 0)
def _send_request_poll_response(self, timeout):
self.requester.send_request()
self.provider.poll(timeout)
return self.requester.poll_response(timeout=timeout)
def setUp(self):
super(TestStatusRequesterAndProvider, self).setUp()
self.response = None
# Set up communications
try:
self._connection = pika.BlockingConnection()
except pika.exceptions.ConnectionClosed:
self.fail(
"Couldn't open connection. Make sure _rmq server is running")
exchange = "{}.{}.status_request".format(self.__class__.__name__,
uuid.uuid4())
self.requester = ProcessStatusRequester(self._connection,
exchange=exchange)
self.controller = ProcessController()
self.provider = ProcessStatusSubscriber(
self._connection,
exchange=exchange,
process_controller=self.controller)
class TestControl(TestCase):
def setUp(self):
super(TestControl, self).setUp()
self._connection = self._create_connection()
self.controller = ProcessController()
self.exchange = "{}.{}.control".format(self.__class__, uuid.uuid4())
self.publisher = ProcessControlPublisher(self._connection,
exchange=self.exchange)
self.subscriber_thread = SubscriberThread(
self._create_connection, self._create_control_subscriber)
self.subscriber_thread.set_poll_time(0.1)
self.subscriber_thread.start()
self.subscriber_thread.wait_till_started()
def tearDown(self):
self.controller.remove_all(timeout=10.)
self.assertEqual(self.controller.get_num_processes(), 0,
"Failed to abort all processes")
super(TestControl, self).tearDown()
self.subscriber_thread.stop()
self.subscriber_thread.join()
self._connection.close()
def test_pause(self):
# Create the process and wait until it is waiting
p = WaitForSignalProcess.new()
self.controller.insert_and_play(p)
wait_until(p, ProcessState.WAITING)
self.assertTrue(p.is_playing())
# Send a message asking the process to pause
self.assertIsNotNone(self.publisher.pause(p.pid, timeout=5.))
self.assertFalse(p.is_playing())
def test_pause_play(self):
# Create the process and wait until it is waiting
p = WaitForSignalProcess.new()
# Play
self.controller.insert_and_play(p)
self.assertTrue(wait_until(p, ProcessState.WAITING, 1))
self.assertTrue(p.is_playing())
# Pause
# Send a message asking the process to pause
self.assertIsNotNone(self.publisher.pause(p.pid, timeout=5.))
self.assertFalse(p.is_playing())
# Now ask it to continue
self.assertIsNotNone(self.publisher.play(p.pid, timeout=5.))
self.assertTrue(p.is_playing())
def test_abort(self):
# Create the process and wait until it is waiting
p = WaitForSignalProcess.new()
self.controller.insert_and_play(p)
self.assertTrue(wait_until(p, ProcessState.WAITING, timeout=2.))
self.assertTrue(p.is_playing())
# Send a message asking the process to abort
self.assertIsNotNone(
self.publisher.abort(p.pid, msg='Farewell', timeout=5.))
self.assertTrue(p.wait(timeout=2.), "Process failed to stop running")
self.assertFalse(p.is_playing())
self.assertTrue(p.has_aborted())
self.assertEqual(p.get_abort_msg(), 'Farewell')
def _create_connection(self):
return pika.BlockingConnection()
def _create_control_subscriber(self, connection):
return ProcessControlSubscriber(connection,
exchange=self.exchange,
process_controller=self.controller)
class TestStatusProvider(TestCase):
def setUp(self):
super(TestStatusProvider, self).setUp()
self._response = None
self._corr_id = None
try:
self._connection = pika.BlockingConnection()
except pika.exceptions.ConnectionClosed:
self.fail(
"Couldn't open connection. Make sure _rmq server is running")
self.channel = self._connection.channel()
# Set up the request exchange
self.request_exchange = '{}.{}.task_control'.format(
self.__class__, uuid.uuid4())
self.channel.exchange_declare(exchange=self.request_exchange,
type='fanout')
# Set up the response queue
result = self.channel.queue_declare(exclusive=True)
self.response_queue = result.method.queue
self.channel.basic_consume(self._on_response,
no_ack=True,
queue=self.response_queue)
self.controller = ProcessController()
self.provider = ProcessStatusSubscriber(
self._connection,
exchange=self.request_exchange,
process_controller=self.controller)
def tearDown(self):
self.controller.abort_all(timeout=10.)
self.assertEqual(self.controller.get_num_processes(), 0,
"Failed to abort all processes")
super(TestStatusProvider, self).tearDown()
self.channel.close()
self._connection.close()
def test_no_processes(self):
response = status_decode(self._send_and_get())
self.assertEqual(len(response[status.PROCS_KEY]), 0)
def test_status(self):
procs = []
for i in range(0, 20):
p = WaitForSignalProcess.new()
procs.append(p)
self.controller.insert_and_play(p)
self.assertTrue(wait_until(procs, ProcessState.WAITING, timeout=2.))
procs_dict = status_decode(self._send_and_get())[status.PROCS_KEY]
self.assertEqual(len(procs_dict), len(procs))
self.assertSetEqual(set([p.pid for p in procs]),
set(procs_dict.keys()))
playing = set([entry['playing'] for entry in procs_dict.itervalues()])
self.assertSetEqual(playing, {True})
self.assertTrue(self.controller.abort_all(timeout=5.),
"Couldn't abort all processes in timeout")
self.controller.remove_all()
response = status_decode(self._send_and_get())
self.assertEqual(len(response[status.PROCS_KEY]), 0)
def _send_and_get(self):
self._send_request()
self._get_response()
return self._response
def _send_request(self):
self._response = None
self._corr_id = str(uuid.uuid4())
self.channel.basic_publish(exchange=self.request_exchange,
routing_key='',
properties=pika.BasicProperties(
reply_to=self.response_queue,
correlation_id=self._corr_id),
body="")
def _get_response(self):
self.provider.poll(1)
self.channel.connection.process_data_events(time_limit=0.1)
def _on_response(self, ch, method, props, body):
if self._corr_id == props.correlation_id:
self._response = body
class ProcessLaunchSubscriber(Subscriber, ProcessListener):
"""
Run tasks as they come form the RabbitMQ task queue
.. warning:: the pika library used is not thread safe and as such the
connection passed to the constructor must be created on the same thread
as the start method is called.
"""
def __init__(self,
connection,
queue=Defaults.TASK_QUEUE,
decoder=json.loads,
process_controller=None):
"""
:param connection: The pika RabbitMQ connection
:type connection: :class:`pika.Connection`
:param queue: The queue name to use
:param decoder: A function to deserialise incoming messages
:param process_controller: The process controller to use, one will be created if None is passed
:type process_controller: :class:`plum.process_controller.ProcessController`
"""
if process_controller is None:
self._controller = ProcessController()
else:
self._controller = process_controller
self._decode = decoder
self._running_processes = {}
self._stopping = False
self._num_processes = 0
# Set up communications
self._channel = connection.channel()
self._channel.basic_qos(prefetch_count=1)
self._channel.queue_declare(queue=queue, durable=True)
self._channel.basic_consume(self._on_launch, queue=queue)
@override
def start(self, poll_time=1.0):
"""
Start polling for tasks. Will block until stop() is called.
.. warning:: Must be called from the same threads where the connection
that was passed into the constructor was created.
"""
while self._channel._consumer_infos:
self.poll(poll_time)
if self._stopping:
self._channel.stop_consuming()
self._stopping = False
@override
def poll(self, time_limit=1.0):
"""
Poll the channel for launch process events
:param time_limit: How long to poll for
:type time_limit: float
:return: The number of launch events consumed
:rtype: int
"""
self._num_processes = 0
self._channel.connection.process_data_events(time_limit=time_limit)
return self._num_processes
@override
def stop(self):
self._stopping = True
self._controller.pause_all()
@override
def shutdown(self):
self._channel.close()
def _on_launch(self, ch, method, properties, body):
"""
Consumer function that processes the launch message.
:param ch: The channel
:param method: The method
:param properties: The message properties
:param body: The message body
"""
self._num_processes += 1
task = self._decode(body)
proc_class = load_class(task['proc_class'])
proc = proc_class.new(inputs=task['inputs'])
proc.add_process_listener(self)
self._running_processes[proc.pid] = _RunningTaskInfo(
proc.pid, ch, method.delivery_tag)
self._controller.insert(proc)
self._controller.play(proc.pid)
# region From ProcessListener
def on_process_done_playing(self, process):
self._process_terminated(process)
# endregion
def _process_terminated(self, process):
"""
A process has finished for whatever reason so clean up.
:param process: The process that finished
:type process: :class:`plum.process.Process`
"""
info = self._running_processes.pop(process.pid)
info.ch.basic_ack(delivery_tag=info.delivery_tag)