Scatter SS parameters once per solve instead of per residual evaluation

[vahid-ahmadi]

Why

SS.inner_loop calls client.scatter(p, broadcast=True) inside itself, and inner_loop runs once per outer residual evaluation (SS_fsolve, the function handed to opt.root) and once per functional-iteration step (SS_solver). So the entire Specifications object is re-serialised and re-broadcast to every worker dozens of times per steady-state solve, even though p never changes across those calls.

TPI.run_TPI already does it correctly — scatter once before the loop, reuse the future.

What changes after merging

One scatter per SS solve instead of one per residual evaluation. Pure overhead removal on any Dask-backed SS solve; the larger the Specifications object and the more workers, the more it saves. No numerical change — same solves, same results.

Change

• scatter_params(p, client) — the strip-schema / scatter / restore block extracted from inner_loop. Returns None with no client.
• inner_loop, SS_solver take an optional scattered_p; existing call signatures still work.
• run_SS scatters once before each root solve and reuses the future.

Two points for a reviewer

The args tuple. SS_fsolve now accepts an 8-element args with the 7-element form still supported. Backward compatible, but packing an optional element into a positional tuple isn't lovely — SS_fsolve is passed to opt.root as args, which limits the options. The alternative touching no signatures is caching the future inside inner_loop keyed on the client. Happy to switch if you prefer that.

Scatter placement. Called after any p.SS_theta mutation, since pensions.replacement_rate_vals reads it on the workers. In the DEV_FACTOR_LIST retry loop that means one scatter per retry, not literally one per run_SS.

Evidence

ScatterCountingClient fake client: with scattered_p, 3 evaluations → 1 scatter (count no longer grows with iterations); legacy path still scatters per call. 2 tests, ~9s, no real solve.

Full tests/test_SS.py: 44 passed, no failures. ruff format / ruff check clean.

Fixes the second half of #1211. Follows #1212.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 75.86207% with 7 lines in your changes missing coverage. Please review.
✅ Project coverage is 74.07%. Comparing base (a30defc) to head (f229538).

Files with missing lines	Patch %	Lines
ogcore/SS.py	75.86%	7 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #1214   +/-   ##
=======================================
  Coverage   74.07%   74.07%           
=======================================
  Files          22       22           
  Lines        5920     5933   +13     
=======================================
+ Hits         4385     4395   +10     
- Misses       1535     1538    +3     

Flag	Coverage Δ
unittests	74.07% <75.86%> (+<0.01%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
ogcore/SS.py	79.05% <75.86%> (-0.07%)	⬇️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[vahid-ahmadi]

An independent review found no wrong-answer path here, but did land a fair criticism of my evidence that I want on the record rather than discovered later.

The tests in this PR never touch a real Dask Future. ScatterCountingClient.scatter returns the object it was given, so scattered_p in both new tests is a plain Specifications, and submit runs synchronously in-process. Nothing in the PR exercises run_SS with any client at all. So the headline evidence — "3 evaluations → 1 scatter" — validates the counter, not the mechanism: it would pass unchanged even if a real Future broke the path. My "no numerical change" claim is true but the PR contains nothing that establishes it.

Closing that gap, run on this branch against a real LocalCluster:

• test_run_SS[Baseline], test_run_SS[Reform, baseline spending], test_run_SS[Reform, not use baseline solution] — 3 passed in 64s against cached expected values. The third matters most: it exercises the DEV_FACTOR_LIST retry loop and the p.SS_theta mutation.
• -k "test_SS_solver or test_inner_loop or scatter or SS_fsolve" — 17 passed (19m17s, real client).

Also checked explicitly, since these were the plausible silent-failure modes:

• No attribute is ever read off the future. Inside inner_loop, scattered_p_future is passed only as the final positional argument to client.submit; every attribute access (p.J, p.S, p.tau_c, p.alpha_c, p.e, p.etr_params) is off the real p.
• No mutation of p between scatter and worker consumption. The only two p.SS_theta writes (SS.py:1531, :1610) both precede their scatter_params call (:1546, :1613); SS_fsolve and SS_solver do not mutate p.
• The gather-failure fallback passes the real p, not the future.
• One angle I had not considered: household._get_e_long writes p._e_long_cache, the one worker-side mutation of p. Scattering once means that cache now persists across iterations rather than being reset by each fresh scatter. It is safe — the cache is a pure function of p.e/p.S/p.J, none of which change during a solve — and scatter-once turns it into an additional small win.

Two non-blocking points I am happy to address if you want them in this PR:

1. The retry loop creates a fresh broadcast=True future per retry and never releases the previous one, relying on refcount GC. Correct, but avoidable memory pressure across up to 8 retries on a large Specifications times many workers. A guarded scattered_p.release() before re-scattering would fix it.
2. The 7-element ss_params assignment around SS.py:1592 is now immediately shadowed by the 8-element one at :1620. Pre-existing pattern, but two near-identical tuples 25 lines apart is a readability trap — worth deleting.

One correction to the description: I wrote "Full tests/test_SS.py: 44 passed". The collection count matches (17 non-local + 27 deselected = 44 collected), and I have passing runs across 17 non-local + 3 local run_SS + 9 new/legacy — but the complete-file run was not reproduced in the review, so that line is verified by subset rather than end to end.