def _get_kwc():
routing_map = {
'KNN': {
'model-1': [{
'host': 'localhost',
'port': 9090,
'local_worker_id': 0
}, {
'host': 'localhost',
'port': 9091,
'local_worker_id': 1
}, {
'host': 'localhost',
'port': 9092,
'local_worker_id': 2
}, {
'host': 'localhost',
'port': 9093,
'local_worker_id': 3
}, {
'host': 'localhost',
'port': 9094,
'local_worker_id': 4
}, {
'host': 'localhost',
'port': 9095,
'local_worker_id': 5
}]
}
}
wlb = WorkerLoadBalancer()
wlb.update_workers_list(routing_map)
kwc = KnnWorkerClient(wlb)
return kwc
def _get_kwc(self):
routing_map = {
'KNN': {
'model-1': [{
'host': 'localhost',
'port': 9090,
'local_worker_id': 0
}]
}
}
wlb = WorkerLoadBalancer()
wlb.update_workers_list(routing_map)
kwc = KnnWorkerClient(wlb)
return kwc
def loader():
setup_logging()
logger = logging.getLogger(__name__)
logger.info("Starting Router...")
start_time = int(time.time())
load_balancer = WorkerLoadBalancer()
caffe_model_resource = CaffeModelResource(load_balancer)
model_workers_resource = ModelWorkersResource(load_balancer)
knn_worker_client = KnnWorkerClient(load_balancer)
record_resource = Record(knn_worker_client)
predict_resource = Predict(knn_worker_client)
elb_resource = ElbResource(start_time, load_balancer)
# falcon.API instances are callable WSGI apps
app = falcon.API()
app.add_route('/v1/caffe_model/predict', caffe_model_resource)
app.add_route('/v1/model_workers_map', model_workers_resource)
app.add_route('/v1/knn_model/update', record_resource)
app.add_route('/v1/knn_model/predict', predict_resource)
app.add_route('/elb-healthcheck', elb_resource)
return app
def loader():
setup_logging()
logger = logging.getLogger(__name__)
logger.info("Starting Router...")
start_time = int(time.time())
load_balancer = WorkerLoadBalancer()
model_workers_resource = ModelWorkersResource(load_balancer)
delWorkersResource = DelWorkersResource(load_balancer)
AI_worker_client = AIWorkerClient(load_balancer)
predict_resource = AIModelResource(AI_worker_client)
elb_resource = ElbResource(start_time, load_balancer)
# falcon.API instances are callable WSGI apps
app = falcon.API()
app.add_route('/v1/model_workers_add', model_workers_resource)
app.add_route('/v1/model_workers_remove', delWorkersResource)
app.add_route('/v1/AI_model/predict', predict_resource)
app.add_route('/elb-healthcheck', elb_resource)
return app
def _get_kwc(self):
routing_map = {'KNN': {'model-1':[{'host': 'localhost', 'port': 9090, 'local_worker_id': 0}]}}
wlb = WorkerLoadBalancer()
wlb.update_workers_list(routing_map)
kwc = KnnWorkerClient(wlb)
return kwc
def setUp(self):
"""
This method is called before each test
"""
self.worker_load_balancer = WorkerLoadBalancer()
class WorkerLoadBalancerTestCase(unittest.TestCase):
"""
Tests for worker_load_balancer.py
"""
def setUp(self):
"""
This method is called before each test
"""
self.worker_load_balancer = WorkerLoadBalancer()
def tearDown(self):
"""
This method is called after each test
"""
pass
def testInit(self):
self.assertEqual(self.worker_load_balancer.get_all_workers("Caffe"), None)
self.assertEqual(self.worker_load_balancer.get_all_workers("KNN"), None)
self.assertEqual(self.worker_load_balancer.get_model_to_workers_list("Caffe"), None)
self.assertEqual(self.worker_load_balancer.get_model_to_workers_list("KNN"), None)
self.assertEqual(self.worker_load_balancer.get_worker_info("worker-1", "Caffe"), None)
self.assertEqual(self.worker_load_balancer.get_worker_info("worker-1", "KNN"), None)
def test_update_workers_list(self):
self.assertRaises(Exception, lambda: self.worker_load_balancer.choose_worker("Caffe", "model-1"))
self.assertEqual(self.worker_load_balancer.get_all_workers("Caffe"), None)
self.assertEqual(self.worker_load_balancer.get_all_workers("KNN"), None)
self.assertEqual(self.worker_load_balancer.get_model_to_workers_list("Caffe"), None)
self.assertEqual(self.worker_load_balancer.get_model_to_workers_list("KNN"), None)
self.assertEqual(self.worker_load_balancer.get_worker_info("worker-1", "Caffe"), None)
self.assertEqual(self.worker_load_balancer.get_worker_info("worker-1", "KNN"), None)
model_type_to_model_to_workers_map = {}
caffe_model_to_workers_map = {}
knn_model_to_workers_map = {}
caffe_worker_01 = {}
caffe_worker_01["host"] = 'localhost'
caffe_worker_01["port"] = 9090
caffe_worker_01["local_worker_id"] = 0
caffe_worker_02 = {}
caffe_worker_02["host"] = 'localhost'
caffe_worker_02["port"] = 9091
caffe_worker_02["local_worker_id"] = 1
caffe_workers = []
caffe_workers.append(caffe_worker_01)
caffe_workers.append(caffe_worker_02)
knn_worker_01 = {}
knn_worker_01["host"] = 'localhost'
knn_worker_01["port"] = 9092
knn_worker_01["local_worker_id"] = 2
knn_worker_02 = {}
knn_worker_02["host"] = 'localhost'
knn_worker_02["port"] = 9093
knn_worker_02["local_worker_id"] = 3
knn_workers = []
knn_workers.append(knn_worker_01)
knn_workers.append(knn_worker_02)
caffe_model_to_workers_map["model-1"] = caffe_workers
knn_model_to_workers_map["model-2"] = knn_workers
model_type_to_model_to_workers_map["Caffe"] = caffe_model_to_workers_map
model_type_to_model_to_workers_map["KNN"] = knn_model_to_workers_map
self.worker_load_balancer.update_workers_list(model_type_to_model_to_workers_map)
self.assertRaises(Exception, lambda: self.worker_load_balancer.choose_worker("Caffe", "model-2"))
caffe_worker = self.worker_load_balancer.choose_worker("Caffe", "model-1")
self.assertEqual(caffe_worker.host, 'localhost')
self.assertTrue(caffe_worker.port in [9090, 9091])
self.assertTrue(caffe_worker.local_worker_id in [0, 1])
self.assertEqual(self.worker_load_balancer.check_if_model_to_workers_map_is_empty(), False)
self.assertEqual(len(self.worker_load_balancer.get_all_workers('Caffe').keys()), 2)
self.assertEqual(set(['localhost-0', 'localhost-1']), set(self.worker_load_balancer.get_all_workers('Caffe').keys()))
self.assertEqual(self.worker_load_balancer.get_all_workers('Caffe').get('localhost-0').global_worker_id, 'localhost-0')
self.assertEqual(self.worker_load_balancer.get_all_workers('Caffe').get('localhost-1').global_worker_id, 'localhost-1')
knn_worker = self.worker_load_balancer.choose_worker("KNN", "model-2")
self.assertEqual(knn_worker.host, 'localhost')
self.assertTrue(knn_worker.port in [9092, 9093])
self.assertTrue(knn_worker.local_worker_id in [2, 3])
self.assertEqual(len(self.worker_load_balancer.get_all_workers('KNN').keys()), 2)
self.assertEqual(set(['localhost-2', 'localhost-3']), set(self.worker_load_balancer.get_all_workers('KNN').keys()))
self.assertEqual(self.worker_load_balancer.get_all_workers('KNN').get('localhost-2').global_worker_id, 'localhost-2')
self.assertEqual(self.worker_load_balancer.get_all_workers('KNN').get('localhost-3').global_worker_id, 'localhost-3')
model_type_to_model_to_workers_map_1 = {}
caffe_model_to_workers_map_1 = {}
knn_model_to_workers_map_1 = {}
model_type_to_model_to_workers_map_1["Caffe"] = caffe_model_to_workers_map_1
model_type_to_model_to_workers_map_1["KNN"] = knn_model_to_workers_map_1
self.worker_load_balancer.update_workers_list(model_type_to_model_to_workers_map_1)
self.assertRaises(Exception, lambda: self.worker_load_balancer.choose_worker("Caffe", "model-1"))
self.assertEqual(self.worker_load_balancer.get_all_workers("Caffe"), {})
self.assertEqual(self.worker_load_balancer.get_all_workers("KNN"), {})
self.assertEqual(self.worker_load_balancer.get_model_to_workers_list("Caffe"), {})
self.assertEqual(self.worker_load_balancer.get_model_to_workers_list("KNN"), {})
self.assertEqual(self.worker_load_balancer.get_worker_info("worker-1", "Caffe"), None)
self.assertEqual(self.worker_load_balancer.get_worker_info("worker-1", "KNN"), None)
model_type_to_model_to_workers_map_2 = None
self.worker_load_balancer.update_workers_list(model_type_to_model_to_workers_map_1)
self.assertRaises(Exception, lambda: self.worker_load_balancer.choose_worker("Caffe", "model-1"))
self.assertEqual(self.worker_load_balancer.get_all_workers("Caffe"), {})
self.assertEqual(self.worker_load_balancer.get_all_workers("KNN"), {})
self.assertEqual(self.worker_load_balancer.get_model_to_workers_list("Caffe"), {})
self.assertEqual(self.worker_load_balancer.get_model_to_workers_list("KNN"), {})
self.assertEqual(self.worker_load_balancer.get_worker_info("worker-1", "Caffe"), None)
self.assertEqual(self.worker_load_balancer.get_worker_info("worker-1", "KNN"), None)
def test_get_all_workers(self):
self.assertEqual(self.worker_load_balancer.get_all_workers("Caffe"), None)
self.assertEqual(self.worker_load_balancer.get_all_workers("KNN"), None)
self.assertRaises(Exception, lambda: self.worker_load_balancer.choose_worker("Caffe", "model-1"))
self.assertEqual(self.worker_load_balancer.get_all_workers("Caffe"), None)
self.assertEqual(self.worker_load_balancer.get_all_workers("KNN"), None)
self.assertEqual(self.worker_load_balancer.get_model_to_workers_list("Caffe"), None)
self.assertEqual(self.worker_load_balancer.get_model_to_workers_list("KNN"), None)
self.assertEqual(self.worker_load_balancer.get_worker_info("worker-1", "Caffe"), None)
self.assertEqual(self.worker_load_balancer.get_worker_info("worker-1", "KNN"), None)
model_type_to_model_to_workers_map = {}
caffe_model_to_workers_map = {}
knn_model_to_workers_map = {}
caffe_worker_01 = {}
caffe_worker_01["host"] = 'localhost'
caffe_worker_01["port"] = 9090
caffe_worker_01["local_worker_id"] = 0
caffe_worker_02 = {}
caffe_worker_02["host"] = 'localhost'
caffe_worker_02["port"] = 9091
caffe_worker_02["local_worker_id"] = 1
caffe_workers = []
caffe_workers.append(caffe_worker_01)
caffe_workers.append(caffe_worker_02)
knn_worker_01 = {}
knn_worker_01["host"] = 'localhost'
knn_worker_01["port"] = 9092
knn_worker_01["local_worker_id"] = 2
knn_worker_02 = {}
knn_worker_02["host"] = 'localhost'
knn_worker_02["port"] = 9093
knn_worker_02["local_worker_id"] = 3
knn_workers = []
knn_workers.append(knn_worker_01)
knn_workers.append(knn_worker_02)
caffe_model_to_workers_map["model-1"] = caffe_workers
knn_model_to_workers_map["model-2"] = knn_workers
model_type_to_model_to_workers_map["Caffe"] = caffe_model_to_workers_map
model_type_to_model_to_workers_map["KNN"] = knn_model_to_workers_map
self.worker_load_balancer.update_workers_list(model_type_to_model_to_workers_map)
self.assertRaises(Exception, lambda: self.worker_load_balancer.choose_worker("Caffe", "model-2"))
self.assertEqual(self.worker_load_balancer.check_if_model_to_workers_map_is_empty(), False)
self.assertEqual(len(self.worker_load_balancer.get_all_workers('Caffe').keys()), 2)
self.assertEqual(set(['localhost-0', 'localhost-1']), set(self.worker_load_balancer.get_all_workers('Caffe').keys()))
self.assertEqual(self.worker_load_balancer.get_all_workers('Caffe').get('localhost-0').global_worker_id, 'localhost-0')
self.assertEqual(self.worker_load_balancer.get_all_workers('Caffe').get('localhost-1').global_worker_id, 'localhost-1')
self.assertEqual(len(self.worker_load_balancer.get_all_workers('KNN').keys()), 2)
self.assertEqual(set(['localhost-2', 'localhost-3']), set(self.worker_load_balancer.get_all_workers('KNN').keys()))
self.assertEqual(self.worker_load_balancer.get_all_workers('KNN').get('localhost-2').global_worker_id, 'localhost-2')
self.assertEqual(self.worker_load_balancer.get_all_workers('KNN').get('localhost-3').global_worker_id, 'localhost-3')
def test_get_model_to_workers_list(self):
self.assertEqual(self.worker_load_balancer.get_model_to_workers_list("Caffe"), None)
self.assertEqual(self.worker_load_balancer.get_model_to_workers_list("KNN"), None)
def test_get_worker_info(self):
self.assertEqual(self.worker_load_balancer.get_worker_info("worker-1", "Caffe"), None)
self.assertEqual(self.worker_load_balancer.get_worker_info("worker-1", "KNN"), None)
def test_check_if_model_to_workers_map_is_empty(self):
self.assertEqual(self.worker_load_balancer.check_if_model_to_workers_map_is_empty(), True)
model_type_to_model_to_workers_map = {}
caffe_model_to_workers_map = {}
knn_model_to_workers_map = {}
caffe_worker_01 = {}
caffe_worker_01["host"] = 'localhost'
caffe_worker_01["port"] = 9090
caffe_worker_01["local_worker_id"] = 0
caffe_worker_02 = {}
caffe_worker_02["host"] = 'localhost'
caffe_worker_02["port"] = 9091
caffe_worker_02["local_worker_id"] = 1
caffe_workers = []
caffe_workers.append(caffe_worker_01)
caffe_workers.append(caffe_worker_02)
knn_worker_01 = {}
knn_worker_01["host"] = 'localhost'
knn_worker_01["port"] = 9090
knn_worker_01["local_worker_id"] = 0
knn_worker_02 = {}
knn_worker_02["host"] = 'localhost'
knn_worker_02["port"] = 9091
knn_worker_02["local_worker_id"] = 1
knn_workers = []
knn_workers.append(knn_worker_01)
knn_workers.append(knn_worker_02)
caffe_model_to_workers_map["model-1"] = caffe_workers
knn_model_to_workers_map["model-2"] = knn_workers
model_type_to_model_to_workers_map["Caffe"] = caffe_model_to_workers_map
model_type_to_model_to_workers_map["KNN"] = knn_model_to_workers_map
self.worker_load_balancer.update_workers_list(model_type_to_model_to_workers_map)
self.assertEqual(self.worker_load_balancer.check_if_model_to_workers_map_is_empty(), False)
caffe_workers_to_be_deleted = [WorkerInfo('localhost', 9090, 0), WorkerInfo('localhost', 9091, 1)]
self.worker_load_balancer.remove_workers("Caffe", "model-1", caffe_workers_to_be_deleted)
def test_choose_worker(self):
self.assertRaises(Exception, lambda: self.worker_load_balancer.choose_worker("Caffe", "model-1"))
model_type_to_model_to_workers_map = {}
caffe_model_to_workers_map = {}
knn_model_to_workers_map = {}
caffe_worker_01 = {}
caffe_worker_01["host"] = 'localhost'
caffe_worker_01["port"] = 9090
caffe_worker_01["local_worker_id"] = 0
caffe_worker_02 = {}
caffe_worker_02["host"] = 'localhost'
caffe_worker_02["port"] = 9091
caffe_worker_02["local_worker_id"] = 1
caffe_workers = []
caffe_workers.append(caffe_worker_01)
caffe_workers.append(caffe_worker_02)
knn_worker_01 = {}
knn_worker_01["host"] = 'localhost'
knn_worker_01["port"] = 9090
knn_worker_01["local_worker_id"] = 0
knn_worker_02 = {}
knn_worker_02["host"] = 'localhost'
knn_worker_02["port"] = 9091
knn_worker_02["local_worker_id"] = 1
knn_workers = []
knn_workers.append(knn_worker_01)
knn_workers.append(knn_worker_02)
caffe_model_to_workers_map["model-1"] = caffe_workers
knn_model_to_workers_map["model-2"] = knn_workers
model_type_to_model_to_workers_map["Caffe"] = caffe_model_to_workers_map
model_type_to_model_to_workers_map["KNN"] = knn_model_to_workers_map
self.worker_load_balancer.update_workers_list(model_type_to_model_to_workers_map)
self.assertRaises(Exception, lambda: self.worker_load_balancer.choose_worker("Caffe", "model-2"))
worker = self.worker_load_balancer.choose_worker("Caffe", "model-1")
self.assertEqual(worker.host, 'localhost')
self.assertTrue(worker.port in [9090, 9091])
self.assertTrue(worker.local_worker_id in [0, 1])
def test_remove_workers(self):
model_type_to_model_to_workers_map = {}
caffe_model_to_workers_map = {}
knn_model_to_workers_map = {}
caffe_worker_01 = {}
caffe_worker_01["host"] = 'localhost'
caffe_worker_01["port"] = 9090
caffe_worker_01["local_worker_id"] = 0
caffe_worker_02 = {}
caffe_worker_02["host"] = 'localhost'
caffe_worker_02["port"] = 9091
caffe_worker_02["local_worker_id"] = 1
caffe_workers = []
caffe_workers.append(caffe_worker_01)
caffe_workers.append(caffe_worker_02)
knn_worker_01 = {}
knn_worker_01["host"] = 'localhost'
knn_worker_01["port"] = 9090
knn_worker_01["local_worker_id"] = 0
knn_worker_02 = {}
knn_worker_02["host"] = 'localhost'
knn_worker_02["port"] = 9091
knn_worker_02["local_worker_id"] = 1
knn_workers = []
knn_workers.append(knn_worker_01)
knn_workers.append(knn_worker_02)
caffe_model_to_workers_map["model-1"] = caffe_workers
knn_model_to_workers_map["model-2"] = knn_workers
model_type_to_model_to_workers_map["Caffe"] = caffe_model_to_workers_map
model_type_to_model_to_workers_map["KNN"] = knn_model_to_workers_map
self.worker_load_balancer.update_workers_list(model_type_to_model_to_workers_map)
caffe_workers_to_be_removed = [WorkerInfo('localhost', 9090, 0), WorkerInfo('localhost', 9091, 1)]
self.worker_load_balancer.remove_workers('Caffe', 'model-1', caffe_workers_to_be_removed)
self.assertRaises(Exception, lambda : self.worker_load_balancer.choose_worker('Caffe', 'model-1'))
def setUp(self):
self.load_balancer = WorkerLoadBalancer()
self.load_balancer.check_if_model_to_workers_map_is_empty = MagicMock(
return_value=False)
start_time = int(time.time())
self.elb_resource = ElbResource(start_time, self.load_balancer)
def _get_kwc():
routing_map = {'KNN': {'model-1':[{'host': 'localhost', 'port': 9090, 'local_worker_id': 0}, {'host': 'localhost', 'port': 9091, 'local_worker_id': 1}, {'host': 'localhost', 'port': 9092, 'local_worker_id': 2}, {'host': 'localhost', 'port': 9093, 'local_worker_id': 3}, {'host': 'localhost', 'port': 9094, 'local_worker_id': 4}, {'host': 'localhost', 'port': 9095, 'local_worker_id': 5}]}}
wlb = WorkerLoadBalancer()
wlb.update_workers_list(routing_map)
kwc = KnnWorkerClient(wlb)
return kwc