[Jammy2211]

Summary

Fixes all seven findings of the post-merge astra review 05 of the search-extensibility A2/A3/A3b PRs (#1679/#1680/#1681), which shipped in 2026.10.11.1 / 2026.10.12.1. One commit per item; every fix has a test that was confirmed failing on untouched main (911fe05) first.

1. Old Fitness pickle silently enabled JIT (F1, P1) — 5402c7619

Finding. Fitness.__setstate__ popped the legacy use_jax_jit flag and defaulted compile=True, so an A1-era eager checkpoint resumed jitted (a likelihood valid only eagerly then raised TracerArrayConversionError).
Fix. The legacy flags map onto compile exactly: vmap -> compiled batched, jit -> compiled scalar, neither -> eager scalar. A pickle predating both flags also dispatched eagerly, so it stays eager.
Tests. test_fitness_objective.py::test_legacy_eager_pickle_stays_eager (hand-built A1 __dict__, use_jax_jit=False, eager-only likelihood evaluates to 50.0 and the restored _call has no __wrapped__), test_legacy_jit_and_vmap_pickles_stay_compiled.

2. Preflight validated the declared config, not the resolved one (F3) — 1f33755f5

Finding. Reproduced both false rejections on main: DynestyStatic(use_jax_jit=False) + eager likelihood -> TracerArrayConversionError; MultiStartAdam(gradient_mode="reverse") + forward-declared analysis + custom_vjp -> "can't apply forward-mode autodiff".
Fix. New preflight.ExecutionSettings (batched, compile, uses_gradient, gradient_mode, fitness_overrides), resolved once by NonLinearSearch.execution_settings(analysis) from the search's fitness_overrides and its gradient_mode override (when honours_gradient_mode). The preflight traces exactly those kinds/modes (an eager scalar execution is not traced; grad kinds always are), and the run(ctx) bridge, Dynesty (use_jax_jit now flows through AbstractDynesty.fitness_overrides) and MultiStartGradient build their Fitness from the same object.
Tests. test_preflight.py: eager Dynesty passes preflight and fit reaches the backend; the reverse override passes while the un-overridden search still fails; Nautilus(use_jax_vmap=False) traces 'scalar'; a jitted search still rejects an untraceable likelihood.

3. .completed written before the archive finalizes; NSS discarded its checkpoint early (F2, P1) — edcf35116

Finding. start_resume_fit wrote .completed right after save_results; the archive (output_search_internal) ran later in post_fit_output. A failed archive write left a "completed" fit whose rerun never retried. NSS also deleted nss_checkpoint.pkl inside run().
Fix. New NonLinearSearch._complete_fit: finalize retained state (archive, then the resume-state discard it supersedes; or remove search_internal) and only then write .completed. Used by the real fit and the test-mode bypass; post_fit_output no longer re-finalizes a fit that just ran (it still zips, and still tidies an already-completed folder). NSS no longer discards in run(); the generic discard runs after the archive. .completed is still inside the zip (zip happens after the marker).
Tests. test_checkpoint.py::test_a_failed_archive_write_leaves_the_fit_incomplete_and_a_rerun_finishes_it (real DirectoryPaths on disk, Drawer whose archive raises OSError("simulated disk full") once: no marker after the failure; the rerun finalizes, archive and .completed are in the zip); nss/test_checkpoint.py::test__run_keeps_its_checkpoint_until_the_result_is_written.

4. Direct _fit built PoolFactory(self) without the analysis (F4) — 86299f6e8

Finding. Reproduced: Nautilus(number_of_cores=2)._fit(model, jax_analysis) chose fit_multiprocessing with is_jax=False, cores 2.
Fix. _pools(analysis=None) builds PoolFactory(self, analysis) when no fit-scoped factory exists; the bridge passes the analysis. The same latent gap was closed where the analysis is in hand: Dynesty's core count, the initializer's n_cores (fitness.analysis) and make_sneaky_pool.
Test. test_fork_rule.py::test_direct_fit_entry_applies_the_fork_rule_to_a_jax_analysis (fork_context patched to raise; asserts the single-process path is taken).

5. FitContext leaked pools on success; ctx.update never saved the checkpoint (F5, F6) — 8f35343de

Finding. Reproduced both: a toy run(ctx) search that calls ctx.pool() got 0 close/join on success; a resumable toy search interrupted after ctx.update had 0 checkpoint saves.
Fix. PoolFactory.close(failed) closes+joins on success and terminates+joins on failure (idempotent, tolerant of pools a backend already shut down); ctx.close calls it on both exits. ctx.update calls resume_state.save(paths, internal) (never finalize) before perform_update; native resume states keep their no-op save; a search without a resume state saves nothing mid-run. docs/design/run_ctx.md member notes updated to match.
Tests. test_fit_context.py: success -> created, close, join; failure -> created, terminate, join; idempotent close; interrupted resumable toy search -> one save, no finalize, file on disk; non-resumable search saves nothing mid-run.

6. Shape-only golden diagnostics (F7) — 2090ca9dd

Finding. Reproduced: replacing any BACKEND_DIAGNOSTIC_KEYS entry with a shape-preserving -999 passed test_samples_info_is_unchanged("nuts") (8/8 mutations passed).
Fix. Diagnostics are compared by value to the fixture (DIAGNOSTIC_RTOL = 0.25, finite, positive), exactly against the installed blackjax estimators on the stored chains, and through deterministic estimator doubles whose outputs must be forwarded unchanged (ownership). The CSV assert_allclose(rtol=1e-12, atol=0) is unchanged.
Tests. test_sentinel_diagnostics_fail_the_golden_comparison[<key>] x8, test_diagnostics_are_the_installed_blackjax_reference, test_diagnostics_forward_the_backend_outputs.

7. Witness: disk-backed checks in a writable environment

Old-era folders were produced by checking out tag 2026.10.7.1 (710f4b3) in a detached scratch worktree (git worktree add <scratchpad>/old_wt 2026.10.7.1) and running a tiny Gaussian fit with that PyAutoFit first on PYTHONPATH (its source __version__ prints the stale 2026.8.17.1 stamp; the code is the tag's). The worktree was removed afterwards. Script: fit.py <dynesty|emcee> <name>, DynestyStatic(nlive=50, iterations_per_full_update=1500) / Emcee(nwalkers=20, nsteps=400); CRASH_AFTER=N hard-kills the process with os._exit(9) after N likelihood calls (no cleanup runs).

(a) Pre-release output folders loaded with the new code

$ PYTHONPATH=<old_wt>:... python fit.py dynesty old_dynesty ; ... emcee old_emcee     # 2026.10.7.1
$ python load_old.py                                                                   # new code, Aggregator.from_directory
 A total of 2 search_outputs and results were found.
old_dynesty | samples 433 | max logL -36.1578 | search DynestyStatic
old_emcee | samples 200 | max logL -36.1788 | search Emcee
   internal file: search_internal.hdf
$ python fit.py dynesty old_dynesty ; python fit.py emcee old_emcee                   # new code, same identifiers
output output/witness/old_dynesty/d75351d847c453ff56e2a64b2d071d6f
old_dynesty - INFO - Fit Already Completed: skipping non-linear search.
is_complete True
likelihood calls this process 0
total samples 433
max log likelihood -36.1578
output output/witness/old_emcee/5db7e684e31400856e9c74a4a1e0521f
old_emcee - INFO - Fit Already Completed: skipping non-linear search.
is_complete True
likelihood calls this process 0
total samples 200
max log likelihood -36.1788

(The old run reported 929 Dynesty samples in-process; samples.csv drops weights < 1e-10, hence 433 on load. Same value on old and new.)

(b) Interrupt and resume (Dynesty and Emcee), new code, and old-era interrupted runs resumed by the new code

=== [dynesty] new code: interrupted run (CRASH_AFTER=4000) ===
kill_new_dynesty - INFO - Starting new Dynesty non-linear search (no previous samples found).
exit=9
=== [dynesty] new code: rerun ===
kill_new_dynesty - INFO - Resuming Dynesty non-linear search (previous samples found).
is_complete True
likelihood calls this process 895
total samples 966
max log likelihood -36.1621
=== [dynesty] OLD code (2026.10.7.1): interrupted run ===
kill_old_dynesty - INFO - Starting new Dynesty non-linear search (no previous samples found).
is_complete True
likelihood calls this process 3620
total samples 899
max log likelihood -36.1496
exit=0
=== [dynesty] new code: resume the old-era interrupted run ===
kill_old_dynesty - INFO - Fit Already Completed: skipping non-linear search.
is_complete True
likelihood calls this process 0
total samples 415
max log likelihood -36.1496
=== [emcee] new code: interrupted run (CRASH_AFTER=4000) ===
kill_new_emcee - INFO - Starting new Emcee non-linear search (no previous samples found).
exit=9
=== [emcee] new code: rerun ===
kill_new_emcee - INFO - Resuming Emcee non-linear search (previous samples found).
is_complete True
likelihood calls this process 4041
total samples 240
max log likelihood -36.1444
=== [emcee] OLD code (2026.10.7.1): interrupted run ===
kill_old_emcee - INFO - Starting new Emcee non-linear search (no previous samples found).
exit=9
=== [emcee] new code: resume the old-era interrupted run ===
kill_old_emcee - INFO - Resuming Emcee non-linear search (previous samples found).
is_complete True
likelihood calls this process 4041
total samples 200
max log likelihood -36.1757

At CRASH_AFTER=4000 the 2026.10.7.1 Dynesty run happened to finish in 3620 calls, so it was repeated with CRASH_AFTER=2500:

=== [dynesty] OLD code (2026.10.7.1): interrupted run (CRASH_AFTER=2500) ===
kill_old2_dynesty - INFO - Starting new Dynesty non-linear search (no previous samples found).
SIMULATED KILL after 2500 likelihood calls
exit=9
=== [dynesty] new code: resume the old-era interrupted run ===
kill_old2_dynesty - INFO - Resuming Dynesty non-linear search (previous samples found).
is_complete True
likelihood calls this process 2210
total samples 970
max log likelihood -36.1438

(c) Zipped restore (unzipped output folder deleted, leaving only <identifier>.zip; rerun with the new code)

=== zipped restore: old_dynesty ===
zip has .completed: True | search_internal entries: []
unzipped folder removed: yes
old_dynesty - INFO - Fit Already Completed: skipping non-linear search.
is_complete True
likelihood calls this process 0
total samples 433
max log likelihood -36.1578
folder restored: yes
=== zipped restore: old_emcee ===
zip has .completed: True | search_internal entries: ['files/search_internal/.start_time', 'files/search_internal/.time', 'files/search_internal/search_internal.hdf']
unzipped folder removed: yes
old_emcee - INFO - Fit Already Completed: skipping non-linear search.
is_complete True
likelihood calls this process 0
total samples 200
max log likelihood -36.1788
folder restored: yes
=== zipped restore: kill_new_dynesty ===
zip has .completed: True | search_internal entries: []
unzipped folder removed: yes
kill_new_dynesty - INFO - Fit Already Completed: skipping non-linear search.
is_complete True
likelihood calls this process 0
total samples 410
max log likelihood -36.1621
folder restored: yes
=== zipped restore: kill_new_emcee ===
zip has .completed: True | search_internal entries: ['files/search_internal/.start_time', 'files/search_internal/.time', 'files/search_internal/search_internal.hdf']
unzipped folder removed: yes
kill_new_emcee - INFO - Fit Already Completed: skipping non-linear search.
is_complete True
likelihood calls this process 0
total samples 240
max log likelihood -36.1444
folder restored: yes

The archive-failure path itself (finding 2) is witnessed on a real on-disk DirectoryPaths by the item-3 test.

Findings that did not reproduce

None: all seven reproduced on 911fe05 (each new test failed on untouched code; item 2's behaviour was also reproduced directly through trace_preflight before the API existed).

Decisions taken (not settled by the plan)

• ExecutionSettings lives in autofit/non_linear/search/preflight.py, exposed as NonLinearSearch.execution_settings(analysis); an eager (compile=False) scalar kind is skipped by the trace, grad and batched kinds are still traced (vmap/grad trace even eagerly).
• ctx.update saves the resume state only (the file a run resumes from), before perform_update, so the checkpoint is never behind the written samples; it does not write the archive mid-run.
• Finalization stays gated on skip_fit_output() as post_fit_output was. Under PYAUTO_SKIP_FIT_OUTPUT=1 NSS now leaves its checkpoint beside .completed (harmless: a completed folder never resumes).
• The _pools(analysis) fix was extended to Dynesty, the initializer and make_sneaky_pool (same latent gap); Zeus's make_pool() is unchanged (no analysis in that signature).
• pytest -n auto does not work on this suite (xdist reports different collections between workers from object-address parametrize ids in test_prior_properties.py, and parametrized fits share output dirs), so both legs ran serially.

Test results

Leg	Command	Result
Full unit suite (jax installed)	python -m pytest test_autofit -q -p no:xdist	3579 passed, 1 skipped, 4 xfailed, 0 failed
unittest-nojax emulation	same, with a sitecustomize on PYTHONPATH whose meta-path finder raises ModuleNotFoundError for jax, jaxlib, jaxnnls, optax, chex, jax_zero_contour, blackjax, nufftax, tensorflow_probability (and importlib.util.find_spec returns None for them, as with the packages uninstalled)	3274 passed, 170 skipped, 2 xfailed, 0 failed

New jax/blackjax tests use pytest.importorskip (or the golden _skip_missing guard).

Sphinx: python -m sphinx -b html docs <out> -w warnings.log -> 30 warnings (baseline docs/sphinx_warning_baseline.txt = 30, not exceeded). python docs/_generate_searches.py --check: up to date.

Closes #1682

🤖 Generated with Claude Code