Update CreateRPC API interface.

PiperOrigin-RevId: 407363453

Author: guanxinq

tensorflow_serving/experimental/tensorflow/ops/remote_predict/BUILD

@@ -26,6 +26,7 @@ cc_library(
         "@org_tensorflow//tensorflow/core:framework_headers_lib",
         "@org_tensorflow//tensorflow/core:protos_all_cc",
         "@org_tensorflow//tensorflow/core/kernels:ops_util_hdrs",
+        "@org_tensorflow//tensorflow/core/platform:statusor",
     ],
     alwayslink = 1,
 )
@@ -96,6 +97,7 @@ tf_kernel_library(
         "@org_tensorflow//tensorflow/core:protos_all_cc",
         "@org_tensorflow//tensorflow/core/kernels:ops_util",
         "@org_tensorflow//tensorflow/core/kernels:split_lib",
+        "@org_tensorflow//tensorflow/core/platform:statusor",
     ],
 )
 
@@ -163,6 +165,7 @@ cc_library(
         "@com_github_grpc_grpc//:grpc++",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/time",
+        "@org_tensorflow//tensorflow/core/platform:statusor",
     ],
 )
 

tensorflow_serving/experimental/tensorflow/ops/remote_predict/kernels/prediction_service_grpc.cc

@@ -41,7 +41,7 @@ PredictionServiceGrpc::PredictionServiceGrpc(
   stub_ = tensorflow::serving::PredictionService::NewStub(channel);
 }
 
-::grpc::ClientContext* PredictionServiceGrpc::CreateRpc(
+StatusOr<::grpc::ClientContext*> PredictionServiceGrpc::CreateRpc(
     absl::Duration max_rpc_deadline) {
   ::grpc::ClientContext* rpc = new ::grpc::ClientContext();
   // TODO(b/159739577): Set deadline as the min value between

tensorflow_serving/experimental/tensorflow/ops/remote_predict/kernels/prediction_service_grpc.h

@@ -18,6 +18,7 @@ limitations under the License.
 
 #include "absl/status/status.h"
 #include "absl/time/time.h"
+#include "tensorflow/core/platform/statusor.h"
 #include "tensorflow_serving/apis/prediction_service.grpc.pb.h"
 
 namespace tensorflow {
@@ -33,7 +34,7 @@ class PredictionServiceGrpc {
     return ::absl::OkStatus();
   }
 
-  ::grpc::ClientContext* CreateRpc(absl::Duration max_rpc_deadline);
+  StatusOr<::grpc::ClientContext*> CreateRpc(absl::Duration max_rpc_deadline);
 
   void Predict(::grpc::ClientContext* rpc, PredictRequest* request,
                PredictResponse* response,

tensorflow_serving/experimental/tensorflow/ops/remote_predict/kernels/prediction_service_grpc_test.cc

@@ -35,7 +35,9 @@ class PredictionServiceGrpcTest : public ::testing::Test {
 
 TEST_F(PredictionServiceGrpcTest, TestSetDeadline) {
   const absl::Duration deadline = absl::Milliseconds(30000);
-  rpc_.reset(grpc_stub_->CreateRpc(deadline));
+  auto rpc_or = grpc_stub_->CreateRpc(deadline);
+  ASSERT_TRUE(rpc_or.ok());
+  rpc_.reset(rpc_or.ValueOrDie());
 
   EXPECT_NEAR(absl::ToDoubleMilliseconds(deadline),
               absl::ToDoubleMilliseconds(absl::FromChrono(rpc_->deadline()) -

tensorflow_serving/experimental/tensorflow/ops/remote_predict/kernels/remote_predict_op_kernel.h

@@ -29,6 +29,7 @@ limitations under the License.
 #include "tensorflow/core/kernels/ops_util.h"
 #include "tensorflow/core/lib/core/threadpool.h"
 #include "tensorflow/core/lib/gtl/cleanup.h"
+#include "tensorflow/core/platform/status.h"
 #include "tensorflow/core/protobuf/named_tensor.pb.h"
 #include "tensorflow_serving/apis/model.pb.h"
 #include "tensorflow_serving/apis/predict.pb.h"
@@ -103,9 +104,17 @@ class RemotePredictOp : public AsyncOpKernel {
 
     PredictResponse* response = new PredictResponse();
 
-    auto rpc = prediction_service_->CreateRpc(
+    auto rpc_or = prediction_service_->CreateRpc(
         absl::Milliseconds(max_rpc_deadline_millis_));
-
+    OP_REQUIRES_ASYNC(context, rpc_or.ok(),
+                      tensorflow::Status(rpc_or.status().code(),
+                                         rpc_or.status().error_message()),
+                      [&]() {
+                        delete request;
+                        delete response;
+                        done();
+                      });
+    auto rpc = rpc_or.ValueOrDie();
     auto callback = [this, context, rpc, request, response,
                      output_tensor_aliases, done](const absl::Status& status) {
       PostProcessResponse(context, response, status, fail_op_on_rpc_error_,

tensorflow_serving/experimental/tensorflow/ops/remote_predict/kernels/remote_predict_op_kernel_test.cc

@@ -20,6 +20,7 @@ limitations under the License.
 #include "tensorflow/cc/ops/const_op.h"
 #include "tensorflow/core/framework/tensor_testutil.h"
 #include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/status.h"
 #include "tensorflow_serving/apis/prediction_service.grpc.pb.h"
 #include "tensorflow_serving/experimental/tensorflow/ops/remote_predict/cc/ops/remote_predict_op.h"
 
@@ -39,7 +40,9 @@ class MockPredictionService {
     return ::absl::OkStatus();
   }
 
-  MockRpc* CreateRpc(absl::Duration max_rpc_deadline) { return new MockRpc; }
+  StatusOr<MockRpc*> CreateRpc(absl::Duration max_rpc_deadline) {
+    return new MockRpc;
+  }
 
   // The model_name in request determines response and/or status.
   void Predict(MockRpc* rpc, PredictRequest* request, PredictResponse* response,