Fix inconsistent caching behaviour in Dataset.map() with multiprocessing #3044
Opened by vlievin (closed)

Fix #3044

1. A rough unit test that fails without the fix. It probably doesn't comply with your code standards, but that just to draft the idea.
2. A one liner fix 

Comments:
Hi ! Thanks for noticing this inconsistence and suggesting a fix :)

If I understand correctly you try to pass the same fingerprint to each processed shard of the dataset. This can be an issue since each shard is actually a different dataset with different data: they shouldn't have the same fingerprint.

Ideally we want the result after `map` to have this fingerprint. The result after `map` is the concatenation of all the processed shards. In this case what we can do is add the `fingerprint` parameter to `concatenate_datasets` to overwrite the fingerprint here if needed:
https://github.com/huggingface/datasets/blob/03b7f123cc17afc517c0aa2f912bbd90cb266185/src/datasets/arrow_dataset.py#L3588-L3590

then you can pass the fingerprint to `concatenate_datasets` here:
https://github.com/huggingface/datasets/blob/03b7f123cc17afc517c0aa2f912bbd90cb266185/src/datasets/arrow_dataset.py#L2044-L2044

Hi @lhoestq, thanks for the pointers! Not having a unique fingerprint per shard was indeed was indeed a problem. 

Let me look into this. I'll be back with a fix soon.

Alright, to clarify about my problem. I using am using `datasets` with large datasets, and want to cache a heavy and non-deterministically fingerprintable function (using `datasets.fingerprint.Hasher`). Using `Dataset.map()` as it is would cause generating a random fingerprint. To circumvent this, I am generating custom deterministic fingerprints, which I pass as an argument to `Dataset.map()`. In that way, a deterministic fingerprint is set, and caching can be used. 

This approach works well when using `num_proc==1`, but not so well when using `num_proc>1`. In both cases, `dataset._fingerprint` is effectively set to `new_fingerprint` at the end of the `.map()` call. However, caching is not used when `num_proc>1`, a non deterministically fingerprintable function and `new_fingerprint != null. The reason is that caching operates within `Dataset._map_single` and `new_fingerprint` is not passed here. 

This pull request implements a quick fix (+unit test) by passing `new_fingerprint=f"{new_fingerprint}-part{rank+1}-{num_proc}"` to each `_map_single` call. Using a separate name for each call makes sure that each worker uses a different cache file (as you mentioned above).

However, this solution still means that using a different value for `num_proc` will require computing new partial cache files. In the long run, performing the caching within `map()` instead of within `_map_single()` would be a cleaner solution.

Hi @vlievin,

If I understand your example correctly, you are trying to use the `new_fingerprint` param to have a deterministic fingerprint of the transform, which is not hashable due to randomness. Any particular reason why you are not using the `cache_file_name` param instead? I did run your example with the `cache_file_name` specified, and it behaves as expected based on the logs. Internally, `new_fingerprint` is needed to inject the calculated fingerprint into a method by the `fingerprint_transform` decorator, which is then used to compute the cache file name in `Dataset._get_cache_file_path` if the user hasn't specified one. 

Hi @lhoestq, I have cleaned up the unit test (incl. styling). It should be ready to merge as such. I am using this branch in my project and everything works fine. 

Hi @mariosasko,  the argument `new_fingerprint` allowed me to deterministically cache my transformation when using `num_proc=1`, so I assumed that was the right way to go. But maybe I have misinterpreted how `new_fingerprint` should be used.

But in any case, `map()` should perform consistently with regards to `num_proc`. In my opinion, the behaviour of `Dataset.map()` should perform the same, and this without requiring the user to input `cache_file_name` when `num_proc>1` is set.
But maybe there is a more elegant way to fix this using `cache_file_name` internally for each `_single_map()` call.

So, I think this is a more high level design decision and I will leave it to the maintainers :) 

Hi @vlievin,

I appreciate your effort, but `new_fingerprint` behaves as described in the `Dataset.map` docs, and we don't have to follow some artificial consistency in regards to `num_proc`:
https://github.com/huggingface/datasets/blob/adc5cec58dd15ee672016086fefdea34b3143e4f/src/datasets/arrow_dataset.py#L1962-L1963

Additionally, to compute the cache file name, you are using a private method (`dset._get_cache_file_path(new_fingerprint)`); prefixed with `_`), so this is a sign you may be doing something wrong because you are relying on the internals. I suggest you use cache_file_name instead and follow the suffix template docs, which explain how to compute file paths of the created cache files when `num_proc > 1`.

Hi @mariosasko, thanks for the pointer regarding the use of the private method in then unit tests. 

Yes, `new_fingerprint` behaves as documented. If you don't think this is an issue, feel free to close this pull request. 


Allowing the users to pass the fingerprint themselves for functions that can't be hashed would be a nice improvements. However I agree that as @mariosasko mentioned this is currently not how we want the API to behave for now - since it has to do with the internals of the library.

Though we can discuss what could be the right way of doing it in https://github.com/huggingface/datasets/issues/3044 if you don't mind !