[RPC] Enable shutdown hook (#308)

19381b51 · Tianqi Chen · GitHub · 48fd9e46 · 19381b51
Commit 19381b51 authored Aug 10, 2017 by Tianqi Chen Committed by GitHub Aug 10, 2017
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
@@ -749,6 +749,9 @@ void RPCSession::ServerLoop() {
  TVMRetValue rv;
  CHECK(HandleUntilReturnEvent(&rv, false, nullptr) == RPCCode::kShutdown);
  LOG(INFO) << "Shutdown...";
+  if (const auto* f = Registry::Get("tvm.contrib.rpc.server.shutdown")) {
+    (*f)();
+  }
  channel_.reset(nullptr);
 }