diff --git a/tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc b/tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc index bcd46a4c06e..c4f2247145c 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::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 3366246afb8..7979e96d3ed 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(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,