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)