fix(datasets) Fix partition inconsistencies across dataset splits

[adamtupper]

Issue

This pull request addresses Issue #5243.

Description

These changes addresses a similar issue in the NaturalIdPartitioner and DirichletParitioner that results in mismatches between dataset splits (e.g., for training, validation, and testing).

• For the NaturalIdPartitioner, natural IDs were mapped to partition IDs based on the ordering of examples in the dataset. This resulted in different natural ID to partition ID mappings if the dataset was ordered differently, or on different subsets of the same dataset.
• For the DirichletPartitioner, a similar issue meant that changing the order of the dataset (or generating partitions for different subsets of the same dataset) led to different label distributions for each partition.

Related issues/PRs

Fixes #5243

Proposal

Sort the unique_natural_ids and unique_classes for the NaturalIdPartitioner and DirichletParitioner, respectively.

Explanation

This fix ensures that partition generation is no longer dependent on the ordering of examples in the dataset. This means that, providing the set of natural IDs/unique classes is the same between dataset splits, the partitions are consistently labeled across the partitioners for the different splits. An example of a scenario where this is important is for personalized FL, where each client has their own unique training and test data.

Checklist

• Implement proposed change
• Write tests
• Update documentation
• Make CI checks pass
• Ping maintainers on Slack (channel #contributions)

Any other comments?

[adamtupper]

Hey everyone!

Just a heads-up on the changes I made to the tests. I adjusted the _dummy_setup methods to shuffle the dataset, which changes the ordering of examples for different partitioners. This introduces some randomness and highlights the issue I mentioned earlier. The downside is that the tests aren't deterministic anymore, but I felt this was a reasonable trade-off compared to making bigger changes to the setup methods.

If anyone has suggestions or a cleaner solution, I'm all ears!

[jafermarq]

Hey everyone!

Just a heads-up on the changes I made to the tests. I adjusted the _dummy_setup methods to shuffle the dataset, which changes the ordering of examples for different partitioners. This introduces some randomness and highlights the issue I mentioned earlier. The downside is that the tests aren't deterministic anymore, but I felt this was a reasonable trade-off compared to making bigger changes to the setup methods.

If anyone has suggestions or a cleaner solution, I'm all ears!

Thanks for highlighting this! I agree that introducing this randomness isn't super ideal so I was wondering if we can keep the construction of the test datasets in _dummy_setup as is and instead create new datasets in the two new test methods introduced. For example, how about using those you presented in your #5243?

[adamtupper]

Thanks for your feedback @jafermarq! Definitely, that seems like a better option. I'll update the tests and rollback the changes to _dummy_setup.

[adamtupper]

@jafermarq I've modified the new consistency tests so that they're deterministic and reverted the changes I made to the setup methods. Thanks for your patience!

Let me know if there's anything else that needs to be improved before merging 😄

[jafermarq]

@adamtupper , many thanks for making those changes!

[jafermarq]

🚀 Thanks for the fixes! @adamtupper