wehub-resource-sync
373 lines
13 KiB
Python
373 lines
13 KiB
Python
"""DataPipe utilities"""
|
|
|
|
import threading
|
|
import time
|
|
|
|
from collections import deque
|
|
from typing import final, List, Set, Type # pylint: disable=no-name-in-module
|
|
|
|
from torch.utils.data import functional_datapipe, IterDataPipe, MapDataPipe
|
|
from torch.utils.data.graph import DataPipe, DataPipeGraph, traverse_dps
|
|
|
|
__all__ = [
|
|
"datapipe_graph_to_adjlist",
|
|
"find_dps",
|
|
"replace_dp",
|
|
"traverse_dps",
|
|
]
|
|
|
|
# Copied from:
|
|
# https://github.com/pytorch/data/blob/88c8bdc6662f37649b7ea5df0bd90a4b24a56876/torchdata/datapipes/iter/util/prefetcher.py#L19-L20
|
|
# Interval between buffer fulfillment checks
|
|
PRODUCER_SLEEP_INTERVAL = 0.0001
|
|
# Interval between checking items availability in buffer
|
|
CONSUMER_SLEEP_INTERVAL = 0.0001
|
|
|
|
|
|
def _get_parents(result_dict, datapipe_graph):
|
|
for k, (v, parents) in datapipe_graph.items():
|
|
if k not in result_dict:
|
|
result_dict[k] = (v, list(parents.keys()))
|
|
_get_parents(result_dict, parents)
|
|
|
|
|
|
def datapipe_graph_to_adjlist(datapipe_graph):
|
|
"""Given a DataPipe graph returned by
|
|
:func:`torch.utils.data.graph.traverse_dps` in DAG form, convert it into
|
|
adjacency list form.
|
|
|
|
Namely, :func:`torch.utils.data.graph.traverse_dps` returns the following
|
|
data structure:
|
|
|
|
.. code::
|
|
|
|
{
|
|
id(datapipe): (
|
|
datapipe,
|
|
{
|
|
id(parent1_of_datapipe): (parent1_of_datapipe, {...}),
|
|
id(parent2_of_datapipe): (parent2_of_datapipe, {...}),
|
|
...
|
|
}
|
|
)
|
|
}
|
|
|
|
We convert it into the following for easier access:
|
|
|
|
.. code::
|
|
|
|
{
|
|
id(datapipe1): (
|
|
datapipe1,
|
|
[id(parent1_of_datapipe1), id(parent2_of_datapipe1), ...]
|
|
),
|
|
id(datapipe2): (
|
|
datapipe2,
|
|
[id(parent1_of_datapipe2), id(parent2_of_datapipe2), ...]
|
|
),
|
|
...
|
|
}
|
|
"""
|
|
|
|
result_dict = {}
|
|
_get_parents(result_dict, datapipe_graph)
|
|
return result_dict
|
|
|
|
|
|
# Copied from:
|
|
# https://github.com/pytorch/data/blob/88c8bdc6662f37649b7ea5df0bd90a4b24a56876/torchdata/dataloader2/graph/utils.py#L16-L35
|
|
def find_dps(graph: DataPipeGraph, dp_type: Type[DataPipe]) -> List[DataPipe]:
|
|
r"""
|
|
Given the graph of DataPipe generated by ``traverse_dps`` function, return DataPipe
|
|
instances with the provided DataPipe type.
|
|
"""
|
|
dps: List[DataPipe] = []
|
|
cache: Set[int] = set()
|
|
|
|
def helper(g) -> None: # pyre-ignore
|
|
for dp_id, (dp, src_graph) in g.items():
|
|
if dp_id in cache:
|
|
continue
|
|
cache.add(dp_id)
|
|
# Please not use `isinstance`, there is a bug.
|
|
if type(dp) is dp_type: # pylint: disable=unidiomatic-typecheck
|
|
dps.append(dp)
|
|
helper(src_graph)
|
|
|
|
helper(graph)
|
|
|
|
return dps
|
|
|
|
|
|
# Copied from:
|
|
# https://github.com/pytorch/data/blob/88c8bdc6662f37649b7ea5df0bd90a4b24a56876/torchdata/dataloader2/graph/utils.py#L82-L97
|
|
# Given the DataPipe needs to be replaced and the expected DataPipe, return a new graph
|
|
def replace_dp(
|
|
graph: DataPipeGraph, old_datapipe: DataPipe, new_datapipe: DataPipe
|
|
) -> DataPipeGraph:
|
|
r"""
|
|
Given the graph of DataPipe generated by ``traverse_dps`` function and the
|
|
DataPipe to be replaced and the new DataPipe, return the new graph of
|
|
DataPipe.
|
|
"""
|
|
assert len(graph) == 1
|
|
|
|
if id(old_datapipe) in graph:
|
|
graph = traverse_dps(new_datapipe)
|
|
|
|
final_datapipe = list(graph.values())[0][0]
|
|
|
|
for recv_dp, send_graph in graph.values():
|
|
_replace_dp(recv_dp, send_graph, old_datapipe, new_datapipe)
|
|
|
|
return traverse_dps(final_datapipe)
|
|
|
|
|
|
# For each `recv_dp`, find if the source_datapipe needs to be replaced by the new one.
|
|
# If found, find where the `old_dp` is located in `recv_dp` and switch it to the `new_dp`
|
|
def _replace_dp(
|
|
recv_dp, send_graph: DataPipeGraph, old_dp: DataPipe, new_dp: DataPipe
|
|
) -> None:
|
|
old_dp_id = id(old_dp)
|
|
for send_id in send_graph:
|
|
if send_id == old_dp_id:
|
|
_assign_attr(recv_dp, old_dp, new_dp, inner_dp=True)
|
|
else:
|
|
send_dp, sub_send_graph = send_graph[send_id]
|
|
_replace_dp(send_dp, sub_send_graph, old_dp, new_dp)
|
|
|
|
|
|
# Recursively re-assign datapipe for the sake of nested data structure
|
|
# `inner_dp` is used to prevent recursive call if we have already met a `DataPipe`
|
|
def _assign_attr(obj, old_dp, new_dp, inner_dp: bool = False):
|
|
if obj is old_dp:
|
|
return new_dp
|
|
elif isinstance(obj, (IterDataPipe, MapDataPipe)):
|
|
# Prevent recursive call for DataPipe
|
|
if not inner_dp:
|
|
return None
|
|
for k in list(obj.__dict__.keys()):
|
|
new_obj = _assign_attr(obj.__dict__[k], old_dp, new_dp)
|
|
if new_obj is not None:
|
|
obj.__dict__[k] = new_obj
|
|
break
|
|
return None
|
|
elif isinstance(obj, dict):
|
|
for k in list(obj.keys()):
|
|
new_obj = _assign_attr(obj[k], old_dp, new_dp)
|
|
if new_obj is not None:
|
|
obj[k] = new_obj
|
|
break
|
|
return None
|
|
# Tuple is immutable, has to re-create a tuple
|
|
elif isinstance(obj, tuple):
|
|
temp_list = []
|
|
flag = False
|
|
for item in obj:
|
|
new_obj = _assign_attr(item, old_dp, new_dp, inner_dp)
|
|
if new_obj is not None:
|
|
flag = True
|
|
temp_list.append(new_dp)
|
|
else:
|
|
temp_list.append(item)
|
|
if flag:
|
|
return tuple(temp_list) # Special case
|
|
else:
|
|
return None
|
|
elif isinstance(obj, list):
|
|
for i in range(len(obj)): # pylint: disable=consider-using-enumerate
|
|
new_obj = _assign_attr(obj[i], old_dp, new_dp, inner_dp)
|
|
if new_obj is not None:
|
|
obj[i] = new_obj
|
|
break
|
|
return None
|
|
elif isinstance(obj, set):
|
|
new_obj = None
|
|
for item in obj:
|
|
if _assign_attr(item, old_dp, new_dp, inner_dp) is not None:
|
|
new_obj = new_dp
|
|
break
|
|
if new_obj is not None:
|
|
obj.remove(old_dp)
|
|
obj.add(new_dp)
|
|
return None
|
|
else:
|
|
return None
|
|
|
|
|
|
class _PrefetchData:
|
|
def __init__(self, source_datapipe, buffer_size: int):
|
|
self.run_prefetcher: bool = True
|
|
self.prefetch_buffer: Deque = deque()
|
|
self.buffer_size: int = buffer_size
|
|
self.source_datapipe = source_datapipe
|
|
self.stop_iteration: bool = False
|
|
self.paused: bool = False
|
|
|
|
|
|
# Copied from:
|
|
# https://github.com/pytorch/data/blob/88c8bdc6662f37649b7ea5df0bd90a4b24a56876/torchdata/datapipes/iter/util/prefetcher.py#L34-L172
|
|
@functional_datapipe("prefetch")
|
|
class PrefetcherIterDataPipe(IterDataPipe):
|
|
r"""
|
|
Prefetches elements from the source DataPipe and puts them into a buffer
|
|
(functional name: ``prefetch``). Prefetching performs the operations (e.g.
|
|
I/O, computations) of the DataPipes up to this one ahead of time and stores
|
|
the result in the buffer, ready to be consumed by the subsequent DataPipe.
|
|
It has no effect aside from getting the sample ready ahead of time.
|
|
|
|
This is used by ``MultiProcessingReadingService`` when the arguments
|
|
``worker_prefetch_cnt`` (for prefetching at each worker process) or
|
|
``main_prefetch_cnt`` (for prefetching at the main loop) are greater than 0.
|
|
|
|
Beyond the built-in use cases, this can be useful to put after I/O DataPipes
|
|
that have expensive I/O operations (e.g. takes a long time to request a file
|
|
from a remote server).
|
|
|
|
Args:
|
|
source_datapipe: IterDataPipe from which samples are prefetched
|
|
buffer_size: the size of the buffer which stores the prefetched samples
|
|
|
|
Example:
|
|
>>> from torchdata.datapipes.iter import IterableWrapper
|
|
>>> dp = IterableWrapper(file_paths).open_files().prefetch(5)
|
|
"""
|
|
|
|
def __init__(self, source_datapipe, buffer_size: int = 10):
|
|
self.source_datapipe = source_datapipe
|
|
if buffer_size <= 0:
|
|
raise ValueError(
|
|
"'buffer_size' is required to be a positive integer."
|
|
)
|
|
self.buffer_size = buffer_size
|
|
self.thread: Optional[threading.Thread] = None
|
|
self.prefetch_data: Optional[_PrefetchData] = None
|
|
|
|
@staticmethod
|
|
def thread_worker(
|
|
prefetch_data: _PrefetchData,
|
|
): # pylint: disable=missing-function-docstring
|
|
itr = iter(prefetch_data.source_datapipe)
|
|
while not prefetch_data.stop_iteration:
|
|
# Run if not paused
|
|
while prefetch_data.run_prefetcher:
|
|
if (
|
|
len(prefetch_data.prefetch_buffer)
|
|
< prefetch_data.buffer_size
|
|
):
|
|
try:
|
|
item = next(itr)
|
|
prefetch_data.prefetch_buffer.append(item)
|
|
except Exception as e: # pylint: disable=broad-except
|
|
prefetch_data.run_prefetcher = False
|
|
prefetch_data.stop_iteration = True
|
|
prefetch_data.prefetch_buffer.append(e)
|
|
else: # Buffer is full, waiting for main thread to consume items
|
|
# TODO: Calculate sleep interval based on previous consumption speed
|
|
time.sleep(PRODUCER_SLEEP_INTERVAL)
|
|
prefetch_data.paused = True
|
|
# Sleep longer when this prefetcher thread is paused
|
|
time.sleep(PRODUCER_SLEEP_INTERVAL * 10)
|
|
|
|
def __iter__(self):
|
|
try:
|
|
prefetch_data = _PrefetchData(
|
|
self.source_datapipe, self.buffer_size
|
|
)
|
|
self.prefetch_data = prefetch_data
|
|
thread = threading.Thread(
|
|
target=PrefetcherIterDataPipe.thread_worker,
|
|
args=(prefetch_data,),
|
|
daemon=True,
|
|
)
|
|
thread.start()
|
|
self.thread = thread
|
|
|
|
while (
|
|
not prefetch_data.stop_iteration
|
|
or len(prefetch_data.prefetch_buffer) > 0
|
|
):
|
|
if len(prefetch_data.prefetch_buffer) > 0:
|
|
data = prefetch_data.prefetch_buffer.popleft()
|
|
if isinstance(data, Exception):
|
|
if isinstance(data, StopIteration):
|
|
break
|
|
raise data
|
|
yield data
|
|
else:
|
|
time.sleep(CONSUMER_SLEEP_INTERVAL)
|
|
finally:
|
|
if "prefetch_data" in locals():
|
|
prefetch_data.run_prefetcher = False
|
|
prefetch_data.stop_iteration = True
|
|
prefetch_data.paused = False
|
|
if "thread" in locals():
|
|
thread.join()
|
|
|
|
def __getstate__(self):
|
|
"""
|
|
Getting state in threading environment requires next operations:
|
|
1) Stopping of the producer thread.
|
|
2) Saving buffer.
|
|
3) Adding lazy restart of producer thread when __next__ is called again
|
|
(this will guarantee that you only change state of the source_datapipe
|
|
after entire state of the graph is saved).
|
|
"""
|
|
# TODO: Update __getstate__ and __setstate__ to support snapshotting and restoration
|
|
return {
|
|
"source_datapipe": self.source_datapipe,
|
|
"buffer_size": self.buffer_size,
|
|
}
|
|
|
|
def __setstate__(self, state):
|
|
self.source_datapipe = state["source_datapipe"]
|
|
self.buffer_size = state["buffer_size"]
|
|
self.thread = None
|
|
|
|
@final
|
|
def reset(self): # pylint: disable=missing-function-docstring
|
|
self.shutdown()
|
|
|
|
def pause(self): # pylint: disable=missing-function-docstring
|
|
if self.thread is not None:
|
|
assert self.prefetch_data is not None
|
|
self.prefetch_data.run_prefetcher = False
|
|
if self.thread.is_alive():
|
|
# Blocking until the thread is paused
|
|
while not self.prefetch_data.paused:
|
|
time.sleep(PRODUCER_SLEEP_INTERVAL * 10)
|
|
|
|
@final
|
|
def resume(self): # pylint: disable=missing-function-docstring
|
|
if (
|
|
self.thread is not None
|
|
and self.prefetch_data is not None
|
|
and (
|
|
not self.prefetch_data.stop_iteration
|
|
or len(self.prefetch_data.prefetch_buffer) > 0
|
|
)
|
|
):
|
|
self.prefetch_data.run_prefetcher = True
|
|
self.prefetch_data.paused = False
|
|
|
|
@final
|
|
def shutdown(self): # pylint: disable=missing-function-docstring
|
|
if hasattr(self, "prefetch_data") and self.prefetch_data is not None:
|
|
self.prefetch_data.run_prefetcher = False
|
|
self.prefetch_data.stop_iteration = True
|
|
self.prefetch_data.paused = False
|
|
self.prefetch_data = None
|
|
if hasattr(self, "thread") and self.thread is not None:
|
|
self.thread.join()
|
|
self.thread = None
|
|
|
|
def __del__(self):
|
|
self.shutdown()
|
|
|
|
def __len__(self) -> int:
|
|
if isinstance(self.source_datapipe, Sized):
|
|
return len(self.source_datapipe)
|
|
raise TypeError(
|
|
f"{type(self).__name__} instance doesn't have valid length"
|
|
)
|