Start spawn workers concurrently (#1553)

[psaikaushik]

Summary:

Add an opt-in spawn_worker_start_parallelism argument to StatefulDataLoader. For the spawn multiprocessing context, worker Process.start() calls can run concurrently up to this bound. Other contexts and the default value of 1 preserve serial startup.

Build all spawn workers before parallel startup, keep registration order deterministic, and clean up every partially started process and queue if a start fails.

Public EuroSAT benchmark

Dataset: TorchVision EuroSAT, 27,000 real JPEG images across 10 classes. It is credential-free, auto-downloadable from a checksum-pinned URL, and its dataset class subclasses ImageFolder. No dataset content is included in this change.

Transforms: ToDtype(torch.float32, scale=True), Resize(256, antialias=True), CenterCrop(224), and ImageNet normalization, using torchvision.io.decode_image.

Public reproduction from the TorchData repository root:

python -m pip install torchvision psutil
python benchmarks/stateful_dataloader/worker_start_benchmark.py --data-root /tmp/torchdata-eurosat --download --download-only
python benchmarks/stateful_dataloader/worker_start_benchmark.py --data-root /tmp/torchdata-eurosat --batch-size 64 --num-workers 8 --max-batches 100 --spawn-worker-start-parallelism 1
python benchmarks/stateful_dataloader/worker_start_benchmark.py --data-root /tmp/torchdata-eurosat --batch-size 64 --num-workers 8 --max-batches 100 --spawn-worker-start-parallelism 8

Primary comparison: five matched trials in fresh Python processes, alternating p1/p8 ordering. Dataset download/extraction occurred before timing; filesystem cache was warm. Host: Linux, Python 3.12.14, PyTorch 2.15.0a0+fb, TorchVision 0.29.0a0+fb, AMD EPYC 9D64 (88 physical/176 logical cores), 251 GiB RAM, CPU-only.

Metric (median; p10–p90)	Serial p1	Parallel p8	Change
Iterator creation	48.43s (47.88–49.06)	6.84s (6.63–6.88)	85.9% lower
Time to first batch	48.62s (48.07–49.24)	7.10s (6.96–7.16)	85.4% lower
Steady-state images/s, batches 2–100	1693.6 (1667.3–1705.5)	1672.4 (1635.5–1680.3)	1.3% lower
Process-tree PSS increase	4884.9 MiB (4865.8–4892.5)	4996.3 MiB (4984.2–5014.2)	2.3% higher

The deterministic first-batch SHA-256 digest was identical in every run. Steady-state throughput is a no-regression control; the optimization targets startup only.

Visualization

Both series use StatefulDataLoader; this is not Regular DataLoader versus StatefulDataLoader. Blue is the default serial worker startup (spawn_worker_start_parallelism=1), and orange changes only that value to num_workers.

The grouped bars are medians, whiskers are p10–p90, and outlined circles are every raw trial. The primary 8-worker slope panel connects the five matched p1/p8 runs. The 6.85x headline is the ratio of the two medians; the 6.79x value in the paired panel is the median of the five per-pair ratios. The checked-in eurosat_spawn_results.json contains every value used in the figure.

Regenerate the plot from the public checkout with:

python -m pip install matplotlib numpy
python benchmarks/stateful_dataloader/plot_worker_start_benchmark.py \
  --results benchmarks/stateful_dataloader/eurosat_spawn_results.json \
  --output /tmp/eurosat_worker_start_benchmark.png

Worker-count scaling, median time to first batch:

Workers	Serial	Parallel	Reduction	Trials
4	24.51s	6.72s	72.6%	3
8	48.62s	7.10s	85.4%	5
16	96.69s	7.40s	92.3%	3

Differential Revision: D118848297

[meta-codesync]

@psaikaushik has exported this pull request. If you are a Meta employee, you can view the originating Diff in D118848297.