[RPC] HandleCtrlC Register for handling Ctrl+C event. (#1652)

* update

* update

* update

* update

* fix lint

python/dgl/distributed/rpc.py

@@ -893,6 +893,11 @@ def fast_pull(name, id_tensor, part_id, service_id,
                                       F.zerocopy_to_dgl_ndarray(local_data))
     return F.zerocopy_from_dgl_ndarray(res_tensor)
 
+def register_ctrl_c():
+    """HandleCtrlC Register for handling Ctrl+C event.
+    """
+    _CAPI_DGLRPCHandleCtrlC()
+
 ############### Some basic services will be defined here #############
 
 CLIENT_REGISTER = 22451

python/dgl/distributed/rpc_server.py

@@ -35,6 +35,8 @@ def start_server(server_id, ip_config, num_clients, server_state, \
     assert num_clients >= 0, 'num_client (%d) cannot be a negative number.' % num_client
     assert max_queue_size > 0, 'queue_size (%d) cannot be a negative number.' % queue_size
     assert net_type in ('socket'), 'net_type (%s) can only be \'socket\'' % net_type
+    # HandleCtrlC Register for handling Ctrl+C event
+    rpc.register_ctrl_c()
     # Register some basic services
     rpc.register_service(rpc.CLIENT_REGISTER,
                          rpc.ClientRegisterRequest,

src/rpc/rpc.cc

@@ -5,6 +5,11 @@
  */
 #include "./rpc.h"
 
+#include <csignal>
+#if defined(__linux__)
+#include <unistd.h>
+#endif
+
 #include <dgl/runtime/container.h>
 #include <dgl/packed_func_ext.h>
 #include <dgl/array.h>
@@ -289,6 +294,27 @@ DGL_REGISTER_GLOBAL("distributed.rpc._CAPI_DGLRPCMessageGetTensors")
   *rv = ret;
 });
 
+#if defined(__linux__)
+/*!
+ * \brief CtrlCHandler, exits if Ctrl+C is pressed
+ * \param s signal
+ */
+void CtrlCHandler(int s) {
+  LOG(INFO) << "\nUser pressed Ctrl+C, Exiting";
+  exit(1);
+}
+
+DGL_REGISTER_GLOBAL("distributed.rpc._CAPI_DGLRPCHandleCtrlC")
+.set_body([] (DGLArgs args, DGLRetValue* rv) {
+  // Ctrl+C handler
+  struct sigaction sigIntHandler;
+  sigIntHandler.sa_handler = CtrlCHandler;
+  sigemptyset(&sigIntHandler.sa_mask);
+  sigIntHandler.sa_flags = 0;
+  sigaction(SIGINT, &sigIntHandler, nullptr);
+});
+#endif
+
 //////////////////////////// ServerState ////////////////////////////
 
 DGL_REGISTER_GLOBAL("distributed.server_state._CAPI_DGLRPCGetServerState")