diff options
| author |  Noah Eisen <ncteisen@google.com> | 2018-08-21 16:58:59 -0700 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2018-08-21 17:02:37 -0700 |
| commit | 95d718a8a41370f31ccb3b32aaac7fd00b0291e4 (patch) | |
| tree | da7a1a2dbe4e7674a9911afda932c537ce0cabdb /tensorflow/core/distributed_runtime | |
| parent | 3cb13feff06ac68a10353313de50dd54f8043ddf (diff) | |
Allow child class of Server to supply custom ChannelArguments
PiperOrigin-RevId: 209685137
Diffstat (limited to 'tensorflow/core/distributed_runtime')
| -rw-r--r-- | tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc | 2 | ||||
| -rw-r--r-- | tensorflow/core/distributed_runtime/rpc/grpc_server_lib.h | 3 |
2 files changed, 5 insertions, 0 deletions
diff --git a/tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc b/tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc index bcd46a4c06..c4f2247145 100644 --- a/tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc +++ b/tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc @@ -190,6 +190,8 @@ Status GrpcServer::Init( builder.SetMaxMessageSize(std::numeric_limits<int32>::max()); builder.SetOption( std::unique_ptr<::grpc::ServerBuilderOption>(new NoReusePortOption)); + // Allow subclasses to specify more args to pass to the gRPC server. + MaybeMutateBuilder(&builder); master_impl_ = CreateMaster(&master_env_); master_service_ = NewGrpcMasterService(master_impl_.get(), config, &builder); worker_impl_ = diff --git a/tensorflow/core/distributed_runtime/rpc/grpc_server_lib.h b/tensorflow/core/distributed_runtime/rpc/grpc_server_lib.h index 3366246afb..7979e96d3e 100644 --- a/tensorflow/core/distributed_runtime/rpc/grpc_server_lib.h +++ b/tensorflow/core/distributed_runtime/rpc/grpc_server_lib.h @@ -59,6 +59,9 @@ typedef std::function<std::unique_ptr<GrpcWorker>(WorkerEnv*)> class GrpcServer : public ServerInterface { protected: GrpcServer(const ServerDef& server_def, Env* env); + // Allow children classes to override this and provide custom args to the + // server before it is constructed. Default behavior is to do nothing. + virtual void MaybeMutateBuilder(::grpc::ServerBuilder* builder) {} public: static Status Create(const ServerDef& server_def, Env* env, |