TensorFlow Serving C++ API Documentation
tfrt_saved_model_warmup_test.cc
1 /* Copyright 2020 Google Inc. 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 #include "tensorflow_serving/servables/tensorflow/tfrt_saved_model_warmup.h"
17 
18 #include <string>
19 
20 #include "google/protobuf/wrappers.pb.h"
21 #include "tensorflow/cc/saved_model/constants.h"
22 #include "tensorflow/cc/saved_model/signature_constants.h"
23 #include "tensorflow/core/lib/core/status_test_util.h"
24 #include "tensorflow/core/platform/test.h"
25 #include "tensorflow/core/tfrt/graph_executor/graph_execution_options.h"
26 #include "tensorflow/core/tfrt/saved_model/saved_model.h"
27 #include "tensorflow/core/tfrt/utils/tensor_util.h"
28 #include "tsl/platform/path.h"
29 #include "tensorflow_serving/apis/classification.pb.h"
30 #include "tensorflow_serving/apis/inference.pb.h"
31 #include "tensorflow_serving/apis/input.pb.h"
32 #include "tensorflow_serving/apis/model.pb.h"
33 #include "tensorflow_serving/apis/predict.pb.h"
34 #include "tensorflow_serving/apis/prediction_log.pb.h"
35 #include "tensorflow_serving/apis/regression.pb.h"
36 #include "tensorflow_serving/servables/tensorflow/saved_model_warmup_test_util.h"
37 #include "tensorflow_serving/servables/tensorflow/test_util/mock_tfrt_saved_model.h"
38 
39 namespace tensorflow {
40 namespace serving {
41 namespace {
42 
43 using ::testing::_;
44 using ::testing::DoAll;
45 using ::testing::Return;
46 using ::testing::ReturnRef;
47 using ::testing::WithArgs;
48 
49 class TFRTSavedModelWarmupOptionsTest : public ::testing::TestWithParam<bool> {
50  public:
51  bool EnableNumRequestIterations() { return GetParam(); }
52 
53  ModelWarmupOptions GetModelWarmupOptions() {
54  ModelWarmupOptions options;
55  if (EnableNumRequestIterations()) {
56  options.mutable_num_request_iterations()->set_value(2);
57  }
58  return options;
59  }
60 
61  int GetNumRequestIterations() {
62  if (EnableNumRequestIterations()) {
63  return 2;
64  }
65  return 1;
66  }
67 };
68 
69 TEST_P(TFRTSavedModelWarmupOptionsTest, MixedWarmupData) {
70  string base_path = io::JoinPath(testing::TmpDir(), "MixedWarmupData");
71  TF_ASSERT_OK(Env::Default()->RecursivelyCreateDir(
72  io::JoinPath(base_path, kSavedModelAssetsExtraDirectory)));
73  string fname = io::JoinPath(base_path, kSavedModelAssetsExtraDirectory,
74  internal::WarmupConsts::kRequestsFileName);
75 
76  int num_warmup_records = 10;
77  std::vector<string> warmup_records;
78  TF_ASSERT_OK(AddMixedWarmupData(&warmup_records));
79  TF_ASSERT_OK(WriteWarmupData(fname, warmup_records, num_warmup_records));
80 
81  std::unique_ptr<test_util::MockSavedModel> saved_model(
82  (new test_util::MockSavedModel()));
83  tfrt::internal::Signature predict_signature;
84  predict_signature.input_names = {kPredictInputs};
85  tfrt::TensorSpec spec(tensorflow::DT_STRING);
86  predict_signature.input_specs = {spec};
87  predict_signature.output_names = {kPredictOutputs};
88  tfrt::FunctionMetadata predict_function_metadata(&predict_signature);
89  EXPECT_CALL(*saved_model, GetFunctionMetadata(kPredictMethodName))
90  .WillRepeatedly(Return(predict_function_metadata));
91 
92  tfrt::internal::Signature classify_signature;
93  classify_signature.input_names = {kClassifyInputs};
94  classify_signature.output_names = {kClassifyOutputClasses,
95  kClassifyOutputScores};
96  tfrt::FunctionMetadata classify_function_metadata(&classify_signature);
97  EXPECT_CALL(*saved_model, GetFunctionMetadata(kClassifyMethodName))
98  .WillRepeatedly(Return(classify_function_metadata));
99 
100  tfrt::internal::Signature regress_signature;
101  regress_signature.input_names = {kRegressInputs};
102  regress_signature.output_names = {kRegressOutputs};
103  tfrt::FunctionMetadata regress_function_metadata(&regress_signature);
104  EXPECT_CALL(*saved_model, GetFunctionMetadata(kRegressMethodName))
105  .WillRepeatedly(Return(regress_function_metadata));
106 
107  MetaGraphDef meta_graph_def;
108  AddSignatures(&meta_graph_def);
109  EXPECT_CALL(*saved_model, GetMetaGraphDef())
110  .WillRepeatedly(ReturnRef(meta_graph_def));
111 
112  Tensor scores(DT_FLOAT, TensorShape({1, 1}));
113  Tensor classes(DT_STRING, TensorShape({1, 1}));
114 
115  EXPECT_CALL(*saved_model, Run(_, ::testing::Eq(kPredictMethodName),
116  ::testing::An<absl::Span<const Tensor>>(), _))
117  .Times(num_warmup_records * GetNumRequestIterations())
118  .WillRepeatedly(
119  DoAll(WithArgs<3>([&](std::vector<Tensor>* output_tensors) {
120  output_tensors->push_back(scores);
121  }),
122  Return(absl::OkStatus())));
123 
124  EXPECT_CALL(*saved_model, Run(_, ::testing::Eq(kRegressMethodName),
125  ::testing::An<absl::Span<const Tensor>>(), _))
126  .Times(num_warmup_records * GetNumRequestIterations())
127  .WillRepeatedly(
128  DoAll(WithArgs<3>([&](std::vector<Tensor>* output_tensors) {
129  output_tensors->push_back(scores);
130  }),
131  Return(absl::OkStatus())));
132 
133  EXPECT_CALL(*saved_model, Run(_, ::testing::Eq(kClassifyMethodName),
134  ::testing::An<absl::Span<const Tensor>>(), _))
135  .Times(num_warmup_records * GetNumRequestIterations())
136  .WillRepeatedly(
137  DoAll(WithArgs<3>([&](std::vector<Tensor>* output_tensors) {
138  output_tensors->push_back(classes);
139  output_tensors->push_back(scores);
140  }),
141  Return(absl::OkStatus())));
142 
143  EXPECT_CALL(*saved_model, RunMultipleSignatures(_, _, _, _))
144  .Times(num_warmup_records * GetNumRequestIterations())
145  .WillRepeatedly(DoAll(
146  WithArgs<3>([&](std::vector<std::vector<Tensor>>* output_tensors) {
147  output_tensors->resize(2);
148  (*output_tensors)[0].push_back(scores);
149  (*output_tensors)[1].push_back(classes);
150  (*output_tensors)[1].push_back(scores);
151  }),
152  Return(absl::OkStatus())));
153 
154  TF_EXPECT_OK(RunSavedModelWarmup(GetModelWarmupOptions(), base_path,
155  /*lazy_init_threshold=*/0,
156  /*skip_warmup_requests_if_initialized=*/true,
157  saved_model.get()));
158 }
159 
160 TEST_P(TFRTSavedModelWarmupOptionsTest, PredictStreamedWarmupData) {
161  std::string base_path =
162  tsl::io::JoinPath(testing::TmpDir(), "PredictStreamedWarmupData");
163  TF_ASSERT_OK(Env::Default()->RecursivelyCreateDir(
164  tsl::io::JoinPath(base_path, kSavedModelAssetsExtraDirectory)));
165  std::string fname =
166  tsl::io::JoinPath(base_path, kSavedModelAssetsExtraDirectory,
167  internal::WarmupConsts::kRequestsFileName);
168 
169  int num_warmup_records = 10;
170  std::vector<std::string> warmup_records;
171  TF_ASSERT_OK(
172  AddToWarmupData(&warmup_records, PredictionLog::kPredictStreamedLog));
173  TF_ASSERT_OK(WriteWarmupData(fname, warmup_records, num_warmup_records));
174 
175  auto saved_model = std::make_unique<test_util::MockSavedModel>();
176 
177  tfrt::internal::Signature signature;
178  signature.input_names = {kPredictInputs};
179  signature.input_specs = {tfrt::TensorSpec(tensorflow::DT_STRING)};
180  tfrt::FunctionMetadata function_metadata(&signature);
181  EXPECT_CALL(*saved_model, GetFunctionMetadata(kPredictMethodName))
182  .WillRepeatedly(Return(function_metadata));
183 
184  MetaGraphDef meta_graph_def;
185  (*meta_graph_def.mutable_signature_def())[kPredictMethodName] =
186  CreateSignatureDef(kPredictMethodName, {kPredictInputs}, {});
187 
188  EXPECT_CALL(*saved_model, GetMetaGraphDef())
189  .WillRepeatedly(ReturnRef(meta_graph_def));
190 
191  EXPECT_CALL(
192  *saved_model,
193  Run(::testing::Field(
194  &tfrt_stub::GraphExecutionRunOptions::streamed_output_callback,
195  ::testing::NotNull()),
196  ::testing::Eq(kPredictMethodName),
197  ::testing::An<absl::Span<const Tensor>>(), _))
198  .Times(num_warmup_records * GetNumRequestIterations())
199  .WillRepeatedly(Return(absl::OkStatus()));
200 
201  TF_EXPECT_OK(RunSavedModelWarmup(GetModelWarmupOptions(), base_path,
202  /*lazy_init_threshold=*/0,
203  /*skip_warmup_requests_if_initialized=*/true,
204  saved_model.get()));
205 }
206 
207 INSTANTIATE_TEST_SUITE_P(WarmupOptions, TFRTSavedModelWarmupOptionsTest,
208  ::testing::Bool());
209 
210 TEST(TFRTSavedModelWarmupTest, UnsupportedLogType) {
211  string base_path = io::JoinPath(testing::TmpDir(), "UnsupportedLogType");
212  TF_ASSERT_OK(Env::Default()->RecursivelyCreateDir(
213  io::JoinPath(base_path, kSavedModelAssetsExtraDirectory)));
214  string fname = io::JoinPath(base_path, kSavedModelAssetsExtraDirectory,
215  internal::WarmupConsts::kRequestsFileName);
216 
217  std::vector<string> warmup_records;
218  // Add unsupported log type
219  PredictionLog prediction_log;
220  TF_ASSERT_OK(
221  PopulatePredictionLog(&prediction_log, PredictionLog::kSessionRunLog));
222  warmup_records.push_back(prediction_log.SerializeAsString());
223  TF_ASSERT_OK(WriteWarmupData(fname, warmup_records, 10));
224 
225  std::unique_ptr<test_util::MockSavedModel> saved_model(
226  (new test_util::MockSavedModel()));
227  MetaGraphDef meta_graph_def;
228  AddSignatures(&meta_graph_def);
229  EXPECT_CALL(*saved_model, GetMetaGraphDef())
230  .WillRepeatedly(ReturnRef(meta_graph_def));
231  const Status status = RunSavedModelWarmup(
232  ModelWarmupOptions(), base_path,
233  /*lazy_init_threshold=*/0,
234  /*skip_warmup_requests_if_initialized=*/true, saved_model.get());
235  ASSERT_FALSE(status.ok());
236  EXPECT_EQ(::tensorflow::error::UNIMPLEMENTED, status.code()) << status;
237  EXPECT_THAT(status.ToString(),
238  ::testing::HasSubstr("Unsupported log_type for warmup"));
239 }
240 
241 TEST(TFRTSavedModelWarmupTest, SkipWarmupRequest) {
242  string base_path = io::JoinPath(testing::TmpDir(), "SkipWarmupRequest");
243  TF_ASSERT_OK(Env::Default()->RecursivelyCreateDir(
244  io::JoinPath(base_path, kSavedModelAssetsExtraDirectory)));
245  string fname = io::JoinPath(base_path, kSavedModelAssetsExtraDirectory,
246  internal::WarmupConsts::kRequestsFileName);
247 
248  int num_warmup_records = 10;
249  std::vector<string> warmup_records;
250  TF_ASSERT_OK(AddMixedWarmupData(
251  &warmup_records, {PredictionLog::kRegressLog, PredictionLog::kClassifyLog,
252  PredictionLog::kPredictLog}));
253  TF_ASSERT_OK(WriteWarmupData(fname, warmup_records, num_warmup_records));
254 
255  std::unique_ptr<test_util::MockSavedModel> saved_model(
256  (new test_util::MockSavedModel()));
257  EXPECT_CALL(*saved_model, GetFunctionMetadata(kPredictMethodName)).Times(0);
258  EXPECT_CALL(*saved_model, GetFunctionMetadata(kClassifyMethodName)).Times(0);
259  EXPECT_CALL(*saved_model, GetFunctionMetadata(kRegressMethodName)).Times(0);
260  MetaGraphDef meta_graph_def;
261  AddSignatures(&meta_graph_def);
262  EXPECT_CALL(*saved_model, GetMetaGraphDef())
263  .WillRepeatedly(ReturnRef(meta_graph_def));
264 
265  TF_EXPECT_OK(RunSavedModelWarmup(ModelWarmupOptions(), base_path,
266  /*lazy_init_threshold=*/10,
267  /*skip_warmup_requests_if_initialized=*/true,
268  saved_model.get()));
269 }
270 
271 } // namespace
272 } // namespace serving
273 } // namespace tensorflow
Generated by doxygen 1.9.1