From cbeb8b407d5d38aeddd103b2f38522dccb59d677 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 17 Aug 2026 15:51:32 +0000 Subject: [PATCH] Speed up parallel model import Parallel model import (`AMICI_IMPORT_NPROCS` > 1) created a brand-new `spawn` worker pool for every single parallelized operation. Since model import performs dozens of them, and every worker of every pool had to import sympy and amici from scratch, the process startup overhead often exceeded the actual work -- for all but the largest models, parallel import was *slower* than serial import. * Create the worker pool lazily and reuse it for all operations, instead of creating and tearing one down per operation. * Use the `forkserver` start method instead of `spawn`. Its workers are forked from a small, single-threaded server process that imports sympy and amici only once, which makes pool creation ~5x cheaper without reintroducing the deadlock potential of forking the (potentially multi-threaded) main process. `spawn` is kept only where `forkserver` is unavailable, i.e. on Windows. * Process small matrices serially, where the inter-process communication overhead outweighs any speed-up. * Compute the jacobian sparsity pattern from the symbols of each row instead of calling `Basic.has` for every (row, variable) pair, which re-traverses the full expression tree every time. This also speeds up serial import. * Validate `AMICI_IMPORT_NPROCS` instead of failing with an opaque error further down, and detect unpicklable functions in `_parallel_applyfunc` before dispatching them. The previous `PicklingError` handler never triggered for the most common case (a lambda), which raises `AttributeError`. Micro-benchmark (4 cores, 12 `smart_jacobian` calls, ~1.5 s to import amici): 25.0 s -> 2.8 s. For a workload large enough for parallelization to pay off, the speed-up over serial import improves from 2.3x to 4.2x on 4 cores. --- CHANGELOG.md | 8 + python/sdist/amici/_symbolic/sympy_utils.py | 202 +++++++++++++++----- python/tests/test_sympy_utils.py | 124 ++++++++++++ 3 files changed, 290 insertions(+), 44 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 77a028f8b7..f1f64b167a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -28,6 +28,14 @@ See also our [versioning policy](https://amici.readthedocs.io/en/latest/versioni * PEtab SciML: implemented additional PyTorch-style layer types (`BatchNorm`, `InstanceNorm`, `AlphaDropout`, `Bilinear`) (#3176). * PEtab SciML: updated support to PEtab v2 (#3165). +* Sped up parallel model import (`AMICI_IMPORT_NPROCS` > 1) considerably. + Previously, a new worker pool was created for every parallelized operation, + and every worker of every pool had to import sympy and amici from scratch. + Now, a single pool is created lazily and reused, and its workers are forked + from a `forkserver` process that performs these imports only once. For small + operations, where the inter-process communication overhead outweighs any + speed-up, processing stays serial. Previously, parallel import could be + slower than serial import for all but the largest models. **Fixes** diff --git a/python/sdist/amici/_symbolic/sympy_utils.py b/python/sdist/amici/_symbolic/sympy_utils.py index 8e2d75d977..4e2c79fd22 100644 --- a/python/sdist/amici/_symbolic/sympy_utils.py +++ b/python/sdist/amici/_symbolic/sympy_utils.py @@ -1,17 +1,22 @@ """Functionality for working with sympy objects.""" +import atexit import contextlib import logging import os +import threading from collections.abc import Callable from functools import wraps from itertools import starmap -from typing import Any +from typing import TYPE_CHECKING, Any import sympy as sp from amici.logging import get_logger, log_execution_time +if TYPE_CHECKING: + from multiprocessing.pool import Pool + logger = get_logger(__name__, logging.ERROR) __all__ = [ @@ -23,6 +28,97 @@ "_piecewise_to_minmax", ] +# Number of matrix elements below which the inter-process communication +# overhead is expected to outweigh any speed-up from parallel processing. +_MIN_PARALLEL_ELEMENTS = 16 + +# The worker pool for parallel model import, the ``(n_procs, pid)`` it was +# created for, and the lock guarding both. Creating a worker pool is expensive +# (every worker has to import sympy and amici), and model import performs many +# parallelizable operations. Therefore, a single pool is created lazily and +# reused for all of them, instead of creating a new one for each operation. +_pool: "Pool | None" = None +_pool_config: tuple[int, int] | None = None +_pool_lock = threading.RLock() + + +def _get_n_procs() -> int: + """Number of processes to be used for model import. + + Controlled via the ``AMICI_IMPORT_NPROCS`` environment variable. + """ + val = os.environ.get("AMICI_IMPORT_NPROCS", "1") + try: + n_procs = int(val) + if n_procs < 1: + raise ValueError + except ValueError: + raise ValueError( + f"Invalid value for AMICI_IMPORT_NPROCS: {val!r}. " + "Must be a positive integer." + ) from None + return n_procs + + +def _get_mp_context(): + """Get the multiprocessing context to be used for model import. + + ``fork`` is not used, since forking a potentially multi-threaded process + may deadlock (see e.g. https://stackoverflow.com/a/66113051). + + ``forkserver`` is used where available, because its workers are forked from + a small, single-threaded server process that imports sympy and amici only + once, whereas every ``spawn`` worker has to import them from scratch. This + makes creating a pool several times cheaper, without being affected by the + problems of forking the main process. Where ``forkserver`` is unavailable + (Windows), ``spawn`` is used. + """ + from multiprocessing import get_all_start_methods, get_context + + if "forkserver" not in get_all_start_methods(): + return get_context("spawn") + + ctx = get_context("forkserver") + # import the modules required by the workers once in the forkserver + # process, instead of once in every worker process + # (unimportable modules are silently ignored by the forkserver) + ctx.set_forkserver_preload(["sympy", __name__]) + return ctx + + +def _get_pool(n_procs: int) -> "Pool": + """Get the (lazily created) worker pool with ``n_procs`` processes.""" + global _pool, _pool_config + + with _pool_lock: + if _pool is not None and _pool_config != (n_procs, os.getpid()): + # either the requested number of processes changed, or we are in a + # forked child process that inherited the parent's unusable pool + _shutdown_pool() + + if _pool is None: + _pool = _get_mp_context().Pool(n_procs) + _pool_config = (n_procs, os.getpid()) + + return _pool + + +def _shutdown_pool() -> None: + """Shut down the worker pool, if any.""" + global _pool, _pool_config + + with _pool_lock: + if _pool is None: + return + if _pool_config[1] == os.getpid(): + # only the process that created the pool may terminate it + _pool.terminate() + _pool.join() + _pool = _pool_config = None + + +atexit.register(_shutdown_pool) + def _custom_pow_eval_derivative(self, s): """ @@ -113,27 +209,39 @@ def smart_jacobian( return sp.MutableSparseMatrix(nrow, ncol, dict()) # preprocess sparsity pattern - elements = ( - (i, j, a, b) - for i, a in enumerate(eq) - for j, b in enumerate(sym_var) - if a.has(b) - ) + if all(b.is_Symbol for b in sym_var): + # `Basic.has` traverses the full expression tree on every call. + # Collecting the symbols of each row once is equivalent, but avoids + # re-traversing every row for each of the `ncol` variables. + symbols_by_row = [a.atoms(sp.Symbol) for a in eq] + elements = ( + (i, j, a, b) + for i, a in enumerate(eq) + for j, b in enumerate(sym_var) + if b in symbols_by_row[i] + ) + else: + elements = ( + (i, j, a, b) + for i, a in enumerate(eq) + for j, b in enumerate(sym_var) + if a.has(b) + ) - if (n_procs := int(os.environ.get("AMICI_IMPORT_NPROCS", 1))) == 1: + if (n_procs := _get_n_procs()) == 1: # serial return sp.MutableSparseMatrix( nrow, ncol, dict(starmap(_jacobian_element, elements)) ) - # parallel - from multiprocessing import get_context + # parallel -- the pool consumes the full iterable anyway + elements = list(elements) + if len(elements) < _MIN_PARALLEL_ELEMENTS: + return sp.MutableSparseMatrix( + nrow, ncol, dict(starmap(_jacobian_element, elements)) + ) - # "spawn" should avoid potential deadlocks occurring with fork - # see e.g. https://stackoverflow.com/a/66113051 - ctx = get_context("spawn") - with ctx.Pool(n_procs) as p: - mapped = p.starmap(_jacobian_element, elements) + mapped = _get_pool(n_procs).starmap(_jacobian_element, elements) return sp.MutableSparseMatrix(nrow, ncol, dict(mapped)) @@ -190,41 +298,47 @@ def _jacobian_element(i, j, eq_i, sym_var_j): def _parallel_applyfunc(obj: sp.Matrix, func: Callable) -> sp.Matrix: """Parallel implementation of sympy's Matrix.applyfunc""" - if (n_procs := int(os.environ.get("AMICI_IMPORT_NPROCS", 1))) == 1: + if (n_procs := _get_n_procs()) == 1: # serial return obj.applyfunc(func) - # parallel - from multiprocessing import get_context - from pickle import PicklingError + from multiprocessing.reduction import ForkingPickler from sympy.matrices.dense import DenseMatrix - # "spawn" should avoid potential deadlocks occurring with fork - # see e.g. https://stackoverflow.com/a/66113051 - ctx = get_context("spawn") - with ctx.Pool(n_procs) as p: - try: - if isinstance(obj, DenseMatrix): - return obj._new(obj.rows, obj.cols, p.map(func, obj)) - elif isinstance(obj, sp.SparseMatrix): - dok = obj.todok() - mapped = p.map(func, dok.values()) - dok = { - k: v - for k, v in zip(dok.keys(), mapped, strict=True) - if v != 0 - } - return obj._new(obj.rows, obj.cols, dok) - else: - raise ValueError(f"Unsupported matrix type {type(obj)}") - except PicklingError as e: - raise ValueError( - f"Couldn't pickle {func}. This is likely because the argument " - "was not a module-level function. Either rewrite the argument " - "to a module-level function or disable parallelization by " - "setting `AMICI_IMPORT_NPROCS=1`." - ) from e + if isinstance(obj, DenseMatrix): + values = list(obj) + dok = None + elif isinstance(obj, sp.SparseMatrix): + dok = obj.todok() + values = list(dok.values()) + else: + raise ValueError(f"Unsupported matrix type {type(obj)}") + + if len(values) < _MIN_PARALLEL_ELEMENTS: + # not worth the inter-process communication overhead + return obj.applyfunc(func) + + try: + # check upfront -- passing `func` to the pool would surface this as an + # opaque error from the pool's task handler thread, with an exception + # type that depends on why exactly `func` is unpicklable + ForkingPickler.dumps(func) + except Exception as e: + raise ValueError( + f"Couldn't pickle {func}. This is likely because the argument " + "was not a module-level function. Either rewrite the argument " + "to a module-level function or disable parallelization by " + "setting `AMICI_IMPORT_NPROCS=1`." + ) from e + + mapped = _get_pool(n_procs).map(func, values) + + if dok is None: + return obj._new(obj.rows, obj.cols, mapped) + + dok = {k: v for k, v in zip(dok.keys(), mapped, strict=True) if v != 0} + return obj._new(obj.rows, obj.cols, dok) def _piecewise_to_minmax( diff --git a/python/tests/test_sympy_utils.py b/python/tests/test_sympy_utils.py index c60d356868..0994467d9e 100644 --- a/python/tests/test_sympy_utils.py +++ b/python/tests/test_sympy_utils.py @@ -1,14 +1,34 @@ """Tests related to the sympy_utils module.""" +import pytest import sympy as sp from amici._symbolic.sympy_utils import ( _custom_pow_eval_derivative, + _get_mp_context, + _get_n_procs, + _get_pool, _monkeypatched, + _parallel_applyfunc, _piecewise_to_minmax, + _shutdown_pool, + smart_jacobian, ) from amici.testing import skip_on_valgrind +@pytest.fixture +def nprocs(request, monkeypatch): + """Set ``AMICI_IMPORT_NPROCS`` and clean up any worker pool afterwards.""" + monkeypatch.setenv("AMICI_IMPORT_NPROCS", str(request.param)) + yield request.param + _shutdown_pool() + + +def _simplify(x): + """Module-level (i.e. picklable) simplification function.""" + return sp.simplify(x) + + @skip_on_valgrind def test_monkeypatch(): t = sp.Symbol("t") @@ -74,3 +94,107 @@ def test_rewrite_piecewise_minmax(): (sp.Min(y, z), True), ) assert replaced == expected + + +@skip_on_valgrind +@pytest.mark.parametrize("nprocs", [1, 2], indirect=True) +def test_smart_jacobian(nprocs): + """Serial and parallel jacobians must agree with sympy's.""" + # enough elements to exceed the serial-processing threshold + x = sp.Matrix(sp.symbols("x0:6")) + eq = sp.Matrix([xi**2 * xj + sp.sin(xi) for xi in x for xj in x]) + expected = eq.jacobian(x) + + actual = smart_jacobian(eq, x) + + assert actual.shape == expected.shape + assert (sp.Matrix(actual) - expected).is_zero_matrix + + # empty input gives an empty matrix, not an error + empty = smart_jacobian(sp.Matrix(0, 1, []), x) + assert empty.shape == (0, len(x)) + + +@skip_on_valgrind +@pytest.mark.parametrize("nprocs", [1, 2], indirect=True) +def test_smart_jacobian_non_symbol_vars(nprocs): + """The sparsity pattern must also be correct for non-``Symbol`` vars.""" + t = sp.Symbol("t") + funcs = sp.Matrix([sp.Function("x")(t), sp.Function("y")(t)]) + eq = sp.Matrix([f**2 for f in funcs] * 10) + + actual = smart_jacobian(eq, funcs) + + assert (sp.Matrix(actual) - eq.jacobian(funcs)).is_zero_matrix + + +@skip_on_valgrind +@pytest.mark.parametrize("nprocs", [1, 2], indirect=True) +@pytest.mark.parametrize("n_elements", [4, 40]) +@pytest.mark.parametrize("sparse", [True, False]) +def test_parallel_applyfunc(nprocs, n_elements, sparse): + """``_parallel_applyfunc`` must match ``Matrix.applyfunc``. + + Tested for dense/sparse matrices and for element counts below and above + the threshold for switching to serial processing. + """ + x = sp.Symbol("x") + entries = [x ** (i + 2) / x for i in range(n_elements)] + if sparse: + obj = sp.MutableSparseMatrix( + n_elements, + n_elements, + {(i, i): entry for i, entry in enumerate(entries)}, + ) + else: + obj = sp.MutableDenseMatrix(n_elements, 1, entries) + + actual = _parallel_applyfunc(obj, _simplify) + + assert type(actual) is type(obj) + assert actual == obj.applyfunc(_simplify) + + +@skip_on_valgrind +@pytest.mark.parametrize("nprocs", [2], indirect=True) +def test_parallel_applyfunc_unpicklable(nprocs): + """Unpicklable functions must give an actionable error message.""" + x = sp.Symbol("x") + obj = sp.MutableDenseMatrix(40, 1, [x**2 / x] * 40) + + with pytest.raises(ValueError, match="Couldn't pickle"): + _parallel_applyfunc(obj, lambda e: sp.simplify(e)) + + +@skip_on_valgrind +@pytest.mark.parametrize("nprocs", [2], indirect=True) +def test_pool_is_reused(nprocs): + """The worker pool must be created once and reused.""" + pool = _get_pool(2) + assert _get_pool(2) is pool + # ... but recreated if a different number of processes is requested + assert _get_pool(3) is not pool + + +@skip_on_valgrind +def test_get_mp_context(): + import multiprocessing + + available = multiprocessing.get_all_start_methods() + expected = "forkserver" if "forkserver" in available else "spawn" + # in particular, never fork the (potentially multi-threaded) main process + assert _get_mp_context().get_start_method() == expected + + +@skip_on_valgrind +def test_get_n_procs(monkeypatch): + monkeypatch.delenv("AMICI_IMPORT_NPROCS", raising=False) + assert _get_n_procs() == 1 + + monkeypatch.setenv("AMICI_IMPORT_NPROCS", "4") + assert _get_n_procs() == 4 + + for invalid in ("0", "-1", "some_string", ""): + monkeypatch.setenv("AMICI_IMPORT_NPROCS", invalid) + with pytest.raises(ValueError, match="AMICI_IMPORT_NPROCS"): + _get_n_procs()