refactor: NSS onto Fitness, trace preflight and x64 check (search-extensibility A3b, stacked on A3)

autofit/non_linear/search/abstract_search.py

@@ -402,6 +402,9 @@ def __init__(
     test_mode_budget = {}
     objective_target = None
     invalid_value = -float("inf")
+    # True for a search whose backend is wrong in float32: `fit` then raises, rather
+    # than warns, when a JAX analysis runs with `jax_enable_x64` off (A3b preflight).
+    requires_fp64 = False
 
     def optimise(
         self,
@@ -899,6 +902,16 @@ def start_resume_fit(self, analysis: Analysis, model: AbstractPriorModel) -> Res
         # smoke runs use `PYAUTO_TEST_MODE>=2`, which returned above.
         cap.check_jax_required(self, analysis)
 
+        # The JAX checks (search-extensibility A3b, `autofit.non_linear.search.preflight`),
+        # after the test-mode bypass and the REQUIRED gate and before any backend state
+        # exists: warn once when x64 is off (raise for `requires_fp64`), then trace the
+        # declared objective kind(s) once with `jax.eval_shape` so a non-traceable
+        # likelihood fails here, naming the search, rather than inside the backend.
+        from autofit.non_linear.search import preflight
+
+        preflight.check_x64(self, analysis)
+        preflight.trace_preflight(self, analysis, model)
+
         # The fit's pools, with the one JAX fork rule applied once (one INFO line when it
         # downgrades) and the effective worker count recorded for `search.summary`.
         self._fit_pools = PoolFactory(self, analysis)
@@ -1458,6 +1471,10 @@ def _fit(self, model: AbstractPriorModel, analysis: Analysis):
 
         from autofit.non_linear.search.fit_context import FitContext
 
+        # The REQUIRED gate, repeated for callers that reach `_fit` without `fit`, before
+        # any `Fitness` or backend state exists (`fit` has already run it).
+        cap.check_jax_required(self, analysis)
+
         fitness = self.make_fitness(
             analysis=analysis, model=model, **self.fitness_overrides(analysis)
         )

autofit/non_linear/search/capabilities.py

@@ -32,6 +32,9 @@
 ``checkpointer``          Where the search archives its internal state when a fit
                           completes (``autofit.non_linear.checkpoint``): ``dill``,
                           ``pickle`` or ``native:<filename>``.
+``requires_fp64``         Whether the backend needs float64: ``fit`` then raises, rather
+                          than warns once, when a JAX analysis runs with
+                          ``jax_enable_x64`` off (``autofit.non_linear.search.preflight``).
 ========================  ==============================================================
 
 ``invalid_value`` records what the backend actually sees and is not yet normalised. For
@@ -176,6 +179,7 @@ def to_dict(self) -> Dict[str, str]:
     "objective_target",
     "invalid_value",
     "checkpointer",
+    "requires_fp64",
 )
 """
 The names of the static capability attributes, in manifest order.

autofit/non_linear/search/nest/dynesty/search/abstract.py

@@ -4,6 +4,7 @@
 import os
 import sys
 from abc import ABC
+from contextlib import ExitStack
 from pathlib import Path
 from typing import Dict, Optional, Tuple, Union, TYPE_CHECKING
 
@@ -55,6 +56,21 @@ def __enter__(self):
     return ForkPool
 
 
+class _SingleCoreRun(Exception):
+    """
+    Control flow inside ``AbstractDynesty._fit``: run dynesty without a pool.
+
+    Raised when the fit is single-core by construction (``number_of_cores <= 1``,
+    ``force_x1_cpu`` or a JAX analysis) and when the multiprocessing pool cannot be
+    created (``RuntimeError`` from its ``__enter__``, e.g. on platforms with weak
+    multiprocessing support), which is then its ``__cause__``. It is a private
+    ``Exception`` subclass, not a ``RuntimeError``, so the ``except`` that catches it
+    never swallows a ``RuntimeError`` raised while dynesty runs (a genuine dynesty
+    error, or JAX's ``XlaRuntimeError``, which subclasses ``RuntimeError``): those
+    propagate instead of silently restarting the run single-core.
+    """
+
+
 DYNESTY_FILENAME = "savestate.save"
 """
 dynesty's checkpoint file, the search's resume state.
@@ -234,8 +250,9 @@ def _fit(
         as on the existing `force_x1_cpu` path.
 
         However, certain operating systems (e.g. Windows) do not support Python multiprocessing particularly well.
-        This can cause Dynesty to crash when a pool is included. If this occurs (raising a `RunTimeException`)
-        a Dynesty object without a pool is created and used instead.
+        If creating the pool raises a `RuntimeError`, a Dynesty object without a pool is created and used instead.
+        Only the pool's creation falls back: a `RuntimeError` raised while Dynesty runs (including JAX's
+        `XlaRuntimeError`) propagates rather than silently restarting the run single-core (`_SingleCoreRun`).
 
         Parameters
         ----------
@@ -283,17 +300,26 @@ def _fit(
         while not finished:
             try:
                 if number_of_cores <= 1 or self.force_x1_cpu or analysis.is_jax:
-                    raise RuntimeError
+                    raise _SingleCoreRun()
 
                 Pool = _fork_pool_cls()
 
-                with Pool(
-                    njobs=number_of_cores,
-                    loglike=fitness,
-                    prior_transform=prior_transform,
-                    logl_args=(model, fitness),
-                    ptform_args=(model,),
-                ) as pool:
+                with ExitStack() as stack:
+                    # Only a failure to *create* the pool falls back to a single-core
+                    # run; a RuntimeError raised while dynesty runs propagates.
+                    try:
+                        pool = stack.enter_context(
+                            Pool(
+                                njobs=number_of_cores,
+                                loglike=fitness,
+                                prior_transform=prior_transform,
+                                logl_args=(model, fitness),
+                                ptform_args=(model,),
+                            )
+                        )
+                    except RuntimeError as error:
+                        raise _SingleCoreRun() from error
+
                     search_internal = self.search_internal_from(
                         model=model,
                         fitness=fitness,
@@ -306,7 +332,7 @@ def _fit(
 
                     checkpoint_exists = True
 
-            except RuntimeError as e:
+            except _SingleCoreRun as e:
                 if not checkpoint_exists:
                     if analysis.is_jax:
                         self.logger.info(
@@ -323,7 +349,7 @@ def _fit(
                     else:
                         self.logger.info(
                             f"""
-                            The Dynesty multiprocessing pool could not be created ({e!r}).
+                            The Dynesty multiprocessing pool could not be created ({e.__cause__!r}).
 
                             A single CPU non-multiprocessing Dynesty run is being performed.
                             """

autofit/non_linear/search/nest/nss/_chunked_nss.py

@@ -25,7 +25,7 @@
 When ``chunk_size`` is None the builder still uses ``jax.vmap`` and is
 bit-identical to upstream. The builder otherwise produces a
 ``blackjax.SamplingAlgorithm`` with the same shape ``blackjax.nss(...)``
-returns, so ``af.NSS._fit`` is a one-line switch.
+returns, so ``af.NSS.run`` is a one-line switch.
 
 See PyAutoFit#1304 for the diagnosis and A100 evidence (jobs 322605 /
 322606 OOM at the same byte counts as before #1303 landed, because the