Fix redis balance (#108)

Author: wangxicoding

docker/Dockerfile

@@ -37,3 +37,4 @@ RUN mkdir -p /tmp/protoc && cd /tmp/protoc && \
 RUN echo "go env -w GO111MODULE=on && go env -w GOPROXY=https://goproxy.io,direct" >> /root/.bashrc
 ENV GO111MODULE=on
 ENV GOPROXY=https://goproxy.io,direct
+

docker/requirements.txt

@@ -5,3 +5,5 @@ flask==1.1.2
 pathlib2==2.3.5
 protobuf==3.8.0
 kubernetes
+redis
+paddle-serving-client

python/paddle_edl/distill/distill_reader.py

@@ -1,13 +1,9 @@
 # -*- coding: utf-8 -*-
 import logging
 import multiprocessing as mps
-import socket
-import time
-import threading
 import os
-
-from contextlib import closing
-from six.moves import queue
+import threading
+import time
 
 from . import distill_worker
 
@@ -17,22 +13,6 @@
 logger = logging.getLogger(__name__)
 
 
-def is_server_alive(server):
-    # FIXME. only for test, need find a better test method
-    if distill_worker._NOP_PREDICT_TEST:
-        return True
-    alive = True
-    ip, port = server.split(":")
-    with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
-        try:
-            s.settimeout(1.5)
-            s.connect((ip, int(port)))
-            s.shutdown(socket.SHUT_RDWR)
-        except:
-            alive = False
-        return alive
-
-
 class ServiceDiscover(object):
     def get_servers(self):
         pass

python/paddle_edl/distill/redis/balance_server.py

@@ -208,7 +208,6 @@ def __init__(self, ip='127.0.0.1', port=9379, table=None):
         self._table = table
         self._handle_func = {
             'register': self._handle_register,
-            'require_service': self._handle_require_service,
             'heartbeat': self._handle_heartbeat
         }
 
@@ -233,26 +232,20 @@ def _handle_register(self, fd, msg):
         }
         self._enqueue_response(fd, msg)
 
-    def _handle_require_service(self, fd, msg):
-        require_num = int(msg['num'])
-        print('fd={}, require_num={}'.format(fd, require_num))
-
-        # for debug
-        if self._table is None:
-            teacher_list = ['127.0.0.1:0001', '127.0.0.1:0002']
-            num = len(teacher_list)
-        else:
-            servers = self._table.get_servers(fd, require_num)
-            teacher_list = servers
-            num = len(teacher_list)
-
-        msg = {'type': 'response_service', 'servers': teacher_list, 'num': num}
-        self._enqueue_response(fd, msg)
-
     def _handle_heartbeat(self, fd, msg):
-        is_update, servers = self._table.is_servers_update(fd)
-        if is_update:
-            msg = {'type': 'servers_change', 'servers': servers}
+        version = 0
+        try:
+            version = int(msg['version'])
+        except KeyError:
+            # compatible old client
+            pass
+        new_version, servers = self._table.is_servers_update(fd, version)
+        if new_version > version:
+            msg = {
+                'type': 'servers_change',
+                'servers': servers,
+                'version': new_version
+            }
         else:
             msg = {'type': 'heartbeat'}
         self._enqueue_response(fd, msg)
@@ -274,7 +267,47 @@ def server_forever(self):
 if __name__ == '__main__':
     from service_table import ServiceTable
 
-    # service_name = 'TestService'
-    table = ServiceTable('127.0.0.1', 6379)  # connect redis ip:port
-    balance_server = BalanceServer('0.0.0.0', 7001, table)  # listen
+    import argparse
+    parser = argparse.ArgumentParser(
+        description='Discovery server with balance')
+    parser.add_argument(
+        '--server',
+        type=str,
+        default='0.0.0.0:7001',
+        help='endpoint of the server, e.g. 127.0.0.1:8888 [default: %(default)s]'
+    )
+    parser.add_argument(
+        '--worker_num',
+        type=int,
+        default=1,
+        help='worker num of server [default: %(default)s]')
+    parser.add_argument(
+        '--db_endpoints',
+        type=str,
+        default='127.0.0.1:6379',
+        help='database endpoints, e.g. 127.0.0.1:2379,127.0.0.1:2380 [default: %(default)s]'
+    )
+    parser.add_argument(
+        '--db_passwd',
+        type=str,
+        default=None,
+        help='detabase password [default: %(default)s]')
+    parser.add_argument(
+        '--db_type',
+        type=str,
+        default='redis',
+        help='database type, only support redis for now [default: %(default)s]')
+
+    args = parser.parse_args()
+    server = args.server
+    worker_num = args.worker_num
+    db_endpoints = args.db_endpoints.split(',')
+
+    redis_ip_port = db_endpoints[0].split(':')
+    server_ip_port = server.split(':')
+
+    table = ServiceTable(redis_ip_port[0],
+                         int(redis_ip_port[1]))  # connect redis ip:port
+    balance_server = BalanceServer(server_ip_port[0],
+                                   int(server_ip_port[1]), table)  # listen
     balance_server.server_forever()

