Add a shared deterministic data-partitioning utility; start by de-duplicating the seven CIFAR-10 Dirichlet splitters

[holgerroth]

Summary

Data partitioning code is re-implemented across the examples instead of living in one place. The CIFAR-10 Dirichlet label-skew partitioner alone exists as seven near-identical copies, and the nvflare package has no partitioning utility at all. This proposes a small, framework-agnostic partitioning module in the package and, as the first step, migrating the CIFAR-10 copies onto it.

Evidence (measured on main)

What	Count
Files under examples/, research/, skills/ named like prepare_data*, *data_utils*, *data_split*	106
Example files defining their own split_* function	34
Example files with their own Dirichlet sampler (np.random.dirichlet)	10
Example files calling torchvision.datasets.CIFAR10(...)	83
Partitioning utilities inside the nvflare package	0

The seven CIFAR-10 Dirichlet partitioners, with token-level similarity of the partition_data body against the examples/advanced/cifar10/pt copy:

File	Similarity	Notes
examples/advanced/cifar10/pt/src/data/cifar10_data_split.py	reference	partition_data(num_sites, alpha, seed), seeds via global np.random.seed
examples/advanced/collab/pt_cifar10/prepare_data.py	1.00	verbatim copy
examples/advanced/job_api/pt/fedavg_script_runner_xsite_val_cifar10.py	0.89
examples/advanced/cifar10/tf/src/data/cifar10_data_utils.py	0.71	_partition_data(num_sites, alpha), no seed parameter
examples/advanced/job_api/tf/src/cifar10_data_split.py	0.70	no seed parameter
examples/tutorials/self-paced-training/part-4_advanced_federated_learning/chapter-7_algorithms_and_workflows/07.2_algorithms/07.2.1_advanced_algos/src/cifar10_data_split.py	0.70	no seed parameter
research/auto-fl-research/data/cifar10_data_split.py	0.69

Same algorithm (per-class Dirichlet proportions with a minimum-shard-size retry loop), but the copies already disagree on whether a seed is a parameter, whether they mutate the global NumPy RNG, and what they return. Fixes made to one copy do not reach the others; the same pattern showed up in review this week, where the shared-cache CIFAR-10 download race and the one-time prepare_data.py convention had to be rediscovered per example.

Other Dirichlet users that could adopt the same helper later: examples/advanced/amplify/src/combine_data.py (and its self-paced copy), examples/advanced/bionemo/downstream/sabdab/prepare_sabdab_data.py, examples/advanced/bionemo/downstream/tap/prepare_tap_data.py, research/fed-bpt/src/data_process.py, research/fedumm/src/common.py.

Proposal

Add one small, NumPy-only module, for example nvflare/app_common/data/partition.py, with unit tests:

def dirichlet_label_partition(labels, num_sites, alpha, seed, min_size=10) -> PartitionResult
def iid_partition(num_samples, num_sites, seed) -> PartitionResult
def quantity_skew_partition(num_samples, num_sites, seed, ...) -> PartitionResult

• PartitionResult carries per-site index arrays plus a per-site class-count summary, and small helpers to write/read the indices as JSON so prepare_data.py scripts and client scripts share one file format.
• Seeding through numpy.random.default_rng(seed); never mutate global RNG state.
• No torch, TensorFlow, or dataset dependencies: dataset download, normalization, and DataLoader construction stay in the examples.
• Document the "prepare once, run offline" convention that hello-lightning, hello-jax, and hello-pt already follow, and point the recipe docs and the conversion skills at the helper so agents stop generating ad-hoc splitters.

First step: CIFAR-10

1. Add the module and tests.
2. Migrate the seven copies above so each example keeps a thin prepare_data.py that calls the helper and writes the per-site index files; the training scripts are unchanged apart from reading those files.
3. Make the TF, Job API, and self-paced copies take an explicit seed like the PT reference, so the two frameworks produce the same partition for the same seed.
4. Add a short docs section under the recipe guide describing the helper and the convention.

Non-goals

• Centralizing dataset loaders or transforms (torchvision, Hugging Face); those belong in the examples.
• Changing the partitioning algorithm; the pilot should reproduce the current PT reference behavior for a given seed.

Acceptance

• nvflare/app_common/data/partition.py with tests covering determinism, per-class balance for the IID case, the minimum-shard retry, and JSON round trip.
• The seven CIFAR-10 files import the helper; git grep np.random.dirichlet examples/ no longer matches them.
• The PT and TF CIFAR-10 examples produce identical partitions for the same seed.