diff --git a/CHANGELOG.md b/CHANGELOG.md index 2e54ca16..4cccec87 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,7 @@ # HEAD +- Homogenize `nthreads` handling across the framework: every function now defaults to `nthreads: int | None = None` and resolves it via `resolve_nthreads()`/`resolve_numba_nthreads()`, removing the previous mix of `0`/`1` literal defaults and duplicated `OMP_NUM_THREADS`-reading code. Thread counts are resolved once at import time (`lbs.NUM_THREADS`, `lbs.NUMBA_NUM_THREADS`, mirroring `lbs.MPI_COMM_WORLD`) instead of on every call. Numba threads can now be tuned independently via `NUMBA_NUM_THREADS`, falling back to `OMP_NUM_THREADS`. **Behavior change**: when neither environment variable is set, threading now defaults to a single thread instead of every hardware thread available to the process, to avoid oversubscription in MPI/multi-process runs; set `OMP_NUM_THREADS` (and, if needed, `NUMBA_NUM_THREADS`) explicitly to use more than one core. Reference PR: [#550](https://github.com/litebird/litebird_sim/pull/550). + - Fix wrong output units when `bandpass_integration=True` in `input_sky` (CMB and dipole were off by a large factor because `pysm3.bandpass_unit_conversion` wasn't told the input unit), remove the resulting dead code path, and document that `cmb_ps_file` spectra are expected in $\mu K_{CMB}^2$. Fixes [#548](https://github.com/litebird/litebird_sim/issues/548). Reference PR: [#549](https://github.com/litebird/litebird_sim/pull/549). - Stop supporting unused MPI processes and drop `MPI_COMM_GRID` class and its attributes. Now throwing an error when `comm_size != n_blocks_det * n_blocks_time`. It fixes [#364](https://github.com/litebird/litebird_sim/issues/364), while incorporating selective changes from [#372](https://github.com/litebird/litebird_sim/pull/372) and [#427](https://github.com/litebird/litebird_sim/pull/427). Reference PR: [#539](https://github.com/litebird/litebird_sim/pull/539). diff --git a/docs/source/beam_convolution.rst b/docs/source/beam_convolution.rst index fa162a41..bed67d6a 100644 --- a/docs/source/beam_convolution.rst +++ b/docs/source/beam_convolution.rst @@ -392,8 +392,7 @@ For a single-task execution, refer to the following example: sim.convolve_sky(sky_alms=alms, beam_alms=blms, convolution_params=Convparams, - pointings_dtype=np.float32, - nthreads = 0) + pointings_dtype=np.float32) API reference diff --git a/docs/source/maps_and_harmonics.rst b/docs/source/maps_and_harmonics.rst index b84ef7eb..abe4ae65 100644 --- a/docs/source/maps_and_harmonics.rst +++ b/docs/source/maps_and_harmonics.rst @@ -107,8 +107,13 @@ These functions handle the complexity of spin-0 (Temperature) vs spin-2 (Polariz * :func:`~litebird_sim.maps_and_harmonics.compute_dl`: Compute $D_{\ell} = \ell(\ell+1)C_{\ell}/(2\pi)$ spectra .. tip:: - All transform functions accept a ``nthreads`` argument. - Setting ``nthreads=0`` (default) uses all available hardware threads, which is optimal for standalone scripts but should be adjusted when running inside an MPI environment. + All transform functions accept an ``nthreads`` argument, which defaults to ``None``. + When left unset, the number of threads is resolved automatically (see :ref:`nthreads_ducc0`): + it honours the ``OMP_NUM_THREADS`` environment variable if set, and otherwise falls back + to ``1`` thread. This makes the default safe to use both in standalone scripts and inside + an MPI environment, but it also means you need to set ``OMP_NUM_THREADS`` yourself if you + want these functions to use more than one core. Pass an explicit ``nthreads`` to override + this for a single call. .. note:: All transform functions support multi-frequency data. When operating on multi-frequency objects, transforms are applied independently to each frequency, and the output maintains the multi-frequency structure. diff --git a/docs/source/mpi.rst b/docs/source/mpi.rst index f4627b08..32d26431 100644 --- a/docs/source/mpi.rst +++ b/docs/source/mpi.rst @@ -30,20 +30,37 @@ Some parts of the LiteBIRD Simulation Framework are able to exploit multiple cores because several of its modules rely on the `Numba `_ library. -If you are running your code on your multi-core laptop, you do not -have to do anything fancy in order to use all the CPUs on your machine: -in its default configuration, the Framework should be able to take -advantage all the available CPU cores. - -However, if you want to tune the way the Framework uses the CPUs, -you can either set the environment variable ``OMP_NUM_THREADS`` -to the number of CPUs to use, or use two parameters in -the constructor of the class :class:`.Simulation`: - -- `numba_num_of_threads`: this is the number of CPUs that Numba will +To use more than one CPU, you must explicitly tell the Framework how +many to use, either by setting the environment variable +``OMP_NUM_THREADS`` to the number of CPUs to use, or by using two +parameters in the constructor of the class :class:`.Simulation`. +The Framework does **not** default to using every CPU core available +on the machine: doing so would be unsafe under MPI, where several +ranks typically share a node and would otherwise all try to grab every +core at once. Leaving thread counts unconfigured makes the Framework +run single-threaded, both for Numba and for ducc0 (see +:ref:`nthreads_ducc0` below); this is a safe default, but it means +that on a single-user laptop or workstation you *do* need to set +``OMP_NUM_THREADS`` (or the equivalent parameters below) yourself if +you want to take advantage of all the available cores. + +- `numba_threads`: this is the number of CPUs that Numba will use for parallel calculations. The parameter defaults to ``None``, - which means that Numba will check how many CPUs are available and will - use all of them. + in which case it is resolved once, following this order of precedence: + + 1. the value of `numba_threads` itself, if you passed one explicitly + (or if it was set in a TOML parameter file, see below); + 2. the environment variable ``NUMBA_NUM_THREADS``, if set; + 3. the environment variable ``OMP_NUM_THREADS``, if set; + 4. ``1``, if none of the above is set. + + This is the same resolution order used for the ``nthreads`` parameters + accepted by the low-level `ducc0 `_-based + functions of the framework (spherical harmonic transforms, beam + convolution, map scanning, ...) — see :ref:`nthreads_ducc0` below. Using + ``NUMBA_NUM_THREADS`` instead of ``OMP_NUM_THREADS`` lets you give Numba a + different thread count than ducc0, if you ever need to; leaving + ``NUMBA_NUM_THREADS`` unset makes the two agree by default. - `numba_threading_layer`: this parameter is a string that specifies which threading library should be used by Numba. The value depends @@ -71,7 +88,7 @@ These parameters can be passed through a TOML parameter file (see # This is file "my_conf.toml" [simulation] random_seed = 12345 - numba_num_of_threads = 32 + numba_threads = 32 numba_threading_layer = "tbb" Both ``tbb`` and ``omp`` require that the relevant library be available on @@ -87,6 +104,42 @@ of running a command like the following: $ module load openmp # OpenMP +.. _nthreads_ducc0: + +Threads and ducc0 +~~~~~~~~~~~~~~~~~~ + +Besides Numba, most of the computationally-heavy, low-level functions of +the framework (spherical harmonic transforms in +:mod:`litebird_sim.maps_and_harmonics`, beam convolution, map scanning, the +map-makers, ...) delegate their parallel work to the +`ducc0 `_ library through an +``nthreads`` parameter. Every such function defaults to ``nthreads=None``, +which is resolved through :func:`litebird_sim.resolve_nthreads` using the +same precedence as ``OMP_NUM_THREADS`` for Numba above: explicit value, +then ``OMP_NUM_THREADS``, then ``1``. The values used by the current +process, resolved once when ``litebird_sim`` is imported, are exposed as +:data:`litebird_sim.NUM_THREADS` (for ducc0) and +:data:`litebird_sim.NUMBA_NUM_THREADS` (for Numba). + +This ``1``-thread fallback is deliberately conservative: it keeps both +serial and MPI runs safe by default (no risk of a rank grabbing every +core on a shared node), at the cost of not using extra cores unless you +ask for them. If you want multithreaded ducc0/Numba execution — whether +on a laptop or per MPI rank on a cluster — set ``OMP_NUM_THREADS`` (and, +if needed, ``NUMBA_NUM_THREADS``) to the number of cores you want to use +before launching your job, e.g. for 8 MPI ranks on a 64-core node, using +8 threads per rank: + +.. code-block:: sh + + $ export OMP_NUM_THREADS=8 + $ mpirun -n 8 python3 my_script.py + +You can always override the resolved value for a single call by passing an +explicit ``nthreads`` argument to the function you are calling. + + MPI ~~~ diff --git a/litebird_sim/__init__.py b/litebird_sim/__init__.py index a426a2d5..603829ca 100644 --- a/litebird_sim/__init__.py +++ b/litebird_sim/__init__.py @@ -42,6 +42,12 @@ EARTH_L2_DISTANCE_KM, NUM_THREADS_ENVVAR, ) +from .utilities import ( + NUM_THREADS, + NUMBA_NUM_THREADS, + resolve_nthreads, + resolve_numba_nthreads, +) from .coordinates import ( DEFAULT_COORDINATE_SYSTEM, DEFAULT_TIME_SCALE, @@ -253,6 +259,11 @@ "SOLAR_VELOCITY_GAL_LON_RAD", "EARTH_L2_DISTANCE_KM", "NUM_THREADS_ENVVAR", + # utilities.py + "NUM_THREADS", + "NUMBA_NUM_THREADS", + "resolve_nthreads", + "resolve_numba_nthreads", # units.py "Units", "UnitUtils", diff --git a/litebird_sim/beam_convolution.py b/litebird_sim/beam_convolution.py index a02a0f5c..266ef11f 100644 --- a/litebird_sim/beam_convolution.py +++ b/litebird_sim/beam_convolution.py @@ -51,7 +51,7 @@ def add_convolved_sky_to_one_detector( mueller_matrix, hwp_angle, convolution_params: BeamConvolutionParameters | None = None, - nthreads: int = 0, + nthreads: int | None = None, ): """ Convolve given sky alms with a detector beam alms and add the result to the TOD of a single detector. @@ -74,9 +74,10 @@ def add_convolved_sky_to_one_detector( convolution_params : BeamConvolutionParameters, optional Parameters controlling the convolution, such as resolution and precision. If None, reasonable defaults are chosen based on the sky and beam properties. - nthreads : int, default=0 - Number of threads to use for convolution. If set to 0, all available CPU cores - will be used. + nthreads : int or None, default=None + Number of threads to use for convolution. If None, resolved via + :func:`.resolve_nthreads` (``OMP_NUM_THREADS``, or all available + threads if unset). Raises ------ @@ -96,6 +97,8 @@ def add_convolved_sky_to_one_detector( - The function modifies `tod_det` in place by adding the convolved signal. """ + nthreads = resolve_nthreads(nthreads) + if not convolution_params: sky_lmax = sky_alms_det.lmax @@ -214,7 +217,7 @@ def add_convolved_sky( convolution_params: BeamConvolutionParameters | None = None, pointings_dtype=np.float64, nside_centering: int | None = None, - nthreads: int = 0, + nthreads: int | None = None, ): """ Convolve a set of sky maps with detector beams and add the resulting signals to the @@ -253,9 +256,10 @@ def add_convolved_sky( nside_centering : int, default=None If set, shifts the detector pointings to the centers of the corresponding HEALPix pixels at the given NSIDE resolution. If None, no centering is applied. - nthreads : int, default=0 + nthreads : int or None, default=None Number of threads to use for convolution and in case for HEALPix operations. - If set to 0, all available CPU cores will be used. + If None, resolved via :func:`.resolve_nthreads` (``OMP_NUM_THREADS``, or all + available threads if unset). Raises ------ @@ -273,6 +277,8 @@ def add_convolved_sky( - The function modifies `tod` in place by adding the convolved signals for all detectors. """ + nthreads = resolve_nthreads(nthreads) + if mueller_hwp is not None: assert tod.shape[0] == mueller_hwp.shape[0] @@ -327,6 +333,7 @@ def add_convolved_sky( output_coordinate_system=coordinates, nside_centering=nside_centering, pointings_dtype=pointings_dtype, + nthreads=nthreads, ) # FIXME: Fix this at some point, ducc wants phi 0 -> 2pi diff --git a/litebird_sim/constants.py b/litebird_sim/constants.py index 226c9107..c1cc688b 100644 --- a/litebird_sim/constants.py +++ b/litebird_sim/constants.py @@ -2,9 +2,13 @@ from astropy.constants import c as c_light from astropy.constants import h, k_B -# Name of the environment variable used in the convolution +# Environment variables used to size ducc0's and Numba's thread pools. +# NUMBA_NUM_THREADS_ENVVAR takes precedence over NUM_THREADS_ENVVAR when +# resolving the number of threads for Numba, so the two runtimes can be +# sized independently if needed (see resolve_nthreads/resolve_numba_nthreads +# in utilities.py). NUM_THREADS_ENVVAR = "OMP_NUM_THREADS" -NUMBA_NUM_THREADS_ENVVAR = "OMP_NUM_THREADS" +NUMBA_NUM_THREADS_ENVVAR = "NUMBA_NUM_THREADS" ARCMIN_TO_RAD = np.pi / 180 / 60 diff --git a/litebird_sim/grasp2alm.py b/litebird_sim/grasp2alm.py index e7911290..407278aa 100644 --- a/litebird_sim/grasp2alm.py +++ b/litebird_sim/grasp2alm.py @@ -71,6 +71,7 @@ def to_alm( mmax: int, epsilon=1e-8, max_num_of_iterations=20, + nthreads: int | None = None, ) -> np.ndarray: """Converts the beam map to spherical harmonic coefficients. @@ -79,6 +80,8 @@ def to_alm( mmax (`int`): Maximum m value for the spherical harmonic expansion. epsilon (`float`): Precision of the result max_num_of_iterations (`int`): Maximum number of iterations + nthreads (`int` or `None`): Number of threads to use. If None, + resolved via :func:`.resolve_nthreads`. Returns: `numpy.ndarray`: The spherical harmonic coefficients, as a (3, N) array. @@ -88,6 +91,8 @@ def to_alm( """ + nthreads = resolve_nthreads(nthreads) + if not self.map.shape[0] <= 3: raise ValueError( "Error in BeamMap.to_alm: map has more than 3 Stokes parameters" @@ -105,7 +110,7 @@ def to_alm( lmax=lmax, mmax=mmax, spin=0, - nthreads=0, + nthreads=nthreads, maxiter=max_num_of_iterations, epsilon=epsilon, **geom, @@ -123,7 +128,7 @@ def to_alm( lmax=lmax, mmax=mmax, spin=2, - nthreads=0, + nthreads=nthreads, maxiter=max_num_of_iterations, epsilon=epsilon, **geom, diff --git a/litebird_sim/hwp_harmonics/hwp_harmonics.py b/litebird_sim/hwp_harmonics/hwp_harmonics.py index 51484e97..c6cb9ddb 100644 --- a/litebird_sim/hwp_harmonics/hwp_harmonics.py +++ b/litebird_sim/hwp_harmonics/hwp_harmonics.py @@ -1,5 +1,4 @@ import logging -import os import numpy as np import numpy.typing as npt @@ -11,7 +10,6 @@ from litebird_sim.hwp_jones_parameters import HWPJonesParams from ..bandpass_template_module import bandpass_profile -from ..constants import NUM_THREADS_ENVVAR from ..coordinates import CoordinateSystem from ..hwp_non_ideal import HWPFormalism, NonIdealHWP from ..input_sky import SkyInput @@ -20,6 +18,7 @@ from ..pointings_in_obs import ( _get_pointings_array, ) +from ..utilities import resolve_nthreads from .jones_methods import ( compute_signal_for_one_detector as compute_signal_for_one_detector_jones, ) @@ -228,9 +227,7 @@ def fill_tod_with_hwp_harmonics( ) assert maps is not None, "You need to pass input maps to fill_tod." - # Set number of threads - if nthreads is None: - nthreads = int(os.environ.get(NUM_THREADS_ENVVAR, 0)) + nthreads = resolve_nthreads(nthreads) if pointings is None: if hwp_angle is not None: @@ -390,6 +387,7 @@ def fill_tod_with_hwp_harmonics( hwp_angle=cur_hwp_angle, output_coordinate_system=coordinates, pointings_dtype=pointings_dtype, + nthreads=nthreads, ) tod_det = tod[idet, :] diff --git a/litebird_sim/input_sky.py b/litebird_sim/input_sky.py index 0efdc467..9866d521 100644 --- a/litebird_sim/input_sky.py +++ b/litebird_sim/input_sky.py @@ -25,6 +25,7 @@ synthesize_alm, ) from .units import Units, UnitUtils +from .utilities import resolve_nthreads # --- Utility Functions --- @@ -145,7 +146,7 @@ def __init__( # Bandpass bandpass_integration: bool = False, # Parallelism - nthreads: int = 0, # 0 usually means "use all available" in ducc0 + nthreads: int | None = None, # None: resolved via resolve_nthreads() # Components to generate make_cmb: bool = True, make_fg: bool = False, @@ -178,7 +179,7 @@ def __init__( self.bandpass_integration = bandpass_integration self.maxiter = maxiter self.epsilon = epsilon - self.nthreads = nthreads + self.nthreads = resolve_nthreads(nthreads) self.make_cmb = make_cmb self.make_fg = make_fg self.make_dipole = make_dipole @@ -635,7 +636,7 @@ def _dipole_map_values(self) -> np.ndarray: npix = hpx.npix() vec = dh.ang2vec(np.array([[lat, lon]]))[0] - pix_vecs = hpx.pix2vec(np.arange(npix)) + pix_vecs = hpx.pix2vec(np.arange(npix), nthreads=self.params.nthreads) dipole_map_val = np.zeros((3, npix)) dipole_map_val[0] = np.dot(pix_vecs, vec) * amp diff --git a/litebird_sim/mapmaking/common.py b/litebird_sim/mapmaking/common.py index 23d25bc4..5a95e0ee 100644 --- a/litebird_sim/mapmaking/common.py +++ b/litebird_sim/mapmaking/common.py @@ -10,6 +10,7 @@ from litebird_sim.coordinates import CoordinateSystem from litebird_sim.observations import Observation from litebird_sim.pointings_in_obs import _get_pointings_array, _get_pol_angle +from litebird_sim.utilities import resolve_nthreads # The threshold on the conditioning number used to determine if a pixel # was really “seen” or not @@ -161,6 +162,7 @@ def _compute_pixel_indices( hwp_angle=hwp_angle, output_coordinate_system=output_coordinate_system, pointings_dtype=pointings_dtype, + nthreads=nthreads, ) if hmap_generation: polang_all[idet] = curr_pointings_det[:, 2] @@ -193,10 +195,13 @@ def _compute_pixel_indices_single_detector( output_coordinate_system: CoordinateSystem, pointings_dtype=np.float64, hmap_generation: bool = False, + nthreads: int | None = None, ) -> tuple[npt.NDArray, npt.NDArray]: """ Same as _compute_pixel_indices but for a single detector, thus returning only the pixel indices and polarization angles for that detector. """ + nthreads = resolve_nthreads(nthreads) + pixidx = np.empty((num_of_samples), dtype=np.int32) polang = np.empty((num_of_samples), dtype=pointings_dtype) curr_pointings_det, hwp_angle = _get_pointings_array( @@ -205,6 +210,7 @@ def _compute_pixel_indices_single_detector( hwp_angle=hwp_angle, output_coordinate_system=output_coordinate_system, pointings_dtype=pointings_dtype, + nthreads=nthreads, ) if hmap_generation: @@ -215,7 +221,7 @@ def _compute_pixel_indices_single_detector( hwp_angle=hwp_angle, pol_angle_detectors=pol_angle_detector, ) - pixidx = hpx.ang2pix(curr_pointings_det[:, :2]) + pixidx = hpx.ang2pix(curr_pointings_det[:, :2], nthreads=nthreads) if output_coordinate_system == CoordinateSystem.Galactic: # Free curr_pointings_det if the output map is already in Galactic coordinates diff --git a/litebird_sim/mapmaking/h_maps.py b/litebird_sim/mapmaking/h_maps.py index 9d079a30..3cc1fa25 100644 --- a/litebird_sim/mapmaking/h_maps.py +++ b/litebird_sim/mapmaking/h_maps.py @@ -18,6 +18,7 @@ _get_hwp_angle, _normalize_observations_and_pointings, ) +from litebird_sim.utilities import resolve_nthreads from .common import ( _build_mask_detector_split, @@ -240,6 +241,7 @@ def make_h_maps( pointings_dtype=np.float64, save_to_file: bool = True, output_directory: str = "./h_n_maps", + nthreads: int | None = None, ) -> HMapsResult: """Generate complex harmonic maps :math:`h_{n,m}` from observations. @@ -269,10 +271,15 @@ def make_h_maps( :type save_to_file: bool :param output_directory: Output directory for generated HDF5 files. :type output_directory: str + :param nthreads: Number of threads to use for ducc0 calls. If None, + resolved via :func:`.resolve_nthreads`. + :type nthreads: int | None :returns: Result container with all computed maps and metadata. :rtype: HnMapResult """ + nthreads = resolve_nthreads(nthreads) + assert ( n_m_couples.shape[1] == 2 ), """the n,m couples should be passed in this shape: array([[n1, m1], @@ -331,6 +338,7 @@ def make_h_maps( output_coordinate_system=output_coordinate_system, pointings_dtype=pointings_dtype, hmap_generation=True, + nthreads=nthreads, ) log.info( f"Pixel indices and angles for detector {all_dets_list[idet]} computed, now building nobs matrices" diff --git a/litebird_sim/maps_and_harmonics.py b/litebird_sim/maps_and_harmonics.py index 6668dfe1..c37decfb 100644 --- a/litebird_sim/maps_and_harmonics.py +++ b/litebird_sim/maps_and_harmonics.py @@ -12,6 +12,7 @@ from .coordinates import ECL_TO_GAL_EULER, GAL_TO_ECL_EULER, CoordinateSystem from .units import Units +from .utilities import resolve_nthreads # ====================================================================== # SphericalHarmonics @@ -2369,7 +2370,7 @@ def interpolate_alm( locations: np.ndarray, *, epsilon: float | None = None, - nthreads: int = 0, + nthreads: int | None = None, ) -> np.ndarray | tuple[np.ndarray, np.ndarray, np.ndarray]: r"""Interpolate spherical-harmonic coefficients at arbitrary positions using :func:`ducc0.sht.synthesis_general`. @@ -2407,9 +2408,9 @@ def interpolate_alm( * complex64 → ``1e-6`` * complex128 → ``1e-13`` - nthreads : int, optional - Number of threads for ducc. If 0 (default), ducc uses the - number of hardware threads. + nthreads : int or None, optional + Number of threads for ducc. If None, resolved via + :func:`.resolve_nthreads`. Returns ------- @@ -2426,6 +2427,8 @@ def interpolate_alm( If the input shapes are inconsistent or ``nstokes`` is not 1 or 3. """ + nthreads = resolve_nthreads(nthreads) + alm = np.asarray(alms.values) loc = np.asarray(locations, dtype=np.float64) if loc.ndim != 2 or loc.shape[1] != 2: @@ -2531,7 +2534,7 @@ def pixelize_alm( nest: bool = False, lmax: int | None = None, mmax: int | None = None, - nthreads: int = 0, + nthreads: int | None = None, ) -> "HealpixMap": r""" Convert spherical harmonics coefficients to a HEALPix map using @@ -2584,9 +2587,9 @@ def pixelize_alm( ``lmax``/``mmax`` is requested than what is stored in ``alms``, the coefficient array is truncated accordingly. - nthreads : int, optional - Number of threads passed to ducc. If zero (default), ducc chooses - the number of threads. + nthreads : int or None, optional + Number of threads passed to ducc. If None, resolved via + :func:`.resolve_nthreads`. Returns ------- @@ -2605,6 +2608,8 @@ def pixelize_alm( or if requested ``lmax``/``mmax`` are inconsistent with the available coefficients. """ + nthreads = resolve_nthreads(nthreads) + # --- basic checks / effective lmax, mmax ------------------------------- if alms.nstokes not in (1, 3): raise ValueError( @@ -2732,7 +2737,7 @@ def estimate_alm( mmax: int | None = None, maxiter: int | None = None, epsilon: float | None = None, - nthreads: int = 0, + nthreads: int | None = None, ) -> SphericalHarmonics: r""" Estimate spherical harmonic coefficients ($a_{\ell m}$) from a HEALPix map. @@ -2802,9 +2807,9 @@ def estimate_alm( * ``complex64`` output -> ``1e-6`` * ``complex128`` output -> ``1e-13`` - nthreads : int, optional - Number of threads passed to ducc. If zero (default), ducc chooses - the number of threads. + nthreads : int or None, optional + Number of threads passed to ducc. If None, resolved via + :func:`.resolve_nthreads`. Returns ------- @@ -2826,6 +2831,8 @@ def estimate_alm( If the HealpixMap object has unsupported ``nstokes``, or if requested ``lmax``/``mmax`` are inconsistent. """ + nthreads = resolve_nthreads(nthreads) + # --- validate maxiter --------------------------------------------------- if maxiter is not None: if not isinstance(maxiter, int): @@ -3028,7 +3035,7 @@ def rotate_alm( phi: float | None = None, mmax_out: int | None = None, inplace: bool = False, - nthreads: int = 0, + nthreads: int | None = None, ) -> SphericalHarmonics: """ Rotate spherical harmonic coefficients using ducc0. @@ -3055,8 +3062,9 @@ def rotate_alm( If True, modifies the input `alms` object in place. Note: In-place rotation is only possible if `mmax_out` is equal to `alms.mmax` (output size must match input size). - nthreads : int, optional, keyword-only - Number of threads to use for the rotation. Default is 0 (use all available). + nthreads : int or None, optional, keyword-only + Number of threads to use for the rotation. If None, resolved via + :func:`.resolve_nthreads`. Returns ------- @@ -3071,6 +3079,8 @@ def rotate_alm( - If `inplace=True` is requested but `mmax_out` differs from `alms.mmax`. """ + nthreads = resolve_nthreads(nthreads) + # 1. Retrieve Geometry from Input lmax_in = alms.lmax mmax_in = alms.mmax diff --git a/litebird_sim/pointings_in_obs.py b/litebird_sim/pointings_in_obs.py index edef0bc9..35b75e19 100644 --- a/litebird_sim/pointings_in_obs.py +++ b/litebird_sim/pointings_in_obs.py @@ -11,6 +11,7 @@ from .hwp import HWP from .observations import Observation from .scanning import RotQuaternion +from .utilities import resolve_nthreads def prepare_pointings( @@ -228,7 +229,7 @@ def _get_pointings_array( output_coordinate_system: CoordinateSystem, nside_centering: int | None = None, pointings_dtype=np.float64, - nthreads: int = 0, + nthreads: int | None = None, ) -> tuple[np.ndarray, np.ndarray | None]: """Compute the pointings (θ, φ) and HWP angle for a given detector. @@ -247,9 +248,9 @@ def _get_pointings_array( If provided, the pointings will be aligned to the center of the HEALPix pixel. pointings_dtype : np.dtype, optional Data type for computed pointings and angles. Default is `np.float64`. - nthreads : int, optional + nthreads : int or None, optional Number of threads to use for HEALPix operations when centering pointings. - Default is 0 (use all available threads). + If None, resolved via :func:`.resolve_nthreads`. Returns ------- @@ -259,6 +260,8 @@ def _get_pointings_array( N_cols is typically 2 ([θ, φ]) or 3 ([θ, φ, ψ]). - `hwp_angle` is either the provided array or the one computed by the callable. """ + nthreads = resolve_nthreads(nthreads) + if isinstance(pointings, np.ndarray): curr_pointings_det = pointings[detector_idx, :, :] computed_hwp_angle = None diff --git a/litebird_sim/scan_map.py b/litebird_sim/scan_map.py index a5379beb..97573cf5 100644 --- a/litebird_sim/scan_map.py +++ b/litebird_sim/scan_map.py @@ -121,7 +121,7 @@ def scan_map( input_names: str | None = None, interpolation: str | None = "", pointings_dtype=np.float64, - nthreads: int = 0, + nthreads: int | None = None, ): """ Scan a sky map and fill time-ordered data (TOD) based on detector observations. @@ -190,9 +190,9 @@ def scan_map( pointings_dtype : dtype, optional Data type for pointings generated on the fly. - nthreads : int, default=0 - Number of threads to use for convolution. If set to 0, all available CPU cores - will be used. + nthreads : int or None, default=None + Number of threads to use for convolution. If None, resolved via + :func:`.resolve_nthreads`. Raises ------ @@ -210,6 +210,8 @@ def scan_map( experiments. """ + nthreads = resolve_nthreads(nthreads) + n_detectors = tod.shape[0] if type(pointings) is np.ndarray: @@ -262,6 +264,7 @@ def scan_map( hwp_angle=hwp_angle, output_coordinate_system=coordinates, pointings_dtype=pointings_dtype, + nthreads=nthreads, ) # ---------------------------------------------------------- diff --git a/litebird_sim/simulations.py b/litebird_sim/simulations.py index 10ee60ff..a275b8da 100644 --- a/litebird_sim/simulations.py +++ b/litebird_sim/simulations.py @@ -33,7 +33,7 @@ add_convolved_sky_to_observations, ) from .beam_synthesis import generate_gauss_beam_alms -from .constants import NUMBA_NUM_THREADS_ENVVAR +from .constants import NUMBA_NUM_THREADS_ENVVAR # noqa: F401 (re-exported) from .coordinates import CoordinateSystem from .detectors import UUID, DetectorInfo, FreqChannelInfo, InstrumentInfo from .dipole import DipoleType, add_dipole_to_observations @@ -75,6 +75,7 @@ from .seeding import RNGHierarchy from .spacecraft import SpacecraftOrbit, spacecraft_pos_and_vel from .units import Units +from .utilities import resolve_numba_nthreads from .version import ( __author__ as litebird_sim_author, ) @@ -415,9 +416,10 @@ def __init__( else: self.imo = Imo() - if not numba_threads and NUMBA_NUM_THREADS_ENVVAR in os.environ: - numba_threads = int(os.environ[NUMBA_NUM_THREADS_ENVVAR]) - + # Resolved below, after _init_missing_params() has had a chance to + # fill this in from a parameter file: explicit argument takes + # precedence, then the parameter file, then resolve_numba_nthreads() + # (NUMBA_NUM_THREADS/OMP_NUM_THREADS/hardware default). self.numba_threads = numba_threads self.numba_threading_layer = numba_threading_layer @@ -442,8 +444,8 @@ def __init__( self._init_missing_params() - if self.numba_threads: - numba.set_num_threads(self.numba_threads) + self.numba_threads = resolve_numba_nthreads(self.numba_threads) + numba.set_num_threads(self.numba_threads) if self.numba_threading_layer: numba.config.THREADING_LAYER = self.numba_threading_layer # type: ignore[attr-defined] diff --git a/litebird_sim/utilities.py b/litebird_sim/utilities.py index 05d2cbdc..c636951f 100644 --- a/litebird_sim/utilities.py +++ b/litebird_sim/utilities.py @@ -1,16 +1,58 @@ import os -from .constants import NUM_THREADS_ENVVAR +from .constants import NUM_THREADS_ENVVAR, NUMBA_NUM_THREADS_ENVVAR + + +def _compute_nthreads() -> int: + if NUM_THREADS_ENVVAR in os.environ: + return int(os.environ[NUM_THREADS_ENVVAR]) + # No explicit thread count was requested. Default to a single thread + # rather than every hardware thread available to the process: MPI/OpenMP + # jobs are expected to set NUM_THREADS_ENVVAR (OMP_NUM_THREADS) per rank, + # so this fallback only matters for un-configured runs (e.g. a laptop), + # where using every core by default would otherwise silently starve + # other processes on the machine. + return 1 + + +def _compute_numba_nthreads() -> int: + if NUMBA_NUM_THREADS_ENVVAR in os.environ: + return int(os.environ[NUMBA_NUM_THREADS_ENVVAR]) + return _compute_nthreads() + + +# Resolved once at import time (each MPI process reads its own environment), +# instead of re-reading the environment on every call. Mirrors the +# lbs.MPI_COMM_WORLD pattern: a single, process-wide value computed once and +# reused everywhere `nthreads` isn't explicitly overridden. +NUM_THREADS = _compute_nthreads() +NUMBA_NUM_THREADS = _compute_numba_nthreads() def resolve_nthreads(nthreads: int | None) -> int: - """Resolve thread count from an explicit value or environment. + """Resolve the number of threads to use for ducc0/general parallel code. - If ``nthreads`` is not ``None``, return it unchanged. Otherwise read - :data:`NUM_THREADS_ENVVAR` and fall back to 0 when the variable is unset. + If ``nthreads`` is given explicitly, return it unchanged. Otherwise + return :data:`NUM_THREADS`, which was resolved once at import time from + :data:`.constants.NUM_THREADS_ENVVAR` (``OMP_NUM_THREADS``), falling back + to ``1`` if the environment variable is not set. """ if nthreads is not None: return nthreads + return NUM_THREADS - return int(os.environ.get(NUM_THREADS_ENVVAR, 0)) + +def resolve_numba_nthreads(nthreads: int | None) -> int: + """Resolve the number of threads to use for Numba. + + If ``nthreads`` is given explicitly, return it unchanged. Otherwise + return :data:`NUMBA_NUM_THREADS`, resolved once at import time from + :data:`.constants.NUMBA_NUM_THREADS_ENVVAR` (``NUMBA_NUM_THREADS``) if + set, falling back to :data:`NUM_THREADS` (``OMP_NUM_THREADS``) so Numba + and ducc0 agree by default while still allowing independent tuning. + """ + + if nthreads is not None: + return nthreads + return NUMBA_NUM_THREADS