fix: replace ProcessPoolExecutor with ThreadPoolExecutor (#552) - #566
Open
parallelArchitect wants to merge 1 commit into
Open
fix: replace ProcessPoolExecutor with ThreadPoolExecutor (#552)#566parallelArchitect wants to merge 1 commit into
parallelArchitect wants to merge 1 commit into
Conversation
Process-based parallelism duplicates zarr-backed data into each worker via copy-on-write, causing severe memory inflation on large volumes (35-45 GB RSS for a 750 MB dataset reported in mehta-lab#552). Thread-based parallelism shares memory across workers — no copying. numpy, scipy, and PyTorch all release the GIL during compute (FFTs, Tikhonov solves), so threads achieve genuine parallelism on the compute-heavy parts of the reconstruction pipeline. The spawn context (mp.get_context('spawn')) was required to avoid tensorstore fork-safety hazards (C++ background threads holding locks that deadlock in forked children, see google/tensorstore#61). Threads share the existing process state, so this hazard disappears entirely. torch.set_num_threads(1) / set_num_interop_threads(1) were set to prevent thread explosion across spawned processes. With a single shared process, PyTorch manages its thread pool correctly without manual limiting. --num-processes / --num_processes retained as deprecated CLI aliases mapping to --num-threads for backward compatibility.
Contributor
|
The PoolExecutor has been very carefully chosen to work with our biahub reconstruction pipeline. I advise not changing the default. We could offer an option to slot in a different executor if needed for certain applications. |
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.
Fixes #552
Problem
Process-based parallelism duplicates zarr-backed data into each worker
via copy-on-write. On large volumes this causes severe memory inflation —
#552 reports 35–45 GB RSS for a 750 MB dataset with 4 workers.
Two additional hazards with the process-based approach:
tensorstore fork-safety: tensorstore runs internal C++ threads
that are not fork-safe. A forked worker inherits locked mutexes from
threads that don't exist in the child, causing deadlocks or segfaults
(Segfault when using PyTorch Dataloader with multiple workers google/tensorstore#61). The
spawncontext avoided forking but addedsignificant startup overhead per worker.
PyTorch thread limiting:
torch.set_num_threads(1)andset_num_interop_threads(1)were required to prevent thread explosionacross spawned processes. This artificially constrained PyTorch's
internal parallelism.
Fix
Replace
ProcessPoolExecutorwithThreadPoolExecutor. Threads sharememory — no copying, no fork-safety hazards, no spawn overhead.
numpy, scipy, and PyTorch all release the GIL during compute (FFTs,
Tikhonov solves), so threads achieve genuine parallelism on the
compute-heavy parts of the reconstruction pipeline.
Changes
waveorder/cli/parsing.py: renameprocesses_option→threads_option,replace
--num_processeswith--num-threads. Deprecated aliases--num-processes/--num_processesretained for backward compatibility.waveorder/cli/apply_inverse_transfer_function.py: swap executor,remove spawn context and torch thread limiting.
waveorder/cli/reconstruct.py: update to usethreads_optionandnum_threads.waveorder/cli/utils.py: renamenum_processes→num_threadsinestimate_resources.waveorder/calib/calibration_workers.py: update keyword argument.Testing
200 tests pass locally (cli_tests, api_tests, models).
waveorder/waveorder_simulator.pyusesProcessPoolExecutorfor adifferent purpose (CPU-bound Jones matrix computation, no zarr I/O) and
is intentionally left unchanged.