From 8e0a7bfbd311e1313e7b09d73058c4077c2ba152 Mon Sep 17 00:00:00 2001 From: xiefangqi Date: Wed, 19 Aug 2020 09:24:11 +0800 Subject: [PATCH] fix tfreader WaitPost register failed --- .../dataset/engine/datasetops/source/tf_reader_op.cc | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/mindspore/ccsrc/minddata/dataset/engine/datasetops/source/tf_reader_op.cc b/mindspore/ccsrc/minddata/dataset/engine/datasetops/source/tf_reader_op.cc index 98c3622415..e58e1819ef 100644 --- a/mindspore/ccsrc/minddata/dataset/engine/datasetops/source/tf_reader_op.cc +++ b/mindspore/ccsrc/minddata/dataset/engine/datasetops/source/tf_reader_op.cc @@ -238,6 +238,9 @@ Status TFReaderOp::CalculateNumRowsPerShard() { Status TFReaderOp::operator()() { RETURN_IF_NOT_OK(CalculateNumRowsPerShard()); + // Put here to avoid register failed when Worker_Entry thread exits unexpected + RETURN_IF_NOT_OK(io_block_queue_wait_post_.Register(tree_->AllTasks())); + // launch one thread, responsible for filling mIOBlockQueue RETURN_IF_NOT_OK(tree_->LaunchWorkers(1, std::bind(&TFReaderOp::WaitToFillIOBlockQueue, this))); @@ -250,8 +253,6 @@ Status TFReaderOp::operator()() { // so workers have to be kept alive until the end of the program TaskManager::FindMe()->Post(); - RETURN_IF_NOT_OK(io_block_queue_wait_post_.Register(tree_->AllTasks())); - NotifyToFillIOBlockQueue(); while (!finished_reading_dataset_) { int64_t buffer_id = 0;