diff --git a/.github/workflows/build-docs.yml b/.github/workflows/build-docs.yml index 372ef276..9265a06b 100644 --- a/.github/workflows/build-docs.yml +++ b/.github/workflows/build-docs.yml @@ -32,7 +32,7 @@ jobs: - name: Install Python dependencies run: | python -m pip install --upgrade pip - python -m pip install pathos pygraphviz neuron cloud-volume k3d scikit-image open3d + python -m pip install joblib pathos pygraphviz neuron cloud-volume k3d scikit-image open3d python -m pip install -e .[test-notebook,all,docs,flybrains,cloud-volume] - name: Download test data run: | diff --git a/.github/workflows/test-tutorials.yml b/.github/workflows/test-tutorials.yml index 503363c6..a4fe2c7f 100644 --- a/.github/workflows/test-tutorials.yml +++ b/.github/workflows/test-tutorials.yml @@ -29,7 +29,7 @@ jobs: - name: Install dependencies run: | python -m pip install --upgrade pip - python -m pip install pathos pygraphviz neuron cloud-volume k3d scikit-image open3d + python -m pip install joblib pathos pygraphviz neuron cloud-volume k3d scikit-image open3d python -m pip install -e .[test-notebook,all] - name: Download test data run: | diff --git a/docs/examples/6_misc/tutorial_misc_00_multiprocess.py b/docs/examples/6_misc/tutorial_misc_00_multiprocess.py index 1b552b93..5cc4f211 100644 --- a/docs/examples/6_misc/tutorial_misc_00_multiprocess.py +++ b/docs/examples/6_misc/tutorial_misc_00_multiprocess.py @@ -5,16 +5,24 @@ This notebook will show you how to use parallel processing with `navis`. By default, most {{ navis }} functions use only a single thread/process (although some third-party functions -used under the hood might). Distributing expensive computations across multiple cores can speed things up considerable. +used under the hood might use more). Distributing expensive computations across multiple cores can speed things +up considerable. Many {{ navis }} functions natively support parallel processing. This notebook will illustrate various ways -to use parallelism. Before we get start: {{ navis }} uses `pathos` for multiprocessing - if you installed -{{ navis }} with `pip install navis[all]` you should be all set. If not, you can install `pathos` separately: +to use parallelism. Before we get started: by default, {{ navis }} uses `joblib` as backend for multiprocessing. +However, you can also use `pathos` as an alternative. If you installed {{ navis }} with `pip install navis[all]` +you should be all set to use either. If not, you can install the packages separately: ```shell +pip install joblib tqdm-joblib -U pip install pathos -U ``` +!!! tip + Parallel processing incurs overhead: we have to spawn additional processes and move data between the main + & the worker processes. If you have fast single-core performance and/or small tasks, that overhead might outweigh + the benefits of parallelism. See also additional notes at the bottom of this tutorial. + ## Running {{ navis }} functions in parallel Since version `0.6.0` many {{ navis }} functions accept a `parallel=True` and an (optional) `n_cores` parameter: @@ -36,7 +44,7 @@ def time_func(func, *args, **kwargs): # %% # !!! important # This documentation is built on Github Actions where the number of cores can be as low as 2. The speedup on -# your machine should be more pronounced than what you see below. That said: parallel processing has some +# your machine should be more pronounced than the times you see below. That said: parallel processing has some # overhead and for small tasks the overhead can be larger than the speed-up. # %% @@ -87,6 +95,17 @@ def time_func(func, *args, **kwargs): # the number of available CPU cores. However, doing so will likely over-subscribe your CPU and end up # slowing things down. +# %% +# Additional parameters for controlling parallelism: +# - `backend`: either "auto" (default), "joblib", or "pathos". This determines which parallel processing +# backend to use. "auto" will pick "joblib" if available, otherwise "pathos". Note: `joblib` is compatible +# with Dask to run on clusters. See the joblib documentation for details on how to set that up and then +# use `backend="joblib:dask"` in {{ navis }}. +# - `chunksize`: either "auto" (default) or an integer. This determines the number of neurons +# that will be processed per worker in each batch. "auto" will pick a chunksize that tries to balance +# load across workers while minimizing overhead. You can also set a fixed chunksize +# (e.g. `chunksize=10`). + # %% # ## Parallelizing generic functions # @@ -114,5 +133,10 @@ def my_func(x): nl.apply, my_func, parallel=True ) +""" +## A note on free-threading - +With version 3.13, Python introduced a free-threading build where the Global Interpreter Lock (GIL) is removed. +In theory, this would allow true multi-threading in Python and make parallel processing with threads much more efficient. +However, as of late 2025, many of the libraries used by {{ navis }} (e.g. `igraph`) do not yet support free-threading. +""" \ No newline at end of file diff --git a/docs/installation.md b/docs/installation.md index 28f97275..a72d3db0 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -123,10 +123,22 @@ directly, are listed below: --- + #### `joblib`: [joblib](https://joblib.readthedocs.io) + + Joblib is a multiprocessing library. {{ navis }} uses it to parallelize functions + across lists of neurons. Please note that you will also need to install `tqdm-joblib` + for progress bars to work with joblib. + + ``` shell + pip install joblib tqdm-joblib + ``` + + --- + #### `pathos`: [pathos](https://github.com/uqfoundation/pathos) - Pathos is a multiprocessing library. {{ navis }} uses it to parallelize functions - across lists of neurons. + Pathos is a multiprocessing library. {{ navis }} uses it as an alternative backend + to parallelize functions across lists of neurons. ``` shell pip install pathos diff --git a/navis/core/core_utils.py b/navis/core/core_utils.py index 3a7ec803..9caafa5a 100644 --- a/navis/core/core_utils.py +++ b/navis/core/core_utils.py @@ -11,10 +11,11 @@ # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the # GNU General Public License for more details. -import functools -import numbers import os import pint +import numbers +import warnings +import functools import pandas as pd import numpy as np @@ -26,17 +27,32 @@ from .. import config, graph, utils, core +try: + from joblib import Parallel, delayed +except ModuleNotFoundError: + Parallel = None + +if Parallel is not None: + try: + from tqdm_joblib import tqdm_joblib + except ModuleNotFoundError: + raise ModuleNotFoundError( + "navis relies on `tqdm-joblib` for joblib progress bars! " + "Please make sure tqdm-joblib is installed:\n" + " pip3 install tqdm-joblib -U" + ) try: - #from pathos.multiprocessing import ProcessingPool + # from pathos.multiprocessing import ProcessingPool # pathos' ProcessingPool apparently ignores chunksize # (see https://stackoverflow.com/questions/55611806/how-to-set-chunk-size-when-using-pathos-processingpools-map) import pathos + ProcessingPool = pathos.pools._ProcessPool except ModuleNotFoundError: ProcessingPool = None -__all__ = ['make_dotprops', 'to_neuron_space'] +__all__ = ["make_dotprops", "to_neuron_space"] # Set up logging logger = config.get_logger(__name__) @@ -44,6 +60,7 @@ def temp_property(func): """Check if neuron is stale. Clear cached temporary attributes if it is.""" + @functools.wraps(func) def wrapper(*args, **kwargs): self = args[0] @@ -52,11 +69,13 @@ def wrapper(*args, **kwargs): if self.is_stale: self._clear_temp_attr() return func(*args, **kwargs) + return wrapper def add_units(compact=True, power=1): """Add neuron units (if present) to output of function.""" + def outer(func): @functools.wraps(func) def wrapper(*args, **kwargs): @@ -317,10 +336,11 @@ def make_dotprops( return make_using(points=x, alpha=alpha, vect=vect, **properties) -def to_neuron_space(units: Union[int, float, pint.Quantity, pint.Unit], - neuron: core.BaseNeuron, - on_error: Union[Literal['ignore'], - Literal['raise']] = 'raise'): +def to_neuron_space( + units: Union[int, float, pint.Quantity, pint.Unit], + neuron: core.BaseNeuron, + on_error: Union[Literal["ignore"], Literal["raise"]] = "raise", +): """Convert units to match neuron space. Note that trying to convert units for non-isometric neurons will fail. @@ -365,9 +385,8 @@ def to_neuron_space(units: Union[int, float, pint.Quantity, pint.Unit], [0.125, 0.125, 0.125] """ - utils.eval_param(on_error, name='on_error', - allowed_values=('ignore', 'raise')) - utils.eval_param(neuron, name='neuron', allowed_types=(core.BaseNeuron, )) + utils.eval_param(on_error, name="on_error", allowed_values=("ignore", "raise")) + utils.eval_param(neuron, name="neuron", allowed_types=(core.BaseNeuron,)) # If string, convert to units if isinstance(units, str): @@ -377,16 +396,20 @@ def to_neuron_space(units: Union[int, float, pint.Quantity, pint.Unit], return units if neuron.units.dimensionless: - if on_error == 'raise': - raise ValueError(f'Unable to convert "{str(units)}": Neuron units ' - 'unknown or dimensionless.') + if on_error == "raise": + raise ValueError( + f'Unable to convert "{str(units)}": Neuron units ' + "unknown or dimensionless." + ) else: return units if not neuron.is_isometric: - if on_error == 'raise': - raise ValueError(f'Unable to convert "{str(units)}": neuron is not ' - 'isometric ({neuron.units}).') + if on_error == "raise": + raise ValueError( + f'Unable to convert "{str(units)}": neuron is not ' + "isometric ({neuron.units})." + ) else: return units @@ -414,28 +437,52 @@ class NeuronProcessor: neuron. """ - def __init__(self, - nl: 'core.NeuronList', - function: Callable, - parallel: bool = False, - n_cores: int = os.cpu_count() // 2, - chunksize: int = 1, - progress: bool = True, - warn_inplace: bool = True, - omit_failures: bool = False, - exclude_zip: list = [], - desc: Optional[str] = None): + def __init__( + self, + nl: "core.NeuronList", + function: Callable, + parallel: bool = False, + n_cores: int = os.cpu_count() // 2, + chunksize: Union[int, Literal["auto"]] = "auto", + progress: bool = True, + warn_inplace: bool = True, + omit_failures: bool = False, + exclude_zip: list = [], + desc: Optional[str] = None, + backend: Literal["joblib", "pathos", "auto"] = "auto", + ): if utils.is_iterable(function): if len(function) != len(nl): - raise ValueError('Number of functions must match neurons.') + raise ValueError("Number of functions must match neurons.") self.funcs = function self.function = function[0] elif callable(function): self.funcs = [function] * len(nl) self.function = function else: - raise TypeError('Expected `function` to be callable or list ' - f'thereof, got "{type(function)}"') + raise TypeError( + "Expected `function` to be callable or list " + f'thereof, got "{type(function)}"' + ) + + utils.eval_param( + value=backend.split(":")[0], + name="backend", + allowed_values=("joblib", "pathos", "auto"), + ) + + if backend == "auto": + if Parallel is not None: + backend = "joblib" + elif ProcessingPool is not None: + backend = "pathos" + else: + raise ModuleNotFoundError( + "navis relies on `joblib` or `pathos` for multiprocessing!" + "Please install either and try again:\n" + " pip3 install joblib -U\n" + " pip3 install pathos -U" + ) self.nl = nl self.desc = desc @@ -446,14 +493,15 @@ def __init__(self, self.warn_inplace = warn_inplace self.exclude_zip = exclude_zip self.omit_failures = omit_failures + self.backend = backend # This makes sure that help and name match the functions being called functools.update_wrapper(self, self.function) def __call__(self, *args, **kwargs): # Explicitly providing these parameters overwrites defaults - parallel = kwargs.pop('parallel', self.parallel) - n_cores = kwargs.pop('n_cores', self.n_cores) + parallel = kwargs.pop("parallel", self.parallel) + n_cores = kwargs.pop("n_cores", self.n_cores) # We will check, for each argument, if it matches the number of # functions to run. If they it does, we will zip the values @@ -484,54 +532,111 @@ def __call__(self, *args, **kwargs): level = logger.getEffectiveLevel() if level < 30: - logger.setLevel('WARNING') + logger.setLevel("WARNING") # Apply function if parallel: - if not ProcessingPool: - raise ModuleNotFoundError( - 'navis relies on pathos for multiprocessing!' - 'Please install pathos and try again:\n' - ' pip3 install pathos -U' + if self.warn_inplace and kwargs.get("inplace", False): + logger.warning("`inplace=True` does not work with multiprocessing ") + + if self.backend.startswith("joblib"): + if Parallel is None: + raise ModuleNotFoundError( + "`joblib` backend for multiprocessing not found!" + "Please install `joblib` and try again:\n" + " pip3 install joblib -U" ) - if self.warn_inplace and kwargs.get('inplace', False): - logger.warning('`inplace=True` does not work with ' - 'multiprocessing ') - - with ProcessingPool(n_cores) as pool: - combinations = list(zip(self.funcs, - parsed_args, - parsed_kwargs)) - chunksize = kwargs.pop('chunksize', self.chunksize) # max(int(len(combinations) / 100), 1) - if not self.omit_failures: wrapper = _call else: wrapper = _try_call - res = list(config.tqdm(pool.imap(wrapper, - combinations, - chunksize=chunksize), - total=len(combinations), - desc=self.desc, - disable=config.pbar_hide or not self.progress, - leave=config.pbar_leave)) + if ":" in self.backend: + backend = self.backend.split(":")[1] + else: + backend = "loky" + + combinations = zip(self.funcs, parsed_args, parsed_kwargs) + total = min(len(self.funcs), len(parsed_args), len(parsed_kwargs)) + + with tqdm_joblib( + desc=self.desc, + total=total, + disable=config.pbar_hide or not self.progress, + leave=config.pbar_leave, + ) as _: + # Suppress the Loky warnings about workers stopping to work + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message=".*worker stopped while some jobs.*", + category=UserWarning, + ) + res = Parallel( + n_jobs=n_cores, + backend=backend, + batch_size=self.chunksize, + )(delayed(wrapper)(x) for x in combinations) + elif self.backend.startswith("pathos"): + if ProcessingPool is None: + raise ModuleNotFoundError( + "`pathos` backend for multiprocessing not found!" + "Please install `pathos` and try again:\n" + " pip3 install pathos -U" + ) + + with ProcessingPool(n_cores) as pool: + combinations = zip(self.funcs, parsed_args, parsed_kwargs) + total = min(len(self.funcs), len(parsed_args), len(parsed_kwargs)) + + if self.chunksize == "auto": + chunksize = max(int(total / (n_cores * 10)), 1) + else: + chunksize = self.chunksize + + if not self.omit_failures: + wrapper = _call + else: + wrapper = _try_call + + res = list( + config.tqdm( + pool.imap(wrapper, combinations, chunksize=chunksize), + total=total, + desc=self.desc, + disable=config.pbar_hide or not self.progress, + leave=config.pbar_leave, + ) + ) + else: + raise ValueError( + f'Unknown backend "{self.backend}" for multiprocessing.' + ) else: res = [] - for i, n in enumerate(config.tqdm(self.nl, desc=self.desc, - disable=(config.pbar_hide - or not self.progress - or len(self.nl) <= 1), - leave=config.pbar_leave)): + for i, n in enumerate( + config.tqdm( + self.nl, + desc=self.desc, + disable=( + config.pbar_hide or not self.progress or len(self.nl) <= 1 + ), + leave=config.pbar_leave, + ) + ): try: res.append(self.funcs[i](*parsed_args[i], **parsed_kwargs[i])) except BaseException as e: if self.omit_failures: - res.append(FailedRun(func=self.funcs[i], - args=parsed_args[i], - kwargs=parsed_kwargs[i], - exception=e)) + res.append( + FailedRun( + func=self.funcs[i], + args=parsed_args[i], + kwargs=parsed_kwargs[i], + exception=e, + ) + ) else: raise @@ -541,11 +646,15 @@ def __call__(self, *args, **kwargs): failed = np.array([isinstance(r, FailedRun) for r in res]) res = [r for r in res if not isinstance(r, FailedRun)] if any(failed): - logger.warn(f'{sum(failed)} of {len(self.funcs)} runs failed. ' - 'Set logging to debug (`navis.set_loggers("DEBUG")`) ' - 'or repeat with `omit_failures=False` for details.') + logger.warn( + f"{sum(failed)} of {len(self.funcs)} runs failed. " + 'Set logging to debug (`navis.set_loggers("DEBUG")`) ' + "or repeat with `omit_failures=False` for details." + ) failed_ids = self.nl.id[np.where(failed)].astype(str) - logger.debug(f'The following IDs failed to complete: {", ".join(failed_ids)}') + logger.debug( + f"The following IDs failed to complete: {', '.join(failed_ids)}" + ) # If result is a list of neurons, combine them back into a single list is_neuron = [isinstance(r, (core.NeuronList, core.BaseNeuron)) for r in res] @@ -575,7 +684,8 @@ def _try_call(x: Sequence): class FailedRun: """Class representing a failed run.""" - def __init__(self, func, args, kwargs, exception='NA'): + + def __init__(self, func, args, kwargs, exception="NA"): self.args = args self.func = func self.kwargs = kwargs @@ -585,5 +695,7 @@ def __repr__(self): return self.__str__() def __str__(self): - return (f'Failed run(function={self.func}, args={self.args}, ' - f'kwargs={self.kwargs}, exception={self.exception})') + return ( + f"Failed run(function={self.func}, args={self.args}, " + f"kwargs={self.kwargs}, exception={self.exception})" + ) diff --git a/navis/core/neuronlist.py b/navis/core/neuronlist.py index 19fc20ec..c8413393 100644 --- a/navis/core/neuronlist.py +++ b/navis/core/neuronlist.py @@ -592,6 +592,8 @@ def apply(self, func: Callable, *, parallel: bool = False, + backend: str = "auto", + chunksize: Union[int, str] = "auto", n_cores: int = os.cpu_count() // 2, omit_failures: bool = False, **kwargs): @@ -636,6 +638,8 @@ def apply(self, proc = NeuronProcessor(self, func, parallel=parallel, + backend=backend, + chunksize=chunksize, n_cores=n_cores, omit_failures=omit_failures, desc=f'Apply {func.__name__}') diff --git a/navis/utils/decorators.py b/navis/utils/decorators.py index 121df55c..17493a7c 100644 --- a/navis/utils/decorators.py +++ b/navis/utils/decorators.py @@ -175,8 +175,6 @@ def wrapper(*args, **kwargs): kwargs["inplace"] = True # Prepare processor - n_cores = kwargs.pop("n_cores", os.cpu_count() // 2) - chunksize = kwargs.pop("chunksize", 1) excl = list(kwargs.keys()) + list(range(1, len(args) + 1)) proc = core.NeuronProcessor( nl, @@ -185,10 +183,11 @@ def wrapper(*args, **kwargs): desc=desc, warn_inplace=False, progress=kwargs.pop("progress", True), + backend=kwargs.pop("backend", "auto"), omit_failures=kwargs.pop("omit_failures", False), - chunksize=chunksize, + chunksize=kwargs.pop("chunksize", "auto"), exclude_zip=excl, - n_cores=n_cores, + n_cores=kwargs.pop("n_cores", os.cpu_count() // 2), ) # Apply function res = proc(nl, *args, **kwargs) @@ -222,8 +221,8 @@ def map_neuronlist_df( ): """Decorate function to run on all neurons in the NeuronList. - This version of the decorator is meant for functions that return a - DataFrame. This decorator will add a `neuron` column with the respective + This version of the parallelization decorator is meant for functions that return + a DataFrame. This decorator will add a `neuron` column with the respective neuron's ID and will then concatenate the dataframes. Parameters @@ -295,8 +294,6 @@ def wrapper(*args, **kwargs): _ = kwargs.pop(nl_key) # Prepare processor - n_cores = kwargs.pop("n_cores", os.cpu_count() // 2) - chunksize = kwargs.pop("chunksize", 1) excl = list(kwargs.keys()) + list(range(1, len(args) + 1)) proc = core.NeuronProcessor( nl, @@ -306,9 +303,10 @@ def wrapper(*args, **kwargs): warn_inplace=False, progress=kwargs.pop("progress", True), omit_failures=kwargs.pop("omit_failures", False), - chunksize=chunksize, + chunksize=kwargs.pop("chunksize", "auto"), exclude_zip=excl, - n_cores=n_cores, + backend=kwargs.pop("backend", "auto"), + n_cores=kwargs.pop("n_cores", os.cpu_count() // 2), ) # Apply function res = proc(nl, *args, **kwargs) @@ -362,23 +360,27 @@ def map_neuronlist_update_docstring(func, allow_parallel): msg = "" if allow_parallel: - msg += dedent(f"""\ + msg += dedent( + f"""\ parallel :{" " * (offset - 10)}bool {" " * (offset - 10)}If True and input is NeuronList, use parallel {" " * (offset - 10)}processing. Requires `pathos`. n_cores : {" " * (offset - 10)}int, optional {" " * (offset - 10)}Numbers of cores to use if `parallel=True`. {" " * (offset - 10)}Defaults to half the available cores. - """) + """ + ) - msg += dedent(f"""\ + msg += dedent( + f"""\ progress :{" " * (offset - 10)}bool {" " * (offset - 10)}Whether to show a progress bar. Overruled by {" " * (offset - 10)}`navis.set_pbars`. omit_failures :{" " * (offset - 15)}bool {" " * (offset - 15)}If True will omit failures instead of raising {" " * (offset - 15)}an exception. Ignored if input is single neuron. - """) + """ + ) # Insert new docstring lines.insert(lastp, indent(msg, wspaces)) diff --git a/requirements.txt b/requirements.txt index 167d1731..7a9551ca 100644 --- a/requirements.txt +++ b/requirements.txt @@ -23,6 +23,11 @@ rdata>=0.5 igraph!=0.10.0,!=0.10.1 skeletor>=1.0.0 +# extra: joblib + +joblib~=1.5.2 +tqdm-joblib~=0.0.5 + pathos>=0.2.7 #extra: pathos Shapely>=1.6.0 #extra: shapely @@ -75,7 +80,7 @@ Shapely>=1.6.0 flake8 wheel mypy -pytest +pytest!=9.0.0 pytest-env pytest-xvfb pytest-timeout