Merge pull request #28 from kirilg/branch_118386107

Upstream changes from internal

Author: Christopher Olston

tensorflow_serving/core/BUILD

@@ -310,6 +310,7 @@ cc_test(
         ":eager_load_policy",
         ":servable_state",
         "//tensorflow_serving/core/test_util:dynamic_manager_test_util",
+        "//tensorflow_serving/core/test_util:fake_loader",
         "//tensorflow_serving/core/test_util:mock_loader",
         "//tensorflow_serving/core/test_util:test_main",
         "//tensorflow_serving/util:any_ptr",

tensorflow_serving/core/availability_helpers.cc

@@ -26,8 +26,9 @@ namespace tensorflow {
 namespace serving {
 namespace {
 
-bool ServablesAvailable(Manager* const manager,
-                        const std::vector<ServableRequest>& servables) {
+// Implements the test logic of WaitUntilServablesAvailableForRequests().
+bool ServablesAvailableForRequests(
+    const std::vector<ServableRequest>& servables, Manager* const manager) {
   // The subset of 'servables' that require a specific version
   std::set<ServableId> query_specific_servables;
   // The subset of 'servables' that do not require a specific version
@@ -69,12 +70,21 @@ bool ServablesAvailable(Manager* const manager,
 
 }  // namespace
 
-void WaitUntilServablesAvailable(
-    Manager* const manager, const std::vector<ServableRequest>& servables) {
-  while (!ServablesAvailable(manager, servables)) {
+void WaitUntilServablesAvailableForRequests(
+    const std::vector<ServableRequest>& servables, Manager* const manager) {
+  while (!ServablesAvailableForRequests(servables, manager)) {
     Env::Default()->SleepForMicroseconds(50 * 1000 /* 50 ms */);
   }
 }
 
+void WaitUntilServablesAvailable(const std::set<ServableId>& servables,
+                                 Manager* const manager) {
+  std::vector<ServableRequest> requests;
+  for (const ServableId& id : servables) {
+    requests.push_back(ServableRequest::FromId(id));
+  }
+  return WaitUntilServablesAvailableForRequests(requests, manager);
+}
+
 }  // namespace serving
 }  // namespace tensorflow

tensorflow_serving/core/availability_helpers.h

@@ -41,8 +41,13 @@ namespace serving {
 // polling is 50ms.  We call ListAvailableServableIds on the manager, which may
 // have side-effects for certain manager implementations, e.g. causing servables
 // to be loaded.
-void WaitUntilServablesAvailable(Manager* manager,
-                                 const std::vector<ServableRequest>& servables);
+void WaitUntilServablesAvailableForRequests(
+    const std::vector<ServableRequest>& servables, Manager* manager);
+
+// Like WaitUntilServablesAvailableForRequests(), but taking a set of servable
+// ids (and hence waits for the specific versions specified in the ids).
+void WaitUntilServablesAvailable(const std::set<ServableId>& servables,
+                                 Manager* manager);
 
 }  // namespace serving
 }  // namespace tensorflow

tensorflow_serving/core/availability_helpers_test.cc

@@ -66,25 +66,20 @@ class FakeManager : public Manager {
   condition_variable list_available_servable_ids_condition_;
 };
 
-TEST(AvailabilityHelpersTest, Available) {
-  std::vector<ServableRequest> available_servables_query;
-  available_servables_query.emplace_back(
-      ServableRequest::Specific("servable0", 0));
+TEST(AvailabilityHelpersTest, SpecificAvailable) {
   FakeManager fake_manager;
   fake_manager.set_available_servable_ids({{"servable0", 0}});
-  WaitUntilServablesAvailable(&fake_manager, available_servables_query);
+  WaitUntilServablesAvailable({{"servable0", 0}}, &fake_manager);
 }
 
 TEST(AvailabilityHelpersTest, SpecificNotAvailable) {
   FakeManager fake_manager;
   fake_manager.set_available_servable_ids({{"servable0", 0}});
-  const std::vector<ServableRequest> available_servables_query = {
-      ServableRequest::Specific("servable0", 0),
-      ServableRequest::Specific("servable1", 0)};
   Notification finished;
   std::unique_ptr<Thread> wait(
       Env::Default()->StartThread({}, "WaitUntilServablesAvailable", [&]() {
-        WaitUntilServablesAvailable(&fake_manager, available_servables_query);
+        WaitUntilServablesAvailable({{"servable0", 0}, {"servable1", 0}},
+                                    &fake_manager);
         finished.Notify();
       }));
   // Waiting for 2 calls ensures that we waited at least once for the servables
@@ -98,13 +93,13 @@ TEST(AvailabilityHelpersTest, SpecificNotAvailable) {
 TEST(AvailabilityHelpersTest, LatestNotAvailable) {
   FakeManager fake_manager;
   fake_manager.set_available_servable_ids({{"servable0", 0}});
-  const std::vector<ServableRequest> available_servables_query = {
-      ServableRequest::Specific("servable0", 0),
-      ServableRequest::Latest("servable1")};
   Notification finished;
   std::unique_ptr<Thread> wait(
       Env::Default()->StartThread({}, "WaitUntilServablesAvailable", [&]() {
-        WaitUntilServablesAvailable(&fake_manager, available_servables_query);
+        WaitUntilServablesAvailableForRequests(
+            {ServableRequest::Specific("servable0", 0),
+             ServableRequest::Latest("servable1")},
+            &fake_manager);
         finished.Notify();
       }));
   // Waiting for 2 calls ensures that we waited at least once for the servables
@@ -116,21 +111,19 @@ TEST(AvailabilityHelpersTest, LatestNotAvailable) {
 }
 
 TEST(AvailabilityHelpersTest, LatestVersion) {
-  std::vector<ServableRequest> available_servables_query;
-  available_servables_query.emplace_back(ServableRequest::Latest("servable0"));
   FakeManager fake_manager;
   fake_manager.set_available_servable_ids({{"servable0", 123}});
-  WaitUntilServablesAvailable(&fake_manager, available_servables_query);
+  WaitUntilServablesAvailableForRequests({ServableRequest::Latest("servable0")},
+                                         &fake_manager);
 }
 
 TEST(AvailabilityHelpersTest, LatestAndExactVersion) {
-  std::vector<ServableRequest> available_servables_query;
-  available_servables_query.emplace_back(ServableRequest::Latest("servable0"));
-  available_servables_query.emplace_back(
-      ServableRequest::Specific("servable1", 1));
   FakeManager fake_manager;
   fake_manager.set_available_servable_ids({{"servable0", 0}, {"servable1", 1}});
-  WaitUntilServablesAvailable(&fake_manager, available_servables_query);
+  WaitUntilServablesAvailableForRequests(
+      {ServableRequest::Latest("servable0"),
+       ServableRequest::Specific("servable1", 1)},
+      &fake_manager);
 }
 
 }  // namespace

tensorflow_serving/core/dynamic_manager_test.cc

@@ -28,6 +28,7 @@ limitations under the License.
 #include "tensorflow_serving/core/eager_load_policy.h"
 #include "tensorflow_serving/core/servable_state.h"
 #include "tensorflow_serving/core/test_util/dynamic_manager_test_util.h"
+#include "tensorflow_serving/core/test_util/fake_loader.h"
 #include "tensorflow_serving/core/test_util/mock_loader.h"
 #include "tensorflow_serving/util/any_ptr.h"
 #include "tensorflow_serving/util/event_bus.h"
@@ -36,40 +37,14 @@ namespace tensorflow {
 namespace serving {
 
 using ::testing::_;
-using ::testing::Eq;
 using ::testing::Invoke;
 using ::testing::NiceMock;
 using ::testing::Return;
 using ::testing::UnorderedElementsAreArray;
+using test_util::FakeLoader;
 
 namespace {
 
-class FakeLoader : public ResourceUnsafeLoader {
- public:
-  explicit FakeLoader(int64 servable, const bool errors_on_load = false)
-      : servable_(servable), errors_on_load_(errors_on_load) {}
-  ~FakeLoader() override { --num_fake_loaders_; }
-
-  Status Load() override {
-    if (errors_on_load_) {
-      return errors::Internal("Error on load.");
-    } else {
-      return Status::OK();
-    }
-  }
-  void Unload() override {}
-  AnyPtr servable() override { return AnyPtr(&servable_); }
-  static int num_fake_loaders() { return num_fake_loaders_; }
-
- private:
-  // Counts the number of Fakeloader objects alive.
-  static int num_fake_loaders_;
-  int64 servable_;
-  // Whether to return an error when Load is called.
-  const bool errors_on_load_;
-};
-int FakeLoader::num_fake_loaders_ = 0;
-
 constexpr char kServableName[] = "kServableName";
 constexpr char kServableName2[] = "kServableName2";
 constexpr int kNumVersionsPerServable = 2;
@@ -227,7 +202,8 @@ TEST_F(DynamicManagerTest, ListAvailableServableIds) {
   // so never moves to a loaded state.
   std::vector<ServableData<std::unique_ptr<Loader>>> aspired_versions;
   const ServableId id = {kServableName, 7};
-  std::unique_ptr<Loader> loader(new FakeLoader(7, true /* errors_on_load */));
+  std::unique_ptr<Loader> loader(
+      new FakeLoader(7, errors::Internal("Error on load.")));
   aspired_versions.push_back({id, std::move(loader)});
   manager_->GetAspiredVersionsCallback()(kServableName,
                                          std::move(aspired_versions));
@@ -496,7 +472,8 @@ TEST_F(DynamicManagerTest, EventBusErroneousVersion) {
 TEST_F(DynamicManagerTest, EventBusErrorOnLoad) {
   std::vector<ServableData<std::unique_ptr<Loader>>> aspired_versions;
   const ServableId id = {kServableName, 7};
-  std::unique_ptr<Loader> loader(new FakeLoader(7, true /* errors_on_load */));
+  std::unique_ptr<Loader> loader(
+      new FakeLoader(7, errors::Internal("Error on load.")));
   aspired_versions.push_back({id, std::move(loader)});
   manager_->GetAspiredVersionsCallback()(kServableName,
                                          std::move(aspired_versions));
@@ -613,7 +590,8 @@ TEST_F(DynamicManagerTest, RetryOnLoadErrorFinallyFails) {
   std::vector<ServableData<std::unique_ptr<Loader>>> aspired_versions;
   const ServableId id = {kServableName, 7};
   // We always fail.
-  std::unique_ptr<Loader> loader(new FakeLoader(7, true /* errors_on_load */));
+  std::unique_ptr<Loader> loader(
+      new FakeLoader(7, errors::Internal("Error on load.")));
   aspired_versions.push_back({id, std::move(loader)});
   manager_->GetAspiredVersionsCallback()(kServableName,
                                          std::move(aspired_versions));

tensorflow_serving/core/manager.h

@@ -38,6 +38,7 @@ struct ServableRequest {
   // Initialization factories, for convenience and readability.
   static ServableRequest Specific(const string& name, const int64 version);
   static ServableRequest Latest(const string& name);
+  static ServableRequest FromId(const ServableId& id);
 
   // The name of a servable stream.
   string name;
@@ -89,6 +90,11 @@ inline ServableRequest ServableRequest::Latest(const string& name) {
   return ServableRequest{name, nullopt};
 }
 
+inline ServableRequest ServableRequest::FromId(const ServableId& id) {
+  DCHECK_GE(id.version, 0);
+  return Specific(id.name, id.version);
+}
+
 inline string ServableRequest::DebugString() const {
   if (version) {
     return strings::StrCat("Specific(", name, ", ", version.value(), ")");

tensorflow_serving/core/manager_test.cc

@@ -169,6 +169,24 @@ TEST(ServableHandleTest, PointerOps) {
   EXPECT_EQ(&servables[0].member, &handles[0]->member);
 }
 
+TEST(ServableRequestTest, Specific) {
+  const auto request = ServableRequest::Specific("servable", 7);
+  EXPECT_EQ("servable", request.name);
+  EXPECT_EQ(7, *request.version);
+}
+
+TEST(ServableRequestTest, Latest) {
+  const auto request = ServableRequest::Latest("servable");
+  EXPECT_EQ("servable", request.name);
+  EXPECT_FALSE(request.version);
+}
+
+TEST(ServableRequestTest, FromId) {
+  const auto request = ServableRequest::FromId({"servable", 7});
+  EXPECT_EQ("servable", request.name);
+  EXPECT_EQ(7, *request.version);
+}
+
 }  // namespace
 
 }  // namespace serving

tensorflow_serving/core/test_util/BUILD

@@ -49,6 +49,18 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "fake_loader",
+    testonly = 1,
+    srcs = ["fake_loader.cc"],
+    hdrs = ["fake_loader.h"],
+    deps = [
+        "@tf//tensorflow/core:lib",
+        "//tensorflow_serving/core:loader",
+        "//tensorflow_serving/util:any_ptr",
+    ],
+)
+
 cc_library(
     name = "mock_storage_path_target",
     testonly = 1,

tensorflow_serving/core/test_util/fake_loader.cc

@@ -0,0 +1,64 @@
+/* Copyright 2016 Google Inc. 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_serving/core/test_util/fake_loader.h"
+
+#include "tensorflow/core/lib/core/errors.h"
+
+namespace tensorflow {
+namespace serving {
+namespace test_util {
+
+thread_local bool FakeLoader::was_deleted_in_this_thread_;
+int FakeLoader::num_fake_loaders_ = 0;
+mutex FakeLoader::num_fake_loaders_mu_(LINKER_INITIALIZED);
+
+FakeLoader::FakeLoader(int64 servable, const Status load_status)
+    : servable_(servable), load_status_(load_status) {
+  was_deleted_in_this_thread_ = false;
+  {
+    mutex_lock l(num_fake_loaders_mu_);
+    ++num_fake_loaders_;
+  }
+}
+
+FakeLoader::~FakeLoader() {
+  {
+    mutex_lock l(num_fake_loaders_mu_);
+    --num_fake_loaders_;
+  }
+  was_deleted_in_this_thread_ = true;
+}
+
+Status FakeLoader::load_status() { return load_status_; }
+
+Status FakeLoader::Load() { return load_status_; }
+
+void FakeLoader::Unload() {}
+
+AnyPtr FakeLoader::servable() { return AnyPtr(&servable_); }
+
+bool FakeLoader::was_deleted_in_this_thread() {
+  return was_deleted_in_this_thread_;
+}
+
+int FakeLoader::num_fake_loaders() {
+  mutex_lock l(num_fake_loaders_mu_);
+  return num_fake_loaders_;
+}
+
+}  // namespace test_util
+}  // namespace serving
+}  // namespace tensorflow