diff --git a/autofit/non_linear/search/nest/nautilus/search.py b/autofit/non_linear/search/nest/nautilus/search.py index 4ae265cc6..2cb9dce56 100644 --- a/autofit/non_linear/search/nest/nautilus/search.py +++ b/autofit/non_linear/search/nest/nautilus/search.py @@ -2,6 +2,7 @@ import logging import os import sys +from contextlib import nullcontext from pathlib import Path from typing import Dict, Optional, Tuple @@ -327,7 +328,18 @@ def fit_multiprocessing(self, fitness, model, analysis): # 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: + # + # For a single core no pool may be created at all: nautilus treats + # pool=None (and the int 1) as fully serial, whereas a Pool(1) object + # forces every likelihood call into a forked worker — which deadlocks + # in XLA compilation when the likelihood touches JAX, since a forked + # child of a JAX-initialized parent cannot compile. + if self.number_of_cores <= 1: + pool_context = nullcontext(None) + else: + pool_context = fork_context().Pool(self.number_of_cores) + + with pool_context as pool: search_internal = self.sampler_cls( prior=PriorVectorized(model=model), likelihood=fitness.call_wrap, diff --git a/test_autofit/non_linear/search/nest/test_nautilus.py b/test_autofit/non_linear/search/nest/test_nautilus.py index 5ac5ff357..ca217876b 100644 --- a/test_autofit/non_linear/search/nest/test_nautilus.py +++ b/test_autofit/non_linear/search/nest/test_nautilus.py @@ -1,3 +1,4 @@ +import numpy as np import pytest import autofit as af @@ -51,3 +52,36 @@ def test__test_mode(): search.apply_test_mode() assert search.n_like_max == 1 + + +def test__single_core_builds_no_pool(monkeypatch): + """ + number_of_cores=1 must not construct a multiprocessing pool: nautilus + treats pool=None as fully serial, whereas a Pool(1) object forces every + likelihood call into a forked worker — which deadlocks in XLA compilation + when the likelihood touches JAX (#1442). + """ + from autofit.non_linear.search.nest.nautilus import search as nautilus_search + + def no_fork_context(): + raise AssertionError( + "fork_context must not be used when number_of_cores == 1" + ) + + monkeypatch.setattr(nautilus_search, "fork_context", no_fork_context) + monkeypatch.setenv("PYAUTO_TEST_MODE", "1") + + model = af.Model(af.ex.Gaussian) + analysis = af.ex.Analysis( + data=np.full(100, 5.0), + noise_map=np.full(100, 1.0), + ) + + search = af.Nautilus( + name="nautilus_single_core", + unique_tag="single_core_no_pool_test", + n_live=10, + number_of_cores=1, + ) + + search.fit(model=model, analysis=analysis)