1"parallel processing utilities, chiefly `run_maybe_parallel`"
  2
  3from __future__ import annotations
  4
  5import multiprocessing
  6import functools
  7from typing import (
  8    Any,
  9    Callable,
 10    Iterable,
 11    Literal,
 12    Optional,
 13    Tuple,
 14    TypeVar,
 15    Dict,
 16    List,
 17    Union,
 18    Protocol,
 19)
 20
 21# for no tqdm fallback
 22from muutils.spinner import SpinnerContext
 23from muutils.validate_type import get_fn_allowed_kwargs
 24
 25
 26InputType = TypeVar("InputType")
 27OutputType = TypeVar("OutputType")
 28# typevars for our iterable and map
 29
 30
 31class ProgressBarFunction(Protocol):
 32    "a protocol for a progress bar function"
 33
 34    def __call__(self, iterable: Iterable, **kwargs: Any) -> Iterable: ...
 35
 36
 37ProgressBarOption = Literal["tqdm", "spinner", "none", None]
 38# type for the progress bar option
 39
 40
 41DEFAULT_PBAR_FN: ProgressBarOption
 42# default progress bar function
 43
 44try:
 45    # use tqdm if it's available
 46    import tqdm  # type: ignore[import-untyped]
 47
 48    DEFAULT_PBAR_FN = "tqdm"
 49
 50except ImportError:
 51    # use progress bar as fallback
 52    DEFAULT_PBAR_FN = "spinner"
 53
 54
 55def spinner_fn_wrap(x: Iterable, **kwargs) -> List:
 56    "spinner wrapper"
 57    spinnercontext_allowed_kwargs: set[str] = get_fn_allowed_kwargs(
 58        SpinnerContext.__init__
 59    )
 60    mapped_kwargs: dict = {
 61        k: v for k, v in kwargs.items() if k in spinnercontext_allowed_kwargs
 62    }
 63    if "desc" in kwargs and "message" not in mapped_kwargs:
 64        mapped_kwargs["message"] = kwargs["desc"]
 65
 66    if "message" not in mapped_kwargs and "total" in kwargs:
 67        mapped_kwargs["message"] = f"Processing {kwargs['total']} items"
 68
 69    with SpinnerContext(**mapped_kwargs):
 70        output = list(x)
 71
 72    return output
 73
 74
 75def map_kwargs_for_tqdm(kwargs: dict) -> dict:
 76    "map kwargs for tqdm, cant wrap because the pbar dissapears?"
 77    tqdm_allowed_kwargs: set[str] = get_fn_allowed_kwargs(tqdm.tqdm.__init__)
 78    mapped_kwargs: dict = {k: v for k, v in kwargs.items() if k in tqdm_allowed_kwargs}
 79
 80    if "desc" not in kwargs:
 81        if "message" in kwargs:
 82            mapped_kwargs["desc"] = kwargs["message"]
 83
 84        elif "total" in kwargs:
 85            mapped_kwargs["desc"] = f"Processing {kwargs.get('total')} items"
 86    return mapped_kwargs
 87
 88
 89def no_progress_fn_wrap(x: Iterable, **kwargs) -> Iterable:
 90    "fallback to no progress bar"
 91    return x
 92
 93
 94def set_up_progress_bar_fn(
 95    pbar: Union[ProgressBarFunction, ProgressBarOption],
 96    pbar_kwargs: Optional[Dict[str, Any]] = None,
 97    **extra_kwargs,
 98) -> Tuple[ProgressBarFunction, dict]:
 99    """set up the progress bar function and its kwargs
100
101    # Parameters:
102     - `pbar : Union[ProgressBarFunction, ProgressBarOption]`
103       progress bar function or option. if a function, we return as-is. if a string, we figure out which progress bar to use
104     - `pbar_kwargs : Optional[Dict[str, Any]]`
105       kwargs passed to the progress bar function (default to `None`)
106       (defaults to `None`)
107
108    # Returns:
109     - `Tuple[ProgressBarFunction, dict]`
110         a tuple of the progress bar function and its kwargs
111
112    # Raises:
113     - `ValueError` : if `pbar` is not one of the valid options
114    """
115    pbar_fn: ProgressBarFunction
116
117    if pbar_kwargs is None:
118        pbar_kwargs = dict()
119
120    pbar_kwargs = {**extra_kwargs, **pbar_kwargs}
121
122    # dont use a progress bar if `pbar` is None or "none", or if `disable` is set to True in `pbar_kwargs`
123    if (pbar is None) or (pbar == "none") or pbar_kwargs.get("disable", False):
124        pbar_fn = no_progress_fn_wrap  # type: ignore[assignment]
125
126    # if `pbar` is a different string, figure out which progress bar to use
127    elif isinstance(pbar, str):
128        if pbar == "tqdm":
129            pbar_fn = tqdm.tqdm
130            pbar_kwargs = map_kwargs_for_tqdm(pbar_kwargs)
131        elif pbar == "spinner":
132            pbar_fn = functools.partial(spinner_fn_wrap, **pbar_kwargs)
133            pbar_kwargs = dict()
134        else:
135            raise ValueError(
136                f"`pbar` must be either 'tqdm' or 'spinner' if `str`, or a valid callable, got {type(pbar) = } {pbar = }"
137            )
138    else:
139        # the default value is a callable which will resolve to tqdm if available or spinner as a fallback. we pass kwargs to this
140        pbar_fn = pbar
141
142    return pbar_fn, pbar_kwargs
143
144
145# TODO: if `parallel` is a negative int, use `multiprocessing.cpu_count() + parallel` to determine the number of processes
146def run_maybe_parallel(
147    func: Callable[[InputType], OutputType],
148    iterable: Iterable[InputType],
149    parallel: Union[bool, int],
150    pbar_kwargs: Optional[Dict[str, Any]] = None,
151    chunksize: Optional[int] = None,
152    keep_ordered: bool = True,
153    use_multiprocess: bool = False,
154    pbar: Union[ProgressBarFunction, ProgressBarOption] = DEFAULT_PBAR_FN,
155) -> List[OutputType]:
156    """a function to make it easier to sometimes parallelize an operation
157
158    - if `parallel` is `False`, then the function will run in serial, running `map(func, iterable)`
159    - if `parallel` is `True`, then the function will run in parallel, running in parallel with the maximum number of processes
160    - if `parallel` is an `int`, it must be greater than 1, and the function will run in parallel with the number of processes specified by `parallel`
161
162    the maximum number of processes is given by the `min(len(iterable), multiprocessing.cpu_count())`
163
164    # Parameters:
165     - `func : Callable[[InputType], OutputType]`
166       function passed to either `map` or `Pool.imap`
167     - `iterable : Iterable[InputType]`
168       iterable passed to either `map` or `Pool.imap`
169     - `parallel : bool | int`
170       whether to run in parallel, and how many processes to use
171     - `pbar_kwargs : Dict[str, Any]`
172       kwargs passed to the progress bar function
173
174    # Returns:
175     - `List[OutputType]`
176       a list of the output of `func` for each element in `iterable`
177
178    # Raises:
179     - `ValueError` : if `parallel` is not a boolean or an integer greater than 1
180     - `ValueError` : if `use_multiprocess=True` and `parallel=False`
181     - `ImportError` : if `use_multiprocess=True` and `multiprocess` is not available
182    """
183
184    # number of inputs in iterable
185    n_inputs: int = len(iterable)  # type: ignore[arg-type]
186    if n_inputs == 0:
187        # Return immediately if there is no input
188        return list()
189
190    # which progress bar to use
191    pbar_fn: ProgressBarFunction
192    pbar_kwargs_processed: dict
193    pbar_fn, pbar_kwargs_processed = set_up_progress_bar_fn(
194        pbar=pbar,
195        pbar_kwargs=pbar_kwargs,
196        # extra kwargs
197        total=n_inputs,
198    )
199
200    # number of processes
201    num_processes: int
202    if isinstance(parallel, bool):
203        num_processes = multiprocessing.cpu_count() if parallel else 1
204    elif isinstance(parallel, int):
205        if parallel < 2:
206            raise ValueError(
207                f"`parallel` must be a boolean, or be an integer greater than 1, got {type(parallel) = } {parallel = }"
208            )
209        num_processes = parallel
210    else:
211        raise ValueError(
212            f"The 'parallel' parameter must be a boolean or an integer, got {type(parallel) = } {parallel = }"
213        )
214
215    # make sure we don't have more processes than iterable, and don't bother with parallel if there's only one process
216    num_processes = min(num_processes, n_inputs)
217    mp = multiprocessing
218    if num_processes == 1:
219        parallel = False
220
221    if use_multiprocess:
222        if not parallel:
223            raise ValueError("`use_multiprocess=True` requires `parallel=True`")
224
225        try:
226            import multiprocess  # type: ignore[import-untyped]
227        except ImportError as e:
228            raise ImportError(
229                "`use_multiprocess=True` requires the `multiprocess` package -- this is mostly useful when you need to pickle a lambda. install muutils with `pip install muutils[multiprocess]` or just do `pip install multiprocess`"
230            ) from e
231
232        mp = multiprocess
233
234    # set up the map function -- maybe its parallel, maybe it's just `map`
235    do_map: Callable[
236        [Callable[[InputType], OutputType], Iterable[InputType]],
237        Iterable[OutputType],
238    ]
239    if parallel:
240        # use `mp.Pool` since we might want to use `multiprocess` instead of `multiprocessing`
241        pool = mp.Pool(num_processes)
242
243        # use `imap` if we want to keep the order, otherwise use `imap_unordered`
244        if keep_ordered:
245            do_map = pool.imap
246        else:
247            do_map = pool.imap_unordered
248
249        # figure out a smart chunksize if one is not given
250        chunksize_int: int
251        if chunksize is None:
252            chunksize_int = max(1, n_inputs // num_processes)
253        else:
254            chunksize_int = chunksize
255
256        # set the chunksize
257        do_map = functools.partial(do_map, chunksize=chunksize_int)  # type: ignore
258
259    else:
260        do_map = map
261
262    # run the map function with a progress bar
263    output: List[OutputType] = list(
264        pbar_fn(
265            do_map(
266                func,
267                iterable,
268            ),
269            **pbar_kwargs_processed,
270        )
271    )
272
273    # close the pool if we used one
274    if parallel:
275        pool.close()
276        pool.join()
277
278    # return the output as a list
279    return output