diff --git a/autofit/non_linear/search/abstract_search.py b/autofit/non_linear/search/abstract_search.py index a1b88883f..4c8744c84 100644 --- a/autofit/non_linear/search/abstract_search.py +++ b/autofit/non_linear/search/abstract_search.py @@ -1037,6 +1037,82 @@ 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). + + 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 + ---------- + 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") + 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/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 1df2b6d57..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,78 +156,39 @@ def __init__( self.logger.debug(f"Creating {self.optax_method} MultiStartGradient Search") - def _steps_in_chunk(self, steps_remaining: int) -> int: + def _is_final_boundary(self, converged: bool, total_steps: int) -> bool: """ - 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. + 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. """ - 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 bool(converged) or total_steps >= self.n_steps - return int(min(cadence, steps_remaining)) - - def _check_step_count(self, value, name: str): + @staticmethod + def _stop_reason_on_resume(stop_reason): """ - 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. + 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. """ - 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." - ) + return stop_reason if stop_reason == "converged" else None def _fit( self, @@ -348,13 +309,15 @@ def batched_value_and_grad(params): # ``resurrect`` is on); seeded independently of the broad-start draw. resurrect_rng = np.random.default_rng(1) + 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 # (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,16 +390,29 @@ 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 - self.perform_update( - model=model, - analysis=analysis, - during_analysis=not is_final, - 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. + if not self._is_final_boundary(converged=converged, total_steps=total_steps): + 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 818396550..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 @@ -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,105 @@ def samples_for(n_starts, total_steps): # The aggregator summary is computed without an IndexError. assert samples.summary() is not None + + +@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 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 + ) + + +@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()) + + 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. + 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__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..df57dde29 --- /dev/null +++ b/test_autofit/non_linear/search/test_steps_until_full_update.py @@ -0,0 +1,126 @@ +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__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 + + 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]) +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", + [af.MultiStartAdam, af.Emcee, af.Zeus, af.BFGS, af.BlackJAXNUTS], +) +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 + 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(search_cls._fit) + + 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