dataloader_iter.py 31.0 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19
# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import six
import sys
import time
import signal
20
import numbers
21 22 23 24 25
import logging
import itertools
import threading
import numpy as np
import multiprocessing
26
from collections import namedtuple
27
from paddle.fluid.framework import _set_expected_place, _current_expected_place, set_flags
28 29

# NOTE: queue has a different name in python2 and python3
T
tianshuo78520a 已提交
30
import queue
31

32 33
import paddle
from .. import core, layers
34
from ..framework import in_dygraph_mode, _in_eager_mode
35
from ..multiprocess_utils import _set_SIGCHLD_handler, MP_STATUS_CHECK_INTERVAL, CleanupFuncRegistrar
36
from .fetcher import _IterableDatasetFetcher, _MapDatasetFetcher
37
from .batch_sampler import _InfiniteIterableSampler
38 39
from .collate import default_collate_fn, default_convert_fn
from .worker import ParentWatchDog, get_worker_info, _worker_loop, \
K
Kaipeng Deng 已提交
40 41
        _DatasetKind, _IterableDatasetStopIteration, _WorkerException, \
        _ResumeIteration
42
from .flat import _flatten_batch, _restore_batch
43 44

__all__ = ['get_worker_info']
45

46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75
# NOTE: fix `terminate called without an active exception`
# if for loop break and program exit immediately(with no model
# layers processing) after iterate **the first few data** in
# distributed lauch mode, distributed launch will call
# terminate() to kill main process on each devices, but thread
# is still iterating to fullfill blocking queue caches, which
# may cause thread error `terminate called without an active
# exception` for terminate is a strong singal and `__del__`
# of DataLoader may not be called, so we add a global link to
# the last DataLoader instance to call `__del__` to clean up
# resources
# NOTE: cannot simply as `__del__` to CleanupFuncRegistrar,
# for this will remain a link to each DataLoader instance in
# global, and will precludes GC to auto collect DataLoader
# instance and will cause memory leak
_loader = None


def _clear_loader():
    global _loader
    if _loader is not None:
        try:
            _loader.__del__()
            del _loader
        except:
            pass


CleanupFuncRegistrar.register(_clear_loader)

76

77 78 79 80 81 82 83 84 85 86 87 88 89 90 91
class _DataLoaderIterBase(object):
    """
    Iterator implement of DataLoader, will load and feed mini-batch
    data by setting in given dataloader.

    Args:
        loader(instance of DataLoader): instance of `fluid.io.DataLoader`
    """

    def __init__(self, loader):
        self._dataset = loader.dataset
        self._feed_list = loader.feed_list or []
        self._places = loader.places
        self._return_list = loader.return_list
        self._batch_sampler = loader.batch_sampler
92
        self._drop_last = loader.drop_last
93
        self._auto_collate_batch = loader.auto_collate_batch
94 95 96
        self._num_workers = loader.num_workers
        self._use_buffer_reader = loader.use_buffer_reader
        self._use_shared_memory = loader.use_shared_memory
97
        self._timeout = loader.timeout if loader.timeout > 0 else MP_STATUS_CHECK_INTERVAL
98
        self._worker_init_fn = loader.worker_init_fn
99
        self._dataset_kind = loader.dataset_kind
100
        self._pin_memory = loader.pin_memory
101

K
Kaipeng Deng 已提交
102
        self._sampler_iter = iter(self._index_sampler)
103 104 105
        if self._auto_collate_batch:
            self._collate_fn = loader.collate_fn or default_collate_fn
        else:
106
            self._collate_fn = loader.collate_fn or default_convert_fn
107

108 109 110 111 112 113 114 115 116
        # LoDTensorBlockingQueue instance for create_py_reader and a thread
        # to put mini-batch data to self._blocking_queue, mini-batch data
        # will be get from:
        # 1. multi-process mode: get data from workers' result queue
        # 2. single-process mode: read mini-batch data in main process
        self._blocking_queue = None
        self._thread = None
        self._thread_done_event = threading.Event()

K
Kaipeng Deng 已提交
117 118 119 120 121 122 123 124 125 126
    @property
    def _index_sampler(self):
        if self._auto_collate_batch:
            return self._batch_sampler
        else:
            if self._dataset_kind == _DatasetKind.MAP:
                return list(range(len(self._dataset)))
            else:
                return _InfiniteIterableSampler(self._dataset, 1)

127 128 129 130 131 132
    def __iter__(self):
        return self

    def __len__(self):
        return len(self._batch_sampler)

133 134 135 136 137 138 139 140 141 142
    def _exit_thread_expectedly(self):
        self._thread_done_event.set()
        if self._blocking_queue:
            self._blocking_queue.close()

    def _exit_thread_unexpectedly(self):
        self._thread_done_event.set()
        if self._blocking_queue:
            self._blocking_queue.kill()

143 144 145 146 147 148 149 150 151 152

class _DataLoaderIterSingleProcess(_DataLoaderIterBase):
    """
    Single process implement of DataLoaderIter, loading data from
    loader.data in main process
    """

    def __init__(self, loader):
        super(_DataLoaderIterSingleProcess, self).__init__(loader)

153
        self._dataset_fetcher = _DatasetKind.create_fetcher(
154
            self._dataset_kind, self._dataset, self._auto_collate_batch,
155
            self._collate_fn, self._drop_last)
156

157 158 159 160 161 162 163 164
        # NOTE: _structrue_infos used to record the data structure of
        # batch to restore batch structure after reading Tensor
        # from blocking_queue in single-process mode. Note that
        # only single process is used in single-process mode, we
        # can record the data structure sequencely in a list without
        # recording the send and recv index
        self._structure_infos = []

165 166 167
        # NOTE: len(self._places) batch data compose as an output
        # iteration, set blocking_queue can cache 2 iteration datas
        # at most here
168
        self._blocking_queue_capacity = 1 * len(self._places)
169 170

        self._init_thread()
171 172 173 174
        self._shutdown = False

        global _loader
        _loader = self
175 176 177 178 179 180 181 182

    def _init_thread(self):
        self._var_names = [v.name for v in self._feed_list]
        self._shapes = [v.shape for v in self._feed_list]
        self._dtypes = [v.dtype for v in self._feed_list]
        self._need_check_feed = [
            v.desc.need_check_feed() for v in self._feed_list
        ]
183
        # if only 1 place, do not need to keep order
184
        self._blocking_queue = core.init_lod_tensor_blocking_queue(
185 186
            core.Variable(), self._blocking_queue_capacity,
            len(self._places) > 1)
187 188
        self._reader = core.create_py_reader(
            self._blocking_queue, self._var_names, self._shapes, self._dtypes,
189 190
            self._need_check_feed, self._places, self._use_buffer_reader, True,
            self._pin_memory)
191

192 193
        self._thread = threading.Thread(
            target=self._thread_loop, args=(_current_expected_place(), ))
194 195 196
        self._thread.daemon = True
        self._thread.start()

197
    def _thread_loop(self, legacy_expected_place):
198 199 200 201 202 203 204 205 206 207 208 209 210
        #NOTE(zhiqiu): Set the expected place for new thread as the same as father thread,
        # and it will call platform::SetDeviceId() in c++ internally.
        # If we do not set cudaDeviceId in new thread, the default cudaDeviceId will be 0,
        # Which may cost hundreds of MB of GPU memory on CUDAPlace(0) if calling some cuda 
        # APIs in this thread.
        _set_expected_place(legacy_expected_place)

        while not self._thread_done_event.is_set():
            try:
                indices = next(self._sampler_iter)

                # read data from dataset in mini-batch
                # with paddle.fluid.dygraph.guard(place=paddle.CPUPlace()):
211
                # read data from dataset in mini-batch
212 213 214 215 216 217 218 219 220 221 222
                batch = self._dataset_fetcher.fetch(indices,
                                                    self._thread_done_event)
            except StopIteration:
                self._exit_thread_expectedly()
                return

            if batch is None or self._thread_done_event.is_set(): break

            # flat batch and record structure infos
            batch, structure = _flatten_batch(batch)
            self._structure_infos.append(structure)
223

224
            if self._thread_done_event.is_set(): break
225

226
            try:
227 228 229
                # pack as LoDTensorArray
                array = core.LoDTensorArray()
                for slot in batch:
K
Kaipeng Deng 已提交
230 231 232
                    if isinstance(slot, paddle.Tensor):
                        slot = slot.value().get_tensor()
                    elif not isinstance(slot, core.LoDTensor):
233 234 235 236 237 238
                        tmp = core.LoDTensor()
                        tmp.set(slot, core.CPUPlace())
                        slot = tmp

                    array.append(slot)

239
                if self._thread_done_event.is_set(): break
240

241 242 243 244
                try:
                    self._blocking_queue.push(array)
                except:
                    self._exit_thread_expectedly()
245

246 247 248 249 250
            except:
                self._exit_thread_unexpectedly()
                six.reraise(*sys.exc_info())

        self._exit_thread_expectedly()
251 252 253 254

    def __next__(self):
        try:
            if in_dygraph_mode():
255
                if _in_eager_mode():
256
                    data = core.eager.read_next_tensor_list(
257 258 259
                        self._reader.read_next_list()[0])
                else:
                    data = self._reader.read_next_var_list()
260
                data = _restore_batch(data, self._structure_infos.pop(0))
261 262
            else:
                if self._return_list:
263
                    data = self._reader.read_next_list()
264 265
                    for i in range(len(data)):
                        data[i] = data[i]._move_to_list()
266 267 268 269 270 271 272
                    data = [
                        _restore_batch(d, s)
                        for d, s in zip(data, self._structure_infos[:len(
                            self._places)])
                    ]
                    self._structure_infos = self._structure_infos[len(
                        self._places):]
273 274 275 276
                    # static graph organized data on multi-device with list, if
                    # place number is 1, there is only 1 device, extra the data
                    # from list for devices to be compatible with dygraph mode
                    if len(self._places) == 1:
277
                        data = data[0]
278
                else:
279 280 281
                    data = self._reader.read_next()

            return data
282
        except StopIteration:
283
            self._reader.shutdown()
284
            self._try_shutdown_all()
285 286
            six.reraise(*sys.exc_info())

287 288 289
    def _shutdown_thread(self):
        if self._thread:
            self._thread_done_event.set()
290 291 292 293 294 295 296 297 298 299 300
            # NOTE: we wait for _thread exit for 3 seconds, if
            #       thread not exit normally, force kill it
            for _ in range(3):
                if self._thread.is_alive():
                    time.sleep(1)
                else:
                    break
            else:
                if self._thread is not threading.current_thread():
                    self._thread.join()

301
            self._thread = None
302

303 304 305 306
    # python2 compatibility
    def next(self):
        return self.__next__()

307 308 309 310 311 312 313 314 315 316 317 318 319 320 321
    def _try_shutdown_all(self):
        if not self._shutdown:
            try:
                # # _blocking_queue in keep order mode holds sub-threads
                # # need to release thread resources on unexpected exit
                if self._blocking_queue:
                    self._blocking_queue.close()
                    self._blocking_queue = None
                # NOTE: blocking queue should be closed firstly for
                # blocking queue read may hang and _thread_done_event
                # cannot be checked
                self._shutdown_thread()
            finally:
                self._shutdown = True

322
    def __del__(self):
323
        self._try_shutdown_all()
324

325 326 327 328 329

class _DataLoaderIterMultiProcess(_DataLoaderIterBase):
    def __init__(self, loader):
        super(_DataLoaderIterMultiProcess, self).__init__(loader)

K
Kaipeng Deng 已提交
330 331 332
        self._persistent_workers = loader._persistent_workers
        self._resume_worker_cnt = 0

333 334 335 336 337 338 339 340
        assert self._num_workers > 0,  "Multi-process DataLoader " \
                    "invalid num_workers({})".format(self._num_workers)

        # subprocess wrokers' result queue
        self._data_queue = None

        # data get from _data_queue will be reordered by _rcvd_idx
        # for data order keeping, data index not equal _rcvd_idx 
341
        # will be cached in _task_infos
342 343 344
        self._send_idx = 0
        self._rcvd_idx = 0
        self._batches_outstanding = 0
345
        self._task_infos = {}
346
        self._structure_infos = []
347 348 349 350 351 352 353 354 355 356

        # indices outstand as _outstanding_capacity at first, and
        # blocking_queue capacity is also _outstanding_capacity.
        # _outstanding_capacity here to make sure each indices_queue
        # has at least 2 indices, and outstanding batch cached
        # output data for at least 2 iterations(Note that len(_places)
        # batches will be composed as an iteration output)
        self._outstanding_capacity = 2 * max(self._num_workers,
                                             len(self._places))

357 358 359
        # see _try_put_indices
        self._thread_lock = threading.Lock()

360
        # init workers and indices queues and put 2 indices in each indices queue
361 362 363 364
        self._init_workers()
        for _ in range(self._outstanding_capacity):
            self._try_put_indices()

365 366 367
        self._init_thread()
        self._shutdown = False

368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386
    def _init_workers(self):
        # multiprocess worker and indice queue list initial as empty
        self._workers = []
        self._worker_status = []
        self._indices_queues = []
        self._workers_idx_cycle = itertools.cycle(range(self._num_workers))

        # create data_queue for workers
        self._data_queue = multiprocessing.Queue()

        # event for workers and thread, thread event is only need 
        # in multi-processing mode
        self._workers_done_event = multiprocessing.Event()
        self._thread_done_event = threading.Event()

        for i in range(self._num_workers):
            indices_queue = multiprocessing.Queue()
            self._indices_queues.append(indices_queue)
            worker = multiprocessing.Process(
387
                target=_worker_loop,
388 389
                args=(self._dataset, self._dataset_kind, indices_queue,
                      self._data_queue, self._workers_done_event,
390
                      self._auto_collate_batch, self._collate_fn,
391 392
                      self._drop_last, self._worker_init_fn, i,
                      self._num_workers, self._use_shared_memory))
393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417
            worker.daemon = True
            worker.start()
            self._workers.append(worker)
            self._worker_status.append(True)

        core._set_process_pids(id(self), tuple(w.pid for w in self._workers))
        _set_SIGCHLD_handler()

    def _clear_and_remove_data_queue(self):
        if self._data_queue is not None:
            while True:
                try:
                    self._data_queue.get_nowait()
                except:
                    self._data_queue.cancel_join_thread()
                    self._data_queue.close()
                    break

    def _init_thread(self):
        self._var_names = [v.name for v in self._feed_list]
        self._shapes = [v.shape for v in self._feed_list]
        self._dtypes = [v.dtype for v in self._feed_list]
        self._need_check_feed = [
            v.desc.need_check_feed() for v in self._feed_list
        ]
418
        # if only 1 place, do not need to keep order
419
        self._blocking_queue = core.init_lod_tensor_blocking_queue(
420
            core.Variable(), self._outstanding_capacity, len(self._places) > 1)
421 422
        self._reader = core.create_py_reader(
            self._blocking_queue, self._var_names, self._shapes, self._dtypes,
423 424
            self._need_check_feed, self._places, self._use_buffer_reader, True,
            self._pin_memory)
425 426

        self._thread_done_event = threading.Event()
K
Kaipeng Deng 已提交
427
        # thread event is only need in multi-processing mode
428 429
        self._thread = threading.Thread(
            target=self._thread_loop, args=(_current_expected_place(), ))
430 431 432
        self._thread.daemon = True
        self._thread.start()

K
Kaipeng Deng 已提交
433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450
    def _reset(self):
        # resume iteration in following steps
        # 1. Resume workers, clear worker caches
        # put _ResumeIteration to all worker as resume iteration flag
        with self._thread_lock:
            self._resume_worker_cnt = self._num_workers
            for worker_id in range(self._num_workers):
                self._indices_queues[worker_id].put(_ResumeIteration())
                self._batches_outstanding += 1
        # all flag will be check in _thread_loop, simply wait here
        while self._resume_worker_cnt > 0:
            time.sleep(0.5)

        # 2. clear blocking_queue caches
        # in order not to restart the thread, we just clear
        # the blocking_queue cachees instead of recreating one
        while self._blocking_queue.size() >= len(self._places):
            if in_dygraph_mode():
451
                if _in_eager_mode():
452
                    data = core.eager.read_next_tensor_list(
453 454 455
                        self._reader.read_next_list()[0])
                else:
                    self._reader.read_next_var_list()
K
Kaipeng Deng 已提交
456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479
            elif self._return_list:
                self._reader.read_next_list()
            else:
                data = self._reader.read_next()

        # 3. reset all states
        self._send_idx = 0
        self._rcvd_idx = 0
        self._batches_outstanding = 0
        self._task_infos = {}
        self._structure_infos = []

        # set all worker status available
        self._worker_status = [True] * self._num_workers

        # 4. reset _sampler_iter and put prefetch indices to start next epoch
        # init workers and indices queues and put 2 indices in each indices queue
        self._sampler_iter = iter(self._index_sampler)
        for _ in range(self._outstanding_capacity):
            self._try_put_indices()

    def _shutdown_worker(self, worker_id, shutdown=False):
        if self._worker_status[worker_id] or (self._persistent_workers and
                                              shutdown):
480 481 482
            self._indices_queues[worker_id].put(None)
            self._worker_status[worker_id] = False

483
    def _try_shutdown_all(self, timeout=None):
484 485 486 487 488 489 490 491 492 493
        if not self._shutdown:
            try:
                self._exit_thread_expectedly()
                self._clear_and_remove_data_queue()

                # set _workers_done_event should be set before put None
                # to indices_queue, workers wll exit on reading None from
                # indices_queue
                self._workers_done_event.set()
                for i in range(self._num_workers):
K
Kaipeng Deng 已提交
494
                    self._shutdown_worker(i, shutdown=True)
495

496 497 498 499 500 501
                if not self._shutdown:
                    for w in self._workers:
                        w.join(timeout)
                    for q in self._indices_queues:
                        q.cancel_join_thread()
                        q.close()
502 503 504 505
            finally:
                core._erase_process_pids(id(self))
                self._shutdown = True

506 507 508 509 510 511 512 513
    def _thread_loop(self, legacy_expected_place):
        #NOTE(zhiqiu): Set the expected place for new thread as the same as father thread,
        # and it will call platform::SetDeviceId() in c++ internally.
        # If we do not set cudaDeviceId in new thread, the default cudaDeviceId will be 0,
        # Which may cost hundreds of MB of GPU memory on CUDAPlace(0) if calling some cuda 
        # APIs in this thread.
        _set_expected_place(legacy_expected_place)

514 515 516 517 518 519
        while not self._thread_done_event.is_set():
            batch = self._get_data()
            if not self._thread_done_event.is_set():
                if batch is None:
                    self._exit_thread_expectedly()
                else:
K
Kaipeng Deng 已提交
520 521 522 523
                    if isinstance(batch, _ResumeIteration):
                        assert self._resume_worker_cnt > 0
                        self._resume_worker_cnt -= 1
                        continue
524 525 526 527 528 529 530 531 532 533
                    try:
                        # pack as LoDTensorArray
                        array = core.LoDTensorArray()
                        if self._use_shared_memory:
                            for tensor in batch:
                                array.append(tensor)
                        else:
                            # LoDTensor not in shared memory is not
                            # serializable, cannot be create in workers
                            for slot in batch:
K
Kaipeng Deng 已提交
534 535 536
                                if isinstance(slot, paddle.Tensor):
                                    slot = slot.value().get_tensor()
                                elif not isinstance(slot, core.LoDTensor):
537 538 539 540 541 542 543
                                    tmp = core.LoDTensor()
                                    tmp.set(slot, core.CPUPlace())
                                    slot = tmp
                                array.append(slot)

                        if not self._blocking_queue.push(array):
                            self._blocking_queue.close()
K
Kaipeng Deng 已提交
544
                    except Exception as e:
545 546 547 548 549 550 551
                        self._exit_thread_unexpectedly()
                        six.reraise(*sys.exc_info())
                    finally:
                        self._rcvd_idx += 1

    def _get_data(self):
        while not self._thread_done_event.is_set():
552 553 554 555 556 557
            # For IterableDataset, batch indices is generated infinitely
            # for each worker to raise StopIteration, but a StopIteration
            # raising process will discard a batch indices which is count
            # in _send_idx but will not increase _rcvd_idx, so we check 
            # whether the worker is still alive here to skip the discarded
            # batch indices and increase _rcvd_idx
558 559 560
            if self._dataset_kind == _DatasetKind.ITER:
                while self._rcvd_idx < self._send_idx:
                    info = self._task_infos[self._rcvd_idx]
561
                    if len(info) == 3 or self._worker_status[info[0]]:
562 563 564 565 566
                        break
                    del self._task_infos[self._rcvd_idx]
                    self._rcvd_idx += 1
                    self._batches_outstanding -= 1
                else:
567 568 569 570 571 572 573 574
                    # NOTE: when _rcvd_idx catch up _send_idx, which means
                    #       one of following:
                    #       1. all 2 * num_workers batches have been loaded
                    #          and stored in _blocking_queue
                    #       2. all data drained
                    #       we need to let _thread blocking at _data_queue
                    #       get_data to inoccupy CPU, otherwise may occupy
                    #       CPU time for model running
K
Kaipeng Deng 已提交
575 576 577 578 579 580 581 582 583
                    # NOTE: in persistent workers mode, do not check data
                    #       drained here, simply let it go to _data_queue
                    #       reading to get _ResumeIteration
                    if not self._persistent_workers:
                        # NOTE: _rcvd_idx and _send_idx only record batches among
                        #       workers, if batches among workers drained, there
                        #       may also be data in blocking queue
                        if self._batches_outstanding < len(self._places):
                            return None
584 585

            if self._rcvd_idx in self._task_infos and \
586 587 588 589
                    len(self._task_infos[self._rcvd_idx]) == 3:
                info = self._task_infos.pop(self._rcvd_idx)
                self._structure_infos.append(info[2])
                return info[1]
590

591 592 593
            try:
                # [ avoid hang ]: main process may blocking at _reader.read_next when
                # KeyboardInterrupt, we do following tradeoff:
594
                # 1. get data with timeout, MP_STATUS_CHECK_INTERVAL(5s) as timeout
595 596 597 598 599 600 601
                #    default, if KeyboardInterrupt blocking, failed workers will be
                #    checked and raise RuntimeError to quit DataLoader in timeout
                #    exception handling.
                # 2. if get data timeout and check workers all alive, continue to
                #    get data again
                data = self._data_queue.get(timeout=self._timeout)
            except Exception as e:
602 603 604 605 606
                # check if thread done event set when waiting data
                if self._thread_done_event.is_set():
                    continue

                # check failed workers
607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627
                failed_workers = []
                for i, w in enumerate(self._workers):
                    if self._worker_status[i] and not w.is_alive():
                        failed_workers.append(w)
                        self._shutdown_worker(i)
                if len(failed_workers) > 0:
                    self._exit_thread_unexpectedly()
                    pids = ', '.join(str(w.pid) for w in failed_workers)
                    raise RuntimeError("DataLoader {} workers exit unexpectedly, " \
                                "pids: {}".format(len(failed_workers), pids))

                # get(timeout) will call _poll(timeout) and may raise IOError
                if isinstance(e, queue.Empty) or isinstance(e, IOError):
                    # continue on timeout to keep getting data from queue
                    continue

                self._exit_thread_unexpectedly()
                logging.error("DataLoader reader thread failed({}) to read data from " \
                              "workers' result queue.".format(e))
                six.reraise(*sys.exc_info())
            else:
628 629 630 631 632 633 634
                if self._dataset_kind == _DatasetKind.ITER and isinstance(
                        data, _IterableDatasetStopIteration):
                    # if a worker get StopIteraion, we shutdown this worker,
                    # note that this batch indices to trigger StopIteration
                    # is discard, outstanding batch number should be decrease
                    # and another indices should be put for other workers
                    # may still working.
K
Kaipeng Deng 已提交
635 636 637 638 639
                    if self._persistent_workers:
                        self._worker_status[data.worker_id] = False
                    else:
                        self._shutdown_worker(data.worker_id)
                        self._batches_outstanding -= 1
640 641 642
                    self._try_put_indices()
                    continue

643
                idx, batch, structure = data
K
Kaipeng Deng 已提交
644 645 646 647 648

                if isinstance(idx, _ResumeIteration) and batch is None \
                        and structure is None:
                    return idx

649 650 651 652
                if isinstance(batch, _WorkerException):
                    self._exit_thread_unexpectedly()
                    batch.reraise()

653
                if idx == self._rcvd_idx:
654
                    del self._task_infos[idx]
655
                    self._structure_infos.append(structure)
656 657
                    return batch
                else:
658
                    self._task_infos[idx] += (batch, structure)
659 660 661
                    continue

    def _try_put_indices(self):
662
        assert self._batches_outstanding <= self._outstanding_capacity, \
663
                    "too many indices have been put to queue"
664 665 666 667 668 669 670 671 672 673 674 675 676 677
        # In multi-process mode for IterableDataset, _try_put_indices will
        # be called both in main process(for our implement has blocking queue,
        # and blocking queue read is in main process) and thread, which may
        # cause error following error
        #   1. "ValueError: generator already executing" in next(self._sampler_iter)
        #   2. re-enter in increase _send_idx
        # add a lock for threading save, for _try_put_indices is only a slight
        # function which is not in data reading pipeline, this lock almost no
        # influence on performance
        with self._thread_lock:
            try:
                indices = next(self._sampler_iter)
            except StopIteration:
                return
678

679 680 681 682 683 684
            for i in range(self._num_workers):
                worker_idx = next(self._workers_idx_cycle)
                if self._worker_status[worker_idx]:
                    break
            else:
                return
685

686 687 688 689
            self._indices_queues[worker_idx].put((self._send_idx, indices))
            self._task_infos[self._send_idx] = (worker_idx, )
            self._batches_outstanding += 1
            self._send_idx += 1
690 691 692 693

    def __del__(self):
        self._try_shutdown_all()

694 695 696
    def _shutdown_on_exit(self):
        self._try_shutdown_all(1)

697 698 699 700 701 702 703 704 705 706
    def __next__(self):
        try:
            # _batches_outstanding here record the total batch data number
            # in 'from after _try_put_indices to beforeoutput data', this
            # value should be _outstanding_capacity if data is not drained,
            # if _batches_outstanding is less than _places number, there are
            # no enough data to generate next output, close blocking_queue and
            # set _thread_done_event here, py_reader will raise StopIteration,
            # end workers and indices_queues in StopIteration handling
            if self._batches_outstanding < len(self._places):
K
Kaipeng Deng 已提交
707 708 709 710 711
                if self._persistent_workers:
                    raise StopIteration
                else:
                    self._thread_done_event.set()
                    self._blocking_queue.close()
712 713

            if in_dygraph_mode():
714
                if _in_eager_mode():
715
                    data = core.eager.read_next_tensor_list(
716 717 718
                        self._reader.read_next_list()[0])
                else:
                    data = self._reader.read_next_var_list()
719
                data = _restore_batch(data, self._structure_infos.pop(0))
720 721 722
            else:
                if self._return_list:
                    data = self._reader.read_next_list()
723 724
                    for i in range(len(data)):
                        data[i] = data[i]._move_to_list()
725 726 727 728 729 730 731
                    data = [
                        _restore_batch(d, s)
                        for d, s in zip(data, self._structure_infos[:len(
                            self._places)])
                    ]
                    self._structure_infos = self._structure_infos[len(
                        self._places):]
732 733 734 735 736 737 738 739 740 741
                    # static graph organized data on multi-device with list, if
                    # place number is 1, there is only 1 device, extra the data
                    # from list for devices to be compatible with dygraph mode
                    if len(self._places) == 1:
                        data = data[0]
                else:
                    data = self._reader.read_next()
            self._on_output_batch()
            return data
        except StopIteration:
K
Kaipeng Deng 已提交
742 743 744
            if not self._persistent_workers:
                self._reader.shutdown()
                self._try_shutdown_all()
745 746 747 748 749 750 751 752 753 754
            six.reraise(*sys.exc_info())

    # python2 compatibility
    def next(self):
        return self.__next__()

    def _on_output_batch(self):
        for _ in range(len(self._places)):
            self._batches_outstanding -= 1
            self._try_put_indices()