Skip to content

Parallel

Helpers shared by the samplers for optional parallel evaluation. See Parallel evaluation for the narrative version; you rarely need to touch these directly, since n_workers and backend on the samplers cover the normal cases.


WorkerPool

A pool held open for the lifetime of a sampler run.

Creating a process pool is expensive - on Windows each worker re-imports the calling module - so a pool rebuilt for every batch of likelihood evaluations can easily cost more than it saves. This opens one pool and reuses it for every batch, then shuts it down at the end.

Use it as a context manager::

with WorkerPool(n_workers=4, backend="auto", func=problem.log_likelihood) as pool:
    values = pool.map(problem.log_likelihood, particles)

With n_workers == 1 no pool is created at all and map is a plain list comprehension, so the serial path stays free of overhead.

Source code in mcmckit/core/parallel.py
class WorkerPool:
    """A pool held open for the lifetime of a sampler run.

    Creating a process pool is expensive - on Windows each worker re-imports the
    calling module - so a pool rebuilt for every batch of likelihood evaluations
    can easily cost more than it saves. This opens one pool and reuses it for
    every batch, then shuts it down at the end.

    Use it as a context manager::

        with WorkerPool(n_workers=4, backend="auto", func=problem.log_likelihood) as pool:
            values = pool.map(problem.log_likelihood, particles)

    With ``n_workers == 1`` no pool is created at all and ``map`` is a plain
    list comprehension, so the serial path stays free of overhead.
    """

    def __init__(self, n_workers=1, backend="auto", func=None, limit_blas=True):
        self.n_workers = resolve_n_workers(n_workers)
        self._requested_backend = backend
        self.backend = self._choose_backend(backend, func)
        self.limit_blas = limit_blas
        self._executor = None

    def _choose_backend(self, backend, func):
        if self.n_workers == 1:
            return "serial"
        if backend not in ("auto", "process", "thread"):
            raise ValueError(
                f"backend must be 'auto', 'process' or 'thread', got {backend!r}"
            )
        if backend == "auto":
            if func is None or check_picklable(func):
                return "process"
            raise ValueError(
                "n_workers > 1 needs a picklable log-likelihood so it can be "
                "sent to worker processes, and this one cannot be pickled. It "
                "is probably a lambda, a closure, or defined inside another "
                "function.\n\n"
                "Move it to module level:\n"
                "    def log_likelihood(theta):\n"
                "        ...\n"
                "    problem = mc.Problem(prior=log_prior, likelihood=log_likelihood)\n\n"
                "Only if the likelihood is genuinely thread-safe and releases "
                "the GIL may you pass backend='thread' instead. Do NOT do that "
                "for a solver holding global state: OpenSeesPy keeps one global "
                "model domain, and running it in threads corrupts the model or "
                "crashes the interpreter."
            )
        if backend == "process" and func is not None and not check_picklable(func):
            raise ValueError(
                "backend='process' needs a picklable log-likelihood, and this "
                "one cannot be pickled. It is probably a lambda, a closure, or "
                "defined inside another function. Move it to module level:\n"
                "    def log_likelihood(theta):\n"
                "        ...\n"
                "    problem = mc.Problem(prior=log_prior, likelihood=log_likelihood)"
            )
        return backend

    # -- context manager ------------------------------------------------

    def __enter__(self):
        if self.backend == "process":
            self._executor = ProcessPoolExecutor(
                max_workers=self.n_workers,
                initializer=limit_blas_threads if self.limit_blas else None,
            )
        elif self.backend == "thread":
            self._executor = ThreadPoolExecutor(max_workers=self.n_workers)
        return self

    def __exit__(self, *exc):
        if self._executor is not None:
            self._executor.shutdown(wait=True)
            self._executor = None
        return False

    # -- work -----------------------------------------------------------

    def map(self, func, items):
        """Apply ``func`` to every item, in order.

        Falls back to serial evaluation if the pool is not open, so a sampler
        can call this whether or not it is inside the context manager.
        """
        items = list(items)
        if self._executor is None:
            return [func(x) for x in items]
        return list(self._executor.map(func, items))

    def __repr__(self):
        return f"WorkerPool(n_workers={self.n_workers}, backend={self.backend!r})"

Methods:

map

map(func, items)

Apply func to every item, in order.

Falls back to serial evaluation if the pool is not open, so a sampler can call this whether or not it is inside the context manager.

Source code in mcmckit/core/parallel.py
def map(self, func, items):
    """Apply ``func`` to every item, in order.

    Falls back to serial evaluation if the pool is not open, so a sampler
    can call this whether or not it is inside the context manager.
    """
    items = list(items)
    if self._executor is None:
        return [func(x) for x in items]
    return list(self._executor.map(func, items))

resolve_n_workers

resolve_n_workers(n_workers) -> int

Normalise n_workers; -1 means "one per core".

Source code in mcmckit/core/parallel.py
def resolve_n_workers(n_workers) -> int:
    """Normalise ``n_workers``; ``-1`` means "one per core"."""
    if n_workers is None:
        return 1
    n_workers = int(n_workers)
    if n_workers == -1:
        return os.cpu_count() or 1
    if n_workers < 1:
        raise ValueError(f"n_workers must be >= 1, or -1 for all cores, got {n_workers}")
    return n_workers

check_picklable

check_picklable(func) -> bool

Whether func survives a pickle round trip.

Process workers receive the function by pickling it. Lambdas, closures and functions defined inside another function cannot be pickled, and the error raised deep inside the executor is unhelpful, so callers check up front.

Source code in mcmckit/core/parallel.py
def check_picklable(func) -> bool:
    """Whether ``func`` survives a pickle round trip.

    Process workers receive the function by pickling it. Lambdas, closures and
    functions defined inside another function cannot be pickled, and the error
    raised deep inside the executor is unhelpful, so callers check up front.
    """
    try:
        pickle.loads(pickle.dumps(func))
        return True
    except Exception:
        return False

limit_blas_threads

limit_blas_threads()

Pin each worker process to a single BLAS thread.

NumPy and SciPy already spread a single large dot or svd over every core. Running N such workers in parallel therefore asks for N x cores threads, and the resulting contention can make the parallel run slower than the serial one. Since the parallelism here is across likelihood evaluations, one BLAS thread per worker is the right split.

Called as the process-pool initialiser, before the worker imports NumPy.

Source code in mcmckit/core/parallel.py
def limit_blas_threads():
    """Pin each worker process to a single BLAS thread.

    NumPy and SciPy already spread a single large ``dot`` or ``svd`` over every
    core. Running N such workers in parallel therefore asks for N x cores
    threads, and the resulting contention can make the parallel run *slower*
    than the serial one. Since the parallelism here is across likelihood
    evaluations, one BLAS thread per worker is the right split.

    Called as the process-pool initialiser, before the worker imports NumPy.
    """
    for var in _BLAS_THREAD_VARS:
        os.environ[var] = "1"