Add silent mode to rpc server and rpc tracker (#1268)

python/tvm/contrib/rpc/base.py

@@ -120,7 +120,7 @@ def random_key(prefix, cmap=None):
         return prefix + str(random.random())
 
 
-def connect_with_retry(addr, timeout=60, retry_period=5):
+def connect_with_retry(addr, timeout=60, retry_period=5, silent=False):
     """Connect to a TPC address with retry
 
     This function is only reliable to short period of server restart.
@@ -135,6 +135,9 @@ def connect_with_retry(addr, timeout=60, retry_period=5):
 
     retry_period : float
          Number of seconds before we retry again.
+
+    silent: bool
+        whether run in silent mode
     """
     tstart = time.time()
     while True:
@@ -149,6 +152,7 @@ def connect_with_retry(addr, timeout=60, retry_period=5):
             if period > timeout:
                 raise RuntimeError(
                     "Failed to connect to server %s" % str(addr))
+            if not silent:
                 logging.info("Cannot connect to tracker%s, retry in %g secs...",
                              str(addr), retry_period)
             time.sleep(retry_period)

python/tvm/contrib/rpc/proxy.py

@@ -536,7 +536,7 @@ def websocket_proxy_server(url, key=""):
     def _connect(key):
         conn = yield websocket.websocket_connect(url)
         on_message = create_on_message(conn)
-        temp = _server_env(None)
+        temp = _server_env(None, None)
         # Start connecton
         conn.write_message(struct.pack('<i', base.RPC_MAGIC), binary=True)
         key = "server:" + key

python/tvm/contrib/rpc/tracker.py

@@ -309,7 +309,6 @@ class TrackerServerHandler(object):
 def _tracker_server(listen_sock, stop_key):
     handler = TrackerServerHandler(listen_sock, stop_key)
     handler.run()
-    logging.info("Tracker Stop signal received, terminating...")
 
 
 class Tracker(object):
@@ -327,11 +326,19 @@ class Tracker(object):
 
     port_end : int, optional
         The end TCP port to search
+
+    silent: bool, optional
+        Whether run in silent mode
     """
     def __init__(self,
                  host,
                  port=9190,
-                 port_end=9199):
+                 port_end=9199,
+                 silent=False):
+        self.logger = logging.getLogger("RPCTracker")
+        if silent:
+            self.logger.disabled = True
+
         sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         self.port = None
         self.stop_key = base.random_key("tracker")
@@ -347,7 +354,7 @@ class Tracker(object):
                     raise sock_err
         if not self.port:
             raise ValueError("cannot bind to any port in [%d, %d)" % (port, port_end))
-        logging.info("RPCTracker: bind to %s:%d", host, self.port)
+        self.logger.info("bind to %s:%d", host, self.port)
         sock.listen(1)
         self.proc = multiprocessing.Process(
             target=_tracker_server, args=(sock, self.stop_key))
@@ -373,7 +380,7 @@ class Tracker(object):
                 self._stop_tracker()
                 self.proc.join(1)
             if self.proc.is_alive():
-                logging.info("Terminating Tracker Server...")
+                self.logger.info("Terminating Tracker Server...")
                 self.proc.terminate()
             self.proc = None
 

python/tvm/exec/rpc_server.py

@@ -27,7 +27,8 @@ def main(args):
                         key=args.key,
                         tracker_addr=tracker_addr,
                         load_library=args.load_library,
-                        custom_addr=args.custom_addr)
+                        custom_addr=args.custom_addr,
+                        silent=args.silent)
     server.proc.join()
 
 
@@ -51,6 +52,8 @@ if __name__ == "__main__":
                          and ROCM compilers.")
     parser.add_argument('--custom-addr', type=str,
                         help="Custom IP Address to Report to RPC Tracker")
+    parser.add_argument('--silent', action='store_true',
+                        help="Whether run in silent mode.")
 
     parser.set_defaults(fork=True)
     args = parser.parse_args()
@@ -62,6 +65,7 @@ if __name__ == "__main__":
             )
         multiprocessing.set_start_method('spawn')
     else:
-        logging.info("If you are running ROCM/Metal, \
-        fork with cause compiler internal error. Try to launch with arg ```--no-fork```")
+        if not args.silent:
+            logging.info("If you are running ROCM/Metal, fork will cause "
+                         "compiler internal error. Try to launch with arg ```--no-fork```")
     main(args)

python/tvm/exec/rpc_tracker.py

@@ -11,7 +11,8 @@ from ..contrib.rpc.tracker import Tracker
 
 def main(args):
     """Main funciton"""
-    tracker = Tracker(args.host, port=args.port)
+    tracker = Tracker(args.host, port=args.port, port_end=args.port_end,
+                      silent=args.silent)
     tracker.proc.join()
 
 
@@ -21,10 +22,15 @@ if __name__ == "__main__":
                         help='the hostname of the tracker')
     parser.add_argument('--port', type=int, default=9190,
                         help='The port of the PRC')
+    parser.add_argument('--port-end', type=int, default=9199,
+                        help='The end search port of the PRC')
     parser.add_argument('--no-fork', dest='fork', action='store_false',
                         help="Use spawn mode to avoid fork. This option \
                          is able to avoid potential fork problems with Metal, OpenCL \
                          and ROCM compilers.")
+    parser.add_argument('--silent', action='store_true',
+                        help="Whether run in silent mode.")
+
     parser.set_defaults(fork=True)
     args = parser.parse_args()
     logging.basicConfig(level=logging.INFO)
@@ -35,6 +41,7 @@ if __name__ == "__main__":
             )
         multiprocessing.set_start_method('spawn')
     else:
-        logging.info("If you are running ROCM/Metal, \
-        fork with cause compiler internal error. Try to launch with arg ```--no-fork```")
+        if not args.silent:
+            logging.info("If you are running ROCM/Metal, fork will cause "
+                         "compiler internal error. Try to launch with arg ```--no-fork```")
     main(args)

src/runtime/rpc/rpc_server_env.cc

@@ -20,7 +20,6 @@ TVM_REGISTER_GLOBAL("tvm.contrib.rpc.server.upload").
 set_body([](TVMArgs args, TVMRetValue *rv) {
     std::string file_name = RPCGetPath(args[0]);
     std::string data = args[1];
-    LOG(INFO) << "Upload " << file_name << "... nbytes=" << data.length();
     SaveBinaryToFile(file_name, data);
   });