add exclusive mode for rpc server (#941)

This commit is contained in:
Lianmin Zheng 2018-03-01 00:54:58 +08:00 коммит произвёл Tianqi Chen
Родитель 589a26515f
Коммит 07cc21f106
2 изменённых файлов: 20 добавлений и 6 удалений

Просмотреть файл

@ -74,10 +74,16 @@ def _recvall(sock, nbytes):
return b"".join(res)
def _listen_loop(sock):
def _listen_loop(sock, exclusive):
"""Lisenting loop"""
last_proc = None
while True:
conn, addr = sock.accept()
if last_proc and last_proc.is_alive() and exclusive:
logging.info("Kill last call")
last_proc.terminate()
logging.info("RPCServer: connection from %s", addr)
magic = struct.unpack("@i", _recvall(conn, 4))[0]
if magic != RPC_MAGIC:
@ -90,9 +96,11 @@ def _listen_loop(sock):
else:
conn.sendall(struct.pack("@i", RPC_MAGIC))
logging.info("Connection from %s", addr)
process = multiprocessing.Process(target=_serve_loop, args=(conn, addr))
process.deamon = True
process.start()
last_proc = process
# close from our side.
conn.close()
@ -158,6 +166,11 @@ class Server(object):
This is recommended to switch on if we want to do local RPC demonstration
for GPU devices to avoid fork safety issues.
exclusive : bool, optional
If this is enabled, the server will kill old connection
when new connection comes. This can make sure the current call
monopolize the hardware resource.
key : str, optional
The key used to identify the server in Proxy connection.
"""
@ -167,6 +180,7 @@ class Server(object):
port_end=9199,
is_proxy=False,
use_popen=False,
exclusive=False,
key=""):
self.host = host
self.port = port
@ -201,7 +215,7 @@ class Server(object):
sock.listen(1)
self.sock = sock
self.proc = multiprocessing.Process(
target=_listen_loop, args=(self.sock,))
target=_listen_loop, args=(self.sock, exclusive))
self.proc.deamon = True
self.proc.start()
else:
@ -210,8 +224,6 @@ class Server(object):
self.proc.deamon = True
self.proc.start()
def terminate(self):
"""Terminate the server process"""
if self.proc:
@ -222,7 +234,6 @@ class Server(object):
self.terminate()
class RPCSession(object):
"""RPC Client session module

Просмотреть файл

@ -21,6 +21,9 @@ def main():
help="Whether to load executor runtime")
parser.add_argument('--load-library', type=str, default="",
help="Additional library to load")
parser.add_argument('--exclusive', action='store_true',
help="If this is enabled, the server will kill old connection"
"when new connection comes")
args = parser.parse_args()
logging.basicConfig(level=logging.INFO)
@ -35,7 +38,7 @@ def main():
libs.append(ctypes.CDLL(file_name, ctypes.RTLD_GLOBAL))
logging.info("Load additional library %s", file_name)
server = rpc.Server(args.host, args.port, args.port_end)
server = rpc.Server(args.host, args.port, args.port_end, exclusive=args.exclusive)
server.libs += libs
server.proc.join()