From 9a2f823f4b05bffede48b6a43335785fc2ebec17 Mon Sep 17 00:00:00 2001 From: vahid-ahmadi Date: Fri, 28 Aug 2026 18:12:13 +0100 Subject: [PATCH] Guard the Dask scatter in run_TPI so serial runs work run_TPI called client.scatter(p, broadcast=True) unconditionally, so run_TPI(p, client=None) raised AttributeError before the TPI loop was entered -- making the serial fallback inside the loop unreachable in exactly the case it exists for. SS.inner_loop already guards the equivalent block with `if client:`; this matches that. Adds a fast regression test that seeds baseline SS results from the cached test_io_data pickles and monkeypatches TPI.inner_loop to raise a sentinel, asserting only that execution reaches the first serial household solve. No SS or TPI solve is performed, so it runs in ~1s and needs no Dask cluster. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01XCKMb1aicxYaeUC1us2nvF --- ogcore/TPI.py | 33 +++++++++++++++++---------------- tests/test_TPI.py | 47 +++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 64 insertions(+), 16 deletions(-) diff --git a/ogcore/TPI.py b/ogcore/TPI.py index 2032d42c3..45b401d7f 100644 --- a/ogcore/TPI.py +++ b/ogcore/TPI.py @@ -955,26 +955,27 @@ def run_TPI(p, client=None): trust_radius_max = getattr(p, "TPI_trust_radius_max", 10.0) prev_accel_dist = np.inf - # Before scattering, temporarily remove unpicklable schema objects - schema_backup = {} - for attr in ["_defaults_schema", "_validator_schema", "sel"]: - if hasattr(p, attr): - schema_backup[attr] = getattr(p, attr) + if client: + # Before scattering, temporarily remove unpicklable schema objects + schema_backup = {} + for attr in ["_defaults_schema", "_validator_schema", "sel"]: + if hasattr(p, attr): + schema_backup[attr] = getattr(p, attr) + try: + delattr(p, attr) + except Exception: + pass + + # Scatter the parameters + scattered_p_future = client.scatter(p, broadcast=True) + + # Restore the schema objects (they're not needed by workers anyway) + for attr, value in schema_backup.items(): try: - delattr(p, attr) + setattr(p, attr, value) except Exception: pass - # Scatter the parameters - scattered_p_future = client.scatter(p, broadcast=True) - - # Restore the schema objects (they're not needed by workers anyway) - for attr, value in schema_backup.items(): - try: - setattr(p, attr, value) - except Exception: - pass - # TPI loop while (TPIiter < p.maxiter) and (TPIdist >= p.mindist_TPI): outer_loop_vars = (r_p, r, w, p_m, BQ, RM, TR, theta) diff --git a/tests/test_TPI.py b/tests/test_TPI.py index 2a2b680d1..85384b724 100644 --- a/tests/test_TPI.py +++ b/tests/test_TPI.py @@ -9,6 +9,7 @@ - test_run_TPI_full_run(), 11 parameterizations, local only - test_run_TPI(), 2 parameterizations, local only - test_run_TPI_extra(), 8 parameterizations, local only + - test_run_TPI_serial_no_client(), 1 parameterization """ import multiprocessing @@ -1206,3 +1207,49 @@ def test_run_TPI_extra(baseline, param_updates, filename, tmpdir, dask_client): rtol=1e-04, atol=1e-04, ) + + +class _ReachedTPILoop(Exception): + """Sentinel raised in place of the household inner loop.""" + + +def test_run_TPI_serial_no_client(tmpdir, monkeypatch): + """ + Regression test: TPI.run_TPI(p, client=None) must reach the TPI loop. + + run_TPI used to call ``client.scatter(p, broadcast=True)`` + unconditionally, so passing ``client=None`` raised an + ``AttributeError`` before the TPI loop was ever entered, making the + serial fallback inside the loop unreachable. This test does not + solve a transition path: it seeds the baseline SS results from the + cached pickles in ``test_io_data`` and monkeypatches + ``TPI.inner_loop`` to raise a sentinel, so it asserts only that + execution gets as far as the first serial household solve. + """ + # Seed cached baseline SS results so no SS solve is needed + old_baseline_dir = os.path.join(CUR_PATH, "test_io_data", "OUTPUT2") + ss_vars = utils.safe_read_pickle( + os.path.join(old_baseline_dir, "SS", "SS_vars.pkl") + ) + ss_vars_new = {SS_VAR_NAME_MAPPING[k]: v for k, v in ss_vars.items()} + baseline_dir = os.path.join(tmpdir, "baseline") + utils.mkdirs(os.path.join(baseline_dir, "SS")) + with open(os.path.join(baseline_dir, "SS", "SS_vars.pkl"), "wb") as f: + pickle.dump(ss_vars_new, f) + + p = Specifications( + baseline=True, + baseline_dir=baseline_dir, + output_base=baseline_dir, + num_workers=1, + ) + p.update_specifications(TEST_PARAM_DICT.copy()) + p.maxiter = 1 + + def mock_inner_loop(*args, **kwargs): + raise _ReachedTPILoop() + + monkeypatch.setattr(TPI, "inner_loop", mock_inner_loop) + + with pytest.raises(_ReachedTPILoop): + TPI.run_TPI(p, client=None)