Implementation:Iterative Dvc Utils Threadpool
| Knowledge Sources | |
|---|---|
| Domains | Concurrency, Performance |
| Last Updated | 2026-02-10 10:00 GMT |
Overview
Extended ThreadPoolExecutor with lazy unordered mapping and error-aware cancellation support, provided by the DVC library.
Description
The ThreadPoolExecutor class in dvc/utils/threadpool.py extends Python's standard concurrent.futures.ThreadPoolExecutor with two key enhancements designed for DVC's data transfer and processing workloads.
The first enhancement is imap_unordered(fn, *iterables), a lazy unordered map operation that processes items from one or more iterables using the thread pool while maintaining bounded memory usage. Unlike the standard map() method, which eagerly creates futures for all items, imap_unordered maintains a rolling window of at most max_workers * 5 pending futures at any time. It works as follows:
- The input iterables are zipped together into an iterator of argument tuples.
- An initial batch of
max_workers * 5futures is submitted. - The method enters a loop where it waits for any future to complete (using
futures.waitwithFIRST_COMPLETED). - As each future completes, its result is yielded to the caller, and new futures are submitted to replace the completed ones (maintaining the buffer size).
- The loop terminates when all futures have been processed and the input iterator is exhausted.
This design provides backpressure: if the consumer processes results slowly, the producer does not accumulate an unbounded number of pending futures. The factor of 5 ensures enough concurrency to keep all worker threads busy even when individual task durations vary.
The second enhancement is cancel_on_error support. The constructor accepts a cancel_on_error boolean parameter. When set to True and the executor is used as a context manager, if an exception propagates out of the with block, the __exit__ method calls self.shutdown(wait=True, cancel_futures=True) to cancel all pending futures before shutting down. When False (or when no exception occurs), the executor shuts down normally, waiting for all submitted futures to complete.
Usage
Import ThreadPoolExecutor as a drop-in replacement for concurrent.futures.ThreadPoolExecutor when you need lazy iteration over concurrent results or error-aware cleanup. It is used throughout DVC for parallel file transfers, hash computations, and remote storage operations where processing thousands of files requires bounded memory usage.
Code Reference
Source Location
- Repository: DVC
- File:
dvc/utils/threadpool.py - Lines: 42 total
- Class: ThreadPoolExecutor (L9-41)
Signature
class ThreadPoolExecutor(futures.ThreadPoolExecutor):
def __init__(
self,
max_workers: Optional[int] = None,
cancel_on_error: bool = False,
**kwargs,
):
...
def imap_unordered(
self, fn: Callable[..., _T], *iterables: Iterable[Any]
) -> Iterator[_T]:
...
def __exit__(self, exc_type, exc_val, exc_tb):
...
Import
from dvc.utils.threadpool import ThreadPoolExecutor
I/O Contract
ThreadPoolExecutor.__init__
| Name | Type | Required | Description |
|---|---|---|---|
| max_workers | Optional[int] | No | Maximum number of worker threads. Defaults to None (Python's default, which is min(32, os.cpu_count() + 4) in Python 3.8+).
|
| cancel_on_error | bool | No | If True, pending futures are cancelled when the context manager exits with an exception. Defaults to False. |
| **kwargs | Any | No | Additional keyword arguments passed to the standard ThreadPoolExecutor (e.g., thread_name_prefix, initializer). |
ThreadPoolExecutor.imap_unordered
| Name | Type | Required | Description |
|---|---|---|---|
| fn | Callable[..., T] | Yes | The function to apply to each set of arguments. |
| *iterables | Iterable[Any] | Yes | One or more iterables whose elements are zipped and passed as positional arguments to fn. |
| Name | Type | Description |
|---|---|---|
| return | Iterator[T] | A lazy iterator yielding results as they complete, in no guaranteed order. Maintains a buffer of at most max_workers * 5 pending futures.
|
Exceptions:
- Any exception raised by fn inside a worker thread is re-raised in the calling thread when the corresponding future's result is accessed via yield fut.result().
ThreadPoolExecutor.__exit__
| Name | Type | Description |
|---|---|---|
| return | bool | Always returns False, meaning exceptions are not suppressed and will propagate to the caller. |
Usage Examples
Basic Parallel Processing
from dvc.utils.threadpool import ThreadPoolExecutor
def process_file(path):
# Simulate file processing
return f"processed: {path}"
files = [f"data/file_{i}.csv" for i in range(1000)]
with ThreadPoolExecutor(max_workers=8) as executor:
for result in executor.imap_unordered(process_file, files):
print(result)
# Only ~40 futures (8 * 5) are pending at any time
Multiple Iterables
from dvc.utils.threadpool import ThreadPoolExecutor
def upload(local_path, remote_path):
# Simulate file upload
return f"{local_path} -> {remote_path}"
local_paths = ["file1.bin", "file2.bin", "file3.bin"]
remote_paths = ["s3://bucket/file1.bin", "s3://bucket/file2.bin", "s3://bucket/file3.bin"]
with ThreadPoolExecutor(max_workers=4) as executor:
for result in executor.imap_unordered(upload, local_paths, remote_paths):
print(result)
Error-Aware Cancellation
from dvc.utils.threadpool import ThreadPoolExecutor
def download(url):
if "error" in url:
raise ConnectionError(f"Failed to download {url}")
return f"downloaded: {url}"
urls = [f"https://example.com/file_{i}" for i in range(100)]
try:
with ThreadPoolExecutor(max_workers=4, cancel_on_error=True) as executor:
for result in executor.imap_unordered(download, urls):
print(result)
except ConnectionError as e:
print(f"Transfer failed: {e}")
# All pending futures are automatically cancelled
Internal Details
Buffer Size Strategy
The buffer size of max_workers * 5 is a deliberate trade-off between memory usage and throughput. A buffer of exactly max_workers would mean that all workers could be idle while waiting for new submissions after completing a batch. The 5x multiplier ensures that there are always enough pending tasks to keep workers busy even when individual task durations vary significantly, while still providing effective backpressure for very large input streams (e.g., tens of thousands of files).
Comparison with Standard map()
| Feature | Standard map() | imap_unordered() |
|---|---|---|
| Future creation | Eager (all at once) | Lazy (rolling window) |
| Result ordering | Preserved | Not preserved |
| Memory usage | O(n) futures | O(max_workers * 5) futures |
| Backpressure | None | Yes (consumer-driven) |