Prevent .map from using multiprocessing when loading from cache
State: closed | Author: thomasw21 (CONTRIBUTOR)

## 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.

Comments: I'm guessing tests are failling, because this was pushed before https://github.com/huggingface/datasets/pull/2779 was merged? cc @albertvillanova 

Hi @thomasw21, yes you are right: those failing tests were fixed with #2779.

Would you mind to merge current upstream master branch and push again?
```
git checkout sequential_map_when_cached
git fetch upstream master
git merge upstream/master
git push -u origin sequential_map_when_cached
```

Thanks for working on this ! I'm sure we can figure something out ;)

Currently `map` starts a process to apply the map function on each shard. If the shard has already been processed, then the process that has been spawned loads the processed shard from the cache and returns it.

I think we should be able to simply not start a process if a shard is already processed and cached.
This way:
- you won't need to specify `sequential=True`
- it won't create new processes if the dataset is already processed and cached
- it will properly reload each processed shard that is cached

To know if we have to start a new process for a shard you can use the function `update_fingerprint` from fingerprint.py to know the expected fingerprint of the processed shard.
Then, if the shard has already been processed, there will be a cache file named `cached-<new_fingerprint>.arrow` and you can load it with
```
Dataset.from_file(path_to_cache_file, info=self.info, split=self.split)
```

Let me know if that makes sense !

Yes that makes total sense, I tried to initially do that, except the way fingerprint is handled doesn't allow to easily manipulate such a field. Typically the fingerprinting is handled in `@fingerprint_transform` which has a bunch of params that aren't quite easy to extract. Those params are used to manipulate args, kwargs in fancy ways in order to finally obtain a dictionary used for fingerprint. I could duplicate everything, but this look like a very risky thing to do. I'll take a look if I can make something work with `inspect` if I can make a very simple wrapper.

A much more simpler solution I think is adding an optional `shard: Optional[int] = None` parameter. If None, we use the number of proc as the number of shards, otherwise we pass down the expected number of shards and use either sequential/multiprocessing (with arbitrary number of workers) to load the shards? This would allow the weird case where one wants a large number of shards with a limited amount of processes. Not the smartest thing to do, but it's not an absurd behaviour. Would this be acceptable?

@lhoestq friendly ping as I feel it's up for review.

The CI error is unrelated to the changes of this PR - it looks like an SSL issue with conda