Skip to content

Parallelize unit plots in export_report - #4757

Open
JESUSROYETH wants to merge 1 commit into
SpikeInterface:mainfrom
JESUSROYETH:perf-parallel-export-report
Open

Parallelize unit plots in export_report#4757
JESUSROYETH wants to merge 1 commit into
SpikeInterface:mainfrom
JESUSROYETH:perf-parallel-export-report

Conversation

@JESUSROYETH

@JESUSROYETH JESUSROYETH commented Sep 3, 2026

Copy link
Copy Markdown
Contributor

Description

export_report renders each unit summary in serial, so the time grows directly with the number of units. This change uses the existing job kwargs to render those independent figures with a process pool instead. The default stays at n_jobs=1.

For disk-backed analyzers, each worker loads the analyzer once, using the real format and backend_options instead of guessing from the path. This prevents a binary_folder whose path happens to end in .zarr from being misdetected.

Template percentiles get prepared before the pool starts, so there are no concurrent writes. Also, numerical thread pools follow max_threads_per_worker.

In-memory analyzers still run in parallel with fork on Linux, but with a spawn process they fall back to serial because the analyzer holds weak references that can't be pickled. A disk-backed analyzer also falls back to serial when an extension is only in memory or has different parameters from the saved version: a worker reloading from disk would not see the same extension state. Read-only binary_folder analyzers fall back as well. An explicit mp_context="fork" is rejected on Windows and falls back to serial on macOS.

Zarr analyzers stay serial because their extension arrays would otherwise be materialized once per worker, which can multiply memory use for large dense analyzers. Saving as binary_folder enables the parallel path without that replication.

Performance

I measured the public function with numerical threads pinned to one. Baseline and candidate order alternated, and setup stayed outside the timed call.

Workload Runs Serial median [range] Parallel median [range] Result
Public MEArec, 10 units / 32 channels, 4 workers 5 3.130 s [3.036, 3.453] 1.984 s [1.971, 2.056] 1.578x / 36.62% lower
Generated, 64 units / 64 channels, 4 workers 5 19.768 s [19.653, 19.870] 6.161 s [6.156, 6.222] 3.209x / 68.83% lower
Generated, 64 units / 64 channels, 12 workers 3 19.743 s [19.658, 19.934] 3.315 s [3.313, 3.331] 5.955x / 83.21% lower
Generated, 256 units / 384 channels, 10 workers 4 77.309 s [73.184, 78.837] 14.847 s [13.318, 15.564] 5.207x / 80.80% lower
Final local code, 24 units / 32 channels, 4 workers 5 4.922 s [4.776, 5.288] 3.255 s [2.873, 3.290] 1.512x / 33.88% lower

The same 16-unit workload was also measured on 8-vCPU Intel, AMD and ARM machines:

CPU Runs Serial median [range] Parallel median [range] Result
Intel Sapphire Rapids 3 5.160 s [5.136, 5.409] 1.938 s [1.648, 1.943] 2.663x / 62.45% lower
AMD Rome/Milan 3 6.345 s [6.327, 6.596] 2.398 s [2.024, 2.400] 2.646x / 62.21% lower
ARM Ampere Altra 2 7.332 s [7.094, 7.569] 2.512 s [2.339, 2.686] 2.918x / 65.73% lower

Across the controlled comparisons, the PNG count and bit-for-bit hash stayed the same. The ARM row has two measured repetitions instead of three, so that limit is explicit here. For the 256-unit case, the export itself saved 62.462 seconds at the median; against a one-hour sorting run, that would reduce the combined sorting-plus-export wall time by about 1.70%, so the larger export-stage speedup should not be confused with the total workflow speedup.

Validation:

  • Linux pytest -m exporters -q: 22 passed, 1391 deselected.
  • Linux pytest -m core -q with Pillow and Matplotlib unavailable: 339 passed, 4 skipped, 1067 deselected.
  • Windows Server 2022 / Python 3.10 pytest -m exporters -q: 20 passed, 3 skipped, 1387 deselected.
  • The regression tests use a real spawn process pool with a sparse binary_folder analyzer and exercise the in-memory Linux fork fast path. They also check the Zarr, in-memory spawn, lazy/mismatched-extension and read-only fallbacks.
  • A serial and a parallel export of the same analyzer produce byte-identical PNGs, including for a two-segment analyzer. A separate worker-init test confirms the reload gets the analyzer's real format and backend_options.
  • The 11 new test cases fail on the matching main revision and pass with this patch; the pre-existing serial test passes on both.
  • The fork tests run only on Linux. Windows rejects that context, while macOS falls back to serial with an actionable warning.
  • Black, whitespace and git diff --check are clean.

I also ran the exporter stage by hand on Windows Server 2022 after fixing two Windows-specific assumptions in the tests. I did not run macOS by hand, so project CI is still the real macOS check.

Fixes #2771.

@JESUSROYETH
JESUSROYETH marked this pull request as draft September 3, 2026 16:54
@JoeZiminski

JoeZiminski commented Sep 3, 2026

Copy link
Copy Markdown
Contributor

Thank a lot for this @JESUSROYETH.

I'm in two minds on this one, because the export report step is not particularly long relative to the sorting itself (e.g. 1< minute after a sorting that takes 1 hour) and it only needs to run once. I'm not sure it is worth the additional complexity in the code. That being said, it is common to have hundreds of units with 384 channels over say 10 cores - I wonder the export time in this case + the performance speed up?

@JESUSROYETH
JESUSROYETH force-pushed the perf-parallel-export-report branch from 3a1f076 to e3532a0 Compare September 3, 2026 18:14
@JESUSROYETH

Copy link
Copy Markdown
Contributor Author

Thanks @JoeZiminski, your concern is valid. I measured the case you suggested using a generated sparse binary_folder analyzer with 256 units, 384 channels and 10 workers. Across four alternating serial/parallel pairs, the serial export took 77.309 s median [73.184, 78.837], while the parallel export took 14.847 s [13.318, 15.564]. This is 5.207x faster for export_report, saving 62.462 s.

For a sorting that takes one hour, however, the combined sorting + report time only drops by about 1.70%. So the export-stage improvement is real, but your point about the total workflow is also correct. The default remains serial, although that does not remove the maintenance cost of the parallel path.

This large benchmark used generated data at the requested scale, so I do not want to overstate it as a real-data end-to-end result. Given the one-time saving and the additional process-safety code, I agree the complexity question is valid. I am happy to simplify the scope, or close the PR if the team considers the maintenance cost is not justified ..

@JESUSROYETH
JESUSROYETH marked this pull request as ready for review September 3, 2026 18:22
@alejoe91 alejoe91 added exporters Related to exporters module performance Performance issues/improvements labels Sep 4, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

exporters Related to exporters module performance Performance issues/improvements

Projects

None yet

Development

Successfully merging this pull request may close these issues.

export_report is very slow

3 participants