GitHub issue #2774: Prevent .map from using multiprocessing when loading from cache
State: closed | Created: 2021-08-09T12:11:38Z

## Context

On our setup, we use different setup to train vs proprocessing datasets. Usually we are able to obtain a high number of cpus to preprocess, which allows us to use `num_proc` however we can't use as many during training phase. Currently if we use `num_proc={whatever the preprocessing value was}` we load from cache, but we get:

```
Traceback (most recent call last):
  File "lib/python3.8/site-packages/multiprocess/pool.py", line 131, in worker
    put((job, i, result))
  File "lib/python3.8/site-packages/multiprocess/queues.py", line 371, in put
    self._writer.send_bytes(obj)
  File "lib/python3.8/site-packages/multiprocess/connection.py", line 203, in send_bytes
    self._send_bytes(m[offset:offset + size])
  File "lib/python3.8/site-packages/multiprocess/connection.py", line 414, in _send_bytes
    self._send(header + buf)
  File "lib/python3.8/site-packages/multiprocess/connection.py", line 371, in _send
    n = write(self._handle, buf)
BrokenPipeError: [Errno 32] Broken pipe
```

Our current guess, is that we're spawning too many processes compared to the number of cpus available, and it's running OOM.  Also we're loading this in DDP setting which means that for each gpu, I need to spawn a high number of processes to match the preprocessing fingerprint.

Instead what we suggest:
 - Allow loading shard sequentially, sharing the same fingerprint as the multiprocessed one, in order to leverage multiprocessing when we actually generate the cache, and remove it when loading from cache.

## Current issues

~I'm having a hard time making fingerprints match. For some reason, the multiprocessing and the sequential version generate two different hash.~

**EDIT**: Turns out multiprocessing and sequential have different `transform` value for fingerprinting (check `fingerprint_transform`) when running `_map_single`:
 - sequential : `datasets.arrow_dataset.Dataset._map_single`
 - multiprocessing: `datasets.arrow_dataset._map_single`
 
 This discrepancy is caused by multiprocessing pickling the transformer function, it doesn't seem to keep the `Dataset` hierarchy. I'm still unclear on why `func.__qual_name__` isn't handled correctly in multiprocessing. But replacing `__qualname__` by `__name__` fixes the issue.

## What was done

~We try to prevent the usage of multiprocessing when loading a dataset. Instead we load all cached shards sequentially.~

I couldn't find a nice way to obtain the cached_file_name and check they all exist before deciding to use the multiprocessing flow or not. Instead I expose an optional boolean `sequential` in `map` method.

## TODO
 - [x] Check that the multiprocessed version and the sequential version output the same output
 - [x] Check that sequential can load multiprocessed
 - [x] Check that multiprocessed can load sequential
 
 ## Test

```python
from datasets import load_dataset
from multiprocessing import Pool
import random

def process(batch, rng):
  length = len(batch["text"])
  return {**batch, "processed_text": [f"PROCESSED {rng.random()}" for _ in range(length)]}

dataset = load_dataset("stas/openwebtext-10k", split="train")
print(dataset.column_names)
print(type(dataset))

rng = random.Random(42)
dataset1 = dataset.map(process, batched=True, batch_size=50, num_proc=4, fn_kwargs={"rng": rng})

# This one should be loaded from cache
rng = random.Random(42)
dataset2 = dataset.map(process, batched=True, batch_size=50, num_proc=4, fn_kwargs={"rng": rng}, sequential=True)

# Just to check that the random generator was correct
print(dataset1[-1]["processed_text"])
print(dataset2[-1]["processed_text"])
```
 
 ## Other solutions

I chose to load everything sequentially, but we can probably find a way to load shards in parallel using another number of workers (essentially this would be an argument not used for fingerprinting, allowing to allow `m` shards using `n` processes, which would be very useful when same dataset have to be loaded on two different setup, and we still want to leverage cache).

Also we can use a env variable similarly to `TOKENIZERS_PARALLELISM` as this seems generally setup related (though this changes slightly if we use multiprocessing).
cc @lhoestq (since I had asked you previously on `num_proc` being used for fingerprinting). Don't know if this is acceptable.