def setUp(self):
sqs = boto3.resource('sqs')
aws_task_request_queue_name = 'task-request-test-' + str(random.randint(1, 100))
self._aws_task_request_queue = sqs.create_queue(QueueName=aws_task_request_queue_name)
logging.info('create sqs queue: ' + self._aws_task_request_queue.url)
self._local_task_request_queue = Queue.Queue(10)
self._task_fetcher = TaskFetcher(self._aws_task_request_queue,self._local_task_request_queue)
self._task_fetcher.start()
class Master(object):
'''
control all workers
'''
def __init__(self,params):
self._params = params
self._local_task_request_queue = Queue.Queue(256)
self._local_task_completion_queue = Queue.Queue(1024)
sqs = boto3.resource('sqs')
self._aws_task_request_queue = sqs.get_queue_by_name(QueueName=self._params['aws_task_request_queue_name'])
self._aws_task_completion_queue = sqs.get_queue_by_name(QueueName=self._params['aws_task_completion_queue_name'])
self._workers = [];
self._task_fetcher = None
self._output_pusher = None
self._prepare_worker(self._params['num_of_worker'])
self._prepare_task_fetcher()
self._prepare_output_pusher()
def _prepare_worker(self,num_of_workers=1):
for num in range(0,num_of_workers):
worker = Worker(str(num),self._local_task_request_queue,self._local_task_completion_queue)
self._workers.append(worker)
def _prepare_task_fetcher(self):
self._task_fetcher = TaskFetcher(self._aws_task_request_queue,self._local_task_request_queue)
def _prepare_output_pusher(self,):
self._output_pusher = OutputPusher(self._local_task_completion_queue,self._aws_task_completion_queue)
def start(self):
self._task_fetcher.start()
#self._task_fetcher.join()
self._output_pusher.start()
#self._output_pusher.join()
for worker in self._workers:
worker.start()
#worker.join()
def stop(self):
self._task_fetcher.stop()
self._output_pusher.stop()
for worker in self._workers:
worker.stop()
def _prepare_task_fetcher(self):
self._task_fetcher = TaskFetcher(self._aws_task_request_queue,self._local_task_request_queue)
class TaskFetcherTest(unittest.TestCase):
def setUp(self):
sqs = boto3.resource('sqs')
aws_task_request_queue_name = 'task-request-test-' + str(random.randint(1, 100))
self._aws_task_request_queue = sqs.create_queue(QueueName=aws_task_request_queue_name)
logging.info('create sqs queue: ' + self._aws_task_request_queue.url)
self._local_task_request_queue = Queue.Queue(10)
self._task_fetcher = TaskFetcher(self._aws_task_request_queue,self._local_task_request_queue)
self._task_fetcher.start()
def tearDown(self):
self._task_fetcher.stop()
self._task_fetcher.join()
if self._aws_task_request_queue is not None:
self._aws_task_request_queue.delete()
def test_should_get_task_given_task_request_queue_when_send_task_request(self):
# aws queue is created
# given task request to aws task request queue
art_task_request = {'_id':'id12345',
'contentImageUrl':'https://www.baidu.com/img/bd_logo1.png',
'style':'starry_night',
'params':{
"ratio":1e5,
"num_iters":10,
"length":512,
"verbose":False
}}
self._aws_task_request_queue.send_message(MessageBody=json.dumps(art_task_request))
art_task = self._local_task_request_queue.get(True, 30)
self.assertIsNotNone(art_task)
self.assertEquals('id12345',art_task._id)
def test_should_get_task_with_style_image_given_task_request_queue_when_send_task_request(self):
# given task request to aws task request queue
art_task_request = {'_id':'id12345',
'contentImageUrl':'https://www.baidu.com/img/bd_logo1.png',
'style':'starry_night',
'params':{
"ratio":1e5,
"num_iters":10,
"length":512,
"verbose":False
}}
self._aws_task_request_queue.send_message(MessageBody=json.dumps(art_task_request))
art_task = self._local_task_request_queue.get(True, 30)
self.assertIsNotNone(art_task.style_image)
self.assertEquals(numpy.ndarray,type(art_task.style_image))
def test_should_get_task_with_content_image_given_task_request_queue_when_send_task_request(self):
# given task request to aws task request queue
art_task_request = {'_id':'id12345',
'contentImageUrl':'https://www.baidu.com/img/bd_logo1.png',
'style':'starry_night',
'params':{
"ratio":1e5,
"num_iters":10,
"length":512,
"verbose":False
}}
self._aws_task_request_queue.send_message(MessageBody=json.dumps(art_task_request))
art_task = self._local_task_request_queue.get(True, 30)
self.assertIsNotNone(art_task.content_image)
self.assertEquals(numpy.ndarray,type(art_task.content_image))
