Merge commit for internal changes

Author: kirilg

tensorflow_serving/batching/BUILD

@@ -53,7 +53,7 @@ cc_library(
     deps = [
         "@tf//tensorflow/core:lib",
         ":batch_scheduler",
-        ":retrier",
+        ":batch_scheduler_retrier",
         "//tensorflow_serving/util:optional",
     ],
 )
@@ -74,8 +74,8 @@ cc_test(
 )
 
 cc_library(
-    name = "retrier",
-    hdrs = ["retrier.h"],
+    name = "batch_scheduler_retrier",
+    hdrs = ["batch_scheduler_retrier.h"],
     visibility = ["//visibility:public"],
     deps = [
         "@tf//tensorflow/core:lib",
@@ -84,15 +84,15 @@ cc_library(
 )
 
 cc_test(
-    name = "retrier_test",
+    name = "batch_scheduler_retrier_test",
     srcs = [
-        "retrier_test.cc",
+        "batch_scheduler_retrier_test.cc",
     ],
     deps = [
         "@tf//tensorflow/core:lib",
         "@tf//tensorflow/core:protos_all_cc",
         "@tf//tensorflow/core:test",
-        ":retrier",
+        ":batch_scheduler_retrier",
         "//tensorflow_serving/core/test_util:test_main",
         "//tensorflow_serving/test_util:fake_clock_env",
     ],

tensorflow_serving/batching/batch_scheduler.h

@@ -16,7 +16,7 @@ limitations under the License.
 // Abstractions for processing small tasks in a batched fashion, to reduce
 // processing times and costs that can be amortized across multiple tasks.
 //
-// The core class is BatchScheduler, which groups Tasks into Batches.
+// The core class is BatchScheduler, which groups tasks into batches.
 //
 // BatchScheduler encapsulates logic for aggregating multiple tasks into a
 // batch, and kicking off processing of a batch on a thread pool it manages.
@@ -42,7 +42,6 @@ limitations under the License.
 
 namespace tensorflow {
 namespace serving {
-namespace batching {
 
 // The abstract superclass for a unit of work to be done as part of a batch.
 //
@@ -53,26 +52,27 @@ namespace batching {
 //  (d) a place to store the output data, upon success.
 //
 // Items (b), (c) and (d) are typically non-owned pointers to data homed
-// elsewhere, because a Task's ownership gets transferred to a BatchScheduler
+// elsewhere, because a task's ownership gets transferred to a BatchScheduler
 // (see below) and it may be deleted as soon as it is done executing.
-class Task {
+class BatchTask {
  public:
-  virtual ~Task() = default;
+  virtual ~BatchTask() = default;
 
   // Returns the size of the task, in terms of how much it contributes to the
   // size of a batch. (A batch's size is the sum of its task sizes.)
   virtual size_t size() const = 0;
 };
 
-// A thread-safe collection of Tasks, to be executed together in some fashion.
+// A thread-safe collection of BatchTasks, to be executed together in some
+// fashion.
 //
 // At a given time, a batch is either "open" or "closed": an open batch can
 // accept new tasks; a closed one cannot. A batch is monotonic: initially it is
 // open and tasks can be added to it; then it is closed and its set of tasks
 // remains fixed for the remainder of its life. A closed batch cannot be re-
 // opened. Tasks can never be removed from a batch.
 //
-// Type parameter TaskType must be a subclass of Task.
+// Type parameter TaskType must be a subclass of BatchTask.
 template <typename TaskType>
 class Batch {
  public:
@@ -127,7 +127,7 @@ class Batch {
 // and processes each batch on a pool of "batch threads" that it manages. The
 // actual logic for processing a batch is accomplished via a callback.
 //
-// Type parameter TaskType must be a subclass of Task.
+// Type parameter TaskType must be a subclass of BatchTask.
 template <typename TaskType>
 class BatchScheduler {
  public:
@@ -252,7 +252,6 @@ void Batch<TaskType>::Close() {
   closed_.Notify();
 }
 
-}  // namespace batching
 }  // namespace serving
 }  // namespace tensorflow
 

tensorflow_serving/batching/retrier.h …rving/batching/batch_scheduler_retrier.htensorflow_serving/batching/retrier.h renamed to tensorflow_serving/batching/batch_scheduler_retrier.h tensorflow_serving/batching/retrier.h renamed to tensorflow_serving/batching/batch_scheduler_retrier.h

@@ -13,8 +13,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#ifndef TENSORFLOW_SERVING_BATCHING_RETRIER_H_
-#define TENSORFLOW_SERVING_BATCHING_RETRIER_H_
+#ifndef TENSORFLOW_SERVING_BATCHING_BATCH_SCHEDULER_RETRIER_H_
+#define TENSORFLOW_SERVING_BATCHING_BATCH_SCHEDULER_RETRIER_H_
 
 #include <stddef.h>
 #include <cstddef>
@@ -27,14 +27,13 @@ limitations under the License.
 
 namespace tensorflow {
 namespace serving {
-namespace batching {
 
 // A wrapper around another BatchScheduler that automatically retries
 // Schedule() requests. Returns an UNAVAILABLE error only after retry attempts
 // have failed (based on parameters that govern the maximum number of retries
 // and the retry time interval).
 template <typename TaskType>
-class Retrier : public BatchScheduler<TaskType> {
+class BatchSchedulerRetrier : public BatchScheduler<TaskType> {
  public:
   struct Options {
     // The maximum amount of time to spend retrying 'wrapped_->Schedule()'
@@ -47,33 +46,33 @@ class Retrier : public BatchScheduler<TaskType> {
     // The environment to use for time and sleeping.
     Env* env = Env::Default();
   };
-  static Status Create(const Options& options,
-                       std::unique_ptr<BatchScheduler<TaskType>> wrapped,
-                       std::unique_ptr<Retrier<TaskType>>* result);
+  static Status Create(
+      const Options& options, std::unique_ptr<BatchScheduler<TaskType>> wrapped,
+      std::unique_ptr<BatchSchedulerRetrier<TaskType>>* result);
 
-  ~Retrier() override = default;
+  ~BatchSchedulerRetrier() override = default;
 
   Status Schedule(std::unique_ptr<TaskType>* task) override;
   size_t NumEnqueuedTasks() const override;
   size_t SchedulingCapacity() const override;
 
  private:
-  Retrier(const Options& options,
-          std::unique_ptr<BatchScheduler<TaskType>> wrapped);
+  BatchSchedulerRetrier(const Options& options,
+                        std::unique_ptr<BatchScheduler<TaskType>> wrapped);
 
   const Options options_;
   std::unique_ptr<BatchScheduler<TaskType>> wrapped_;
 
-  TF_DISALLOW_COPY_AND_ASSIGN(Retrier);
+  TF_DISALLOW_COPY_AND_ASSIGN(BatchSchedulerRetrier);
 };
 
 //////////
 // Implementation details follow. API users need not read.
 
 template <typename TaskType>
-Status Retrier<TaskType>::Create(
+Status BatchSchedulerRetrier<TaskType>::Create(
     const Options& options, std::unique_ptr<BatchScheduler<TaskType>> wrapped,
-    std::unique_ptr<Retrier<TaskType>>* result) {
+    std::unique_ptr<BatchSchedulerRetrier<TaskType>>* result) {
   if (options.max_time_micros < 0) {
     return errors::InvalidArgument("max_time_micros must be non-negative; was ",
                                    options.max_time_micros);
@@ -83,12 +82,13 @@ Status Retrier<TaskType>::Create(
         "retry_delay_micros must be non-negative; was ",
         options.retry_delay_micros);
   }
-  result->reset(new Retrier(options, std::move(wrapped)));
+  result->reset(new BatchSchedulerRetrier(options, std::move(wrapped)));
   return Status::OK();
 }
 
 template <typename TaskType>
-Status Retrier<TaskType>::Schedule(std::unique_ptr<TaskType>* task) {
+Status BatchSchedulerRetrier<TaskType>::Schedule(
+    std::unique_ptr<TaskType>* task) {
   Status status;
 
   const uint64 start_time_micros = options_.env->NowMicros();
@@ -112,22 +112,21 @@ Status Retrier<TaskType>::Schedule(std::unique_ptr<TaskType>* task) {
 }
 
 template <typename TaskType>
-size_t Retrier<TaskType>::NumEnqueuedTasks() const {
+size_t BatchSchedulerRetrier<TaskType>::NumEnqueuedTasks() const {
   return wrapped_->NumEnqueuedTasks();
 }
 
 template <typename TaskType>
-size_t Retrier<TaskType>::SchedulingCapacity() const {
+size_t BatchSchedulerRetrier<TaskType>::SchedulingCapacity() const {
   return wrapped_->SchedulingCapacity();
 }
 
 template <typename TaskType>
-Retrier<TaskType>::Retrier(const Options& options,
-                           std::unique_ptr<BatchScheduler<TaskType>> wrapped)
+BatchSchedulerRetrier<TaskType>::BatchSchedulerRetrier(
+    const Options& options, std::unique_ptr<BatchScheduler<TaskType>> wrapped)
     : options_(options), wrapped_(std::move(wrapped)) {}
 
-}  // namespace batching
 }  // namespace serving
 }  // namespace tensorflow
 
-#endif  // TENSORFLOW_SERVING_BATCHING_RETRIER_H_
+#endif  // TENSORFLOW_SERVING_BATCHING_BATCH_SCHEDULER_RETRIER_H_

…sorflow_serving/batching/retrier_test.cc …batching/batch_scheduler_retrier_test.cctensorflow_serving/batching/retrier_test.cc renamed to tensorflow_serving/batching/batch_scheduler_retrier_test.cc tensorflow_serving/batching/retrier_test.cc renamed to tensorflow_serving/batching/batch_scheduler_retrier_test.cc

@@ -13,7 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include "tensorflow_serving/batching/retrier.h"
+#include "tensorflow_serving/batching/batch_scheduler_retrier.h"
 
 #include <algorithm>
 #include <limits>
@@ -29,10 +29,9 @@ limitations under the License.
 
 namespace tensorflow {
 namespace serving {
-namespace batching {
 namespace {
 
-class FakeTask : public Task {
+class FakeTask : public BatchTask {
  public:
   FakeTask() = default;
   ~FakeTask() override = default;
@@ -101,23 +100,23 @@ class StubbornScheduler : public BatchScheduler<FakeTask> {
   TF_DISALLOW_COPY_AND_ASSIGN(StubbornScheduler);
 };
 
-TEST(RetrierTest, ConstMethodsForwardToWrappedScheduler) {
+TEST(BatchSchedulerRetrierTest, ConstMethodsForwardToWrappedScheduler) {
   auto broken_scheduler = std::unique_ptr<BrokenScheduler>(new BrokenScheduler);
-  Retrier<FakeTask>::Options options;
-  std::unique_ptr<Retrier<FakeTask>> retrier;
-  TF_CHECK_OK(Retrier<FakeTask>::Create(options, std::move(broken_scheduler),
-                                        &retrier));
+  BatchSchedulerRetrier<FakeTask>::Options options;
+  std::unique_ptr<BatchSchedulerRetrier<FakeTask>> retrier;
+  TF_CHECK_OK(BatchSchedulerRetrier<FakeTask>::Create(
+      options, std::move(broken_scheduler), &retrier));
   EXPECT_EQ(7, retrier->NumEnqueuedTasks());
   EXPECT_EQ(42, retrier->SchedulingCapacity());
 }
 
-TEST(RetrierTest, PermanentFailure) {
+TEST(BatchSchedulerRetrierTest, PermanentFailure) {
   auto broken_scheduler = std::unique_ptr<BrokenScheduler>(new BrokenScheduler);
   auto broken_scheduler_ptr = broken_scheduler.get();
-  Retrier<FakeTask>::Options options;
-  std::unique_ptr<Retrier<FakeTask>> retrier;
-  TF_CHECK_OK(Retrier<FakeTask>::Create(options, std::move(broken_scheduler),
-                                        &retrier));
+  BatchSchedulerRetrier<FakeTask>::Options options;
+  std::unique_ptr<BatchSchedulerRetrier<FakeTask>> retrier;
+  TF_CHECK_OK(BatchSchedulerRetrier<FakeTask>::Create(
+      options, std::move(broken_scheduler), &retrier));
   auto task = std::unique_ptr<FakeTask>(new FakeTask);
   Status status = retrier->Schedule(&task);
   ASSERT_FALSE(status.ok());
@@ -126,7 +125,7 @@ TEST(RetrierTest, PermanentFailure) {
   EXPECT_EQ(1, broken_scheduler_ptr->num_submit_calls());
 }
 
-TEST(RetrierTest, MaxTime) {
+TEST(BatchSchedulerRetrierTest, MaxTime) {
   for (int num_attempts_to_succeed = 1; num_attempts_to_succeed < 3;
        ++num_attempts_to_succeed) {
     for (int max_attempts = 1; max_attempts < 5; ++max_attempts) {
@@ -135,12 +134,12 @@ TEST(RetrierTest, MaxTime) {
       auto stubborn_scheduler = std::unique_ptr<StubbornScheduler>(
           new StubbornScheduler(num_attempts_to_succeed));
       auto stubborn_scheduler_ptr = stubborn_scheduler.get();
-      Retrier<FakeTask>::Options options;
+      BatchSchedulerRetrier<FakeTask>::Options options;
       options.retry_delay_micros = 1;
       options.max_time_micros = max_attempts;
       options.env = &env;
-      std::unique_ptr<Retrier<FakeTask>> retrier;
-      TF_CHECK_OK(Retrier<FakeTask>::Create(
+      std::unique_ptr<BatchSchedulerRetrier<FakeTask>> retrier;
+      TF_CHECK_OK(BatchSchedulerRetrier<FakeTask>::Create(
           options, std::move(stubborn_scheduler), &retrier));
 
       const bool expect_success = max_attempts >= num_attempts_to_succeed;
@@ -171,20 +170,20 @@ TEST(RetrierTest, MaxTime) {
   }
 }
 
-TEST(RetrierTest, RetryDelay) {
+TEST(BatchSchedulerRetrierTest, RetryDelay) {
   test_util::FakeClockEnv env(Env::Default());
 
   const int num_attempts_to_succeed = 3;
   auto stubborn_scheduler = std::unique_ptr<StubbornScheduler>(
       new StubbornScheduler(num_attempts_to_succeed));
   auto stubborn_scheduler_ptr = stubborn_scheduler.get();
-  Retrier<FakeTask>::Options options;
+  BatchSchedulerRetrier<FakeTask>::Options options;
   options.retry_delay_micros = 7;
   options.max_time_micros = 100;
   options.env = &env;
-  std::unique_ptr<Retrier<FakeTask>> retrier;
-  TF_CHECK_OK(Retrier<FakeTask>::Create(options, std::move(stubborn_scheduler),
-                                        &retrier));
+  std::unique_ptr<BatchSchedulerRetrier<FakeTask>> retrier;
+  TF_CHECK_OK(BatchSchedulerRetrier<FakeTask>::Create(
+      options, std::move(stubborn_scheduler), &retrier));
 
   Notification done;
   std::unique_ptr<Thread> run_retrier(Env::Default()->StartThread(
@@ -205,6 +204,5 @@ TEST(RetrierTest, RetryDelay) {
 }
 
 }  // namespace
-}  // namespace batching
 }  // namespace serving
 }  // namespace tensorflow

tensorflow_serving/batching/batch_scheduler_test.cc

@@ -21,11 +21,10 @@ limitations under the License.
 
 namespace tensorflow {
 namespace serving {
-namespace batching {
 
 namespace {
 
-class FakeTask : public Task {
+class FakeTask : public BatchTask {
  public:
   explicit FakeTask(size_t size) : size_(size) {}
 
@@ -111,6 +110,5 @@ TEST(BatchTest, DeletionBlocksUntilClosed) {
   deleted.WaitForNotification();
 }
 
-}  // namespace batching
 }  // namespace serving
 }  // namespace tensorflow

tensorflow_serving/batching/streaming_batch_scheduler.cc

@@ -17,7 +17,6 @@ limitations under the License.
 
 namespace tensorflow {
 namespace serving {
-namespace batching {
 
 namespace internal {
 
@@ -86,6 +85,5 @@ void SingleTaskScheduler::ThreadLogic() {
 
 }  // namespace internal
 
-}  // namespace batching
 }  // namespace serving
 }  // namespace tensorflow