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)