Auto-pick non-fork mp_ctx when sampling with JAX

JAX is not fork-safe and can deadlock under multiprocessing's fork start
method. When a JAX backend is detected and the user has not specified an
`mp_ctx`, fall back to `forkserver` (or `spawn`). Warn instead of switch
if the user explicitly asked for fork.

Applies to both `pm.sample` and `pm.sample_smc`.

Author: ricardoV94

pymc/sampling/mcmc.py

@@ -827,7 +827,10 @@ def sample(
     if chains is None:
         chains = max(2, cores)
 
-    mp_ctx = _initialize_multiprocessing_context(mp_ctx, quiet=quiet)
+    compile_kwargs = resolve_backend_compile_kwargs(backend, compile_kwargs)
+    mp_ctx = _initialize_multiprocessing_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
     )
@@ -870,8 +873,6 @@ def sample(
         )
     )
 
-    compile_kwargs = resolve_backend_compile_kwargs(backend, compile_kwargs)
-
     if nuts_sampler is None:
         # Try to use nutpie by default if no setting is clearly at odds.
         # Requires all model variables, numba or jax preference,

pymc/sampling/parallel.py

@@ -19,6 +19,7 @@
 import platform
 import time
 import traceback
+import warnings
 
 from collections import namedtuple
 from collections.abc import Sequence
@@ -28,6 +29,8 @@
 import cloudpickle
 import numpy as np
 
+from pytensor.compile import get_mode
+from pytensor.link.jax.linker import JAXLinker
 from rich.theme import Theme
 from threadpoolctl import threadpool_limits
 
@@ -78,8 +81,14 @@ def rebuild_exc(exc, tb):
 
 
 def _initialize_multiprocessing_context(
-    mp_ctx: str | multiprocessing.context.BaseContext | None, quiet: bool = False
+    mp_ctx: str | multiprocessing.context.BaseContext | None,
+    *,
+    mode=None,
+    quiet: bool = False,
 ) -> multiprocessing.context.BaseContext:
+    user_specified = mp_ctx is not None
+    jax_mode = mode is not None and isinstance(get_mode(mode).linker, JAXLinker)
+
     if mp_ctx is None or isinstance(mp_ctx, str):
         # Closes issue https://github.com/pymc-devs/pymc/issues/3849
         # Related issue https://github.com/pymc-devs/pymc/issues/5339
@@ -96,6 +105,21 @@ def _initialize_multiprocessing_context(
 
         mp_ctx = multiprocessing.get_context(mp_ctx)
 
+    if jax_mode and mp_ctx.get_start_method() == "fork":
+        if user_specified:
+            warnings.warn(
+                "Using a JAX backend with multiprocessing start method 'fork' is unsafe "
+                "and may deadlock. Consider passing `mp_ctx='forkserver'` or `mp_ctx='spawn'`.",
+                UserWarning,
+                stacklevel=2,
+            )
+        else:
+            # JAX is not fork-safe: pick a non-fork default when user didn't specify.
+            new_method = (
+                "forkserver" if "forkserver" in multiprocessing.get_all_start_methods() else "spawn"
+            )
+            mp_ctx = multiprocessing.get_context(new_method)
+
     return mp_ctx
 
 

pymc/smc/sampling.py

@@ -179,15 +179,16 @@ def sample_smc(
     else:
         cores = min(chains, cores)
 
-    kernel_kwargs["compile_kwargs"] = resolve_backend_compile_kwargs(backend, compile_kwargs)
+    compile_kwargs = resolve_backend_compile_kwargs(backend, compile_kwargs)
+    kernel_kwargs["compile_kwargs"] = compile_kwargs
 
     random_seed = _get_seeds_per_chain(random_state=random_seed, chains=chains)
 
     model = modelcontext(model)
 
     logger.info("Initializing SMC sampler...")
 
-    mp_ctx = _initialize_multiprocessing_context(mp_ctx)
+    mp_ctx = _initialize_multiprocessing_context(mp_ctx, mode=compile_kwargs.get("mode"))
     joined_blas_limiter, cores, num_blas_cores_per_worker = setup_cores_blas_cores(
         blas_cores, chains, cores, mp_ctx
     )

tests/sampling/test_parallel.py

@@ -41,6 +41,42 @@ def test_context():
             pm.sample(tune=2, draws=2, chains=2, cores=2, mp_ctx=ctx)
 
 
+class TestMpCtxJaxSwitch:
+    def test_switches_default_away_from_fork_under_jax(self):
+        if ps._initialize_multiprocessing_context(None).get_start_method() != "fork":
+            pytest.skip("platform default is not fork")
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter("always")
+            ctx = ps._initialize_multiprocessing_context(None, mode="JAX")
+        assert ctx.get_start_method() != "fork"
+        assert not any("JAX backend" in str(x.message) for x in w)
+
+    def test_warns_when_user_explicitly_picks_fork_under_jax(self):
+        if "fork" not in multiprocessing.get_all_start_methods():
+            pytest.skip("fork start method not available on this platform")
+        with pytest.warns(UserWarning, match="JAX backend"):
+            ctx = ps._initialize_multiprocessing_context("fork", mode="JAX")
+        assert ctx.get_start_method() == "fork"
+
+    def test_leaves_non_jax_default_alone(self):
+        expected = ps._initialize_multiprocessing_context(None).get_start_method()
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter("always")
+            ctx = ps._initialize_multiprocessing_context(None, mode="NUMBA")
+        assert ctx.get_start_method() == expected
+        assert not any("JAX backend" in str(x.message) for x in w)
+
+    @pytest.mark.parametrize("explicit", ["spawn", "forkserver"])
+    def test_no_warn_when_user_picks_safe_method(self, explicit):
+        if explicit not in multiprocessing.get_all_start_methods():
+            pytest.skip(f"{explicit} start method not available on this platform")
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter("always")
+            ctx = ps._initialize_multiprocessing_context(explicit, mode="JAX")
+        assert ctx.get_start_method() == explicit
+        assert not any("JAX backend" in str(x.message) for x in w)
+
+
 class NoUnpickle:
     def __getstate__(self):
         return self.__dict__.copy()