From e6279c53fc2a231ba47616249d60fe8ac754699c Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 31 Jul 2026 15:34:05 +0100 Subject: [PATCH] fix: pin multiprocessing pools to the fork start method Python 3.14 changed the default start method on Linux from "fork" to "forkserver", which corrupts the model instances dynesty pool workers receive (factor-graph fits fail with "'Gaussian' object is not iterable") and silently degrades other pool creation to single-CPU fallbacks. A new autofit.non_linear.parallel.fork_context helper reproduces the pre-3.14 default on every platform (fork on POSIX except macOS, platform default on macOS/Windows) and is applied at every pool/process creation site: the dynesty pool (via a subclass, upstream hardcodes the default context), make_pool, SneakyPool/SneakierPool, the Process/Queue layer, the parallel EP optimiser, and nautilus (passed a pool object instead of an int so it no longer builds its own default-context pool). The dynesty single-CPU fallback log now includes the exception that caused it instead of unconditionally blaming the operating system. Fixes #1437 Co-Authored-By: Claude Fable 5 --- .../expectation_propagation/optimiser.py | 3 +- autofit/non_linear/parallel/__init__.py | 1 + autofit/non_linear/parallel/context.py | 33 ++++++ autofit/non_linear/parallel/process.py | 8 +- autofit/non_linear/parallel/sneaky.py | 5 +- autofit/non_linear/search/abstract_search.py | 4 +- .../search/nest/dynesty/search/abstract.py | 39 ++++++- .../non_linear/search/nest/nautilus/search.py | 57 +++++----- test_autofit/non_linear/test_fork_context.py | 102 ++++++++++++++++++ 9 files changed, 214 insertions(+), 38 deletions(-) create mode 100644 autofit/non_linear/parallel/context.py create mode 100644 test_autofit/non_linear/test_fork_context.py diff --git a/autofit/graphical/expectation_propagation/optimiser.py b/autofit/graphical/expectation_propagation/optimiser.py index bfbe292d9..3370283a1 100644 --- a/autofit/graphical/expectation_propagation/optimiser.py +++ b/autofit/graphical/expectation_propagation/optimiser.py @@ -13,6 +13,7 @@ from autofit.graphical.mean_field import Status, MeanField, FactorApproximation from autofit.graphical.utils import StatusFlag, LogWarnings from autofit.mapper.identifier import Identifier +from autofit.non_linear.parallel import fork_context from autofit.non_linear.paths import DirectoryPaths from autofit.non_linear.paths.abstract import AbstractPaths from autofit.tools.util import IntervalCounter @@ -460,7 +461,7 @@ def __init__( paths=paths, updater=updater, ) - self.pool = multiprocessing.Pool(n_cores - 1) + self.pool = fork_context().Pool(n_cores - 1) def run( self, diff --git a/autofit/non_linear/parallel/__init__.py b/autofit/non_linear/parallel/__init__.py index ade7cf9db..a856b7989 100644 --- a/autofit/non_linear/parallel/__init__.py +++ b/autofit/non_linear/parallel/__init__.py @@ -1,3 +1,4 @@ +from .context import fork_context from .process import AbstractJob from .process import AbstractJobResult from .process import Process diff --git a/autofit/non_linear/parallel/context.py b/autofit/non_linear/parallel/context.py new file mode 100644 index 000000000..781e6e4dc --- /dev/null +++ b/autofit/non_linear/parallel/context.py @@ -0,0 +1,33 @@ +import multiprocessing +import sys + + +def fork_context(): + """ + The multiprocessing context used for every pool and process PyAutoFit + creates. + + Python 3.14 changed the default start method on Linux (and other POSIX + platforms except macOS) from "fork" to "forkserver". PyAutoFit's + parallelism relies on fork semantics: fitness functions, models and the + state of user scripts (which run at module level without an + ``if __name__ == "__main__"`` guard) are inherited by worker processes + rather than pickled or re-imported. Under "forkserver" workers receive + corrupted model instances + (https://github.com/PyAutoLabs/PyAutoFit/issues/1437), so the "fork" + context is pinned explicitly. + + macOS is deliberately excluded: its default has been "spawn" since Python + 3.8 and forking a process whose threads hold ObjC/CoreFoundation state can + abort, so pinning "fork" there would introduce new behaviour rather than + restore old behaviour. This helper reproduces the pre-3.14 default on + every platform. + + Returns + ------- + The "fork" multiprocessing context on POSIX platforms other than macOS, + else the platform default ("spawn" on Windows and macOS). + """ + if sys.platform != "darwin" and "fork" in multiprocessing.get_all_start_methods(): + return multiprocessing.get_context("fork") + return multiprocessing.get_context() diff --git a/autofit/non_linear/parallel/process.py b/autofit/non_linear/parallel/process.py index b034611a5..a2721b2ce 100644 --- a/autofit/non_linear/parallel/process.py +++ b/autofit/non_linear/parallel/process.py @@ -5,6 +5,8 @@ from itertools import count from typing import Iterable +from .context import fork_context + logger = logging.getLogger( __name__ ) @@ -42,7 +44,7 @@ def perform(self, *args): """ -class Process(multiprocessing.Process): +class Process(fork_context().Process): def __init__( self, @@ -71,7 +73,7 @@ def __init__( self.logger.info("created") self.job_queue = job_queue - self.queue = multiprocessing.Queue() + self.queue = fork_context().Queue() self.initializer = initializer self.initargs = initargs @@ -146,7 +148,7 @@ def run_jobs( "The number of cores available must be at least 2 for parallel to run" ) - job_queue = multiprocessing.Queue() + job_queue = fork_context().Queue() processes = [ cls( diff --git a/autofit/non_linear/parallel/sneaky.py b/autofit/non_linear/parallel/sneaky.py index ad9682126..78c8ca567 100644 --- a/autofit/non_linear/parallel/sneaky.py +++ b/autofit/non_linear/parallel/sneaky.py @@ -8,6 +8,7 @@ from autonerves import conf from autofit.non_linear.paths.abstract import AbstractPaths +from .context import fork_context from .process import AbstractJob, Process, StopCommand logger = logging.getLogger(__name__) @@ -140,7 +141,7 @@ def __init__( super().__init__( name, - job_queue=mp.Queue(), + job_queue=fork_context().Queue(), initializer=initializer, initargs=initargs, job_args=job_args, @@ -391,7 +392,7 @@ def __enter__(self): logger.info("... using multiprocessing") - self.pool = mp.Pool( + self.pool = fork_context().Pool( processes=self.processes, ) diff --git a/autofit/non_linear/search/abstract_search.py b/autofit/non_linear/search/abstract_search.py index 819a97444..fcb9e45af 100644 --- a/autofit/non_linear/search/abstract_search.py +++ b/autofit/non_linear/search/abstract_search.py @@ -35,7 +35,7 @@ from autofit.mapper.model import ModelInstance from autofit.non_linear.initializer import Initializer from autofit.non_linear.fitness import Fitness -from autofit.non_linear.parallel import SneakyPool, SneakierPool +from autofit.non_linear.parallel import SneakyPool, SneakierPool, fork_context from autofit.non_linear.paths.abstract import AbstractPaths from autofit.non_linear.paths.database import DatabasePaths from autofit.non_linear.paths.directory import DirectoryPaths @@ -1316,7 +1316,7 @@ def make_pool(self): identify a 'master core' (the one whose id value is lowest) which handles model result output, visualization, etc.""" self.logger.info("...using pool") - return mp.Pool(processes=self.number_of_cores) + return fork_context().Pool(processes=self.number_of_cores) @check_cores def make_sneaky_pool(self, fitness: Fitness) -> Optional[SneakyPool]: diff --git a/autofit/non_linear/search/nest/dynesty/search/abstract.py b/autofit/non_linear/search/nest/dynesty/search/abstract.py index f7c54db8f..91e6301f9 100644 --- a/autofit/non_linear/search/nest/dynesty/search/abstract.py +++ b/autofit/non_linear/search/nest/dynesty/search/abstract.py @@ -20,6 +20,37 @@ logger = logging.getLogger(__name__) +def _fork_pool_cls(): + """ + dynesty's `Pool` pinned to the "fork" start method via + `autofit.non_linear.parallel.fork_context` — upstream hardcodes the default + multiprocessing context, which Python 3.14 changed to "forkserver" on Linux + (see that helper's docstring). `__enter__` mirrors `dynesty.pool.Pool.__enter__` + exactly apart from the context. + """ + from dynesty import pool as dynesty_pool + + from autofit.non_linear.parallel import fork_context + + class ForkPool(dynesty_pool.Pool): + def __enter__(self): + initargs = ( + self.loglike_0, + self.prior_transform_0, + self.logl_args or (), + self.logl_kwargs or {}, + self.ptform_args or (), + self.ptform_kwargs or {}, + ) + self.pool = fork_context().Pool( + self.njobs, dynesty_pool.initializer, initargs + ) + dynesty_pool.initializer(*initargs) + return self + + return ForkPool + + def prior_transform(cube, model): phys_cube = model.vector_from_unit_vector( unit_vector=cube, @@ -209,7 +240,7 @@ def _fit( if self.force_x1_cpu or analysis._use_jax: raise RuntimeError - from dynesty.pool import Pool + Pool = _fork_pool_cls() with Pool( njobs=self.number_of_cores, @@ -230,7 +261,7 @@ def _fit( checkpoint_exists = True - except RuntimeError: + except RuntimeError as e: if not checkpoint_exists: if getattr(analysis, "_use_jax", False): self.logger.info( @@ -242,8 +273,8 @@ def _fit( ) else: self.logger.info( - """ - Your operating system does not support Python multiprocessing. + f""" + The Dynesty multiprocessing pool could not be created ({e!r}). A single CPU non-multiprocessing Dynesty run is being performed. """ diff --git a/autofit/non_linear/search/nest/nautilus/search.py b/autofit/non_linear/search/nest/nautilus/search.py index 24ca5b490..4ae265cc6 100644 --- a/autofit/non_linear/search/nest/nautilus/search.py +++ b/autofit/non_linear/search/nest/nautilus/search.py @@ -10,6 +10,7 @@ from autofit.mapper.prior_model.abstract import AbstractPriorModel from autofit.mapper.prior.vectorized import PriorVectorized from autofit.non_linear.fitness import Fitness +from autofit.non_linear.parallel import fork_context from autofit.non_linear.paths.null import NullPaths from autofit.non_linear.search.nest import abstract_nest from autofit.non_linear.samples.sample import Sample @@ -323,34 +324,38 @@ def fit_multiprocessing(self, fitness, model, analysis): Contains the data and the log likelihood function which fits an instance of the model to the data, returning the log likelihood the search maximizes. """ - search_internal = self.sampler_cls( - prior=PriorVectorized(model=model), - likelihood=fitness.call_wrap, - n_dim=model.prior_count, - filepath=self.checkpoint_file, - pool=self.number_of_cores, - n_live=self.n_live, - n_update=self.n_update, - enlarge_per_dim=self.enlarge_per_dim, - n_points_min=self.n_points_min, - split_threshold=self.split_threshold, - n_networks=self.n_networks, - n_batch=self.n_batch, - n_like_new_bound=self.n_like_new_bound, - vectorized=self.vectorized, - seed=self.seed, - ) + # A pool object is passed rather than pool= so the pool uses the + # "fork" start method (see autofit.non_linear.parallel.fork_context) — + # nautilus builds its internal pools from the default context. + with fork_context().Pool(self.number_of_cores) as pool: + search_internal = self.sampler_cls( + prior=PriorVectorized(model=model), + likelihood=fitness.call_wrap, + n_dim=model.prior_count, + filepath=self.checkpoint_file, + pool=pool, + n_live=self.n_live, + n_update=self.n_update, + enlarge_per_dim=self.enlarge_per_dim, + n_points_min=self.n_points_min, + split_threshold=self.split_threshold, + n_networks=self.n_networks, + n_batch=self.n_batch, + n_like_new_bound=self.n_like_new_bound, + vectorized=self.vectorized, + seed=self.seed, + ) - search_internal = self.call_search( - search_internal=search_internal, - model=model, - analysis=analysis, - fitness=fitness - ) + search_internal = self.call_search( + search_internal=search_internal, + model=model, + analysis=analysis, + fitness=fitness + ) - # Nautilus creates its own multiprocessing.Pool internally when pool=N. - # Close them here so their finalizers don't fire at interpreter shutdown - # (after pickle has been torn down, causing AttributeError on Pool.__del__). + # Drop the pool references so their finalizers don't fire at interpreter + # shutdown (after pickle has been torn down, causing AttributeError on + # Pool.__del__). for pool_attr in ("pool_l", "pool_s"): pool = getattr(search_internal, pool_attr, None) if pool is not None: diff --git a/test_autofit/non_linear/test_fork_context.py b/test_autofit/non_linear/test_fork_context.py new file mode 100644 index 000000000..c276f0512 --- /dev/null +++ b/test_autofit/non_linear/test_fork_context.py @@ -0,0 +1,102 @@ +import multiprocessing +import sys + +import numpy as np +import pytest + +import autofit as af +from autofit.non_linear.fitness import Fitness +from autofit.non_linear.parallel import fork_context +from autofit.non_linear.parallel.process import Process +from autofit.non_linear.search.nest.dynesty.search.abstract import ( + _fork_pool_cls, + prior_transform, +) + +pins_fork = ( + sys.platform != "darwin" + and "fork" in multiprocessing.get_all_start_methods() +) + +requires_fork = pytest.mark.skipif( + not pins_fork, + reason="fork start method is not pinned on this platform", +) + + +@requires_fork +def test_fork_context_is_fork(): + assert fork_context().get_start_method() == "fork" + + +def test_fork_context_valid_everywhere(): + assert fork_context().get_start_method() in multiprocessing.get_all_start_methods() + + +@requires_fork +def test_process_class_is_fork_bound(): + assert Process._start_method == "fork" + + +@pytest.fixture(name="factor_graph_and_fitness") +def make_factor_graph_and_fitness(): + data = np.full(100, 5.0) + noise_map = np.full(100, 1.0) + + model = af.Collection( + gaussian=af.Model(af.ex.Gaussian), + exponential=af.Model(af.ex.Exponential), + ) + + analysis_factor_list = [ + af.AnalysisFactor( + prior_model=model.copy(), + analysis=af.ex.Analysis(data=data, noise_map=noise_map), + ) + for _ in range(2) + ] + + factor_graph = af.FactorGraphModel(*analysis_factor_list) + global_prior_model = factor_graph.global_prior_model + + fitness = Fitness( + model=global_prior_model, + analysis=factor_graph, + paths=None, + fom_is_log_likelihood=True, + resample_figure_of_merit=-1.0e99, + ) + + return factor_graph, global_prior_model, fitness + + +@requires_fork +def test_factor_graph_instance_shape_in_pool_worker(factor_graph_and_fitness): + """ + A pool worker must see the same nested instance structure as the main + process. On Python 3.14 the default start method became "forkserver", + under which workers received a flattened instance (a bare Gaussian in + place of the per-factor Collection) — the failure documented in issue + #1437. The fork-pinned dynesty pool preserves the structure. + """ + _, global_prior_model, fitness = factor_graph_and_fitness + + unit_vector = np.full(global_prior_model.prior_count, 0.5) + in_process_vector = global_prior_model.vector_from_unit_vector( + unit_vector=unit_vector + ) + in_process_likelihood = fitness(in_process_vector) + + ForkPool = _fork_pool_cls() + + with ForkPool( + njobs=2, + loglike=fitness, + prior_transform=prior_transform, + logl_args=(global_prior_model, fitness), + ptform_args=(global_prior_model,), + ) as pool: + worker_vector = list(pool.map(pool.prior_transform, [unit_vector]))[0] + worker_likelihood = list(pool.map(pool.loglike, [worker_vector]))[0] + + assert worker_likelihood == pytest.approx(in_process_likelihood, rel=1e-10)