Source code for mantidimaging.core.parallel.utility

# Copyright (C) 2024 ISIS Rutherford Appleton Laboratory UKRI
# SPDX - License - Identifier: GPL-3.0-or-later
from __future__ import annotations

import os
from logging import getLogger
from multiprocessing import shared_memory
from typing import TYPE_CHECKING
from import Callable

import numpy as np

from mantidimaging.core.utility.memory_usage import system_free_memory
from mantidimaging.core.utility.progress_reporting import Progress
from mantidimaging.core.utility.size_calculator import full_size_KB, full_size_bytes
from mantidimaging.core.parallel import manager as pm

    from functools import partial
    import numpy.typing as npt
    from multiprocessing.shared_memory import SharedMemory

LOG = getLogger(__name__)

[docs] def enough_memory(shape, dtype): return full_size_KB(shape=shape, dtype=dtype) < system_free_memory().kb()
[docs] def create_array(shape: tuple[int, ...], dtype: npt.DTypeLike = np.float32) -> SharedArray: """ Create an array in shared memory :param shape: Shape of the array :param dtype: Dtype of the array :return: The created SharedArray """ if not enough_memory(shape, dtype): raise RuntimeError( "The machine does not have enough physical memory available to allocate space for this data.") return _create_shared_array(shape, dtype)
def _create_shared_array(shape: tuple[int, ...], dtype: npt.DTypeLike = np.float32) -> SharedArray: size = full_size_bytes(shape, dtype)'Requested shared array with shape={shape}, size={size}, dtype={dtype}') name = pm.generate_mi_shared_mem_name() mem = shared_memory.SharedMemory(name=name, create=True, size=size) return _read_array_from_shared_memory(shape, dtype, mem, True) def _read_array_from_shared_memory(shape: tuple[int, ...], dtype: npt.DTypeLike, mem: SharedMemory, free_mem_on_delete: bool) -> SharedArray: array: np.ndarray = np.ndarray(shape, dtype=dtype, buffer=mem.buf) return SharedArray(array, mem, free_mem_on_del=free_mem_on_delete)
[docs] def copy_into_shared_memory(array: np.ndarray) -> SharedArray: shared_array = create_array(array.shape, array.dtype) shared_array.array[:] = array[:] return shared_array
[docs] def calculate_chunksize(cores): """ TODO possible proper calculation of chunksize, although best performance has been with 1 From performance tests, the chunksize doesn't seem to make much of a difference, but having larger chunks usually led to slower performance: Shape: (50,512,512) 1 chunk 3.06s 2 chunks 3.05s 3 chunks 3.07s 4 chunks 3.06s 5 chunks 3.16s 6 chunks 3.06s 7 chunks 3.058s 8 chunks 3.25s 9 chunks 3.45s """ return 1
[docs] def multiprocessing_necessary(shape: int, is_shared_data: bool) -> bool: # This environment variable will be present when running PYDEVD from PyCharm # and that has the bug that multiprocessing Pools can never finish `.join()` ing # thus never actually finish their processing. if 'PYDEVD_LOAD_VALUES_ASYNC' in os.environ and 'PYTEST_CURRENT_TEST' not in os.environ:"Debugging environment variable 'PYDEVD_LOAD_VALUES_ASYNC' found. Running synchronously on 1 core") return False if not is_shared_data:"Not all of the data uses shared memory") return False elif shape <= 10:"Shape under 10") return False"Multiprocessing required") return True
[docs] def execute_impl(img_num: int, partial_func: partial, is_shared_data: bool, progress: Progress, msg: str): task_name = f"{msg}" progress = Progress.ensure_instance(progress, num_steps=img_num, task_name=task_name) indices_list = range(img_num) if multiprocessing_necessary(img_num, is_shared_data) and pm.pool:"Running async on {pm.cores} cores") # Using _ in the for _ enumerate is slightly faster, because the tuple from enumerate isn't unpacked, # and thus some time is saved # Using imap here seems to be the best choice: # - imap_unordered gives the images back in random order # - map and map_async do not improve speed performance for _ in pm.pool.imap(partial_func, indices_list, chunksize=calculate_chunksize(pm.cores)): progress.update(1, msg) else:"Running synchronously on 1 core") for ind in indices_list: partial_func(ind) progress.update(1, msg) progress.mark_complete()
[docs] def run_compute_func_impl(worker_func: Callable[[int], None], num_operations: int, is_shared_data: bool, progress=None, msg: str = ""): task_name = f"{msg}" progress = Progress.ensure_instance(progress, num_steps=num_operations, task_name=task_name) indices_list = range(num_operations) if multiprocessing_necessary(num_operations, is_shared_data) and pm.pool:"Running async on {pm.cores} cores") for _ in pm.pool.imap(worker_func, indices_list, chunksize=calculate_chunksize(pm.cores)): progress.update(1, msg) else:"Running synchronously on 1 core") for ind in indices_list: worker_func(ind) progress.update(1, msg) progress.mark_complete()
[docs] class SharedArray: def __init__(self, array: np.ndarray, shared_memory: SharedMemory | None, free_mem_on_del: bool = True): self.array = array self._shared_memory = shared_memory self._free_mem_on_del = free_mem_on_del def __del__(self): if self.has_shared_memory: self._shared_memory.close() if self._free_mem_on_del: try: self._shared_memory.unlink() except FileNotFoundError: # Do nothing, memory has already been freed pass @property def has_shared_memory(self) -> bool: return self._shared_memory is not None @property def array_proxy(self) -> SharedArrayProxy: mem_name = if self._shared_memory else None return SharedArrayProxy(mem_name=mem_name, shape=self.array.shape, dtype=self.array.dtype)
[docs] class SharedArrayProxy: def __init__(self, mem_name: str | None, shape: tuple[int, ...], dtype: npt.DTypeLike): self._mem_name = mem_name self._shape = shape self._dtype = dtype self._shared_array: SharedArray | None = None @property def array(self) -> np.ndarray: if self._shared_array is None: mem = shared_memory.SharedMemory(name=self._mem_name) self._shared_array = _read_array_from_shared_memory(self._shape, self._dtype, mem, False) return self._shared_array.array