Fit processors on the training split in set_task (split=PatientSplit)

[solarsys]

Stacked on #1266 (partial processors), so this diff includes its commit. Merge #1266 first; this PR's own change is the last commit, 9ce56ba.

Implements option A + S1 from the design note on fitting processors on the training split, as agreed.

What it does

from pyhealth.datasets import PatientSplit

train, val, test = dataset.set_task(task, split=PatientSplit(ratios=(0.7, 0.1, 0.2), seed=42))
train.fit_split   # {'kind': 'patient', 'ratios': [0.7, 0.1, 0.2], 'seed': 42}

• Fitting: every processor is fitted on the training patients only, vocabularies included, so test-only codes become <unk>.
• Processing: validation and test are processed with those frozen processors.
• Return value: a tuple of SampleDataset.subset views over one processed cache: (train, val, test), or (train, test) for two ratios.
• Determinism: PatientSplit sorts patient IDs before the seeded shuffle, so it gives the same patients on every run and machine.
• Default unchanged: set_task(task) without split fits on all samples and returns one dataset, with the same cache key as before.
• New flag FeatureProcessor.learns_statistics (default False). When a processor with the flag set is fitted on all samples, one warning points to split=. Vocabulary processors stay quiet. Nothing in core sets it yet; Add training-fitted normalization to time-series processors #1231's normalization would.

Streaming (S1)

1. Index pass: unchanged. It builds patient_to_index, which master already keeps in memory and pickles into schema.pkl.
2. Split: from that index. Patients are sorted, shuffled and cut by ratio. The only new memory is one int64 index array per part plus a boolean training mask: about 9 bytes per sample.
3. Fit: each processor reads the same sequential stream through a masked view that skips other patients. There is no random access and nothing is collected, and it decodes exactly what master's fit decodes.
4. Process: unchanged, one streamed pass.
5. Cache: split.npz stores the part indices, written before the samples, whose index.json marks the cache complete. The processed-cache key gains the split spec only when split is given.

Benchmarks

Synthetic MIMIC-III-format data, so no patient data was used. Generated and run by the new examples/benchmark_perf/benchmark_split_set_task.py, with MIMIC3Dataset and MortalityPredictionMIMIC3. Each mode ran in its own process with a fresh cache. Peak memory is the process plus its workers, sampled every 0.1 s. Machine: 10-core Mac, 34 GB.

Patients	Samples	Event rows	Workers	Mode	set_task time	set_task peak memory
20,000	49,884	1.4M	1	no split / split	116 s / 105 s	627 / 768 MiB
100,000	250,144	7.2M	1	no split / split	551 s / 541 s	1,482 / 1,461 MiB
300,000	749,901	21.5M	4	no split / split	243 s / 245 s	5,753 / 5,358 MiB

Whole-pipeline peaks are noisy. The event-table stage, identical code in both modes, peaked at 2,856 vs 4,234 MiB at 20k. So I also isolated SampleBuilder.fit on the real 100k task output, 3 repetitions each in one process:

250,144 samples, 83,146 patients	No split	Split
Fit time (median)	46.1 s	45.2 s
Peak Python allocations during fit	91.9 MiB	85.4 MiB
patient_to_index (exists on master)	8.6 MiB, 36 bytes/sample	same
Added by S1 (index arrays + mask)	—	2.1 MiB, ~9 bytes/sample

For laptops, extrapolating linearly:

• Patient index: the one master already holds is about 36 MB per million samples.
• S1: adds about 9 MB per million samples.
• The real peak in these runs comes from the event-table build and the task workers, which this PR doesn't change.

Correctness at scale (300k, 4 workers). I reloaded the split from cache (3.2 s), streamed all 749,901 samples, and compared each part's streamed patients with its index:

• The parts match their indices exactly: 174,902 / 24,986 / 49,973 patients.
• Zero patients appear in more than one part.
• The sample counts add up to the unsplit total at every scale.

To run the same on a laptop:

python examples/benchmark_perf/benchmark_split_set_task.py --patients 100000            # 1 worker
python examples/benchmark_perf/benchmark_split_set_task.py --patients 300000 --num-workers 4 --workdir /tmp/bench

Tests

New tests/core/test_split_set_task.py (14 tests) runs on a config-based dataset where test patients' values are shifted by +100 and carry a test-only code. It covers:

• parts that are disjoint and complete;
• the mean-learning processor's mean equalling the training mean;
• test-only codes mapping to <unk>;
• fit_split recorded;
• a rerun reusing the cache and the same split;
• another seed giving another split and cache;
• the two-part split;
• the unsplit cache key unchanged;
• supplied processors kept;
• the warning (and none for vocabulary-only);
• invalid ratios;
• fit with a split staying under 5 MiB of peak allocations on a ~50 MB lazy stream (it also fails if len() is called on the stream).

A mutation check confirmed the tests catch leakage: forcing fit on all samples makes 4 of them fail. The set_task signature check in test_caching.py now includes split.

Full core suite: 1,415 tests. The one failure was that pinned-signature test, which is updated here; it and the split tests then pass. tools/check_pr_rules.py passes. The docs get a section on the splitter page and a pointer on the processors page, and there's a new runnable examples/split_aware_set_task_mimic3demo.py.

Not in this PR

• K-fold and temporal split specs (P1-7): they would be further split= values.
• The tabular processor (P1-6): it waits for Add training-fitted normalization to time-series processors #1231.

🤖 Generated with Claude Code

[jhnwu3]

I've added a couple of additional fixes to the problem here

1. Added a Split() abstract class. There are many different types of splits that exist out there that move beyond just a 3 way split. Other uses
2. The doc strings should probably highlight more of the necessary functions here.
3. Bumped up the litdata version, so we no longer have a bug that unintentionally leaks patient data at higher worker counts.