def handle_worker_heartbeat(event): """ Celery event handler for 'worker-heartbeat' events. The event is first parsed and logged. Then the existing Worker objects are searched for one to update. If an existing one is found, it is updated. Otherwise a new Worker entry is created. Logging at the info and debug level is also done. :param event: A celery event to handle. :type event: dict """ event_info = _parse_and_log_event(event) find_worker_criteria = Criteria(filters={'_id': event_info['worker_name']}, fields=('_id', 'last_heartbeat')) find_worker_list = list(resources.filter_workers(find_worker_criteria)) if find_worker_list: Worker.get_collection().find_and_modify( query={'_id': event_info['worker_name']}, update={'$set': {'last_heartbeat': event_info['timestamp']}} ) else: new_worker = Worker(event_info['worker_name'], event_info['timestamp']) msg = _("New worker '%(worker_name)s' discovered") % event_info _logger.info(msg) new_worker.save()
def test__reserve_resource_with_existing_reservation(self): """ Test _reserve_resource() with a resource that has an existing reservation in the database. It should return the queue listed in the database, and increment the reservation counter. """ # Set up a worker with a reservation count of 1 now = datetime.utcnow() worker_1 = Worker(WORKER_1, now) worker_1.save() # Set up a resource reservation, using our worker from above reserved_resource_1 = ReservedResource('resource_1', worker_1.queue_name, 1) reserved_resource_1.save() # This should increase the reserved_resource_1 num_reservations to 2. worker_1's name should # be returned queue = tasks._reserve_resource('resource_1') self.assertEqual(queue, WORKER_1_QUEUE) # Make sure the ReservedResource is correct rrc = ReservedResource.get_collection() self.assertEqual(rrc.count(), 1) rr_1 = rrc.find_one({'_id': reserved_resource_1.name}) self.assertEqual(rr_1['assigned_queue'], WORKER_1_QUEUE) self.assertEqual(rr_1['num_reservations'], 2)
def test_resource_in_resource_map(self): """ Test _release_resource() with a valid resource. This should remove the resource from the database. """ # Set up two workers now = datetime.utcnow() worker_1 = Worker(WORKER_1, now) worker_1.save() worker_2 = Worker(WORKER_2, now) worker_2.save() # Set up two reserved resources reserved_resource_1 = ReservedResource(uuid.uuid4(), worker_1.name, 'resource_1') reserved_resource_1.save() reserved_resource_2 = ReservedResource(uuid.uuid4(), worker_2.name, 'resource_2') reserved_resource_2.save() # This should remove resource_2 from the _resource_map. tasks._release_resource(reserved_resource_2.task_id) # resource_2 should have been removed from the database rrc = ReservedResource.get_collection() self.assertEqual(rrc.count(), 1) rr_1 = rrc.find_one({'_id': reserved_resource_1.task_id}) self.assertEqual(rr_1['worker_name'], reserved_resource_1.worker_name) self.assertEqual(rr_1['resource_id'], 'resource_1')
def test__release_resource_task_count_one(self): """ Test _release_resource() with a resource that has a task count of one. This should remove the resource from the database. """ # Set up two workers now = datetime.utcnow() worker_1 = Worker(WORKER_1, now) worker_1.save() worker_2 = Worker(WORKER_2, now) worker_2.save() # Set up two reserved resources reserved_resource_1 = ReservedResource('resource_1', worker_1.name, 7) reserved_resource_1.save() reserved_resource_2 = ReservedResource('resource_2', worker_2.name, 1) reserved_resource_2.save() # This should remove resource_2 from the _resource_map. tasks._release_resource('resource_2') # resource_2 should have been removed from the database rrc = ReservedResource.get_collection() self.assertEqual(rrc.count(), 1) rr_1 = rrc.find_one({'_id': reserved_resource_1.name}) self.assertEqual(rr_1['assigned_queue'], reserved_resource_1.assigned_queue) self.assertEqual(rr_1['num_reservations'], 7)
def test__release_resource_task_count_two(self): """ Test _release_resource() with a resource that has a task count of two. This should simply decrement the task_count for the resource, but should not remove it from the database. """ # Set up two workers now = datetime.utcnow() worker_1 = Worker(WORKER_1, now) worker_1.save() worker_2 = Worker(WORKER_2, now) worker_2.save() # Set up two resource reservations, using our workers from above reserved_resource_1 = ReservedResource('resource_1', worker_1.name, 7) reserved_resource_1.save() reserved_resource_2 = ReservedResource('resource_2', worker_2.name, 2) reserved_resource_2.save() # This should reduce the reserved_resource_2 num_reservations to 1. tasks._release_resource('resource_2') # Make sure the ReservedResources are also correct rrc = ReservedResource.get_collection() self.assertEqual(rrc.count(), 2) rr_1 = rrc.find_one({'_id': reserved_resource_1.name}) self.assertEqual(rr_1['assigned_queue'], reserved_resource_1.assigned_queue) self.assertEqual(rr_1['num_reservations'], 7) rr_2 = rrc.find_one({'_id': reserved_resource_2.name}) self.assertEqual(rr_2['assigned_queue'], reserved_resource_2.assigned_queue) self.assertEqual(rr_2['num_reservations'], 1)
def test_ignores_queues_that_arent_workers(self): """ It is possible for the assigned_queue in a ReservedResource to reference a queue that is not in the workers collection. This test ensures that this queue is properly ignored, even if it is the most "enticing" choice. """ # Set up three Workers, with the least busy one in the middle so that we can # demonstrate that it did pick the least busy and not the last or first. now = datetime.utcnow() worker_1 = Worker('busy_worker', now) worker_2 = Worker('less_busy_worker', now) worker_3 = Worker('most_busy_worker', now) for worker in (worker_1, worker_2, worker_3): worker.save() # Now we need to make some reservations against these Workers' queues. We'll give worker_1 # 8 reservations, putting it in the middle of busyness. rr_1 = ReservedResource(name='resource_1', assigned_queue=worker_1.queue_name, num_reservations=8) # These next two will give worker_2 a total of 7 reservations, so it should get picked. rr_2 = ReservedResource(name='resource_2', assigned_queue=worker_2.queue_name, num_reservations=3) rr_3 = ReservedResource(name='resource_3', assigned_queue=worker_2.queue_name, num_reservations=4) # These next three will give worker_3 a total of 9 reservations, so it should be the most # busy. rr_4 = ReservedResource(name='resource_4', assigned_queue=worker_3.queue_name, num_reservations=2) rr_5 = ReservedResource(name='resource_5', assigned_queue=worker_3.queue_name, num_reservations=3) rr_6 = ReservedResource(name='resource_6', assigned_queue=worker_3.queue_name, num_reservations=4) # Now we will make a ReservedResource that references a queue that does not correspond to a # Worker and has the lowest reservation count. This RR should be ignored. rr_7 = ReservedResource(name='resource_7', assigned_queue='doesnt_exist', num_reservations=1) for rr in (rr_1, rr_2, rr_3, rr_4, rr_5, rr_6, rr_7): rr.save() worker = resources.get_least_busy_worker() self.assertEqual(type(worker), Worker) self.assertEqual(worker.name, 'less_busy_worker')
def test_post(self, mock_publish, mock_get_worker_for_reservation): """ Test that publish repo group creates a task for a worker. """ # Setup group_id = 'group-1' distributor_id = 'dist-1' mock_publish.apply_async_with_reservation.return_value = AsyncResult( '1234') self.manager.create_repo_group(group_id) self.distributor_manager.add_distributor(group_id, 'dummy-group-distributor', {}, distributor_id=distributor_id) mock_get_worker_for_reservation.return_value = Worker( 'some_queue', datetime.datetime.now()) # Test data = {'id': distributor_id} status, body = self.post( '/v2/repo_groups/%s/actions/publish/' % group_id, data) # Verify self.assertEqual(202, status) self.assertEqual(mock_publish.apply_async_with_reservation.call_count, 1)
def test_get_unreserved_worker_breaks_out_of_loop(self): self.mock_get_worker_for_reservation.side_effect = NoWorkers() self.mock_get_unreserved_worker.return_value = Worker( 'worker1', datetime.utcnow()) tasks._queue_reserved_task('task_name', 'my_task_id', 'my_resource_id', [1, 2], {'a': 2}) self.assertTrue(not self.mock_time.sleep.called)
def test_update_repo_and_plugins(self, distributor_update, mock_get_worker_for_reservation): """ Tests the aggregate call to update a repo and its plugins. """ mock_get_worker_for_reservation.return_value = Worker( 'some_queue', datetime.datetime.now()) self.manager.create_repo('repo-1', 'Original', 'Original Description') importer_manager = manager_factory.repo_importer_manager() distributor_manager = manager_factory.repo_distributor_manager() importer_manager.set_importer('repo-1', 'mock-importer', {'key-i1': 'orig-1'}) distributor_manager.add_distributor('repo-1', 'mock-distributor', {'key-d1': 'orig-1'}, True, distributor_id='dist-1') distributor_manager.add_distributor('repo-1', 'mock-distributor', {'key-d2': 'orig-2'}, True, distributor_id='dist-2') # Test repo_delta = {'display_name': 'Updated'} new_importer_config = {'key-i1': 'updated-1', 'key-i2': 'new-1'} new_distributor_configs = { 'dist-1': { 'key-d1': 'updated-1' }, } # only update one of the two distributors result = self.manager.update_repo_and_plugins('repo-1', repo_delta, new_importer_config, new_distributor_configs) self.assertTrue(isinstance(result, TaskResult)) self.assertEquals(None, result.error) repo = result.return_value # Verify self.assertEqual(repo['id'], 'repo-1') self.assertEqual(repo['display_name'], 'Updated') self.assertEqual(repo['description'], 'Original Description') importer = importer_manager.get_importer('repo-1') self.assertEqual(importer['config'], new_importer_config) dist_1 = distributor_manager.get_distributor('repo-1', 'dist-1') self.assertEqual(dist_1['config'], new_distributor_configs['dist-1']) dist_2 = distributor_manager.get_distributor('repo-1', 'dist-2') self.assertEqual(dist_2['config'], {'key-d2': 'orig-2'}) # There should have been a spawned task for the new distributor config expected_task_id = dispatch.TaskStatus.get_collection().find_one( {'tags': 'pulp:repository_distributor:dist-1'})['task_id'] self.assertEqual(result.spawned_tasks, [{'task_id': expected_task_id}])
def test_post_with_override_config(self, mock_get_worker_for_reservation, mock_uuid, mock_apply_async): # Setup uuid_list = [uuid.uuid4() for i in range(10)] mock_uuid.uuid4.side_effect = copy.deepcopy(uuid_list) expected_async_result = AsyncResult(str(uuid_list[0])) mock_get_worker_for_reservation.return_value = Worker('some_queue', datetime.datetime.now()) upload_id = self.upload_manager.initialize_upload() self.upload_manager.save_data(upload_id, 0, 'string data') repo_manager = manager_factory.repo_manager() repo_manager.create_repo('repo-upload') importer_manager = manager_factory.repo_importer_manager() importer_manager.set_importer('repo-upload', 'dummy-importer', {}) # Test test_override_config = {'key1': 'value1', 'key2': 'value2'} body = { 'upload_id' : upload_id, 'unit_type_id' : 'dummy-type', 'unit_key' : {'name' : 'foo'}, 'unit_metadata' : {'stuff' : 'bar'}, 'override_config': test_override_config, } status, body = self.post('/v2/repositories/repo-upload/actions/import_upload/', body) # Verify self.assertEqual(202, status) assert_body_matches_async_task(body, expected_async_result) exepcted_call_args = ['repo-upload', 'dummy-type', {'name': 'foo'}, {'stuff': 'bar'}, upload_id, test_override_config] self.assertEqual(exepcted_call_args, mock_apply_async.call_args[0][0])
def test_dispatches__release_resource(self): self.mock_get_worker_for_reservation.return_value = Worker( 'worker1', datetime.utcnow()) tasks._queue_reserved_task('task_name', 'my_task_id', 'my_resource_id', [1, 2], {'a': 2}) self.mock__release_resource.apply_async.assert_called_once_with( ('my_task_id', ), routing_key='worker1', exchange='C.dq')
def test_creates_and_saves_reserved_resource(self): self.mock_get_worker_for_reservation.return_value = Worker( 'worker1', datetime.utcnow()) tasks._queue_reserved_task('task_name', 'my_task_id', 'my_resource_id', [1, 2], {'a': 2}) self.mock_reserved_resource.assert_called_once_with( 'my_task_id', 'worker1', 'my_resource_id') self.mock_reserved_resource.return_value.save.assert_called_once_with()
def test__delete_worker(self, logger, cancel, mock_add_consumer): """ Assert that the correct Tasks get canceled when their Worker is deleted, and that the Worker is removed from the database. """ # cause two workers to be added to the database as having workers worker_watcher.handle_worker_heartbeat({ 'timestamp': time.time(), 'type': 'worker-heartbeat', 'hostname': WORKER_1, }) worker_watcher.handle_worker_heartbeat({ 'timestamp': time.time(), 'type': 'worker-heartbeat', 'hostname': WORKER_2, }) # Let's simulate three tasks being assigned to WORKER_2, with two of them being # in an incomplete state and one in a complete state. We will delete WORKER_2, # which should cause the two to get canceled. Let's put task_1 in progress TaskStatusManager.create_task_status('task_1', WORKER_2_QUEUE, state=CALL_RUNNING_STATE) TaskStatusManager.create_task_status('task_2', WORKER_2_QUEUE, state=CALL_WAITING_STATE) # This task shouldn't get canceled because it isn't in an incomplete state TaskStatusManager.create_task_status('task_3', WORKER_2_QUEUE, state=CALL_FINISHED_STATE) # Let's make a task in a worker that is still present just to make sure it isn't touched. TaskStatusManager.create_task_status('task_4', WORKER_1_QUEUE, state=CALL_RUNNING_STATE) # Let's just make sure the setup worked and that we have a Worker with RR2 worker_collection = Worker.get_collection() self.assertEqual(worker_collection.find({'_id': WORKER_2}).count(), 1) # Now let's delete the Worker named WORKER_2 tasks._delete_worker.apply_async(args=(WORKER_2, ), queue=tasks.RESOURCE_MANAGER_QUEUE) # cancel() should have been called twice with task_1 and task_2 as parameters self.assertEqual(cancel.call_count, 2) # Let's build a set out of the two times that cancel was called. We can't know for sure # which order the Tasks got canceled in, but we can assert that the correct two tasks were # canceled (task_3 should not appear in this set). cancel_param_set = set([c[1] for c in cancel.mock_calls]) self.assertEqual(cancel_param_set, set([('task_1', ), ('task_2', )])) # We should have logged that we are canceling the tasks self.assertEqual(logger.call_count, 0) self.assertTrue(WORKER_2 in logger.mock_calls[0][1][0]) self.assertTrue('Canceling the tasks' in logger.mock_calls[0][1][0]) # The Worker should have been deleted self.assertEqual(worker_collection.find({'_id': WORKER_2}).count(), 0) # the Worker for RW1 should remain self.assertEqual(worker_collection.find({'_id': WORKER_1}).count(), 1)
def test__reserve_resource_without_existing_reservation(self): """ Test _reserve_resource() with a resource that does not have an existing reservation in the database. It should find the least busy worker, add a reservation to the database with that worker's queue, and then return the queue name. """ # Set up a worker worker_1 = Worker(WORKER_1, datetime.utcnow()) worker_1.save() queue = tasks._reserve_resource('resource_1') self.assertEqual(queue, WORKER_1_QUEUE) # Make sure the ReservedResource is correct rrc = ReservedResource.get_collection() self.assertEqual(rrc.count(), 1) rr_1 = rrc.find_one({'_id': 'resource_1'}) self.assertEqual(rr_1['assigned_queue'], WORKER_1_QUEUE) self.assertEqual(rr_1['num_reservations'], 1)
def test_picks_least_busy_worker(self): """ Test that the function picks the least busy worker. """ # Set up three Workers, with the least busy one in the middle so that we can # demonstrate that it did pick the least busy and not the last or first. now = datetime.utcnow() worker_1 = Worker('busy_worker', now) worker_2 = Worker('less_busy_worker', now) worker_3 = Worker('most_busy_worker', now) for worker in (worker_1, worker_2, worker_3): worker.save() # Now we need to make some reservations against these Workers' queues. We'll give worker_1 # 8 reservations, putting it in the middle of busyness. rr_1 = ReservedResource(name='resource_1', assigned_queue=worker_1.queue_name, num_reservations=8) # These next two will give worker_2 a total of 7 reservations, so it should get picked. rr_2 = ReservedResource(name='resource_2', assigned_queue=worker_2.queue_name, num_reservations=3) rr_3 = ReservedResource(name='resource_3', assigned_queue=worker_2.queue_name, num_reservations=4) # These next three will give worker_3 a total of 9 reservations, so it should be the most # busy. rr_4 = ReservedResource(name='resource_4', assigned_queue=worker_3.queue_name, num_reservations=2) rr_5 = ReservedResource(name='resource_5', assigned_queue=worker_3.queue_name, num_reservations=3) rr_6 = ReservedResource(name='resource_6', assigned_queue=worker_3.queue_name, num_reservations=4) for rr in (rr_1, rr_2, rr_3, rr_4, rr_5, rr_6): rr.save() worker = resources.get_least_busy_worker() self.assertEqual(type(worker), Worker) self.assertEqual(worker.name, 'less_busy_worker')
def test_post(self, mock_get_worker_for_reservation): # Setup group_id = 'group-1' distributor_id = 'dist-1' self.manager.create_repo_group(group_id) self.distributor_manager.add_distributor(group_id, 'dummy-group-distributor', {}, distributor_id=distributor_id) mock_get_worker_for_reservation.return_value = Worker('some_queue', datetime.datetime.now()) # Test data = {'id' : distributor_id} status, body = self.post('/v2/repo_groups/%s/actions/publish/' % group_id, data)
def GET(self, task_id): try: task = TaskStatus.objects.get(task_id=task_id) except DoesNotExist: raise MissingResource(task_id) task_dict = task_serializer(task) if 'worker_name' in task_dict: queue_name = Worker(task_dict['worker_name'], datetime.now()).queue_name task_dict.update({'queue': queue_name}) return self.ok(task_dict)
def test__delete_worker(self, logger, cancel, mock_add_consumer): """ Assert that the correct Tasks get canceled when their Worker is deleted, and that the Worker is removed from the database. """ # cause two workers to be added to the database as having workers worker_watcher.handle_worker_heartbeat({ 'timestamp': time.time(), 'type': 'worker-heartbeat', 'hostname': WORKER_1, }) worker_watcher.handle_worker_heartbeat({ 'timestamp': time.time(), 'type': 'worker-heartbeat', 'hostname': WORKER_2, }) # Let's simulate three tasks being assigned to WORKER_2, with two of them being # in an incomplete state and one in a complete state. We will delete WORKER_2, # which should cause the two to get canceled. Let's put task_1 in progress TaskStatusManager.create_task_status('task_1', WORKER_2_QUEUE, state=CALL_RUNNING_STATE) TaskStatusManager.create_task_status('task_2', WORKER_2_QUEUE, state=CALL_WAITING_STATE) # This task shouldn't get canceled because it isn't in an incomplete state TaskStatusManager.create_task_status('task_3', WORKER_2_QUEUE, state=CALL_FINISHED_STATE) # Let's make a task in a worker that is still present just to make sure it isn't touched. TaskStatusManager.create_task_status('task_4', WORKER_1_QUEUE, state=CALL_RUNNING_STATE) # Let's just make sure the setup worked and that we have a Worker with RR2 worker_collection = Worker.get_collection() self.assertEqual(worker_collection.find({'_id': WORKER_2}).count(), 1) # Now let's delete the Worker named WORKER_2 tasks._delete_worker.apply_async(args=(WORKER_2,), queue=tasks.RESOURCE_MANAGER_QUEUE) # cancel() should have been called twice with task_1 and task_2 as parameters self.assertEqual(cancel.call_count, 2) # Let's build a set out of the two times that cancel was called. We can't know for sure # which order the Tasks got canceled in, but we can assert that the correct two tasks were # canceled (task_3 should not appear in this set). cancel_param_set = set([c[1] for c in cancel.mock_calls]) self.assertEqual(cancel_param_set, set([('task_1',), ('task_2',)])) # We should have logged that we are canceling the tasks self.assertEqual(logger.call_count, 0) self.assertTrue(WORKER_2 in logger.mock_calls[0][1][0]) self.assertTrue('Canceling the tasks' in logger.mock_calls[0][1][0]) # The Worker should have been deleted self.assertEqual(worker_collection.find({'_id': WORKER_2}).count(), 0) # the Worker for RW1 should remain self.assertEqual(worker_collection.find({'_id': WORKER_1}).count(), 1)
def test_DELETE_completed_celery_task(self): """ Test the DELETE() method raises a TaskComplete exception if the task is already complete. """ task_id = '1234abcd' now = datetime.utcnow() test_worker = Worker('test_worker', now) TaskStatusManager.create_task_status( task_id, test_worker.name, state=constants.CALL_FINISHED_STATE) self.assertRaises(PulpCodedException, self.task_resource.DELETE, task_id)
def test_cancel_after_task_finished(self, logger, revoke): task_id = '1234abcd' now = datetime.utcnow() test_worker = Worker('test_worker', now) TaskStatusManager.create_task_status(task_id, test_worker.name, state=CALL_FINISHED_STATE) self.assertRaises(PulpCodedException, tasks.cancel, task_id) task_status = TaskStatusManager.find_by_task_id(task_id) self.assertEqual(task_status['state'], CALL_FINISHED_STATE)
def GET(self, task_id): task = TaskStatusManager.find_by_task_id(task_id) if task is None: raise MissingResource(task_id) else: link = serialization.link.link_obj('/pulp/api/v2/tasks/%s/' % task_id) task.update(link) task.update(serialization.dispatch.spawned_tasks(task)) if 'worker_name' in task: queue_name = Worker(task['worker_name'], datetime.now()).queue_name task.update({'queue': queue_name}) return self.ok(task)
def test_dispatches_inner_task(self): self.mock_get_worker_for_reservation.return_value = Worker( 'worker1', datetime.utcnow()) tasks._queue_reserved_task('task_name', 'my_task_id', 'my_resource_id', [1, 2], {'a': 2}) apply_async = self.mock_celery.tasks['task_name'].apply_async apply_async.assert_called_once_with(1, 2, a=2, routing_key='worker1', task_id='my_task_id', exchange='C.dq')
def handle_worker_heartbeat(event): """ Celery event handler for 'worker-heartbeat' events. The event is first parsed and logged. If this event is from the resource manager, there is no further processing to be done. Then the existing Worker objects are searched for one to update. If an existing one is found, it is updated. Otherwise a new Worker entry is created. Logging at the info and debug level is also done. :param event: A celery event to handle. :type event: dict """ event_info = _parse_and_log_event(event) # if this is the resource_manager do nothing if _is_resource_manager(event): return find_worker_criteria = Criteria(filters={'_id': event_info['worker_name']}, fields=('_id', 'last_heartbeat', 'num_reservations')) find_worker_list = list(resources.filter_workers(find_worker_criteria)) if find_worker_list: Worker.get_collection().find_and_modify( query={'_id': event_info['worker_name']}, update={'$set': { 'last_heartbeat': event_info['timestamp'] }}) else: new_worker = Worker(event_info['worker_name'], event_info['timestamp']) msg = _("New worker '%(worker_name)s' discovered") % event_info _logger.info(msg) new_worker.save()
def test_DELETE_celery_task(self, revoke): """ Test the DELETE() method with a UUID that does not correspond to a UUID that the coordinator is aware of. This should cause a revoke call to Celery's Controller. """ task_id = '1234abcd' now = datetime.utcnow() test_worker = Worker('test_worker', now) TaskStatusManager.create_task_status(task_id, test_worker.name) self.task_resource.DELETE(task_id) revoke.assert_called_once_with(task_id, terminate=True)
def test_cancel_successful(self, logger, revoke): task_id = '1234abcd' now = datetime.utcnow() test_worker = Worker('test_worker', now) TaskStatusManager.create_task_status(task_id, test_worker.name) tasks.cancel(task_id) revoke.assert_called_once_with(task_id, terminate=True) self.assertEqual(logger.info.call_count, 1) log_msg = logger.info.mock_calls[0][1][0] self.assertTrue(task_id in log_msg) self.assertTrue('Task canceled' in log_msg) task_status = TaskStatusManager.find_by_task_id(task_id) self.assertEqual(task_status['state'], CALL_CANCELED_STATE)
def _delete_worker(name, normal_shutdown=False): """ Delete the Worker with _id name from the database, cancel any associated tasks and reservations If the worker shutdown normally, no message is logged, otherwise an error level message is logged. Default is to assume the worker did not shut down normally. Any resource reservations associated with this worker are cleaned up by this function. Any tasks associated with this worker are explicitly canceled. :param name: The name of the worker you wish to delete. In the database, the _id field is the name. :type name: basestring :param normal_shutdown: True if the worker shutdown normally, False otherwise. Defaults to False. :type normal_shutdown: bool """ if normal_shutdown is False: msg = _( 'The worker named %(name)s is missing. Canceling the tasks in its queue.' ) msg = msg % {'name': name} logger.error(msg) # Delete the worker document worker_list = list( resources.filter_workers(Criteria(filters={'_id': name}))) if len(worker_list) > 0: worker_document = worker_list[0] worker_document.delete() # Delete all reserved_resource documents for the worker ReservedResource.get_collection().remove({'worker_name': name}) # Cancel all of the tasks that were assigned to this worker's queue worker = Worker.from_bson({'_id': name}) for task in TaskStatusManager.find_by_criteria( Criteria( filters={ 'worker_name': worker.name, 'state': { '$in': constants.CALL_INCOMPLETE_STATES } })): cancel(task['task_id'])
def test__release_resource_not_in__resource_map(self): """ Test _release_resource() with a resource that is not in the database. This should be gracefully handled, and result in no changes to the database. """ # Set up two workers worker_1 = Worker(WORKER_1, datetime.utcnow()) worker_1.save() worker_2 = Worker(WORKER_2, datetime.utcnow()) worker_2.save() # Set up two resource reservations, using our workers from above reserved_resource_1 = ReservedResource('resource_1', worker_1.name, 7) reserved_resource_1.save() reserved_resource_2 = ReservedResource('resource_2', worker_2.name, 3) reserved_resource_2.save() # This should not raise any Exception, but should also not alter either the Worker # collection or the ReservedResource collection tasks._release_resource('made_up_resource_id') # Make sure that the workers collection has not been altered worker_collection = Worker.get_collection() self.assertEqual(worker_collection.count(), 2) worker_1 = worker_collection.find_one({'_id': worker_1.name}) self.assertTrue(worker_1) worker_2 = worker_collection.find_one({'_id': worker_2.name}) self.assertTrue(worker_2) # Make sure that the reserved resources collection has not been altered rrc = ReservedResource.get_collection() self.assertEqual(rrc.count(), 2) rr_1 = rrc.find_one({'_id': reserved_resource_1.name}) self.assertEqual(rr_1['assigned_queue'], reserved_resource_1.assigned_queue) self.assertEqual(rr_1['num_reservations'], 7) rr_2 = rrc.find_one({'_id': reserved_resource_2.name}) self.assertEqual(rr_2['assigned_queue'], reserved_resource_2.assigned_queue) self.assertEqual(rr_2['num_reservations'], 3)
def test_filter(self): """ Test a filter operation to make sure the results appear to be correct. """ # Make three workers. We'll filter for two of them. now = datetime.utcnow() kw_1 = Worker('worker_1', now) kw_1.save() kw_2 = Worker('worker_2', now) kw_2.save() kw_3 = Worker('worker_3', now) kw_3.save() criteria = Criteria(filters={'_id': {'$gt': 'worker_1'}}, sort=[('_id', pymongo.ASCENDING)]) workers = resources.filter_workers(criteria) # Let's assert that workers is a generator, and then let's cast it to a list so it's easier # to test that we got the correct instances back. self.assertEqual(type(workers), types.GeneratorType) workers = list(workers) self.assertEqual(all([isinstance(w, Worker) for w in workers]), True) self.assertEqual(workers[0].name, 'worker_2') self.assertEqual(workers[1].name, 'worker_3')
def test_apply_async_task_canceled(self, apply_async): """ Assert that apply_async() honors 'canceled' task status. """ args = [1, 'b', 'iii'] kwargs = {'1': 'for the money', 'tags': ['test_tags']} task_id = 'test_task_id' now = datetime.utcnow() TaskStatusManager.create_task_status(task_id, Worker('test-worker', now), state=CALL_CANCELED_STATE) apply_async.return_value = celery.result.AsyncResult(task_id) task = tasks.Task() task.apply_async(*args, **kwargs) task_status = TaskStatusManager.find_by_task_id(task_id) self.assertEqual(task_status['state'], 'canceled') self.assertEqual(task_status['start_time'], None)
def _delete_worker(name, normal_shutdown=False): """ Delete the Worker with _id name from the database, cancel any associated tasks and reservations If the worker shutdown normally, no message is logged, otherwise an error level message is logged. Default is to assume the worker did not shut down normally. Any resource reservations associated with this worker are cleaned up by this function. Any tasks associated with this worker are explicitly canceled. :param name: The name of the worker you wish to delete. In the database, the _id field is the name. :type name: basestring :param normal_shutdown: True if the worker shutdown normally, False otherwise. Defaults to False. :type normal_shutdown: bool """ if normal_shutdown is False: msg = _('The worker named %(name)s is missing. Canceling the tasks in its queue.') msg = msg % {'name': name} _logger.error(msg) # Delete the worker document worker_list = list(resources.filter_workers(Criteria(filters={'_id': name}))) if len(worker_list) > 0: worker_document = worker_list[0] worker_document.delete() # Delete all reserved_resource documents for the worker ReservedResource.get_collection().remove({'worker_name': name}) # Cancel all of the tasks that were assigned to this worker's queue worker = Worker.from_bson({'_id': name}) for task_status in TaskStatus.objects(worker_name=worker.name, state__in=constants.CALL_INCOMPLETE_STATES): cancel(task_status['task_id'])
def test_DELETE_doesnt_cancel_spawned_celery_task(self, revoke): """ Test the DELETE() which should cause a revoke call to Celery's Controller. This also tests that the spawned tasks are canceled as well. """ task_id = '1234abcd' spawned_task_id = 'spawned_task' spawned_by_spawned_task_id = 'spawned_by_spawned_task' now = datetime.utcnow() test_worker = Worker('test_worker', now) TaskStatusManager.create_task_status(task_id, test_worker.queue_name) TaskStatusManager.create_task_status(spawned_task_id, test_worker.queue_name) TaskStatusManager.create_task_status(spawned_by_spawned_task_id, test_worker.queue_name) TaskStatusManager.update_task_status( task_id, delta={'spawned_tasks': [spawned_task_id]}) TaskStatusManager.update_task_status( spawned_task_id, delta={'spawned_tasks': [spawned_by_spawned_task_id]}) self.task_resource.DELETE(task_id) self.assertEqual(revoke.call_count, 1) revoke.assert_called_once_with(task_id, terminate=True)
def test_resource_not_in_resource_map(self): """ Test _release_resource() with a resource that is not in the database. This should be gracefully handled, and result in no changes to the database. """ # Set up two workers worker_1 = Worker(WORKER_1, datetime.utcnow()) worker_1.save() worker_2 = Worker(WORKER_2, datetime.utcnow()) worker_2.save() # Set up two resource reservations, using our workers from above reserved_resource_1 = ReservedResource(uuid.uuid4(), worker_1.name, 'resource_1') reserved_resource_1.save() reserved_resource_2 = ReservedResource(uuid.uuid4(), worker_2.name, 'resource_2') reserved_resource_2.save() # This should not raise any Exception, but should also not alter either the Worker # collection or the ReservedResource collection tasks._release_resource('made_up_resource_id') # Make sure that the workers collection has not been altered worker_collection = Worker.get_collection() self.assertEqual(worker_collection.count(), 2) worker_1 = worker_collection.find_one({'_id': worker_1.name}) self.assertTrue(worker_1) worker_2 = worker_collection.find_one({'_id': worker_2.name}) self.assertTrue(worker_2) # Make sure that the reserved resources collection has not been altered rrc = ReservedResource.get_collection() self.assertEqual(rrc.count(), 2) rr_1 = rrc.find_one({'_id': reserved_resource_1.task_id}) self.assertEqual(rr_1['worker_name'], reserved_resource_1.worker_name) self.assertEqual(rr_1['resource_id'], 'resource_1') rr_2 = rrc.find_one({'_id': reserved_resource_2.task_id}) self.assertEqual(rr_2['worker_name'], reserved_resource_2.worker_name) self.assertEqual(rr_2['resource_id'], 'resource_2')
def tearDown(self): Worker.get_collection().remove() ReservedResource.get_collection().remove() TaskStatus.objects().delete()
def tearDown(self): Worker.get_collection().remove() ReservedResource.get_collection().remove() TaskStatus.get_collection().remove()
def test_filter(self): """ Test a filter operation to make sure the results appear to be correct. """ # Make three workers. We'll filter for two of them. now = datetime.utcnow() kw_1 = Worker('worker_1', now) kw_1.save() kw_2 = Worker('worker_2', now) kw_2.save() kw_3 = Worker('worker_3', now) kw_3.save() criteria = Criteria(filters={'_id': { '$gt': 'worker_1' }}, sort=[('_id', pymongo.ASCENDING)]) workers = resources.filter_workers(criteria) # Let's assert that workers is a generator, and then let's cast it to a list so it's easier # to test that we got the correct instances back. self.assertEqual(type(workers), types.GeneratorType) workers = list(workers) self.assertEqual(all([isinstance(w, Worker) for w in workers]), True) self.assertEqual(workers[0].name, 'worker_2') self.assertEqual(workers[1].name, 'worker_3')