From fc3f3c3ea9a27587d3b566e7933714ec5d2ae000 Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Fri, 10 Jul 2026 14:36:55 +0200 Subject: [PATCH 1/3] Install jax from pip --- conda-envs/environment-alternative-backends.yml | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/conda-envs/environment-alternative-backends.yml b/conda-envs/environment-alternative-backends.yml index 195c5ad92c..7379c60aa0 100644 --- a/conda-envs/environment-alternative-backends.yml +++ b/conda-envs/environment-alternative-backends.yml @@ -13,13 +13,11 @@ dependencies: - numba # nutpie is installed from git in the CI workflow (see tests.yml) so we can track # the arviz 1.0 -compatible branch until an upstream release ships. -# Jaxlib version must not be greater than jax version! -- jax>=0.4.28 -- jaxlib>=0.4.28 +# The JAX stack (jax, jaxlib, numpyro, blackjax) is installed from PyPI in the pip +# section below rather than conda-forge -- see the note there. - libblas=*=*mkl - mkl-service - numpy>=1.25.0 -- numpyro>=0.8.0 - pandas>=0.24.0 - pip - pytensor>=3.1.2,<3.2 @@ -39,7 +37,15 @@ dependencies: - types-cachetools # blackjax now does not pull in fastprogress by default - fastprogress>=1.0.0 -- blackjax>=1.5 - pip: - numdifftools>=0.9.40 - mcbackend>=0.4.0 + # conda-forge builds jaxlib from source with its own XLA/LLVM toolchain, and its + # 0.10.x CPU build segfaults in XLA codegen while compiling some graphs + # (DirichletMultinomial, MvStudentT draws) on the CI runners. The official PyPI + # wheels compile the same graphs fine, so install the whole JAX stack from PyPI. + # Jaxlib version must not be greater than jax version! + - jax>=0.4.28 + - jaxlib>=0.4.28 + - numpyro>=0.8.0 + - blackjax>=1.5 From 45f032d93ca367e3ffbd47f76a4778a4b035a1c5 Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Fri, 10 Jul 2026 13:47:08 +0200 Subject: [PATCH 2/3] Don't import matplotlib to understand red --- pymc/progress_bar/rich_progress.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/pymc/progress_bar/rich_progress.py b/pymc/progress_bar/rich_progress.py index a98df51879..7b99b4aeb2 100644 --- a/pymc/progress_bar/rich_progress.py +++ b/pymc/progress_bar/rich_progress.py @@ -149,11 +149,10 @@ def __rich_console__(self, console: Console, options: ConsoleOptions) -> RenderR class CustomBarColumn(BarColumn): """Bar column that recolors on divergences and renders a separator marker.""" - def __init__(self, *args, failing_color: str = "red", **kwargs): - from matplotlib.colors import to_rgb - - self.failing_color = failing_color - self.failing_rgb = [int(x * 255) for x in to_rgb(self.failing_color)] + def __init__(self, *args, failing_style: Style | None = None, **kwargs): + self.failing_style = ( + failing_style if failing_style is not None else Style.parse("rgb(214,39,40)") + ) super().__init__(*args, **kwargs) @@ -163,8 +162,8 @@ def __init__(self, *args, failing_color: str = "red", **kwargs): def callbacks(self, task: Task): """Update bar color based on failure state.""" if task.fields["failing"]: - self.complete_style = Style.parse("rgb({},{},{})".format(*self.failing_rgb)) - self.finished_style = Style.parse("rgb({},{},{})".format(*self.failing_rgb)) + self.complete_style = self.failing_style + self.finished_style = self.failing_style else: self.complete_style = self.default_complete_style self.finished_style = self.default_finished_style @@ -264,7 +263,6 @@ def _create_progress_bar( CustomBarColumn( bar_width=None, table_column=Column("Progress", ratio=2), - failing_color="tab:red", complete_style=Style.parse("rgb(31,119,180)"), finished_style=Style.parse("rgb(31,119,180)"), ), From f83b2aa51df8e4b7837b5e9ae83f678d02876c9f Mon Sep 17 00:00:00 2001 From: Ricardo Vieira Date: Fri, 10 Jul 2026 13:04:34 +0200 Subject: [PATCH 3/3] Allow free-threaded parallel sampling --- .github/workflows/tests.yml | 54 ++++++- pymc/model/core.py | 21 +++ pymc/sampling/mcmc.py | 209 ++++++++++++++++++++++--- pymc/sampling/parallel.py | 34 ++++ pymc/step_methods/arraystep.py | 42 +++++ pymc/step_methods/compound.py | 35 +++++ pymc/step_methods/hmc/base_hmc.py | 33 ++++ scripts/check_all_tests_are_covered.py | 5 +- setup.py | 1 + tests/conftest.py | 16 ++ tests/sampling/test_parallel.py | 47 ++++++ tests/step_methods/test_compound.py | 81 ++++++++++ 12 files changed, 555 insertions(+), 23 deletions(-) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index a4169733ff..53ef248925 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -345,6 +345,57 @@ jobs: name: Alternative backends py${{ matrix.python-version }} linker=${{ matrix.linker }} fail_ci_if_error: false + free_threading: + needs: changes + if: ${{ needs.changes.outputs.changes == 'true' }} + strategy: + matrix: + # A small, sampling-focused subset that exercises the thread-based parallel + # sampling path used on free-threaded (no-GIL) builds. The GIL-stays-disabled + # check is a pytest_sessionfinish hook in tests/conftest.py. + test-subset: + - | + tests/sampling/test_parallel.py + tests/step_methods/test_compound.py + fail-fast: false + runs-on: ubuntu-latest + env: + TEST_SUBSET: ${{ matrix.test-subset }} + # Disable the C backend (cxx=""). Its single-phase-init C extensions (e.g. + # lazylinker_ext) re-enable the GIL on a free-threaded build; with no C + # compiler PyTensor uses the free-threading-safe numba/Python backend. + PYTENSOR_FLAGS: cxx= + # Intentionally leave PYTHON_GIL unset. Setting it to 0 would force the GIL + # off and hide a dependency that fails to support free-threading; leaving it + # unset lets such a dependency re-enable the GIL, which the conftest + # pytest_sessionfinish hook asserts against. + defaults: + run: + shell: bash -leo pipefail {0} + steps: + - uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0 + with: + persist-credentials: false + - uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0 + with: + python-version: "3.14t" + - name: Install pymc and test dependencies + run: | + python -m pip install --upgrade pip + # Install from PyPI (free-threaded wheels), not conda-forge. numba is + # included as the free-threading-safe PyTensor backend. + pip install -e . numba pytest pytest-cov + # Free-threaded sampling needs the (still unreleased) PyTensor + # free-threading support (pymc-devs/pytensor#2286); track its branch + # until it is released. + pip install "pytensor @ git+https://github.com/ricardoV94/pytensor@free_threaded" + python --version + python -c "import sys, sysconfig; print('Py_GIL_DISABLED =', sysconfig.get_config_var('Py_GIL_DISABLED'), '| gil_enabled =', sys._is_gil_enabled())" + pip list + - name: Run tests + run: | + python -m pytest -vv --durations=25 $TEST_SUBSET + float32: needs: changes if: ${{ needs.changes.outputs.changes == 'true' }} @@ -395,7 +446,7 @@ jobs: all_tests: if: ${{ always() }} runs-on: ubuntu-latest - needs: [changes, ubuntu, windows, macos, alternative_backends, float32] + needs: [changes, ubuntu, windows, macos, alternative_backends, free_threading, float32] steps: - name: Check build matrix status if: ${{ needs.changes.outputs.changes == 'true' && @@ -403,5 +454,6 @@ jobs: needs.windows.result != 'success' || needs.macos.result != 'success' || needs.alternative_backends.result != 'success' || + needs.free_threading.result != 'success' || needs.float32.result != 'success' ) }} run: exit 1 diff --git a/pymc/model/core.py b/pymc/model/core.py index 676c5d6f67..82f7ed384d 100644 --- a/pymc/model/core.py +++ b/pymc/model/core.py @@ -13,6 +13,7 @@ # limitations under the License. from __future__ import annotations +import copy import functools import sys import threading @@ -275,6 +276,26 @@ def get_extra_values(self): return {var.name: self._extra_vars_shared[var.name].get_value() for var in self._extra_vars} + def fork(self) -> ValueGradFunction: + """Return an independent copy for concurrent (multi-thread) use. + + The copy shares the compiled code and any read-only inputs (e.g. observed + data) with the original, but gets its own input/output storage and its own + ``extra_vars`` shared variables. This makes it safe to call and to + ``set_extra_values`` on from a different thread than the original. + """ + new = copy.copy(self) + swap = {} + new._extra_vars_shared = {} + for name, shared in self._extra_vars_shared.items(): + value = shared.get_value(borrow=False) + fresh = pytensor.shared(value, shared.name, shape=value.shape) + new._extra_vars_shared[name] = fresh + swap[shared] = fresh + new._pytensor_function = self._pytensor_function.copy(share_memory=False, swap=swap or None) + new._extra_are_set = False + return new + def __call__(self, grad_vars, *, extra_vars=None): if extra_vars is not None: self.set_extra_values(extra_vars) diff --git a/pymc/sampling/mcmc.py b/pymc/sampling/mcmc.py index 3b2cb464e2..ffce797a0a 100644 --- a/pymc/sampling/mcmc.py +++ b/pymc/sampling/mcmc.py @@ -21,10 +21,12 @@ import pickle import re import sys +import threading import time import warnings from collections.abc import Callable, Iterator, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor from typing import ( TYPE_CHECKING, Any, @@ -67,7 +69,13 @@ default_progress_theme, ) from pymc.pytensorf import resolve_backend_compile_kwargs -from pymc.sampling.parallel import Draw, _cpu_count, _initialize_multiprocessing_context +from pymc.sampling.parallel import ( + THREAD_CONTEXT, + Draw, + _cpu_count, + _initialize_multiprocessing_context, + _initialize_parallel_context, +) from pymc.sampling.population import _sample_population from pymc.stats.convergence import ( log_warning_stats, @@ -745,9 +753,11 @@ def sample( the ``draw.chain`` argument can be used to determine which of the active chains the sample is drawn from. Sampling can be interrupted by throwing a ``KeyboardInterrupt`` in the callback. - mp_ctx : multiprocessing.context.BaseContent - A multiprocessing context for parallel sampling. - See multiprocessing documentation for details. + mp_ctx : multiprocessing.context.BaseContext or "thread", optional + A multiprocessing context for parallel sampling; see the multiprocessing + documentation for details. Pass ``"thread"`` to run chains as threads in a + single process instead, which only runs in parallel on a free-threaded + (no-GIL) interpreter and is selected automatically there when unset. model : Model (optional if in ``with`` context) Model to sample from. The model needs to have free random variables. backend: str, optional. @@ -878,9 +888,9 @@ def sample( chains = max(2, cores) compile_kwargs = resolve_backend_compile_kwargs(backend, compile_kwargs) - mp_ctx = _initialize_multiprocessing_context( - mp_ctx, mode=compile_kwargs.get("mode"), quiet=quiet - ) + # `mp_ctx` may resolve to the THREAD_CONTEXT sentinel (free-threaded builds, or + # an explicit mp_ctx="thread"): run the chains as threads instead of processes. + mp_ctx = _initialize_parallel_context(mp_ctx, mode=compile_kwargs.get("mode"), quiet=quiet) joined_blas_limiter, cores, num_blas_cores_per_worker = setup_cores_blas_cores( blas_cores, chains, cores, mp_ctx ) @@ -1107,6 +1117,11 @@ def sample( ) parallel = cores > 1 and chains > 1 and not has_population_samplers + # When the parallel context resolved to threads (free-threaded build, or an + # explicit mp_ctx="thread"), run the chains as threads: one shared model + + # compiled code, no pickling. Only attempted when the step method supports + # fork(); otherwise we fall back to multiprocessing below. + threaded = parallel and mp_ctx == THREAD_CONTEXT # At some point it was decided that PyMC should not set a global seed by default, # unless the user specified a seed. This is a symptom of the fact that PyMC samplers # are built around global seeding. This branch makes sure we maintain this unspoken @@ -1121,13 +1136,34 @@ def sample( sample_args["rngs"] = rngs t_start = time.time() + sampled = False with _quiet_logging(quiet): - if parallel: + if threaded: + if not quiet: + _log.info(f"Threaded sampling ({chains} chains in {cores} threads)") + _print_step_hierarchy(step) + try: + with joined_blas_limiter(): + _thread_sample(**sample_args) + sampled = True + except NotImplementedError as e: + if not quiet: + _log.warning( + f"Step method does not support threaded sampling ({e}); " + "falling back to multiprocessing." + ) + # Resolve a real process context in place of the "thread" sentinel. + mp_ctx = _initialize_multiprocessing_context( + None, mode=compile_kwargs.get("mode"), quiet=quiet + ) + parallel_args["mp_ctx"] = mp_ctx + if parallel and not sampled: if not quiet: _log.info(f"Multiprocess sampling ({chains} chains in {cores} jobs)") _print_step_hierarchy(step) try: _mp_sample(**sample_args, **parallel_args) + sampled = True except pickle.PickleError: if not quiet: _log.warning("Could not pickle model, sampling singlethreaded.") @@ -1140,7 +1176,7 @@ def sample( _log.warning("Could not pickle model, sampling singlethreaded.") _log.debug("Pickling error:", exc_info=True) parallel = False - if not parallel: + if not sampled: if has_population_samplers: if not quiet: _log.info(f"Population sampling ({chains} chains)") @@ -1198,11 +1234,12 @@ def setup_cores_blas_cores( num_blas_cores_per_worker = None elif isinstance(blas_cores, int): - if mp_ctx.get_start_method() == "fork": + if mp_ctx != THREAD_CONTEXT and mp_ctx.get_start_method() == "fork": # https://github.com/pymc-devs/pymc/issues/7354 joined_blas_limiter = contextlib.nullcontext else: - + # Threads share one process, so a single process-wide BLAS limiter + # (rather than a per-worker split) applies to all chains. def joined_blas_limiter(): return threadpool_limits(limits=blas_cores) @@ -1420,6 +1457,111 @@ def _sample_many( return +def _thread_sample( + *, + draws: int, + chains: int, + cores: int, + traces: Sequence[IBaseTrace], + start: Sequence[PointType], + rngs: Sequence[np.random.Generator], + step: Step, + callback: SamplingIteratorCallback | None = None, + **kwargs, +): + """Sample all chains concurrently as threads (free-threaded / no-GIL builds). + + Each chain runs on an independent ``step.fork(rng)`` so chains never share + mutable compiled-function storage, adaptation state, or RNG. This is the + in-process counterpart of :func:`_mp_sample`: it shares one model, one set of + compiled thunks, and the observed data across chains instead of duplicating + them per process. + + At most ``cores`` chains run at once (via a thread pool of that size), matching + the multiprocessing path. ``Ctrl-C`` sets a stop flag that the workers check + between draws and unwind cooperatively -- worker threads never receive the + ``KeyboardInterrupt`` themselves. + + Raises ``NotImplementedError`` if the step method has not implemented + ``fork``; the caller catches this and falls back to multiprocessing. + + Parameters + ---------- + draws : int + Total number of iterations per chain, including tuning. + chains : int + Number of chains to sample. + cores : int + Maximum number of chains to sample concurrently. + traces : list of backends + One trace per chain; each is written only by its own thread. + start : list + Starting point for each chain. + rngs : list of Generators + One independent Generator per chain. + step : Step + The step method for chain 0; the other chains use ``step.fork(rng)``. + """ + tune = kwargs.get("tune", 0) + model = kwargs.get("model", None) + + # Fork one independent step per chain up front, in the main thread, so all + # compilation/copying happens before any worker starts -- and so a step that + # can't fork raises here (not inside a thread) for a clean fallback. + steps = [step] + [step.fork(rngs[i]) for i in range(1, chains)] + + progress_manager = MCMCProgressBarManager( + step_method=step, + chains=chains, + draws=draws - tune, + tune=tune, + progressbar=kwargs.get("progressbar", True), + progressbar_theme=kwargs.get("progressbar_theme", default_progress_theme), + ) + + errors: list[BaseException | None] = [None] * chains + progress_lock = threading.Lock() + stop_event = threading.Event() + + def _worker(i: int) -> None: + if stop_event.is_set(): + return + try: + _sample( + chain=i, + rng=rngs[i], + start=start[i], + draws=draws, + step=steps[i], + trace=traces[i], + tune=tune, + model=model, + callback=callback, + progress_manager=progress_manager, + progress_lock=progress_lock, + stop_event=stop_event, + ) + except BaseException as e: + errors[i] = e + + with progress_manager, ThreadPoolExecutor(max_workers=cores) as executor: + futures = [executor.submit(_worker, i) for i in range(chains)] + try: + for future in futures: + future.result() + except KeyboardInterrupt: + # Stop running workers at their next draw and drop chains that haven't + # started; the pool's shutdown waits for the in-flight ones to unwind. + stop_event.set() + for future in futures: + future.cancel() + + for e in errors: + if e is not None: + raise e + return + + def _sample( *, chain: int, @@ -1432,11 +1574,15 @@ def _sample( model: Model | None = None, callback=None, progress_manager: MCMCProgressBarManager, + progress_lock: threading.Lock | None = None, + stop_event: threading.Event | None = None, **kwargs, ) -> None: - """Sample one chain (singleprocess). + """Sample one chain (single process, in the calling thread). - Multiple step methods are supported via compound step methods. + Multiple step methods are supported via compound step methods. Used both by + the sequential sampler (:func:`_sample_many`) and, once per thread, by the + threaded sampler (:func:`_thread_sample`). Parameters ---------- @@ -1458,7 +1604,23 @@ def _sample( PyMC model. If None, the model is taken from the current context. progress_manager: ProgressBarManager Helper class used to handle progress bar styling and updates + progress_lock : threading.Lock, optional + Serializes progress-bar updates when several chains share one + ``progress_manager`` across threads. ``None`` (the sequential path) skips + locking entirely. + stop_event : threading.Event, optional + Cooperative cancellation flag checked between draws; when set, the chain + stops early without marking its bar complete. Used by threaded sampling to + unwind all chains on ``Ctrl-C``. """ + + def update(**kwargs): + if progress_lock is None: + progress_manager.update(**kwargs) + else: + with progress_lock: + progress_manager.update(**kwargs) + sampling_gen = _iter_sample( draws=draws, step=step, @@ -1472,17 +1634,22 @@ def _sample( ) try: for it, stats in enumerate(sampling_gen): - progress_manager.update( - chain_idx=chain, is_last=False, draw=it, stats=stats, tuning=it < tune - ) - - if not progress_manager.combined_progress or chain == progress_manager.chains - 1: - progress_manager.update( - chain_idx=chain, is_last=True, draw=it, stats=stats, tuning=False - ) + update(chain_idx=chain, is_last=False, draw=it, stats=stats, tuning=it < tune) + if stop_event is not None and stop_event.is_set(): + break + else: + # Reached only when the chain drew every sample (no early stop). In + # combined-progress mode a single bar tracks all chains, so only the + # last chain marks it complete; otherwise each chain finalizes its own. + if not progress_manager.combined_progress or chain == progress_manager.chains - 1: + update(chain_idx=chain, is_last=True, draw=it, stats=stats, tuning=False) except KeyboardInterrupt: pass + finally: + # Close the generator (and thus the trace) deterministically, whether we + # exhausted it, broke out on the stop flag, or errored. + sampling_gen.close() def _iter_sample( diff --git a/pymc/sampling/parallel.py b/pymc/sampling/parallel.py index 6d89149293..e8c5a04171 100644 --- a/pymc/sampling/parallel.py +++ b/pymc/sampling/parallel.py @@ -18,6 +18,7 @@ import multiprocessing import multiprocessing.sharedctypes import platform +import sys import time import traceback import warnings @@ -85,6 +86,39 @@ def rebuild_exc(exc, tb): return exc +#: Sentinel returned by `_initialize_parallel_context` to request thread-based +#: (rather than process-based) multi-chain sampling. +THREAD_CONTEXT = "thread" + + +def _initialize_parallel_context( + mp_ctx: str | multiprocessing.context.BaseContext | None, + *, + mode=None, + quiet: bool = False, +) -> multiprocessing.context.BaseContext | str: + """Resolve how to parallelize chains: threads or a multiprocessing context. + + Returns the :data:`THREAD_CONTEXT` sentinel to run chains as threads -- when + the user explicitly passes ``mp_ctx="thread"``, or, on a free-threaded (no-GIL) + interpreter, when they have not requested a specific context. Either way we + respect it: the request was explicit, or we concluded threading was available. + Otherwise it composes :func:`_initialize_multiprocessing_context` for a real + process context. + """ + gil_enabled = getattr(sys, "_is_gil_enabled", lambda: True)() + if mp_ctx == THREAD_CONTEXT or (mp_ctx is None and not gil_enabled): + if mp_ctx == THREAD_CONTEXT and gil_enabled and not quiet: + warnings.warn( + "mp_ctx='thread' was requested but the interpreter's GIL is " + "enabled, so chains will not actually run in parallel.", + UserWarning, + stacklevel=2, + ) + return THREAD_CONTEXT + return _initialize_multiprocessing_context(mp_ctx, mode=mode, quiet=quiet) + + def _initialize_multiprocessing_context( mp_ctx: str | multiprocessing.context.BaseContext | None, *, diff --git a/pymc/step_methods/arraystep.py b/pymc/step_methods/arraystep.py index 0c20e09a47..f8dcdca73e 100644 --- a/pymc/step_methods/arraystep.py +++ b/pymc/step_methods/arraystep.py @@ -12,11 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +import copy + from abc import abstractmethod from collections.abc import Callable from typing import cast import numpy as np +import pytensor from pymc.blocking import DictToArrayBijection, PointType, RaveledVars, StatsType from pymc.model import modelcontext @@ -104,6 +107,45 @@ def __init__(self, vars, shared, blocked=True, rng: RandomGenerator = None): self.blocked = blocked self.rng = get_random_generator(rng) + def fork(self, rng: RandomGenerator) -> "ArrayStepShared": + """Return an independent copy of this step for a new chain/thread. + + Gives the copy its own ``self.shared`` variables and its own storage for + every compiled pytensor function it holds (via ``Function.copy(swap=...)``, + so the compiled code is shared but not the mutable containers). Mutable + tuning/adaptation state is cloned independently through the + ``sampling_state`` round-trip. + + The copy is *not* seeded from ``rng`` here: ``_iter_sample`` calls + ``set_rng`` on every chain's step in all sampling modes, and ``set_rng`` + advances the passed generator (via ``spawn``). Seeding here as well would + consume the chain's shared RNG twice and desynchronize threaded chains + from the sequential/multiprocessing results. + + A subclass whose extra per-chain state is not captured by ``self.shared``, + a compiled ``Function`` attribute, or ``sampling_state`` must override or + re-disable ``fork`` (raise ``NotImplementedError``) so threaded sampling + falls back to multiprocessing. + """ + # Fresh shared containers for the extra-var shareds; swap map repoints the + # copied compiled functions to read this chain's private containers. + swap = {} + new_shared = {} + for name, shared_var in self.shared.items(): + fresh = pytensor.shared(shared_var.get_value(borrow=False), shared_var.name) + swap[shared_var] = fresh + new_shared[name] = fresh + + new = copy.copy(self) + new.shared = new_shared + for attr, value in vars(self).items(): + if isinstance(value, pytensor.compile.Function): + setattr(new, attr, value.copy(share_memory=False, swap=swap or None)) + + # Independent mutable tuning + RNG state (deep-copied, pytensor-free). + new.sampling_state = self.sampling_state + return new + def step(self, point: PointType) -> tuple[PointType, StatsType]: full_point = None if self.shared: diff --git a/pymc/step_methods/compound.py b/pymc/step_methods/compound.py index 556b81b549..ed040da6fb 100644 --- a/pymc/step_methods/compound.py +++ b/pymc/step_methods/compound.py @@ -233,6 +233,26 @@ def stop_tuning(self): def set_rng(self, rng: RandomGenerator): self.rng = get_random_generator(rng, copy=False) + def fork(self, rng: RandomGenerator) -> "BlockedStep": + """Return an independent copy of this step for a new chain in a thread. + + Threaded (free-threaded / no-GIL) multi-chain sampling runs one chain + per thread sharing a single model. Each chain needs a step with its own + mutable state -- compiled-function storage, shared variables, tuning and + adaptation state, and RNG -- or the chains race. ``fork`` produces such + an independent copy, ideally sharing read-only compiled code. + + The base implementation raises ``NotImplementedError``. A step type opts + in to threaded sampling by overriding ``fork``; until then, a threaded + sampler must catch this and fall back to multiprocessing. Failing loudly + (rather than a generic best-effort copy) keeps unknown or third-party + step methods correct-by-default and lets ``fork`` be added incrementally. + """ + raise NotImplementedError( + f"{type(self).__name__} does not implement fork() for threaded sampling. " + "Fall back to multiprocessing for this step method." + ) + def flat_statname(sampler_idx: int, sname: str) -> str: """Get the flat-stats name for a samplers stat.""" @@ -318,6 +338,21 @@ def set_rng(self, rng: RandomGenerator): for method, _rng in zip(self.methods, _rngs): method.set_rng(_rng) + def fork(self, rng: RandomGenerator) -> "CompoundStep": + """Fork every child step for a new chain/thread. + + Propagates ``NotImplementedError`` if any child cannot fork, so the + caller falls back to multiprocessing for the whole compound step rather + than forking only some children. + + Children are not seeded from ``rng`` here; ``_iter_sample`` calls + ``set_rng`` on the compound step once per chain, which spawns and + distributes an independent generator to each child. Seeding during fork + as well would consume the chain's RNG twice (see ``BlockedStep.fork``). + """ + forked = [method.fork(rng) for method in self.methods] + return CompoundStep(forked) + def _progressbar_config(self, n_chains=1): from functools import reduce diff --git a/pymc/step_methods/hmc/base_hmc.py b/pymc/step_methods/hmc/base_hmc.py index c3e6d75e5c..083e82acd3 100644 --- a/pymc/step_methods/hmc/base_hmc.py +++ b/pymc/step_methods/hmc/base_hmc.py @@ -14,6 +14,7 @@ from __future__ import annotations +import copy import logging import time @@ -300,3 +301,35 @@ def reset(self, start=None): def set_rng(self, rng: RandomGenerator): self.rng = get_random_generator(rng, copy=False) self.potential.set_rng(self.rng.spawn(1)[0]) + + def fork(self, rng: RandomGenerator) -> BaseHMC: + """Return an independent copy of this step for a new chain/thread. + + Shares the compiled logp/dlogp code and read-only data, but gives the + copy its own function storage, its own ``extra_vars`` shared variables, + and fresh adaptation state (mass matrix, step size). Intended to be + called *before* tuning, so the deep-copied adaptation state starts from + its initial values. + """ + # Detach the compiled function (and the integrator that references it) + # before deep-copying: `copy.deepcopy` handles the pure-numpy adaptation + # state fine, but must not touch the pytensor Function (its composite + # deepcopy is unsafe for concurrent use). We re-attach an independent + # `fork()` of the function afterwards. + logp_dlogp_func = self._logp_dlogp_func + integrator = self.integrator + del self._logp_dlogp_func + del self.integrator + try: + new = copy.deepcopy(self) + finally: + self._logp_dlogp_func = logp_dlogp_func + self.integrator = integrator + + forked_func = logp_dlogp_func.fork() + new._logp_dlogp_func = forked_func + # Rewire the shared-input dict to the fork's fresh shared variables so + # that `ArrayStepShared.step` sets extra values on this chain's copies. + new.shared = dict(forked_func._extra_vars_shared) + new.integrator = integration.CpuLeapfrogIntegrator(new.potential, forked_func) + return new diff --git a/scripts/check_all_tests_are_covered.py b/scripts/check_all_tests_are_covered.py index 59e5bec157..c10f35c063 100644 --- a/scripts/check_all_tests_are_covered.py +++ b/scripts/check_all_tests_are_covered.py @@ -41,7 +41,10 @@ def from_yaml(): wfname = wf.rstrip(".yml") wfdef = yaml.safe_load(open(Path(".github", "workflows", wf))) for jobname, jobdef in wfdef["jobs"].items(): - if jobname in ("float32", "all_tests"): + # `free_threading` runs a small curated subset (already covered by the + # other jobs) on a no-GIL interpreter and has no `linker` matrix, so it + # is excluded from the once-per-OS/linker coverage accounting. + if jobname in ("float32", "all_tests", "free_threading"): continue runs_on = jobdef.get("runs-on", "unknown") floatX = "float32" if jobname == "float32" else "float64" diff --git a/setup.py b/setup.py index ee4024a6ba..21265de633 100755 --- a/setup.py +++ b/setup.py @@ -33,6 +33,7 @@ "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", "Programming Language :: Python :: 3.14", + "Programming Language :: Python :: Free Threading :: 1 - Unstable", "License :: OSI Approved :: Apache Software License", "Intended Audience :: Science/Research", "Topic :: Scientific/Engineering", diff --git a/tests/conftest.py b/tests/conftest.py index 7ded384563..7c6a9c4dbc 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -11,6 +11,8 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import sys +import sysconfig import warnings import numpy as np @@ -20,6 +22,20 @@ from numba.core.errors import NumbaPerformanceWarning, NumbaWarning +def pytest_sessionfinish(session, exitstatus): + # On a free-threaded build, importing a C extension that has not opted into + # free-threading (via the Py_mod_gil slot) makes CPython re-enable the GIL for + # the rest of the process. PyTensor's default (numba) backend must not do this. + # Checked after the whole session so it is independent of collection order. + if sysconfig.get_config_var("Py_GIL_DISABLED") == 1 and sys._is_gil_enabled(): + print( # noqa: T201 + "\nERROR: the GIL was re-enabled during the test session; a C extension " + "lacking free-threading support (missing Py_MOD_GIL_NOT_USED) was imported.", + file=sys.stderr, + ) + session.exitstatus = pytest.ExitCode.TESTS_FAILED + + @pytest.fixture(scope="function", autouse=True) def pytensor_config(): config = pytensor.config.change_flags(on_opt_error="raise") diff --git a/tests/sampling/test_parallel.py b/tests/sampling/test_parallel.py index 84e816054e..12d7c62a83 100644 --- a/tests/sampling/test_parallel.py +++ b/tests/sampling/test_parallel.py @@ -15,6 +15,7 @@ import os import platform import sys +import sysconfig import warnings import cloudpickle @@ -27,6 +28,7 @@ from pytensor.tensor.type import TensorType import pymc as pm +import pymc.sampling.mcmc as mcmc import pymc.sampling.parallel as ps from pymc.pytensorf import floatX @@ -285,3 +287,48 @@ def test_sampling_with_random_generator_matches(cores): post2 = pm.sample(random_seed=np.random.default_rng(42), **kwargs).posterior assert post1.equals(post2), (post1["x"].mean().item(), post2["x"].mean().item()) + + +@pytest.mark.skipif( + sysconfig.get_config_var("Py_GIL_DISABLED") != 1, + reason="requires a free-threaded (no-GIL) build", +) +def test_thread_parallel_sampling_matches_sequential(monkeypatch): + # On a free-threaded build `mp_ctx="thread"` runs chains as threads (the path this + # test targets); `nuts_sampler="pymc"` keeps it on PyMC's stepper instead of + # dispatching to nutpie. The conftest `pytest_sessionfinish` hook separately guards + # that no import re-enabled the GIL during the run. + # + # Spy on `_thread_sample`, recording success only *after* it returns, so the threaded + # path is proven to have run and not silently fallen back to multiprocessing (which + # would produce the same draws and hide a regression). + thread_sampled = False + real_thread_sample = mcmc._thread_sample + + def spy(**kwargs): + nonlocal thread_sampled + result = real_thread_sample(**kwargs) + thread_sampled = True + return result + + monkeypatch.setattr(mcmc, "_thread_sample", spy) + + kwargs = { + "draws": 50, + "tune": 50, + "chains": 2, + "nuts_sampler": "pymc", + "progressbar": False, + "compute_convergence_checks": False, + "random_seed": 1, + } + with pm.Model(): + pm.Normal("x", 0.0, 1.0) + pm.Normal("y", pm.Normal("mu", 0.0, 1.0), 1.0, shape=3) + threaded = pm.sample(cores=2, mp_ctx="thread", **kwargs) + sequential = pm.sample(cores=1, **kwargs) + + assert thread_sampled, "sampling did not run through the threaded path" + # Same seed => each chain's RNG is seeded once and identically regardless of + # execution mode, so threaded sampling reproduces sequential draws bit-for-bit. + assert threaded.posterior.equals(sequential.posterior) diff --git a/tests/step_methods/test_compound.py b/tests/step_methods/test_compound.py index 4452f80d0d..136dbaa7ca 100644 --- a/tests/step_methods/test_compound.py +++ b/tests/step_methods/test_compound.py @@ -12,11 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. +import numpy as np import pytensor import pytest import pymc as pm +from pymc.model import modelcontext from pymc.step_methods import ( NUTS, CompoundStep, @@ -26,16 +28,95 @@ Slice, ) from pymc.step_methods.compound import ( + BlockedStep, StatsBijection, flatten_steps, get_stats_dtypes_shapes_from_steps, infer_warn_stats_info, ) from pymc.testing import fast_unstable_sampling_mode +from pymc.util import get_random_generator, get_value_vars_from_user_vars from tests.helpers import StepMethodTester from tests.models import simple_2model_continuous +class NonForkableStep(BlockedStep): + """A minimal custom step method that samples but never implements ``fork``. + + Stands in for a third-party / user-defined sampler: it can drive ``pm.sample`` + (a plain random-walk Metropolis), but because it does not override ``fork`` it + hits the ``BlockedStep.fork`` default that raises ``NotImplementedError`` -- + which is how threaded sampling detects it must fall back to multiprocessing. + """ + + name = "nonforkable" + stats_dtypes_shapes: dict = {} + + def __init__(self, vars, model=None, rng=None): + model = modelcontext(model) + self.vars = get_value_vars_from_user_vars(vars, model) + self.var_names = [v.name for v in self.vars] + self._logp = model.compile_logp() + self.rng = get_random_generator(rng) + self.tune = False + + def set_rng(self, rng): + self.rng = get_random_generator(rng, copy=False) + + def step(self, point): + proposal = dict(point) + for name in self.var_names: + proposal[name] = point[name] + self.rng.normal(scale=0.3, size=np.shape(point[name])) + if np.log(self.rng.uniform()) < float(self._logp(proposal)) - float(self._logp(point)): + return proposal, [] + return dict(point), [] + + +def test_fork_fallback_for_non_forkable_step(monkeypatch, caplog): + """A step method that never implements ``fork`` must fail loudly so threaded + sampling falls back to multiprocessing instead of running an unsafe shared + step across threads. + + Covers: the base ``fork`` raising ``NotImplementedError``, ``CompoundStep`` + propagating it (rather than forking only some children), and ``pm.sample`` + completing via the process path when the threaded branch is forced on. + """ + import logging + + # The step itself, and any CompoundStep containing it, refuse to fork. + with pm.Model(): + x = pm.Normal("x", 0, 1) + y = pm.Normal("y", 0, 1) + with pytest.raises(NotImplementedError, match="NonForkableStep"): + NonForkableStep([x]).fork(np.random.default_rng(0)) + with pytensor.config.change_flags(mode=fast_unstable_sampling_mode): + compound = CompoundStep([NUTS([y]), NonForkableStep([x])]) + with pytest.raises(NotImplementedError, match="NonForkableStep"): + compound.fork(np.random.default_rng(0)) + + # Force the free-threaded dispatch branch regardless of the actual build; + # sampling must still complete by falling back to multiprocessing. + # `sys._is_gil_enabled` only exists on Python 3.13+; create it if absent so + # the free-threaded branch is exercised on GIL builds too. + monkeypatch.setattr("pymc.sampling.mcmc.sys._is_gil_enabled", lambda: False, raising=False) + with pm.Model(): + x = pm.Normal("x", 2.0, 1.0) + step = NonForkableStep([x]) + with caplog.at_level(logging.WARNING, logger="pymc"): + idata = pm.sample( + step=step, + draws=20, + tune=20, + chains=2, + cores=2, + progressbar=False, + compute_convergence_checks=False, + random_seed=1, + ) + assert "does not support threaded sampling" in caplog.text + assert idata.posterior["x"].shape == (2, 20) + + def test_stepmethods_do_not_require_tune_stat(): step_types = pm.step_methods.STEP_METHODS assert len(step_types) > 5