Homogenize nthreads handling across the framework - #550
Open
paganol wants to merge 6 commits into
Open
Conversation
Standardize every nthreads parameter on `int | None = None` +
resolve_nthreads(), removing the previous mix of 0/1 literal defaults
and duplicated OMP_NUM_THREADS-reading code:
- utilities.py: resolve NUM_THREADS/NUMBA_NUM_THREADS once at import
time (mirrors lbs.MPI_COMM_WORLD) instead of re-reading the
environment on every call. resolve_nthreads()/resolve_numba_nthreads()
become cheap lookups against the cached value.
- constants.py: NUMBA_NUM_THREADS_ENVVAR now points at its own
NUMBA_NUM_THREADS variable (was accidentally aliased to
OMP_NUM_THREADS), so Numba and ducc0 can be sized independently;
OMP_NUM_THREADS remains the shared fallback.
- simulations.py: Simulation.numba_threads now resolves via
resolve_numba_nthreads(), after _init_missing_params() so a
parameter-file value still takes precedence over the environment.
- hwp_harmonics.py: replaced an inline reimplementation of
resolve_nthreads() with the shared helper.
- beam_convolution.py, maps_and_harmonics.py, scan_map.py,
pointings_in_obs.py, input_sky.py: replaced literal nthreads=0
defaults with None + resolve_nthreads(), and threaded already-resolved
nthreads values into internal _get_pointings_array() calls that were
silently falling back to ducc0's own default instead of the caller's
resolved value.
- grasp2alm.py: PseudoAnalysisSHT.to_alm() gained an nthreads parameter
(was hardcoded to 0 in two direct ducc0 calls).
- mueller_convolver.py intentionally left untouched: it's vendored
from upstream ("maintained externally... to make updating from
outside sources easier") and its only call site already passes an
explicit resolved value.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
- mpi.rst: fix constructor parameter name (numba_num_of_threads -> numba_threads, matching the actual Simulation.__init__ signature, including in the TOML example), document the NUMBA_NUM_THREADS/OMP_NUM_THREADS precedence, and add a new "Threads and ducc0" section explaining resolve_nthreads(), lbs.NUM_THREADS/lbs.NUMBA_NUM_THREADS, and the MPI thread-contention risk of leaving OMP_NUM_THREADS unset with multiple ranks per node. - maps_and_harmonics.rst: replace the stale "nthreads=0 (default)" tip with a description of the new None + resolve_nthreads() default. - beam_convolution.rst: drop the now-unnecessary explicit nthreads=0 from the convolve_sky() example. Verified via `sphinx -b doctest` that the example's expected output is unchanged. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Coverage report
Click to see where and how coverage changed
This report was generated by python-coverage-comment-action |
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
anand-avinash
left a comment
anand-avinash
left a comment
Contributor
There was a problem hiding this comment.
Everything looks good except few missed threading opportunity and addressing the contention.
| def _compute_nthreads() -> int: | ||
| if NUM_THREADS_ENVVAR in os.environ: | ||
| return int(os.environ[NUM_THREADS_ENVVAR]) | ||
| return _hardware_threads() |
Contributor
There was a problem hiding this comment.
This will lead to using all available processes for threading, causing contention. Would be better if we simply set it to return 1, when no envvar is set.
Member
Author
There was a problem hiding this comment.
Hi @anand-avinash, I made some test here and it seems to work fine on clusters. Can we keep as it is and invest more on the documentation? misc.available_hardware_threads of ducc0 seems to work very well. I've also checked whether numba/ducc0 could conflict directly, but I could not find any case.
Two more direct ducc0.healpix calls were missing nthreads entirely, found by auditing every ducc0 call site in input_sky.py and mapmaking/common.py: - input_sky.py: BeamCommon._dipole_map_values() called hpx.pix2vec() without nthreads; now passes self.params.nthreads. - mapmaking/common.py: _compute_pixel_indices_single_detector() had no nthreads parameter at all, so its hpx.ang2pix() call and its _get_pointings_array() call both used ducc0's/that function's own default instead of the caller's resolved value. Added nthreads: int | None = None, resolved via resolve_nthreads(). - mapmaking/h_maps.py: make_h_maps() (the only caller of _compute_pixel_indices_single_detector) had no nthreads parameter either; added one and threaded it through. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Resolves conflicts with PR #549 (fix input sky unit error), which landed on master after this branch was opened. Both changes touch litebird_sim/input_sky.py in non-overlapping regions and merged cleanly; only CHANGELOG.md needed manual conflict resolution (kept both entries). Verified: full test suite (289 passed, 6 skipped) + test_mpi_n4.py + ruff check, all green.
#544) Per Avinash's review comment on #550: with no OMP_NUM_THREADS (or NUMBA_NUM_THREADS) set, resolve_nthreads()/resolve_numba_nthreads() previously fell back to ducc0.misc.available_hardware_threads() -- every hardware thread available to the process. That risks oversubscription on unconfigured multi-rank MPI/OpenMP jobs, which is the more consequential failure mode: HPC jobs are expected to set OMP_NUM_THREADS per rank anyway, so this fallback mainly mattered for laptop/workstation runs, where using every core by default is a much smaller risk than one rank silently grabbing a whole shared node. - utilities.py: drop the ducc0.misc-based _hardware_threads() helper; _compute_nthreads() now falls back to 1 when NUM_THREADS_ENVVAR is unset. _compute_numba_nthreads() is unaffected in structure (still falls back to _compute_nthreads()), so it inherits the same change. - mpi.rst, maps_and_harmonics.rst: updated to describe the new fallback and drop the now-obsolete MPI-oversubscription warning (the new default is safe by construction); the Multithreading intro now says explicitly that the Framework does not use every core by default and that OMP_NUM_THREADS must be set for that. - CHANGELOG.md: flagged as a behavior change. Verified: full test suite (289 passed) + test_mpi_n4.py + ruff check, all green; doctests in the touched docs pages still pass unchanged (their expected output was already thread-count independent). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
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
2 participants
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 #544.
Standardizes every
nthreadsparameter onint | None = None+resolve_nthreads(), removing the previous mix of0/1literal defaults and duplicatedOMP_NUM_THREADS-reading code that motivated the issue.Summary
utilities.py:NUM_THREADS/NUMBA_NUM_THREADSare now resolved once at import time (mirrorslbs.MPI_COMM_WORLD) instead of re-reading the environment on every call.resolve_nthreads()/resolve_numba_nthreads()become cheap lookups against the cached value, and both globals are exposed at the top-levellbsnamespace.constants.py:NUMBA_NUM_THREADS_ENVVARwas accidentally aliased to"OMP_NUM_THREADS"(identical toNUM_THREADS_ENVVAR) — fixed to point at its ownNUMBA_NUM_THREADSvariable, so Numba and ducc0 thread counts can be tuned independently, withOMP_NUM_THREADSas the shared fallback whenNUMBA_NUM_THREADSisn't set.simulations.py:Simulation.numba_threadsnow resolves viaresolve_numba_nthreads(), called after_init_missing_params()so anumba_threadsvalue from a parameter file still takes precedence over the environment (precedence: explicit constructor arg > parameter file >NUMBA_NUM_THREADS>OMP_NUM_THREADS> hardware default).hwp_harmonics.py: replaced an inline reimplementation ofresolve_nthreads()'s logic with the shared helper.beam_convolution.py,maps_and_harmonics.py(interpolate_alm,pixelize_alm,estimate_alm,rotate_alm),scan_map.py,pointings_in_obs.py,input_sky.py: replaced literalnthreads=0defaults withNone+resolve_nthreads().nthreadsvalues into internal_get_pointings_array()calls inbeam_convolution.py,scan_map.py,hwp_harmonics.py, andmapmaking/common.pythat had a resolved value in scope but weren't passing it through, silently falling back to ducc0's own default instead.grasp2alm.py:PseudoAnalysisSHT.to_alm()gained annthreadsparameter (was hardcoded to0in two direct ducc0 calls, with no way to override).mueller_convolver.pyintentionally left untouched: it's vendored from upstream ("maintained externally... to make updating from outside sources easier") and its only call site already passes an explicit resolved value.Testing
test_mpi_n4.pypassed.OMP_NUM_THREADSonly,OMP_NUM_THREADS+NUMBA_NUM_THREADStogether) andSimulationprecedence (explicit arg > parameter-file value > env > hardware default).ruff checkclean.Test plan
pytest test/(non-MPI)pytest test/test_mpi_n4.pyresolve_nthreads/resolve_numba_nthreadsSimulation(numba_threads=...)vs parameter-filenumba_threadsvs env vars