Streaming from s3 hangs when num_workers > 1

[awaelchli]

🐛 Bug

The PR #237 broke streaming from S3.

To Reproduce

from litdata.streaming import StreamingDataLoader, StreamingDataset, TokensLoader

if __name__ == "__main__":
    train_dataset = StreamingDataset(
        # input_dir="s3://tinyllama-template/slimpajama/train",
        input_dir="/teamspace/s3_connections/tinyllama-template/slimpajama/train",
        item_loader=TokensLoader(block_size=128),
    )
    train_dataloader = StreamingDataLoader(train_dataset, shuffle=True, batch_size=1, num_workers=4)

    for batch in train_dataloader:  # hangs
        print(batch)
        break

The next() call on the iterator hangs.
Only hangs with num_workers > 1.

Expected behavior

Returns data.
Before the commit #237 it works.

Environment

• PyTorch Version (e.g., 1.0): 2.3
• OS (e.g., Linux): Linux
• How you installed PyTorch (conda, pip, source): pip
• Build command you used (if compiling from source):
• Python version: 3.11
• CUDA/cuDNN version:
• GPU models and configuration:
• Any other relevant information:

Additional context