Iterating over DataLoader based on HF datasets is stuck forever

### Describe the bug

I am using Amazon Sagemaker notebook (Amazon Linux 2) with python 3.10 based Conda environment.
I have a dataset in parquet format locally. When I try to iterate over it, the loader is stuck forever. Note that the same code is working for python 3.6 based conda environment seamlessly. What should be my next steps here?

### Steps to reproduce the bug

```
train_dataset = load_dataset(
    "parquet", data_files = {'train': tr_data_path + '*.parquet'}, 
    split = 'train', 
    collate_fn = streaming_data_collate_fn,
    streaming = True
).with_format('torch')

train_dataloader = DataLoader(train_dataset, batch_size = 2, num_workers = 0)

t = time.time()
iter_ = 0
for batch in train_dataloader:
    iter_ += 1
    
    if iter_ == 1000:
        break
        
print (time.time() - t) 
```

### Expected behavior

The snippet should work normally and load the next batch of data.

### Environment info

datasets: '2.14.0'
pyarrow: '12.0.0'
torch: '2.0.0'
Python: 3.10.10 | packaged by conda-forge | (main, Mar 24 2023, 20:08:06) [GCC 11.3.0]

!uname -r
5.10.178-162.673.amzn2.x86_64

When the process starts to hang, can you interrupt it with CTRL + C and paste the error stack trace here? 

Thanks @mariosasko for your prompt response, here's the stack trace:

```
KeyboardInterrupt                         Traceback (most recent call last)
Cell In[12], line 4
      2 t = time.time()
      3 iter_ = 0
----> 4 for batch in train_dataloader:
      5     #batch_proc = streaming_obj.collect_streaming_data_batch(batch)
      6     iter_ += 1
      8     if iter_ == 1:

File ~/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/torch/utils/data/dataloader.py:634, in _BaseDataLoaderIter.__next__(self)
    631 if self._sampler_iter is None:
    632     # TODO(https://github.com/pytorch/pytorch/issues/76750)
    633     self._reset()  # type: ignore[call-arg]
--> 634 data = self._next_data()
    635 self._num_yielded += 1
    636 if self._dataset_kind == _DatasetKind.Iterable and \
    637         self._IterableDataset_len_called is not None and \
    638         self._num_yielded > self._IterableDataset_len_called:

File ~/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/torch/utils/data/dataloader.py:678, in _SingleProcessDataLoaderIter._next_data(self)
    676 def _next_data(self):
    677     index = self._next_index()  # may raise StopIteration
--> 678     data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
    679     if self._pin_memory:
    680         data = _utils.pin_memory.pin_memory(data, self._pin_memory_device)

File ~/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py:32, in _IterableDatasetFetcher.fetch(self, possibly_batched_index)
     30 for _ in possibly_batched_index:
     31     try:
---> 32         data.append(next(self.dataset_iter))
     33     except StopIteration:
     34         self.ended = True

File ~/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/datasets/iterable_dataset.py:1353, in IterableDataset.__iter__(self)
   1350         yield formatter.format_row(pa_table)
   1351     return
-> 1353 for key, example in ex_iterable:
   1354     if self.features:
   1355         # `IterableDataset` automatically fills missing columns with None.
   1356         # This is done with `_apply_feature_types_on_example`.
   1357         example = _apply_feature_types_on_example(
   1358             example, self.features, token_per_repo_id=self._token_per_repo_id
   1359         )

File ~/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/datasets/iterable_dataset.py:956, in BufferShuffledExamplesIterable.__iter__(self)
    954 # this is the shuffle buffer that we keep in memory
    955 mem_buffer = []
--> 956 for x in self.ex_iterable:
    957     if len(mem_buffer) == buffer_size:  # if the buffer is full, pick and example from it
    958         i = next(indices_iterator)

File ~/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/datasets/iterable_dataset.py:296, in ShuffledDataSourcesArrowExamplesIterable.__iter__(self)
    294 for key, pa_table in self.generate_tables_fn(**kwargs_with_shuffled_shards):
    295     for pa_subtable in pa_table.to_reader(max_chunksize=config.ARROW_READER_BATCH_SIZE_IN_DATASET_ITER):
--> 296         formatted_batch = formatter.format_batch(pa_subtable)
    297         for example in _batch_to_examples(formatted_batch):
    298             yield key, example

File ~/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/datasets/formatting/formatting.py:448, in PythonFormatter.format_batch(self, pa_table)
    446 if self.lazy:
    447     return LazyBatch(pa_table, self)
--> 448 batch = self.python_arrow_extractor().extract_batch(pa_table)
    449 batch = self.python_features_decoder.decode_batch(batch)
    450 return batch

File ~/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/datasets/formatting/formatting.py:150, in PythonArrowExtractor.extract_batch(self, pa_table)
    149 def extract_batch(self, pa_table: pa.Table) -> dict:
--> 150     return pa_table.to_pydict()

KeyboardInterrupt: 
```


Update: If i let it run, it eventually fails with:

