feat: per-step progress logging for the multi-start gradient searches

[Jammy2211]

Overview

af.MultiStartProdigy / MultiStartAdam / MultiStartADABelief / MultiStartLion emit exactly two log lines for an entire run — "Starting new <rule> MultiStartGradient search (N starts, ...)" and, however many hours later, "<rule> MultiStartGradient sampling complete." Nothing in between, so a user cannot tell a live 300-step fit from a hung one.

Both standard autofit progress channels are inert here: iterations_per_full_update defaults to 1e99, so the whole run is one chunk whose single perform_update is (correctly) suppressed by _is_final_boundary; and Fitness.call's quick-update counter never fires because fitness.call is traced inside jax.jit(jax.vmap(...)), so its Python-side counting runs once at trace time. Other searches sidestep this by delegating to their sampler's own progress bar (emcee progress=True, dynesty print_progress); the gradient searches own their step loop and have no such fallback.

The auto-convergence arc (#1406 / #1409) already ships fom_history in samples_info plus the figure_of_merit_vs_iteration plotter — but that is the post-hoc view. You can see the descent after the run, never during it. This issue closes that gap.

Plan

• Add a cadence-controlled progress log line to AbstractMultiStartGradient's step loop, gated on the existing silence flag.
• Always log the first step regardless of cadence — that line is what tells the user the long XLA compile finished and stepping has begun, which is the real "is it hung?" moment.
• Report what is already on the host each step: step/ceiling, best log posterior, improvement since the last line, and live-start count.
• For Prodigy, additionally report its self-estimated step scale d (min/median/max across starts) — invisible today, and the one thing that says whether it is still ramping.
• Keep it a plain constructor kwarg — no new packaged config key — so no workspace can KeyError on it.
• Split the logic into pure, numpy-only seams so the library suite can test it without JAX, plus a source-level wiring guard.

Detailed implementation plan

Affected Repositories

• PyAutoFit (primary, and only)

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoFit	main	clean (0 ahead / 0 behind origin)

No worktree claim on PyAutoFit — the three tasks in active.md claim PyAutoLens, autolens_workspace_test and PyAutoMind.

Suggested branch: feature/multi-start-gradient-progress-logging
Worktree root: ~/Code/PyAutoLabs-wt/multi-start-gradient-progress-logging/
Work Classification: Library

Implementation Steps

1. __init__ (autofit/non_linear/search/mle/multi_start_gradient/search.py) — new iterations_per_log: int = 10, stored as self.iterations_per_log. Name follows the house iterations_per_*_update convention. Carried through to_dict/from_dict automatically, like resurrect / batch_size.

2. New pure seam _should_log_progress(self, total_steps, converged) -> bool — true on total_steps == 1, on total_steps % self.iterations_per_log == 0, or on converged; false when self.silence.

3. New pure seam _progress_message(self, total_steps, best_fom, previous_fom, n_alive, estim_lr=None) -> str — formats e.g.

   prodigy step 30/300 | best log_post 31787.8438 | improved 2.11 | alive 46/48 | d 1.2e-02 / 3.0e-02 / 4.8e-02
   

   log_post = -0.5 * best_fom, the same convention as samples_via_internal_from (search.py:628). Non-finite best_fom (no basin found yet) renders as best log_post —. estim_lr=None drops the d field entirely.

4. Wire into the step loop — after fom_history.append(best_fom), before the resurrect block: call the two seams and self.logger.info(...). estim_lr comes from optax.tree_utils.tree_get(opt_state, "estim_lr"), pulled only on a logging step, converted via np.asarray.

   Verified against the installed optax that tree_get reaches estim_lr through the apply_if_finite wrapper: returns a per-start (n_starts,) array for prodigy (initialised at 1e-6, ramps up), and None for adam — so Adam / ADABelief / Lion get the generic line with no special-casing.

Cost

The loop already forces a device sync every step at search.py:361 (np.isfinite(np.asarray(foms))), so the fom / alive / step fields are free. The estim_lr pull is an extra (n_starts,) device→host transfer on 1-in-iterations_per_log steps, which forces the optimizer-state sync marginally earlier than the next iteration would anyway. Not literally zero, but negligible and constant — and it adds no traced operation, so no XLA recompile.

Tests (test_autofit/non_linear/search/mle/test_multi_start_gradient.py, numpy-only — no JAX)

• Parametrised _should_log_progress: first step, cadence hits/misses, converged=True off-cadence, silence=True suppresses everything.
• _progress_message: with and without estim_lr; the -0.5 * sign convention; non-finite best_fom.
• Round-trip: iterations_per_log survives to_dict / from_dict.
• Source-level wiring guard in the style of the existing test__fit_uses_both_seams_and_keeps_the_converged_loop_guard, pinning both call sites so the seams cannot be silently orphaned.

Beyond the unit suite, the numpy-only tests cannot prove the line fires in a live JAX run — verify by running the searches/mle.py smoke script in autofit_workspace with a Prodigy search and pasting the real output.

Key Files

• autofit/non_linear/search/mle/multi_start_gradient/search.py — AbstractMultiStartGradient.__init__ + _fit step loop; the two new seams.
• test_autofit/non_linear/search/mle/test_multi_start_gradient.py — the numpy-only seam tests + wiring guard.

Scope notes

• The knob lands on AbstractMultiStartGradient, so all four subclasses get the line, not Prodigy alone. That is the right fix — the silence is a base-class problem — but it is a shared-search change, not a Prodigy-only tweak.
• Library-only; no autofit_workspace_test PR. Both auto-convergence phases shipped a jax_assertions/ validation script, so there is precedent for a second leg. Deliberately declined here: a permanent JAX regression net for a log line's wording is a maintenance cost out of proportion to the risk, and the source-level wiring guard already catches the orphaning failure mode.

Sizing

Brain sized this large (score 8); overridden to small, no phase split. Every point came from prompt prose — 521 words (+3), the words dynesty / emcee / gradient / jax / sampler (+3), jax / vmap (+1), memory-context-required (+1) — while repos_affected=1 and architectural_risk=[] contributed nothing. The heuristic penalised a well-researched prompt rather than measuring scope.

Original Prompt

Click to expand starting prompt

# MultiStartGradient runs silently for their whole duration — no per-step progress

Type: feature
Target: autofit
Repos:
- PyAutoFit
Difficulty: small
Autonomy: supervised
Priority: medium
Status: formalised

## Original request (verbatim)

> might not be feasible but if this could have a verbose update would be nice:
> 2026-07-30 19:45:31,924 - start_here - INFO - Starting new prodigy MultiStartGradient search (48 starts, no previous samples found)
> e.g. can Prodigy display its updates a bit more so a user knows it progressing

## Problem

`AbstractMultiStartGradient._fit` emits exactly two log lines for an entire run:
`"Starting new <rule> MultiStartGradient search (N starts, ...)"` and, however many
hours later, `"<rule> MultiStartGradient sampling complete."` Nothing in between.
A user cannot tell a live 300-step fit from a hung one.

Both of the standard autofit progress channels are inert for this search:

- `iterations_per_full_update` defaults to `1e99` (`autofit/config/general.yaml:3`),
  so `_steps_until_full_update` returns the whole budget — one chunk, one
  `perform_update`, which `_is_final_boundary` then correctly suppresses as
  duplicated work. Zero intermediate output by construction.
- `Fitness.call`'s quick-update counter never fires: `fitness.call` is traced inside
  `jax.jit(jax.vmap(...))` (`search.py:260-261`), so the Python-side counting runs
  once at trace time, never per step.

Other searches sidestep this by delegating to their sampler's own progress bar
(emcee `progress=True`, dynesty `print_progress=not self.silence`). The gradient
searches own their step loop, so they have no such fallback.

## Why it is cheap

The step loop (`search.py:358-400`) is plain Python and **already forces a device
sync every step** at line 361 (`np.isfinite(np.asarray(foms))`). Everything a
progress line wants is already materialised on the host each step — `total_steps`
vs `self.n_steps`, `best_fom` (best log posterior is `-0.5 * best_fom`), this
step's `foms_np[best_index]`, `alive.sum()` live starts, `n_resurrections`. No
extra sync, no recompile, no new device work.

## Ask

Add a cadence-controlled progress log line to `AbstractMultiStartGradient`'s step
loop, gated on the existing `silence` flag.

Prodigy-specific extra: `optax.contrib.prodigy`'s state carries `estim_lr` — its
`d`, the self-estimated step scale — confirmed against the installed optax
(`ProdigyState(exp_avg, exp_avg_sq, grad_sum, params0, estim_lr,
numerator_weighted, count)`). It is reachable through the `apply_if_finite`
wrapper via `optax.tree_utils.tree_get(opt_state, "estim_lr")`, the same helper
the resume path already uses at `search.py:298`. State is vmapped per-start, so
this gives min/median/max across the 48 starts — the one genuinely
Prodigy-specific diagnostic that is currently invisible (has `d` finished ramping,
or is it still climbing?). Adam/ADABelief/Lion have no such field; `tree_get`
returns `None` and they get the generic line.

## Constraints / decisions to make

- **Log line, not a tqdm bar.** tqdm is installed only transitively (not declared
  in `pyproject.toml`), a bar's ETA would mislead because auto-convergence makes
  `n_steps` a ceiling rather than a target, and progress lines survive SLURM/HPC
  log capture where bars do not.
- The knob lands on `AbstractMultiStartGradient`, so it applies to all four
  subclasses (Adam / ADABelief / Lion / Prodigy), not Prodigy alone. That is the
  right scope — the silence is a base-class problem — but it is a shared-search
  change, not a Prodigy-only tweak.
- Must not perturb numerics, add a device sync, or trigger an XLA recompile.
- Cadence default should be sane for both a 300-step CPU run and a long GPU run.

[Jammy2211]

Scope extension: the compile wait is in scope too

Human-approved addition to the plan above.

An adjacent in-flight task (draft/bug/autofit/sampler_cli_output_actual_numbers.md) covers the JAX-compile-wait message — "JAX jit compiling likelihood function, could take seconds or minutes..." — scoped to autofit/non_linear/fitness.py's _jit / _vmap / _grad wrappers.

Those wrappers are exactly what the multi-start gradient searches do not use. _fit builds its transforms directly off the plain method:

_value_and_grad = jax.value_and_grad(fitness.call)   # search.py:260
_vmapped = jax.jit(jax.vmap(_value_and_grad))        # search.py:261

fitness.call, not fitness._jit. So that task's fix would ship without ever reaching MultiStartProdigy / MultiStartAdam / MultiStartADABelief / MultiStartLion, leaving the longest unexplained wait in a real run still silent. No file collision between the two tasks — the gap is one of coverage, not of merge.

Added to this issue's scope

There are two distinct compiles that block before any progress can be reported, and both need a notice:

1. The single-point objective, compiled on the first call inside _broad_starts (search.py:312-316, via value_and_grad_single=jax.jit(_value_and_grad)). Log before drawing the starts, naming n_starts.
2. The vmapped (or batch_size-chunked) objective, compiled on the first batched_value_and_grad(params) call in the step loop (search.py:359). Log before the first evaluation only — a one-shot, not per-step.

Wording should follow the sibling task's so the two read consistently, e.g.:

JAX: compiling the vmapped objective (48 starts) — could take seconds or minutes...

Both notices are gated on the existing silence flag, like the step line. Together with the always-log-first-step rule already in the plan, this gives the user an unbroken trail: compile started → compile finished, stepping → step 10, 20, 30… → converged.

Test impact

The compile notices are emitted from _fit, which the numpy-only suite cannot run — so they are covered the same way as the existing seams: pure formatting helpers unit-tested directly, plus the source-level wiring guard extended to pin the new call sites.

[Jammy2211]

Smoke tests — autofit_workspace

PR #1435 is open. Workspace-impact analysis came out (iii) no workspace impact: the change is additive with a default, no symbol is removed or renamed, and no script references the new knob. The searches guides document these searches but not their sibling knobs (batch_size, resurrect, convergence are all absent), so adding iterations_per_log there would break that precedent.

Workspace	Passed	Failed	Total
autofit_workspace	8	0	8

PASS  overview/overview_1_the_basics.py
PASS  searches/mcmc.py
PASS  searches/nest.py
PASS  searches/mle.py          <- the MultiStart gradient guide
PASS  cookbooks/model.py
PASS  cookbooks/result.py
PASS  cookbooks/samples.py
PASS  cookbooks/latent_variables.py

What this run does and does not prove

Stated explicitly so the evidence is not over-read:

• Does prove: inserting iterations_per_log into the constructor breaks no downstream script. Every caller uses keyword arguments (checked across all workspaces), so the new parameter shifts nothing positionally, and searches/mle.py constructs and runs MultiStartAdam end-to-end.
• Does not prove: that the progress lines fire. The smoke profile sets PYAUTO_DISABLE_JAX=1 and PYAUTO_TEST_MODE=2, and mle.py carries no ENV: jax release token — so the sampler is bypassed (TEST MODE 2 ... Skipping sampler, calling likelihood function once) and _fit never executes.

The step loop is therefore validated separately, by the real JAX MultiStartProdigy run recorded in the PR body, plus 22 NumPy-only unit tests covering the cadence rule, message formatting, sign convention, non-finite edges, compile notices, dict round-trip and a source-level wiring guard.

Coverage bounded

Only autofit_workspace was run, not all six workspaces. The changed code lives in a search that only the autofit-level guides construct directly, and the change is log-output-only. Naming the bound rather than implying full coverage.

Gate

• Library suite: 1614 passed, 1 skipped
• Heart: YELLOW (score 75, no RED) on three pre-existing reasons unrelated to this change — manifest drift (tenant firewall vs repos.yaml), test run status unknown (no report.json), and release validation stale across 5 repos. Human-acknowledged before ship.
• Diff: pure additions, 373 insertions, 0 deletions.

[Jammy2211]

Shipped

PR #1435 merged (09a134d8). The multi-start gradient searches now report themselves while they run.

Before, a run emitted two log lines hours apart. Now:

JAX: jit compiling the single-point objective used to draw 16 broad starts, could take seconds or minutes...
Starting new prodigy MultiStartGradient search (16 starts, no previous samples found).
JAX: jit compiling the vmapped objective over all 16 starts at once, could take seconds or minutes...
prodigy step 1/120   | best log_post -3335.7342 | alive 16/16 | d 1.00e-06 / 1.00e-06 / 1.00e-06
prodigy step 40/120  | best log_post -1743.4656 | gained 1574.6372 | alive 15/16 | d 2.56e+00 / 4.86e+00 / 5.07e+00
prodigy step 120/120 | best log_post   -58.8587 | gained    0.0647 | alive 16/16 | d 2.56e+00 / 7.70e+00 / 1.19e+01

d is Prodigy's estim_lr — its self-estimated step scale, per-start. Watching it ramp 1e-06 → 1.19e+01 and then flatten is the one learning-rate-free diagnostic that had no way of reaching the user.

Gate

• Library suite 1614 passed, 1 skipped; 22 new NumPy-only unit tests
• CI green on both matrix legs — unittest (3.12), unittest (3.13), docs-build
• autofit_workspace smoke 8/8
• Heart YELLOW (score 75, no RED) on three pre-existing unrelated reasons, human-acknowledged
• Diff pure additions: 373 insertions, 0 deletions

Traps worth keeping

• The smoke test cannot validate this. PYAUTO_DISABLE_JAX=1 + PYAUTO_TEST_MODE=2 are smoke-profile defaults and searches/mle.py carries no ENV: jax token, so the sampler is bypassed and _fit never runs. A green smoke here proves the new kwarg breaks no caller — nothing more. The step loop needed a real JAX run.
• These searches bypass Fitness._jit / _vmap / _grad entirely, building jax.value_and_grad(fitness.call) directly (search.py:260-261). Any framework-level JAX-compile or per-iteration instrumentation added to those wrappers will silently miss all four multi-start searches. That is why the compile notices had to live in _fit, and it is reported on fix: sampler CLI prints variable name, JAX compile msg fires early #1434.
• No new packaged config key. A key in the shared updates: block would imply every search honours it and would KeyError in any workspace shadowing that section — the trap from feat(autofit): multi-start gradient convergence results contract (phase 2) #1409. A plain constructor kwarg avoids both.
• optax.tree_utils.tree_get(opt_state, "estim_lr") reaches through the apply_if_finite wrapper, returning a per-start (n_starts,) array for prodigy and None for adam. Verified against the installed optax before being designed in, rather than assumed.
• Sizing heuristic misfires on well-researched prompts. The Brain scored this large (8) and wanted a phase split; every point came from prompt prose (521 words +3, the words dynesty/emcee/gradient/jax/sampler +3, jax/vmap +1, memory-context +1) while repos_affected=1 and architectural_risk=[] contributed nothing. fix: sampler CLI prints variable name, JAX compile msg fires early #1434 hit the identical misfire the same evening.

Concurrency

Ran in parallel with #1434 under a human-authorized concurrent claim on PyAutoFit — disjoint files (multi_start_gradient/search.py here, abstract_search.py + fitness.py there), both branched off a50ba95b0, which had not moved at merge.