1 /**
2 * Copyright 2021 Huawei Technologies Co., Ltd
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16 #include "minddata/dataset/engine/datasetops/source/mappable_leaf_op.h"
17 #include "utils/ms_utils.h"
18 #include "minddata/dataset/core/config_manager.h"
19 #include "minddata/dataset/engine/datasetops/source/sampler/sequential_sampler.h"
20 #include "minddata/dataset/engine/db_connector.h"
21 #include "minddata/dataset/engine/execution_tree.h"
22
23 namespace mindspore {
24 namespace dataset {
MappableLeafOp(int32_t num_wkrs,int32_t queue_size,std::shared_ptr<SamplerRT> sampler)25 MappableLeafOp::MappableLeafOp(int32_t num_wkrs, int32_t queue_size, std::shared_ptr<SamplerRT> sampler)
26 : ParallelOp(num_wkrs, queue_size, std::move(sampler)) {}
27
28 // Main logic, Register Queue with TaskGroup, launch all threads and do the functor's work
operator ()()29 Status MappableLeafOp::operator()() {
30 RETURN_IF_NOT_OK(LaunchThreadsAndInitOp());
31 TensorRow sample_row;
32 RETURN_IF_NOT_OK(sampler_->GetNextSample(&sample_row));
33 int64_t row_cnt = 0;
34 while (true) { // each iteration is 1 epoch, breaks when IsLastIteration() is true
35 while (sample_row.eoe() == false) {
36 std::shared_ptr<Tensor> sample_ids = sample_row[0];
37 for (auto itr = sample_ids->begin<int64_t>(); itr != sample_ids->end<int64_t>(); ++itr) {
38 if ((*itr) >= num_rows_) {
39 MS_LOG(WARNING) << "Skipping sample with ID: " << *itr << " since it is out of bound: " << num_rows_;
40 continue; // index out of bound, skipping
41 }
42 RETURN_IF_NOT_OK(
43 io_block_queues_[row_cnt++ % num_workers_]->Add(std::make_unique<IOBlock>(*itr, IOBlock::kDeIoBlockNone)));
44 }
45 RETURN_IF_NOT_OK(sampler_->GetNextSample(&sample_row));
46 }
47 if (IsLastIteration()) {
48 std::unique_ptr<IOBlock> eoe_block = std::make_unique<IOBlock>(IOBlock::kDeIoBlockFlagEoe);
49 std::unique_ptr<IOBlock> eof_block = std::make_unique<IOBlock>(IOBlock::kDeIoBlockFlagEof);
50 RETURN_IF_NOT_OK(io_block_queues_[(row_cnt++) % num_workers_]->Add(std::move(eoe_block)));
51 RETURN_IF_NOT_OK(io_block_queues_[(row_cnt++) % num_workers_]->Add(std::move(eof_block)));
52 for (int32_t i = 0; i < num_workers_; ++i) {
53 RETURN_IF_NOT_OK(
54 io_block_queues_[i]->Add(std::make_unique<IOBlock>(std::vector<int64_t>(), IOBlock::kDeIoBlockNone)));
55 }
56 return Status::OK();
57 } else { // not the last repeat.
58 RETURN_IF_NOT_OK(
59 io_block_queues_[(row_cnt++) % num_workers_]->Add(std::make_unique<IOBlock>(IOBlock::kDeIoBlockFlagEoe)));
60 }
61
62 if (epoch_sync_flag_) {
63 // If epoch_sync_flag_ is set, then master thread sleeps until all the worker threads have finished their job for
64 // the current epoch.
65 RETURN_IF_NOT_OK(WaitForWorkers());
66 }
67 // If not the last repeat, self-reset and go to loop again.
68 if (!IsLastIteration()) {
69 RETURN_IF_NOT_OK(Reset());
70 RETURN_IF_NOT_OK(sampler_->GetNextSample(&sample_row));
71 }
72 UpdateRepeatAndEpochCounter();
73 }
74 }
75
76 // Reset Sampler and wakeup Master thread (functor)
Reset()77 Status MappableLeafOp::Reset() {
78 MS_LOG(DEBUG) << Name() << " performing a self-reset.";
79 RETURN_IF_NOT_OK(sampler_->ResetSampler());
80 return Status::OK();
81 }
82
83 // hand shake with Sampler, allow Sampler to call RandomAccessOp's functions to get NumRows
InitSampler()84 Status MappableLeafOp::InitSampler() {
85 RETURN_IF_NOT_OK(sampler_->HandshakeRandomAccessOp(this));
86 return Status::OK();
87 }
88
89 // contains the main logic of pulling a IOBlock from IOBlockQueue, load a row and push the row to out_connector_
90 // IMPORTANT: 1 IOBlock produces 1 row
WorkerEntry(int32_t worker_id)91 Status MappableLeafOp::WorkerEntry(int32_t worker_id) {
92 TaskManager::FindMe()->Post();
93 std::unique_ptr<IOBlock> io_block;
94 RETURN_IF_NOT_OK(io_block_queues_[worker_id]->PopFront(&io_block));
95 while (io_block != nullptr) {
96 if (io_block->wait()) {
97 // Sync io_block is a signal that master thread wants us to pause and sync with other workers.
98 // The last guy who comes to this sync point should reset the counter and wake up the master thread.
99 if (++num_workers_paused_ == num_workers_) {
100 wait_for_workers_post_.Set();
101 }
102 } else if (io_block->eoe()) {
103 RETURN_IF_NOT_OK(out_connector_->SendEOE(worker_id));
104 } else if (io_block->eof()) {
105 RETURN_IF_NOT_OK(out_connector_->SendEOF(worker_id));
106 } else {
107 std::vector<int64_t> keys;
108 RETURN_IF_NOT_OK(io_block->GetKeys(&keys));
109 if (keys.empty()) return Status::OK(); // empty key is a quit signal for workers
110 TensorRow trow;
111 RETURN_IF_NOT_OK(this->LoadTensorRow(keys[0], &trow));
112 RETURN_IF_NOT_OK(out_connector_->Add(std::move(trow), worker_id));
113 }
114 RETURN_IF_NOT_OK(io_block_queues_[worker_id]->PopFront(&io_block));
115 }
116 RETURN_STATUS_UNEXPECTED("[Internal ERROR] Unexpected nullptr received in worker.");
117 }
118 } // namespace dataset
119 } // namespace mindspore
120