TensorFlow Serving C++ API Documentation
resource_tracker_test.cc
1 /* Copyright 2016 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/resources/resource_tracker.h"
17 
18 #include <algorithm>
19 #include <memory>
20 #include <string>
21 
22 #include <gmock/gmock.h>
23 #include <gtest/gtest.h>
24 #include "tensorflow/core/lib/core/status.h"
25 #include "tensorflow/core/lib/core/status_test_util.h"
26 #include "tensorflow/core/platform/types.h"
27 #include "tensorflow_serving/core/test_util/mock_loader.h"
28 #include "tensorflow_serving/resources/resources.pb.h"
29 #include "tensorflow_serving/test_util/test_util.h"
30 #include "tensorflow_serving/util/any_ptr.h"
31 
32 using ::tensorflow::serving::test_util::CreateProto;
33 using ::tensorflow::serving::test_util::EqualsProto;
34 using ::testing::_;
35 using ::testing::Invoke;
36 using ::testing::NiceMock;
37 
38 namespace tensorflow {
39 namespace serving {
40 namespace {
41 
42 class ResourceTrackerTest : public ::testing::Test {
43  protected:
44  ResourceTrackerTest()
45  : total_resources_(
46  CreateProto<ResourceAllocation>("resource_quantities { "
47  " resource { "
48  " device: 'main' "
49  " device_instance { value: 0 } "
50  " kind: 'ram' "
51  " } "
52  " quantity: 16 "
53  "} "
54  "resource_quantities { "
55  " resource { "
56  " device: 'gpu' "
57  " device_instance { value: 0 } "
58  " kind: 'ram' "
59  " } "
60  " quantity: 16 "
61  "} "
62  "resource_quantities { "
63  " resource { "
64  " device: 'gpu' "
65  " device_instance { value: 1 } "
66  " kind: 'ram' "
67  " } "
68  " quantity: 16 "
69  "} ")) {
70  std::unique_ptr<ResourceUtil> util(
71  new ResourceUtil({{{"main", 1}, {"gpu", 2}}}));
72  TF_CHECK_OK(ResourceTracker::Create(
73  total_resources_, std::unique_ptr<ResourceUtil>(std::move(util)),
74  &tracker_));
75 
76  loader_0_.reset(new NiceMock<test_util::MockLoader>);
77  ON_CALL(*loader_0_, EstimateResources(_))
78  .WillByDefault(Invoke([](ResourceAllocation* estimate) {
79  *estimate = CreateProto<ResourceAllocation>(
80  "resource_quantities { "
81  " resource { "
82  " device: 'main' "
83  " kind: 'ram' "
84  " } "
85  " quantity: 1 "
86  "} "
87  "resource_quantities { "
88  " resource { "
89  " device: 'gpu' "
90  " device_instance { value: 0 } "
91  " kind: 'ram' "
92  " } "
93  " quantity: 3 "
94  "} ");
95  return Status();
96  }));
97 
98  loader_1_.reset(new NiceMock<test_util::MockLoader>);
99  ON_CALL(*loader_1_, EstimateResources(_))
100  .WillByDefault(Invoke([](ResourceAllocation* estimate) {
101  *estimate = CreateProto<ResourceAllocation>(
102  "resource_quantities { "
103  " resource { "
104  " device: 'main' "
105  " kind: 'ram' "
106  " } "
107  " quantity: 5 "
108  "} "
109  "resource_quantities { "
110  " resource { "
111  " device: 'gpu' "
112  " device_instance { value: 0 } "
113  " kind: 'ram' "
114  " } "
115  " quantity: 7 "
116  "} ");
117  return Status();
118  }));
119 
120  loader_2_.reset(new NiceMock<test_util::MockLoader>);
121  ON_CALL(*loader_2_, EstimateResources(_))
122  .WillByDefault(Invoke([](ResourceAllocation* estimate) {
123  *estimate = CreateProto<ResourceAllocation>(
124  "resource_quantities { "
125  " resource { "
126  " device: 'main' "
127  " kind: 'ram' "
128  " } "
129  " quantity: 15 "
130  "} ");
131  return Status();
132  }));
133 
134  loader_3_.reset(new NiceMock<test_util::MockLoader>);
135  ON_CALL(*loader_3_, EstimateResources(_))
136  .WillByDefault(Invoke([](ResourceAllocation* estimate) {
137  *estimate = CreateProto<ResourceAllocation>(
138  "resource_quantities { "
139  " resource { "
140  " device: 'gpu' "
141  " kind: 'ram' "
142  " } "
143  " quantity: 12 "
144  "} ");
145  return Status();
146  }));
147 
148  invalid_resources_loader_.reset(new NiceMock<test_util::MockLoader>);
149  ON_CALL(*invalid_resources_loader_, EstimateResources(_))
150  .WillByDefault(Invoke([](ResourceAllocation* estimate) {
151  *estimate = CreateProto<ResourceAllocation>(
152  "resource_quantities { "
153  " resource { "
154  " device: 'bogus_device' "
155  " device_instance { value: 0 } "
156  " kind: 'ram' "
157  " } "
158  " quantity: 4 "
159  "} ");
160  return Status();
161  }));
162 
163  // Disallow calls to Load()/Unload().
164  for (auto* loader : {loader_0_.get(), loader_1_.get(), loader_2_.get(),
165  loader_3_.get(), invalid_resources_loader_.get()}) {
166  EXPECT_CALL(*loader, Load()).Times(0);
167  EXPECT_CALL(*loader, Unload()).Times(0);
168  }
169  }
170 
171  // The original total-resources valued provided to 'tracker_'.
172  const ResourceAllocation total_resources_;
173 
174  // The object under testing.
175  std::unique_ptr<ResourceTracker> tracker_;
176 
177  // Some mock loaders with specific resource estimates (see the constructor).
178  std::unique_ptr<test_util::MockLoader> loader_0_;
179  std::unique_ptr<test_util::MockLoader> loader_1_;
180  std::unique_ptr<test_util::MockLoader> loader_2_;
181  std::unique_ptr<test_util::MockLoader> loader_3_;
182  std::unique_ptr<test_util::MockLoader> invalid_resources_loader_;
183 };
184 
185 TEST_F(ResourceTrackerTest, UnboundTotalResources) {
186  std::unique_ptr<ResourceUtil> util(
187  new ResourceUtil({{{"main", 1}, {"gpu", 2}}}));
188  std::unique_ptr<ResourceTracker> tracker;
189  const auto unbound_resources = CreateProto<ResourceAllocation>(
190  "resource_quantities { "
191  " resource { "
192  " device: 'gpu' "
193  " kind: 'ram' "
194  " } "
195  " quantity: 12 "
196  "} ");
197  EXPECT_FALSE(ResourceTracker::Create(
198  unbound_resources,
199  std::unique_ptr<ResourceUtil>(std::move(util)), &tracker)
200  .ok());
201 }
202 
203 TEST_F(ResourceTrackerTest, UnnormalizedTotalResources) {
204  std::unique_ptr<ResourceUtil> util(
205  new ResourceUtil({{{"main", 1}, {"gpu", 2}}}));
206  std::unique_ptr<ResourceTracker> tracker;
207  const auto unnormalized_resources = CreateProto<ResourceAllocation>(
208  "resource_quantities { "
209  " resource { "
210  " device: 'main' "
211  " kind: 'ram' "
212  " } "
213  " quantity: 12 "
214  "} ");
215  TF_ASSERT_OK(ResourceTracker::Create(
216  unnormalized_resources, std::unique_ptr<ResourceUtil>(std::move(util)),
217  &tracker));
218  // The total_resources proto should get normalized.
219  EXPECT_THAT(tracker->total_resources(),
220  EqualsProto("resource_quantities { "
221  " resource { "
222  " device: 'main' "
223  " device_instance { value: 0 } "
224  " kind: 'ram' "
225  " } "
226  " quantity: 12 "
227  "} "));
228 }
229 
230 TEST_F(ResourceTrackerTest, RecomputeUsedResources) {
231  // Verify the initial state.
232  EXPECT_THAT(tracker_->used_resources(), EqualsProto(""));
233  EXPECT_THAT(tracker_->total_resources(), EqualsProto(total_resources_));
234 
235  // Recompute used resources for {loader_0_, loader_1_, loader_3_}.
236  TF_ASSERT_OK(tracker_->RecomputeUsedResources(
237  {loader_0_.get(), loader_1_.get(), loader_3_.get()}));
238  EXPECT_THAT(tracker_->used_resources(),
239  EqualsProto("resource_quantities { "
240  " resource { "
241  " device: 'main' "
242  " device_instance { value: 0 } "
243  " kind: 'ram' "
244  " } "
245  " quantity: 6 "
246  "} "
247  "resource_quantities { "
248  " resource { "
249  " device: 'gpu' "
250  " device_instance { value: 0 } "
251  " kind: 'ram' "
252  " } "
253  " quantity: 10 "
254  "} "
255  "resource_quantities { "
256  " resource { "
257  " device: 'gpu' "
258  " kind: 'ram' "
259  " } "
260  " quantity: 12 "
261  "} "));
262  EXPECT_THAT(tracker_->total_resources(), EqualsProto(total_resources_));
263 
264  // Recompute used resources for just {loader_0_}.
265  TF_ASSERT_OK(tracker_->RecomputeUsedResources({loader_0_.get()}));
266  EXPECT_THAT(tracker_->used_resources(),
267  EqualsProto("resource_quantities { "
268  " resource { "
269  " device: 'main' "
270  " device_instance { value: 0 } "
271  " kind: 'ram' "
272  " } "
273  " quantity: 1 "
274  "} "
275  "resource_quantities { "
276  " resource { "
277  " device: 'gpu' "
278  " device_instance { value: 0 } "
279  " kind: 'ram' "
280  " } "
281  " quantity: 3 "
282  "} "));
283  EXPECT_THAT(tracker_->total_resources(), EqualsProto(total_resources_));
284 }
285 
286 TEST_F(ResourceTrackerTest, ReserveResourcesSuccessWithUsedResourcesBound) {
287  TF_ASSERT_OK(tracker_->RecomputeUsedResources({loader_0_.get()}));
288  EXPECT_THAT(tracker_->used_resources(),
289  EqualsProto("resource_quantities { "
290  " resource { "
291  " device: 'main' "
292  " device_instance { value: 0 } "
293  " kind: 'ram' "
294  " } "
295  " quantity: 1 "
296  "} "
297  "resource_quantities { "
298  " resource { "
299  " device: 'gpu' "
300  " device_instance { value: 0 } "
301  " kind: 'ram' "
302  " } "
303  " quantity: 3 "
304  "} "));
305 
306  // If just loader_0_ is loaded, loader_2_ should also fit.
307  bool success;
308  TF_ASSERT_OK(tracker_->ReserveResources(*loader_2_, &success));
309  EXPECT_TRUE(success);
310  EXPECT_THAT(tracker_->used_resources(),
311  EqualsProto("resource_quantities { "
312  " resource { "
313  " device: 'main' "
314  " device_instance { value: 0 } "
315  " kind: 'ram' "
316  " } "
317  " quantity: 16 "
318  "} "
319  "resource_quantities { "
320  " resource { "
321  " device: 'gpu' "
322  " device_instance { value: 0 } "
323  " kind: 'ram' "
324  " } "
325  " quantity: 3 "
326  "} "));
327  EXPECT_THAT(tracker_->total_resources(), EqualsProto(total_resources_));
328 }
329 
330 TEST_F(ResourceTrackerTest, ReserveResourcesFailureWithUsedResourcesBound) {
331  TF_ASSERT_OK(tracker_->RecomputeUsedResources({loader_1_.get()}));
332  EXPECT_THAT(tracker_->used_resources(),
333  EqualsProto("resource_quantities { "
334  " resource { "
335  " device: 'main' "
336  " device_instance { value: 0 } "
337  " kind: 'ram' "
338  " } "
339  " quantity: 5 "
340  "} "
341  "resource_quantities { "
342  " resource { "
343  " device: 'gpu' "
344  " device_instance { value: 0 } "
345  " kind: 'ram' "
346  " } "
347  " quantity: 7 "
348  "} "));
349 
350  // If loader_1_ is loaded, there isn't room for loader_2_.
351  bool success;
352  TF_ASSERT_OK(tracker_->ReserveResources(*loader_2_, &success));
353  EXPECT_FALSE(success);
354  // The used resources should remain unchanged (i.e. only reflect loader_1_).
355  EXPECT_THAT(tracker_->used_resources(),
356  EqualsProto("resource_quantities { "
357  " resource { "
358  " device: 'main' "
359  " device_instance { value: 0 } "
360  " kind: 'ram' "
361  " } "
362  " quantity: 5 "
363  "} "
364  "resource_quantities { "
365  " resource { "
366  " device: 'gpu' "
367  " device_instance { value: 0 } "
368  " kind: 'ram' "
369  " } "
370  " quantity: 7 "
371  "} "));
372  EXPECT_THAT(tracker_->total_resources(), EqualsProto(total_resources_));
373 }
374 
375 TEST_F(ResourceTrackerTest, ReserveResourcesSuccessWithUsedResourcesUnbound) {
376  TF_ASSERT_OK(tracker_->RecomputeUsedResources({loader_3_.get()}));
377  EXPECT_THAT(tracker_->used_resources(), EqualsProto("resource_quantities { "
378  " resource { "
379  " device: 'gpu' "
380  " kind: 'ram' "
381  " } "
382  " quantity: 12 "
383  "} "));
384 
385  // If just loader_3_ is loaded, loader_0_ should also fit.
386  bool success;
387  TF_ASSERT_OK(tracker_->ReserveResources(*loader_0_, &success));
388  EXPECT_TRUE(success);
389  EXPECT_THAT(tracker_->used_resources(),
390  EqualsProto("resource_quantities { "
391  " resource { "
392  " device: 'gpu' "
393  " kind: 'ram' "
394  " } "
395  " quantity: 12 "
396  "} "
397  "resource_quantities { "
398  " resource { "
399  " device: 'main' "
400  " device_instance { value: 0 } "
401  " kind: 'ram' "
402  " } "
403  " quantity: 1 "
404  "} "
405  "resource_quantities { "
406  " resource { "
407  " device: 'gpu' "
408  " device_instance { value: 0 } "
409  " kind: 'ram' "
410  " } "
411  " quantity: 3 "
412  "} "));
413  EXPECT_THAT(tracker_->total_resources(), EqualsProto(total_resources_));
414 }
415 
416 TEST_F(ResourceTrackerTest, ReserveResourcesFailureWithUsedResourcesUnbound) {
417  TF_ASSERT_OK(tracker_->RecomputeUsedResources({loader_3_.get()}));
418  EXPECT_THAT(tracker_->used_resources(), EqualsProto("resource_quantities { "
419  " resource { "
420  " device: 'gpu' "
421  " kind: 'ram' "
422  " } "
423  " quantity: 12 "
424  "} "));
425 
426  // If loader_3_ is loaded, there isn't room for loader_1_, or for another
427  // copy of loader_3_.
428  bool success;
429  TF_ASSERT_OK(tracker_->ReserveResources(*loader_1_, &success));
430  EXPECT_FALSE(success);
431  TF_ASSERT_OK(tracker_->ReserveResources(*loader_3_, &success));
432  EXPECT_FALSE(success);
433  // The used resources should remain unchanged (i.e. only reflect a single copy
434  // of loader_3_).
435  EXPECT_THAT(tracker_->used_resources(), EqualsProto("resource_quantities { "
436  " resource { "
437  " device: 'gpu' "
438  " kind: 'ram' "
439  " } "
440  " quantity: 12 "
441  "} "));
442  EXPECT_THAT(tracker_->total_resources(), EqualsProto(total_resources_));
443 }
444 
445 TEST_F(ResourceTrackerTest, InvalidResourceEstimate) {
446  bool success;
447  EXPECT_FALSE(
448  tracker_->ReserveResources(*invalid_resources_loader_, &success).ok());
449  EXPECT_FALSE(tracker_
450  ->RecomputeUsedResources(
451  {loader_0_.get(), invalid_resources_loader_.get()})
452  .ok());
453 }
454 
455 } // namespace
456 } // namespace serving
457 } // namespace tensorflow
Generated by doxygen 1.9.1