From 8543c545e48d2f93f125d540db4e82446c17b61c Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 30 Jul 2026 22:17:06 +0100 Subject: [PATCH] feat: per-step progress logging for the multi-start gradient searches The MultiStartGradient searches emitted two log lines for an entire run -- "Starting new ..." and "... sampling complete" -- leaving a user unable to tell a live fit from a hung one. Neither of the framework's progress channels reaches them: iterations_per_full_update makes the whole run a single chunk whose lone perform_update is (correctly) suppressed as duplicated work, and Fitness's quick-update counter is Python state mutated inside fitness.call, which these searches trace under jax.jit(jax.vmap(...)) -- so it runs once at trace time and never again. The samplers delegate to their own library's progress bar; a search that owns its step loop has to report for itself. Adds iterations_per_log (default 10) on AbstractMultiStartGradient, inherited by all four concrete searches. The line reports step, best log posterior, gain since the last line and live-start count, plus Prodigy's estim_lr -- its self-estimated step scale -- where the rule carries one. The first step always logs, since that line is what tells the user the XLA compile finished and stepping has begun. Two JAX-compile notices cover the single-point objective _broad_starts calls per draw and the vmapped/chunked one the step loop calls. Neither is reachable by the Fitness._jit/_vmap/_grad notices, because this search builds its transforms straight off fitness.call. All output is gated on the existing silence flag. Numerically inert: no traced operation is added, and the loop already forced a device sync every step. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Unu8t9xqScV93XinRmx1Do --- .../search/mle/multi_start_gradient/search.py | 202 ++++++++++++++++++ .../search/mle/test_multi_start_gradient.py | 171 +++++++++++++++ 2 files changed, 373 insertions(+) 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 31e84d405..1cb6d902f 100644 --- a/autofit/non_linear/search/mle/multi_start_gradient/search.py +++ b/autofit/non_linear/search/mle/multi_start_gradient/search.py @@ -42,6 +42,7 @@ def __init__( start_upper_limit: float = 0.85, resurrect: bool = False, convergence: Optional[MultiStartGradientConvergence] = None, + iterations_per_log: int = 10, initializer: Optional[AbstractInitializer] = None, iterations_per_full_update: int = None, iterations_per_quick_update: int = None, @@ -134,6 +135,29 @@ def __init__( when ``resurrect=True`` (the pixelized regime, whose best-fom climbs in breakthrough jumps that a plateau check would false-stop on); there the search leans on the ``n_steps`` ceiling. + iterations_per_log + How many steps pass between progress lines on the search log. The + step loop is otherwise entirely silent between "Starting new ..." and + "... sampling complete", which on a long fit leaves a user unable to + tell a live search from a hung one. + + Neither of the framework's usual progress channels reaches this + search, which is why it needs its own: ``iterations_per_full_update`` + defaults to the never-sentinel, so the whole run is a single chunk + whose lone ``perform_update`` is (correctly) suppressed as duplicated + work; and ``Fitness``'s quick-update counter is Python state mutated + inside ``fitness.call``, which is traced under ``jax.jit(jax.vmap(…))`` + here — it runs once at trace time and never again. The samplers dodge + the problem entirely by delegating to their own library's progress bar + (emcee ``progress=True``, dynesty ``print_progress``); a search that + owns its step loop has to report for itself. + + The first step always logs regardless of this cadence, because that + line is what tells the user the XLA compile finished and stepping has + begun — the longest unexplained wait in a real run. A log line rather + than a progress bar: ``n_steps`` is a ceiling that auto-convergence + routinely stops short of, so a bar's ETA would mislead, and lines + survive SLURM/HPC log capture where bars do not. """ super().__init__( @@ -162,6 +186,15 @@ def __init__( convergence if convergence is not None else MultiStartGradientConvergence() ) + # Rejected here rather than clamped, and at construction rather than at + # the first step: a sub-1 cadence makes ``total_steps % cadence`` either + # raise (0) or log every step (negative, since the modulo is never 0), + # and a user who mistyped the knob is better served by an immediate error + # than by a schedule they never asked for. Mirrors the reasoning on + # ``AbstractSearch._check_step_count``, which it reuses. + self._check_step_count(iterations_per_log, "iterations_per_log") + self.iterations_per_log = int(iterations_per_log) + self.logger.debug(f"Creating {self.optax_method} MultiStartGradient Search") def _is_final_boundary(self, converged: bool, total_steps: int) -> bool: @@ -198,6 +231,128 @@ def _stop_reason_on_resume(stop_reason): """ return stop_reason if stop_reason == "converged" else None + def _should_log_progress(self, total_steps: int, converged: bool) -> bool: + """ + Whether the step just completed should emit a progress line. + + Three ways to qualify: the first step (the compile is over and the search + is genuinely moving — the single most useful line in the run), every + ``iterations_per_log``-th step thereafter, and the step that converged + (which lands off-cadence far more often than not, and is the one step a + user most wants to see). + + ``silence`` suppresses all of them, matching how the samplers gate their + own progress output (e.g. Dynesty's ``print_progress=not self.silence``). + + A named predicate rather than an inline condition so the cadence rule can + be tested directly; ``_fit`` needs jax + optax + a JAX-traceable + ``Analysis`` and cannot be driven from the NumPy-only library suite. + """ + if self.silence: + return False + + return ( + bool(converged) + or total_steps <= 1 + or total_steps % self.iterations_per_log == 0 + ) + + def _compile_message(self, batched: bool) -> str: + """ + The notice logged immediately before a step that will trigger an XLA + compile, so the ensuing wait is explained rather than silent. + + Two distinct compiles block a fresh run, hence the ``batched`` switch: + the single-point ``value_and_grad`` that ``_broad_starts`` calls per draw, + and the vmapped (or ``batch_size``-chunked) one the step loop calls. Both + can run to minutes on a multi-band objective, and neither is reachable by + the ``Fitness._jit`` / ``_vmap`` / ``_grad`` compile notices — this search + builds its transforms straight off ``fitness.call``. + + Phrased to match the sampler-wide compile message rather than inventing a + second dialect for the same event. + """ + if not batched: + return ( + "JAX: jit compiling the single-point objective used to draw " + f"{self.n_starts} broad starts, could take seconds or minutes..." + ) + + over = ( + f"batches of {self.batch_size} starts" + if self.batch_size is not None + else f"all {self.n_starts} starts at once" + ) + + return ( + f"JAX: jit compiling the vmapped objective over {over}, " + "could take seconds or minutes..." + ) + + def _progress_message( + self, + total_steps: int, + best_fom: float, + previous_fom: Optional[float], + n_alive: int, + estim_lr=None, + ) -> str: + """ + The one-line progress report for a single step. + + Compact and pipe-delimited rather than prose: unlike the one-shot notices + this line repeats tens of times per run, and sentences would bury the + numbers that carry the signal. + + Reports the log posterior, not the raw figure of merit. ``best_fom`` is + ``-2 * log_posterior`` (what the search minimises), so the sign-and-halve + conversion here is the same one ``samples_via_internal_from`` applies — + a user comparing the log to their results should not have to do it in + their head. + + Parameters + ---------- + total_steps + Steps completed, reported against the ``n_steps`` ceiling. + best_fom + The global best figure-of-merit. Non-finite until some start finds a + finite basin, which is reported plainly rather than as ``inf``. + previous_fom + The global best at the previous progress line, used for the gain + since. ``None`` on the first line, and skipped while either end of + the comparison is non-finite (a first finite basin is an arrival, not + an improvement, and ``inf - x`` would print as ``inf``). + n_alive + Starts whose objective is currently finite, against ``n_starts`` — + the population-health signal that a falling best-fom alone hides. + estim_lr + Per-start self-estimated step scale, for the learning-rate-free rules + that carry one (Prodigy's ``d``). ``None`` for the fixed-rate Adam + family, whose optimizer state has no such field; the field is then + omitted rather than rendered empty. + """ + parts = [f"{self.optax_method} step {total_steps}/{self.n_steps}"] + + if np.isfinite(best_fom): + parts.append(f"best log_post {-0.5 * best_fom:.4f}") + + if previous_fom is not None and np.isfinite(previous_fom): + # fom is -2*log_posterior, so a fom drop of d is a log-posterior + # gain of d/2. + parts.append(f"gained {0.5 * (previous_fom - best_fom):.4f}") + else: + parts.append("best log_post — (no finite basin yet)") + + parts.append(f"alive {n_alive}/{self.n_starts}") + + if estim_lr is not None: + lr = np.asarray(estim_lr, dtype=float) + parts.append( + f"d {np.nanmin(lr):.2e} / {np.nanmedian(lr):.2e} / {np.nanmax(lr):.2e}" + ) + + return " | ".join(parts) + def _fit( self, model: AbstractPriorModel, @@ -309,6 +464,9 @@ def batched_value_and_grad(params): except (FileNotFoundError, TypeError, KeyError): + if not self.silence: + self.logger.info(self._compile_message(batched=False)) + params = self._broad_starts( model=model, value_and_grad_single=jax.jit(_value_and_grad), @@ -337,6 +495,18 @@ def batched_value_and_grad(params): stop_reason = self._stop_reason_on_resume(stop_reason) + # The vmapped objective compiles on its first call below, not here, so + # the notice is emitted at the call site the first time round. Resumed + # runs pay it too (a fresh process has no live trace cache), which is why + # the flag is set here rather than only on the fresh path. + awaiting_compile = True + + # The global best at the previous progress line, so each line can report + # the gain since the last one rather than since the previous *step* — + # over a cadence of ``iterations_per_log`` steps that is the number a + # user can actually act on. + previous_logged_fom = 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). @@ -356,6 +526,11 @@ def batched_value_and_grad(params): converged = False for _ in range(iterations): + if awaiting_compile: + if not self.silence: + self.logger.info(self._compile_message(batched=True)) + awaiting_compile = False + foms, grads = batched_value_and_grad(params) alive = np.isfinite(np.asarray(foms)) @@ -397,6 +572,33 @@ def batched_value_and_grad(params): fom_history ): converged = True + + # Logged before the ``converged`` break so the terminating step + # reports itself; it is the one step a user most wants to see and + # it lands off-cadence more often than not. + if self._should_log_progress( + total_steps=total_steps, converged=converged + ): + # Prodigy et al. carry a self-estimated step scale; the Adam + # family's state has no such field and ``tree_get`` returns + # None, which the message renders by omission. Read only on a + # logging step: it is an (n_starts,) device->host copy, which + # merely pulls forward the sync the next iteration's + # ``np.asarray(foms)`` would force anyway. + estim_lr = optax.tree_utils.tree_get(opt_state, "estim_lr") + + self.logger.info( + self._progress_message( + total_steps=total_steps, + best_fom=best_fom, + previous_fom=previous_logged_fom, + n_alive=int(alive.sum()), + estim_lr=estim_lr, + ) + ) + previous_logged_fom = best_fom + + if converged: break if converged: 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 f2d8e98d7..359aa931e 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 @@ -469,6 +469,177 @@ def test__stop_reason_on_resume(restored, expected): assert af.MultiStartAdam._stop_reason_on_resume(restored) == expected +@pytest.mark.parametrize( + "total_steps, converged, silence, should_log", + [ + (1, False, False, True), # first step: the compile is over, we are moving + (2, False, False, False), # off-cadence + (10, False, False, True), # on-cadence + (20, False, False, True), + (23, False, False, False), + (23, True, False, True), # converged off-cadence still reports + (10, False, True, False), # silence suppresses a cadence hit + (1, False, True, False), # ...and the first step + (23, True, True, False), # ...and convergence + ], +) +def test__should_log_progress(total_steps, converged, silence, should_log): + """The cadence rule. The first step always logs because that line is what + tells a user the XLA compile finished and stepping has begun — the longest + unexplained wait in a real run. The converged step always logs because it + lands off-cadence more often than not.""" + search = af.MultiStartAdam(n_steps=300, iterations_per_log=10, silence=silence) + + assert ( + search._should_log_progress(total_steps=total_steps, converged=converged) + is should_log + ) + + +@pytest.mark.parametrize("bad", [0, -1, 2.5]) +def test__iterations_per_log__rejects_a_cadence_that_cannot_schedule(bad): + """Rejected at construction, not clamped: ``total_steps % 0`` raises and a + negative cadence is never 0, so it would log every step. A user who mistyped + the knob is better served by an error than by a schedule they never asked + for.""" + with pytest.raises(ValueError, match="iterations_per_log"): + af.MultiStartAdam(iterations_per_log=bad) + + +def test__progress_message__reports_log_posterior_not_the_raw_fom(): + """``best_fom`` is ``-2 * log_posterior`` (what the search minimises), so the + line must apply the same sign-and-halve conversion as + ``samples_via_internal_from`` — otherwise the number in the log cannot be + compared with the number in the results.""" + search = af.MultiStartAdam(n_starts=48, n_steps=300) + + message = search._progress_message( + total_steps=30, + best_fom=-63575.6876, + previous_fom=-63571.4676, + n_alive=46, + ) + + assert "adam step 30/300" in message + assert "best log_post 31787.8438" in message + # a fom drop of 4.22 is a log-posterior gain of 2.11 + assert "gained 2.1100" in message + assert "alive 46/48" in message + # no learning-rate-free state was supplied, so the field is omitted entirely + assert " d " not in message + + +def test__progress_message__includes_the_prodigy_step_scale_when_present(): + """Prodigy's per-start ``estim_lr`` (its ``d``) is the one genuinely + learning-rate-free diagnostic, and is otherwise invisible: min/median/max + across starts says whether it has finished ramping or is still climbing.""" + search = af.MultiStartProdigy(n_starts=4, n_steps=300) + + message = search._progress_message( + total_steps=10, + best_fom=-100.0, + previous_fom=-90.0, + n_alive=4, + estim_lr=np.array([1.0e-2, 2.0e-2, 4.0e-2, 8.0e-2]), + ) + + assert "prodigy step 10/300" in message + assert "d 1.00e-02 / 3.00e-02 / 8.00e-02" in message + + +@pytest.mark.parametrize( + "best_fom, previous_fom, expect_gain", + [ + (np.inf, None, False), # nothing has found a finite basin yet + (-100.0, None, False), # first line: no previous to compare against + (-100.0, np.inf, False), # arrival at a first basin is not a "gain" + (-100.0, -90.0, True), + ], +) +def test__progress_message__non_finite_and_first_line_edges( + best_fom, previous_fom, expect_gain +): + """``inf - x`` would print as ``inf`` and a bare ``inf`` best-fom would print + as a figure of merit the user cannot act on, so both are reported plainly + instead.""" + search = af.MultiStartAdam(n_starts=8, n_steps=300) + + message = search._progress_message( + total_steps=10, + best_fom=best_fom, + previous_fom=previous_fom, + n_alive=8, + ) + + assert ("gained" in message) is expect_gain + assert "inf" not in message + + if not np.isfinite(best_fom): + assert "no finite basin yet" in message + + +@pytest.mark.parametrize( + "batch_size, expected", + [ + (None, "all 16 starts at once"), + (4, "batches of 4 starts"), + ], +) +def test__compile_message__names_what_is_being_compiled(batch_size, expected): + """Two distinct compiles block a fresh run — the single-point objective + ``_broad_starts`` calls per draw, and the vmapped/chunked one the step loop + calls. Neither is reachable by the ``Fitness._jit`` / ``_vmap`` / ``_grad`` + notices, because this search builds its transforms straight off + ``fitness.call``.""" + search = af.MultiStartProdigy(n_starts=16, batch_size=batch_size) + + single = search._compile_message(batched=False) + assert "single-point objective" in single + assert "16 broad starts" in single + assert "could take seconds or minutes" in single + + batched = search._compile_message(batched=True) + assert "vmapped objective" in batched + assert expected in batched + assert "could take seconds or minutes" in batched + + +def test__dict_round_trip__iterations_per_log(): + """The cadence must survive serialisation so a resumed search keeps + reporting at the rate the user chose.""" + restored = from_dict(to_dict(af.MultiStartProdigy(iterations_per_log=25))) + + assert isinstance(restored, af.MultiStartProdigy) + assert restored.iterations_per_log == 25 + + +def test__fit_wires_the_progress_and_compile_seams(): + """Wiring guard for the three progress seams, 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. + """ + source = " ".join(inspect.getsource(af.MultiStartAdam._fit).split()) + + # the cadence gate and the line it guards + assert ( + "if self._should_log_progress( total_steps=total_steps, converged=converged )" + in source + ) + assert "self._progress_message(" in source + + # both compile notices, each emitted before the call that triggers its compile + assert "self._compile_message(batched=False)" in source + assert "self._compile_message(batched=True)" in source + + # the step scale is read from the optimizer state, not recomputed + assert 'optax.tree_utils.tree_get(opt_state, "estim_lr")' in source + + # the converged step must still be able to report before the loop breaks: + # setting `converged` and breaking must not be fused back together. + assert "converged = True break" not in source + + 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.