fixed numba attempting to parallelize in already parallelized processes#337
Merged
aaronleesander merged 8 commits intomainfrom Feb 19, 2026
Merged
fixed numba attempting to parallelize in already parallelized processes#337aaronleesander merged 8 commits intomainfrom
aaronleesander merged 8 commits intomainfrom
Conversation
Contributor
|
Warning Rate limit exceeded
⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. 📝 WalkthroughWalkthroughAdds OS-aware multiprocessing context selection and Numba thread-control to the simulator worker initialization; introduces a changelog entry for a new Changes
Sequence Diagram(s)sequenceDiagram
participant Main as Main Process
participant Context as Multiprocessing Context
participant Worker as Worker Process
participant Numba as Numba Runtime
participant TPCTL as threadpoolctl
rect rgba(135,206,235,0.5)
Main->>Context: _get_parallel_context() (select fork/spawn)
Context-->>Main: Context object
Main->>Context: spawn Worker(s) with THREAD_ENV_VARS (incl. NUMBA_NUM_THREADS)
end
rect rgba(144,238,144,0.5)
Context->>Worker: start process -> run _worker_init(env, n_threads)
Worker->>Numba: set_num_threads(n_threads) (if available)
Worker->>TPCTL: limit threadpools (threadpoolctl)
TPCTL-->>Worker: threadpool limits applied
Numba-->>Worker: numba threads capped
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/mqt/yaqs/simulator.py (1)
402-419:⚠️ Potential issue | 🟡 MinorStale section header contradicts the new behavior.
Lines 402–405 still read
"spawn" CONTEXTand warn that"fork"can hang/crash, but the function now deliberately returns"fork"on POSIX. Update the header comment to match the new intent.Suggested fix
-# 8) MULTIPROCESS "spawn" CONTEXT -# On Linux, using "fork" with heavy numerical libs can hang/crash due to -# non-fork-safe OpenMP/BLAS state. "spawn" is the robust cross-platform choice. +# 8) MULTIPROCESS CONTEXT SELECTION +# On Windows, "spawn" is the only safe option. +# On Linux, "fork" avoids the OOM overhead of "spawn" and is safe provided +# that threading libraries are capped to 1 thread before/after forking.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/mqt/yaqs/simulator.py` around lines 402 - 419, The top comment block is outdated and says "spawn" CONTEXT and warns that "fork" can hang/crash while _get_parallel_context now returns "fork" on POSIX; update the header and docstring to reflect the current intent: state that we prefer "fork" on POSIX for performance (with safeguards to cap OpenMP/BLAS threads) and use "spawn" on Windows, and adjust any wording that warns against "fork" to instead note the precautions taken (e.g., capping threading libraries before forking). Ensure references to multiprocessing.get_context and os.name remain correct.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/mqt/yaqs/simulator.py`:
- Around line 407-419: The _get_parallel_context function currently returns
multiprocessing.get_context("fork") for all POSIX systems which is unsafe on
macOS; change the logic so that "fork" is only used when running on Linux (e.g.,
check sys.platform == "linux"), and default to
multiprocessing.get_context("spawn") for macOS and other non-Linux platforms;
also correct the docstring typo "caped" to "capped". Ensure references to the
function name _get_parallel_context and the multiprocessing.get_context calls
are updated accordingly.
In `@tests/test_simulator.py`:
- Line 40: The import line brings in _get_parallel_context and _worker_init but
has an unnecessary noqa comment suppressing PLC2701; remove the trailing " #
noqa: PLC2701" from the import statement so the import remains unchanged but the
no-op suppression is deleted (target the import of _get_parallel_context and
_worker_init in tests/test_simulator.py).
- Around line 86-100: The test leaves THREAD_ENV_VARS entries and
NUMBA_NUM_THREADS in the process environment because the finally block contains
a no-op; update the cleanup to either use pytest.monkeypatch to setenv/undo or
explicitly save and restore all affected variables: in tests/test_simulator.py
replace the no-op with code that restores original values for each name in
THREAD_ENV_VARS (pop new keys or reset to saved values) and remove any permanent
os.environ.setdefault usage in _limit_worker_threads by changing it to set only
for the test (use os.environ.setdefault only when the caller guarantees
restoration) or have _limit_worker_threads accept an env-mutation hook so tests
can isolate changes; ensure _worker_init/_limit_worker_threads and
NUMBA_NUM_THREADS are fully reverted in the finally block.
---
Outside diff comments:
In `@src/mqt/yaqs/simulator.py`:
- Around line 402-419: The top comment block is outdated and says "spawn"
CONTEXT and warns that "fork" can hang/crash while _get_parallel_context now
returns "fork" on POSIX; update the header and docstring to reflect the current
intent: state that we prefer "fork" on POSIX for performance (with safeguards to
cap OpenMP/BLAS threads) and use "spawn" on Windows, and adjust any wording that
warns against "fork" to instead note the precautions taken (e.g., capping
threading libraries before forking). Ensure references to
multiprocessing.get_context and os.name remain correct.
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/mqt/yaqs/simulator.py (1)
403-421:⚠️ Potential issue | 🟡 MinorStale section comment contradicts the implementation.
Lines 404–406 still say
"fork" with heavy numerical libs can hang/crashand"spawn" is the robust cross-platform choice, but the function now deliberately uses"fork"on Linux. The docstring (lines 411–413) is correct; the section header should be updated to match.Suggested fix
# --------------------------------------------------------------------------- -# 8) MULTIPROCESS "spawn" CONTEXT -# On Linux, using "fork" with heavy numerical libs can hang/crash due to -# non-fork-safe OpenMP/BLAS state. "spawn" is the robust cross-platform choice. +# 8) MULTIPROCESS CONTEXT SELECTION +# On Linux, "fork" is used for speed (safe when thread pools are capped). +# On Windows and macOS, "spawn" is the standard/required choice. # ---------------------------------------------------------------------------🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/mqt/yaqs/simulator.py` around lines 403 - 421, The section header comment above _get_parallel_context is stale and contradicts the implementation: update the top block comment (lines describing "MULTIPROCESS 'spawn' CONTEXT" and "spawn is the robust cross-platform choice") to reflect the current behavior that uses "fork" on Linux and "spawn" on Windows/macOS; keep the existing docstring in _get_parallel_context as-is, and make the header explain that we prefer "fork" on Linux when safe and "spawn" otherwise to avoid confusion.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/mqt/yaqs/simulator.py`:
- Around line 386-392: The inline comment before attempting to import numba is
truncated and the exception handling is inconsistent with _worker_init: complete
the comment sentence to clearly state that numba threading must be set before
threadpoolctl limits backend threads, and update _worker_init to catch both
ImportError and AttributeError (matching the try/except around
numba.set_num_threads) so both places handle the same edge-case Numba builds;
ensure you reference the numba.set_num_threads call and the _worker_init
function when making these changes.
In `@tests/test_simulator.py`:
- Around line 104-117: The finally block that restores environment variables
mutates os.environ while iterating it (for key in os.environ) which raises
RuntimeError; fix by iterating over a static list of keys (e.g., for key in
list(os.environ)) or, better, replace the entire manual cleanup with
pytest.MonkeyPatch to setenv/undo changes; update the code that references
env_snapshot and the deletion loop in the finally block (and keep
numba.set_num_threads(original_numba_threads) as-is) so you restore added keys
by iterating list(os.environ) or let MonkeyPatch handle env isolation
automatically.
---
Outside diff comments:
In `@src/mqt/yaqs/simulator.py`:
- Around line 403-421: The section header comment above _get_parallel_context is
stale and contradicts the implementation: update the top block comment (lines
describing "MULTIPROCESS 'spawn' CONTEXT" and "spawn is the robust
cross-platform choice") to reflect the current behavior that uses "fork" on
Linux and "spawn" on Windows/macOS; keep the existing docstring in
_get_parallel_context as-is, and make the header explain that we prefer "fork"
on Linux when safe and "spawn" otherwise to avoid confusion.
---
Duplicate comments:
In `@tests/test_simulator.py`:
- Line 40: The import line in tests/test_simulator.py includes an unnecessary
noqa directive ("# noqa: PLC2701") for the symbols _get_parallel_context and
_worker_init; inspect the project's Ruff configuration to see if PL (PLC2701)
rules are enabled and if so keep the directive, otherwise remove the "# noqa:
PLC2701" from the import statement; ensure this change targets the import that
references _get_parallel_context and _worker_init so linting reflects the actual
project config.
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
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
1 participant
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.
Description
This pull request fixes numba from using multiple threads despite the larger simulation already being maximally parallelized such as in a trajectory simulation.
This led to crashes for certain problems on certain operating systems.
Additionally, this now allows numba to use multiple threads when the simulation is deterministic or serialized.
Finally, Linux now uses fork rather than spawn which seems to stabilize simulations (despite expectation of the opposite).
This should be monitored in the future.
Checklist: