Add multiprocessing
State: closed | Author: lhoestq (MEMBER)

Adding multiprocessing to `.map`

It works in 3 steps:
- shard the dataset in `num_proc` shards
- spawn one process per shard and call `map` on them
- concatenate the resulting datasets

Example of usage:

```python
from nlp import load_dataset

dataset = load_dataset("squad", split="train")

def function(x):
    return {"lowered": x.lower()}

processed = d.map(
    function,
    input_columns=["context"],
    num_proc=4,
    cache_file_name="playground/tmp.arrow",
    load_from_cache_file=False
)
```

Here it writes 4 files  depending on the process rank:
- `playground/tmp_00000_of_00004.arrow`
- `playground/tmp_00001_of_00004.arrow`
- `playground/tmp_00002_of_00004.arrow`
- `playground/tmp_00003_of_00004.arrow`

The suffix format can be specified by the user.

If the `cache_file_name` is not specified, it writes into separated files depending on the fingerprint, as usual.

I still need to:
- write tests for this
- try to improve the logging (currently it shows 4 progress bars, but if one finishes before the others, then the following messages are written over the progress bars)


Comments: Logging looks like

```
Done writing 21900 indices in 3854400 bytes .
Process #0 will write at playground/tmp_00000_of_00004.arrow
Done writing 21900 indices in 3854400 bytes .
Process #1 will write at playground/tmp_00001_of_00004.arrow
Done writing 21900 indices in 3854400 bytes .
Process #2 will write at playground/tmp_00002_of_00004.arrow
Done writing 21899 indices in 3854224 bytes .
Process #3 will write at playground/tmp_00003_of_00004.arrow
Spawning 4 processes
#3: 100%|████████████████████████████████████████████████| 21899/21899 [00:02<00:00, 8027.41ex/s]
#0: 100%|████████████████████████████████████████████████| 21900/21900 [00:02<00:00, 7982.87ex/s]
#1: 100%|████████████████████████████████████████████████| 21900/21900 [00:02<00:00, 7923.89ex/s]
#2: 100%|████████████████████████████████████████████████| 21900/21900 [00:02<00:00, 7920.04ex/s]
Concatenating 4 shards from multiprocessing
```

I added tests and improved logging.
Both `map` and `filter` support multiprocessing

A bit strange that the benchmarks on map/filter are worth than `master`.
(maybe because they are not done on the same machine)

The benchmark also got worse in other PRs (see [here](https://github.com/huggingface/nlp/pull/550#commitcomment-41931609) for example, where we have 16sec for `map fast-tokenizer batched` and 18 sec for `map identity`)

Hi,

when I use the multiprocessing in ```.map```:
```
dataset = load_dataset("text", data_files=file_path, split="train")
dataset = dataset.map(lambda ex: tokenizer(ex["text"], add_special_tokens=True,
                                      truncation=True, max_length=args.block_size), batched=True, num_proc=16)
dataset.set_format(type='torch', columns=['input_ids'])
```
I get the following error:
```
Traceback (most recent call last):
  File "src/run.py", line 373, in <module>
    main()
  File "src/run.py", line 295, in main
    get_dataset(data_args, tokenizer=tokenizer, cache_dir=model_args.cache_dir) if training_args.do_train else None
  File "src/run.py", line 153, in get_dataset
    dataset = dataset.map(lambda ex: tokenizer(ex["text"], add_special_tokens=True,
  File "/root/miniconda3/envs/py3.8/lib/python3.8/site-packages/datasets/arrow_dataset.py", line 1287, in map
    transformed_shards = [r.get() for r in results]
  File "/root/miniconda3/envs/py3.8/lib/python3.8/site-packages/datasets/arrow_dataset.py", line 1287, in <listcomp>
    transformed_shards = [r.get() for r in results]
  File "/root/miniconda3/envs/py3.8/lib/python3.8/multiprocessing/pool.py", line 771, in get
    raise self._value
    put(task)
  File "/root/miniconda3/envs/py3.8/lib/python3.8/multiprocessing/connection.py", line 206, in send
    self._send_bytes(_ForkingPickler.dumps(obj))
  File "/root/miniconda3/envs/py3.8/lib/python3.8/multiprocessing/reduction.py", line 51, in dumps
    cls(buf, protocol).dump(obj)
AttributeError: Can't pickle local object 'get_dataset.<locals>.<lambda>'
```
I think you should use [pathos](https://github.com/uqfoundation/pathos) to pickle the lambda function and some others!
I change the 30 line of src/datasets/arrow_dataset.py as following:
```
# 30 line: from multiprocessing import Pool, RLock
import pathos
from pathos.multiprocessing import Pool
from multiprocessing import RLock
```
and it works!

That's very cool indeed !
Shall we condiser adding this dependency @thomwolf ?

We already use `dill` so that's definitely a very interesting option indeed!

it gets stuck on debian 9 when num_proc > 1


Are you using a tokenizer ?
Did you try to set `TOKENIZERS_PARALLELISM=false` ?

Feel free to discuss it in #620 , we're discussing this issue

I set `TOKENIZERS_PARALLELISM=false`. Just the warning went away. The program was still stuck