Persist IterableDataset epoch in workers
Opened by lhoestq (closed)

Use shared memory for the IterableDataset epoch.

This way calling `ds.set_epoch()` in the main process will update the epoch in the DataLoader workers as well.
This is useful especially because the epoch is used to compute the `effective_seed` used for shuffling.

I used torch's shared memory in case users want to send dataset copies without shared memory using pickle. I also find it easier to use than using `multiprocessing.shared_memory` than requires unlinking only in the main process, or `mp.Value` that is not picklable.

close https://github.com/huggingface/datasets/issues/6673

cc @rwightman 

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/datasets/pr_6710). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

<details>
<summary>Show benchmarks</summary>

PyArrow==8.0.0

<details>
<summary>Show updated benchmarks!</summary>

### Benchmark: benchmark_array_xd.json

| metric | read_batch_formatted_as_numpy after write_array2d | read_batch_formatted_as_numpy after write_flattened_sequence | read_batch_formatted_as_numpy after write_nested_sequence | read_batch_unformated after write_array2d | read_batch_unformated after write_flattened_sequence | read_batch_unformated after write_nested_sequence | read_col_formatted_as_numpy after write_array2d | read_col_formatted_as_numpy after write_flattened_sequence | read_col_formatted_as_numpy after write_nested_sequence | read_col_unformated after write_array2d | read_col_unformated after write_flattened_sequence | read_col_unformated after write_nested_sequence | read_formatted_as_numpy after write_array2d | read_formatted_as_numpy after write_flattened_sequence | read_formatted_as_numpy after write_nested_sequence | read_unformated after write_array2d | read_unformated after write_flattened_sequence | read_unformated after write_nested_sequence | write_array2d | write_flattened_sequence | write_nested_sequence |
|--------|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 0.005283 / 0.011353 (-0.006070) | 0.003866 / 0.011008 (-0.007142) | 0.063124 / 0.038508 (0.024616) | 0.030240 / 0.023109 (0.007131) | 0.232855 / 0.275898 (-0.043043) | 0.257538 / 0.323480 (-0.065942) | 0.004165 / 0.007986 (-0.003820) | 0.002826 / 0.004328 (-0.001502) | 0.049735 / 0.004250 (0.045485) | 0.045297 / 0.037052 (0.008244) | 0.251831 / 0.258489 (-0.006658) | 0.277812 / 0.293841 (-0.016029) | 0.030004 / 0.128546 (-0.098542) | 0.012319 / 0.075646 (-0.063328) | 0.206881 / 0.419271 (-0.212391) | 0.036561 / 0.043533 (-0.006972) | 0.234364 / 0.255139 (-0.020775) | 0.258316 / 0.283200 (-0.024884) | 0.017815 / 0.141683 (-0.123867) | 1.114111 / 1.452155 (-0.338043) | 1.165428 / 1.492716 (-0.327288) |

### Benchmark: benchmark_getitem\_100B.json

| metric | get_batch_of\_1024\_random_rows | get_batch_of\_1024\_rows | get_first_row | get_last_row |
|--------|---|---|---|---|
| new / old (diff) | 0.099302 / 0.018006 (0.081296) | 0.309195 / 0.000490 (0.308705) | 0.000261 / 0.000200 (0.000061) | 0.000044 / 0.000054 (-0.000010) |

### Benchmark: benchmark_indices_mapping.json

| metric | select | shard | shuffle | sort | train_test_split |
|--------|---|---|---|---|---|
| new / old (diff) | 0.018765 / 0.037411 (-0.018646) | 0.063123 / 0.014526 (0.048597) | 0.075437 / 0.176557 (-0.101119) | 0.122570 / 0.737135 (-0.614566) | 0.076637 / 0.296338 (-0.219702) |

### Benchmark: benchmark_iterating.json

