Implement MPI shared-memory allocation for pointing quaternions to mitigate out-of-memory issues

.github/workflows/tests.yml

@@ -116,6 +116,7 @@ jobs:
               echo "Running MPI test ($MPI) with $proc processes"
               PYTHONPATH=. mpiexec -n $proc uv run python -m pytest ./test/test_mpi.py
               PYTHONPATH=. mpiexec -n $proc uv run python -m pytest ./test/test_detector_blocks.py
+              PYTHONPATH=. mpiexec -n $proc uv run python -m pytest ./test/test_shared_memory_pointings.py
             done
             echo "Running MPI test ($MPI) with 4 processes"
             PYTHONPATH=. mpiexec --map-by :OVERSUBSCRIBE -n 4 uv run python -m pytest ./test/test_mpi_n4.py
@@ -131,6 +132,7 @@ jobs:
               echo "Running MPI test ($MPI) with $proc processes"
               PYTHONPATH=. mpiexec -n $proc uv run python -m pytest ./test/test_mpi.py
               PYTHONPATH=. mpiexec -n $proc uv run python -m pytest ./test/test_detector_blocks.py
+              PYTHONPATH=. mpiexec -n $proc uv run python -m pytest ./test/test_shared_memory_pointings.py
             done
             echo "Running MPI test ($MPI) with 4 processes"
             PYTHONPATH=. mpiexec --map-by :OVERSUBSCRIBE -n 4 uv run python -m pytest ./test/test_mpi_n4.py

CHANGELOG.md

@@ -1,5 +1,7 @@
 # HEAD
 
+-   Introduced MPI shared-memory allocator for pointing quaternions (`SharedMemoryManager`) to drastically reduce memory usage and prevent OOM issues on multi-core HPC nodes (like [#489](https://github.com/litebird/litebird_sim/issues/489)). Added `Simulation.set_scanning_strategy_shmem()` and `Simulation.prepare_pointings_shmem()` to leverage this functionality.
+
 -   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).

litebird_sim/observations.py

@@ -15,9 +15,11 @@
 from .input_sky import SkyInput
 from .mpi import _SerialMpiCommunicator
 from .pointings import DEFAULT_INTERNAL_BUFFER_SIZE_FOR_POINTINGS_MB, PointingProvider
-from .scanning import RotQuaternion
+from .scanning import RotQuaternion, SharedRotQuaternion
 from .units import Units
 from .utilities import resolve_nthreads
+from .quaternions import normalize_quaternions
+from .shared_memory import SharedMemoryManager
 
 
 @dataclass
@@ -983,6 +985,87 @@ def prepare_pointings(
             self.hwp = hwp
             self.has_hwp = True
 
+    def prepare_pointings_shmem(
+        self,
+        instrument: InstrumentInfo,
+        spin2ecliptic_quats: RotQuaternion,
+        hwp: HWP | None = None,
+        maximum_internal_buffer_mem_mb: float = DEFAULT_INTERNAL_BUFFER_SIZE_FOR_POINTINGS_MB,
+    ) -> None:
+        """Prepare quaternion-based pointing and HWP information using MPI shared memory.
+
+        This functions similarly to `prepare_pointings` but uses the provided
+        `SharedMemoryManager` to allocate the underlying pointing quaternion arrays
+        in node-shared memory, optimizing the global memory usage.
+        """
+
+        assert (maximum_internal_buffer_mem_mb > 0) or (
+            maximum_internal_buffer_mem_mb == -1
+        ), (
+            "Invalid value for maximum_internal_buffer_mem_mb ({val}), it must either be -1 or a positive number".format(
+                val=maximum_internal_buffer_mem_mb
+            )
+        )
+
+        if not hasattr(self.comm_time_block, "Split_type"):
+            raise RuntimeError(
+                "MPI is required for shared memory pointings, but the current time block communicator is not a valid MPI communicator."
+            )
+
+        self.shared_memory_manager = SharedMemoryManager(base_comm=self.comm_time_block)
+
+        n_quats = spin2ecliptic_quats.quats.shape[0]
+        dtype = spin2ecliptic_quats.quats.dtype
+
+        # Allocate flat 1D array on node shared memory
+        flat_array, _ = self.shared_memory_manager.alloc_shared_node(
+            size=n_quats * 4,
+            dtype=dtype,
+        )
+
+        # Create a 2D view
+        shared_quats_view = flat_array.reshape(n_quats, 4)
+
+        # Only the node root computes the multiplication and fills the array
+        if self.shared_memory_manager.node_rank == self.shared_memory_manager.node_root:
+            from .quaternions import multiply_quaternions_list_x_one
+
+            multiply_quaternions_list_x_one(
+                spin2ecliptic_quats.quats,
+                instrument.bore2spin_quat.quats[0],
+                shared_quats_view,
+            )
+            normalize_quaternions(shared_quats_view)
+
+        # Synchronize all ranks on this node
+        self.shared_memory_manager.fence_comm_all(self.shared_memory_manager.node_comm)
+
+        # Wrap into SharedRotQuaternion (which does not normalize)
+        shared_rot_quat = SharedRotQuaternion(
+            quats=shared_quats_view,
+            start_time=spin2ecliptic_quats.start_time,
+            sampling_rate_hz=spin2ecliptic_quats.sampling_rate_hz,
+        )
+
+        pointing_provider = PointingProvider(
+            bore2ecliptic_quats=shared_rot_quat,
+            hwp=hwp,
+            maximum_internal_buffer_mem_mb=maximum_internal_buffer_mem_mb,
+        )
+
+        self.pointing_provider = pointing_provider
+
+        # If the hwp object is passed and is not initialised in the observations, it gets applied to all detectors
+        if hwp is None:
+            assert self.no_mueller_hwp() or self.no_jones_hwp(), (
+                "Some detectors have been initialized with a mueller_hwp or jones_hwp,"
+                "but no HWP object has been passed to prepare_shared_pointings."
+            )
+            self.has_hwp = False
+        else:
+            self.hwp = hwp
+            self.has_hwp = True
+
     def get_pointings(
         self,
         detector_idx: int | list[int] | str = "all",
@@ -1270,7 +1353,7 @@ def _set_mpi_subcommunicators(self):
         self.comm_det_block = _SerialMpiCommunicator()
         self.comm_time_block = _SerialMpiCommunicator()
 
-        if self.comm and self.comm.size > 1:
+        if hasattr(self.comm, "Split"):
             det_color = self.comm.rank // self.n_blocks_time
             time_color = self.comm.rank % self.n_blocks_time
             self.comm_det_block = self.comm.Split(det_color)

litebird_sim/pointings_in_obs.py

@@ -65,6 +65,29 @@ def prepare_pointings(
         )
 
 
+def prepare_pointings_shmem(
+    observations: Observation | list[Observation],
+    instrument: InstrumentInfo,
+    spin2ecliptic_quats: RotQuaternion,
+    hwp: HWP | None = None,
+) -> None:
+    """Initialize pointing and HWP information using MPI shared memory.
+
+    This acts just like `prepare_pointings`, but relies on `SharedMemoryManager`
+    to allocate the underlying pointing quaternion arrays into MPI node-shared
+    memory, thereby preventing redundant memory allocation on multi-core nodes.
+    """
+    if isinstance(observations, Observation):
+        obs_list = [observations]
+    else:
+        obs_list = observations
+
+    for cur_obs in obs_list:
+        cur_obs.prepare_pointings_shmem(
+            instrument=instrument, spin2ecliptic_quats=spin2ecliptic_quats, hwp=hwp
+        )
+
+
 def precompute_pointings(
     observations: Observation | list[Observation],
     pointings_dtype=np.float64,

litebird_sim/scanning.py

@@ -636,6 +636,50 @@ def is_close_to(self, other: "RotQuaternion") -> bool:
         return np.allclose(self.quats, other.quats)
 
 
+class SharedRotQuaternion(RotQuaternion):
+    """A version of RotQuaternion that wraps an existing shared-memory NumPy array.
+
+    This class is used to avoid copying or modifying (normalizing) the underlying
+    array, as that is expected to be handled explicitly by the
+    node root process of the shared memory communicator.
+    """
+
+    def __init__(
+        self,
+        quats: npt.NDArray,
+        start_time: float | astropy.time.Time | None = None,
+        sampling_rate_hz: float | None = None,
+    ):
+        """
+        Create a new instance of a time-dependent quaternion
+
+        If both `start_time` and `sampling_freq_hz` are ``None``, the quaternion
+        is assumed to be constant in time.
+
+        :param quats: Either a 4-element NumPy array, or another instance
+            of :class:`TimeDependentQuaternion`
+        :param start_time: the start time, either a floating point number
+            or an ``astropy.time.Time`` object
+        :param sampling_rate_hz: the sampling frequency
+        """
+        # Directly use the provided array without reshaping or copying
+        self.quats = quats
+
+        if self.quats.shape[0] > 1:
+            assert start_time is not None, (
+                "You must specify start_time if the quaternion is not constant"
+            )
+            assert sampling_rate_hz is not None, (
+                "You must specify sampling_rate_hz if the quaternion is not constant"
+            )
+
+        self.start_time = start_time
+        self.sampling_rate_hz = sampling_rate_hz
+
+        # We DO NOT normalize quaternions here.
+        # The node root is expected to normalize them after writing to shared memory.
+
+
 # This is an Abstract Base Class (ABC)
 class ScanningStrategy(ABC):
     """A class that simulate a scanning strategy