Only use hardware traces on a single subgraph since tracing more than that

isn't supported

PiperOrigin-RevId: 166364134

3 files changed, 8 insertions, 10 deletions

diff --git a/tensorflow/cc/training/queue_runner.cc b/tensorflow/cc/training/queue_runner.cc
index 5aaaa116cf..b6f240350b 100644
--- a/tensorflow/cc/training/queue_runner.cc
+++ b/tensorflow/cc/training/queue_runner.cc
@@ -88,7 +88,7 @@ QueueRunner::~QueueRunner() {
 Status QueueRunner::Start(Session* sess) { return Start(sess, 0); }
 
 Status QueueRunner::StartAndCollectCostGraph(Session* sess,
-                                             const RunOptions* run_options) {
+                                             const RunOptions& run_options) {
   SetRunArgumentsAndCostGraph(run_options);
   return Start(sess, 0);
 }
@@ -121,7 +121,7 @@ Status QueueRunner::Start(Session* sess, int wait_for) {
 }
 
 Status QueueRunner::StartAndCollectCostGraph(Session* session, int wait_for_ms,
-                                             const RunOptions* run_options) {
+                                             const RunOptions& run_options) {
   SetRunArgumentsAndCostGraph(run_options);
   return Start(session, wait_for_ms);
 }
@@ -214,15 +214,13 @@ Status QueueRunner::ExportCostGraph(CostGraphDef* cost_graph) const {
   return Status::OK();
 }
 
-void QueueRunner::SetRunArgumentsAndCostGraph(const RunOptions* run_options) {
+void QueueRunner::SetRunArgumentsAndCostGraph(const RunOptions& run_options) {
   cg_mu_.reset(new mutex());
   {
     mutex_lock l(*cg_mu_);
     cost_graph_.reset(new CostGraphDef());
   }
-  if (run_options) {
-    run_options_ = *run_options;
-  }
+  run_options_ = run_options;
 }
 
 Status QueueRunner::RealRun(Session* sess, const string& op,
diff --git a/tensorflow/cc/training/queue_runner.h b/tensorflow/cc/training/queue_runner.h
index 71ed44c9c6..2d34500323 100644
--- a/tensorflow/cc/training/queue_runner.h
+++ b/tensorflow/cc/training/queue_runner.h
@@ -62,13 +62,13 @@ class QueueRunner : public RunnerInterface {
   /// Starts the queue runner with the given session and sets the run arguments
   /// for sess->Run. It also collects and stores the cost model.
   Status StartAndCollectCostGraph(Session* sess,
-                                  const RunOptions* run_options = nullptr);
+                                  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.
   Status Start(Session* sess, int wait_for_ms);
   Status StartAndCollectCostGraph(Session* session, int wait_for_ms,
-                                  const RunOptions* run_options = nullptr);
+                                  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
@@ -105,7 +105,7 @@ class QueueRunner : public RunnerInterface {
 
   bool IsRunning() const override { return !stopped_; }
 
-  void SetRunArgumentsAndCostGraph(const RunOptions* run_options);
+  void SetRunArgumentsAndCostGraph(const RunOptions& run_options);
 
   Status RealRun(Session* sess, const string& op, bool update_costs);
 
diff --git a/tensorflow/cc/training/queue_runner_test.cc b/tensorflow/cc/training/queue_runner_test.cc
index da2fc03b6c..075bc8a55a 100644
--- a/tensorflow/cc/training/queue_runner_test.cc
+++ b/tensorflow/cc/training/queue_runner_test.cc
@@ -373,7 +373,7 @@ TEST(QueueRunnerTest, RunMetaDataTest) {
   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));
+  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.