| metric | read 5000 | read 50000 | read_batch 50000 10 | read_batch 50000 100 | read_batch 50000 1000 | read_formatted numpy 5000 | read_formatted pandas 5000 | read_formatted tensorflow 5000 | read_formatted torch 5000 | read_formatted_batch numpy 5000 10 | read_formatted_batch numpy 5000 1000 | shuffled read 5000 | shuffled read 50000 | shuffled read_batch 50000 10 | shuffled read_batch 50000 100 | shuffled read_batch 50000 1000 | shuffled read_formatted numpy 5000 | shuffled read_formatted_batch numpy 5000 10 | shuffled read_formatted_batch numpy 5000 1000 |
|--------|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 0.289965 / 0.215209 (0.074756) | 2.839053 / 2.077655 (0.761398) | 1.503463 / 1.504120 (-0.000657) | 1.390833 / 1.541195 (-0.150361) | 1.401918 / 1.468490 (-0.066572) | 0.711000 / 4.584777 (-3.873777) | 2.325513 / 3.745712 (-1.420199) | 2.831630 / 5.269862 (-2.438231) | 1.908370 / 4.565676 (-2.657307) | 0.077867 / 0.424275 (-0.346408) | 0.005509 / 0.007607 (-0.002098) | 0.336494 / 0.226044 (0.110450) | 3.358587 / 2.268929 (1.089658) | 1.901067 / 55.444624 (-53.543558) | 1.590130 / 6.876477 (-5.286347) | 1.753850 / 2.142072 (-0.388223) | 0.792458 / 4.805227 (-4.012769) | 0.135584 / 6.500664 (-6.365080) | 0.042028 / 0.075469 (-0.033441) |

### Benchmark: benchmark_map_filter.json

| metric | filter | map fast-tokenizer batched | map identity | map identity batched | map no-op batched | map no-op batched numpy | map no-op batched pandas | map no-op batched pytorch | map no-op batched tensorflow |
|--------|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 0.966162 / 1.841788 (-0.875625) | 11.705310 / 8.074308 (3.631002) | 9.158842 / 10.191392 (-1.032550) | 0.128793 / 0.680424 (-0.551631) | 0.014422 / 0.534201 (-0.519779) | 0.299009 / 0.579283 (-0.280274) | 0.262873 / 0.434364 (-0.171491) | 0.340836 / 0.540337 (-0.199501) | 0.464440 / 1.386936 (-0.922496) |

</details>
PyArrow==latest

<details>
<summary>Show updated benchmarks!</summary>

### Benchmark: benchmark_array_xd.json

| metric | read_batch_formatted_as_numpy after write_array2d | read_batch_formatted_as_numpy after write_flattened_sequence | read_batch_formatted_as_numpy after write_nested_sequence | read_batch_unformated after write_array2d | read_batch_unformated after write_flattened_sequence | read_batch_unformated after write_nested_sequence | read_col_formatted_as_numpy after write_array2d | read_col_formatted_as_numpy after write_flattened_sequence | read_col_formatted_as_numpy after write_nested_sequence | read_col_unformated after write_array2d | read_col_unformated after write_flattened_sequence | read_col_unformated after write_nested_sequence | read_formatted_as_numpy after write_array2d | read_formatted_as_numpy after write_flattened_sequence | read_formatted_as_numpy after write_nested_sequence | read_unformated after write_array2d | read_unformated after write_flattened_sequence | read_unformated after write_nested_sequence | write_array2d | write_flattened_sequence | write_nested_sequence |
|--------|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 0.005951 / 0.011353 (-0.005402) | 0.003984 / 0.011008 (-0.007024) | 0.051432 / 0.038508 (0.012924) | 0.033223 / 0.023109 (0.010113) | 0.263972 / 0.275898 (-0.011926) | 0.289060 / 0.323480 (-0.034420) | 0.004446 / 0.007986 (-0.003540) | 0.002891 / 0.004328 (-0.001438) | 0.049347 / 0.004250 (0.045096) | 0.041191 / 0.037052 (0.004138) | 0.278334 / 0.258489 (0.019844) | 0.314065 / 0.293841 (0.020224) | 0.032020 / 0.128546 (-0.096526) | 0.012472 / 0.075646 (-0.063174) | 0.061288 / 0.419271 (-0.357984) | 0.033489 / 0.043533 (-0.010044) | 0.266831 / 0.255139 (0.011692) | 0.283008 / 0.283200 (-0.000192) | 0.018491 / 0.141683 (-0.123192) | 1.133634 / 1.452155 (-0.318521) | 1.154627 / 1.492716 (-0.338089) |

### Benchmark: benchmark_getitem\_100B.json

| metric | get_batch_of\_1024\_random_rows | get_batch_of\_1024\_rows | get_first_row | get_last_row |
|--------|---|---|---|---|
| new / old (diff) | 0.101831 / 0.018006 (0.083825) | 0.317942 / 0.000490 (0.317452) | 0.000217 / 0.000200 (0.000018) | 0.000056 / 0.000054 (0.000002) |

### Benchmark: benchmark_indices_mapping.json

| metric | select | shard | shuffle | sort | train_test_split |
|--------|---|---|---|---|---|
| new / old (diff) | 0.022608 / 0.037411 (-0.014803) | 0.076776 / 0.014526 (0.062250) | 0.088686 / 0.176557 (-0.087870) | 0.129092 / 0.737135 (-0.608044) | 0.090780 / 0.296338 (-0.205558) |

### Benchmark: benchmark_iterating.json

| metric | read 5000 | read 50000 | read_batch 50000 10 | read_batch 50000 100 | read_batch 50000 1000 | read_formatted numpy 5000 | read_formatted pandas 5000 | read_formatted tensorflow 5000 | read_formatted torch 5000 | read_formatted_batch numpy 5000 10 | read_formatted_batch numpy 5000 1000 | shuffled read 5000 | shuffled read 50000 | shuffled read_batch 50000 10 | shuffled read_batch 50000 100 | shuffled read_batch 50000 1000 | shuffled read_formatted numpy 5000 | shuffled read_formatted_batch numpy 5000 10 | shuffled read_formatted_batch numpy 5000 1000 |
|--------|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 0.286762 / 0.215209 (0.071553) | 2.824307 / 2.077655 (0.746652) | 1.547215 / 1.504120 (0.043095) | 1.424522 / 1.541195 (-0.116673) | 1.446414 / 1.468490 (-0.022076) | 0.723683 / 4.584777 (-3.861094) | 0.974129 / 3.745712 (-2.771583) | 2.952552 / 5.269862 (-2.317309) | 1.903663 / 4.565676 (-2.662013) | 0.078786 / 0.424275 (-0.345489) | 0.005130 / 0.007607 (-0.002477) | 0.338925 / 0.226044 (0.112881) | 3.378557 / 2.268929 (1.109629) | 1.892951 / 55.444624 (-53.551674) | 1.599844 / 6.876477 (-5.276633) | 1.611963 / 2.142072 (-0.530109) | 0.793614 / 4.805227 (-4.011613) | 0.133795 / 6.500664 (-6.366869) | 0.040777 / 0.075469 (-0.034692) |

### Benchmark: benchmark_map_filter.json

| metric | filter | map fast-tokenizer batched | map identity | map identity batched | map no-op batched | map no-op batched numpy | map no-op batched pandas | map no-op batched pytorch | map no-op batched tensorflow |
|--------|---|---|---|---|---|---|---|---|---|
| new / old (diff) | 1.001391 / 1.841788 (-0.840397) | 12.166811 / 8.074308 (4.092503) | 10.588180 / 10.191392 (0.396788) | 0.141609 / 0.680424 (-0.538815) | 0.020941 / 0.534201 (-0.513260) | 0.340149 / 0.579283 (-0.239134) | 0.122988 / 0.434364 (-0.311376) | 0.339747 / 0.540337 (-0.200591) | 0.434338 / 1.386936 (-0.952598) |

</details>
</details>

![](https://cml.dev/watermark.png#4ba47a35820069e7de9f374479d06b2a7935767e "CML watermark")