Dataloader is slow with iterdatapipes and shuffle that has large in-memory fields (because traverse_dps is slow)

🐛 Describe the bug

Hello, I found that a standard DataLoader takes unreasonably long to construct itself and to load the first batch if there is a filed in a dataset that takes long to pickle (e.g. an in-memory dataset with panda frame and strings).

This happens only if shuffle is true and the datapipe is an IterDataPipe. The dataloader calls apply_shuffle_settings which in turn calls traverse_dps, then _list_connected_datapipes which eventually pickles all object fields in a dataset. I was not able to comprehend why would one need to pickle datafields to build a datapipe graph.

Here is a reproduction:

import torch
import time
import pandas as pd
import numpy as np
from torchdata.datapipes.map import MapDataPipe


class TabularDataset(MapDataPipe):
    """ Make MapDataPipe so that shuffle is fast"""
    def __init__(self, data):
        self._data = data

    def __getitem__(self, idx: int):
        return  self._data['complex_objects'].iloc[idx]

    def __len__(self):
        return len(self._data)
    

# Large panda frames with strings take a long time to pickle
n_rows = 1000000
data = pd.DataFrame({
    'complex_objects': [
        np.array(["1", "2", "3", "4"]) 
        if r % 2 == 0 else np.array(["4", "3", "2", "1"])
        for r in range(n_rows)
    ]
})

dataset = TabularDataset(data).shuffle()
start_time = time.time()
dataloader = torch.utils.data.DataLoader(
    dataset, shuffle=True, batch_size=2, num_workers=0, collate_fn=lambda x: x
)
print(f"Dataloader creation time: {time.time() - start_time:.2f}s")
for b in dataloader:
    assert len(b) == 2
    print(b)
    print(f"Time to get the first batch: {time.time() - start_time:.2f}s")
    break

This code prints out:

Dataloader creation time: 7.44s
[array(['4', '3', '2', '1'], dtype='<U1'), array(['4', '3', '2', '1'], dtype='<U1')]
Time to get the first batch: 22.85s

In my case, I have even slower datapipe that takes 5 minutes to pickle.

A workaround

One can make a datapipe that doesn't take long to pickle (e.g. with lambdas):

class TabularDatasetIndirect(MapDataPipe):
    """ Indirect access to data"""
    def __init__(self, data):
        self._data_f = lambda: data

    def __getitem__(self, idx: int):
        return  self._data_f()['complex_objects'].iloc[idx]

    def __len__(self):
        return len(self._data_f())

In that case, the printout is

Dataloader creation time: 0.16s
[array(['1', '2', '3', '4'], dtype='<U1'), array(['4', '3', '2', '1'], dtype='<U1')]
Time to get the first batch: 0.75s

which is a 30x speedup in this case.

Versions

PyTorch version: 2.1.2
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.6 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: 13.0.1-++20220120110924+75e33f71c2da-1~exp1~20220120231001.58
CMake version: version 3.16.3
Libc version: glibc-2.31

Python version: 3.9.5 (default, Nov 23 2021, 15:27:38)  [GCC 9.3.0] (64-bit runtime)
Python platform: Linux-5.15.0-1038-azure-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 11.8.89
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA A100-SXM4-80GB
GPU 1: NVIDIA A100-SXM4-80GB
GPU 2: NVIDIA A100-SXM4-80GB
GPU 3: NVIDIA A100-SXM4-80GB
GPU 4: NVIDIA A100-SXM4-80GB
GPU 5: NVIDIA A100-SXM4-80GB
GPU 6: NVIDIA A100-SXM4-80GB
GPU 7: NVIDIA A100-SXM4-80GB

Nvidia driver version: 535.129.03
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.9.5
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.9.5
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.9.5
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.9.5
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.9.5
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.9.5
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.9.5
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                    x86_64
CPU op-mode(s):                  32-bit, 64-bit
Byte Order:                      Little Endian
Address sizes:                   48 bits physical, 48 bits virtual
CPU(s):                          96
On-line CPU(s) list:             0-95
Thread(s) per core:              1
Core(s) per socket:              48
Socket(s):                       2
NUMA node(s):                    4
Vendor ID:                       AuthenticAMD
CPU family:                      23
Model:                           49
Model name:                      AMD EPYC 7V12 64-Core Processor
Stepping:                        0
CPU MHz:                         2445.440
BogoMIPS:                        4890.88
Hypervisor vendor:               Microsoft
Virtualization type:             full
L1d cache:                       3 MiB
L1i cache:                       3 MiB
L2 cache:                        48 MiB
L3 cache:                        384 MiB
NUMA node0 CPU(s):               0-23
NUMA node1 CPU(s):               24-47
NUMA node2 CPU(s):               48-71
NUMA node3 CPU(s):               72-95
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Mmio stale data:   Not affected
Vulnerability Retbleed:          Mitigation; untrained return thunk; SMT disabled
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; Retpolines, STIBP disabled, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl tsc_reliable nonstop_tsc cpuid extd_apicid aperfmperf pni pclmulqdq ssse3 fma cx16 sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm cmp_legacy cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw topoext perfctr_core ssbd vmmcall fsgsbase bmi1 avx2 smep bmi2 rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves clzero xsaveerptr rdpru arat umip rdpid

Versions of relevant libraries:
[pip3] efficientnet-pytorch==0.7.1
[pip3] numpy==1.25.2
[pip3] onnx==1.15.0
[pip3] open-clip-torch==2.19.0
[pip3] pytorch-lightning==2.1.2
[pip3] torch==2.1.2
[pip3] torch-fidelity==0.3.0
[pip3] torchdata==0.7.1
[pip3] torchmetrics==1.3.0.post0
[pip3] torchvision==0.16.1
[pip3] triton==2.2.0
[conda] Could not collect

cc @SsnL @VitalyFedyunin @ejguan @dzhulgakov

No repository file or test is named. Start with DataLoader's apply_shuffle_settings path and trace traverse_dps into _list_connected_datapipes, then run the supplied reproduction to compare construction and first-batch timings; done means the large in-memory fields no longer cause the reported traversal slowdown while shuffle behavior remains correct.