From 3d3885c1c18173e7aeb2af0b07ae7c67acd6d142 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Mon, 27 Jul 2026 16:08:46 +0100 Subject: [PATCH 1/3] fix: three cadence/update bugs from the PR#1421 review Part 1 - Emcee and BlackJAX NUTS crashed on a real cadence, the same defect PR#1421 fixed for MultiStart: emcee's sample() does range(iterations) with no cast, and jax.random.split rejects a float. Both verified by probing the callee rather than reading the call shape. The plan said fix the producer (the float() coercion in AbstractSearch.__init__), on the grounds that Python ints hold 1e99 fine. Rejected after auditing the consumers: int(1e99) is a 99-digit integer, and writing that into every saved search.json in place of a readable inf-like sentinel is a bad trade. Instead the conversion moves up to one shared, validated AbstractSearch._steps_until_full_update, so it is derived once rather than re-derived (and forgotten) per search - which is exactly how this class of bug shipped three times. MultiStart's local _steps_in_chunk is folded into it. Part 2 - MultiStart emitted during_analysis=False on its last chunk while start_resume_fit emits the final update unconditionally anyway, so the whole final pass (sample output, latents, visualization, profiling) ran twice on every search. Emcee, BFGS and Nautilus all pass True unconditionally; MultiStart was the outlier. Now it matches them. Part 3 - a resumed run inherited the previous run's stop_reason, so raising n_steps to extend a finished search left a stale max_steps in every intermediate checkpoint and a still-running search reported itself finished. Cleared on entry, with converged preserved because the loop guard uses it to refuse resuming a converged search. Co-Authored-By: Claude Opus 5 --- autofit/non_linear/search/abstract_search.py | 77 ++++++++ .../search/mcmc/blackjax/nuts/search.py | 4 +- .../non_linear/search/mcmc/emcee/search.py | 7 +- .../search/mle/multi_start_gradient/search.py | 100 +++------- .../search/mle/test_multi_start_gradient.py | 178 ++++++++---------- .../search/test_steps_until_full_update.py | 124 ++++++++++++ 6 files changed, 305 insertions(+), 185 deletions(-) create mode 100644 test_autofit/non_linear/search/test_steps_until_full_update.py diff --git a/autofit/non_linear/search/abstract_search.py b/autofit/non_linear/search/abstract_search.py index a1b88883f..db3003d45 100644 --- a/autofit/non_linear/search/abstract_search.py +++ b/autofit/non_linear/search/abstract_search.py @@ -1037,6 +1037,83 @@ def output_search_internal(self, search_internal): obj=search_internal, ) + def _steps_until_full_update(self, iterations_remaining: int) -> int: + """ + How many iterations to run before the next ``perform_update`` + checkpoint, given how much of the budget is left. + + Searches that chunk their run around ``iterations_per_full_update`` all + need this same number, and they all need it as an ``int`` — most feed it + to something that ultimately does ``range(...)``. But + ``iterations_per_full_update`` is stored as a **float** by + ``__init__``, because the packaged default is the inf-like ``1e99`` + meaning "one chunk, checkpoint only at the end". That float is kept + deliberately: ``int(1e99)`` is a 99-digit integer, and writing *that* + into every saved ``search.json`` in place of a readable ``1e99`` + sentinel would be a poor trade for a conversion each caller can do at + the point of use. + + So the conversion lives here, once, instead of being re-derived (and + occasionally forgotten) per search. Forgetting it is not hypothetical: + it crashed the MultiStart gradient search for every user who set a real + cadence (PyAutoFit#1420), and left the same latent crash in Emcee and + BlackJAX NUTS (PyAutoFit#1422). + + A falsy ``iterations_per_full_update`` means "no intermediate + checkpoint" and yields the whole remaining budget in one chunk. + + Parameters + ---------- + iterations_remaining + Iterations left in this search's budget. Validated, because the + returned chunk can only be a usable positive whole number if this + is one — searches store their budget under different names + (``n_steps``, ``nsteps``, ``num_samples``, ``maxiter``) and none of + them validates it. + + Raises + ------ + ValueError + If ``iterations_remaining``, or an explicitly supplied + ``iterations_per_full_update``, is not a whole number of at + least 1. + """ + self._check_step_count(iterations_remaining, "iterations_remaining") + + if not self.iterations_per_full_update: + return int(iterations_remaining) + + self._check_step_count( + self.iterations_per_full_update, "iterations_per_full_update" + ) + return int(min(self.iterations_per_full_update, iterations_remaining)) + + def _check_step_count(self, value, name: str): + """ + Reject an iteration count that cannot describe a whole number of + iterations. ``float`` values are allowed — the packaged config default + ``1e99`` is one — provided they are integral. + + Rejecting rather than clamping is deliberate: a value below 1 truncates + to a zero-length chunk, so the enclosing ``while`` loop makes no + progress and spins forever re-running ``perform_update``. Clamping it to + 1 would hide the mistake and silently run a schedule the user never + asked for, and a silent hang on a cluster is far more expensive than an + error at the first chunk boundary. + """ + try: + is_whole = value == int(value) + except (TypeError, ValueError, OverflowError): + is_whole = False + + if not is_whole or value < 1: + raise ValueError( + f"{type(self).__name__}: `{name}` must be a whole number of " + f"iterations and at least 1, but was {value!r}. A fractional " + "or sub-1 value gives a zero-length chunk, which never advances " + "the search and would hang it rather than fail." + ) + @property def _updater(self): # The cached ``SearchUpdater`` must be invalidated whenever diff --git a/autofit/non_linear/search/mcmc/blackjax/nuts/search.py b/autofit/non_linear/search/mcmc/blackjax/nuts/search.py index 3f1bfb729..005b514fa 100644 --- a/autofit/non_linear/search/mcmc/blackjax/nuts/search.py +++ b/autofit/non_linear/search/mcmc/blackjax/nuts/search.py @@ -288,7 +288,9 @@ def run_chunk(rng_key, initial_state, n_steps): iterations_remaining = self.num_samples while iterations_remaining > 0: - chunk_n = min(self.iterations_per_full_update, iterations_remaining) + # ``run_chunk`` scans ``chunk_n`` steps and the key split below sizes + # itself from it, so this must be an ``int`` (PyAutoFit#1422). + chunk_n = self._steps_until_full_update(iterations_remaining) rng_key, sample_key = jax.random.split(rng_key) states, infos = run_chunk(sample_key, state, chunk_n) diff --git a/autofit/non_linear/search/mcmc/emcee/search.py b/autofit/non_linear/search/mcmc/emcee/search.py index a1c3d7412..02e26cb93 100644 --- a/autofit/non_linear/search/mcmc/emcee/search.py +++ b/autofit/non_linear/search/mcmc/emcee/search.py @@ -200,10 +200,9 @@ def _fit(self, model: AbstractPriorModel, analysis): iterations_remaining = self.nsteps while iterations_remaining > 0: - if self.iterations_per_full_update > iterations_remaining: - iterations = iterations_remaining - else: - iterations = self.iterations_per_full_update + # ``emcee``'s ``sample`` does ``range(iterations)`` internally and + # never casts, so this must be an ``int`` (PyAutoFit#1422). + iterations = self._steps_until_full_update(iterations_remaining) for sample in search_internal.sample( initial_state=state, 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 1df2b6d57..e474423b2 100644 --- a/autofit/non_linear/search/mle/multi_start_gradient/search.py +++ b/autofit/non_linear/search/mle/multi_start_gradient/search.py @@ -156,79 +156,6 @@ def __init__( self.logger.debug(f"Creating {self.optax_method} MultiStartGradient Search") - def _steps_in_chunk(self, steps_remaining: int) -> int: - """ - The number of gradient steps to run before the next ``perform_update`` - checkpoint boundary, given how much of the ``n_steps`` budget is left. - - ``AbstractSearch.__init__`` stores ``iterations_per_full_update`` as a - **float**, because the packaged default is the inf-like ``1e99`` (a - single chunk, i.e. checkpoint only at the end). ``range`` requires an - ``int``, so the chunk size is cast here at the consumer rather than in - the shared float coercion, which every other search relies on. - - Without the cast the loop only survives on the default: ``min(1e99, - steps_remaining)`` returns the ``int`` operand, so the float never - reaches ``range``. A user-supplied cadence *below* the remaining budget - (e.g. ``iterations_per_full_update=50`` with ``n_steps=3000``) does - reach it, and raises ``TypeError: 'float' object cannot be interpreted - as an integer``. - - Both the cadence and ``n_steps`` are validated as whole numbers of at - least one step, rather than being coerced into something plausible. A - value below 1 (or a fractional one) truncates to ``range(0)``: no step - runs, ``total_steps`` never advances, and the enclosing ``while`` loop - spins forever re-running ``perform_update`` — a silent hang that on a - cluster burns the whole allocation. Clamping such a value to 1 would - hide the mistake instead of reporting it, so it raises here. Once both - are validated, ``steps_remaining`` is an ``int`` of at least 1 whenever - the loop is entered, so the returned chunk is at least 1 and never - overshoots the remaining budget. - - Parameters - ---------- - steps_remaining - The number of steps left in the ``n_steps`` budget. - - Raises - ------ - ValueError - If ``n_steps``, or an explicitly supplied - ``iterations_per_full_update``, is not a whole number of at least 1. - """ - self._check_step_count(self.n_steps, "n_steps") - - # A falsy cadence means "no intermediate checkpoint": one chunk covering - # the whole budget. Only a value the user actually supplied is validated. - if self.iterations_per_full_update: - self._check_step_count( - self.iterations_per_full_update, "iterations_per_full_update" - ) - cadence = self.iterations_per_full_update - else: - cadence = self.n_steps - - return int(min(cadence, steps_remaining)) - - def _check_step_count(self, value, name: str): - """ - Reject a step count that cannot describe a whole number of gradient - steps. ``float`` values are allowed — the packaged config default - ``1e99`` is one — provided they are integral. - """ - try: - is_whole = value == int(value) - except (TypeError, ValueError, OverflowError): - is_whole = False - - if not is_whole or value < 1: - raise ValueError( - f"{type(self).__name__}: `{name}` must be a whole number of " - f"gradient steps and at least 1, but was {value!r}. A fractional " - "or sub-1 value gives a zero-length step chunk, which never " - "advances the search and would hang it rather than fail." - ) - def _fit( self, model: AbstractPriorModel, @@ -348,13 +275,23 @@ def batched_value_and_grad(params): # ``resurrect`` is on); seeded independently of the broad-start draw. resurrect_rng = np.random.default_rng(1) + # A resumed run inherits the previous run's ``stop_reason``. ``"converged"`` + # must survive — it short-circuits the loop below so a converged search is + # never resumed into more steps. Any other reason describes a run that is + # now over: keeping it would report a stale ``"max_steps"`` in every + # intermediate checkpoint of a search that is, in fact, still running + # (raising ``n_steps`` is the documented way to extend a budget). It is + # re-derived at each chunk boundary below, so clear it here. + if stop_reason != "converged": + stop_reason = None + # ``n_steps`` is the hard ceiling / max budget; ``stop_reason`` becomes # ``"converged"`` if the auto-convergence check stops the search early # (this also short-circuits the loop on a resumed, already-converged run). while total_steps < self.n_steps and stop_reason != "converged": steps_remaining = self.n_steps - total_steps - iterations = self._steps_in_chunk(steps_remaining) + iterations = self._steps_until_full_update(steps_remaining) # Convergence is assessed every step (``fom_history`` updates every # step) rather than only at the ``iterations_per_full_update`` boundary, @@ -427,13 +364,20 @@ def batched_value_and_grad(params): } self.paths.save_search_internal(obj=search_internal) - # A converged (or ceiling-reached) boundary is the final update, so it - # runs with ``during_analysis=False``; intermediate boundaries do not. - is_final = converged or total_steps >= self.n_steps + # Always an *intermediate* update, even on the last chunk. The final + # ``during_analysis=False`` update is ``start_resume_fit``'s job, and + # it runs unconditionally once ``_fit`` returns — emitting one here + # too made the whole final pass (sample output, latents, + # visualization, profiling) run twice on every search. Every other + # search passes ``during_analysis=True`` unconditionally here for + # exactly this reason (Emcee, BFGS, Nautilus); MultiStart was the + # outlier. The ``stop_reason``/``converged`` outcome still reaches the + # final samples, because that update is built from the + # ``search_internal`` dict returned below. self.perform_update( model=model, analysis=analysis, - during_analysis=not is_final, + during_analysis=True, fitness=fitness, search_internal=search_internal, ) diff --git a/test_autofit/non_linear/search/mle/test_multi_start_gradient.py b/test_autofit/non_linear/search/mle/test_multi_start_gradient.py index 818396550..3658b7c40 100644 --- a/test_autofit/non_linear/search/mle/test_multi_start_gradient.py +++ b/test_autofit/non_linear/search/mle/test_multi_start_gradient.py @@ -5,6 +5,7 @@ import autofit as af from autofit import example +from autofit.non_linear.search import abstract_search from autonerves.dictable import from_dict, to_dict # The MultiStart gradient searches are JAX-native at fit time, but their @@ -138,108 +139,6 @@ def test__convergence_default_is_on_and_carried(): assert cls(convergence=custom).convergence is custom -def test__steps_in_chunk__real_cadence_is_an_int_range_can_consume(): - """The step loop does ``for _ in range(...)`` over the chunk size, but - ``iterations_per_full_update`` is stored as a float by ``AbstractSearch``. - - A cadence *below* the remaining budget is the case that reaches ``range`` - (with the ``1e99`` default, ``min`` returns the int ``steps_remaining`` - instead), and it used to raise ``TypeError: 'float' object cannot be - interpreted as an integer`` on step-loop entry. - """ - search = af.MultiStartAdam(n_steps=3000, iterations_per_full_update=50) - - # The knob really is stored as a float — this is what made the crash latent. - assert isinstance(search.iterations_per_full_update, float) - - iterations = search._steps_in_chunk(steps_remaining=3000) - - assert iterations == 50 - assert isinstance(iterations, int) - # The actual failing operation in ``_fit``. - assert len(list(range(iterations))) == 50 - - -def test__steps_in_chunk__clamps_to_the_remaining_budget(): - # Near the end of the budget the chunk shrinks to what is left, so the loop - # never overshoots ``n_steps``. - search = af.MultiStartAdam(n_steps=3000, iterations_per_full_update=50) - - assert search._steps_in_chunk(steps_remaining=20) == 20 - assert isinstance(search._steps_in_chunk(steps_remaining=20), int) - - -def test__steps_in_chunk__default_cadence_is_a_single_chunk(): - # The packaged default is the inf-like 1e99: one chunk covering the whole - # budget (checkpoint only at the end), which is why the crash never fired - # on a default run. - search = af.MultiStartAdam(n_steps=300) - - assert search.iterations_per_full_update == pytest.approx(1e99) - - iterations = search._steps_in_chunk(steps_remaining=300) - - assert iterations == 300 - assert isinstance(iterations, int) - - -def test__fit_step_loop_takes_its_chunk_size_from_the_helper(): - """Wiring guard: every other test here exercises ``_steps_in_chunk`` - directly, so all of them would still pass if ``_fit`` went back to computing - the chunk inline — which is the exact regression this fix is about. - - ``_fit`` cannot be executed from this suite (it needs jax, optax and a - JAX-traceable ``Analysis``, and the library suite is NumPy-only), so the - call site is asserted at the source level instead. - """ - source = inspect.getsource(af.MultiStartAdam._fit) - - assert "self._steps_in_chunk(" in source - # the un-cast inline form the crash came from must not come back - assert "self.iterations_per_full_update or self.n_steps" not in source - - -@pytest.mark.parametrize("cadence", [0.5, 50.9, -5]) -def test__steps_in_chunk__unusable_cadence_raises_rather_than_being_clamped(cadence): - """A cadence below 1 (or a fractional one) truncates to ``range(0)``: no - step runs, ``total_steps`` never advances, and the step loop's enclosing - ``while`` spins forever re-running ``perform_update``. - - Clamping such a value to 1 would hide the mistake and silently run a - cadence the user did not ask for, so it is rejected instead — a hang on a - cluster is far more expensive than an error at the first chunk boundary. - """ - search = af.MultiStartAdam(n_steps=300, iterations_per_full_update=cadence) - - with pytest.raises(ValueError, match="iterations_per_full_update"): - search._steps_in_chunk(steps_remaining=300) - - -@pytest.mark.parametrize("n_steps", [2.5, 0, -10]) -def test__steps_in_chunk__unusable_n_steps_raises(n_steps): - """``steps_remaining`` is only guaranteed to be an ``int`` of at least 1 - because ``n_steps`` is one. A fractional ``n_steps`` leaves a fractional - remainder (e.g. 0.5) that truncates to a zero-length chunk, so the budget - is validated too rather than assumed from its type annotation.""" - search = af.MultiStartAdam(n_steps=n_steps) - - with pytest.raises(ValueError, match="n_steps"): - search._steps_in_chunk(steps_remaining=1) - - -def test__steps_in_chunk__falsy_cadence_falls_back_to_n_steps(): - # A falsy cadence means "no checkpoint boundary": fall back to the full - # ``n_steps`` budget rather than a zero-length chunk that would spin the - # while-loop forever. ``__init__`` cannot produce this (a falsy argument - # resolves from config to 1e99), so it is set directly — which is exactly - # what a caller overwriting the attribute post-construction does. - search = af.MultiStartAdam(n_steps=120) - search.iterations_per_full_update = 0.0 - - assert search._steps_in_chunk(steps_remaining=120) == 120 - assert isinstance(search._steps_in_chunk(steps_remaining=120), int) - - def test__check_if_converged__plateau_stops_climbing_does_not(): # rtol/atol both zero: converges only on an exactly-flat window. convergence = af.MultiStartGradientConvergence( @@ -486,3 +385,78 @@ def samples_for(n_starts, total_steps): # The aggregator summary is computed without an IndexError. assert samples.summary() is not None + + +def test__fit_never_emits_the_final_update_itself(): + """``_fit`` must only ever emit *intermediate* updates. + + ``start_resume_fit`` performs the single ``during_analysis=False`` update + unconditionally once ``_fit`` returns (``abstract_search.py``), so a + ``during_analysis=False`` in the step loop made the whole final pass — + sample output, latent computation, visualization, profiling — run twice on + every search. Every other chunked search passes ``during_analysis=True`` + unconditionally for this reason; MultiStart was the outlier. + + ``_fit`` needs jax + optax + a JAX-traceable ``Analysis`` and cannot run + from this NumPy-only suite, so the call is asserted at the source level. + """ + source = inspect.getsource(af.MultiStartAdam._fit) + + assert "during_analysis=True" in source + assert "during_analysis=not is_final" not in source + assert "is_final" not in source + + # The framework's single final update is still there and unconditional — + # without it this change would remove the final update rather than + # de-duplicate it. ``start_resume_fit`` is decorated without + # ``functools.wraps``, so its body is sliced out of the module source rather + # than reached through the bound method. + module = inspect.getsource(abstract_search) + body = module[module.index(" def start_resume_fit(") :] + body = body[: body.index("\n def ", 1)] + assert "during_analysis=False" in body + + +def test__fit_clears_a_stale_stop_reason_but_keeps_converged(): + """A resumed run inherits the previous run's ``stop_reason``. + + Resuming a search that finished with ``"max_steps"`` under a raised + ``n_steps`` (the documented way to extend a budget) used to leave that stale + reason in every intermediate checkpoint, so a search that was still running + reported itself finished to the aggregator and any results inspector. + + ``"converged"`` must survive the clear, because the loop guard uses it to + refuse resuming a converged search into more steps. + """ + source = inspect.getsource(af.MultiStartAdam._fit) + + assert 'if stop_reason != "converged":' in source + assert "stop_reason = None" in source + # the guard that "converged" protects is still the loop condition + assert 'stop_reason != "converged"' in source + + +def test__samples_info__reports_a_cleared_stop_reason_as_unfinished(): + """The observable half of the fix: a checkpoint written mid-run carries no + stop reason, so nothing downstream reads it as a finished search.""" + model = af.Model(example.Gaussian) + best_params = np.asarray(model.vector_from_unit_vector([0.5] * model.prior_count)) + + search = af.MultiStartAdam(n_starts=1, n_steps=600) + samples = search.samples_via_internal_from( + model=model, + search_internal={ + "params": np.stack([best_params]), + "best_params": best_params, + "best_fom": -2.0, + "fom_history": np.asarray([-2.0] * 300), + "total_steps": 300, + "n_resurrections": 0, + # what an intermediate checkpoint of a resumed, still-running search + # now writes — previously this said "max_steps". + "stop_reason": None, + }, + ) + + assert samples.samples_info["stop_reason"] is None + assert samples.samples_info["converged"] is False diff --git a/test_autofit/non_linear/search/test_steps_until_full_update.py b/test_autofit/non_linear/search/test_steps_until_full_update.py new file mode 100644 index 000000000..1c16dda33 --- /dev/null +++ b/test_autofit/non_linear/search/test_steps_until_full_update.py @@ -0,0 +1,124 @@ +import inspect + +import pytest + +import autofit as af + +# ``_steps_until_full_update`` is the one place the float +# ``iterations_per_full_update`` becomes the ``int`` chunk size that chunked +# searches feed to ``range`` (or to a JAX scan / key split). It lives on +# ``AbstractSearch`` because every chunked search needs it and forgetting it is +# a real, shipped bug class: it crashed MultiStart for anyone setting a real +# cadence (PyAutoFit#1420) and stayed latent in Emcee and BlackJAX NUTS +# (PyAutoFit#1422). +# +# NumPy-only, like the rest of the library suite — the searches' ``_fit`` bodies +# need jax/optax/emcee and are exercised in the workspace test repos. + +pytestmark = pytest.mark.filterwarnings("ignore::FutureWarning") + + +def test__real_cadence_is_an_int_range_can_consume(): + """A cadence *below* the remaining budget is the case that reaches + ``range``. With the ``1e99`` default ``min`` returns the int + ``iterations_remaining`` instead, which is why this stayed latent for every + user of the default until someone set a real checkpoint cadence.""" + search = af.MultiStartAdam(n_steps=3000, iterations_per_full_update=50) + + # The knob really is stored as a float — this is what made the crash latent. + assert isinstance(search.iterations_per_full_update, float) + + iterations = search._steps_until_full_update(iterations_remaining=3000) + + assert iterations == 50 + assert isinstance(iterations, int) + # The operation that used to raise TypeError. + assert len(list(range(iterations))) == 50 + + +def test__clamps_to_the_remaining_budget(): + # Near the end of the budget the chunk shrinks to what is left, so a search + # never overshoots its ceiling. + search = af.MultiStartAdam(n_steps=3000, iterations_per_full_update=50) + + assert search._steps_until_full_update(iterations_remaining=20) == 20 + assert isinstance(search._steps_until_full_update(iterations_remaining=20), int) + + +def test__default_cadence_is_a_single_chunk(): + # The packaged default is the inf-like 1e99: one chunk covering the whole + # budget, i.e. checkpoint only at the end. + search = af.MultiStartAdam(n_steps=300) + + assert search.iterations_per_full_update == pytest.approx(1e99) + + iterations = search._steps_until_full_update(iterations_remaining=300) + + assert iterations == 300 + assert isinstance(iterations, int) + + +def test__falsy_cadence_means_no_intermediate_checkpoint(): + # A falsy cadence is "no boundary": run the remaining budget in one chunk, + # rather than a zero-length chunk that would never advance. + search = af.MultiStartAdam(n_steps=120) + search.iterations_per_full_update = 0.0 + + assert search._steps_until_full_update(iterations_remaining=120) == 120 + assert isinstance(search._steps_until_full_update(iterations_remaining=120), int) + + +@pytest.mark.parametrize("cadence", [0.5, 50.9, -5]) +def test__unusable_cadence_raises_rather_than_being_clamped(cadence): + """A cadence below 1 (or a fractional one) truncates to a zero-length + chunk: nothing runs, the counter never advances, and the enclosing ``while`` + spins forever re-running ``perform_update``. + + Clamping to 1 would hide the mistake and silently run a schedule the user + never asked for — a silent hang on a cluster is far more expensive than an + error at the first chunk boundary. + """ + search = af.MultiStartAdam(n_steps=300, iterations_per_full_update=cadence) + + with pytest.raises(ValueError, match="iterations_per_full_update"): + search._steps_until_full_update(iterations_remaining=300) + + +@pytest.mark.parametrize("remaining", [2.5, 0, -10]) +def test__unusable_remaining_budget_raises(remaining): + """The chunk can only be a usable positive whole number if the remaining + budget is one. Searches keep that budget under different names (``n_steps``, + ``nsteps``, ``num_samples``, ``maxiter``) and none of them validates it, so + it is validated here at the point of use — a fractional budget leaves a + fractional remainder that truncates to a zero-length chunk.""" + search = af.MultiStartAdam(n_steps=300, iterations_per_full_update=50) + + with pytest.raises(ValueError, match="iterations_remaining"): + search._steps_until_full_update(iterations_remaining=remaining) + + +@pytest.mark.parametrize( + "search_cls, method_name", + [ + (af.MultiStartAdam, "_fit"), + (af.Emcee, "_fit"), + ], +) +def test__chunked_searches_take_their_chunk_size_from_the_helper( + search_cls, method_name +): + """Wiring guard. Every test above drives ``_steps_until_full_update`` + directly, so all of them would still pass if a search went back to computing + its chunk inline — which is precisely the regression this helper exists to + prevent, and precisely what shipped in PyAutoFit#1420. + + These ``_fit`` bodies cannot be executed from this suite (jax/optax/emcee + plus a real ``Analysis``), so the call site is asserted at the source level. + """ + source = inspect.getsource(getattr(search_cls, method_name)) + + assert "self._steps_until_full_update(" in source + # the un-cast inline forms the crashes came from must not come back + assert "self.iterations_per_full_update or self.n_steps" not in source + assert "min(self.iterations_per_full_update" not in source + assert "iterations = self.iterations_per_full_update" not in source From ac2e6e0e90a312220300d703f967f804b8cbb8ed Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Mon, 27 Jul 2026 16:22:41 +0100 Subject: [PATCH 2/3] fix: act on the adversarial review of the first commit Three findings, all confirmed before acting. 1. The Part 2 fix was cosmetic. Flipping during_analysis to True does not deduplicate the final pass: SearchUpdater.update rebuilds the samples, recomputes the summary and re-runs likelihood profiling on every call regardless of the flag, and the visualize gate keys off paths.is_complete, which is not written until after _fit returns. The in-loop update is now skipped outright at a terminal boundary. On the default single-chunk cadence _fit emits no update at all, leaving exactly one final pass. 2. Zeus was NOT safe, just not crashing. It casts iterations internally, but PyAutoFit adds the uncast float to total_iterations, so a fractional cadence drifts the bookkeeping away from the samples zeus actually drew (nsteps=100, cadence=50.9 finishes with 99). BFGS also bypassed the validation the helper advertises for maxiter. Both now use the shared helper, so every chunked search derives its chunk one way. This supersedes the issue's "do not touch zeus/bfgs" note, which rested on my own earlier finding that they were fine. 3. The falsy-cadence branch reintroduced the silent behaviour the validation exists to remove: a stored 0 meant "never checkpoint". 1e99 is already that sentinel and needs no special case, so 0 is a misconfiguration - reachable through the HPC override, which assigns the config value with no "or" fallback - and now raises. Also fixes the test guards the review showed were brittle or tautological: BlackJAX, Zeus and BFGS added to the wiring guard; the loop-guard assertion now pins the while line rather than being satisfied by the clearing if; source assertions normalise whitespace. Co-Authored-By: Claude Opus 5 --- autofit/non_linear/search/abstract_search.py | 11 +++-- autofit/non_linear/search/mcmc/zeus/search.py | 10 +++-- autofit/non_linear/search/mle/bfgs/search.py | 5 ++- .../search/mle/multi_start_gradient/search.py | 41 +++++++++++-------- .../search/mle/test_multi_start_gradient.py | 39 ++++++++++-------- .../search/test_steps_until_full_update.py | 30 +++++++------- 6 files changed, 76 insertions(+), 60 deletions(-) diff --git a/autofit/non_linear/search/abstract_search.py b/autofit/non_linear/search/abstract_search.py index db3003d45..4c8744c84 100644 --- a/autofit/non_linear/search/abstract_search.py +++ b/autofit/non_linear/search/abstract_search.py @@ -1059,8 +1059,11 @@ def _steps_until_full_update(self, iterations_remaining: int) -> int: cadence (PyAutoFit#1420), and left the same latent crash in Emcee and BlackJAX NUTS (PyAutoFit#1422). - A falsy ``iterations_per_full_update`` means "no intermediate - checkpoint" and yields the whole remaining budget in one chunk. + There is no second "no intermediate checkpoint" sentinel: the config + default ``1e99`` already means that, and it flows through the ``min`` + below without a special case. A stored ``0`` is therefore a + misconfiguration — most plausibly an HPC override — and is rejected + rather than silently reinterpreted as "checkpoint never". Parameters ---------- @@ -1079,10 +1082,6 @@ def _steps_until_full_update(self, iterations_remaining: int) -> int: least 1. """ self._check_step_count(iterations_remaining, "iterations_remaining") - - if not self.iterations_per_full_update: - return int(iterations_remaining) - self._check_step_count( self.iterations_per_full_update, "iterations_per_full_update" ) diff --git a/autofit/non_linear/search/mcmc/zeus/search.py b/autofit/non_linear/search/mcmc/zeus/search.py index 52133bb70..352b47c20 100644 --- a/autofit/non_linear/search/mcmc/zeus/search.py +++ b/autofit/non_linear/search/mcmc/zeus/search.py @@ -236,10 +236,12 @@ def _fit(self, model: AbstractPriorModel, analysis): iterations_remaining = self.nsteps while iterations_remaining > 0: - if self.iterations_per_full_update > iterations_remaining: - iterations = iterations_remaining - else: - iterations = self.iterations_per_full_update + # zeus casts internally (``self.nsteps = int(iterations)``), so a + # float never crashes it — but ``total_iterations`` below is + # incremented by this value, so an uncast float would drift the + # bookkeeping away from the samples zeus actually drew + # (PyAutoFit#1422). + iterations = self._steps_until_full_update(iterations_remaining) for sample in search_internal.sample( start=state, diff --git a/autofit/non_linear/search/mle/bfgs/search.py b/autofit/non_linear/search/mle/bfgs/search.py index 74ab8aafc..da76d5a08 100644 --- a/autofit/non_linear/search/mle/bfgs/search.py +++ b/autofit/non_linear/search/mle/bfgs/search.py @@ -168,7 +168,10 @@ def _fit( while total_iterations < self.maxiter: iterations_remaining = self.maxiter - total_iterations - iterations = min(self.iterations_per_full_update, iterations_remaining) + # SciPy tolerates an integral float ``maxiter``, but a fractional one + # would quietly acquire ceiling semantics instead of being rejected + # (PyAutoFit#1422). + iterations = self._steps_until_full_update(iterations_remaining) if iterations > 0: options = dict(self.options) 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 e474423b2..38cee8262 100644 --- a/autofit/non_linear/search/mle/multi_start_gradient/search.py +++ b/autofit/non_linear/search/mle/multi_start_gradient/search.py @@ -364,23 +364,30 @@ def batched_value_and_grad(params): } self.paths.save_search_internal(obj=search_internal) - # Always an *intermediate* update, even on the last chunk. The final - # ``during_analysis=False`` update is ``start_resume_fit``'s job, and - # it runs unconditionally once ``_fit`` returns — emitting one here - # too made the whole final pass (sample output, latents, - # visualization, profiling) run twice on every search. Every other - # search passes ``during_analysis=True`` unconditionally here for - # exactly this reason (Emcee, BFGS, Nautilus); MultiStart was the - # outlier. The ``stop_reason``/``converged`` outcome still reaches the - # final samples, because that update is built from the - # ``search_internal`` dict returned below. - self.perform_update( - model=model, - analysis=analysis, - during_analysis=True, - fitness=fitness, - search_internal=search_internal, - ) + # Update only at *intermediate* boundaries. ``start_resume_fit`` + # performs the final update unconditionally the moment ``_fit`` + # returns, so any update emitted here at a terminal boundary is pure + # duplicated work — and flipping ``during_analysis`` does not avoid + # it: ``SearchUpdater.update`` rebuilds the samples, recomputes the + # summary and re-runs likelihood profiling on *every* call, + # regardless of the flag. The previous ``during_analysis=not + # is_final`` was worse still, running the whole final pass twice + # including final visualization and latents. + # + # Nothing is lost by skipping: the checkpoint is + # ``save_search_internal`` above, and the ``search_internal`` dict + # this loop just built (carrying ``stop_reason``, ``converged`` and + # ``fom_history``) is what ``_fit`` returns and what that final + # update is computed from. + is_final = converged or total_steps >= self.n_steps + if not is_final: + self.perform_update( + model=model, + analysis=analysis, + during_analysis=True, + fitness=fitness, + search_internal=search_internal, + ) if converged: self.logger.info( diff --git a/test_autofit/non_linear/search/mle/test_multi_start_gradient.py b/test_autofit/non_linear/search/mle/test_multi_start_gradient.py index 3658b7c40..a93de1faa 100644 --- a/test_autofit/non_linear/search/mle/test_multi_start_gradient.py +++ b/test_autofit/non_linear/search/mle/test_multi_start_gradient.py @@ -387,24 +387,27 @@ def samples_for(n_starts, total_steps): assert samples.summary() is not None -def test__fit_never_emits_the_final_update_itself(): - """``_fit`` must only ever emit *intermediate* updates. +def test__fit_skips_its_update_at_a_terminal_boundary(): + """``start_resume_fit`` performs the final update unconditionally the moment + ``_fit`` returns, so any update ``_fit`` emits at a terminal boundary is + duplicated work. - ``start_resume_fit`` performs the single ``during_analysis=False`` update - unconditionally once ``_fit`` returns (``abstract_search.py``), so a - ``during_analysis=False`` in the step loop made the whole final pass — - sample output, latent computation, visualization, profiling — run twice on - every search. Every other chunked search passes ``during_analysis=True`` - unconditionally for this reason; MultiStart was the outlier. + Flipping ``during_analysis`` is *not* enough to avoid it: ``SearchUpdater`` + rebuilds the samples, recomputes the summary and re-runs likelihood + profiling on every call regardless of the flag. The update has to be skipped + outright, which is what the ``if not is_final`` guard does. ``_fit`` needs jax + optax + a JAX-traceable ``Analysis`` and cannot run - from this NumPy-only suite, so the call is asserted at the source level. + from this NumPy-only suite, so the guard is asserted at the source level. + Whitespace is normalised so reformatting does not cause a false failure. """ - source = inspect.getsource(af.MultiStartAdam._fit) + source = " ".join(inspect.getsource(af.MultiStartAdam._fit).split()) + # the update is guarded, and is intermediate-flavoured when it does run + assert "if not is_final: self.perform_update(" in source assert "during_analysis=True" in source + # the form that ran the whole final pass twice must not come back assert "during_analysis=not is_final" not in source - assert "is_final" not in source # The framework's single final update is still there and unconditional — # without it this change would remove the final update rather than @@ -425,15 +428,15 @@ def test__fit_clears_a_stale_stop_reason_but_keeps_converged(): reason in every intermediate checkpoint, so a search that was still running reported itself finished to the aggregator and any results inspector. - ``"converged"`` must survive the clear, because the loop guard uses it to - refuse resuming a converged search into more steps. + ``"converged"`` must survive the clear, because the ``while`` guard uses it + to refuse resuming a converged search into more steps — so that guard is + asserted on the ``while`` line specifically. Asserting the bare substring + would be satisfied by the clearing ``if`` itself and would pin nothing. """ - source = inspect.getsource(af.MultiStartAdam._fit) + source = " ".join(inspect.getsource(af.MultiStartAdam._fit).split()) - assert 'if stop_reason != "converged":' in source - assert "stop_reason = None" in source - # the guard that "converged" protects is still the loop condition - assert 'stop_reason != "converged"' in source + assert 'if stop_reason != "converged": stop_reason = None' in source + assert 'while total_steps < self.n_steps and stop_reason != "converged":' in source def test__samples_info__reports_a_cleared_stop_reason_as_unfinished(): diff --git a/test_autofit/non_linear/search/test_steps_until_full_update.py b/test_autofit/non_linear/search/test_steps_until_full_update.py index 1c16dda33..df57dde29 100644 --- a/test_autofit/non_linear/search/test_steps_until_full_update.py +++ b/test_autofit/non_linear/search/test_steps_until_full_update.py @@ -58,14 +58,21 @@ def test__default_cadence_is_a_single_chunk(): assert isinstance(iterations, int) -def test__falsy_cadence_means_no_intermediate_checkpoint(): - # A falsy cadence is "no boundary": run the remaining budget in one chunk, - # rather than a zero-length chunk that would never advance. +def test__zero_cadence_raises_rather_than_meaning_never_checkpoint(): + """``1e99`` is already the "no intermediate checkpoint" sentinel and flows + through the ``min`` without a special case, so a stored ``0`` is a + misconfiguration, not a second sentinel. + + ``__init__`` cannot produce it from a public argument (``x or config`` + replaces a falsy one), but the HPC branch assigns the config value with no + such fallback — so an HPC cadence of ``0`` would otherwise silently disable + checkpointing for the whole run instead of failing. + """ search = af.MultiStartAdam(n_steps=120) search.iterations_per_full_update = 0.0 - assert search._steps_until_full_update(iterations_remaining=120) == 120 - assert isinstance(search._steps_until_full_update(iterations_remaining=120), int) + with pytest.raises(ValueError, match="iterations_per_full_update"): + search._steps_until_full_update(iterations_remaining=120) @pytest.mark.parametrize("cadence", [0.5, 50.9, -5]) @@ -98,15 +105,10 @@ def test__unusable_remaining_budget_raises(remaining): @pytest.mark.parametrize( - "search_cls, method_name", - [ - (af.MultiStartAdam, "_fit"), - (af.Emcee, "_fit"), - ], + "search_cls", + [af.MultiStartAdam, af.Emcee, af.Zeus, af.BFGS, af.BlackJAXNUTS], ) -def test__chunked_searches_take_their_chunk_size_from_the_helper( - search_cls, method_name -): +def test__chunked_searches_take_their_chunk_size_from_the_helper(search_cls): """Wiring guard. Every test above drives ``_steps_until_full_update`` directly, so all of them would still pass if a search went back to computing its chunk inline — which is precisely the regression this helper exists to @@ -115,7 +117,7 @@ def test__chunked_searches_take_their_chunk_size_from_the_helper( These ``_fit`` bodies cannot be executed from this suite (jax/optax/emcee plus a real ``Analysis``), so the call site is asserted at the source level. """ - source = inspect.getsource(getattr(search_cls, method_name)) + source = inspect.getsource(search_cls._fit) assert "self._steps_until_full_update(" in source # the un-cast inline forms the crashes came from must not come back From 886d29d4d769990f06cfc1ab96a56598d13648e0 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Mon, 27 Jul 2026 16:35:32 +0100 Subject: [PATCH 3/3] test: pin the two MultiStart rules behaviourally, not by source string Second review pass confirmed the production fixes (it ran real JAX fits: a one-chunk run now emits zero in-loop updates and exactly one final update; a two-chunk run emits [True, False]). What it did not accept were the regression tests: source-level string assertions can pass while a semantically-equivalent regression is reintroduced - assigning is_final False, or adding a second unconditional update after the guarded one. So extract the two rules the loop was applying inline into named seams - _is_final_boundary and _stop_reason_on_resume - and test those directly and exhaustively, the same move that made the cadence fix testable without JAX. The source assertion shrinks to a wiring guard over the call sites, with the behaviour pinned by the parametrised tests instead. Known residual, accepted rather than papered over: the cross-search wiring guard still cannot prove a search *uses* the value the helper returns, so a contrived regression that calls the helper and discards the result would pass it. Closing that needs the searches' _fit bodies, which need jax/optax/emcee and are out of scope for the NumPy-only library suite; it is covered where those bodies actually run, in the workspace test repos. Co-Authored-By: Claude Opus 5 --- .../search/mle/multi_start_gradient/search.py | 47 +++++++--- .../search/mle/test_multi_start_gradient.py | 94 ++++++++++++------- 2 files changed, 95 insertions(+), 46 deletions(-) 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 38cee8262..cb392e38c 100644 --- a/autofit/non_linear/search/mle/multi_start_gradient/search.py +++ b/autofit/non_linear/search/mle/multi_start_gradient/search.py @@ -156,6 +156,40 @@ def __init__( self.logger.debug(f"Creating {self.optax_method} MultiStartGradient Search") + def _is_final_boundary(self, converged: bool, total_steps: int) -> bool: + """ + Whether the chunk boundary just reached ends the search — either the + auto-convergence check fired, or the ``n_steps`` ceiling was hit. + + ``_fit`` must not emit a ``perform_update`` at a terminal boundary: + ``start_resume_fit`` performs the final update unconditionally the + moment ``_fit`` returns, so one here is duplicated work. Flipping + ``during_analysis`` does not avoid that — ``SearchUpdater.update`` + rebuilds the samples, recomputes the summary and re-runs likelihood + profiling on every call regardless of the flag. + + A named predicate rather than an inline expression so the rule can be + tested directly; ``_fit`` itself needs jax + optax + a JAX-traceable + ``Analysis`` and cannot be driven from the NumPy-only library suite. + """ + return bool(converged) or total_steps >= self.n_steps + + @staticmethod + def _stop_reason_on_resume(stop_reason): + """ + The ``stop_reason`` a resumed run should start from, given the one + restored from its previous ``search_internal``. + + ``"converged"`` survives: the step loop's guard uses it to refuse + resuming a converged search into further steps. Every other reason + describes a run that is now over — keeping it would stamp a stale + ``"max_steps"`` on every intermediate checkpoint of a search that is in + fact still running, which is exactly what raising ``n_steps`` (the + documented way to extend a budget) does. It is re-derived at each chunk + boundary, so it starts clear. + """ + return stop_reason if stop_reason == "converged" else None + def _fit( self, model: AbstractPriorModel, @@ -275,15 +309,7 @@ def batched_value_and_grad(params): # ``resurrect`` is on); seeded independently of the broad-start draw. resurrect_rng = np.random.default_rng(1) - # A resumed run inherits the previous run's ``stop_reason``. ``"converged"`` - # must survive — it short-circuits the loop below so a converged search is - # never resumed into more steps. Any other reason describes a run that is - # now over: keeping it would report a stale ``"max_steps"`` in every - # intermediate checkpoint of a search that is, in fact, still running - # (raising ``n_steps`` is the documented way to extend a budget). It is - # re-derived at each chunk boundary below, so clear it here. - if stop_reason != "converged": - stop_reason = None + stop_reason = self._stop_reason_on_resume(stop_reason) # ``n_steps`` is the hard ceiling / max budget; ``stop_reason`` becomes # ``"converged"`` if the auto-convergence check stops the search early @@ -379,8 +405,7 @@ def batched_value_and_grad(params): # this loop just built (carrying ``stop_reason``, ``converged`` and # ``fom_history``) is what ``_fit`` returns and what that final # update is computed from. - is_final = converged or total_steps >= self.n_steps - if not is_final: + if not self._is_final_boundary(converged=converged, total_steps=total_steps): self.perform_update( model=model, analysis=analysis, diff --git a/test_autofit/non_linear/search/mle/test_multi_start_gradient.py b/test_autofit/non_linear/search/mle/test_multi_start_gradient.py index a93de1faa..460890171 100644 --- a/test_autofit/non_linear/search/mle/test_multi_start_gradient.py +++ b/test_autofit/non_linear/search/mle/test_multi_start_gradient.py @@ -387,58 +387,82 @@ def samples_for(n_starts, total_steps): assert samples.summary() is not None -def test__fit_skips_its_update_at_a_terminal_boundary(): - """``start_resume_fit`` performs the final update unconditionally the moment - ``_fit`` returns, so any update ``_fit`` emits at a terminal boundary is - duplicated work. - - Flipping ``during_analysis`` is *not* enough to avoid it: ``SearchUpdater`` +@pytest.mark.parametrize( + "converged, total_steps, is_final", + [ + (False, 50, False), # ordinary intermediate boundary -> update here + (False, 299, False), # still short of the ceiling + (False, 300, True), # ceiling reached + (False, 301, True), # overshoot is still terminal + (True, 50, True), # early convergence, well short of the ceiling + (True, 300, True), + ], +) +def test__is_final_boundary(converged, total_steps, is_final): + """``_fit`` must not emit a ``perform_update`` at a terminal boundary: + ``start_resume_fit`` performs the final update unconditionally the moment + ``_fit`` returns, so one here is duplicated work — and flipping + ``during_analysis`` does not avoid it, because ``SearchUpdater.update`` rebuilds the samples, recomputes the summary and re-runs likelihood - profiling on every call regardless of the flag. The update has to be skipped - outright, which is what the ``if not is_final`` guard does. + profiling on every call regardless of the flag. + + The rule is tested directly; ``_fit`` needs jax + optax + a JAX-traceable + ``Analysis`` and cannot be driven from this NumPy-only suite. + """ + search = af.MultiStartAdam(n_steps=300) + + assert ( + search._is_final_boundary(converged=converged, total_steps=total_steps) + is is_final + ) - ``_fit`` needs jax + optax + a JAX-traceable ``Analysis`` and cannot run - from this NumPy-only suite, so the guard is asserted at the source level. - Whitespace is normalised so reformatting does not cause a false failure. + +@pytest.mark.parametrize( + "restored, expected", + [ + ("converged", "converged"), # must survive: the loop guard reads it + ("max_steps", None), # stale — the run it described is over + (None, None), + ("some_future_reason", None), + ], +) +def test__stop_reason_on_resume(restored, expected): + """A resumed run inherits the previous run's ``stop_reason``. Keeping a + ``"max_steps"`` stamps every intermediate checkpoint of a search that is + still running as finished — which is what raising ``n_steps`` (the + documented way to extend a budget) produces. ``"converged"`` must survive, + because the step loop's guard uses it to refuse resuming a converged search + into further steps.""" + assert af.MultiStartAdam._stop_reason_on_resume(restored) == expected + + +def test__fit_uses_both_seams_and_keeps_the_converged_loop_guard(): + """Wiring guard for the two seams above, which are otherwise tested in + isolation and would keep passing if ``_fit`` stopped calling them. + + Necessarily a source-level check — ``_fit`` cannot run from this suite. It + pins the call sites and the loop condition; the *behaviour* of each seam is + pinned by the parametrised tests above, so a semantically-equivalent rewrite + of a seam is caught there rather than here. """ source = " ".join(inspect.getsource(af.MultiStartAdam._fit).split()) - # the update is guarded, and is intermediate-flavoured when it does run - assert "if not is_final: self.perform_update(" in source - assert "during_analysis=True" in source + assert "if not self._is_final_boundary(" in source + assert "stop_reason = self._stop_reason_on_resume(stop_reason)" in source + assert 'while total_steps < self.n_steps and stop_reason != "converged":' in source # the form that ran the whole final pass twice must not come back assert "during_analysis=not is_final" not in source # The framework's single final update is still there and unconditional — # without it this change would remove the final update rather than # de-duplicate it. ``start_resume_fit`` is decorated without - # ``functools.wraps``, so its body is sliced out of the module source rather - # than reached through the bound method. + # ``functools.wraps``, so its body is sliced out of the module source. module = inspect.getsource(abstract_search) body = module[module.index(" def start_resume_fit(") :] body = body[: body.index("\n def ", 1)] assert "during_analysis=False" in body -def test__fit_clears_a_stale_stop_reason_but_keeps_converged(): - """A resumed run inherits the previous run's ``stop_reason``. - - Resuming a search that finished with ``"max_steps"`` under a raised - ``n_steps`` (the documented way to extend a budget) used to leave that stale - reason in every intermediate checkpoint, so a search that was still running - reported itself finished to the aggregator and any results inspector. - - ``"converged"`` must survive the clear, because the ``while`` guard uses it - to refuse resuming a converged search into more steps — so that guard is - asserted on the ``while`` line specifically. Asserting the bare substring - would be satisfied by the clearing ``if`` itself and would pin nothing. - """ - source = " ".join(inspect.getsource(af.MultiStartAdam._fit).split()) - - assert 'if stop_reason != "converged": stop_reason = None' in source - assert 'while total_steps < self.n_steps and stop_reason != "converged":' in source - - def test__samples_info__reports_a_cleared_stop_reason_as_unfinished(): """The observable half of the fix: a checkpoint written mid-run carries no stop reason, so nothing downstream reads it as a finished search."""