Parallelize symmetric sorting comparisons - #4810
Open
JESUSROYETH wants to merge 1 commit into
Open
JESUSROYETH wants to merge 1 commit into
JESUSROYETH wants to merge 1 commit into
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What this changes
Symmetric sorting comparisons with the default
countagreement method run the forward and reverse spike-matching scans one after the other. The two scans are independent, and the existing Numba kernel already releases the GIL, so nothing stops them running at the same time. (Thedistanceagreement method already computes both directions in one pass, so it's unaffected.)This adds
n_jobstocompare_two_sorters()and the underlying agreement helpers, at the end of each signature. Withn_jobs>1, the reverse scan runs in a worker thread while the caller runs the forward scan. The default stays at1: for very small sortings the thread setup costs more than it saves, and with only two directions, two jobs is the useful ceiling anyway.n_jobsstill goes throughfix_job_kwargs(), but it doesn't route throughTimeSeriesChunkExecutor— there's no time-series chunking here, just two independent whole-input scans, so it only needs one executor worker plus the caller.GroundTruthComparisonusesensure_symmetry=False, so it's unaffected.MultiSortingComparisondoesn't passn_jobsthrough either: it already parallelises across sorter pairs, and nesting a thread pool inside each worker process would oversubscribe the machine. So today's gain is scoped to a directcompare_two_sorters()/SymmetricSortingComparisoncall.Benchmark
I tested the complete
compare_two_sorters()call on a published Kilosort4 sorting with 713 units / 11.69M spikes against its 650-unit curated view / 10.60M spikes. Results are 7 alternating runs on an Intel i9-13900HX, Python 3.12.3, NumPy 2.3.5 and Numba 0.67.0:n_jobs=1n_jobs=2Each measured run produced the same SHA256 over match counts, agreement scores and both Hungarian mappings.
Validation
mainand pass with this change.compare_multiple_sortersprocess-parallelism timing test, unrelated to this change, also passed separately.ThreadPoolExecutorand the underlyingnogil=TrueNumba kernel are available on macOS/Windows too, but I haven't verified this path there.