Align `extract_partitioning_index` logic with upstream shuffling #60

rjzamora · 2024-05-10T16:34:46Z

Dask modified how partitioning_index is used for shuffling in dask/dask#10705 (included in dask>=2023.12.1). This PR modifies extract_partitioning_index to use the same logic.

TODO:

Testing

rjzamora · 2024-05-10T16:44:51Z

cc @VibhuJawa

ayushdg

Initial review since this is a wip.

I can test the fix out on larger datasets if you feel it's ready for testing.
You might need to amend/update existing commits to include a signoff (commit -s), to pass the DCO check

nemo_curator/utils/fuzzy_dedup_utils/merge_utils.py

Signed-off-by: rjzamora <rzamora217@gmail.com>

This PR adds a new tutorial to demonstrate data curation for PEFT use-cases. Signed-off-by: Mehran Maghoumi <Maghoumi@users.noreply.github.com> Signed-off-by: rjzamora <rzamora217@gmail.com>

Signed-off-by: rjzamora <rzamora217@gmail.com>

rjzamora · 2024-05-13T14:55:48Z

@ayushdg - thanks for the review. Still don't have an evironment to test this myself. I will try to do that later today, but if it's easy for you to test it is very welcome on my end :)

ayushdg · 2024-05-13T23:21:08Z

Can confirm I'm seeing expected results on the larger scale dataset with this fix. I'll run a few more tests on my end but it's generally looking good.
@rjzamora Do you have a small example that checks consistency of the behavior of these two shuffle approaches, that could be added as a unit test with the PR?

* Move PII constants to a seperate file that does not import presidio/spacy and other GPU dependencies Signed-off-by: Ayush Dattagupta <ayushdg95@gmail.com> * Add comment around import, move constant import to global scope Signed-off-by: Ayush Dattagupta <ayushdg95@gmail.com> --------- Signed-off-by: Ayush Dattagupta <ayushdg95@gmail.com>

…itioning_index

rjzamora · 2024-05-14T13:58:03Z

Do you have a small example that checks consistency of the behavior of these two shuffle approaches, that could be added as a unit test with the PR?

@VibhuJawa - Still trying to fix my environment so I can confirm locally, but won't the FuzzyDuplicates tests fail without this fix in place?

ayushdg · 2024-05-14T16:39:04Z

Do you have a small example that checks consistency of the behavior of these two shuffle approaches, that could be added as a unit test with the PR?

@VibhuJawa - Still trying to fix my environment so I can confirm locally, but won't the FuzzyDuplicates tests fail without this fix in place?

In theory the shuffle gives incorrect results, but the dataset/num_partitions here is small enough that it doesn't impact the correctness of final results (duplicate documents detected)

VibhuJawa

Approving, lets merge it in when we add the requested unit test.

Signed-off-by: rjzamora <rzamora217@gmail.com>

tests/test_fuzzy_dedup.py

rjzamora · 2024-05-15T13:35:30Z

tests/test_fuzzy_dedup.py

+        npartitions_right,
+        npartitions_right,


I'm having trouble wrapping my head around the "correct" way to test the batched case. However, my sense is that this test already covers the critical requirement that we are extracting a partitioning index that is consistent with ddf.shuffle.

Yeah. I don't think this specific bug impacts the batched case any differently than the reproducer here so I think it should be good to go.

Signed-off-by: rjzamora <rzamora217@gmail.com>

rjzamora · 2024-05-15T14:00:45Z

nemo_curator/utils/fuzzy_dedup_utils/merge_utils.py

+    from nemo_curator.utils.fuzzy_dedup_utils.shuffle_utils import (
+        rearange_by_column_direct,
+    )


Note that I moved this import into merge_left_to_shuffled_right, because shuffle_utils currently requires cudf/dask_cuda, while other logic in this module does not.

Signed-off-by: rjzamora <rzamora217@gmail.com>

ayushdg

Thanks a lot @rjzamora

ayushdg · 2024-05-15T18:22:23Z

tests/test_fuzzy_dedup.py

+        npartitions_right,
+        npartitions_right,


Yeah. I don't think this specific bug impacts the batched case any differently than the reproducer here so I think it should be good to go.

…DIA#60) * update extract_partitioning_index with compat code Signed-off-by: rjzamora <rzamora217@gmail.com> * [Tutorials] Add a tutorial for PEFT data curation (NVIDIA#45) This PR adds a new tutorial to demonstrate data curation for PEFT use-cases. Signed-off-by: Mehran Maghoumi <Maghoumi@users.noreply.github.com> Signed-off-by: rjzamora <rzamora217@gmail.com> * move compat code to _compat file Signed-off-by: rjzamora <rzamora217@gmail.com> * Only import PII constants during Curator import (NVIDIA#61) * Move PII constants to a seperate file that does not import presidio/spacy and other GPU dependencies Signed-off-by: Ayush Dattagupta <ayushdg95@gmail.com> * Add comment around import, move constant import to global scope Signed-off-by: Ayush Dattagupta <ayushdg95@gmail.com> --------- Signed-off-by: Ayush Dattagupta <ayushdg95@gmail.com> * add unit test Signed-off-by: rjzamora <rzamora217@gmail.com> * add pytest.mark.gpu Signed-off-by: rjzamora <rzamora217@gmail.com> * move extract_partitioning_index import for now Signed-off-by: rjzamora <rzamora217@gmail.com> * test both cudf and pandas Signed-off-by: rjzamora <rzamora217@gmail.com> * spelling Signed-off-by: rjzamora <rzamora217@gmail.com> --------- Signed-off-by: rjzamora <rzamora217@gmail.com> Signed-off-by: Mehran Maghoumi <Maghoumi@users.noreply.github.com> Signed-off-by: Ayush Dattagupta <ayushdg95@gmail.com> Co-authored-by: Mehran Maghoumi <Maghoumi@users.noreply.github.com> Co-authored-by: Ayush Dattagupta <ayushdg95@gmail.com> Signed-off-by: Vibhu Jawa <vibhujawa@gmail.com>

