[RPC] callback option rpc server starts (#1092)

ee514eac · Tianqi Chen · GitHub · 9082a93b · ee514eac
Commit ee514eac authored Apr 08, 2018 by Tianqi Chen Committed by GitHub Apr 08, 2018
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 0 deletions

src/runtime/rpc/rpc_session.cc
+3 -0

No files found.
--- a/src/runtime/rpc/rpc_session.cc
+++ b/src/runtime/rpc/rpc_session.cc
@@ -783,6 +783,9 @@ void RPCSession::Shutdown() {

 void RPCSession::ServerLoop() {
  std::lock_guard<std::recursive_mutex> lock(mutex_);
+  if (const auto* f = Registry::Get("tvm.contrib.rpc.server.start")) {
+    (*f)();
+  }
  TVMRetValue rv;
  CHECK(HandleUntilReturnEvent(&rv, false, nullptr) == RPCCode::kShutdown);
  LOG(INFO) << "Shutdown...";