experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Use a single-threaded thread pool to handle streaming enqueue requests. This can avoid starving other requests.

Browse Source 
PiperOrigin-RevId: 265765740
This commit is contained in:
Xiao Yu 2019-08-27 14:22:23 -07:00 committed by TensorFlower Gardener
parent c6859f92a0
commit 28e24ad2f1
2 changed files with 30 additions and 24 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
tensorflow/core/distributed_runtime/rpc/eager/grpc_eager_service_impl.cc
View File

@ -27,7 +27,9 @@ namespace eager {
GrpcEagerServiceImpl::GrpcEagerServiceImpl( GrpcEagerServiceImpl::GrpcEagerServiceImpl(
const WorkerEnv* env, ::grpc::ServerBuilder* server_builder) const WorkerEnv* env, ::grpc::ServerBuilder* server_builder)
: env_(env), local_impl_(env) { : env_(env),
local_impl_(env),
enqueue_streaming_thread_(env_->env, "enqueue_streaming_thread", 1) {
server_builder->RegisterService(&service_); server_builder->RegisterService(&service_);
cq_ = server_builder->AddCompletionQueue(); cq_ = server_builder->AddCompletionQueue();
} }

50
tensorflow/core/distributed_runtime/rpc/eager/grpc_eager_service_impl.h
View File

@ -77,38 +77,42 @@ class GrpcEagerServiceImpl : public AsyncServiceInterface {
// call. // call.
// StreamingEnqueueHandler gets the request from the `call` and fills the // StreamingEnqueueHandler gets the request from the `call` and fills the
// response (also found in `call`) by invoking the local EagerServiceImpl. // response (also found in `call`) by invoking the local EagerServiceImpl.
// The local EagerServiceImpl is invoked in this thread instead of using a // The local EagerServiceImpl is invoked in a single-threaded thread pool. We
// thread-pool as is done for all other methods above. We do this to preserve // do this to preserve request order. The local service can parallelize based
// request order. The local service can parallelize based on context_id in // on context_id in request if necessary. Remote contexts are created in async
// request if necessary. Remote contexts are created in async mode by default, // mode by default, so the local service impl just puts the request on eager
// so the local service impl just puts the request on eager executor queue. // executor queue.
void StreamingEnqueueHandler( void StreamingEnqueueHandler(
StreamingCall<EnqueueRequest, EnqueueResponse>* call) { StreamingCall<EnqueueRequest, EnqueueResponse>* call) {
// NOTE(fishx): Use the address of StreamingCall as the stream_id since we enqueue_streaming_thread_.Schedule([this, call]() {
// reuse the same StreamingCall for multiple requests in the same streaming // NOTE(fishx): Use the address of StreamingCall as the stream_id since we
// connection. // reuse the same StreamingCall for multiple requests in the same
Status status = // streaming connection.
local_impl_.Enqueue(&call->request(), call->mutable_response(), Status status = local_impl_.Enqueue(
reinterpret_cast<uint64>(static_cast<void*>(call))); &call->request(), call->mutable_response(),
reinterpret_cast<uint64>(static_cast<void*>(call)));
if (status.ok()) { if (status.ok()) {
VLOG(1) << "local_impl_.Enqueue completed successfully"; VLOG(1) << "local_impl_.Enqueue completed successfully";
call->SendResponse(); call->SendResponse();
} else { } else {
VLOG(1) << "local_impl_.Enqueue failed with " << status.ToString() VLOG(1) << "local_impl_.Enqueue failed with " << status.ToString()
<< " on request " << call->request().DebugString(); << " on request " << call->request().DebugString();
call->Finish(ToGrpcStatus(status)); call->Finish(ToGrpcStatus(status));
} }
// We do not tell gRPC to accept a new StreamingEnqueue request because this // We do not tell gRPC to accept a new StreamingEnqueue request because
// method can be called multiple times for a given streaming call. // this method can be called multiple times for a given streaming call.
// The StreamingCall does this per call instead, after a call has been // The StreamingCall does this per call instead, after a call has been
// opened. // opened.
});
} }
const WorkerEnv* const env_; // Not owned. const WorkerEnv* const env_; // Not owned.
EagerServiceImpl local_impl_; EagerServiceImpl local_impl_;
// A single-threaded thread pool to handle streaming enqueue rpc request.
thread::ThreadPool enqueue_streaming_thread_;
std::unique_ptr<::grpc::Alarm> shutdown_alarm_; std::unique_ptr<::grpc::Alarm> shutdown_alarm_;
std::unique_ptr<::grpc::ServerCompletionQueue> cq_; std::unique_ptr<::grpc::ServerCompletionQueue> cq_;