```
RuntimeError                              Traceback (most recent call last)
Cell In[16], line 4
      2 t = time.time()
      3 iter_ = 0
----> 4 for batch in train_dataloader:
      5     #batch_proc = streaming_obj.collect_streaming_data_batch(batch)
      6     iter_ += 1
      8     if iter_ == 1:

File ~/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/torch/utils/data/dataloader.py:634, in _BaseDataLoaderIter.__next__(self)
    631 if self._sampler_iter is None:
    632     # TODO(https://github.com/pytorch/pytorch/issues/76750)
    633     self._reset()  # type: ignore[call-arg]
--> 634 data = self._next_data()
    635 self._num_yielded += 1
    636 if self._dataset_kind == _DatasetKind.Iterable and \
    637         self._IterableDataset_len_called is not None and \
    638         self._num_yielded > self._IterableDataset_len_called:

File ~/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/torch/utils/data/dataloader.py:678, in _SingleProcessDataLoaderIter._next_data(self)
    676 def _next_data(self):
    677     index = self._next_index()  # may raise StopIteration
--> 678     data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
    679     if self._pin_memory:
    680         data = _utils.pin_memory.pin_memory(data, self._pin_memory_device)

File ~/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py:32, in _IterableDatasetFetcher.fetch(self, possibly_batched_index)
     30 for _ in possibly_batched_index:
     31     try:
---> 32         data.append(next(self.dataset_iter))
     33     except StopIteration:
     34         self.ended = True

File ~/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/datasets/iterable_dataset.py:1360, in IterableDataset.__iter__(self)
   1354 if self.features:
   1355     # `IterableDataset` automatically fills missing columns with None.
   1356     # This is done with `_apply_feature_types_on_example`.
   1357     example = _apply_feature_types_on_example(
   1358         example, self.features, token_per_repo_id=self._token_per_repo_id
   1359     )
-> 1360 yield format_dict(example) if format_dict else example

File ~/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/datasets/formatting/torch_formatter.py:85, in TorchFormatter.recursive_tensorize(self, data_struct)
     84 def recursive_tensorize(self, data_struct: dict):
---> 85     return map_nested(self._recursive_tensorize, data_struct, map_list=False)

File ~/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/datasets/utils/py_utils.py:463, in map_nested(function, data_struct, dict_only, map_list, map_tuple, map_numpy, num_proc, parallel_min_length, types, disable_tqdm, desc)
    461     num_proc = 1
    462 if num_proc != -1 and num_proc <= 1 or len(iterable) < parallel_min_length:
--> 463     mapped = [
    464         _single_map_nested((function, obj, types, None, True, None))
    465         for obj in logging.tqdm(iterable, disable=disable_tqdm, desc=desc)
    466     ]
    467 else:
    468     mapped = parallel_map(function, iterable, num_proc, types, disable_tqdm, desc, _single_map_nested)

File ~/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/datasets/utils/py_utils.py:464, in <listcomp>(.0)
    461     num_proc = 1
    462 if num_proc != -1 and num_proc <= 1 or len(iterable) < parallel_min_length:
    463     mapped = [
--> 464         _single_map_nested((function, obj, types, None, True, None))
    465         for obj in logging.tqdm(iterable, disable=disable_tqdm, desc=desc)
    466     ]
    467 else:
    468     mapped = parallel_map(function, iterable, num_proc, types, disable_tqdm, desc, _single_map_nested)

File ~/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/datasets/utils/py_utils.py:366, in _single_map_nested(args)
    364 # Singleton first to spare some computation
    365 if not isinstance(data_struct, dict) and not isinstance(data_struct, types):
--> 366     return function(data_struct)
    368 # Reduce logging to keep things readable in multiprocessing with tqdm
    369 if rank is not None and logging.get_verbosity() < logging.WARNING:

File ~/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/datasets/formatting/torch_formatter.py:82, in TorchFormatter._recursive_tensorize(self, data_struct)
     80 elif isinstance(data_struct, (list, tuple)):
     81     return self._consolidate([self.recursive_tensorize(substruct) for substruct in data_struct])
---> 82 return self._tensorize(data_struct)

File ~/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/datasets/formatting/torch_formatter.py:68, in TorchFormatter._tensorize(self, value)
     66     if isinstance(value, PIL.Image.Image):
     67         value = np.asarray(value)
---> 68 return torch.tensor(value, **{**default_dtype, **self.torch_tensor_kwargs})

RuntimeError: Could not infer dtype of decimal.Decimal
```

PyTorch tensors cannot store `Decimal` objects. Casting the column with decimals to `float` should fix the issue.

I already have cast in collate_fn, in which I perform .astype(float) for each numerical field.
On the same instance, I installed a conda env with python 3.6, and this works well.

Sample:

```
def streaming_data_collate_fn(batch):
    df = pd.DataFrame.from_dict(batch)
    feat_vals = torch.FloatTensor(np.nan_to_num(np.array(df[feats].astype(float))))

```

