From 3bdf3c592472c2b54c513417de8d9b538d3f6078 Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Mon, 8 Oct 2018 08:08:31 -0700 Subject: Make ExecutorState preserve the thread context. PiperOrigin-RevId: 216187878 --- tensorflow/core/common_runtime/executor.cc | 4 ++++ 1 file changed, 4 insertions(+) (limited to 'tensorflow/core') diff --git a/tensorflow/core/common_runtime/executor.cc b/tensorflow/core/common_runtime/executor.cc index 2c48084cab..40ec1502da 100644 --- a/tensorflow/core/common_runtime/executor.cc +++ b/tensorflow/core/common_runtime/executor.cc @@ -54,6 +54,7 @@ limitations under the License. #include "tensorflow/core/lib/hash/hash.h" #include "tensorflow/core/lib/strings/str_util.h" #include "tensorflow/core/lib/strings/stringprintf.h" +#include "tensorflow/core/platform/context.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/macros.h" #include "tensorflow/core/platform/mutex.h" @@ -1240,6 +1241,7 @@ class ExecutorState { StepStatsCollectorInterface* const stats_collector_; const tracing::TraceCollector* const trace_collector_; const tracing::EventCollector* const event_collector_; + Context context_; // QUESTION: Make it a checkpoint::TensorSliceReaderCacheWrapper // instead of a pointer? (avoids having to delete). @@ -1367,6 +1369,7 @@ ExecutorState::ExecutorState(const Executor::Args& args, ExecutorImpl* impl) trace_collector_(tracing::GetTraceCollector()), event_collector_( tracing::GetEventCollector(tracing::EventCategory::kCompute)), + context_(ContextKind::kThread), slice_reader_cache_(new checkpoint::TensorSliceReaderCacheWrapper), call_frame_(args.call_frame), impl_(impl), @@ -1586,6 +1589,7 @@ bool MightTrace(const NodeItem& item, } void ExecutorState::Process(TaggedNode tagged_node, int64 scheduled_nsec) { + WithContext wc(context_); const GraphView& gview = impl_->gview_; TaggedNodeSeq ready; TaggedNodeReadyQueue inline_ready; -- cgit v1.2.3