Rewrite master_service in terms of more up-to-date gRPC APIs

PiperOrigin-RevId: 202544091

7 files changed, 6 insertions, 404 deletions

diff --git a/tensorflow/core/distributed_runtime/BUILD b/tensorflow/core/distributed_runtime/BUILD
index c6db2aec06..8247651c24 100644
--- a/tensorflow/core/distributed_runtime/BUILD
+++ b/tensorflow/core/distributed_runtime/BUILD
@@ -635,12 +635,12 @@ tf_cuda_cc_test(
         "//tensorflow/core:lib",
         "//tensorflow/core:lib_internal",
         "//tensorflow/core:master_proto_cc",
+        "//tensorflow/core:master_service_proto_cc",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core:test",
         "//tensorflow/core:test_main",
         "//tensorflow/core:testlib",
         "//tensorflow/core/distributed_runtime/rpc:grpc_channel",
-        "//tensorflow/core/distributed_runtime/rpc:grpc_master_service_impl",
         "//tensorflow/core/distributed_runtime/rpc:grpc_testlib",
         "//tensorflow/core/distributed_runtime/rpc:grpc_util",
         "//tensorflow/core/distributed_runtime/rpc:grpc_worker_cache",
diff --git a/tensorflow/core/distributed_runtime/master_test.cc b/tensorflow/core/distributed_runtime/master_test.cc
index 62b18a45b1..09e96cbd40 100644
--- a/tensorflow/core/distributed_runtime/master_test.cc
+++ b/tensorflow/core/distributed_runtime/master_test.cc
@@ -21,7 +21,6 @@ limitations under the License.
 #include "grpcpp/grpcpp.h"
 
 #include "tensorflow/core/distributed_runtime/rpc/grpc_channel.h"
-#include "tensorflow/core/distributed_runtime/rpc/grpc_master_service_impl.h"
 #include "tensorflow/core/distributed_runtime/rpc/grpc_testlib.h"
 #include "tensorflow/core/distributed_runtime/rpc/grpc_util.h"
 #include "tensorflow/core/framework/allocator.h"
@@ -38,6 +37,7 @@ limitations under the License.
 #include "tensorflow/core/platform/test.h"
 #include "tensorflow/core/platform/types.h"
 #include "tensorflow/core/protobuf/master.pb.h"
+#include "tensorflow/core/protobuf/master_service.grpc.pb.h"
 
 namespace tensorflow {
 
diff --git a/tensorflow/core/distributed_runtime/rpc/BUILD b/tensorflow/core/distributed_runtime/rpc/BUILD
index 7b19427e4b..382ea336ca 100644
--- a/tensorflow/core/distributed_runtime/rpc/BUILD
+++ b/tensorflow/core/distributed_runtime/rpc/BUILD
@@ -201,11 +201,11 @@ cc_library(
     srcs = ["grpc_remote_master.cc"],
     hdrs = ["grpc_remote_master.h"],
     deps = [
-        ":grpc_master_service_impl",
         ":grpc_util",
         "//tensorflow/core:lib",
         "//tensorflow/core:lib_internal",
         "//tensorflow/core:master_proto_cc",
+        "//tensorflow/core:master_service_proto_cc",
         "//tensorflow/core/distributed_runtime:call_options",
         "//tensorflow/core/distributed_runtime:master_interface",
     ],
@@ -219,11 +219,11 @@ cc_library(
     deps = [
         ":async_service_interface",
         ":grpc_call",
-        ":grpc_master_service_impl",
         ":grpc_util",
         "//tensorflow/core:lib",
         "//tensorflow/core:lib_internal",
         "//tensorflow/core:master_proto_cc",
+        "//tensorflow/core:master_service_proto_cc",
         "//tensorflow/core/distributed_runtime:master",
         "@grpc//:grpc++",
     ],
@@ -231,16 +231,6 @@ cc_library(
 )
 
 cc_library(
-    name = "grpc_master_service_impl",
-    srcs = ["grpc_master_service_impl.cc"],
-    hdrs = ["grpc_master_service_impl.h"],
-    deps = [
-        "//tensorflow/core:master_proto_cc",
-        "@grpc//:grpc++",
-    ],
-)
-
-cc_library(
     name = "rpc_rendezvous_mgr",
     srcs = ["rpc_rendezvous_mgr.cc"],
     hdrs = ["rpc_rendezvous_mgr.h"],
diff --git a/tensorflow/core/distributed_runtime/rpc/grpc_master_service.cc b/tensorflow/core/distributed_runtime/rpc/grpc_master_service.cc
index 127dea2882..2c2c1d484a 100644
--- a/tensorflow/core/distributed_runtime/rpc/grpc_master_service.cc
+++ b/tensorflow/core/distributed_runtime/rpc/grpc_master_service.cc
@@ -36,12 +36,12 @@ limitations under the License.
 #include "tensorflow/core/distributed_runtime/master.h"
 #include "tensorflow/core/distributed_runtime/rpc/async_service_interface.h"
 #include "tensorflow/core/distributed_runtime/rpc/grpc_call.h"
-#include "tensorflow/core/distributed_runtime/rpc/grpc_master_service_impl.h"
 #include "tensorflow/core/distributed_runtime/rpc/grpc_util.h"
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/macros.h"
 #include "tensorflow/core/platform/tracing.h"
 #include "tensorflow/core/protobuf/master.pb.h"
+#include "tensorflow/core/protobuf/master_service.grpc.pb.h"
 
 namespace tensorflow {
 
diff --git a/tensorflow/core/distributed_runtime/rpc/grpc_master_service_impl.cc b/tensorflow/core/distributed_runtime/rpc/grpc_master_service_impl.cc
deleted file mode 100644
index 770a0fcf14..0000000000
--- a/tensorflow/core/distributed_runtime/rpc/grpc_master_service_impl.cc
+++ /dev/null
@@ -1,164 +0,0 @@
-/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License.
-==============================================================================*/
-
-#include "tensorflow/core/distributed_runtime/rpc/grpc_master_service_impl.h"
-
-#include "grpcpp/impl/codegen/async_stream.h"
-#include "grpcpp/impl/codegen/async_unary_call.h"
-#include "grpcpp/impl/codegen/channel_interface.h"
-#include "grpcpp/impl/codegen/client_unary_call.h"
-#include "grpcpp/impl/codegen/method_handler_impl.h"
-#include "grpcpp/impl/codegen/rpc_service_method.h"
-#include "grpcpp/impl/codegen/service_type.h"
-#include "grpcpp/impl/codegen/sync_stream.h"
-
-namespace tensorflow {
-
-namespace grpc {
-
-static const char* grpcMasterService_method_names[] = {
-    "/tensorflow.MasterService/CreateSession",
-    "/tensorflow.MasterService/ExtendSession",
-    "/tensorflow.MasterService/PartialRunSetup",
-    "/tensorflow.MasterService/RunStep",
-    "/tensorflow.MasterService/CloseSession",
-    "/tensorflow.MasterService/ListDevices",
-    "/tensorflow.MasterService/Reset",
-    "/tensorflow.MasterService/MakeCallable",
-    "/tensorflow.MasterService/RunCallable",
-    "/tensorflow.MasterService/ReleaseCallable",
-};
-
-std::unique_ptr<MasterService::Stub> MasterService::NewStub(
-    const std::shared_ptr< ::grpc::ChannelInterface>& channel,
-    const ::grpc::StubOptions& options) {
-  std::unique_ptr<MasterService::Stub> stub(new MasterService::Stub(channel));
-  return stub;
-}
-
-MasterService::Stub::Stub(
-    const std::shared_ptr< ::grpc::ChannelInterface>& channel)
-    : channel_(channel),
-      rpcmethod_CreateSession_(grpcMasterService_method_names[0],
-                               ::grpc::internal::RpcMethod::NORMAL_RPC,
-                               channel),
-      rpcmethod_ExtendSession_(grpcMasterService_method_names[1],
-                               ::grpc::internal::RpcMethod::NORMAL_RPC,
-                               channel),
-      rpcmethod_PartialRunSetup_(grpcMasterService_method_names[2],
-                                 ::grpc::internal::RpcMethod::NORMAL_RPC,
-                                 channel),
-      rpcmethod_RunStep_(grpcMasterService_method_names[3],
-                         ::grpc::internal::RpcMethod::NORMAL_RPC, channel),
-      rpcmethod_CloseSession_(grpcMasterService_method_names[4],
-                              ::grpc::internal::RpcMethod::NORMAL_RPC, channel),
-      rpcmethod_ListDevices_(grpcMasterService_method_names[5],
-                             ::grpc::internal::RpcMethod::NORMAL_RPC, channel),
-      rpcmethod_Reset_(grpcMasterService_method_names[6],
-                       ::grpc::internal::RpcMethod::NORMAL_RPC, channel),
-      rpcmethod_MakeCallable_(grpcMasterService_method_names[7],
-                              ::grpc::internal::RpcMethod::NORMAL_RPC, channel),
-      rpcmethod_RunCallable_(grpcMasterService_method_names[8],
-                             ::grpc::internal::RpcMethod::NORMAL_RPC, channel),
-      rpcmethod_ReleaseCallable_(grpcMasterService_method_names[9],
-                                 ::grpc::internal::RpcMethod::NORMAL_RPC,
-                                 channel) {}
-
-::grpc::Status MasterService::Stub::CreateSession(
-    ::grpc::ClientContext* context, const CreateSessionRequest& request,
-    CreateSessionResponse* response) {
-  return ::grpc::internal::BlockingUnaryCall(
-      channel_.get(), rpcmethod_CreateSession_, context, request, response);
-}
-
-::grpc::Status MasterService::Stub::ExtendSession(
-    ::grpc::ClientContext* context, const ExtendSessionRequest& request,
-    ExtendSessionResponse* response) {
-  return ::grpc::internal::BlockingUnaryCall(
-      channel_.get(), rpcmethod_ExtendSession_, context, request, response);
-}
-
-::grpc::Status MasterService::Stub::PartialRunSetup(
-    ::grpc::ClientContext* context, const PartialRunSetupRequest& request,
-    PartialRunSetupResponse* response) {
-  return ::grpc::internal::BlockingUnaryCall(
-      channel_.get(), rpcmethod_PartialRunSetup_, context, request, response);
-}
-
-::grpc::Status MasterService::Stub::RunStep(::grpc::ClientContext* context,
-                                            const RunStepRequest& request,
-                                            RunStepResponse* response) {
-  return ::grpc::internal::BlockingUnaryCall(channel_.get(), rpcmethod_RunStep_,
-                                             context, request, response);
-}
-
-::grpc::Status MasterService::Stub::CloseSession(
-    ::grpc::ClientContext* context, const CloseSessionRequest& request,
-    CloseSessionResponse* response) {
-  return ::grpc::internal::BlockingUnaryCall(
-      channel_.get(), rpcmethod_CloseSession_, context, request, response);
-}
-
-::grpc::Status MasterService::Stub::ListDevices(
-    ::grpc::ClientContext* context, const ListDevicesRequest& request,
-    ListDevicesResponse* response) {
-  return ::grpc::internal::BlockingUnaryCall(
-      channel_.get(), rpcmethod_ListDevices_, context, request, response);
-}
-
-::grpc::Status MasterService::Stub::Reset(::grpc::ClientContext* context,
-                                          const ResetRequest& request,
-                                          ResetResponse* response) {
-  return ::grpc::internal::BlockingUnaryCall(channel_.get(), rpcmethod_Reset_,
-                                             context, request, response);
-}
-
-::grpc::Status MasterService::Stub::MakeCallable(
-    ::grpc::ClientContext* context, const MakeCallableRequest& request,
-    MakeCallableResponse* response) {
-  return ::grpc::internal::BlockingUnaryCall(
-      channel_.get(), rpcmethod_MakeCallable_, context, request, response);
-}
-
-::grpc::Status MasterService::Stub::RunCallable(
-    ::grpc::ClientContext* context, const RunCallableRequest& request,
-    RunCallableResponse* response) {
-  return ::grpc::internal::BlockingUnaryCall(
-      channel_.get(), rpcmethod_RunCallable_, context, request, response);
-}
-
-::grpc::Status MasterService::Stub::ReleaseCallable(
-    ::grpc::ClientContext* context, const ReleaseCallableRequest& request,
-    ReleaseCallableResponse* response) {
-  return ::grpc::internal::BlockingUnaryCall(
-      channel_.get(), rpcmethod_ReleaseCallable_, context, request, response);
-}
-
-MasterService::AsyncService::AsyncService() {
-  int method_len = sizeof(grpcMasterService_method_names) / 
-                    sizeof(grpcMasterService_method_names[0]);
-  for (int i = 0; i < method_len; ++i) {
-    AddMethod(new ::grpc::internal::RpcServiceMethod(
-        grpcMasterService_method_names[i],
-        ::grpc::internal::RpcMethod::NORMAL_RPC, nullptr));
-    ::grpc::Service::MarkMethodAsync(i);
-  }
-}
-
-MasterService::AsyncService::~AsyncService() {}
-
-}  // namespace grpc
-
-}  // namespace tensorflow
diff --git a/tensorflow/core/distributed_runtime/rpc/grpc_master_service_impl.h b/tensorflow/core/distributed_runtime/rpc/grpc_master_service_impl.h
deleted file mode 100644
index 751f2633e7..0000000000
--- a/tensorflow/core/distributed_runtime/rpc/grpc_master_service_impl.h
+++ /dev/null
@@ -1,224 +0,0 @@
-/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License.
-==============================================================================*/
-
-#ifndef TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_RPC_GRPC_MASTER_SERVICE_IMPL_H_
-#define TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_RPC_GRPC_MASTER_SERVICE_IMPL_H_
-
-#include "grpcpp/impl/codegen/async_stream.h"
-#include "grpcpp/impl/codegen/async_unary_call.h"
-#include "grpcpp/impl/codegen/proto_utils.h"
-#include "grpcpp/impl/codegen/rpc_method.h"
-#include "grpcpp/impl/codegen/service_type.h"
-#include "grpcpp/impl/codegen/status.h"
-#include "grpcpp/impl/codegen/stub_options.h"
-#include "grpcpp/impl/codegen/sync_stream.h"
-
-#include "tensorflow/core/protobuf/master.pb.h"
-
-namespace grpc {
-class CompletionQueue;
-class Channel;
-class RpcService;
-class ServerCompletionQueue;
-class ServerContext;
-}  // namespace grpc
-
-namespace tensorflow {
-
-namespace grpc {
-
-// Implementation of `tensorflow.MasterService`, based on the
-// definition in "//tensorflow/core/protobuf/master_service.proto",
-// and the gRPC generated stub and service classes.
-// See that file for the definition of methods and messages.
-class MasterService final {
- public:
-  class StubInterface {
-   public:
-    virtual ~StubInterface() {}
-    virtual ::grpc::Status CreateSession(::grpc::ClientContext* context,
-                                         const CreateSessionRequest& request,
-                                         CreateSessionResponse* response) = 0;
-    virtual ::grpc::Status ExtendSession(::grpc::ClientContext* context,
-                                         const ExtendSessionRequest& request,
-                                         ExtendSessionResponse* response) = 0;
-    virtual ::grpc::Status PartialRunSetup(
-        ::grpc::ClientContext* context, const PartialRunSetupRequest& request,
-        PartialRunSetupResponse* response) = 0;
-    virtual ::grpc::Status RunStep(::grpc::ClientContext* context,
-                                   const RunStepRequest& request,
-                                   RunStepResponse* response) = 0;
-    virtual ::grpc::Status CloseSession(::grpc::ClientContext* context,
-                                        const CloseSessionRequest& request,
-                                        CloseSessionResponse* response) = 0;
-    virtual ::grpc::Status ListDevices(::grpc::ClientContext* context,
-                                       const ListDevicesRequest& request,
-                                       ListDevicesResponse* response) = 0;
-    virtual ::grpc::Status Reset(::grpc::ClientContext* context,
-                                 const ResetRequest& request,
-                                 ResetResponse* response) = 0;
-    virtual ::grpc::Status MakeCallable(::grpc::ClientContext* context,
-                                        const MakeCallableRequest& request,
-                                        MakeCallableResponse* response) = 0;
-    virtual ::grpc::Status RunCallable(::grpc::ClientContext* context,
-                                       const RunCallableRequest& request,
-                                       RunCallableResponse* response) = 0;
-    virtual ::grpc::Status ReleaseCallable(
-        ::grpc::ClientContext* context, const ReleaseCallableRequest& request,
-        ReleaseCallableResponse* response) = 0;
-  };
-  class Stub final : public StubInterface {
-   public:
-    Stub(const std::shared_ptr< ::grpc::ChannelInterface>& channel);
-    ::grpc::Status CreateSession(::grpc::ClientContext* context,
-                                 const CreateSessionRequest& request,
-                                 CreateSessionResponse* response) override;
-    ::grpc::Status ExtendSession(::grpc::ClientContext* context,
-                                 const ExtendSessionRequest& request,
-                                 ExtendSessionResponse* response) override;
-    ::grpc::Status PartialRunSetup(::grpc::ClientContext* context,
-                                   const PartialRunSetupRequest& request,
-                                   PartialRunSetupResponse* response) override;
-    ::grpc::Status RunStep(::grpc::ClientContext* context,
-                           const RunStepRequest& request,
-                           RunStepResponse* response) override;
-    ::grpc::Status CloseSession(::grpc::ClientContext* context,
-                                const CloseSessionRequest& request,
-                                CloseSessionResponse* response) override;
-    ::grpc::Status ListDevices(::grpc::ClientContext* context,
-                               const ListDevicesRequest& request,
-                               ListDevicesResponse* response) override;
-    ::grpc::Status Reset(::grpc::ClientContext* context,
-                         const ResetRequest& request,
-                         ResetResponse* response) override;
-    ::grpc::Status MakeCallable(::grpc::ClientContext* context,
-                                const MakeCallableRequest& request,
-                                MakeCallableResponse* response) override;
-    ::grpc::Status RunCallable(::grpc::ClientContext* context,
-                               const RunCallableRequest& request,
-                               RunCallableResponse* response) override;
-    ::grpc::Status ReleaseCallable(::grpc::ClientContext* context,
-                                   const ReleaseCallableRequest& request,
-                                   ReleaseCallableResponse* response) override;
-
-   private:
-    std::shared_ptr< ::grpc::ChannelInterface> channel_;
-    const ::grpc::internal::RpcMethod rpcmethod_CreateSession_;
-    const ::grpc::internal::RpcMethod rpcmethod_ExtendSession_;
-    const ::grpc::internal::RpcMethod rpcmethod_PartialRunSetup_;
-    const ::grpc::internal::RpcMethod rpcmethod_RunStep_;
-    const ::grpc::internal::RpcMethod rpcmethod_CloseSession_;
-    const ::grpc::internal::RpcMethod rpcmethod_ListDevices_;
-    const ::grpc::internal::RpcMethod rpcmethod_Reset_;
-    const ::grpc::internal::RpcMethod rpcmethod_MakeCallable_;
-    const ::grpc::internal::RpcMethod rpcmethod_RunCallable_;
-    const ::grpc::internal::RpcMethod rpcmethod_ReleaseCallable_;
-  };
-  static std::unique_ptr<Stub> NewStub(
-      const std::shared_ptr< ::grpc::ChannelInterface>& channel,
-      const ::grpc::StubOptions& options = ::grpc::StubOptions());
-
-  class AsyncService : public ::grpc::Service {
-   public:
-    AsyncService();
-    virtual ~AsyncService();
-    void RequestCreateSession(
-        ::grpc::ServerContext* context, CreateSessionRequest* request,
-        ::grpc::ServerAsyncResponseWriter<CreateSessionResponse>* response,
-        ::grpc::CompletionQueue* new_call_cq,
-        ::grpc::ServerCompletionQueue* notification_cq, void* tag) {
-      ::grpc::Service::RequestAsyncUnary(0, context, request, response,
-                                         new_call_cq, notification_cq, tag);
-    }
-    void RequestExtendSession(
-        ::grpc::ServerContext* context, ExtendSessionRequest* request,
-        ::grpc::ServerAsyncResponseWriter<ExtendSessionResponse>* response,
-        ::grpc::CompletionQueue* new_call_cq,
-        ::grpc::ServerCompletionQueue* notification_cq, void* tag) {
-      ::grpc::Service::RequestAsyncUnary(1, context, request, response,
-                                         new_call_cq, notification_cq, tag);
-    }
-    void RequestPartialRunSetup(
-        ::grpc::ServerContext* context, PartialRunSetupRequest* request,
-        ::grpc::ServerAsyncResponseWriter<PartialRunSetupResponse>* response,
-        ::grpc::CompletionQueue* new_call_cq,
-        ::grpc::ServerCompletionQueue* notification_cq, void* tag) {
-      ::grpc::Service::RequestAsyncUnary(2, context, request, response,
-                                         new_call_cq, notification_cq, tag);
-    }
-    void RequestRunStep(
-        ::grpc::ServerContext* context, RunStepRequest* request,
-        ::grpc::ServerAsyncResponseWriter<RunStepResponse>* response,
-        ::grpc::CompletionQueue* new_call_cq,
-        ::grpc::ServerCompletionQueue* notification_cq, void* tag) {
-      ::grpc::Service::RequestAsyncUnary(3, context, request, response,
-                                         new_call_cq, notification_cq, tag);
-    }
-    void RequestCloseSession(
-        ::grpc::ServerContext* context, CloseSessionRequest* request,
-        ::grpc::ServerAsyncResponseWriter<CloseSessionResponse>* response,
-        ::grpc::CompletionQueue* new_call_cq,
-        ::grpc::ServerCompletionQueue* notification_cq, void* tag) {
-      ::grpc::Service::RequestAsyncUnary(4, context, request, response,
-                                         new_call_cq, notification_cq, tag);
-    }
-    void RequestListDevices(
-        ::grpc::ServerContext* context, ListDevicesRequest* request,
-        ::grpc::ServerAsyncResponseWriter<ListDevicesResponse>* response,
-        ::grpc::CompletionQueue* new_call_cq,
-        ::grpc::ServerCompletionQueue* notification_cq, void* tag) {
-      ::grpc::Service::RequestAsyncUnary(5, context, request, response,
-                                         new_call_cq, notification_cq, tag);
-    }
-    void RequestReset(
-        ::grpc::ServerContext* context, ResetRequest* request,
-        ::grpc::ServerAsyncResponseWriter<ResetResponse>* response,
-        ::grpc::CompletionQueue* new_call_cq,
-        ::grpc::ServerCompletionQueue* notification_cq, void* tag) {
-      ::grpc::Service::RequestAsyncUnary(6, context, request, response,
-                                         new_call_cq, notification_cq, tag);
-    }
-    void RequestMakeCallable(
-        ::grpc::ServerContext* context, MakeCallableRequest* request,
-        ::grpc::ServerAsyncResponseWriter<MakeCallableResponse>* response,
-        ::grpc::CompletionQueue* new_call_cq,
-        ::grpc::ServerCompletionQueue* notification_cq, void* tag) {
-      ::grpc::Service::RequestAsyncUnary(7, context, request, response,
-                                         new_call_cq, notification_cq, tag);
-    }
-    void RequestRunCallable(
-        ::grpc::ServerContext* context, RunCallableRequest* request,
-        ::grpc::ServerAsyncResponseWriter<RunCallableResponse>* response,
-        ::grpc::CompletionQueue* new_call_cq,
-        ::grpc::ServerCompletionQueue* notification_cq, void* tag) {
-      ::grpc::Service::RequestAsyncUnary(8, context, request, response,
-                                         new_call_cq, notification_cq, tag);
-    }
-    void RequestReleaseCallable(
-        ::grpc::ServerContext* context, ReleaseCallableRequest* request,
-        ::grpc::ServerAsyncResponseWriter<ReleaseCallableResponse>* response,
-        ::grpc::CompletionQueue* new_call_cq,
-        ::grpc::ServerCompletionQueue* notification_cq, void* tag) {
-      ::grpc::Service::RequestAsyncUnary(9, context, request, response,
-                                         new_call_cq, notification_cq, tag);
-    }
-  };
-};
-
-}  // namespace grpc
-
-}  // namespace tensorflow
-
-#endif  // TENSORFLOW_CORE_DISTRIBUTED_RUNTIME_RPC_GRPC_MASTER_SERVICE_IMPL_H_
diff --git a/tensorflow/core/distributed_runtime/rpc/grpc_remote_master.cc b/tensorflow/core/distributed_runtime/rpc/grpc_remote_master.cc
index b832a2115c..6c2940553c 100644
--- a/tensorflow/core/distributed_runtime/rpc/grpc_remote_master.cc
+++ b/tensorflow/core/distributed_runtime/rpc/grpc_remote_master.cc
@@ -19,13 +19,13 @@ limitations under the License.
 
 #include "tensorflow/core/distributed_runtime/call_options.h"
 #include "tensorflow/core/distributed_runtime/master_interface.h"
-#include "tensorflow/core/distributed_runtime/rpc/grpc_master_service_impl.h"
 #include "tensorflow/core/distributed_runtime/rpc/grpc_util.h"
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/lib/strings/strcat.h"
 #include "tensorflow/core/platform/tracing.h"
 #include "tensorflow/core/protobuf/master.pb.h"
+#include "tensorflow/core/protobuf/master_service.grpc.pb.h"
 
 namespace tensorflow {