diff --git a/autofit/non_linear/fitness.py b/autofit/non_linear/fitness.py index 8b8cfc313..f7631ee1a 100644 --- a/autofit/non_linear/fitness.py +++ b/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) diff --git a/autofit/non_linear/parallel/pool.py b/autofit/non_linear/parallel/pool.py index c9ad6cd65..05d5fb101 100644 --- a/autofit/non_linear/parallel/pool.py +++ b/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) diff --git a/autofit/non_linear/search/abstract_search.py b/autofit/non_linear/search/abstract_search.py index b5528b393..7823ec64f 100644 --- a/autofit/non_linear/search/abstract_search.py +++ b/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,7 +1511,9 @@ 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( @@ -1484,7 +1521,7 @@ def _fit(self, model: AbstractPriorModel, analysis: Analysis): 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__ diff --git a/autofit/non_linear/search/fit_context.py b/autofit/non_linear/search/fit_context.py index 9f2c29dd6..a557b7ad8 100644 --- a/autofit/non_linear/search/fit_context.py +++ b/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() diff --git a/autofit/non_linear/search/mle/multi_start_gradient/search.py b/autofit/non_linear/search/mle/multi_start_gradient/search.py index 490036fb9..044fad1b2 100644 --- a/autofit/non_linear/search/mle/multi_start_gradient/search.py +++ b/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" diff --git a/autofit/non_linear/search/nest/dynesty/search/abstract.py b/autofit/non_linear/search/nest/dynesty/search/abstract.py index f69232412..44965920e 100644 --- a/autofit/non_linear/search/nest/dynesty/search/abstract.py +++ b/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 diff --git a/autofit/non_linear/search/nest/nss/search.py b/autofit/non_linear/search/nest/nss/search.py index e2b5f5c3b..368f3a9b2 100644 --- a/autofit/non_linear/search/nest/nss/search.py +++ b/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 diff --git a/autofit/non_linear/search/preflight.py b/autofit/non_linear/search/preflight.py index 3b31e0b4d..ff10ed542 100644 --- a/autofit/non_linear/search/preflight.py +++ b/autofit/non_linear/search/preflight.py @@ -11,9 +11,13 @@ environment variable read at ``autonerves`` import, so importing ``jax`` first runs every likelihood in fp32 without a word. Warn once per process; raise for a search that declares ``requires_fp64 = True`` (none does today). -- ``trace_preflight``: one ``jax.eval_shape`` of the search's declared objective kind - (``"batched"`` for a ``batched`` search, else ``"scalar"``) at the prior-median vector, - plus one of its gradient kind when the search declares ``gradient == "uses"``. That is +- ``trace_preflight``: one ``jax.eval_shape`` of the objective kind the search actually + runs (``"batched"`` for a batched execution, else ``"scalar"``) at the prior-median + vector, plus one of its gradient kind (in the gradient mode the search resolves) when + the search declares ``gradient == "uses"``. The kinds and modes come from the search's + ``ExecutionSettings`` (``NonLinearSearch.execution_settings``), the same resolved object + the backend builds its ``Fitness`` from, so an eager (``compile=False``) scalar + execution is not traced at all and a search's ``gradient_mode`` override is honoured. That is one trace and no compile or execution, so it costs seconds at most where the backend's own first call would compile for minutes, and a non-traceable likelihood (an ``np.asarray`` or ``float()`` on a parameter-dependent value, a Python ``if`` on an @@ -36,6 +40,8 @@ import logging import os import warnings +from dataclasses import dataclass, field +from typing import Optional from autofit import exc from autofit.non_linear.objective import ( @@ -131,41 +137,127 @@ def check_x64(search, analysis): warnings.warn(X64_WARNING, UserWarning, stacklevel=3) -def preflight_kinds(search) -> tuple: +@dataclass(frozen=True) +class ExecutionSettings: """ - The objective kinds the trace preflight traces for ``search``: its declared - execution kind, plus the matching gradient kind when it declares - ``gradient == "uses"``. + How a search will execute its objective on one analysis, resolved once from the + search's capability declarations, its runtime options and the analysis: the one + object the trace preflight validates and the backend builds its ``Fitness`` from + (``NonLinearSearch.execution_settings``). + + Attributes + ---------- + batched + Whether the backend evaluates the batched (``jax.vmap``) objective. A search may + resolve this at runtime (``Nautilus(use_jax_vmap=False)`` runs the scalar one). + compile + Whether the scalar objective is jitted. ``False`` is eager execution (the + deprecated ``use_jax_jit=False``), which never traces the scalar likelihood. + uses_gradient + Whether the backend differentiates the objective (``gradient == "uses"``). + gradient_mode + The resolved AD mode: the search's ``gradient_mode`` override when it honours + one, else the analysis's declaration (``autofit.jax.gradient``). + fitness_overrides + The ``Fitness`` keyword arguments the backend passes to ``make_fitness``, + including the resolved ``compile`` / ``batched`` / ``gradient_mode`` where they + differ from the defaults. """ + + batched: bool + compile: bool + uses_gradient: bool + gradient_mode: str + fitness_overrides: dict = field(default_factory=dict) + + @property + def kinds(self) -> tuple: + """ + The objective kinds the backend evaluates: its execution kind, plus the matching + gradient kind when it differentiates. + """ + kinds = [BATCHED if self.batched else SCALAR] + if self.uses_gradient: + kinds.append(BATCHED_VALUE_AND_GRAD if self.batched else VALUE_AND_GRAD) + return tuple(kinds) + + @classmethod + def resolve(cls, search, analysis) -> "ExecutionSettings": + """ + Resolve ``search``'s execution on ``analysis``: its ``fitness_overrides`` + (runtime batching, eager execution, batch size, ...) over its class + declarations, and its ``gradient_mode`` override (when the search honours one) + over the analysis's declared mode. + """ + from autofit.jax.gradient import resolve_gradient_mode + + search_cls = type(search) + overrides = dict(search.fitness_overrides(analysis)) + + override_mode: Optional[str] = overrides.get("gradient_mode") + if override_mode is None and search_cls.honours_gradient_mode: + override_mode = getattr(search, "gradient_mode", None) + if override_mode is not None: + overrides["gradient_mode"] = override_mode + + return cls( + batched=bool(overrides.get("batched", search_cls.batched)), + compile=bool(overrides.get("compile", True)), + uses_gradient=search_cls.gradient == cap.Gradient.USES, + gradient_mode=resolve_gradient_mode(analysis, override=override_mode), + fitness_overrides=overrides, + ) + + +def preflight_kinds(search, analysis=None) -> tuple: + """ + The objective kinds the trace preflight traces for ``search``: the kinds of its + resolved ``ExecutionSettings`` on ``analysis``, or of its class declarations when + no analysis is given. + """ + if analysis is not None: + return search.execution_settings(analysis).kinds + cls = type(search) - kinds = [BATCHED if cls.batched else SCALAR] - if cls.gradient == cap.Gradient.USES: - kinds.append(BATCHED_VALUE_AND_GRAD if cls.batched else VALUE_AND_GRAD) - return tuple(kinds) + return ExecutionSettings( + batched=bool(cls.batched), + compile=True, + uses_gradient=cls.gradient == cap.Gradient.USES, + gradient_mode="reverse", + ).kinds -def _preflight_fitness(search, analysis, model): +def _preflight_fitness(search, analysis, model, settings: ExecutionSettings): """ - A throwaway ``Fitness`` in the search's declared convention, without paths (so no - resume sanity check evaluates the likelihood) and without quick updates. + A throwaway ``Fitness`` built from the search's resolved execution settings, without + paths (so no resume sanity check evaluates the likelihood) and without quick + updates. """ return search.make_fitness( analysis=analysis, model=model, - paths=None, - iterations_per_quick_update=None, - background_quick_update=False, - live_visual_update=False, + **{ + **settings.fitness_overrides, + "paths": None, + "iterations_per_quick_update": None, + "background_quick_update": False, + "live_visual_update": False, + }, ) def trace_preflight(search, analysis, model): """ - Trace the search's declared objective kind(s) once with ``jax.eval_shape`` at the + Trace the objective kind(s) the search resolves for ``analysis`` + (``NonLinearSearch.execution_settings``) once with ``jax.eval_shape`` at the prior-median vector, raising a ``SearchException`` that names the search and chains the tracer error when the likelihood cannot be traced. Then run the numerical probe if ``PYAUTO_JAX_PREFLIGHT_PROBE`` asks for it. + An eager scalar execution (``compile=False``) is never traced by the backend, so its + scalar kind is not traced here either; a gradient kind is, since differentiation + traces even an eager objective. + A no-op for a numpy analysis and when ``PYAUTO_JAX_PREFLIGHT=0``. """ if not analysis.is_jax or not _env_flag(PREFLIGHT_ENV, default=True): @@ -174,11 +266,15 @@ def trace_preflight(search, analysis, model): import jax import jax.numpy as jnp - fitness = _preflight_fitness(search, analysis, model) + settings = search.execution_settings(analysis) + + fitness = _preflight_fitness(search, analysis, model, settings) vector = jnp.asarray(model.physical_values_from_prior_medians) - for kind in preflight_kinds(search): + for kind in settings.kinds: + if kind == SCALAR and not settings.compile: + continue objective = fitness.objective(kind, compile=False) argument = vector[None, :] if kind in (BATCHED, BATCHED_VALUE_AND_GRAD) else vector try: diff --git a/docs/design/run_ctx.md b/docs/design/run_ctx.md index 2b30ff440..5c57a870d 100644 --- a/docs/design/run_ctx.md +++ b/docs/design/run_ctx.md @@ -136,9 +136,13 @@ pins the following details the table left open; each is flagged for review in th `Checkpointer` handle, not the loaded state, because a backend may reopen its own native file (`Nautilus` resumes through `Sampler(filepath=...)`). **`ctx.rng`** is a `SeedSequence` of the search's `seed` attribute, if any, until A4's search-level seed. -- **`ctx.close(failed)`** is the cleanup the lifecycle rules require; on failure it closes - the pools the context built, shuts the quick-update threads down and releases the - compiled objectives. +- **`ctx.close(failed)`** is the cleanup the lifecycle rules require. On every exit it + shuts down the pools the context built (closed and joined on success, terminated and + joined on failure); on failure it also shuts the quick-update threads down and + releases the compiled objectives. +- **`ctx.update(internal)`** saves the checkpoint before the update: the search's + `resume_state` gets `save(internal)` (a no-op for a native file the backend writes + itself), never `finalize`, which is the completion step. - **The bridge is `NonLinearSearch._fit`**: when a search defines `run`, the base `_fit` builds the `Fitness` and the context (after the gates `start_resume_fit` runs) and calls `run(ctx)`, so every caller of `_fit` (and the conformance suite's call counter) sees a diff --git a/test_autofit/non_linear/samples/test_golden_samples.py b/test_autofit/non_linear/samples/test_golden_samples.py index eb18f270f..4430e0e16 100644 --- a/test_autofit/non_linear/samples/test_golden_samples.py +++ b/test_autofit/non_linear/samples/test_golden_samples.py @@ -19,7 +19,15 @@ ``samples_info`` is compared exactly except for the convergence diagnostics BlackJAX itself computes (ESS, bulk/tail ESS, rank-normalised R-hat), whose estimators change -between BlackJAX releases; for those only the shape is pinned. +between BlackJAX releases. Those are pinned three ways, so a wrong value or wrong +wiring cannot pass on shape alone (review 05 finding 7): + +- by value against the fixture, to ``DIAGNOSTIC_RTOL`` (loose enough for an estimator + revision, far tighter than any wiring error), finite and positive; +- exactly against the installed BlackJAX's own calculation on the stored chains + (``test_diagnostics_are_the_installed_blackjax_reference``); +- by ownership, with the BlackJAX estimators replaced by deterministic doubles whose + outputs must be forwarded unchanged (``test_diagnostics_forward_the_backend_outputs``). """ import csv import importlib.util @@ -67,6 +75,32 @@ def reference_mode(monkeypatch): ) +#: Relative tolerance on the BlackJAX diagnostics against the fixture: room for an +#: estimator revision between BlackJAX releases, while a sentinel, a swapped key or a +#: per-parameter mix-up moves them by far more. +DIAGNOSTIC_RTOL = 0.25 + + +def assert_diagnostics_match(info: dict, expected: dict): + """ + Pop every ``BACKEND_DIAGNOSTIC_KEYS`` entry from ``info`` and ``expected`` and + compare them by value: the same shape, finite, positive (ESS and R-hat both are) and + within ``DIAGNOSTIC_RTOL`` of the fixture. + """ + for key in BACKEND_DIAGNOSTIC_KEYS: + if key not in expected: + continue + actual = np.asarray(info.pop(key), dtype=float) + reference = np.asarray(expected.pop(key), dtype=float) + + assert actual.shape == reference.shape, key + assert np.all(np.isfinite(actual)), key + assert np.all(actual > 0.0), key + np.testing.assert_allclose( + actual, reference, rtol=DIAGNOSTIC_RTOL, atol=0.0, err_msg=key + ) + + def _read_csv(path): with open(path, newline="") as f: rows = list(csv.reader(f)) @@ -101,8 +135,102 @@ def test_samples_info_is_unchanged(name): info = cases.samples_info_for_comparison(samples) - for key in BACKEND_DIAGNOSTIC_KEYS: - if key in expected: - assert np.shape(info.pop(key)) == np.shape(expected.pop(key)), key + assert_diagnostics_match(info, expected) assert info == expected + + +# Review 05 finding 7: the diagnostics must be protected by value, not only by shape. + + +def _sentinel(value): + return (np.full(np.shape(value), -999.0)).tolist() if np.ndim(value) else -999.0 + + +@pytest.mark.parametrize("key", BACKEND_DIAGNOSTIC_KEYS) +def test_sentinel_diagnostics_fail_the_golden_comparison(key, monkeypatch): + """ + Replacing any backend diagnostic with a shape-preserving ``-999`` must fail + ``test_samples_info_is_unchanged``. + """ + _skip_missing("nuts") + + original = cases.samples_info_for_comparison + + def mutated(samples): + info = original(samples) + info[key] = _sentinel(info[key]) + return info + + monkeypatch.setattr(cases, "samples_info_for_comparison", mutated) + + with pytest.raises(AssertionError): + test_samples_info_is_unchanged("nuts") + + +def test_diagnostics_are_the_installed_blackjax_reference(): + """ + The NUTS diagnostics equal the installed BlackJAX's own estimators run on the stored + chains, reshaped chain-major exactly as the search documents. + """ + _skip_missing("nuts") + import jax.numpy as jnp + from blackjax import diagnostics + + from autofit.non_linear.search.mcmc.blackjax.nuts.search import ( + split_chain_diagnostics, + ) + + internal = cases.internal_from("nuts") + info = cases.samples_info_for_comparison(cases.samples_from("nuts")) + + chain_major = jnp.asarray(split_chain_diagnostics(internal["positions"])) + + reference = { + "ess_per_param": diagnostics.effective_sample_size(chain_major), + "ess_bulk_per_param": diagnostics.ess_bulk(chain_major), + "ess_tail_per_param": diagnostics.ess_tail(chain_major), + "rhat_per_param": diagnostics.rhat(chain_major), + } + + for key, value in reference.items(): + value = np.atleast_1d(np.asarray(value)) + np.testing.assert_allclose(info[key], value, rtol=1e-12, atol=0.0, err_msg=key) + + assert info["ess_min"] == pytest.approx(min(info["ess_per_param"]), rel=1e-15) + assert info["ess_bulk_min"] == pytest.approx(min(info["ess_bulk_per_param"]), rel=1e-15) + assert info["ess_tail_min"] == pytest.approx(min(info["ess_tail_per_param"]), rel=1e-15) + assert info["rhat_max"] == pytest.approx(max(info["rhat_per_param"]), rel=1e-15) + + +def test_diagnostics_forward_the_backend_outputs(monkeypatch): + """ + Ownership: with each BlackJAX estimator replaced by a deterministic double, the + search forwards its output unchanged to the matching ``samples_info`` key and takes + the documented min / max. + """ + _skip_missing("nuts") + from blackjax import diagnostics + + doubles = { + "effective_sample_size": [11.0, 7.0, 13.0], + "ess_bulk": [21.0, 17.0, 23.0], + "ess_tail": [31.0, 37.0, 33.0], + "rhat": [1.01, 1.07, 1.03], + } + for name, value in doubles.items(): + monkeypatch.setattr( + diagnostics, name, lambda chains, _value=value: np.asarray(_value) + ) + + internal = cases.internal_from("nuts") + info = cases.search_from("nuts").samples_info_from(internal) + + assert info["ess_per_param"] == doubles["effective_sample_size"] + assert info["ess_min"] == 7.0 + assert info["ess_bulk_per_param"] == doubles["ess_bulk"] + assert info["ess_bulk_min"] == 17.0 + assert info["ess_tail_per_param"] == doubles["ess_tail"] + assert info["ess_tail_min"] == 31.0 + assert info["rhat_per_param"] == doubles["rhat"] + assert info["rhat_max"] == 1.07 diff --git a/test_autofit/non_linear/search/nest/nss/test_checkpoint.py b/test_autofit/non_linear/search/nest/nss/test_checkpoint.py index e3789e3dd..2a9b34dca 100644 --- a/test_autofit/non_linear/search/nest/nss/test_checkpoint.py +++ b/test_autofit/non_linear/search/nest/nss/test_checkpoint.py @@ -236,3 +236,93 @@ def test__resume_state_loaded_when_ctx_resume_is_set(tmp_path): resume.load.assert_called_once_with(paths) assert isinstance(internal, nss_search_module._NSSInternal) + + +@requires_nss +def test__run_keeps_its_checkpoint_until_the_result_is_written(tmp_path): + """Review 05 finding 2: ``run(ctx)`` must not delete ``nss_checkpoint.pkl``. + + The resume state is discarded by ``NonLinearSearch._complete_fit`` only after the + archive is finalized, so a failure while the result is processed (after ``run`` + returns) leaves the run resumable. The run below resumes and terminates at once, + exactly as ``test__resume_state_loaded_when_ctx_resume_is_set``. + """ + from types import SimpleNamespace + from unittest.mock import MagicMock + + from autofit.non_linear.paths.null import NullPaths + + sentinel_state = SimpleNamespace( + integrator=SimpleNamespace(logZ=0.0, logZ_live=-100.0), + particles=SimpleNamespace( + position=jnp.zeros((4, 2)), + loglikelihood=jnp.zeros(4), + ), + ) + + # Minimum stub for the model. ``run`` reads ``prior_count`` and draws the + # initial particles through ``vector_from_unit_vector``; the prior closure + # isn't traced unless one_step fires. + mock_model = SimpleNamespace( + prior_count=2, + vector_from_unit_vector=lambda v, xp=None: jnp.asarray([0.0, 0.0]), + log_prior_list_from_vector=lambda **kw: [0.0, 0.0], + ) + + resume = MagicMock() + resume.load.return_value = (sentinel_state, [], jax.random.PRNGKey(0), 17) + + paths = NullPaths() + ctx = SimpleNamespace( + model=mock_model, + paths=paths, + resume=resume, + fitness=SimpleNamespace( + objective=lambda kind, compile=None: (lambda params: 0.0), + analysis=SimpleNamespace(), + ), + ) + + search = af.NSS(n_live=4, num_mcmc_steps=1, num_delete=1, termination=-3.0) + + checkpoint = tmp_path / "nss_checkpoint.pkl" + checkpoint.write_bytes(b"resume state") + discarded = [] + search.resume_state = SimpleNamespace( + path=lambda paths: checkpoint, + discard=lambda paths: discarded.append(paths), + ) + + with patch.object( + nss_search_module, + "_blackjax", + ) as mock_bjax, patch.object( + nss_search_module, + "_nss_finalise", + return_value=SimpleNamespace( + particles=SimpleNamespace( + position=np.zeros((1, 2)), + loglikelihood=np.zeros(1), + ), + update_info=SimpleNamespace( + num_expansions=np.zeros(1, dtype=int), + num_shrink=np.zeros(1, dtype=int), + ), + ), + ), patch.object( + nss_search_module, + "_nss_log_weights", + return_value=jnp.zeros((1, 100)), + ): + mock_bjax.nss.return_value.init.side_effect = AssertionError( + "algo.init called from resume path — expected ctx.resume.load instead." + ) + mock_bjax.nss.return_value.step.side_effect = AssertionError( + "algo.step called even though logZ termination should fire immediately." + ) + + internal = search.run(ctx) + + assert isinstance(internal, nss_search_module._NSSInternal) + assert discarded == [] + assert checkpoint.exists() diff --git a/test_autofit/non_linear/search/test_fit_context.py b/test_autofit/non_linear/search/test_fit_context.py index 85ad6159a..fdc18ad0f 100644 --- a/test_autofit/non_linear/search/test_fit_context.py +++ b/test_autofit/non_linear/search/test_fit_context.py @@ -187,3 +187,152 @@ def ctx(): path.write_bytes(b"state") assert ctx().resume is search.resume_state + + +# Review 05 findings 5 and 6: the frozen contract promises that the context closes the +# pools it built on every exit, and that `ctx.update` saves the checkpoint. + + +class _RecordingPool: + def __init__(self, calls, processes, **kwargs): + self.calls = calls + calls.append("created") + + def close(self): + self.calls.append("close") + + def terminate(self): + self.calls.append("terminate") + + def join(self): + self.calls.append("join") + + +@pytest.fixture(name="pool_calls") +def make_pool_calls(monkeypatch): + from types import SimpleNamespace + + from autofit.non_linear.parallel import pool as pool_module + + calls = [] + monkeypatch.setattr( + pool_module, + "fork_context", + lambda: SimpleNamespace( + Pool=lambda processes, **kwargs: _RecordingPool(calls, processes, **kwargs) + ), + ) + monkeypatch.delenv("PYAUTO_TEST_MODE", raising=False) + return calls + + +class PoolToySearch(ToySearch): + fail = False + + def run(self, ctx): + assert ctx.pool() is not None + internal = super().run(ctx) + if self.fail: + raise RuntimeError("backend failed") + return internal + + +def test_a_successful_run_closes_and_joins_the_pools_the_context_built(pool_calls): + search = PoolToySearch(number_of_cores=2, initializer=af.InitializerPrior()) + search.fit(model=af.Model(af.ex.Gaussian), analysis=_analysis()) + + assert pool_calls == ["created", "close", "join"] + + +def test_a_failed_run_terminates_and_joins_the_pools_the_context_built(pool_calls): + search = PoolToySearch(number_of_cores=2, initializer=af.InitializerPrior()) + search.fail = True + + with pytest.raises(RuntimeError, match="backend failed"): + search.fit(model=af.Model(af.ex.Gaussian), analysis=_analysis()) + + assert pool_calls == ["created", "terminate", "join"] + + +def test_pool_factory_close_is_idempotent(pool_calls): + from autofit.non_linear.parallel import PoolFactory + + factory = PoolFactory(ToySearch(number_of_cores=2)) + factory() + factory.close(failed=False) + factory.close(failed=True) + + assert pool_calls == ["created", "close", "join"] + + +def _recording_resume_state(): + from autofit.non_linear.checkpoint import DillCheckpointer + + class RecordingCheckpointer(DillCheckpointer): + def __init__(self): + super().__init__(filename="toy_resume.dill") + self.saved = [] + self.finalized = [] + + def save(self, paths, obj): + self.saved.append(obj) + super().save(paths, obj) + + def finalize(self, paths, obj=None): + self.finalized.append(obj) + + return RecordingCheckpointer() + + +class UpdatingToySearch(ToySearch): + interrupt = False + + def run(self, ctx): + internal = super().run(ctx) + ctx.update(internal) + if self.interrupt: + raise InterruptedError("killed after the update") + return internal + + +def test_ctx_update_saves_the_resume_state_so_an_interrupted_run_keeps_it( + tmp_path, monkeypatch +): + monkeypatch.delenv("PYAUTO_TEST_MODE", raising=False) + + strategy = _recording_resume_state() + + class Resumable(UpdatingToySearch): + resume_state = strategy + + search = Resumable( + initializer=af.InitializerPrior(), path_prefix=str(tmp_path), name="toy_update" + ) + search.interrupt = True + + with pytest.raises(InterruptedError): + search.fit(model=af.Model(af.ex.Gaussian), analysis=_analysis()) + + assert len(strategy.saved) == 1 + assert len(strategy.saved[0]["log_likelihoods"]) == 6 + # Saved through `save`, never the completion `finalize`, and still on disk to resume. + assert strategy.finalized == [] + assert strategy.exists(search.paths) + + +def test_ctx_update_saves_nothing_for_a_search_without_a_resume_state( + tmp_path, monkeypatch +): + monkeypatch.delenv("PYAUTO_TEST_MODE", raising=False) + + saved = [] + search = UpdatingToySearch( + initializer=af.InitializerPrior(), path_prefix=str(tmp_path), name="toy_plain" + ) + monkeypatch.setattr( + type(search.checkpointer), "save", lambda self, paths, obj: saved.append(obj) + ) + search.fit(model=af.Model(af.ex.Gaussian), analysis=_analysis()) + + # The archive is written once, at completion, not by the mid-run update. + assert len(saved) <= 1 diff --git a/test_autofit/non_linear/search/test_preflight.py b/test_autofit/non_linear/search/test_preflight.py index 08c98495c..e3c14f6a0 100644 --- a/test_autofit/non_linear/search/test_preflight.py +++ b/test_autofit/non_linear/search/test_preflight.py @@ -297,3 +297,85 @@ def test_no_public_search_requires_fp64(): from autofit.non_linear.search.registry import entries assert not any(entry.capabilities["requires_fp64"] for entry in entries()) + + +# Review 05 finding 3: the preflight validates the *resolved* execution settings the +# backend runs (eager execution, the search's gradient override, its actual batching), +# not the class declarations. + + +def test_preflight_accepts_an_eager_likelihood_when_the_search_runs_eagerly( + monkeypatch, +): + monkeypatch.delenv("PYAUTO_TEST_MODE", raising=False) + search = af.DynestyStatic(use_jax_jit=False) + + settings = search.execution_settings(NumpyCastAnalysis()) + assert settings.compile is False + assert settings.fitness_overrides["compile"] is False + + # Eager scalar execution never traces, so the preflight has nothing to reject. + preflight.trace_preflight(search, NumpyCastAnalysis(), _model()) + + reached = [] + monkeypatch.setattr( + search, "_fit", lambda *args, **kwargs: reached.append(True) or 1 / 0 + ) + with pytest.raises(ZeroDivisionError): + search.fit(model=_model(), analysis=NumpyCastAnalysis()) + assert reached == [True] + + +def _custom_vjp_analysis(declared_mode): + @jax.custom_vjp + def reverse_only(x): + return x + + def _fwd(x): + return x, None + + def _bwd(_, g): + return (g,) + + reverse_only.defvjp(_fwd, _bwd) + + class ReverseOnlyAnalysis(af.Analysis): + gradient_mode = declared_mode + + def __init__(self): + super().__init__(use_jax=True) + + def log_likelihood_function(self, instance): + return -0.5 * reverse_only(instance.centre - 50.0) ** 2 + + return ReverseOnlyAnalysis() + + +def test_preflight_honours_the_search_gradient_mode_override(): + analysis = _custom_vjp_analysis("forward") + + # The analysis declares forward mode, which a custom_vjp function cannot run ... + with pytest.raises(exc.SearchException, match="value_and_grad"): + preflight.trace_preflight(af.MultiStartAdam(), analysis, _model()) + + # ... but the search overrides it with reverse mode, which is what the backend runs. + search = af.MultiStartAdam(gradient_mode="reverse") + assert search.execution_settings(analysis).gradient_mode == "reverse" + preflight.trace_preflight(search, analysis, _model()) + + +def test_preflight_follows_the_resolved_batching(): + search = af.Nautilus(use_jax_vmap=False) + settings = search.execution_settings(TraceableAnalysis()) + assert settings.batched is False + assert settings.kinds == ("scalar",) + + with pytest.raises(exc.SearchException, match="'scalar'"): + preflight.trace_preflight(search, NumpyCastAnalysis(), _model()) + + +def test_preflight_still_rejects_an_untraceable_likelihood_of_a_jitted_search(): + with pytest.raises(exc.SearchException, match="TracerArrayConversionError"): + preflight.trace_preflight( + af.DynestyStatic(use_jax_jit=True), NumpyCastAnalysis(), _model() + ) diff --git a/test_autofit/non_linear/test_checkpoint.py b/test_autofit/non_linear/test_checkpoint.py index fe3d207ac..fef4f5d13 100644 --- a/test_autofit/non_linear/test_checkpoint.py +++ b/test_autofit/non_linear/test_checkpoint.py @@ -307,3 +307,54 @@ def test_a_retained_archive_survives_with_search_internal_off( assert "files/search_internal/search_internal.hdf" in _archived(emcee) assert not any("search_internal/" in name for name in _archived(drawer)) + + +class FlakyArchive(DillCheckpointer): + """ + An archive whose first ``finalize`` fails as a full disk would (review 05 finding 2). + """ + + failures = 0 + calls = 0 + + def finalize(self, paths, obj=None): + type(self).calls += 1 + if type(self).failures > 0: + type(self).failures -= 1 + raise OSError("simulated disk full") + super().finalize(paths, obj) + + +class FlakyDrawer(af.Drawer): + checkpointer = FlakyArchive(retain_after_completion=True) + + +def test_a_failed_archive_write_leaves_the_fit_incomplete_and_a_rerun_finishes_it( + paths, monkeypatch +): + """ + ``.completed`` is written only once the archive is finalized: an archive failure + leaves no marker, and the rerun re-finalizes the archive and completes. + """ + import numpy as np + + monkeypatch.setenv("PYAUTO_TEST_MODE", "1") + monkeypatch.delenv("PYAUTO_SKIP_FIT_OUTPUT", raising=False) + FlakyArchive.failures = 1 + FlakyArchive.calls = 0 + + analysis = af.ex.Analysis(data=np.full(10, 5.0), noise_map=np.ones(10)) + + first = FlakyDrawer(name="flaky_archive", unique_tag="checkpoint_flaky", total_draws=5) + with pytest.raises(OSError, match="simulated disk full"): + first.fit(model=af.Model(af.ex.Gaussian), analysis=analysis) + + assert not first.paths.is_complete + + second = FlakyDrawer(name="flaky_archive", unique_tag="checkpoint_flaky", total_draws=5) + second.fit(model=af.Model(af.ex.Gaussian), analysis=analysis) + + assert second.paths.is_complete or Path(second.paths._zip_path).exists() + assert FlakyArchive.calls == 2 + assert "files/search_internal/search_internal.dill" in _archived(second) + assert ".completed" in _archived(second) diff --git a/test_autofit/non_linear/test_fitness_objective.py b/test_autofit/non_linear/test_fitness_objective.py index d01abca41..b32d7cb91 100644 --- a/test_autofit/non_linear/test_fitness_objective.py +++ b/test_autofit/non_linear/test_fitness_objective.py @@ -225,3 +225,69 @@ def test_evaluate_in_chunks_matches_one_call_and_compiles_one_shape(): assert np.allclose(np.asarray(chunked), np.asarray(whole)) # Shapes (2, d) and (5, d): the padded ragged chunk reuses the (2, d) executable. assert batched.__wrapped__._cache_size() == 2 + + +class EagerOnlyAnalysis(af.ex.Analysis): + """ + A JAX analysis whose likelihood is valid eagerly but not under `jax.jit`: it + converts a parameter to a Python float, which raises on a tracer. + """ + + def log_likelihood_function(self, instance, **kwargs): + return float(np.asarray(instance.centre)) * 0.0 + 50.0 + + +def _legacy_state(fitness, use_jax_vmap, use_jax_jit): + """ + The `__dict__` a pre-A2 (A1-era) `Fitness` pickled: `use_jax_vmap` / `use_jax_jit` + instead of `batched` / `compile`. + """ + state = fitness.__getstate__() + state.pop("batched") + state.pop("compile") + state["use_jax_vmap"] = use_jax_vmap + state["use_jax_jit"] = use_jax_jit + return state + + +def _restore(state): + import pickle + + restored = Fitness.__new__(Fitness) + restored.__setstate__(state) + return pickle.loads(pickle.dumps(restored)) + + +def test_legacy_eager_pickle_stays_eager(): + pytest.importorskip("jax") + + model = af.Model(af.ex.Gaussian) + analysis = EagerOnlyAnalysis( + data=np.ones(20), noise_map=np.ones(20) * 0.1, use_jax=True + ) + fitness = Fitness(model=model, analysis=analysis, compile=False) + vector = _vector(fitness) + assert float(fitness._call(vector)) == pytest.approx(50.0) + + restored = _restore(_legacy_state(fitness, use_jax_vmap=False, use_jax_jit=False)) + + assert restored.compile is False + assert restored.batched is False + assert not hasattr(restored._call, "__wrapped__") + assert float(restored._call(vector)) == pytest.approx(50.0) + + +def test_legacy_jit_and_vmap_pickles_stay_compiled(): + pytest.importorskip("jax") + + fitness = _fitness(use_jax=True) + + jitted = _restore(_legacy_state(fitness, use_jax_vmap=False, use_jax_jit=True)) + assert jitted.compile is True + assert jitted.batched is False + assert hasattr(jitted._call, "__wrapped__") + + vmapped = _restore(_legacy_state(fitness, use_jax_vmap=True, use_jax_jit=False)) + assert vmapped.compile is True + assert vmapped.batched is True + assert hasattr(vmapped._call, "__wrapped__") diff --git a/test_autofit/non_linear/test_fork_rule.py b/test_autofit/non_linear/test_fork_rule.py index 5ba302a18..9bba2bb9a 100644 --- a/test_autofit/non_linear/test_fork_rule.py +++ b/test_autofit/non_linear/test_fork_rule.py @@ -110,3 +110,49 @@ def test_analysis_pool_refuses_a_jax_analysis(): with pytest.raises(exc.SearchException, match="JAX analyses never fork"): AnalysisPool([_Analysis(is_jax=True)], n_cores=2) + + +class _BackendReached(Exception): + pass + + +def test_direct_fit_entry_applies_the_fork_rule_to_a_jax_analysis(monkeypatch): + """ + Review 05 finding 4: calling the `run(ctx)` bridge `_fit` directly (outside `fit`, + so no fit-scoped `PoolFactory` exists) must still build the context's factory from + the analysis, so Nautilus with a JAX analysis and two cores takes its single-process + path and nothing forks. + """ + pytest.importorskip("jax") + + def no_fork(): + raise AssertionError("must not fork") + + monkeypatch.setattr(pool_module, "fork_context", no_fork) + + search = af.Nautilus(number_of_cores=2) + chosen = {} + + def x1_cpu(fitness, model, update): + chosen["path"] = "fit_x1_cpu" + raise _BackendReached + + def multiprocessing(fitness, model, pools, update): + chosen["path"] = "fit_multiprocessing" + chosen["is_jax"] = pools.is_jax + chosen["cores"] = pools.number_of_cores + raise _BackendReached + + monkeypatch.setattr(search, "fit_x1_cpu", x1_cpu) + monkeypatch.setattr(search, "fit_multiprocessing", multiprocessing) + + analysis = af.ex.Analysis( + data=np.ones(10), noise_map=np.ones(10), use_jax=True + ) + + with pytest.raises(_BackendReached): + search._fit(model=af.Model(af.ex.Gaussian), analysis=analysis) + + assert chosen == {"path": "fit_x1_cpu"} + assert search._pools(analysis).is_jax + assert search._pools(analysis).number_of_cores == 1