[RPC] Init server_state during get_server_state() (#1559)

* init server_state

* update

* update

python/dgl/distributed/rpc_server.py

@@ -2,6 +2,7 @@
 
 from . import rpc
 from .constants import MAX_QUEUE_SIZE
+from .server_state import get_server_state
 
 def start_server(server_id, ip_config, num_clients, \
     max_queue_size=MAX_QUEUE_SIZE, net_type='socket'):
@@ -69,8 +70,8 @@ def start_server(server_id, ip_config, num_clients, \
         for client_id, _ in client_namebook.items():
             register_res = rpc.ClientRegisterResponse(client_id)
             rpc.send_response(client_id, register_res)
+    server_state = get_server_state()
     # main service loop
-    server_state = None
     while True:
         req, client_id = rpc.recv_request()
         res = req.process_request(server_state)

src/rpc/rpc.cc

@@ -281,7 +281,9 @@ DGL_REGISTER_GLOBAL("distributed.rpc._CAPI_DGLRPCMessageGetTensors")
 DGL_REGISTER_GLOBAL("distributed.server_state._CAPI_DGLRPCGetServerState")
 .set_body([] (DGLArgs args, DGLRetValue* rv) {
   auto st = RPCContext::ThreadLocal()->server_state;
-  CHECK(st) << "Server state has not been initialized.";
+  if (st.get() == nullptr) {
+    RPCContext::ThreadLocal()->server_state = std::make_shared<ServerState>();
+  }
   *rv = st;
 });