rpc_server.py 1.84 KB
Newer Older
1 2 3 4 5
"""Start an RPC server"""
from __future__ import absolute_import

import logging
import argparse
6 7
import os
import ctypes
8
from ..contrib import rpc
9
from .._ffi.libinfo import find_lib_path
10 11 12 13 14 15 16 17

def main():
    """Main funciton"""
    parser = argparse.ArgumentParser()
    parser.add_argument('--host', type=str, default="0.0.0.0",
                        help='the hostname of the server')
    parser.add_argument('--port', type=int, default=9090,
                        help='The port of the PRC')
18
    parser.add_argument('--port-end', type=int, default=9199,
19
                        help='The end search port of the PRC')
20 21
    parser.add_argument('--with-executor', type=bool, default=False,
                        help="Whether to load executor runtime")
22 23
    parser.add_argument('--load-library', type=str, default="",
                        help="Additional library to load")
24 25 26
    parser.add_argument('--exclusive', action='store_true',
                        help="If this is enabled, the server will kill old connection"
                             "when new connection comes")
27
    args = parser.parse_args()
28

29
    logging.basicConfig(level=logging.INFO)
30
    load_library = [lib for lib in args.load_library.split(":") if len(lib) != 0]
31 32 33
    curr_path = os.path.dirname(os.path.abspath(os.path.expanduser(__file__)))
    apps_path = os.path.join(curr_path, "../../../apps/graph_executor/lib/")
    libs = []
34
    if args.with_executor:
35 36 37 38 39
        load_library += ["libtvm_graph_exec.so"]
    for file_name in load_library:
        file_name = find_lib_path(file_name, apps_path)[0]
        libs.append(ctypes.CDLL(file_name, ctypes.RTLD_GLOBAL))
        logging.info("Load additional library %s", file_name)
40

41
    server = rpc.Server(args.host, args.port, args.port_end, exclusive=args.exclusive)
42
    server.libs += libs
43 44 45 46
    server.proc.join()

if __name__ == "__main__":
    main()