Better command line interaction when ModelServerManager enabled (#1409)

Co-authored-by: Lin Ma <[email protected]>

Author: rickyyx

script/model/model_server.py

@@ -25,6 +25,7 @@
 from typing import Dict, Optional, Tuple, List, Any
 import json
 import logging
+import os
 import pprint
 import pickle
 from pathlib import Path
@@ -134,7 +135,10 @@ def __init__(self, end_point: str) -> ModelServer:
         logging.debug(
             f"Python model trying to connect to manager at {end_point}")
         self.socket.connect(f"ipc://{end_point}")
-        logging.debug(f"Python model connected at {end_point}")
+        logging.info(f"Python model connected at {end_point}")
+
+        # If the ModelServer is closing
+        self._closing = False
 
         # Register the exit callback
         atexit.register(self.cleanup_zmq)
@@ -331,6 +335,19 @@ def _infer(self, data: Dict) -> Tuple[List, bool, str]:
 
         return y_pred.tolist(), True, ""
 
+    def _recv(self) -> str:
+        """
+        Receive from the ZMQ socket. This is a blocking call.
+
+        :return: Message paylod
+        """
+        identity = self.socket.recv()
+        _delim = self.socket.recv()
+        payload = self.socket.recv()
+        logging.debug(f"Python recv: {str(identity)}, {str(payload)}")
+
+        return payload.decode("ascii")
+
     def _execute_cmd(self, cmd: Command, data: Dict) -> Tuple[Dict, bool]:
         """
         Execute a command from the ModelServerManager
@@ -382,12 +399,20 @@ def run_loop(self):
         """
 
         while(1):
-            identity = self.socket.recv()
-            _delim = self.socket.recv()
-            payload = self.socket.recv()
-            logging.debug(f"Python recv: {str(identity)}, {str(payload)}")
+            try:
+                payload = self._recv()
+            except UnicodeError as e:
+                logging.warning(f"Failed to decode : {e.reason}")
+                continue
+            except KeyboardInterrupt:
+                if self._closing:
+                    logging.warning("Forced shutting down now.")
+                    os._exit(-1)
+                else:
+                    logging.info("Received KeyboardInterrupt. Ctrl+C again to force shutting down.")
+                    self._closing = True
+                    continue
 
-            payload = payload.decode("ascii")
             send_id, recv_id, msg = self._parse_msg(payload)
             if msg is None:
                 continue

src/self_driving/model_server/model_server_manager.cpp

@@ -69,6 +69,7 @@ ModelServerManager::ModelServerManager(const std::string &model_bin,
       case Callback::NOOP:
         break;
       case Callback::CONNECTED:
+        MODEL_SERVER_LOG_INFO("[PID={}] ModelServer connected", ::getpid());
         connected_ = true;
         break;
       default: