wehub-resource-sync
888 lines
34 KiB
Python
888 lines
34 KiB
Python
# 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 itertools
|
|
import logging
|
|
import os
|
|
import queue
|
|
import sys
|
|
import threading
|
|
import time
|
|
import warnings
|
|
|
|
import numpy as np
|
|
|
|
import paddle
|
|
from paddle import profiler
|
|
from paddle.base.framework import _current_expected_place, _set_expected_place
|
|
from paddle.pir.core import datatype_to_vartype
|
|
from paddle.profiler.timer import benchmark
|
|
from paddle.profiler.utils import in_profiler_mode
|
|
|
|
from ...framework import core, in_dynamic_mode, in_pir_mode
|
|
from ..multiprocess_utils import (
|
|
MP_STATUS_CHECK_INTERVAL,
|
|
CleanupFuncRegistrar,
|
|
_set_SIGCHLD_handler,
|
|
)
|
|
from .batch_sampler import _InfiniteIterableSampler
|
|
from .collate import default_collate_fn, default_convert_fn
|
|
from .flat import _flatten_batch, _restore_batch
|
|
from .worker import (
|
|
_DatasetKind,
|
|
_IterableDatasetStopIteration,
|
|
_ResumeIteration,
|
|
_worker_loop,
|
|
_WorkerException,
|
|
)
|
|
|
|
# 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 launch mode, distributed launch will call
|
|
# terminate() to kill main process on each devices, but thread
|
|
# is still iterating to fulfill blocking queue caches, which
|
|
# may cause thread error `terminate called without an active
|
|
# exception` for terminate is a strong signal 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)
|
|
|
|
|
|
class _DataLoaderIterBase:
|
|
"""
|
|
Iterator implement of DataLoader, will load and feed mini-batch
|
|
data by setting in given dataloader.
|
|
|
|
Args:
|
|
loader(instance of DataLoader): instance of `paddle.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
|
|
self._drop_last = loader.drop_last
|
|
self._auto_collate_batch = loader.auto_collate_batch
|
|
self._num_workers = loader.num_workers
|
|
self._use_buffer_reader = loader.use_buffer_reader
|
|
self._reader_buffer_size = loader.reader_buffer_size
|
|
self._prefetch_factor = loader.prefetch_factor
|
|
self._use_shared_memory = loader.use_shared_memory
|
|
self._timeout = (
|
|
loader.timeout if loader.timeout > 0 else MP_STATUS_CHECK_INTERVAL
|
|
)
|
|
self._worker_init_fn = loader.worker_init_fn
|
|
self._dataset_kind = loader.dataset_kind
|
|
self._pin_memory = loader.pin_memory
|
|
|
|
self._sampler_iter = iter(self._index_sampler)
|
|
if self._auto_collate_batch:
|
|
self._collate_fn = loader.collate_fn or default_collate_fn
|
|
else:
|
|
self._collate_fn = loader.collate_fn or default_convert_fn
|
|
|
|
# DenseTensorBlockingQueue 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()
|
|
|
|
@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)
|
|
|
|
def __iter__(self):
|
|
return self
|
|
|
|
def __next__(self):
|
|
raise NotImplementedError('Should implement `__next__` for a iterator')
|
|
|
|
def __len__(self):
|
|
return len(self._batch_sampler)
|
|
|
|
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()
|
|
|
|
|
|
class _DataLoaderIterSingleProcess(_DataLoaderIterBase):
|
|
"""
|
|
Single process implement of DataLoaderIter, loading data from
|
|
loader.data in main process
|
|
"""
|
|
|
|
def __init__(self, loader):
|
|
super().__init__(loader)
|
|
|
|
self._dataset_fetcher = _DatasetKind.create_fetcher(
|
|
self._dataset_kind,
|
|
self._dataset,
|
|
self._auto_collate_batch,
|
|
self._collate_fn,
|
|
self._drop_last,
|
|
)
|
|
|
|
# NOTE: _structure_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 = []
|
|
|
|
# NOTE: len(self._places) batch data compose as an output
|
|
# iteration, set blocking_queue can cache "self._prefetch_factor" iteration datas
|
|
# at most here
|
|
self._blocking_queue_capacity = self._prefetch_factor * len(
|
|
self._places
|
|
)
|
|
|
|
self._shutdown = False
|
|
try:
|
|
self._init_thread()
|
|
except Exception:
|
|
self._try_shutdown_all()
|
|
raise
|
|
|
|
global _loader
|
|
_loader = self
|
|
|
|
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]
|
|
if in_pir_mode():
|
|
self._need_check_feed = [False for v in self._feed_list]
|
|
self._dtypes = [
|
|
datatype_to_vartype[v.dtype] for v in self._feed_list
|
|
]
|
|
else:
|
|
self._need_check_feed = [
|
|
v.desc.need_check_feed() for v in self._feed_list
|
|
]
|
|
self._dtypes = [v.dtype for v in self._feed_list]
|
|
# if only 1 place, do not need to keep order
|
|
self._blocking_queue = core.init_dense_tensor_blocking_queue(
|
|
core.Variable(),
|
|
self._blocking_queue_capacity,
|
|
len(self._places) > 1,
|
|
)
|
|
self._reader = core.create_py_reader(
|
|
self._blocking_queue,
|
|
self._var_names,
|
|
self._shapes,
|
|
self._dtypes,
|
|
self._need_check_feed,
|
|
self._places,
|
|
self._use_buffer_reader,
|
|
True,
|
|
self._pin_memory,
|
|
self._reader_buffer_size,
|
|
)
|
|
|
|
self._thread = threading.Thread(
|
|
target=self._thread_loop, args=(_current_expected_place(),)
|
|
)
|
|
self._thread.daemon = True
|
|
self._thread.start()
|
|
|
|
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.
|
|
core.set_current_thread_name("Dataloader_" + str(id(self)))
|
|
_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.base.dygraph.guard(place=paddle.CPUPlace()):
|
|
# read data from dataset in mini-batch
|
|
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)
|
|
|
|
if self._thread_done_event.is_set():
|
|
break
|
|
|
|
try:
|
|
# pack as DenseTensorArray
|
|
array = core.DenseTensorArray()
|
|
for slot in batch:
|
|
if isinstance(slot, paddle.Tensor):
|
|
slot = slot.value().get_tensor()
|
|
elif not isinstance(slot, core.DenseTensor):
|
|
tmp = core.DenseTensor()
|
|
tmp.set(slot, core.CPUPlace())
|
|
slot = tmp
|
|
|
|
array.append(slot)
|
|
|
|
if self._thread_done_event.is_set():
|
|
break
|
|
|
|
try:
|
|
self._blocking_queue.push(array)
|
|
except:
|
|
self._exit_thread_expectedly()
|
|
|
|
except Exception as e:
|
|
self._exit_thread_unexpectedly()
|
|
raise e
|
|
|
|
self._exit_thread_expectedly()
|
|
|
|
def __next__(self):
|
|
if in_profiler_mode():
|
|
trace_event = profiler.RecordEvent(
|
|
name="_DataLoaderIterSingleProcess",
|
|
event_type=profiler.TracerEventType.Dataloader,
|
|
)
|
|
trace_event.begin()
|
|
try:
|
|
benchmark().check_if_need_record(self)
|
|
benchmark().before_reader()
|
|
if in_dynamic_mode():
|
|
data = core.eager.read_next_tensor_list(
|
|
self._reader.read_next_list()[0]
|
|
)
|
|
data = _restore_batch(data, self._structure_infos.pop(0))
|
|
else:
|
|
# in static graph mode
|
|
if self._return_list:
|
|
data = self._reader.read_next_list()
|
|
for i in range(len(data)):
|
|
data[i] = data[i]._move_to_list()
|
|
structs = [
|
|
self._structure_infos.pop(0)
|
|
for _ in range(len(self._places))
|
|
]
|
|
data = [_restore_batch(d, s) for d, s in zip(data, structs)]
|
|
# 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()
|
|
benchmark().after_reader()
|
|
|
|
return data
|
|
except StopIteration:
|
|
self._reader.shutdown()
|
|
self._try_shutdown_all()
|
|
raise
|
|
finally:
|
|
if in_profiler_mode():
|
|
trace_event.end()
|
|
|
|
def _shutdown_thread(self):
|
|
if self._thread:
|
|
self._thread_done_event.set()
|
|
# 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()
|
|
|
|
self._thread = None
|
|
|
|
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
|
|
|
|
def __del__(self):
|
|
self._try_shutdown_all()
|
|
|
|
|
|
class _DataLoaderIterMultiProcess(_DataLoaderIterBase):
|
|
def __init__(self, loader):
|
|
super().__init__(loader)
|
|
|
|
self._persistent_workers = loader._persistent_workers
|
|
self._resume_worker_cnt = 0
|
|
|
|
assert self._num_workers > 0, (
|
|
f"Multi-process DataLoader invalid num_workers({self._num_workers})"
|
|
)
|
|
|
|
# subprocess workers' 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
|
|
# will be cached in _task_infos
|
|
self._send_idx = 0
|
|
self._rcvd_idx = 0
|
|
self._batches_outstanding = 0
|
|
self._task_infos = {}
|
|
self._structure_infos = []
|
|
|
|
# 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 "_prefetch_factor" indices, and outstanding batch cached
|
|
# output data for at least "_prefetch_factor" iterations(Note that len(_places)
|
|
# batches will be composed as an iteration output)
|
|
self._outstanding_capacity = self._prefetch_factor * max(
|
|
self._num_workers, len(self._places)
|
|
)
|
|
|
|
# see _try_put_indices
|
|
self._thread_lock = threading.Lock()
|
|
|
|
self._base_seed = np.random.randint(low=0, high=sys.maxsize)
|
|
|
|
# Note(zhangbo): shm_buffer_size is used for MemoryMapAllocationPool.
|
|
# MemoryMapAllocationPool is used to cache and reuse shm, thus reducing munmap in dataloader.
|
|
# For more details, please see: paddle/base/memory/allocation/mmap_allocator.h
|
|
if os.environ.get('FLAGS_use_shm_cache', False) in [
|
|
1,
|
|
'1',
|
|
True,
|
|
'True',
|
|
'true',
|
|
]:
|
|
try:
|
|
self._worker_shm_buffer_size = (2 + 1) * len(self._dataset[0])
|
|
except:
|
|
self._worker_shm_buffer_size = 0
|
|
warnings.warn(
|
|
"Setting the shm cache buffer size to 0, equivalent to not using the shm cache policy."
|
|
)
|
|
else:
|
|
self._worker_shm_buffer_size = 0
|
|
self._main_thread_shm_buffer_size = (
|
|
(self._worker_shm_buffer_size) * 2 * self._num_workers
|
|
)
|
|
|
|
self._shutdown = False
|
|
# init workers and indices queues and put 2 indices in each indices queue
|
|
self._init_workers()
|
|
for _ in range(self._outstanding_capacity):
|
|
self._try_put_indices()
|
|
|
|
try:
|
|
self._init_thread()
|
|
except Exception:
|
|
self._try_shutdown_all()
|
|
raise
|
|
|
|
def _init_workers(self):
|
|
from paddle.incubate import multiprocessing
|
|
|
|
# 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()
|
|
indices_queue.cancel_join_thread()
|
|
self._indices_queues.append(indices_queue)
|
|
worker = multiprocessing.Process(
|
|
target=_worker_loop,
|
|
args=(
|
|
self._dataset,
|
|
self._dataset_kind,
|
|
indices_queue,
|
|
self._data_queue,
|
|
self._workers_done_event,
|
|
self._auto_collate_batch,
|
|
self._collate_fn,
|
|
self._drop_last,
|
|
self._worker_init_fn,
|
|
i,
|
|
self._num_workers,
|
|
self._use_shared_memory,
|
|
self._base_seed,
|
|
self._worker_shm_buffer_size,
|
|
),
|
|
)
|
|
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]
|
|
if in_pir_mode():
|
|
self._need_check_feed = [False for v in self._feed_list]
|
|
self._dtypes = [
|
|
datatype_to_vartype[v.dtype] for v in self._feed_list
|
|
]
|
|
else:
|
|
self._need_check_feed = [
|
|
v.desc.need_check_feed() for v in self._feed_list
|
|
]
|
|
self._dtypes = [v.dtype for v in self._feed_list]
|
|
# if only 1 place, do not need to keep order
|
|
self._blocking_queue = core.init_dense_tensor_blocking_queue(
|
|
core.Variable(), self._outstanding_capacity, len(self._places) > 1
|
|
)
|
|
core._set_max_memory_map_allocation_pool_size(
|
|
self._main_thread_shm_buffer_size
|
|
)
|
|
self._reader = core.create_py_reader(
|
|
self._blocking_queue,
|
|
self._var_names,
|
|
self._shapes,
|
|
self._dtypes,
|
|
self._need_check_feed,
|
|
self._places,
|
|
self._use_buffer_reader,
|
|
True,
|
|
self._pin_memory,
|
|
self._reader_buffer_size,
|
|
)
|
|
|
|
self._thread_done_event = threading.Event()
|
|
# thread event is only need in multi-processing mode
|
|
self._thread = threading.Thread(
|
|
target=self._thread_loop, args=(_current_expected_place(),)
|
|
)
|
|
self._thread.daemon = True
|
|
self._thread.start()
|
|
|
|
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_dynamic_mode():
|
|
data = core.eager.read_next_tensor_list(
|
|
self._reader.read_next_list()[0]
|
|
)
|
|
else:
|
|
if 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 worker_id < len(self._worker_status) and (
|
|
self._worker_status[worker_id]
|
|
or self._persistent_workers
|
|
and shutdown
|
|
):
|
|
self._indices_queues[worker_id].put(None)
|
|
self._worker_status[worker_id] = False
|
|
|
|
def _try_shutdown_all(self, timeout=None):
|
|
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 will exit on reading None from
|
|
# indices_queue
|
|
self._workers_done_event.set()
|
|
for i in range(self._num_workers):
|
|
self._shutdown_worker(i, shutdown=True)
|
|
|
|
if not self._shutdown:
|
|
for w in self._workers:
|
|
w.join(timeout)
|
|
for q in self._indices_queues:
|
|
q.cancel_join_thread()
|
|
q.close()
|
|
finally:
|
|
core._erase_process_pids(id(self))
|
|
self._shutdown = True
|
|
|
|
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.
|
|
core.set_current_thread_name("Dataloader_" + str(id(self)))
|
|
_set_expected_place(legacy_expected_place)
|
|
|
|
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:
|
|
if isinstance(batch, _ResumeIteration):
|
|
assert self._resume_worker_cnt > 0
|
|
self._resume_worker_cnt -= 1
|
|
continue
|
|
try:
|
|
# pack as DenseTensorArray
|
|
array = core.DenseTensorArray()
|
|
if self._use_shared_memory:
|
|
for tensor in batch:
|
|
array.append(tensor)
|
|
else:
|
|
# DenseTensor not in shared memory is not
|
|
# serializable, cannot be create in workers
|
|
for slot in batch:
|
|
if isinstance(slot, paddle.Tensor):
|
|
slot = slot.get_tensor()
|
|
elif not isinstance(slot, core.DenseTensor):
|
|
tmp = core.DenseTensor()
|
|
tmp.set(slot, core.CPUPlace())
|
|
slot = tmp
|
|
array.append(slot)
|
|
|
|
if not self._blocking_queue.push(array):
|
|
self._blocking_queue.close()
|
|
except Exception as e:
|
|
self._exit_thread_unexpectedly()
|
|
raise e
|
|
finally:
|
|
self._rcvd_idx += 1
|
|
|
|
def _get_data(self):
|
|
while not self._thread_done_event.is_set():
|
|
# 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
|
|
if self._dataset_kind == _DatasetKind.ITER:
|
|
while self._rcvd_idx < self._send_idx:
|
|
info = self._task_infos[self._rcvd_idx]
|
|
if len(info) == 3 or self._worker_status[info[0]]:
|
|
break
|
|
del self._task_infos[self._rcvd_idx]
|
|
self._rcvd_idx += 1
|
|
self._batches_outstanding -= 1
|
|
else:
|
|
# 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
|
|
# 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
|
|
|
|
if (
|
|
self._rcvd_idx in self._task_infos
|
|
and 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]
|
|
|
|
try:
|
|
# [ avoid hang ]: main process may blocking at _reader.read_next when
|
|
# KeyboardInterrupt, we do following tradeoff:
|
|
# 1. get data with timeout, MP_STATUS_CHECK_INTERVAL(5s) as timeout
|
|
# 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:
|
|
# check if thread done event set when waiting data
|
|
if self._thread_done_event.is_set():
|
|
continue
|
|
|
|
# check failed workers
|
|
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)
|
|
logging.warning(
|
|
f"DataLoader {len(failed_workers)} workers exit unexpectedly, "
|
|
f"pids: {pids}"
|
|
)
|
|
return
|
|
|
|
# get(timeout) will call _poll(timeout) and may raise IOError
|
|
if isinstance(e, (IOError, queue.Empty)):
|
|
# continue on timeout to keep getting data from queue
|
|
continue
|
|
|
|
self._exit_thread_unexpectedly()
|
|
logging.error(
|
|
f"DataLoader reader thread failed({e}) to read data from "
|
|
"workers' result queue."
|
|
)
|
|
raise e
|
|
else:
|
|
if self._dataset_kind == _DatasetKind.ITER and isinstance(
|
|
data, _IterableDatasetStopIteration
|
|
):
|
|
# if a worker get StopIteration, 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.
|
|
if self._persistent_workers:
|
|
self._worker_status[data.worker_id] = False
|
|
else:
|
|
self._shutdown_worker(data.worker_id)
|
|
self._batches_outstanding -= 1
|
|
self._try_put_indices()
|
|
continue
|
|
|
|
idx, batch, structure = data
|
|
|
|
if (
|
|
isinstance(idx, _ResumeIteration)
|
|
and batch is None
|
|
and structure is None
|
|
):
|
|
return idx
|
|
|
|
if isinstance(batch, _WorkerException):
|
|
self._exit_thread_unexpectedly()
|
|
batch.reraise()
|
|
|
|
if idx == self._rcvd_idx:
|
|
if idx in self._task_infos:
|
|
del self._task_infos[idx]
|
|
self._structure_infos.append(structure)
|
|
return batch
|
|
else:
|
|
self._task_infos[idx] += (batch, structure)
|
|
continue
|
|
|
|
def _try_put_indices(self):
|
|
assert self._batches_outstanding <= self._outstanding_capacity, (
|
|
"too many indices have been put to queue"
|
|
)
|
|
# 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
|
|
|
|
for i in range(self._num_workers):
|
|
worker_idx = next(self._workers_idx_cycle)
|
|
if self._worker_status[worker_idx]:
|
|
break
|
|
else:
|
|
return
|
|
|
|
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
|
|
|
|
def __del__(self):
|
|
self._try_shutdown_all()
|
|
|
|
def _shutdown_on_exit(self):
|
|
self._try_shutdown_all(1)
|
|
|
|
def __next__(self):
|
|
if in_profiler_mode():
|
|
trace_event = profiler.RecordEvent(
|
|
name="_DataLoaderIterMultiProcess",
|
|
event_type=profiler.TracerEventType.Dataloader,
|
|
)
|
|
trace_event.begin()
|
|
try:
|
|
benchmark().check_if_need_record(self)
|
|
benchmark().before_reader()
|
|
# _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):
|
|
if self._persistent_workers:
|
|
raise StopIteration
|
|
else:
|
|
self._thread_done_event.set()
|
|
self._blocking_queue.close()
|
|
|
|
if in_dynamic_mode():
|
|
data = core.eager.read_next_tensor_list(
|
|
self._reader.read_next_list()[0]
|
|
)
|
|
data = _restore_batch(data, self._structure_infos.pop(0))
|
|
else:
|
|
if self._return_list:
|
|
data = self._reader.read_next_list()
|
|
for i in range(len(data)):
|
|
data[i] = data[i]._move_to_list()
|
|
structs = [
|
|
self._structure_infos.pop(0)
|
|
for _ in range(len(self._places))
|
|
]
|
|
data = [_restore_batch(d, s) for d, s in zip(data, structs)]
|
|
# 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()
|
|
benchmark().after_reader()
|
|
return data
|
|
except StopIteration:
|
|
if not self._persistent_workers:
|
|
self._reader.shutdown()
|
|
self._try_shutdown_all()
|
|
raise
|
|
finally:
|
|
if in_profiler_mode():
|
|
trace_event.end()
|
|
|
|
def _on_output_batch(self):
|
|
for _ in range(len(self._places)):
|
|
self._batches_outstanding -= 1
|
|
self._try_put_indices()
|