Merge commit for internal changes

Author: kirilg

WORKSPACE

@@ -4,7 +4,7 @@ local_repository(
 )
 
 load('//tensorflow/tensorflow:workspace.bzl', 'tf_workspace')
-tf_workspace("tensorflow/")
+tf_workspace("tensorflow/", "@tf")
 
 # ===== gRPC dependencies =====
 

tensorflow

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

tensorflow_serving/core/BUILD

@@ -356,7 +356,10 @@ void DynamicManager::SetAspiredVersions(
           loader = version.ConsumeDataOrDie();
         }
         std::shared_ptr<LoaderHarness> harness =
-            std::make_shared<LoaderHarness>(version.id(), std::move(loader));
+            std::make_shared<LoaderHarness>(
+                version.id(), std::move(loader),
+                LoaderHarness::Options{options_.max_num_load_tries,
+                                       options_.load_retry_interval_micros});
         if (!version.status().ok()) {
           LOG(ERROR) << "Version error: " << version.status().ToString();
           harness->Error(version.status());

tensorflow_serving/core/dynamic_manager.cc

@@ -79,6 +79,13 @@ class DynamicManager final : public Manager,
     // EventBus to publish servable state changes. This is optional, if unset,
     // we don't publish.
     EventBus<ServableState>* servable_event_bus = nullptr;
+
+    // Maximum number of times we try to load a servable before we give up.
+    int max_num_load_tries = 5;
+
+    // The interval, in microseconds, between each servable load retry.
+    // Default: 1min.
+    int64 load_retry_interval_micros = 1 * 60 * 1000 * 1000;
   };
 
   explicit DynamicManager(Options options);

tensorflow_serving/core/dynamic_manager.h

@@ -28,25 +28,25 @@ 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/mock_loader.h"
 #include "tensorflow_serving/util/any_ptr.h"
 #include "tensorflow_serving/util/event_bus.h"
 
 namespace tensorflow {
 namespace serving {
 
 using ::testing::Eq;
+using ::testing::Invoke;
+using ::testing::NiceMock;
+using ::testing::Return;
 using ::testing::UnorderedElementsAreArray;
 
 namespace {
 
 class FakeLoader : public Loader {
  public:
-  explicit FakeLoader(int64 servable) : FakeLoader(servable, false) {}
-
-  FakeLoader(int64 servable, const bool errors_on_load)
-      : servable_(servable), errors_on_load_(errors_on_load) {
-    ++num_fake_loaders_;
-  }
+  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 {
@@ -83,13 +83,15 @@ class DynamicManagerTest : public ::testing::Test {
           LOG(INFO) << "Published state: " << state.DebugString();
           last_published_servable_state_ = state;
         });
-    DynamicManager::Options options;
     // The state manager thread won't be run automatically.
-    options.manage_state_interval_micros = -1;
-    options.env = Env::Default();
-    options.version_policy.reset(new EagerLoadPolicy());
-    options.servable_event_bus = servable_event_bus_.get();
-    manager_.reset(new DynamicManager(std::move(options)));
+    dynamic_manager_options_.manage_state_interval_micros = -1;
+    dynamic_manager_options_.env = Env::Default();
+    dynamic_manager_options_.version_policy.reset(new EagerLoadPolicy());
+    dynamic_manager_options_.servable_event_bus = servable_event_bus_.get();
+    dynamic_manager_options_.max_num_load_tries = 2;
+    dynamic_manager_options_.load_retry_interval_micros = 0;
+    // dynamic_manager_options_.load_retry_interval_micros = 0;
+    manager_.reset(new DynamicManager(std::move(dynamic_manager_options_)));
   }
 
   // Creates an aspired-versions entry with 'id' and a FakeLoader whose servable
@@ -138,6 +140,7 @@ class DynamicManagerTest : public ::testing::Test {
   std::unique_ptr<EventBus<ServableState>::Subscription>
       servable_state_subscription_;
   ServableState last_published_servable_state_;
+  DynamicManager::Options dynamic_manager_options_;
   std::unique_ptr<DynamicManager> manager_;
 };
 
@@ -578,6 +581,51 @@ TEST_F(DynamicManagerTest, NoEventBus) {
                                          std::move(aspired_versions));
 }
 
+TEST_F(DynamicManagerTest, RetryOnLoadErrorFinallySucceeds) {
+  std::vector<ServableData<std::unique_ptr<Loader>>> aspired_versions;
+
+  test_util::MockLoader* loader = new NiceMock<test_util::MockLoader>;
+  // Prevents it being changed without our knowledge.
+  CHECK_EQ(dynamic_manager_options_.max_num_load_tries, 2);
+  // We succeed on the last load, before the manager gives up.
+  EXPECT_CALL(*loader, Load())
+      .WillOnce(Return(errors::Internal("Error on load.")))
+      .WillOnce(Return(Status::OK()));
+
+  const ServableId id = {kServableName, 7};
+  aspired_versions.push_back({id, std::unique_ptr<Loader>(loader)});
+  manager_->GetAspiredVersionsCallback()(kServableName,
+                                         std::move(aspired_versions));
+
+  RunManageState();
+
+  const ServableState available_state = {
+      {kServableName, 7},
+      true,
+      ServableState::ManagerState::kAvailable,
+      Status::OK()};
+  EXPECT_THAT(last_published_servable_state_,
+              EqualsServableState(available_state));
+}
+
+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 */));
+  aspired_versions.push_back({id, std::move(loader)});
+  manager_->GetAspiredVersionsCallback()(kServableName,
+                                         std::move(aspired_versions));
+
+  RunManageState();
+
+  const ServableState error_state = {{kServableName, 7},
+                                     true,
+                                     ServableState::ManagerState::kEnd,
+                                     errors::Internal("Error on load.")};
+  EXPECT_THAT(last_published_servable_state_, EqualsServableState(error_state));
+}
+
 }  // namespace
 }  // namespace serving
 }  // namespace tensorflow

