def setUp(self):
sqs = boto3.resource('sqs')
aws_task_completion_queue_name = 'task-completion-test-' + str(random.randint(1, 100))
self._aws_task_completion_queue = sqs.create_queue(QueueName=aws_task_completion_queue_name)
logging.info('create sqs queue: ' + self._aws_task_completion_queue.url)
self._local_task_output_queue = Queue.Queue(10)
self._output_pusher = OutputPusher(self._local_task_output_queue,self._aws_task_completion_queue)
self._output_pusher.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_output_pusher(self,):
self._output_pusher = OutputPusher(self._local_task_completion_queue,self._aws_task_completion_queue)
class OutputPusherTest(unittest.TestCase):
def setUp(self):
sqs = boto3.resource('sqs')
aws_task_completion_queue_name = 'task-completion-test-' + str(random.randint(1, 100))
self._aws_task_completion_queue = sqs.create_queue(QueueName=aws_task_completion_queue_name)
logging.info('create sqs queue: ' + self._aws_task_completion_queue.url)
self._local_task_output_queue = Queue.Queue(10)
self._output_pusher = OutputPusher(self._local_task_output_queue,self._aws_task_completion_queue)
self._output_pusher.start()
def tearDown(self):
self._output_pusher.stop()
self._output_pusher.join()
if self._aws_task_completion_queue is not None:
self._aws_task_completion_queue.delete()
def test_should_send_completion_given_aws_queue_when_completed(self):
#construct
params = {
"ratio":1e5,
"num_iters":1,
"length":512,
"verbose":False
}
art_task = ArtTask('id12345',params)
art_task.content_image = caffe.io.load_image('test-resource/content-image/face-1.jpg')
art_task.style_image = caffe.io.load_image('resource/style-image/starry_night.jpg')
art_task.output_image = caffe.io.load_image('test-resource/content-image/face-1.jpg')
self._local_task_output_queue.put_nowait(art_task)
time.sleep(30)
# check
messages = self._aws_task_completion_queue.receive_messages()
self.assertEqual(1, len(messages))
completion_message = messages[0]
completion = json.loads(completion_message.body)
self.assertEquals('id12345',completion['_id'])
self.assertIsNotNone(completion['outputImageUrl'])
def test_should_upload_output_image_given_aws_queue_when_completed(self):
#construct
params = {
"ratio":1e5,
"num_iters":1,
"length":512,
"verbose":False
}
art_task = ArtTask('id'+str(random.randint(1000000, 9999999)),params)
art_task.content_image = caffe.io.load_image('test-resource/content-image/face-1.jpg')
art_task.style_image = caffe.io.load_image('resource/style-image/starry_night.jpg')
art_task.output_image = caffe.io.load_image('test-resource/content-image/face-1.jpg')
self._local_task_output_queue.put_nowait(art_task)
time.sleep(30)
# verify
s3_bucket = os.environ['S3_IMAGE_BUCKET']
output_image_url = 'http://'+s3_bucket+'.s3.amazonaws.com/' + datetime.date.today().strftime('%Y%m%d') + '/'+art_task._id+'/'+art_task._id+'-output.jpg'
output_image = caffe.io.load_image(output_image_url)
self.assertIsNotNone(output_image)
logging.info('verify uploaded image pass')
