fix: A2/A3/A3b post-merge fixes from astra review 05 (search-extensibility)

autofit/non_linear/fitness.py

@@ -960,12 +960,14 @@ def __setstate__(self, state):
         self.__dict__.setdefault("gradient_mode", None)
         # `Fitness` objects pickled before the objective factory (search-extensibility A2)
         # carry `use_jax_vmap` / `use_jax_jit` instead of `batched` / `compile`. Their
-        # dispatch maps onto the factory: vmap -> the batched objective; jit or not, a JAX
-        # scalar objective is now jitted lazily.
+        # dispatch is preserved exactly: vmap -> the (jitted) batched objective; jit -> the
+        # jitted scalar objective; neither -> the eager scalar objective. A pickle that
+        # predates both flags dispatched eagerly too, so it stays eager.
         legacy_vmap = self.__dict__.pop("use_jax_vmap", None)
-        self.__dict__.pop("use_jax_jit", None)
+        legacy_jit = self.__dict__.pop("use_jax_jit", None)
+        if "compile" not in self.__dict__:
+            self.__dict__["compile"] = bool(legacy_vmap) or bool(legacy_jit)
         self.__dict__.setdefault("batched", bool(legacy_vmap))
-        self.__dict__.setdefault("compile", True)
         self._objectives = {}
         self._call = self.objective(BATCHED if self.batched else SCALAR)
 

autofit/non_linear/parallel/pool.py

@@ -222,17 +222,26 @@ def sneaky(self, fitness):
         self._created.append(pool)
         return pool
 
-    def close(self):
+    def close(self, failed: bool = True):
         """
-        Terminate every pool this factory built, after a failed fit. A successful fit's
-        pools are closed by the search that used them, exactly as before.
+        Shut down every pool this factory built, once: ``close`` then ``join`` after a
+        successful fit (the workers finish and exit), ``terminate`` then ``join`` after a
+        failed one. Idempotent, and safe on a pool its backend already shut down.
         """
-        for pool in self._created:
-            terminate = getattr(pool, "terminate", None)
-            if terminate is None:
+        created, self._created = self._created, []
+        for pool in created:
+            for name in ("terminate",) if failed else ("close",):
+                method = getattr(pool, name, None)
+                if method is None:
+                    continue
+                try:
+                    method()
+                except Exception:
+                    logger.debug("Pool %s after a fit raised", name, exc_info=True)
+            join = getattr(pool, "join", None)
+            if join is None:
                 continue
             try:
-                terminate()
+                join()
             except Exception:
-                logger.debug("Pool termination after a failed fit raised", exc_info=True)
-        self._created = []
+                logger.debug("Pool join after a fit raised", exc_info=True)

autofit/non_linear/search/abstract_search.py

@@ -744,7 +744,11 @@ class represented by model M and gives a score for their fitness.
                     info=info,
                 )
 