`collate_fn` is applied after the `torch` formatting step, so I think the only option when working with an `IterableDataset` is to remove the `with_format` call and perform the conversion from Python values to PyTorch tensors in `collate_fn`. The standard `Dataset` supports `with_format("numpy")`, which should make this conversion faster.

Thanks! 
Python 3.10 conda-env: After replacing with_format("torch") with with_format("numpy"), the error went away. However, it was still taking over 2 minutes to load a very small batch of 64 samples with num_workers set to 32. Once I removed with_format call altogether, it is finishing in 11 seconds.

Python 3.6 based conda-env: When I switch the kernel , neither of the above work, and with_format("torch") is the only thing that works, and executes in 1.6 seconds.

I feel something else is also amiss here.

Can you share the `datasets` and `torch` versions installed in these conda envs?

> Once I removed with_format call altogether, it is finishing in 11 seconds.

Hmm, that's surprising. What are your dataset's `.features`?

Python 3.6: 
datasets.__version__ 2.4.0
torch.__version__ 1.10.1+cu102

Python 3.10:
datasets.__version__ 2.14.0
torch.__version__ 2.0.0

Anonymized features are of the form (subset shown here):
{
'string_feature_i': Value(dtype='string', id=None),
'numerical_feature_i': Value(dtype='decimal128(38, 0)', id=None),
'numerical_feature_series_i': Sequence(feature=Value(dtype='float64', id=None), length=-1, id=None),
}


There is no output from .features in python 3.6 kernel BTW.

One more thing, in python 3.10 based kernel, interestingly increasing num_workers seem to be increasing the runtime of iterating I was trying out. In python 3.10 kernel execution, I do not even see multiple CPU cores spiking unlike in 3.6.

512 batch size on 32 workers executes in 2.4 seconds on python 3.6 kernel, while it takes ~118 seconds on 3.10!

**Update**: It seems the latency part is more of a multiprocessing issue with torch and some host specific issue, and I had to scourge through relevant pytorch issues, when I stumbled across these threads:
1. https://github.com/pytorch/pytorch/issues/102494
2. https://github.com/pytorch/pytorch/issues/102269
3. https://github.com/pytorch/pytorch/issues/99625

Out of the suggested solutions, the one that worked in my case was:
```
os.environ['KMP_AFFINITY'] = "disabled"
```
It is working for now, though I have no clue why, just I hope it does not get stuck when I do actual model training, will update by tomorrow.




I'm facing a similar situation in the local VS Code. 

Datasets version 2.14.4
Torch 2.0.1+cu118

Same code runs without issues in Colab

```
from datasets import load_dataset

dataset = load_dataset("Supermaxman/esa-hubble", streaming=True)
sample = next(iter(dataset["train"]))
```

is stuck for minutes. If I interrupt, I get

```
---------------------------------------------------------------------------
KeyboardInterrupt                         Traceback (most recent call last)
Cell In[5], line 5
      1 from datasets import load_dataset
      3 dataset = load_dataset("Supermaxman/esa-hubble", streaming=True)
----> 5 sample = next(iter(dataset["train"]))
      6 print(sample["text"])
      7 sample["image"]

File [~/miniconda3/envs/book/lib/python3.10/site-packages/datasets/iterable_dataset.py:1353](https://file+.vscode-resource.vscode-cdn.net/home/osanseviero/Desktop/workspace/genai/nbs/~/miniconda3/envs/book/lib/python3.10/site-packages/datasets/iterable_dataset.py:1353), in IterableDataset.__iter__(self)
   1350         yield formatter.format_row(pa_table)
   1351     return
-> 1353 for key, example in ex_iterable:
   1354     if self.features:
   1355         # `IterableDataset` automatically fills missing columns with None.
   1356         # This is done with `_apply_feature_types_on_example`.
   1357         example = _apply_feature_types_on_example(
   1358             example, self.features, token_per_repo_id=self._token_per_repo_id
   1359         )

File [~/miniconda3/envs/book/lib/python3.10/site-packages/datasets/iterable_dataset.py:255](https://file+.vscode-resource.vscode-cdn.net/home/osanseviero/Desktop/workspace/genai/nbs/~/miniconda3/envs/book/lib/python3.10/site-packages/datasets/iterable_dataset.py:255), in ArrowExamplesIterable.__iter__(self)
    253 def __iter__(self):
    254     formatter = PythonFormatter()
--> 255     for key, pa_table in self.generate_tables_fn(**self.kwargs):
    256         for pa_subtable in pa_table.to_reader(max_chunksize=config.ARROW_READER_BATCH_SIZE_IN_DATASET_ITER):
...
-> 1130         return self._sslobj.read(len, buffer)
   1131     else:
   1132         return self._sslobj.read(len)
```

@osanseviero I assume the `self._sslobj.read(len, buffer)` line comes from the built-in `ssl` module, so this probably has something to do with your network. Please open a new issue with the full stack trace in case you haven't resolved this yet.

Thank you reporting this and sharing the solution, I ran into this as well!