python/paddle_edl/distill/redis/client.py

@@ -46,6 +46,7 @@ def __init__(self, endpoints, service_name, require_num, token=None):
         self._balance_list = endpoints
         self.teacher_list = []
         self._is_update = False
+        self._version = 0
 
         self._thread = None
         self.client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
@@ -86,22 +87,10 @@ def _register(self, require_num=1):
         servers = msg['servers']
         return servers
 
-    def _require(self, require_num):
-        # require service
-        msg = {'type': 'require_service', 'num': require_num}
-        self._send_msg(msg)
-
-        # get servers
-        msg = self._recv_msg()
-        assert msg['type'] == 'response_service'
-        response_num = msg['num']
-        servers = msg['servers']
-        return servers
-
     def _heartbeat(self):
         while self._need_stop is False:
             time.sleep(2)
-            msg = {'type': 'heartbeat'}
+            msg = {'type': 'heartbeat', 'version': self._version}
             self._send_msg(msg)
             msg = self._recv_msg()
 
@@ -110,11 +99,15 @@ def _heartbeat(self):
                 continue
             elif msg['type'] == 'servers_change':
                 self._is_update = True
+                try:
+                    self._version = int(msg['version'])
+                except KeyError:
+                    # compatible with old balance server
+                    pass
                 self.teacher_list = msg['servers']