-        if not self.paths.is_complete:
+        # `start_resume_fit` finalizes the archive itself, before it writes
+        # `.completed`; a completed fit's `post_fit_output` still tidies the folder.
+        resumed_or_started = not self.paths.is_complete
+
+        if resumed_or_started:
             result = self.start_resume_fit(
                 analysis=analysis,
                 model=model,
@@ -762,6 +766,7 @@ class represented by model M and gives a score for their fitness.
 
             self.post_fit_output(
                 search_internal=result.search_internal,
+                finalize_search_internal=not resumed_or_started,
             )
 
         gc.collect()
@@ -963,10 +968,25 @@ def start_resume_fit(self, analysis: Analysis, model: AbstractPriorModel) -> Res
 
         model.unfreeze()
 
-        self.paths.completed()
+        self._complete_fit(search_internal=search_internal)
 
         return result
 
+    def _complete_fit(self, search_internal):
+        """
+        Finalize the fit's retained internal state, then mark the fit complete.
+
+        The order is the transaction: ``.completed`` is written only once the archive
+        (and the discard of the resume state it supersedes) has succeeded. An archive
+        write that fails (a full disk) therefore leaves no marker, the resume state is
+        still on disk, and the next run resumes and finalizes again, rather than
+        loading a "completed" fit whose archive never landed.
+        """
+        if not skip_fit_output():
+            self.finalize_search_internal(search_internal=search_internal)
+
+        self.paths.completed()
+
     def _test_mode_valid_parameter_vector(
         self,
         model: AbstractPriorModel,
@@ -1192,7 +1212,21 @@ def result_via_completed_fit(
 
         return result
 
-    def post_fit_output(self, search_internal):
+    def finalize_search_internal(self, search_internal):
+        """
+        Archive ``search_internal`` (``output_search_internal``) when the search retains
+        its internal state, otherwise remove the ``search_internal`` folder.
+
+        Run before ``.completed`` is written (``_complete_fit``), so a failure here
+        leaves the fit resumable.
+        """
+        if not self.retains_search_internal:
+            self.logger.info("Removing search internal folder.")
+            self.paths.remove_search_internal()
+        elif search_internal is not None:
+            self.output_search_internal(search_internal=search_internal)
+
+    def post_fit_output(self, search_internal, finalize_search_internal: bool = True):
         """
         Cleans up the output folderds after a completed non-linear search.
 
@@ -1207,12 +1241,13 @@ def post_fit_output(self, search_internal):
         ----------
         search_internal
             The internal search.
+        finalize_search_internal
+            Whether to archive (or remove) the internal state here. ``False`` when the
+            fit just ran, since ``start_resume_fit`` finalized it before writing
+            ``.completed``.
         """
-        if not self.retains_search_internal:
-            self.logger.info("Removing search internal folder.")
-            self.paths.remove_search_internal()
-        elif search_internal is not None:
-            self.output_search_internal(search_internal=search_internal)
+        if finalize_search_internal:
+            self.finalize_search_internal(search_internal=search_internal)
 
         if not self.disable_output:
             self.logger.info("Removing all files except for .zip file")
@@ -1346,7 +1381,7 @@ def _fit_bypass_test_mode(
         # Mark the fit complete, exactly as start_resume_fit does — a bypassed
         # fit must be resumable (paths.is_complete -> result_via_completed_fit
         # on the next run), or every rerun re-bypasses the whole pipeline.
-        self.paths.completed()
+        self._complete_fit(search_internal=None)
 
         return result
 
@@ -1476,15 +1511,17 @@ def _fit(self, model: AbstractPriorModel, analysis: Analysis):
         cap.check_jax_required(self, analysis)
 
         fitness = self.make_fitness(
-            analysis=analysis, model=model, **self.fitness_overrides(analysis)
+            analysis=analysis,
+            model=model,
+            **self.execution_settings(analysis).fitness_overrides,
         )
 
         ctx = FitContext(
             search=self,
             model=model,
             analysis=analysis,
             fitness=fitness,
-            pool=self._pools(),
+            pool=self._pools(analysis),
             test_mode_level=test_mode_level(),
         )
 
@@ -1523,6 +1560,20 @@ def fitness_overrides(self, analysis: Analysis) -> dict:
         """
         return {}
 
+    def execution_settings(self, analysis: Analysis):
+        """
+        How this search executes its objective on ``analysis``, resolved once: batching,
+        eager or jitted execution, gradient use and the resolved gradient mode, plus the
+        ``Fitness`` keyword arguments the backend builds its ``Fitness`` with
+        (``autofit.non_linear.search.preflight.ExecutionSettings``).
+
+        The JAX trace preflight validates exactly this object and the backend builds its
+        ``Fitness`` from it, so the two cannot disagree.
+        """
+        from autofit.non_linear.search.preflight import ExecutionSettings
+
+        return ExecutionSettings.resolve(self, analysis)
+
     def make_fitness(
         self, analysis: Analysis, model: AbstractPriorModel, **overrides
     ) -> Fitness:
@@ -1897,7 +1948,7 @@ def start_points(
                 model=model,
                 fitness=fitness,
                 paths=self.paths,
-                n_cores=self._pools().number_of_cores,
+                n_cores=self._pools(fitness.analysis).number_of_cores,
                 **initializer_kwargs,
             )
         )
@@ -2003,13 +2054,14 @@ def load_search_internal(self):
         """
         return self.paths.load_search_internal()
 
-    def _pools(self) -> PoolFactory:
+    def _pools(self, analysis: Optional[Analysis] = None) -> PoolFactory:
         """
         The current fit's `PoolFactory` (the JAX fork rule applied to its analysis), or,
-        outside a fit, one built from `number_of_cores` alone.
+        outside a fit, one built for `analysis` (so a direct `_fit` caller still gets
+        the fork rule), or from `number_of_cores` alone when no analysis is given.
         """
         pools = self.__dict__.get("_fit_pools")
-        return pools if pools is not None else PoolFactory(self)
+        return pools if pools is not None else PoolFactory(self, analysis)
 
     @check_cores
     def make_pool(self):
@@ -2041,7 +2093,7 @@ def make_sneaky_pool(self, fitness: Fitness) -> Optional[SneakyPool]:
         An implementation of a multiprocessing pool
         """
 
-        return self._pools().sneaky(fitness)
+        return self._pools(getattr(fitness, "analysis", None)).sneaky(fitness)
 
     def __eq__(self, other):
         return isinstance(other, NonLinearSearch) and self.__dict__ == other.__dict__

autofit/non_linear/search/fit_context.py

@@ -121,7 +121,19 @@ def update(self, internal: Any):
 
         The backend is paused for the duration of the call (the update is synchronous),
         so ``internal`` is read as it stands at this point of the run.
+
+        The checkpoint is saved first: the search's resume state (``ctx.resume``'s
+        strategy, ``NonLinearSearch.resume_state``) gets ``save(internal)``, so an
+        interrupted run resumes from the last update. A native resume state (one the
+        backend writes itself) keeps its no-op ``save``; a search without a resume
+        state is not resumable and saves nothing until its archive is finalized at
+        completion. ``finalize`` is never called here, since that is the completion step
+        that may discard the resume state.
         """
+        resume_state = self._search.resume_state
+        if resume_state is not None:
+            resume_state.save(self.paths, internal)
+
         return self._search.perform_update(
             model=self.model,
             analysis=self._analysis,
@@ -132,15 +144,19 @@ def update(self, internal: Any):
 
     def close(self, failed: bool):
         """
-        Clean up after ``run``. On failure: close the pools the context built, shut down
-        the quick-update threads and release the compiled objectives. On success the
-        fitness is still needed for the final update, which shuts its threads down.
+        Clean up after ``run``, on every exit. The pools the context built are shut
+        down on both paths (closed and joined on success, terminated and joined on
+        failure; ``PoolFactory.close``). On failure the quick-update threads are also
+        shut down and the compiled objectives released; on success the fitness is still
+        needed for the final update, which shuts its threads down.
         """
+        close_pools = getattr(self.pool, "close", None)
+        if close_pools is not None:
+            close_pools(failed=failed)
+
         if not failed:
             return
 
-        self.pool.close()
-
         shutdown = getattr(self.fitness, "shutdown_quick_update", None)
         if shutdown is not None:
             shutdown()

autofit/non_linear/search/mle/multi_start_gradient/search.py

@@ -1003,10 +1003,15 @@ def _fit(
             else self.iterations_per_quick_update
         )
 
+        settings = self.execution_settings(analysis)
+
         fitness = self.make_fitness(
             analysis=analysis,
             model=model,
-            iterations_per_quick_update=iterations_per_quick_update,
+            **{
+                **settings.fitness_overrides,
+                "iterations_per_quick_update": iterations_per_quick_update,
+            },
         )
 
         # -2 * log_posterior, to MINIMIZE. value_and_grad over every start.
@@ -1031,7 +1036,7 @@ def _fit(
         # analysis's declaration and this search's override, and used for both
         # objectives below. Forward and reverse give the same ``(value, grad)``;
         # everything downstream of the two builders is mode-agnostic.
-        gradient_mode = self._resolved_gradient_mode(analysis)
+        gradient_mode = settings.gradient_mode
         if not self.silence:
             source = (
                 "search override"

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

@@ -211,6 +211,14 @@ def __init__(
 
         self.logger.debug(f"Creating {self.__class__.__name__} Search")
 
+    def fitness_overrides(self, analysis) -> dict:
+        """
+        A JAX analysis is jitted lazily by the scalar objective `__call__` dispatches to;
+        the deprecated `use_jax_jit=False` keeps its old meaning (eager JAX) for one
+        release.
+        """
+        return {"compile": getattr(self, "use_jax_jit", None) is not False}
+
     @property
     def search_kwargs(self) -> Dict:
         """Shared search kwargs passed to both Static and Dynamic Dynesty samplers."""
@@ -271,10 +279,7 @@ def _fit(
         fitness = self.make_fitness(
             analysis=analysis,
             model=model,
-            # A JAX analysis is jitted lazily by the scalar objective `__call__`
-            # dispatches to; the deprecated `use_jax_jit=False` keeps its old meaning
-            # (eager JAX) for one release.
-            compile=self.use_jax_jit is not False,
+            **self.execution_settings(analysis).fitness_overrides,
         )
 
         if not isinstance(self.paths, NullPaths):
@@ -293,7 +298,7 @@ def _fit(
 
         # The fit's effective core count: the JAX fork rule makes it 1 for a JAX
         # analysis (`autofit.non_linear.parallel.pool`).
-        number_of_cores = self._pools().number_of_cores
+        number_of_cores = self._pools(analysis).number_of_cores
 
         finished = False
 

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

@@ -536,18 +536,10 @@ def _safe_ess(log_w_mean):
             seed=int(self.seed),
         )
 
-        if checkpoint_path is not None and checkpoint_path.exists():
-            try:
-                self.resume_state.discard(ctx.paths)
-            except OSError as exc:
-                self.logger.warning(
-                    "Failed to delete completed-run checkpoint %s: %s. The "
-                    "next fresh af.NSS fit at this path will attempt to resume "
-                    "from it — delete manually if that is not desired.",
-                    checkpoint_path,
-                    exc,
-                )
-
+        # The checkpoint is kept until the result exists: `start_resume_fit` discards
+        # this resume state only after the archive is finalized and before it writes
+        # `.completed` (`NonLinearSearch._complete_fit`), so a failure while the result
+        # is processed leaves the run resumable.
         return search_internal
 
     @property