perf: make proposal clustering linear in SubgraphSampler

[anna-grim]

Problem

Whole-brain split inference stalls at Inference: 0%| | 0/2304627 and never emits a batch.

SubgraphSampler.set_proposal_clusters accumulated its visited set by rebinding:

visited = visited.union(cluster)

set.union() allocates a new set containing everything seen so far, once per cluster. With ~2.3M proposals forming ~1.44M clusters, that is roughly 1.6e12 element copies.

This runs in SubgraphSampler.__init__, before the first batch is yielded, and the sampler is constructed inside the dataset's producer thread — so the progress bar sits at 0 for the entire time with no indication of work. It also recurs on every inference round, since the sampler is rebuilt each round.

It went unnoticed because at block-level scales (tens of thousands of proposals) it costs a second or two.

Measurements

The pattern in isolation, union vs. in-place update:

n	visited.union(...)	visited.update(...)
10k	0.30 s	0.0023 s
20k	1.39 s	0.0045 s
40k	6.35 s	0.0097 s
80k	27.05 s	0.0200 s

Clean quadratic — 4x the time for 2x the input. Extrapolated to 2,304,627 proposals: ~6 hours per inference round.

Fix

visited.update(cluster)

Equivalent, since visited is only read by the proposal not in visited membership test — contents are identical, only the allocation behavior changes.

Full set_proposal_clusters after the fix, benchmarked against a synthetic proposal graph with realistic cluster sizes:

n=   50,001   0.14s
n=  200,000   0.79s
n=  800,000   3.39s
n=2,300,000   7.29s   <- whole-brain scale

Linear. ~7 s instead of ~6 h.

Notes

• I swept the rest of src/ for the same accumulate-by-copy pattern. The only other .union in a comparable position is proposal_graph.py:333, which is a single call outside a loop and is fine.
• Not addressed here: merge_proposals calls rx.graph_has_path per above-threshold proposal, each a fresh traversal of the full graph. That is the next likely bottleneck at whole-brain scale, but it is pre-existing behavior and separate from this stall.
• No tests — tests/ contains only __init__.py.

🤖 Generated with Claude Code