map_parallel

Complexity

Total Complexity

26

Size/Duplication

Total Lines	162
Duplicated Lines	28.4 %

Importance

Changes

0

7 Functions

Rating	Name	Duplication	Size	Complexity
A	_map_parallel_dask()	0	17	3
A	_map_parallel_multithreading()	0	14	3
A	_map_parallel_multiprocessing()	0	14	3
A	_starmap_parallel_mpi_simple()	0	25	4
A	_map_parallel_mpi()	0	9	3
A	map_parallel()	29	29	5
A	starmap_parallel()	17	17	5

from __future__ import annotations

__version__ = '0.1.1'

from itertools import starmap
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from typing import Optional, Callable, Dict
    from collections.abc import Iterable


def _map_parallel_multiprocessing(
    f: Callable,
    *args: Iterable,
    processes: Optional[int] = None,
    return_results: bool = True,
) -> list:
    from concurrent.futures import ProcessPoolExecutor

    with ProcessPoolExecutor(max_workers=processes) as process_pool_executor:
        res = process_pool_executor.map(f, *args)
        if return_results:
            return list(res)
        else:
            return []


def _map_parallel_multithreading(
    f: Callable,
    *args: Iterable,
    processes: Optional[int] = None,
    return_results: bool = True,
) -> list:
    from concurrent.futures import ThreadPoolExecutor

    with ThreadPoolExecutor(max_workers=processes) as thread_pool_executor:
        res = thread_pool_executor.map(f, *args)
        if return_results:
            return list(res)
        else:
            return []


def _map_parallel_dask(
    f: Callable,
    *args: Iterable,
    processes: Optional[int] = None,
    return_results: bool = True,
) -> list:
    from dask.distributed import Client
    from dask.distributed import LocalCluster

    cluster = LocalCluster(n_workers=processes, dashboard_address=None)
    client = Client(cluster)
    if return_results:
        return [future.result() for future in client.map(f, *args)]
    else:
        for future in client.map(f, *args):
            future.result()
        return []


def _map_parallel_mpi(f: Callable, *args: Iterable, return_results: bool = True, **kwargs) -> list:
    from mpi4py.futures import MPIPoolExecutor

    with MPIPoolExecutor() as mpi_pool_executor:
        res = mpi_pool_executor.map(f, *args)
        if return_results:
            return list(res)
        else:
            return []


def _starmap_parallel_mpi_simple(
    f: Callable,
    args: Iterable[Iterable],
    return_results: bool = True,
    **kwargs,
) -> list:
    from mpi4py import MPI

    comm = MPI.COMM_WORLD
    size = comm.Get_size()
    rank = comm.Get_rank()

    args_list = list(args)
    if args_list:
        n = len(args_list)
        start = (rank * n) // size
        end = ((rank + 1) * n) // size
        local_args = args_list[start:end]
        res = list(starmap(f, local_args))

        if return_results:
            res = comm.gather(res, root=0)
            if rank == 0:
                return sum(res, [])
    return []


_map_parallel_func: Dict[str, Callable] = {

The variable Callable does not seem to be defined in case TYPE_CHECKING on line 8 is False. Are you sure this can never be the case?

    'multiprocessing': _map_parallel_multiprocessing,
    'multithreading': _map_parallel_multithreading,
    'dask': _map_parallel_dask,
    'mpi': _map_parallel_mpi,
}

_starmap_parallel_func: Dict[str, Callable] = {
    'mpi_simple': _starmap_parallel_mpi_simple,
}


def map_parallel(

This code seems to be duplicated in your project.

    f: Callable,
    *args: Iterable,
    processes: Optional[int] = None,
    mode: str = 'multiprocessing',
    return_results: bool = True,
) -> list:
    '''equiv to `map(f, *args)` but in parallel

    :param str mode: backend for parallelization

        - multiprocessing: using multiprocessing from standard library
        - multithreading: using multithreading from standard library
        - dask: using dask.distributed
        - mpi: using mpi4py.futures. May not work depending on your MPI vendor
        - mpi_simple: using mpi4py with simple scheduling that divides works into equal chunks
        - serial: using map
    :param int processes: no. of parallel processes

        (in the case of mpi, it is determined by mpiexec/mpirun args)

    :param bool return_results: (Only affects mode == 'mpi_simple') if True, return results in rank 0.
    '''
    if processes is None or processes > 1:
        if mode in _map_parallel_func:
            return _map_parallel_func[mode](f, *args, processes=processes, return_results=return_results)
        elif mode in _starmap_parallel_func:
            return _starmap_parallel_func[mode](f, zip(*args), processes=processes, return_results=return_results)
    return list(map(f, *args))


def starmap_parallel(

This code seems to be duplicated in your project.

    f: Callable,
    args: Iterable[Iterable],
    processes: Optional[int] = None,
    mode: str = 'multiprocessing',
    return_results: bool = True,
) -> list:
    '''equiv to `starmap(f, args)` but in parallel

    See docstring from :func:`~map_parallel.map_parallel`
    '''
    if processes is None or processes > 1:
        if mode in _map_parallel_func:
            return _map_parallel_func[mode](f, *zip(*args), processes=processes, return_results=return_results)
        elif mode in _starmap_parallel_func:
            return _starmap_parallel_func[mode](f, args, processes=processes, return_results=return_results)
    return list(starmap(f, args))