Speed up parallel model import

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**
 

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(