diff --git a/src/runtime/disco/process_session.cc b/src/runtime/disco/process_session.cc index 6474db479e94..6687a64e7f85 100644 --- a/src/runtime/disco/process_session.cc +++ b/src/runtime/disco/process_session.cc @@ -48,11 +48,21 @@ class DiscoPipeMessageQueue : private dmlc::Stream, private DiscoProtocol(num_args); + type_codes = ArenaAlloc(num_args); + TVMArgsSetter setter(values, type_codes); + setter(0, static_cast(DiscoAction::kShutDown)); + setter(1, 0); + } else { + RPCReference::RecvPackedSeq(&values, &type_codes, &num_args, this); + } return TVMArgs(values, type_codes, num_args); } @@ -62,18 +72,38 @@ class DiscoPipeMessageQueue : private dmlc::Stream, private DiscoProtocolRecycleAll(); RPCCode code = RPCCode::kReturn; this->Read(&code); + return false; } size_t Read(void* data, size_t size) final {