Guard the Dask scatter in run_TPI so serial (client=None) runs work

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)

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)