Home | History | Annotate | Download | only in rpc
      1 /* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
      2 
      3 Licensed under the Apache License, Version 2.0 (the "License");
      4 you may not use this file except in compliance with the License.
      5 You may obtain a copy of the License at
      6 
      7     http://www.apache.org/licenses/LICENSE-2.0
      8 
      9 Unless required by applicable law or agreed to in writing, software
     10 distributed under the License is distributed on an "AS IS" BASIS,
     11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
     12 See the License for the specific language governing permissions and
     13 limitations under the License.
     14 ==============================================================================*/
     15 
     16 #ifndef TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_RPC_GRPC_MASTER_SERVICE_IMPL_H_
     17 #define TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_RPC_GRPC_MASTER_SERVICE_IMPL_H_
     18 
     19 #include "grpc++/impl/codegen/async_stream.h"
     20 #include "grpc++/impl/codegen/async_unary_call.h"
     21 #include "grpc++/impl/codegen/proto_utils.h"
     22 #include "grpc++/impl/codegen/rpc_method.h"
     23 #include "grpc++/impl/codegen/service_type.h"
     24 #include "grpc++/impl/codegen/status.h"
     25 #include "grpc++/impl/codegen/stub_options.h"
     26 #include "grpc++/impl/codegen/sync_stream.h"
     27 
     28 #include "tensorflow/core/distributed_runtime/rpc/grpc_serialization_traits.h"
     29 #include "tensorflow/core/protobuf/master.pb.h"
     30 
     31 // Contains potentially large GraphDef.
     32 TF_GRPC_ALLOW_UNLIMITED_MESSAGE_SIZE(tensorflow::CreateSessionRequest);
     33 // Contains potentially large GraphDef.
     34 TF_GRPC_ALLOW_UNLIMITED_MESSAGE_SIZE(tensorflow::ExtendSessionRequest);
     35 // Contains potentially large TensorProto.
     36 TF_GRPC_ALLOW_UNLIMITED_MESSAGE_SIZE(tensorflow::RunStepRequest);
     37 // Contains potentially large StepStats, TensorProto.
     38 TF_GRPC_ALLOW_UNLIMITED_MESSAGE_SIZE(tensorflow::RunStepResponse);
     39 
     40 namespace grpc {
     41 class CompletionQueue;
     42 class Channel;
     43 class RpcService;
     44 class ServerCompletionQueue;
     45 class ServerContext;
     46 }  // namespace grpc
     47 
     48 namespace tensorflow {
     49 
     50 namespace grpc {
     51 
     52 // Implementation of `tensorflow.MasterService`, based on the
     53 // definition in "//tensorflow/core/protobuf/master_service.proto",
     54 // and the gRPC generated stub and service classes.
     55 // See that file for the definition of methods and messages.
     56 class MasterService final {
     57  public:
     58   class StubInterface {
     59    public:
     60     virtual ~StubInterface() {}
     61     virtual ::grpc::Status CreateSession(::grpc::ClientContext* context,
     62                                          const CreateSessionRequest& request,
     63                                          CreateSessionResponse* response) = 0;
     64     virtual ::grpc::Status ExtendSession(::grpc::ClientContext* context,
     65                                          const ExtendSessionRequest& request,
     66                                          ExtendSessionResponse* response) = 0;
     67     virtual ::grpc::Status PartialRunSetup(
     68         ::grpc::ClientContext* context, const PartialRunSetupRequest& request,
     69         PartialRunSetupResponse* response) = 0;
     70     virtual ::grpc::Status RunStep(::grpc::ClientContext* context,
     71                                    const RunStepRequest& request,
     72                                    RunStepResponse* response) = 0;
     73     virtual ::grpc::Status CloseSession(::grpc::ClientContext* context,
     74                                         const CloseSessionRequest& request,
     75                                         CloseSessionResponse* response) = 0;
     76     virtual ::grpc::Status ListDevices(::grpc::ClientContext* context,
     77                                        const ListDevicesRequest& request,
     78                                        ListDevicesResponse* response) = 0;
     79     virtual ::grpc::Status Reset(::grpc::ClientContext* context,
     80                                  const ResetRequest& request,
     81                                  ResetResponse* response) = 0;
     82   };
     83   class Stub final : public StubInterface {
     84    public:
     85     Stub(const std::shared_ptr< ::grpc::ChannelInterface>& channel);
     86     ::grpc::Status CreateSession(::grpc::ClientContext* context,
     87                                  const CreateSessionRequest& request,
     88                                  CreateSessionResponse* response) override;
     89     ::grpc::Status ExtendSession(::grpc::ClientContext* context,
     90                                  const ExtendSessionRequest& request,
     91                                  ExtendSessionResponse* response) override;
     92     ::grpc::Status PartialRunSetup(::grpc::ClientContext* context,
     93                                    const PartialRunSetupRequest& request,
     94                                    PartialRunSetupResponse* response) override;
     95     ::grpc::Status RunStep(::grpc::ClientContext* context,
     96                            const RunStepRequest& request,
     97                            RunStepResponse* response) override;
     98     ::grpc::Status CloseSession(::grpc::ClientContext* context,
     99                                 const CloseSessionRequest& request,
    100                                 CloseSessionResponse* response) override;
    101     ::grpc::Status ListDevices(::grpc::ClientContext* context,
    102                                const ListDevicesRequest& request,
    103                                ListDevicesResponse* response) override;
    104     ::grpc::Status Reset(::grpc::ClientContext* context,
    105                          const ResetRequest& request,
    106                          ResetResponse* response) override;
    107 
    108    private:
    109     std::shared_ptr< ::grpc::ChannelInterface> channel_;
    110     const ::grpc::internal::RpcMethod rpcmethod_CreateSession_;
    111     const ::grpc::internal::RpcMethod rpcmethod_ExtendSession_;
    112     const ::grpc::internal::RpcMethod rpcmethod_PartialRunSetup_;
    113     const ::grpc::internal::RpcMethod rpcmethod_RunStep_;
    114     const ::grpc::internal::RpcMethod rpcmethod_CloseSession_;
    115     const ::grpc::internal::RpcMethod rpcmethod_ListDevices_;
    116     const ::grpc::internal::RpcMethod rpcmethod_Reset_;
    117   };
    118   static std::unique_ptr<Stub> NewStub(
    119       const std::shared_ptr< ::grpc::ChannelInterface>& channel,
    120       const ::grpc::StubOptions& options = ::grpc::StubOptions());
    121 
    122   class AsyncService : public ::grpc::Service {
    123    public:
    124     AsyncService();
    125     virtual ~AsyncService();
    126     void RequestCreateSession(
    127         ::grpc::ServerContext* context, CreateSessionRequest* request,
    128         ::grpc::ServerAsyncResponseWriter<CreateSessionResponse>* response,
    129         ::grpc::CompletionQueue* new_call_cq,
    130         ::grpc::ServerCompletionQueue* notification_cq, void* tag) {
    131       ::grpc::Service::RequestAsyncUnary(0, context, request, response,
    132                                          new_call_cq, notification_cq, tag);
    133     }
    134     void RequestExtendSession(
    135         ::grpc::ServerContext* context, ExtendSessionRequest* request,
    136         ::grpc::ServerAsyncResponseWriter<ExtendSessionResponse>* response,
    137         ::grpc::CompletionQueue* new_call_cq,
    138         ::grpc::ServerCompletionQueue* notification_cq, void* tag) {
    139       ::grpc::Service::RequestAsyncUnary(1, context, request, response,
    140                                          new_call_cq, notification_cq, tag);
    141     }
    142     void RequestPartialRunSetup(
    143         ::grpc::ServerContext* context, PartialRunSetupRequest* request,
    144         ::grpc::ServerAsyncResponseWriter<PartialRunSetupResponse>* response,
    145         ::grpc::CompletionQueue* new_call_cq,
    146         ::grpc::ServerCompletionQueue* notification_cq, void* tag) {
    147       ::grpc::Service::RequestAsyncUnary(2, context, request, response,
    148                                          new_call_cq, notification_cq, tag);
    149     }
    150     void RequestRunStep(
    151         ::grpc::ServerContext* context, RunStepRequest* request,
    152         ::grpc::ServerAsyncResponseWriter<RunStepResponse>* response,
    153         ::grpc::CompletionQueue* new_call_cq,
    154         ::grpc::ServerCompletionQueue* notification_cq, void* tag) {
    155       ::grpc::Service::RequestAsyncUnary(3, context, request, response,
    156                                          new_call_cq, notification_cq, tag);
    157     }
    158     void RequestCloseSession(
    159         ::grpc::ServerContext* context, CloseSessionRequest* request,
    160         ::grpc::ServerAsyncResponseWriter<CloseSessionResponse>* response,
    161         ::grpc::CompletionQueue* new_call_cq,
    162         ::grpc::ServerCompletionQueue* notification_cq, void* tag) {
    163       ::grpc::Service::RequestAsyncUnary(4, context, request, response,
    164                                          new_call_cq, notification_cq, tag);
    165     }
    166     void RequestListDevices(
    167         ::grpc::ServerContext* context, ListDevicesRequest* request,
    168         ::grpc::ServerAsyncResponseWriter<ListDevicesResponse>* response,
    169         ::grpc::CompletionQueue* new_call_cq,
    170         ::grpc::ServerCompletionQueue* notification_cq, void* tag) {
    171       ::grpc::Service::RequestAsyncUnary(5, context, request, response,
    172                                          new_call_cq, notification_cq, tag);
    173     }
    174     void RequestReset(
    175         ::grpc::ServerContext* context, ResetRequest* request,
    176         ::grpc::ServerAsyncResponseWriter<ResetResponse>* response,
    177         ::grpc::CompletionQueue* new_call_cq,
    178         ::grpc::ServerCompletionQueue* notification_cq, void* tag) {
    179       ::grpc::Service::RequestAsyncUnary(6, context, request, response,
    180                                          new_call_cq, notification_cq, tag);
    181     }
    182   };
    183 };
    184 
    185 }  // namespace grpc
    186 
    187 }  // namespace tensorflow
    188 
    189 #endif  // TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_RPC_GRPC_MASTER_SERVICE_IMPL_H_
    190