diff options
| author |  Derek Murray <mrry@google.com> | 2016-03-22 21:37:35 -0800 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2016-03-22 23:09:22 -0700 |
| commit | 4cec6830bc0aaed6cc089c06c437eae1f240758b (patch) | |
| tree | b01b0c4a433273d8c5fbff16ecd5f8d7c511a094 /tensorflow/core/distributed_runtime/rpc/grpc_channel.h | |
| parent | d56b1968633a544a17c63b216262acecf9b5793f (diff) | |
Add extension points to the gRPC Session and Server code.
This enables a subclass of GrpcServer or GrpcSession to override the
server credential and channel creation logic, allowing the use of
different security mechanisms.
Change: 117900179
Diffstat (limited to 'tensorflow/core/distributed_runtime/rpc/grpc_channel.h')
| -rw-r--r-- | tensorflow/core/distributed_runtime/rpc/grpc_channel.h | 7 |
1 files changed, 5 insertions, 2 deletions
diff --git a/tensorflow/core/distributed_runtime/rpc/grpc_channel.h b/tensorflow/core/distributed_runtime/rpc/grpc_channel.h index 2083894323..c09dcac674 100644 --- a/tensorflow/core/distributed_runtime/rpc/grpc_channel.h +++ b/tensorflow/core/distributed_runtime/rpc/grpc_channel.h @@ -73,7 +73,10 @@ class GrpcChannelCache { virtual string TranslateTask(const string& task) = 0; }; -GrpcChannelCache* NewGrpcChannelCache(const GrpcChannelSpec& p); +typedef std::function<SharedGrpcChannelPtr(string)> ChannelCreationFunction; + +GrpcChannelCache* NewGrpcChannelCache(const GrpcChannelSpec& p, + ChannelCreationFunction channel_func); // Below here are internal-only functions. @@ -87,7 +90,7 @@ SharedGrpcChannelPtr NewHostPortGrpcChannel(const string& target); // The caller takes ownership of the returned object. GrpcChannelCache* NewHostPortsGrpcChannelCache( const string& job_id, const std::vector<string>& host_ports, - int tasks_per_replica); + int tasks_per_replica, ChannelCreationFunction channel_func); // Returns a ChannelCache that is the union of a number of other ChannelCaches. GrpcChannelCache* NewMultiGrpcChannelCache( |