GitHub issue: Add concurrent loading of shards to datasets.load_from_disk 
Author: kkoutini (CONTRIBUTOR)
State: closed | Created: 2023-12-01 13:13:53

In some file systems (like luster), memory mapping arrow files takes time. This can be accelerated by performing the mmap in parallel on processes or threads.

- Threads seem to be faster than processes when gathering the list of tables from the workers (see https://github.com/huggingface/datasets/issues/2252).
- I'm not sure if using threads would respect the `IN_MEMORY_MAX_SIZE` config.
- I'm not sure if we need to expose num_proc from `BaseReader.read` to `DatasetBuilder.as_dataset`. Since ` DatasetBuilder.as_dataset` is used in many places beside `load_dataset`.

### Tests on luster file system (on a shared partial node):
Loading 1231 shards of ~2GBs.
The files were pre-loaded in another process before the script runs (couldn't get a fresh node).
```python
import logging
from time import perf_counter

import datasets
logger = datasets.logging.get_logger(__name__)
datasets.logging.set_verbosity_info()
logging.basicConfig(level=logging.DEBUG, format="%(message)s")

class catchtime:
    # context to measure loading time: https://stackoverflow.com/questions/33987060/python-context-manager-that-measures-time
    def __init__(self, debug_print="Time", logger=logger):
        self.debug_print = debug_print
        self.logger = logger

    def __enter__(self):
        self.start = perf_counter()
        return self

    def __exit__(self, type, value, traceback):
        self.time = perf_counter() - self.start
        readout = f"{self.debug_print}: {self.time:.3f} seconds"
        self.logger.info(readout)

dataset_path=""

# warmup
with catchtime("Loading in parallel", logger=logger):
   ds = datasets.load_from_disk(dataset_path,num_proc=16)

# num_proc=16
with catchtime("Loading in parallel", logger=logger):
   ds = datasets.load_from_disk(dataset_path,num_proc=16)
# num_proc=32
with catchtime("Loading in parallel", logger=logger):
   ds = datasets.load_from_disk(dataset_path,num_proc=32)
# num_proc=1
with catchtime("Loading in conseq", logger=logger):
   ds = datasets.load_from_disk(dataset_path,num_proc=1)

```

#### Run 1
```
open file: .../dataset_dict.json
Loading the dataset from disk using 16 threads: 100%|██████████| 1231/1231 [01:28<00:00, 13.96shards/s]
Loading in parallel: 88.690 seconds
open file: .../dataset_dict.json
Loading the dataset from disk using 16 threads: 100%|██████████| 1231/1231 [01:48<00:00, 11.31shards/s]
Loading in parallel: 109.339 seconds
open file: .../dataset_dict.json
Loading the dataset from disk using 32 threads: 100%|██████████| 1231/1231 [01:06<00:00, 18.56shards/s]
Loading in parallel: 66.931 seconds
open file: .../dataset_dict.json
Loading the dataset from disk: 100%|██████████| 1231/1231 [05:09<00:00,  3.98shards/s]
Loading in conseq: 309.792 seconds
```
#### Run 2

```
open file: .../dataset_dict.json
Loading the dataset from disk using 16 threads: 100%|██████████| 1231/1231 [01:38<00:00, 12.53shards/s]
Loading in parallel: 98.831 seconds
open file: .../dataset_dict.json
Loading the dataset from disk using 16 threads: 100%|██████████| 1231/1231 [02:01<00:00, 10.16shards/s]
Loading in parallel: 121.669 seconds
open file: .../dataset_dict.json
Loading the dataset from disk using 32 threads: 100%|██████████| 1231/1231 [01:07<00:00, 18.18shards/s]
Loading in parallel: 68.192 seconds
open file: .../dataset_dict.json
Loading the dataset from disk: 100%|██████████| 1231/1231 [05:19<00:00,  3.86shards/s]
Loading in conseq: 319.759 seconds
```
#### Run 3
```
open file: .../dataset_dict.json
Loading the dataset from disk using 16 threads: 100%|██████████| 1231/1231 [01:36<00:00, 12.74shards/s]
Loading in parallel: 96.936 seconds
open file: .../dataset_dict.json
Loading the dataset from disk using 16 threads: 100%|██████████| 1231/1231 [02:00<00:00, 10.24shards/s]
Loading in parallel: 120.761 seconds
open file: .../dataset_dict.json
Loading the dataset from disk using 32 threads: 100%|██████████| 1231/1231 [01:08<00:00, 18.04shards/s]
Loading in parallel: 68.666 seconds
open file: .../dataset_dict.json
Loading the dataset from disk: 100%|██████████| 1231/1231 [05:35<00:00,  3.67shards/s]
Loading in conseq: 335.777 seconds
```


fix #2252