-                sys.stderr.write('servers_change: ' + str(msg['servers']) +
-                                 '\n')
-                # Todo
-                pass
+                sys.stderr.write(
+                    '[INFO] service change version={} teachers={}\n'.format(
+                        self._version, str(self.teacher_list)))
 
     def get_teacher_list(self):
         '''

python/paddle_edl/distill/redis/server_register.py

@@ -88,22 +88,49 @@ def start(self, daemon=False):
     import sys
     from redis_store import RedisStore
 
-    ip = '127.0.0.1'
-    port = 5458
-    service_name = 'TestService'
-    service_name = 'DistillService'
-
-    if len(sys.argv) == 2:
-        port = int(sys.argv[1])
-    elif len(sys.argv) >= 3:
-        ip = sys.argv[1]
-        port = int(sys.argv[2])
-        if len(sys.argv) == 4:
-            service_name = sys.argv[3]
-
-    print('register {}:{} service_name={}'.format(ip, port, service_name))
-
-    #store = RedisStore('127.0.0.1', 6379)
-    store = RedisStore('10.255.100.13', 6379)
-    register = ServerRegister(ip, port, service_name, store)
+    import argparse
+    parser = argparse.ArgumentParser(description='Server Register')
+    parser.add_argument(
+        '--db_endpoints',
+        type=str,
+        default='127.0.0.1:6379',
+        help='database endpoints, e.g. 127.0.0.1:6379 [default: %(default)s]')
+    parser.add_argument(
+        '--db_passwd',
+        type=str,
+        default=None,
+        help='detabase password [default: %(default)s]')
+    parser.add_argument(
+        '--db_type',
+        type=str,
+        default='redis',
+        help='database type, only support redis for now [default: %(default)s]')
+    parser.add_argument(
+        '--service_name',
+        type=str,
+        help='service name where the server is located',
+        required=True)
+    parser.add_argument(
+        '--server',
+        type=str,
+        help='endpoint of the server, e.g. 127.0.0.1:8888',
+        required=True)
+    # TODO. service_token
+    parser.add_argument(
+        '--service_token',
+        type=str,
+        default=None,
+        help='service token, which the same can register [default: %(default)s]'
+    )
+
+    args = parser.parse_args()
+    server = args.server
+    db_endpoints = args.db_endpoints.split(',')
+
+    redis_ip_port = db_endpoints[0].split(':')
+    server_ip_port = server.split(':')
+
+    store = RedisStore(redis_ip_port[0], int(redis_ip_port[1]))
+    register = ServerRegister(server_ip_port[0],
+                              int(server_ip_port[1]), args.service_name, store)
     register.start()

python/paddle_edl/distill/redis/service_table.py

@@ -12,12 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
 import threading
 import time
 import sys
-import heapq
 from redis_store import RedisStore
 
+logging.basicConfig(
+    level=logging.DEBUG,
+    format="[%(levelname)s %(asctime)s %(filename)s:%(lineno)d] %(message)s")
+
 
 class ServiceTable(object):
     def __init__(self, ip='127.0.0.1', port=6379, passwd=None,
@@ -43,15 +47,16 @@ def __init__(self, ip='127.0.0.1', port=6379, passwd=None,
         # {fd: set(servers), }
         self._fd_to_servers = {}
         self._server_to_fds = {}
-        self._fd_to_update = {}
+        self._fd_to_version = {}
         self._fd_to_max_num = {}
 
-    def is_servers_update(self, fd):
-        ret = self._fd_to_update[fd]
-        self._fd_to_update[fd] = False
-        if ret is True:
-            return ret, list(self._fd_to_servers[fd])
-        return ret, []
+    def is_servers_update(self, fd, version):
+        new_version = self._fd_to_version[fd]
+
+        if new_version > version:  # is update
+            return new_version, list(self._fd_to_servers[fd])
+        else:
+            return new_version, None
 
     def get_servers(self, fd, num):
         if fd not in self._fd_to_service_name:
@@ -62,13 +67,12 @@ def get_servers(self, fd, num):
         if service_name not in self._service_name_to_servers or \
            self._service_name_to_update[service_name] is True:
             self._refresh_service(service_name)
-            self._fd_to_update[fd] = False
 
         return list(self._fd_to_servers[fd])
 
     def add_service_name(self, fd, service_name, num):
         self._fd_to_servers[fd] = set()
-        self._fd_to_update[fd] = False
+        self._fd_to_version[fd] = 0
         self._fd_to_max_num[fd] = num
         print('fd={}, service_name={}, max_num={}'.format(fd, service_name,
                                                           num))
@@ -103,7 +107,7 @@ def rm_service_name(self, fd):
         del self._fd_to_service_name[fd]
 
         del self._fd_to_max_num[fd]
-        del self._fd_to_update[fd]
+        del self._fd_to_version[fd]
 
         for server in self._fd_to_servers[fd]:
             self._server_to_fds[server].remove(fd)
@@ -163,7 +167,7 @@ def _refresh_service(self, service_name):
         with self._mutex:
             if service_name not in self._service_name_to_fds:
                 for fd in update_fd:
-                    self._fd_to_update[fd] = True
+                    self._fd_to_version[fd] += 1
                 return
             fd_num = len(self._service_name_to_fds[service_name])
 
@@ -172,7 +176,7 @@ def _refresh_service(self, service_name):
         if server_num == 0:
             print('service={} server_num=0'.format(service_name))
             for fd in update_fd:
-                self._fd_to_update[fd] = True
+                self._fd_to_version[fd] += 1
             return
         # assume: fd_num=3, server_num=97
         # assign: {fd0:32, fd1:32, fd2:32}
@@ -193,24 +197,23 @@ def _refresh_service(self, service_name):
                 self._fd_to_servers[fd].remove(server)
                 update_fd.add(fd)
                 print('pop fd={} server={}'.format(fd, server))
-            # add heap
-            # heapq.heappush(server_conn, (len(self._server_to_fds[server]), server))
-
-            # Todo. use ReadWrite Lock
         try:
-            # fd greed connect with server
-            for fd in self._service_name_to_fds[service_name]:
+            fds = self._service_name_to_fds[service_name]
+            for fd in fds:
                 max_connect = min(fd_max_connect, self._fd_to_max_num[fd])
-                print('fd={} max_connect={}'.format(fd, max_connect))
+                logging.info('fd={} max_connect={}'.format(fd, max_connect))
                 if fd not in self._fd_to_servers:
                     self._fd_to_servers[fd] = set()
                 # limit connect of fd
                 while len(self._fd_to_servers[fd]) > max_connect:
                     server = self._fd_to_servers[fd].pop()
                     self._server_to_fds[server].remove(fd)
                     update_fd.add(fd)
-                    print('pop1 fd={} server={}'.format(fd, server))
+                    logging.info('pop1 fd={} server={}'.format(fd, server))
 
+            # fd greed connect with server
+            for fd in fds:
+                max_connect = min(fd_max_connect, self._fd_to_max_num[fd])
                 for server in servers:
                     if len(self._fd_to_servers[fd]) >= max_connect:
                         break
@@ -222,12 +225,12 @@ def _refresh_service(self, service_name):
                     self._fd_to_servers[fd].add(server)
                     self._server_to_fds[server].add(fd)
                     update_fd.add(fd)
-                    print('add fd={} server={}'.format(fd, server))
+                    logging.info('add fd={} server={}'.format(fd, server))
         except Exception, e:
             sys.stderr.write(str(e) + '\n')
 
         for fd in update_fd:
-            self._fd_to_update[fd] = True
+            self._fd_to_version[fd] += 1
 
     def _refresh(self):
         while True: