Dataloader Hangs on Exit

🐛 Describe the bug

I am running into an issue where the dataloaders hang for 5 seconds per dataloader worker on exit. I have persistent_workers=True and pin_memory=True.

Here is a minimal example to reproduce the bug:

import torch
import torch.utils.data
import torchvision
import time
import atexit

torch.cuda.set_device(0)

def atexit_hook():
  print(str(time.time()) + " atexit ran")

atexit.register(atexit_hook)

num_workers = 4  # I tried setting this on [1, 2, 4, 8, 16]

dataloader = torch.utils.data.DataLoader(
    torchvision.datasets.MNIST(
            "~/datasets",
            train=True,
            download=True,
            transform=torchvision.transforms.Compose([torchvision.transforms.ToTensor()]),
    ),
    pin_memory=True,
    batch_size=1024,
    persistent_workers=True,
    num_workers=num_workers,
)

for _ in dataloader:
    break

print(str(time.time()) + " finishing")

I varied the number of dataloader workers on [1, 2, 4, 8, 16]:

# Test bench

wget -O dl_bug_1.py https://gist.githubusercontent.com/ravi-mosaicml/3c552c3f6d4a6abf9621fa56353f0fbf/raw/0bf88f6b392b90f1b22223d8dfc9b8771850de39/dl_bug.py
wget -O dl_bug_2.py https://gist.githubusercontent.com/ravi-mosaicml/c277fe6ea49fd43de1db1687a5b868f8/raw/1a6ac44299504097694b9720db2df98f0e96bb70/dl_bug_2.py
wget -O dl_bug_4.py https://gist.githubusercontent.com/ravi-mosaicml/c0a1f2268d47f3e27d8e360adf604d22/raw/29190a9054d0874e082066807b4247aa280dc076/dl_bug_4.py
wget -O dl_bug_8.py https://gist.githubusercontent.com/ravi-mosaicml/901ad8423421aa5157acc52ac9b17889/raw/44454a6b52f6dc7b26d7bdf3d42bf289163fefc0/dl_bug_8.py
wget -O dl_bug_16.py https://gist.githubusercontent.com/ravi-mosaicml/225dfb4045cd847b7e095f83c6ed4e6a/raw/b1cf3e70dc339daf073f1c926c7c8e67f0c5948c/dl_bug_16.py
python dl_bug_1.py
python dl_bug_2.py
python dl_bug_4.py
python dl_bug_8.py
python dl_bug_16.py

Output:

+ python dl_bug_1.py
...
1658340836.637308 finishing
1658340841.6431446 atexit ran
+ python dl_bug_2.py
1658340844.0869184 finishing
1658340854.0987492 atexit ran
+ python dl_bug_4.py
1658340856.7868223 finishing
1658340876.8072598 atexit ran
+ python dl_bug_8.py
1658340879.3098278 finishing
1658340919.3564277 atexit ran
+ python dl_bug_16.py
1658340921.9546304 finishing
1658341002.0392313 atexit ran

When analyzing the results, there appears to be a near-perfect 5 second delay per dataloader worker between when the python script "finished" and when atexit ran. The python process terminated immediately after the "atexit ran" lines were printed.

The expected behavior would be that there is a small (constant) overhead at the end of training to shut down the dataloader workers. It should not scale linearly (at least not linearly in any meaningful way, since processes can be killed in parallel) with the number of DL workers.

Versions

Collecting environment information...
PyTorch version: 1.12.0+cu113
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A
OS: Ubuntu 20.04.4 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.31
Python version: 3.9.13 (main, May 23 2022, 22:01:06)  [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.4.0-120-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 11.3.109
GPU models and configuration: GPU 0: NVIDIA A100-SXM4-40GB
Nvidia driver version: 510.47.03
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.2.0
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True
Versions of relevant libraries:
[pip3] numpy==1.23.1
[pip3] torch==1.12.0+cu113
[pip3] torchvision==0.13.0+cu113
[conda] Could not collect

cc @SsnL @VitalyFedyunin @ejguan @NivekT

Start at the DataLoader shutdown path, focusing on the interaction between persistent_workers and pin_memory. Run the minimal reproduction and compare the exit timings for different worker counts. Done means worker shutdown no longer adds roughly five seconds per worker, with regression coverage for the reported behavior.