ayushdg requested changes May 10, 2024

View reviewed changes

nemo_curator/utils/fuzzy_dedup_utils/merge_utils.py Outdated Show resolved Hide resolved

rjzamora and others added 3 commits May 13, 2024 07:07

update extract_partitioning_index with compat code

68af3f9

Signed-off-by: rjzamora <rzamora217@gmail.com>

[Tutorials] Add a tutorial for PEFT data curation (NVIDIA#45)

644739b

This PR adds a new tutorial to demonstrate data curation for PEFT use-cases. Signed-off-by: Mehran Maghoumi <Maghoumi@users.noreply.github.com> Signed-off-by: rjzamora <rzamora217@gmail.com>

move compat code to _compat file

1f28a35

Signed-off-by: rjzamora <rzamora217@gmail.com>

rjzamora force-pushed the revise-extract_partitioning_index branch from d453b6e to 1f28a35 Compare May 13, 2024 14:08

ayushdg and others added 2 commits May 14, 2024 06:38

Merge remote-tracking branch 'upstream/main' into revise-extract_part…

33064e8

…itioning_index

VibhuJawa approved these changes May 15, 2024

View reviewed changes

add unit test

5da12ce

Signed-off-by: rjzamora <rzamora217@gmail.com>

rjzamora marked this pull request as ready for review May 15, 2024 03:58

rjzamora commented May 15, 2024

View reviewed changes

tests/test_fuzzy_dedup.py Outdated Show resolved Hide resolved

rjzamora commented May 15, 2024

View reviewed changes

rjzamora added 2 commits May 15, 2024 06:49

add pytest.mark.gpu

25c1eb2

Signed-off-by: rjzamora <rzamora217@gmail.com>

move extract_partitioning_index import for now

647406f

Signed-off-by: rjzamora <rzamora217@gmail.com>

rjzamora force-pushed the revise-extract_partitioning_index branch from b0e41ab to 647406f Compare May 15, 2024 13:50

test both cudf and pandas

b48a37c

Signed-off-by: rjzamora <rzamora217@gmail.com>

rjzamora commented May 15, 2024

View reviewed changes

spelling

9ef4aa3

Signed-off-by: rjzamora <rzamora217@gmail.com>

ayushdg approved these changes May 15, 2024

View reviewed changes

ayushdg merged commit ecd4f4b into NVIDIA:main May 15, 2024
3 checks passed

rjzamora deleted the revise-extract_partitioning_index branch May 15, 2024 19:14

rjzamora mentioned this pull request Jul 3, 2024

Explicit-comms shuffle produces different partitioning than "tasks" rapidsai/dask-cuda#1355

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Align `extract_partitioning_index` logic with upstream shuffling #60

Align `extract_partitioning_index` logic with upstream shuffling #60

rjzamora commented May 10, 2024 •

edited

Loading

rjzamora commented May 10, 2024

ayushdg left a comment •

edited

Loading

rjzamora commented May 13, 2024 •

edited

Loading

ayushdg commented May 13, 2024

rjzamora commented May 14, 2024

ayushdg commented May 14, 2024 •

edited

Loading

VibhuJawa left a comment

rjzamora May 15, 2024

ayushdg May 15, 2024

rjzamora May 15, 2024

ayushdg left a comment

ayushdg May 15, 2024

Align extract_partitioning_index logic with upstream shuffling #60

Align extract_partitioning_index logic with upstream shuffling #60

Conversation

rjzamora commented May 10, 2024 • edited Loading

rjzamora commented May 10, 2024

ayushdg left a comment • edited Loading

Choose a reason for hiding this comment

rjzamora commented May 13, 2024 • edited Loading

ayushdg commented May 13, 2024

rjzamora commented May 14, 2024

ayushdg commented May 14, 2024 • edited Loading

VibhuJawa left a comment

Choose a reason for hiding this comment

rjzamora May 15, 2024

Choose a reason for hiding this comment

ayushdg May 15, 2024

Choose a reason for hiding this comment

rjzamora May 15, 2024

Choose a reason for hiding this comment

ayushdg left a comment

Choose a reason for hiding this comment

ayushdg May 15, 2024

Choose a reason for hiding this comment

Align `extract_partitioning_index` logic with upstream shuffling #60

Align `extract_partitioning_index` logic with upstream shuffling #60

rjzamora commented May 10, 2024 •

edited

Loading

ayushdg left a comment •

edited

Loading

rjzamora commented May 13, 2024 •

edited

Loading

ayushdg commented May 14, 2024 •

edited

Loading