StreamingDataset and StreamingDataLoader hangs after cancelling once

[Aceticia]

🐛 Bug

I have an optimized dataset on a slurm HPC that I optimized with zstd compression, and 128MB chunk size, and the latest litdata. I can successfully train a model with a StreamingDataset and a StreamingDataLoader right after creating the optimized dataset with fabric. But after the first run is cancelled, I find that the dataset will always get stuck at a specific item. Explicitly setting the cache directory to a new, empty directory doesn't seem to help.

I tried optimizing again, and interestingly the same behavior happens again. i.e., I can train it once, and cancelling it causes all further usages stuck.

The way I use the dataset is:

fabric.setup()
train_loader = StreamingDataLoader(
    StreamingDataset(
        str(dataset_path),
        cache_dir=str(local_cache_dir),
        shuffle=True,
    ),
    num_workers=data_config_dict["num_workers"],
    batch_size=data_config_dict["train_batch_size"],
    drop_last=True,
    pin_memory=True,
)
train_loader = fabric.setup_dataloader(train_loader)
for batch in train_loader:
    ...

Any ideas on what other things to test?

To Reproduce

I tried to simply manually iterate over all samples with a python loop without the DataLoader, this seems to finish running. However, if I use a StreamingDataLoader and manually iterate through it, it will indeed be stuck at some point. By stuck, I mean getting stuck for more than 30 minutes. If I interrupt at this point, the trace looks like this:

  File "<stdin>", line 1, in <module>
  File "/gpfs/data/oermannlab/users/xl3942/.conda/envs/neuro/lib/python3.12/site-packages/tqdm/std.py", line 1181, in __iter__
    for obj in iterable:
               ^^^^^^^^
  File "/gpfs/data/oermannlab/users/xl3942/.conda/envs/neuro/lib/python3.12/site-packages/litdata/streaming/dataloader.py", line 622, in __iter__
    for batch in super().__iter__():
                 ^^^^^^^^^^^^^^^^^^
  File "/gpfs/data/oermannlab/users/xl3942/.conda/envs/neuro/lib/python3.12/site-packages/torch/utils/data/dataloader.py", line 701, in __next__
    data = self._next_data()
           ^^^^^^^^^^^^^^^^^
  File "/gpfs/data/oermannlab/users/xl3942/.conda/envs/neuro/lib/python3.12/site-packages/torch/utils/data/dataloader.py", line 757, in _next_data
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/gpfs/data/oermannlab/users/xl3942/.conda/envs/neuro/lib/python3.12/site-packages/torch/utils/data/_utils/fetch.py", line 33, in fetch
    data.append(next(self.dataset_iter))
                ^^^^^^^^^^^^^^^^^^^^^^^
  File "/gpfs/data/oermannlab/users/xl3942/.conda/envs/neuro/lib/python3.12/site-packages/litdata/streaming/dataset.py", line 381, in __next__
    data = self.__getitem__(
           ^^^^^^^^^^^^^^^^^
  File "/gpfs/data/oermannlab/users/xl3942/.conda/envs/neuro/lib/python3.12/site-packages/litdata/streaming/dataset.py", line 349, in __getitem__
    return self.cache[index]
           ~~~~~~~~~~^^^^^^^
  File "/gpfs/data/oermannlab/users/xl3942/.conda/envs/neuro/lib/python3.12/site-packages/litdata/streaming/cache.py", line 145, in __getitem__
    return self._reader.read(index)
           ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/gpfs/data/oermannlab/users/xl3942/.conda/envs/neuro/lib/python3.12/site-packages/litdata/streaming/reader.py", line 287, in read
    item = self._item_loader.load_item_from_chunk(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/gpfs/data/oermannlab/users/xl3942/.conda/envs/neuro/lib/python3.12/site-packages/litdata/streaming/item_loader.py", line 145, in load_item_from_chunk
    sleep(0.1)

or this:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/gpfs/data/oermannlab/users/xl3942/.conda/envs/neuro/lib/python3.12/site-packages/tqdm/std.py", line 1181, in __iter__
    for obj in iterable:
               ^^^^^^^^
  File "/gpfs/data/oermannlab/users/xl3942/.conda/envs/neuro/lib/python3.12/site-packages/litdata/streaming/dataloader.py", line 622, in __iter__
    for batch in super().__iter__():
                 ^^^^^^^^^^^^^^^^^^
  File "/gpfs/data/oermannlab/users/xl3942/.conda/envs/neuro/lib/python3.12/site-packages/torch/utils/data/dataloader.py", line 701, in __next__
    data = self._next_data()
           ^^^^^^^^^^^^^^^^^
  File "/gpfs/data/oermannlab/users/xl3942/.conda/envs/neuro/lib/python3.12/site-packages/torch/utils/data/dataloader.py", line 1448, in _next_data
    idx, data = self._get_data()
                ^^^^^^^^^^^^^^^^
  File "/gpfs/data/oermannlab/users/xl3942/.conda/envs/neuro/lib/python3.12/site-packages/torch/utils/data/dataloader.py", line 1412, in _get_data
    success, data = self._try_get_data()
                    ^^^^^^^^^^^^^^^^^^^^
  File "/gpfs/data/oermannlab/users/xl3942/.conda/envs/neuro/lib/python3.12/site-packages/torch/utils/data/dataloader.py", line 1243, in _try_get_data
    data = self._data_queue.get(timeout=timeout)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/gpfs/data/oermannlab/users/xl3942/.conda/envs/neuro/lib/python3.12/multiprocessing/queues.py", line 113, in get
    if not self._poll(timeout):
           ^^^^^^^^^^^^^^^^^^^
  File "/gpfs/data/oermannlab/users/xl3942/.conda/envs/neuro/lib/python3.12/multiprocessing/connection.py", line 257, in poll
    return self._poll(timeout)
           ^^^^^^^^^^^^^^^^^^^
  File "/gpfs/data/oermannlab/users/xl3942/.conda/envs/neuro/lib/python3.12/multiprocessing/connection.py", line 440, in _poll
    r = wait([self], timeout)
        ^^^^^^^^^^^^^^^^^^^^^
  File "/gpfs/data/oermannlab/users/xl3942/.conda/envs/neuro/lib/python3.12/multiprocessing/connection.py", line 1136, in wait
    ready = selector.select(timeout)
            ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/gpfs/data/oermannlab/users/xl3942/.conda/envs/neuro/lib/python3.12/selectors.py", line 415, in select
    fd_event_list = self._selector.poll(timeout)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
KeyboardInterrupt

Expected behavior

The second time of using the dataset should load normally.

Additional context

[tchaton]

Hey @Aceticia. Could you check the number of batches on each rank ? I need to add debug log to get the bottom of this.

[Aceticia]

All ranks have the same number of batches

Rank 5 has 75120 batches
Rank 2 has 75120 batches
Rank 4 has 75120 batches
Rank 0 has 75120 batches
Rank 3 has 75120 batches
Rank 4 has 75120 batches
Rank 0 has 75120 batches
Rank 6 has 75120 batches
Rank 1 has 75120 batches
Rank 7 has 75120 batches
Rank 2 has 75120 batches
Rank 3 has 75120 batches
Rank 4 has 75120 batches
Rank 6 has 75120 batches
Rank 7 has 75120 batches
Rank 3 has 75120 batches
Rank 5 has 75120 batches
Rank 1 has 75120 batches
Rank 5 has 75120 batches
Rank 0 has 75120 batches
Rank 6 has 75120 batches
Rank 7 has 75120 batches
Rank 1 has 75120 batches
Rank 2 has 75120 batches

[tchaton]

Great, so I think one chunk gets deleted by a rank and another is looking for it.

Would you be free next week to jump on a debugging call ?

Best,
T.C

[Aceticia]

I've rebuilt the whole thing without compression, 512MB chunk size and it now hangs out of the box. I'm not sure if it's due to deleted chunks, since I also tried assigning a rank-specific cache dir:

rank_this = fabric.global_rank
dataset_this = StreamingDataset(
    str(dataset_path / f"partition_{i}_of_512"),
    cache_dir=str(local_cache_dir / f"partition_{i}_of_512_{rank_this}"),
)

But I guess this only prevents the race between ranks, not between workers in the same rank. I can call Monday from 2PM to 5PM EST.

[Aceticia]

I experimented with using num_workers=0 and indeed it no longer gets stuck - this does seems like a race condition between workers.

[tchaton]

Hey @Aceticia. The bug is most likeky coming from there: https://github.com/Lightning-AI/litdata/blob/main/src/litdata/streaming/dataset.py#L274

If a worker delete the chunk locally that's meant to be used by another worker, then it will hang there: https://github.com/Lightning-AI/litdata/blob/main/src/litdata/streaming/item_loader.py#L146 waiting for the chunk to exist in the filesystem.

[tchaton]

Hey @Aceticia Let me know if you have some bandwidth to get on a call. My email is thomas@lightning.ai. Send me an invitation if you do.

[tchaton]

Hey @Aceticia Could you try this PR: https://github.com/Lightning-AI/litdata/pull/456/files ? I am trying to add a force download mechanism when a chunk is missing.

[Aceticia]

Hi @tchaton , I found that turning off compression fixes the hang in my case. This tracks with your current hypothesis about download race between workers. This also tells me that this is probably the same bug as these other hang bugs since now the hang is fixed without the "download" process (I'm on a local cluster) - so I'll close this issue. Best luck to yall on fixing this!

[tchaton]

Hey @Aceticia I think it would be good to understand what is happening. When using compression, it should make things faster when downloading from a remote source.

Are you still up for a debugging call ?

[Aceticia]

Sure, I can still get into a call later tomorrow. I'm not sure if I will be able to reproduce the issues though, since I no longer have the compression built dataset anymore. I deleted it since the uncompressed version is very large.