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 6c8ea48b4..1df2b6d57 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,79 @@ 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, @@ -281,9 +354,7 @@ def batched_value_and_grad(params): while total_steps < self.n_steps and stop_reason != "converged": steps_remaining = self.n_steps - total_steps - iterations = min( - self.iterations_per_full_update or self.n_steps, steps_remaining - ) + iterations = self._steps_in_chunk(steps_remaining) # Convergence is assessed every step (``fom_history`` updates every # step) rather than only at the ``iterations_per_full_update`` boundary, 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 59ea78630..818396550 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 @@ -1,3 +1,5 @@ +import inspect + import numpy as np import pytest @@ -136,6 +138,108 @@ 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(