tensorflow_serving/core/dynamic_manager_test.cc

@@ -17,12 +17,19 @@ limitations under the License.
 
 #include <algorithm>
 
+#include "tensorflow/core/platform/env.h"
+
 namespace tensorflow {
 namespace serving {
 
 LoaderHarness::LoaderHarness(const ServableId& id,
                              std::unique_ptr<Loader> loader)
-    : id_(id), loader_(std::move(loader)) {
+    : LoaderHarness(id, std::move(loader), Options()) {}
+
+LoaderHarness::LoaderHarness(const ServableId& id,
+                             std::unique_ptr<Loader> loader,
+                             const Options& options)
+    : id_(id), loader_(std::move(loader)), options_(options) {
   VLOG(1) << "New aspired servable version " << id_;
 }
 
@@ -48,7 +55,26 @@ Status LoaderHarness::Load() {
     state_ = kLoading;
     VLOG(1) << "Loading servable version " << id_;
   }
-  const Status status = loader_->Load();
+
+  const Status status = [&]() {
+    Status load_status;
+    int num_tries = 0;
+    do {
+      if (num_tries > 0) {
+        if (options_.load_retry_interval_micros > 0) {
+          Env::Default()->SleepForMicroseconds(
+              options_.load_retry_interval_micros);
+        }
+        LOG(INFO) << "Retrying load on servable version: " << id_
+                  << " retry: " << num_tries;
+      }
+      load_status = loader_->Load();
+      ++num_tries;
+    } while (is_aspired() && !load_status.ok() &&
+             num_tries < options_.max_num_load_tries);
+    return load_status;
+  }();
+
   {
     mutex_lock l(mu_);
     DCHECK_EQ(kLoading, state_);

tensorflow_serving/core/loader_harness.cc

@@ -72,7 +72,17 @@ class LoaderHarness final {
     kError
   };
 
+  struct Options {
+    // Maximum number of times we try to load a servable before we give up.
+    int max_num_load_tries;
+
+    // The interval, in microseconds, between each servable load retry.
+    int64 load_retry_interval_micros;
+  };
+
   LoaderHarness(const ServableId& id, std::unique_ptr<Loader> loader);
+  LoaderHarness(const ServableId& id, std::unique_ptr<Loader> loader,
+                const Options& options);
 
   // Legal to destruct iff current state is kNew|kDisabled|kError.
   // Check-fails if violated.
@@ -91,9 +101,14 @@ class LoaderHarness final {
   // Returns the current overall state snapshot of the underlying Servable.
   ServableStateSnapshot loader_state_snapshot() const LOCKS_EXCLUDED(mu_);
 
-  // Transitions to kLoading, delegates to Servable::Load(), then
-  // transitions either to kReady if Load() succeeds, or to kError if
-  // Load() fails. This call may take a long time.
+  // Transitions to kLoading, delegates to Servable::Load(), then transitions
+  // either to kReady if Load() succeeds, or to kError if Load() fails. This
+  // call may take a long time.
+  //
+  // We retry the Servable::Load() according to the options set during
+  // construction of this class. We stop retrying and give up if 1. we have
+  // reached max_num_load_tries or, 2. if is_aspired is set to false.
+  //
   // Legal to call iff current state is kNew. Check-fails if violated.
   Status Load() LOCKS_EXCLUDED(mu_);
 
@@ -140,6 +155,7 @@ class LoaderHarness final {
 
   const ServableId id_;
   const std::unique_ptr<Loader> loader_;
+  const Options options_;
   mutable mutex mu_;
   State state_ GUARDED_BY(mu_) = kNew;
   bool is_aspired_ GUARDED_BY(mu_) = true;

tensorflow_serving/core/loader_harness.h

@@ -28,7 +28,9 @@ limitations under the License.
 
 using ::testing::HasSubstr;
 using ::testing::InvokeWithoutArgs;
+using ::testing::InSequence;
 using ::testing::IsNull;
+using ::testing::NiceMock;
 using ::testing::Return;
 using ::testing::ReturnRef;
 using ::testing::StrictMock;
@@ -160,6 +162,60 @@ TEST(ServableVersionHarnessTest, LoadError) {
   EXPECT_EQ(LoaderHarness::kError, harness.state());
 }
 
+TEST(ServableVersionHarnessTest, RetryOnLoadErrorFinallySucceeds) {
+  test_util::MockLoader* loader = new NiceMock<test_util::MockLoader>;
+  LoaderHarness harness(
+      ServableId{"test", 0}, std::unique_ptr<Loader>(loader),
+      {2 /* max_num_load_tries */, 1 /* load_retry_interval_micros */});
+
+  EXPECT_CALL(*loader, Load())
+      .WillOnce(InvokeWithoutArgs(
+          []() { return Status(error::UNKNOWN, "test load error"); }))
+      .WillOnce(InvokeWithoutArgs([]() { return Status::OK(); }));
+  const Status status = harness.Load();
+  TF_EXPECT_OK(status);
+
+  QuiesceAndUnload(&harness);
+}
+
+// Tests cancelling a load by setting is_aspired to false,
+TEST(ServableVersionHarnessTest, RetryOnLoadErrorCancelledLoad) {
+  test_util::MockLoader* loader = new NiceMock<test_util::MockLoader>;
+  LoaderHarness harness(ServableId{"test", 0}, std::unique_ptr<Loader>(loader),
+                        {10 /* max_num_load_tries */, -1});
+
+  Notification load_called;
+  Notification load_should_return;
+  EXPECT_CALL(*loader, Load())
+      .WillOnce(InvokeWithoutArgs([&load_called, &load_should_return]() {
+        load_called.Notify();
+        load_should_return.WaitForNotification();
+        return Status(error::UNKNOWN, "test load error");
+      }))
+      .WillRepeatedly(InvokeWithoutArgs([]() { return Status::OK(); }));
+  std::unique_ptr<Thread> test_thread(
+      Env::Default()->StartThread(ThreadOptions(), "test", [&harness]() {
+        const Status status = harness.Load();
+        EXPECT_THAT(status.error_message(), HasSubstr("test load error"));
+      }));
+  load_called.WaitForNotification();
+  harness.set_is_aspired(false);
+  load_should_return.Notify();
+}
+
+TEST(ServableVersionHarnessTest, RetryOnLoadErrorFinallyFails) {
+  test_util::MockLoader* loader = new NiceMock<test_util::MockLoader>;
+  LoaderHarness harness(ServableId{"test", 0}, std::unique_ptr<Loader>(loader),
+                        {2 /* max_num_load_tries */, -1});
+
+  EXPECT_CALL(*loader, Load())
+      .Times(2)
+      .WillRepeatedly(InvokeWithoutArgs(
+          []() { return Status(error::UNKNOWN, "test load error"); }));
+  const Status status = harness.Load();
+  EXPECT_THAT(status.error_message(), HasSubstr("test load error"));
+}
+
 TEST(LoaderHarnessTest, ExternallySignalledError) {
   LoaderHarness harness(ServableId{"test", 0}, nullptr);
   EXPECT_EQ(LoaderHarness::State::kNew, harness.state());

tensorflow_serving/core/loader_harness_test.cc

@@ -80,8 +80,20 @@ cc_library(
     testonly = 1,
     srcs = ["dynamic_manager_test_util.cc"],
     hdrs = ["dynamic_manager_test_util.h"],
-    visibility = ["//tensorflow_serving:internal"],
     deps = [
         "//tensorflow_serving/core:dynamic_manager",
     ],
 )
+
+cc_library(
+    name = "availability_test_util",
+    testonly = 1,
+    srcs = ["availability_test_util.cc"],
+    hdrs = ["availability_test_util.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "@tf//tensorflow/core:lib",
+        "//external:gtest",
+        "//tensorflow_serving/core:servable_state_monitor",
+    ],
+)