Skip to content

fix: replace ProcessPoolExecutor with ThreadPoolExecutor (#552) - #566

Open
parallelArchitect wants to merge 1 commit into
mehta-lab:mainfrom
parallelArchitect:fix-threaded-inverse-transfer
Open

fix: replace ProcessPoolExecutor with ThreadPoolExecutor (#552)#566
parallelArchitect wants to merge 1 commit into
mehta-lab:mainfrom
parallelArchitect:fix-threaded-inverse-transfer

Conversation

@parallelArchitect

Copy link
Copy Markdown

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:

  1. 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 spawn context avoided forking but added
    significant startup overhead per worker.

  2. PyTorch thread limiting: torch.set_num_threads(1) and
    set_num_interop_threads(1) were required to prevent thread explosion
    across spawned processes. This artificially constrained PyTorch's
    internal parallelism.

Fix

Replace ProcessPoolExecutor with ThreadPoolExecutor. Threads share
memory — 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: rename processes_optionthreads_option,
    replace --num_processes with --num-threads. Deprecated aliases
    --num-processes / --num_processes retained 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 use threads_option and
    num_threads.
  • waveorder/cli/utils.py: rename num_processesnum_threads in
    estimate_resources.
  • waveorder/calib/calibration_workers.py: update keyword argument.

Testing

200 tests pass locally (cli_tests, api_tests, models).
waveorder/waveorder_simulator.py uses ProcessPoolExecutor for a
different purpose (CPU-bound Jones matrix computation, no zarr I/O) and
is intentionally left unchanged.

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.
@talonchandler
talonchandler requested a review from srivarra July 31, 2026 22:06
@ieivanov

ieivanov commented Aug 5, 2026

Copy link
Copy Markdown
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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Replace mp.Pool with ThreadPoolExecutor in apply_inverse_transfer_function_single_position

2 participants