diff --git a/dvc/data/transfer.py b/dvc/data/transfer.py index 69a9c0c75..1b56a307f 100644 --- a/dvc/data/transfer.py +++ b/dvc/data/transfer.py @@ -5,8 +5,8 @@ from typing import TYPE_CHECKING, Any, Callable, Iterable, Optional from funcy import split +from dvc.objects.executors import ThreadPoolExecutor from dvc.progress import Tqdm -from dvc.utils.threadpool import ThreadPoolExecutor if TYPE_CHECKING: from dvc.objects.db import ObjectDB diff --git a/dvc/objects/executors.py b/dvc/objects/executors.py new file mode 100644 index 000000000..dde15ba6e --- /dev/null +++ b/dvc/objects/executors.py @@ -0,0 +1,74 @@ +import queue +import sys +from concurrent import futures +from itertools import islice +from typing import Any, Callable, Iterable, Iterator, Set, TypeVar + +_T = TypeVar("_T") + + +class ThreadPoolExecutor(futures.ThreadPoolExecutor): + _max_workers: int + + def __init__( + self, max_workers: int = None, cancel_on_error: bool = False, **kwargs + ): + super().__init__(max_workers=max_workers, **kwargs) + self._cancel_on_error = cancel_on_error + + @property + def max_workers(self) -> int: + return self._max_workers + + def imap_unordered( + self, fn: Callable[..., _T], *iterables: Iterable[Any] + ) -> Iterator[_T]: + """Lazier version of map that does not preserve ordering of results. + + It does not create all the futures at once to reduce memory usage. + """ + + def create_taskset(n: int) -> Set[futures.Future]: + return {self.submit(fn, *args) for args in islice(it, n)} + + it = zip(*iterables) + tasks = create_taskset(self.max_workers * 5) + while tasks: + done, tasks = futures.wait( + tasks, return_when=futures.FIRST_COMPLETED + ) + for fut in done: + yield fut.result() + tasks.update(create_taskset(len(done))) + + def shutdown(self, wait=True, *, cancel_futures=False): + if sys.version_info > (3, 9): + # pylint: disable=unexpected-keyword-arg + return super().shutdown(wait=wait, cancel_futures=cancel_futures) + + with self._shutdown_lock: + self._shutdown = True + if cancel_futures: + # Drain all work items from the queue, and then cancel their + # associated futures. + while True: + try: + work_item = self._work_queue.get_nowait() + except queue.Empty: + break + if work_item is not None: + work_item.future.cancel() + + # Send a wake-up to prevent threads calling + # _work_queue.get(block=True) from permanently blocking. + self._work_queue.put(None) + if wait: + for t in self._threads: + t.join() + + def __exit__(self, exc_type, exc_val, exc_tb): + if self._cancel_on_error: + self.shutdown(wait=True, cancel_futures=exc_val is not None) + else: + self.shutdown(wait=True) + return False diff --git a/dvc/objects/fs/base.py b/dvc/objects/fs/base.py index 47b528c42..22efd61fa 100644 --- a/dvc/objects/fs/base.py +++ b/dvc/objects/fs/base.py @@ -20,8 +20,8 @@ from typing import ( from funcy import cached_property from dvc.exceptions import DvcException -from dvc.utils.threadpool import ThreadPoolExecutor +from ..executors import ThreadPoolExecutor from ._callback import DEFAULT_CALLBACK, FsspecCallback if TYPE_CHECKING: