Esempio n. 1
0
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
Esempio n. 2
0
 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
Esempio n. 3
0
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
Esempio n. 4
0
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)
Esempio n. 9
0
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