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')