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

isn't supported

PiperOrigin-RevId: 166364134
This commit is contained in:
Benoit Steiner 2017-08-24 10:55:08 -07:00 committed by TensorFlower Gardener
parent a320bf1a78
commit 09fea258a3
4 changed files with 16 additions and 11 deletions

View File

@ -88,7 +88,7 @@ QueueRunner::~QueueRunner() {
Status QueueRunner::Start(Session* sess) { return Start(sess, 0); } Status QueueRunner::Start(Session* sess) { return Start(sess, 0); }
Status QueueRunner::StartAndCollectCostGraph(Session* sess, Status QueueRunner::StartAndCollectCostGraph(Session* sess,
const RunOptions* run_options) { const RunOptions& run_options) {
SetRunArgumentsAndCostGraph(run_options); SetRunArgumentsAndCostGraph(run_options);
return Start(sess, 0); 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, Status QueueRunner::StartAndCollectCostGraph(Session* session, int wait_for_ms,
const RunOptions* run_options) { const RunOptions& run_options) {
SetRunArgumentsAndCostGraph(run_options); SetRunArgumentsAndCostGraph(run_options);
return Start(session, wait_for_ms); return Start(session, wait_for_ms);
} }
@ -214,15 +214,13 @@ Status QueueRunner::ExportCostGraph(CostGraphDef* cost_graph) const {
return Status::OK(); return Status::OK();
} }
void QueueRunner::SetRunArgumentsAndCostGraph(const RunOptions* run_options) { void QueueRunner::SetRunArgumentsAndCostGraph(const RunOptions& run_options) {
cg_mu_.reset(new mutex()); cg_mu_.reset(new mutex());
{ {
mutex_lock l(*cg_mu_); mutex_lock l(*cg_mu_);
cost_graph_.reset(new CostGraphDef()); cost_graph_.reset(new CostGraphDef());
} }
if (run_options) { run_options_ = run_options;
run_options_ = *run_options;
}
} }
Status QueueRunner::RealRun(Session* sess, const string& op, Status QueueRunner::RealRun(Session* sess, const string& op,

View File

@ -62,13 +62,13 @@ class QueueRunner : public RunnerInterface {
/// Starts the queue runner with the given session and sets the run arguments /// Starts the queue runner with the given session and sets the run arguments
/// for sess->Run. It also collects and stores the cost model. /// for sess->Run. It also collects and stores the cost model.
Status StartAndCollectCostGraph(Session* sess, 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 /// 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. /// specified time (in milliseconds) for the queues to start to fill up.
Status Start(Session* sess, int wait_for_ms); Status Start(Session* sess, int wait_for_ms);
Status StartAndCollectCostGraph(Session* session, 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 /// 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 /// 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_; } 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); Status RealRun(Session* sess, const string& op, bool update_costs);

View File

@ -373,7 +373,7 @@ TEST(QueueRunnerTest, RunMetaDataTest) {
std::unique_ptr<QueueRunner> qr; std::unique_ptr<QueueRunner> qr;
TF_EXPECT_OK(QueueRunner::New(queue_runner_def, &qr)); TF_EXPECT_OK(QueueRunner::New(queue_runner_def, &qr));
RunOptions run_options; 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 // 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. // race condition where we close the queue before it was populated.

View File

@ -156,12 +156,19 @@ Status SingleMachine::Run(const GraphDef& graph_def,
// Also clear the timeline to save memory // Also clear the timeline to save memory
init_metadata_.clear_step_stats(); init_metadata_.clear_step_stats();
} }
// We can have at most one hardware trace. Use it for the main graph, and
// downgrade tracing of the queue runners to a software trace.
RunOptions queue_options = run_options_;
if (queue_options.trace_level() >= RunOptions::HARDWARE_TRACE) {
queue_options.set_trace_level(RunOptions::SOFTWARE_TRACE);
}
for (size_t i = 0; i < queue_runner_defs_.size(); ++i) { for (size_t i = 0; i < queue_runner_defs_.size(); ++i) {
std::unique_ptr<QueueRunner> queue_runner; std::unique_ptr<QueueRunner> queue_runner;
TF_RETURN_IF_ERROR(QueueRunner::New(queue_runner_defs_[i], TF_RETURN_IF_ERROR(QueueRunner::New(queue_runner_defs_[i],
coordinator_.get(), &queue_runner)); coordinator_.get(), &queue_runner));
TF_RETURN_IF_ERROR(queue_runner->StartAndCollectCostGraph( TF_RETURN_IF_ERROR(queue_runner->StartAndCollectCostGraph(
session_.get(), &run_options_)); session_.get(), queue_options));
TF_RETURN_IF_ERROR( TF_RETURN_IF_ERROR(
coordinator_->RegisterRunner(std::move(queue_runner))); coordinator_->RegisterRunner(std::move(queue_runner)));
TF_RETURN_IF_ERROR(coordinator_->GetStatus()); TF_RETURN_IF_ERROR(coordinator_->GetStatus());