From 6fff5dea63e1833f2883cb8a61edba0582536295 Mon Sep 17 00:00:00 2001
From: Andrew Audibert <aaudibert@google.com>
Date: Fri, 17 Jan 2020 16:56:39 -0800
Subject: [PATCH] Wait for the stats thread to terminate in interleave
destructor.
PiperOrigin-RevId: 290367959
Change-Id: Ief9fe5a194d87575ceb68057ffd497cf854520c6
---
.../core/kernels/data/parallel_interleave_dataset_op.cc | 4 +++-
1 file changed, 3 insertions(+), 1 deletion(-)
diff --git a/tensorflow/core/kernels/data/parallel_interleave_dataset_op.cc b/tensorflow/core/kernels/data/parallel_interleave_dataset_op.cc
index 5e4f6567eb0..39fe1ca7425 100644
--- a/tensorflow/core/kernels/data/parallel_interleave_dataset_op.cc
+++ b/tensorflow/core/kernels/data/parallel_interleave_dataset_op.cc
@@ -494,6 +494,7 @@ class ParallelInterleaveDatasetOp::Dataset : public DatasetBase {
IncrementOutstandingThreads();
thread_pool_->Schedule([this]() { WorkerManagerThread(); });
if (ctx_->stats_aggregator()) {
+ IncrementOutstandingThreads();
thread_pool_->Schedule([this]() { StatsThread(); });
}
threads_initialized_ = true;
@@ -983,7 +984,8 @@ class ParallelInterleaveDatasetOp::Dataset : public DatasetBase {
l, std::chrono::milliseconds(kStatsReportingPeriodMillis));
}
if (cancelled_) {
- break;
+ DecrementOutstandingThreads();
+ return;
}
num_current_active_workers = num_current_active_workers_;
num_current_workers = num_current_workers_;