fix: MultiStart gradient step loop crashes on a real iterations_per_full_update cadence

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,

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(