chore: import upstream snapshot with attribution

tensorflow/cc/training/coordinator.cc

@@ -0,0 +1,140 @@
+/* Copyright 2016 The TensorFlow Authors. 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/cc/training/coordinator.h"
+
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "xla/tsl/protobuf/error_codes.pb.h"
+#include "tensorflow/core/framework/cost_graph.pb.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/platform/status.h"
+
+namespace tensorflow {
+
+Coordinator::Coordinator() : Coordinator(std::vector<error::Code>()) {}
+
+Coordinator::Coordinator(const std::vector<error::Code>& clean_stop_errors)
+    : should_stop_(false) {
+  if (clean_stop_errors.empty()) {
+    clean_stop_errors_.insert(error::OUT_OF_RANGE);
+  } else {
+    for (const auto& code : clean_stop_errors) {
+      clean_stop_errors_.insert(static_cast<int>(code));
+    }
+  }
+}
+
+Coordinator::~Coordinator() {
+  RequestStop().IgnoreError();
+  Join().IgnoreError();
+}
+
+absl::Status Coordinator::RegisterRunner(
+    std::unique_ptr<RunnerInterface> runner) {
+  {
+    mutex_lock l(mu_);
+    if (should_stop_) {
+      return absl::Status(absl::StatusCode::kFailedPrecondition,
+                          "The coordinator has been stopped.");
+    }
+  }
+  mutex_lock l(runners_lock_);
+  runners_.push_back(std::move(runner));
+  return absl::OkStatus();
+}
+
+bool Coordinator::AllRunnersStopped() {
+  mutex_lock l(runners_lock_);
+  for (const auto& runner : runners_) {
+    if (runner->IsRunning()) {
+      return false;
+    }
+  }
+  return true;
+}
+
+absl::Status Coordinator::RequestStop() {
+  mutex_lock l(mu_);
+  if (should_stop_) {
+    return absl::Status(absl::StatusCode::kFailedPrecondition,
+                        "The Coordinator is not running.");
+  }
+  should_stop_ = true;
+  wait_for_stop_.notify_all();
+  return absl::OkStatus();
+}
+
+bool Coordinator::ShouldStop() {
+  mutex_lock l(mu_);
+  return should_stop_;
+}
+
+absl::Status Coordinator::Join() {
+  // TODO(yuefengz): deal with stragglers.
+  {
+    mutex_lock l(mu_);
+    if (!should_stop_) {
+      return absl::Status(absl::StatusCode::kFailedPrecondition,
+                          "Joining coordinator without requesting to stop.");
+    }
+  }
+
+  {
+    mutex_lock l(runners_lock_);
+    for (const auto& t : runners_) {
+      ReportStatus(t->Join());
+    }
+    runners_.clear();
+  }
+  return GetStatus();
+}
+
+void Coordinator::ReportStatus(const absl::Status& status) {
+  mutex_lock l(status_lock_);
+  if (status.ok() || !status_.ok() ||
+      clean_stop_errors_.count(static_cast<int>(status.code())) > 0) {
+    return;
+  }
+  status_ = status;
+}
+
+absl::Status Coordinator::GetStatus() {
+  mutex_lock l(status_lock_);
+  return status_;
+}
+
+void Coordinator::WaitForStop() {
+  mutex_lock l(mu_);
+  while (!should_stop_) {
+    wait_for_stop_.wait(l);
+  }
+}
+
+absl::Status Coordinator::ExportCostGraph(CostGraphDef* cost_graph) const {
+  mutex_lock l(runners_lock_);
+  for (auto& t : runners_) {
+    absl::Status s = t->ExportCostGraph(cost_graph);
+    if (!s.ok()) {
+      return s;
+    }
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace tensorflow

tensorflow/cc/training/coordinator.h

@@ -0,0 +1,136 @@
+/* Copyright 2016 The TensorFlow Authors. 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CC_TRAINING_COORDINATOR_H_
+#define TENSORFLOW_CC_TRAINING_COORDINATOR_H_
+
+#include <atomic>
+#include <memory>
+#include <unordered_set>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "xla/tsl/protobuf/error_codes.pb.h"
+#include "tensorflow/core/framework/cost_graph.pb.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/platform/macros.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/protobuf/config.pb.h"
+#include "tensorflow/core/protobuf/error_codes.pb.h"
+#include "tsl/platform/thread_annotations.h"
+
+namespace tensorflow {
+
+/// The abstract interface for runners which must implement the Join and the
+/// IsRunning function.
+class RunnerInterface {
+ public:
+  virtual ~RunnerInterface() = default;
+  virtual absl::Status Join() = 0;
+  virtual absl::Status ExportCostGraph(CostGraphDef* cost_graph) const {
+    return absl::Status(absl::StatusCode::kInvalidArgument,
+                        "No cost model to export.");
+  }
+  /// Returns true iff the runner is running, i.e. if it is trying to populate
+  /// its queue.
+  virtual bool IsRunning() const = 0;
+};
+
+/// Coordinator class manages the termination of a collection of QueueRunners.
+/// Without a coordinator, QueueRunners have to be joined in a specific order;
+/// otherwise the QueueRunner::Join() could sometimes hang. The
+/// Coordinator::RequestStop() plays the key role which notifies all running
+/// threads under a coordinator to stop. This function could be called by any
+/// thread or any client.
+/// Usage, in the client:
+///   Coordinator coord;
+///   std::unique_ptr<QueueRunner> qr(&coord, ...);
+///   qr.Start(session);
+///   coord.RegisterRunner(std::move(qr));
+///   /// do some work
+///   TF_CHECK_OK(coord.Join());
+/// In each thread of QueueRunner, the coordinator needs to be used as:
+///   void Run() {
+///     while (!coord->ShouldStop()) {
+///       /// do some work
+///       if (error) {
+///         coord->RequestStop();
+///         coord->ReportStatus(error_status);
+///       }
+///     }
+///   }
+class Coordinator {
+ public:
+  Coordinator();
+
+  /// Constructor with a list of error codes which would not be taken as errors
+  /// in status reporting.
+  Coordinator(const std::vector<error::Code>& clean_stop_errors);
+
+  /// In the destructor, RequestStop() and Join() would be called.
+  ~Coordinator();
+
+  /// Registers a runner, i.e. a unit of running threads which is usually a
+  /// QueueRunner. It takes the ownership of runner to avoid lifecycle-related
+  /// problems. Note, the coordinator would not start these threads; they are
+  /// supposed to be in running state when they are registered here.
+  absl::Status RegisterRunner(std::unique_ptr<RunnerInterface> runner);
+
+  /// Returns true iff all the registered runners have been stopped.
+  bool AllRunnersStopped();
+
+  /// Requests all running threads to stop.
+  absl::Status RequestStop();
+
+  /// Returns true if its RequestStop() has been called.
+  bool ShouldStop();
+
+  /// Joins all threads, returns OK or the first reported and unexpected status.
+  absl::Status Join();
+
+  /// Reports status to the coordinator. This is usually called by threads.
+  void ReportStatus(const absl::Status& status);
+
+  /// Returns the latest status.
+  absl::Status GetStatus();
+
+  /// Returns immediately if the coordinator is stopped or blocks until
+  /// RequestStop() is called.
+  void WaitForStop();
+
+  // Returns the cost graph from stored run metadata in registered runners.
+  absl::Status ExportCostGraph(CostGraphDef* cost_graph) const;
+
+ private:
+  std::unordered_set<int> clean_stop_errors_;
+  condition_variable wait_for_stop_;
+
+  mutex mu_;
+  bool should_stop_ TF_GUARDED_BY(mu_);
+
+  mutex status_lock_;
+  absl::Status status_ TF_GUARDED_BY(status_lock_);
+
+  mutable mutex runners_lock_;
+  std::vector<std::unique_ptr<RunnerInterface>> runners_
+      TF_GUARDED_BY(runners_lock_);
+
+  Coordinator(const Coordinator&) = delete;
+  void operator=(const Coordinator&) = delete;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_TRAINING_COORDINATOR_H_

tensorflow/cc/training/coordinator_test.cc

@@ -0,0 +1,236 @@
+/* Copyright 2016 The TensorFlow Authors. 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/cc/training/coordinator.h"
+
+#include <atomic>
+#include <functional>
+#include <memory>
+#include <utility>
+
+#include "absl/status/status.h"
+#include "absl/synchronization/notification.h"
+#include "xla/tsl/lib/core/status_test_util.h"
+#include "xla/tsl/protobuf/error_codes.pb.h"
+#include "tensorflow/core/platform/blocking_counter.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/platform/threadpool.h"
+#include "tensorflow/core/protobuf/error_codes.pb.h"
+
+namespace tensorflow {
+namespace {
+
+using error::Code;
+
+void WaitForStopThread(Coordinator* coord, absl::Notification* about_to_wait,
+                       absl::Notification* done) {
+  about_to_wait->Notify();
+  coord->WaitForStop();
+  done->Notify();
+}
+
+TEST(CoordinatorTest, TestStopAndWaitOnStop) {
+  Coordinator coord;
+  EXPECT_EQ(coord.ShouldStop(), false);
+
+  absl::Notification about_to_wait;
+  absl::Notification done;
+  Env::Default()->SchedClosure(
+      std::bind(&WaitForStopThread, &coord, &about_to_wait, &done));
+  about_to_wait.WaitForNotification();
+  Env::Default()->SleepForMicroseconds(1000 * 1000);
+  EXPECT_FALSE(done.HasBeenNotified());
+
+  TF_EXPECT_OK(coord.RequestStop());
+  done.WaitForNotification();
+  EXPECT_TRUE(coord.ShouldStop());
+}
+
+class MockQueueRunner : public RunnerInterface {
+ public:
+  explicit MockQueueRunner(Coordinator* coord) {
+    coord_ = coord;
+    join_counter_ = nullptr;
+    thread_pool_ =
+        std::make_unique<thread::ThreadPool>(Env::Default(), "test-pool", 10);
+    stopped_ = false;
+  }
+
+  MockQueueRunner(Coordinator* coord, int* join_counter)
+      : MockQueueRunner(coord) {
+    join_counter_ = join_counter;
+  }
+
+  void StartCounting(std::atomic<int>* counter, int until,
+                     absl::Notification* start = nullptr) {
+    thread_pool_->Schedule(
+        std::bind(&MockQueueRunner::CountThread, this, counter, until, start));
+  }
+
+  void StartSettingStatus(const absl::Status& status, BlockingCounter* counter,
+                          absl::Notification* start) {
+    thread_pool_->Schedule(std::bind(&MockQueueRunner::SetStatusThread, this,
+                                     status, counter, start));
+  }
+
+  absl::Status Join() override {
+    if (join_counter_ != nullptr) {
+      (*join_counter_)++;
+    }
+    thread_pool_.reset();
+    return status_;
+  }
+
+  absl::Status GetStatus() { return status_; }
+
+  void SetStatus(const absl::Status& status) { status_ = status; }
+
+  bool IsRunning() const override { return !stopped_; };
+
+  void Stop() { stopped_ = true; }
+
+ private:
+  void CountThread(std::atomic<int>* counter, int until,
+                   absl::Notification* start) {
+    if (start != nullptr) start->WaitForNotification();
+    while (!coord_->ShouldStop() && counter->load() < until) {
+      (*counter)++;
+      Env::Default()->SleepForMicroseconds(10 * 1000);
+    }
+    coord_->RequestStop().IgnoreError();
+  }
+  void SetStatusThread(const absl::Status& status, BlockingCounter* counter,
+                       absl::Notification* start) {
+    start->WaitForNotification();
+    SetStatus(status);
+    counter->DecrementCount();
+  }
+  std::unique_ptr<thread::ThreadPool> thread_pool_;
+  absl::Status status_;
+  Coordinator* coord_;
+  int* join_counter_;
+  bool stopped_;
+};
+
+TEST(CoordinatorTest, TestRealStop) {
+  std::atomic<int> counter(0);
+  Coordinator coord;
+
+  std::unique_ptr<MockQueueRunner> qr1 =
+      std::make_unique<MockQueueRunner>(&coord);
+  qr1->StartCounting(&counter, 100);
+  TF_ASSERT_OK(coord.RegisterRunner(std::move(qr1)));
+
+  std::unique_ptr<MockQueueRunner> qr2 =
+      std::make_unique<MockQueueRunner>(&coord);
+  qr2->StartCounting(&counter, 100);
+  TF_ASSERT_OK(coord.RegisterRunner(std::move(qr2)));
+
+  // Wait until the counting has started
+  while (counter.load() == 0)
+    ;
+  TF_EXPECT_OK(coord.RequestStop());
+
+  int temp_counter = counter.load();
+  Env::Default()->SleepForMicroseconds(1000 * 1000);
+  EXPECT_EQ(temp_counter, counter.load());
+  TF_EXPECT_OK(coord.Join());
+}
+
+TEST(CoordinatorTest, TestRequestStop) {
+  Coordinator coord;
+  std::atomic<int> counter(0);
+  absl::Notification start;
+  std::unique_ptr<MockQueueRunner> qr;
+  for (int i = 0; i < 10; i++) {
+    qr = std::make_unique<MockQueueRunner>(&coord);
+    qr->StartCounting(&counter, 10, &start);
+    TF_ASSERT_OK(coord.RegisterRunner(std::move(qr)));
+  }
+  start.Notify();
+
+  coord.WaitForStop();
+  EXPECT_EQ(coord.ShouldStop(), true);
+  EXPECT_EQ(counter.load(), 10);
+  TF_EXPECT_OK(coord.Join());
+}
+
+TEST(CoordinatorTest, TestJoin) {
+  Coordinator coord;
+  int join_counter = 0;
+  std::unique_ptr<MockQueueRunner> qr1 =
+      std::make_unique<MockQueueRunner>(&coord, &join_counter);
+  TF_ASSERT_OK(coord.RegisterRunner(std::move(qr1)));
+  std::unique_ptr<MockQueueRunner> qr2 =
+      std::make_unique<MockQueueRunner>(&coord, &join_counter);
+  TF_ASSERT_OK(coord.RegisterRunner(std::move(qr2)));
+
+  TF_EXPECT_OK(coord.RequestStop());
+  TF_EXPECT_OK(coord.Join());
+  EXPECT_EQ(join_counter, 2);
+}
+
+TEST(CoordinatorTest, StatusReporting) {
+  Coordinator coord({Code::CANCELLED, Code::OUT_OF_RANGE});
+  absl::Notification start;
+  BlockingCounter counter(3);
+
+  std::unique_ptr<MockQueueRunner> qr1 =
+      std::make_unique<MockQueueRunner>(&coord);
+  qr1->StartSettingStatus(absl::Status(absl::StatusCode::kCancelled, ""),
+                          &counter, &start);
+  TF_ASSERT_OK(coord.RegisterRunner(std::move(qr1)));
+
+  std::unique_ptr<MockQueueRunner> qr2 =
+      std::make_unique<MockQueueRunner>(&coord);
+  qr2->StartSettingStatus(absl::Status(absl::StatusCode::kInvalidArgument, ""),
+                          &counter, &start);
+  TF_ASSERT_OK(coord.RegisterRunner(std::move(qr2)));
+
+  std::unique_ptr<MockQueueRunner> qr3 =
+      std::make_unique<MockQueueRunner>(&coord);
+  qr3->StartSettingStatus(absl::Status(absl::StatusCode::kOutOfRange, ""),
+                          &counter, &start);
+  TF_ASSERT_OK(coord.RegisterRunner(std::move(qr3)));
+
+  start.Notify();
+  counter.Wait();
+  TF_EXPECT_OK(coord.RequestStop());
+  EXPECT_EQ(coord.Join().code(), absl::StatusCode::kInvalidArgument);
+}
+
+TEST(CoordinatorTest, JoinWithoutStop) {
+  Coordinator coord;
+  std::unique_ptr<MockQueueRunner> qr =
+      std::make_unique<MockQueueRunner>(&coord);
+  TF_ASSERT_OK(coord.RegisterRunner(std::move(qr)));
+
+  EXPECT_EQ(coord.Join().code(), Code::FAILED_PRECONDITION);
+}
+
+TEST(CoordinatorTest, AllRunnersStopped) {
+  Coordinator coord;
+  MockQueueRunner* qr = new MockQueueRunner(&coord);
+  TF_ASSERT_OK(coord.RegisterRunner(std::unique_ptr<RunnerInterface>(qr)));
+
+  EXPECT_FALSE(coord.AllRunnersStopped());
+  qr->Stop();
+  EXPECT_TRUE(coord.AllRunnersStopped());
+}
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/cc/training/queue_runner.cc

@@ -0,0 +1,263 @@
+/* Copyright 2016 The TensorFlow Authors. 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/cc/training/queue_runner.h"
+
+#include <chrono>
+#include <cstddef>
+#include <cstdint>
+#include <functional>
+#include <limits>
+#include <memory>
+#include <string>
+
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "tensorflow/cc/training/coordinator.h"
+#include "xla/tsl/protobuf/error_codes.pb.h"
+#include "tensorflow/core/framework/cost_graph.pb.h"
+#include "tensorflow/core/framework/ops_util.h"
+#include "tensorflow/core/platform/blocking_counter.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/threadpool.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/protobuf/config.pb.h"
+#include "tensorflow/core/protobuf/queue_runner.pb.h"
+#include "tensorflow/core/public/session.h"
+
+namespace tensorflow {
+
+absl::Status QueueRunner::New(const QueueRunnerDef& queue_runner_def,
+                              std::unique_ptr<QueueRunner>* result) {
+  result->reset(new QueueRunner());
+  return (*result)->Init(queue_runner_def);
+}
+
+absl::Status QueueRunner::New(const QueueRunnerDef& queue_runner_def,
+                              Coordinator* coord,
+                              std::unique_ptr<QueueRunner>* result) {
+  result->reset(new QueueRunner());
+  (*result)->coord_ = coord;
+  return (*result)->Init(queue_runner_def);
+}
+
+void QueueRunner::AddErrorCallback(
+    const std::function<void(absl::Status)>& cb) {
+  mutex_lock l(cb_mu_);
+  callbacks_.push_back(cb);
+}
+
+void QueueRunner::ClearErrorCallbacks() {
+  mutex_lock l(cb_mu_);
+  callbacks_.clear();
+}
+
+absl::Status QueueRunner::Init(const QueueRunnerDef& queue_runner_def) {
+  queue_name_ = queue_runner_def.queue_name();
+  enqueue_op_names_.clear();
+  enqueue_op_names_.insert(enqueue_op_names_.end(),
+                           queue_runner_def.enqueue_op_name().begin(),
+                           queue_runner_def.enqueue_op_name().end());
+  size_t op_names_size = enqueue_op_names_.size();
+  if (op_names_size > std::numeric_limits<int32_t>::max()) {
+    return absl::Status(absl::StatusCode::kInvalidArgument,
+                        "Enqueue ops to run cannot exceed kint32max");
+  }
+  runs_ = static_cast<int>(op_names_size);
+  if (runs_ == 0) {
+    return absl::Status(absl::StatusCode::kInvalidArgument,
+                        "Empty enqueue ops to run.");
+  }
+  close_op_name_ = queue_runner_def.close_op_name();
+  cancel_op_name_ = queue_runner_def.cancel_op_name();
+  if (queue_runner_def.queue_closed_exception_types_size() == 0) {
+    queue_closed_exception_types_.insert(error::OUT_OF_RANGE);
+  } else {
+    for (const auto& code : queue_runner_def.queue_closed_exception_types()) {
+      queue_closed_exception_types_.insert(static_cast<int>(code));
+    }
+  }
+
+  int nthreads = runs_;
+  if (coord_) {
+    // One more thread to call Stop()
+    nthreads++;
+  }
+  thread_pool_ = std::make_unique<thread::ThreadPool>(
+      Env::Default(), SanitizeThreadSuffix(queue_name_), nthreads);
+
+  return absl::OkStatus();
+}
+
+QueueRunner::~QueueRunner() {
+  // Cannot run Stop() here because the session might already be closed or
+  // destroyed.
+  Join().IgnoreError();
+}
+
+absl::Status QueueRunner::Start(Session* sess) { return Start(sess, 0); }
+
+absl::Status QueueRunner::StartAndCollectCostGraph(
+    Session* sess, const RunOptions& run_options) {
+  SetRunArgumentsAndCostGraph(run_options);
+  return Start(sess, 0);
+}
+
+absl::Status QueueRunner::Start(Session* sess, int wait_for) {
+  counter_ = std::make_unique<BlockingCounter>(runs_);
+  for (const std::string& enqueue_op : enqueue_op_names_) {
+    thread_pool_->Schedule(
+        std::bind(&QueueRunner::Run, this, sess, enqueue_op));
+  }
+  if (coord_) {
+    thread_pool_->Schedule(std::bind(&QueueRunner::Stop, this, sess));
+  }
+  // Wait for up to 'wait_for' milliseconds.
+  if (wait_for > 0) {
+    if (!counter_->WaitFor(std::chrono::milliseconds(wait_for))) {
+      return absl::Status(absl::StatusCode::kDeadlineExceeded,
+                          "Queues not fed before the timeout");
+    }
+    // Check the status of the queue runner as well as the result of the enqueue
+    // operations.
+    mutex_lock l(mu_);
+    if (!enqueue_status_.ok()) {
+      return enqueue_status_;
+    } else {
+      return status_;
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status QueueRunner::StartAndCollectCostGraph(
+    Session* session, int wait_for_ms, const RunOptions& run_options) {
+  SetRunArgumentsAndCostGraph(run_options);
+  return Start(session, wait_for_ms);
+}
+
+void QueueRunner::Stop(Session* sess) {
+  if (coord_ != nullptr) {
+    coord_->WaitForStop();
+  }
+  if (!cancel_op_name_.empty()) {
+    UpdateStatus(RealRun(sess, cancel_op_name_, false));
+  }
+  stopped_ = true;
+}
+
+absl::Status QueueRunner::Join() {
+  thread_pool_.reset();
+  mutex_lock l(mu_);
+  return status_;
+}
+
+void QueueRunner::UpdateStatus(const absl::Status& status) {
+  {
+    mutex_lock l(mu_);
+    if (!status_.ok() || status.ok() || IsQueueClosed(status)) {
+      return;
+    }
+    status_ = status;
+  }
+  if (coord_) {
+    coord_->ReportStatus(status);
+  }
+  mutex_lock l(cb_mu_);
+  for (auto& cb : callbacks_) {
+    cb(status);
+  }
+}
+
+void QueueRunner::Run(Session* sess, const std::string& enqueue_op) {
+  bool first_iteration = true;
+  absl::Status status;
+  while (status.ok()) {
+    if (coord_ && coord_->ShouldStop()) {
+      break;
+    }
+    status = RealRun(sess, enqueue_op, true);
+    if (first_iteration) {
+      if (!status.ok()) {
+        mutex_lock l(mu_);
+        enqueue_status_ = status;
+      }
+      counter_->DecrementCount();
+      first_iteration = false;
+    }
+  }
+  bool last_run = false;
+  {
+    mutex_lock l(mu_);
+    runs_--;
+    last_run = (runs_ == 0);
+  }
+
+  // Close the queue unless the coordinator is shutting down since the cancel op
+  // will be run anyway in this case.
+  if (IsQueueClosed(status) && (!coord_ || !coord_->ShouldStop())) {
+    if (last_run && !close_op_name_.empty()) {
+      UpdateStatus(RealRun(sess, close_op_name_, false));
+    }
+  } else if (!status.ok()) {
+    LOG(ERROR) << "Queue runner thread got a failure status: " << status;
+    UpdateStatus(status);
+    if (coord_) {
+      coord_->RequestStop().IgnoreError();
+    }
+  }
+}
+
+absl::Status QueueRunner::GetStatus() {
+  mutex_lock l(mu_);
+  return status_;
+}
+
+absl::Status QueueRunner::ExportCostGraph(CostGraphDef* cost_graph) const {
+  if (!cg_mu_) {
+    return absl::Status(absl::StatusCode::kFailedPrecondition,
+                        "This QueueRunner doesn't collect a cost graph.");
+  }
+  mutex_lock l(*cg_mu_);
+  cost_graph->MergeFrom(*cost_graph_);
+  return absl::OkStatus();
+}
+
+void QueueRunner::SetRunArgumentsAndCostGraph(const RunOptions& run_options) {
+  cg_mu_ = std::make_unique<mutex>();
+  {
+    mutex_lock l(*cg_mu_);
+    cost_graph_ = std::make_unique<CostGraphDef>();
+  }
+  run_options_ = run_options;
+}
+
+absl::Status QueueRunner::RealRun(Session* sess, const std::string& op,
+                                  bool update_costs) {
+  absl::Status s;
+  if (update_costs && cg_mu_) {
+    RunMetadata metadata;
+    s = sess->Run(run_options_, {}, {}, {op}, nullptr, &metadata);
+    mutex_lock l(*cg_mu_);
+    cost_graph_->Swap(metadata.mutable_cost_graph());
+  } else {
+    s = sess->Run({}, {}, {op}, nullptr);
+  }
+  return s;
+}
+
+}  // namespace tensorflow

tensorflow/cc/training/queue_runner.h

@@ -0,0 +1,145 @@
+/* Copyright 2016 The TensorFlow Authors. 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CC_TRAINING_QUEUE_RUNNER_H_
+#define TENSORFLOW_CC_TRAINING_QUEUE_RUNNER_H_
+
+#include <memory>
+#include <string>
+#include <unordered_set>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "tensorflow/cc/training/coordinator.h"
+#include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/lib/core/threadpool.h"
+#include "tensorflow/core/platform/blocking_counter.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/protobuf/config.pb.h"
+#include "tensorflow/core/protobuf/error_codes.pb.h"
+#include "tensorflow/core/protobuf/queue_runner.pb.h"
+#include "tensorflow/core/public/session.h"
+#include "tsl/platform/thread_annotations.h"
+
+namespace tensorflow {
+
+/// QueueRunner class imitates the behavior of the python version of QueueRunner
+/// which creates a thread for each enqueue op, runs close op on completion.
+class QueueRunner : public RunnerInterface {
+ public:
+  /// Creates a new QueueRunner from proto.
+  // TODO(yuefengz): we may want to initialize from queues and ops in the
+  // future.
+  static absl::Status New(const QueueRunnerDef& queue_runner_def,
+                          std::unique_ptr<QueueRunner>* result);
+
+  /// Creates a new QueueRunner with a coordinator, see coordinator.h for usage.
+  static absl::Status New(const QueueRunnerDef& queue_runner_def,
+                          Coordinator* coord,
+                          std::unique_ptr<QueueRunner>* result);
+
+  /// Adds a callback that the queue runner will call when it detects an error.
+  void AddErrorCallback(const std::function<void(absl::Status)>& cb);
+
+  /// Delete the previously registered callbacks.
+  void ClearErrorCallbacks();
+
+  /// The destructor would join all the threads.
+  ~QueueRunner() override;
+
+  /// Starts the queue runner with the given session.
+  absl::Status Start(Session* sess);
+
+  /// Starts the queue runner with the given session and sets the run arguments
+  /// for sess->Run. It also collects and stores the cost model.
+  absl::Status StartAndCollectCostGraph(
+      Session* sess, const RunOptions& run_options = RunOptions());
+
+  /// Starts the queue runner with the given session, and wait for up to the
+  /// specified time (in milliseconds) for the queues to start to fill up.
+  absl::Status Start(Session* sess, int wait_for_ms);
+  absl::Status StartAndCollectCostGraph(
+      Session* session, int wait_for_ms,
+      const RunOptions& run_options = RunOptions());
+
+  /// Requests to stop and runs the cancel op. It would be called in a separate
+  /// thread when coordinator is set. If there is no coordinator it should be
+  /// called before calling Join.
+  void Stop(Session* sess);
+
+  /// Joins all the threads. Returns okay if all threads run successfully;
+  /// otherwise returns the first captured failure status.
+  absl::Status Join() final;
+
+  /// Returns the latest status.
+  absl::Status GetStatus();
+
+  // Returns the stored cost model.
+  absl::Status ExportCostGraph(CostGraphDef* cost_graph) const override;
+
+ private:
+  QueueRunner() : coord_(nullptr), stopped_(false), cg_mu_(nullptr) {}
+
+  // Initializes the instance with the QueueRunnerDef proto.
+  absl::Status Init(const QueueRunnerDef& queue_runner_def);
+
+  // The Run function for each thread.
+  void Run(Session* sess, const std::string& enqueue_op);
+
+  // Updates the internal status; it only keeps OK or the first unexpected error
+  // status.
+  void UpdateStatus(const absl::Status& status);
+
+  bool IsQueueClosed(absl::Status status) const {
+    return queue_closed_exception_types_.count(
+               static_cast<int>(status.code())) > 0;
+  }
+
+  bool IsRunning() const override { return !stopped_; }
+
+  void SetRunArgumentsAndCostGraph(const RunOptions& run_options);
+
+  absl::Status RealRun(Session* sess, const std::string& op, bool update_costs);
+
+  std::string queue_name_;
+  std::vector<std::string> enqueue_op_names_;
+  std::string close_op_name_;
+  std::string cancel_op_name_;
+  // code::Code casted to int to avoid a hash function.
+  std::unordered_set<int> queue_closed_exception_types_;
+
+  std::unique_ptr<thread::ThreadPool> thread_pool_;
+  mutex mu_;
+  int runs_ = 0;
+  absl::Status status_ TF_GUARDED_BY(mu_);
+  absl::Status enqueue_status_ TF_GUARDED_BY(mu_);
+  std::unique_ptr<BlockingCounter> counter_;
+
+  Coordinator* coord_;
+
+  std::atomic<bool> stopped_;
+
+  mutex cb_mu_;
+  std::vector<std::function<void(absl::Status)>> callbacks_;
+
+  mutable std::unique_ptr<mutex> cg_mu_;
+  std::unique_ptr<CostGraphDef> cost_graph_ TF_GUARDED_BY(cg_mu_);
+  RunOptions run_options_;
+};
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CC_TRAINING_QUEUE_RUNNER_H_

tensorflow/cc/training/queue_runner_test.cc

@@ -0,0 +1,426 @@
+/* Copyright 2016 The TensorFlow Authors. 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/cc/training/queue_runner.h"
+
+#include <functional>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/synchronization/notification.h"
+#include "tensorflow/cc/framework/ops.h"
+#include "tensorflow/cc/framework/scope.h"
+#include "tensorflow/cc/ops/const_op.h"
+#include "tensorflow/cc/ops/data_flow_ops.h"
+#include "tensorflow/cc/ops/math_ops.h"
+#include "tensorflow/cc/ops/random_ops.h"
+#include "tensorflow/cc/ops/state_ops.h"
+#include "tensorflow/cc/training/coordinator.h"
+#include "xla/tsl/lib/core/status_test_util.h"
+#include "xla/tsl/platform/status.h"
+#include "xla/tsl/protobuf/error_codes.pb.h"
+#include "tensorflow/core/framework/cost_graph.pb.h"
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/protobuf/config.pb.h"
+#include "tensorflow/core/protobuf/error_codes.pb.h"
+#include "tensorflow/core/protobuf/queue_runner.pb.h"
+#include "tensorflow/core/public/session.h"
+#include "tensorflow/core/public/session_options.h"
+
+namespace tensorflow {
+namespace {
+
+using error::Code;
+using ops::Assign;
+using ops::Const;
+using ops::CountUpTo;
+using ops::FIFOQueue;
+using ops::QueueClose;
+using ops::QueueDequeue;
+using ops::QueueEnqueue;
+using ops::RandomNormal;
+using ops::Square;
+using ops::Variable;
+
+constexpr char kAssignOpName[] = "assign";
+constexpr char kCancelOp0[] = "cancel0";
+constexpr char kCancelOp1[] = "cancel1";
+constexpr char kCloseOp0[] = "close0";
+constexpr char kCloseOp1[] = "close1";
+constexpr char kCountUpToOpName[] = "count";
+constexpr char kDequeueOp0[] = "dequeue0";
+constexpr char kDequeueOp1[] = "dequeue1";
+constexpr char kEnqueueOp0[] = "enqueue0";
+constexpr char kEnqueueOp1[] = "enqueue1";
+constexpr char kIllegalOpName1[] = "would fail";
+constexpr char kIllegalOpName2[] = "fail again";
+constexpr char kQueueName[] = "unit_test";
+constexpr char kQueueName0[] = "q0";
+constexpr char kQueueName1[] = "q1";
+constexpr char kSquareOpName[] = "square";
+constexpr char kVarOpName[] = "var";
+
+GraphDef BuildSimpleGraph() {
+  Scope root = Scope::NewRootScope();
+  auto init_value = Const(root, 0);
+  auto var = Variable(root.WithOpName(kVarOpName), TensorShape({}),
+                      DataType::DT_INT32);
+  auto assign = Assign(root.WithOpName(kAssignOpName), var, init_value);
+  auto count = CountUpTo(root.WithOpName(kCountUpToOpName), var, 10);
+  Square(root.WithOpName(kSquareOpName), var);  // NOLINT
+
+  GraphDef graph_def;
+  TF_EXPECT_OK(root.ToGraphDef(&graph_def));
+  return graph_def;
+}
+
+QueueRunnerDef BuildQueueRunnerDef(
+    const std::string& queue_name, const std::vector<std::string>& enqueue_ops,
+    const std::string& close_op, const std::string& cancel_op,
+    const std::vector<Code>& queue_closed_error_codes) {
+  QueueRunnerDef queue_runner_def;
+  *queue_runner_def.mutable_queue_name() = queue_name;
+  for (const std::string& enqueue_op : enqueue_ops) {
+    *queue_runner_def.mutable_enqueue_op_name()->Add() = enqueue_op;
+  }
+  *queue_runner_def.mutable_close_op_name() = close_op;
+  *queue_runner_def.mutable_cancel_op_name() = cancel_op;
+  for (const auto& error_code : queue_closed_error_codes) {
+    *queue_runner_def.mutable_queue_closed_exception_types()->Add() =
+        error_code;
+  }
+  return queue_runner_def;
+}
+
+std::unique_ptr<Session> BuildSessionAndInitVariable(
+    const GraphDef& graph_def) {
+  SessionOptions options;
+  std::unique_ptr<Session> session(NewSession(options));
+  TF_CHECK_OK(session->Create(graph_def));
+
+  TF_CHECK_OK(session->Run({}, {}, {kAssignOpName}, nullptr));
+  return session;
+}
+
+TEST(QueueRunnerTest, BasicTest) {
+  GraphDef graph_def = BuildSimpleGraph();
+  auto session = BuildSessionAndInitVariable(graph_def);
+
+  QueueRunnerDef queue_runner_def = BuildQueueRunnerDef(
+      kQueueName, {kCountUpToOpName}, kSquareOpName, "", {});
+
+  std::unique_ptr<QueueRunner> qr;
+  TF_EXPECT_OK(QueueRunner::New(queue_runner_def, &qr));
+  TF_CHECK_OK(qr->Start(session.get()));
+  TF_EXPECT_OK(qr->Join());
+
+  std::vector<Tensor> outputs;
+  TF_EXPECT_OK(session->Run({}, {kSquareOpName}, {}, &outputs));
+  int square_value = *outputs[0].scalar<int>().data();
+  EXPECT_EQ(square_value, 100);
+}
+
+TEST(QueueRunnerTest, QueueClosedCode) {
+  GraphDef graph_def = BuildSimpleGraph();
+  auto session = BuildSessionAndInitVariable(graph_def);
+
+  // Start two queues so that multiple threads are in Run.
+  QueueRunnerDef queue_runner_def = BuildQueueRunnerDef(
+      kQueueName, {kCountUpToOpName, kCountUpToOpName}, kSquareOpName, "",
+      {Code::OUT_OF_RANGE, Code::CANCELLED});
+
+  std::unique_ptr<QueueRunner> qr;
+  TF_EXPECT_OK(QueueRunner::New(queue_runner_def, &qr));
+  TF_EXPECT_OK(qr->Start(session.get()));
+  TF_EXPECT_OK(qr->Join());
+
+  std::vector<Tensor> outputs;
+  TF_EXPECT_OK(session->Run({}, {kSquareOpName}, {}, &outputs));
+  int square_value = *outputs[0].scalar<int>().data();
+  EXPECT_EQ(square_value, 100);
+}
+
+TEST(QueueRunnerTest, QueueCloseFails) {
+  GraphDef graph_def = BuildSimpleGraph();
+  auto session = BuildSessionAndInitVariable(graph_def);
+
+  QueueRunnerDef queue_runner_def =
+      BuildQueueRunnerDef(kQueueName, {kCountUpToOpName}, kIllegalOpName1, "",
+                          {Code::OUT_OF_RANGE});
+
+  std::unique_ptr<QueueRunner> qr;
+  TF_EXPECT_OK(QueueRunner::New(queue_runner_def, &qr));
+  TF_EXPECT_OK(qr->Start(session.get()));
+  auto status = qr->Join();
+  EXPECT_EQ(status.code(), Code::NOT_FOUND) << status;
+}
+
+TEST(QueueRunnerTest, CatchErrorInJoin) {
+  GraphDef graph_def = BuildSimpleGraph();
+  auto session = BuildSessionAndInitVariable(graph_def);
+
+  QueueRunnerDef queue_runner_def = BuildQueueRunnerDef(
+      kQueueName, {kIllegalOpName1, kIllegalOpName2}, kCountUpToOpName, "", {});
+
+  std::unique_ptr<QueueRunner> qr;
+  TF_EXPECT_OK(QueueRunner::New(queue_runner_def, &qr));
+  TF_EXPECT_OK(qr->Start(session.get()));
+  EXPECT_EQ(qr->Join().code(), Code::NOT_FOUND);
+}
+
+GraphDef BuildDoubleQueueGraph() {
+  Scope root = Scope::NewRootScope();
+  auto q0 = FIFOQueue(root.WithOpName(kQueueName0), {DataType::DT_INT32});
+  auto ten = Const(root, 10);
+  auto enqueue0 = QueueEnqueue(root.WithOpName(kEnqueueOp0), q0, {ten});
+  auto close0 = QueueClose(root.WithOpName(kCloseOp0), q0);
+  auto cancel0 = QueueClose(root.WithOpName(kCancelOp0), q0,
+                            QueueClose::CancelPendingEnqueues(true));
+  auto q1 = FIFOQueue(root.WithOpName(kQueueName1), {DataType::DT_INT32},
+                      FIFOQueue::Capacity(3));
+  auto dequeue0 =
+      QueueDequeue(root.WithOpName(kDequeueOp0), q0, {DataType::DT_INT32});
+  auto enqueue1 = QueueEnqueue(root.WithOpName(kEnqueueOp1), q1, {dequeue0[0]});
+  auto dequeue1 =
+      QueueDequeue(root.WithOpName(kDequeueOp1), q1, {DataType::DT_INT32});
+  auto close1 = QueueClose(root.WithOpName(kCloseOp1), q1);
+  auto cancel1 = QueueClose(root.WithOpName(kCancelOp1), q1,
+                            QueueClose::CancelPendingEnqueues(true));
+
+  GraphDef graph_def;
+  TF_EXPECT_OK(root.ToGraphDef(&graph_def));
+  return graph_def;
+}
+
+TEST(QueueRunnerTest, RealEnqueueDequeue) {
+  auto graph_def = BuildDoubleQueueGraph();
+
+  SessionOptions options;
+  std::unique_ptr<Session> session(NewSession(options));
+  TF_CHECK_OK(session->Create(graph_def));
+
+  QueueRunnerDef queue_runner_def =
+      BuildQueueRunnerDef(kQueueName, {kEnqueueOp1}, kCloseOp1, "", {});
+  std::unique_ptr<QueueRunner> qr;
+  TF_EXPECT_OK(QueueRunner::New(queue_runner_def, &qr));
+  TF_CHECK_OK(qr->Start(session.get()));
+
+  TF_EXPECT_OK(session->Run({}, {}, {kEnqueueOp0}, nullptr));
+  TF_EXPECT_OK(session->Run({}, {}, {kEnqueueOp0}, nullptr));
+  // Closing queue 0 would also close the queue runner.
+  TF_EXPECT_OK(session->Run({}, {}, {kCloseOp0}, nullptr));
+
+  TF_EXPECT_OK(qr->Join());
+  std::vector<Tensor> dq1;
+  TF_EXPECT_OK(session->Run({}, {kDequeueOp1}, {}, &dq1));
+  EXPECT_EQ(*dq1[0].scalar<int>().data(), 10);
+  std::vector<Tensor> dq2;
+  TF_EXPECT_OK(session->Run({}, {kDequeueOp1}, {}, &dq2));
+  EXPECT_EQ(*dq2[0].scalar<int>().data(), 10);
+
+  EXPECT_EQ(session->Run({}, {kDequeueOp1}, {}, nullptr).code(),
+            Code::OUT_OF_RANGE);
+}
+
+void JoinThread(QueueRunner* queue_runner, bool* join_succeeded,
+                absl::Notification* join_done) {
+  EXPECT_EQ(queue_runner->Join().code(), Code::CANCELLED);
+  *join_succeeded = true;
+  join_done->Notify();
+}
+
+TEST(QueueRunnerTest, SessionCloseCancelPendingEnqueue) {
+  auto graph_def = BuildDoubleQueueGraph();
+
+  SessionOptions options;
+  std::unique_ptr<Session> session(NewSession(options));
+  TF_CHECK_OK(session->Create(graph_def));
+
+  QueueRunnerDef queue_runner_def = BuildQueueRunnerDef(
+      kQueueName1, {kEnqueueOp1}, kCloseOp1, kCancelOp1, {});
+  std::unique_ptr<QueueRunner> qr;
+  TF_EXPECT_OK(QueueRunner::New(queue_runner_def, &qr));
+  TF_CHECK_OK(qr->Start(session.get()));
+
+  TF_EXPECT_OK(session->Run({}, {}, {kEnqueueOp0}, nullptr));
+
+  std::vector<Tensor> dq1;
+  TF_EXPECT_OK(session->Run({}, {kDequeueOp1}, {}, &dq1));
+  EXPECT_EQ(*dq1[0].scalar<int>().data(), 10);
+
+  // The expected behavior is the QueueRunner::Join() call is blocked until
+  // Session::Close() is called.
+  bool join_succeeded = false;
+  absl::Notification join_done;
+  Env::Default()->SchedClosure(
+      std::bind(&JoinThread, qr.get(), &join_succeeded, &join_done));
+
+  Env::Default()->SleepForMicroseconds(10000000);
+  EXPECT_EQ(join_succeeded, false);
+
+  // Closing the session is required to cancel pending enqueue nodes.
+  TF_EXPECT_OK(session->Close());
+
+  join_done.WaitForNotification();
+  EXPECT_EQ(join_succeeded, true);
+}
+
+TEST(QueueRunnerTest, EmptyEnqueueOps) {
+  QueueRunnerDef queue_runner_def =
+      BuildQueueRunnerDef(kQueueName, {}, kCountUpToOpName, "", {});
+
+  std::unique_ptr<QueueRunner> qr;
+  EXPECT_EQ(QueueRunner::New(queue_runner_def, &qr).code(),
+            Code::INVALID_ARGUMENT);
+}
+
+TEST(QueueRunnerTest, StartTimeout) {
+  GraphDef graph_def = BuildDoubleQueueGraph();
+  SessionOptions options;
+  std::unique_ptr<Session> session(NewSession(options));
+  TF_CHECK_OK(session->Create(graph_def));
+
+  QueueRunnerDef queue_runner_def = BuildQueueRunnerDef(
+      kQueueName1, {kEnqueueOp1}, kCloseOp1, kCancelOp1, {});
+
+  std::unique_ptr<QueueRunner> qr;
+  TF_EXPECT_OK(QueueRunner::New(queue_runner_def, &qr));
+  // This will timeout since queue0 is not fed and queue1 is fetching data from
+  // queue0.
+  EXPECT_EQ(qr->Start(session.get(), 1).code(), Code::DEADLINE_EXCEEDED);
+  TF_EXPECT_OK(session->Close());
+}
+
+TEST(QueueRunnerTest, TestCoordinatorStop) {
+  auto graph_def = BuildDoubleQueueGraph();
+  SessionOptions options;
+  std::unique_ptr<Session> session(NewSession(options));
+  TF_CHECK_OK(session->Create(graph_def));
+
+  QueueRunnerDef queue_runner0 =
+      BuildQueueRunnerDef(kQueueName0, {kEnqueueOp0}, kCloseOp0, kCancelOp0,
+                          {Code::OUT_OF_RANGE, Code::CANCELLED});
+  QueueRunnerDef queue_runner1 =
+      BuildQueueRunnerDef(kQueueName1, {kEnqueueOp1}, kCloseOp1, kCancelOp1,
+                          {Code::OUT_OF_RANGE, Code::CANCELLED});
+
+  Coordinator coord;
+  std::unique_ptr<QueueRunner> qr0;
+  TF_EXPECT_OK(QueueRunner::New(queue_runner0, &coord, &qr0));
+  TF_CHECK_OK(qr0->Start(session.get()));
+  std::unique_ptr<QueueRunner> qr1;
+  TF_EXPECT_OK(QueueRunner::New(queue_runner1, &coord, &qr1));
+  TF_CHECK_OK(qr1->Start(session.get()));
+
+  TF_EXPECT_OK(coord.RegisterRunner(std::move(qr0)));
+  TF_EXPECT_OK(coord.RegisterRunner(std::move(qr1)));
+
+  std::vector<Tensor> dq;
+  TF_EXPECT_OK(session->Run({}, {kDequeueOp1}, {}, &dq));
+  EXPECT_EQ(*dq[0].scalar<int>().data(), 10);
+
+  TF_EXPECT_OK(coord.RequestStop());
+  TF_EXPECT_OK(coord.Join());
+}
+
+TEST(QueueRunnerTest, CallbackCalledOnError) {
+  GraphDef graph_def = BuildSimpleGraph();
+  auto session = BuildSessionAndInitVariable(graph_def);
+
+  QueueRunnerDef queue_runner_def = BuildQueueRunnerDef(
+      kQueueName, {kIllegalOpName1, kIllegalOpName2}, kCountUpToOpName, "", {});
+
+  std::unique_ptr<QueueRunner> qr;
+  TF_EXPECT_OK(QueueRunner::New(queue_runner_def, &qr));
+  bool error_caught = false;
+  qr->AddErrorCallback(
+      [&error_caught](const absl::Status&) { error_caught = true; });
+  TF_EXPECT_OK(qr->Start(session.get()));
+  EXPECT_FALSE(qr->Join().ok());
+  EXPECT_TRUE(error_caught);
+}
+
+TEST(QueueRunnerTest, RunMetaDataTest) {
+  Scope root = Scope::NewRootScope();
+  auto q0 = FIFOQueue(root.WithOpName(kQueueName), {DataType::DT_FLOAT});
+  Output rnd = RandomNormal(root.WithOpName("rnd"), {1, 1}, DataType::DT_FLOAT);
+  Output square = Square(root.WithOpName(kSquareOpName), rnd);
+  auto enqueue0 = QueueEnqueue(root.WithOpName(kEnqueueOp0), q0, {square});
+  auto close0 = QueueClose(root.WithOpName(kCloseOp0), q0);
+  auto cancel0 = QueueClose(root.WithOpName(kCancelOp0), q0,
+                            QueueClose::CancelPendingEnqueues(true));
+  auto dequeue0 =
+      QueueDequeue(root.WithOpName(kDequeueOp0), q0, {DataType::DT_FLOAT});
+
+  GraphDef graph_def;
+  TF_EXPECT_OK(root.ToGraphDef(&graph_def));
+  for (auto& node : *graph_def.mutable_node()) {
+    node.set_device("/cpu:0");
+  }
+  SessionOptions sess_options;
+  sess_options.config.mutable_graph_options()->set_build_cost_model(1);
+  std::unique_ptr<Session> session(NewSession(sess_options));
+
+  TF_CHECK_OK(session->Create(graph_def));
+
+  QueueRunnerDef queue_runner_def =
+      BuildQueueRunnerDef(kQueueName, {kEnqueueOp0}, kCloseOp0, kCancelOp0, {});
+  std::unique_ptr<QueueRunner> qr;
+  TF_EXPECT_OK(QueueRunner::New(queue_runner_def, &qr));
+  RunOptions run_options;
+  TF_CHECK_OK(qr->StartAndCollectCostGraph(session.get(), run_options));
+
+  // Make sure there was at least one element enqueued in q0: this prevents a
+  // race condition where we close the queue before it was populated.
+  std::vector<Tensor> dq0;
+  TF_EXPECT_OK(session->Run({}, {kDequeueOp0}, {}, &dq0));
+  // Second call to run dequeue op is to make sure the cost graph has been
+  // stored.
+  TF_EXPECT_OK(session->Run({}, {kDequeueOp0}, {}, &dq0));
+
+  CostGraphDef cost_graph;
+  TF_CHECK_OK(qr->ExportCostGraph(&cost_graph));
+  EXPECT_TRUE(cost_graph.node_size() > 0);
+
+  qr->Stop(session.get());
+}
+
+TEST(QueueRunnerTest, NoRunMetaDataTest) {
+  GraphDef graph_def = BuildSimpleGraph();
+  auto session = BuildSessionAndInitVariable(graph_def);
+
+  QueueRunnerDef queue_runner_def = BuildQueueRunnerDef(
+      kQueueName, {kCountUpToOpName}, kSquareOpName, "", {});
+  std::unique_ptr<QueueRunner> qr;
+  TF_EXPECT_OK(QueueRunner::New(queue_runner_def, &qr));
+  TF_CHECK_OK(qr->Start(session.get()));
+
+  TF_EXPECT_OK(qr->Join());
+  CostGraphDef cost_graph;
+  EXPECT_EQ(qr->ExportCostGraph(&cost_graph).code(),
+            error::FAILED_PRECONDITION);
+}
+
+}  // namespace
+}  // namespace tensorflow