diff --git a/ogcore/SS.py b/ogcore/SS.py index 12a050b27..bbed8aea7 100644 --- a/ogcore/SS.py +++ b/ogcore/SS.py @@ -220,7 +220,51 @@ def solve_for_j( ) -def inner_loop(outer_loop_vars, p, client): +def scatter_params(p, client): + """ + Scatter the model parameters object to the Dask workers once. + + The Specifications object does not change over the course of a + steady-state solve, so it only needs to be serialized and + broadcast to the workers a single time. The resulting Future can + then be reused across all calls to `inner_loop`. + + Args: + p (OG-Core Specifications object): model parameters + client (Dask client object): client + + Returns: + scattered_p (Dask Future or None): future pointing to `p` on + the workers, or None if there is no client + + """ + if not client: + return None + + # 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 = 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 + + return scattered_p + + +def inner_loop(outer_loop_vars, p, client, scattered_p=None): """ This function solves for the inner loop of the SS. That is, given the guesses of the outer loop variables (r, w, TR, factor) this @@ -242,6 +286,9 @@ def inner_loop(outer_loop_vars, p, client): factor (scalar): scaling factor converting model units to dollars p (OG-Core Specifications object): model parameters client (Dask client object): client + scattered_p (Dask Future or None): future pointing to the + model parameters already scattered to the workers. If + None and a client is provided, `p` is scattered here. Returns: (tuple): results from household solution: @@ -290,25 +337,12 @@ def inner_loop(outer_loop_vars, p, client): # from dask.base import dask_sizeof 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: - setattr(p, attr, value) - except Exception: - pass + # Scatter the parameters only if they have not already been + # scattered by the caller (run_SS scatters once per solve). + if scattered_p is None: + scattered_p_future = scatter_params(p, client) + else: + scattered_p_future = scattered_p # Launch in parallel with submit (or map) futures = [] @@ -684,6 +718,7 @@ def SS_solver( p, client, fsolve_flag=False, + scattered_p=None, ): """ Solves for the steady state distribution of capital, labor, as well @@ -702,6 +737,9 @@ def SS_solver( factor (scalar): scaling factor converting model units to dollars p (OG-Core Specifications object): model parameters client (Dask client object): client + fsolve_flag (bool): flag for whether solution came from fsolve + scattered_p (Dask Future or None): future pointing to the model + parameters already scattered to the Dask workers Returns: output (dictionary): dictionary with steady state solution @@ -759,7 +797,7 @@ def SS_solver( new_factor, new_BQ, average_income_model, - ) = inner_loop(outer_loop_vars, p, client) + ) = inner_loop(outer_loop_vars, p, client, scattered_p) # update guesses for next iteration bmat = utils.convex_combo(new_bmat, bmat, nu_ss) @@ -1265,13 +1303,29 @@ def SS_fsolve(guesses, *args): factor_ss (scalar): scaling factor converting model units to dollars p (OG-Core Specifications object): model parameters client (Dask client object): client + scattered_p (Dask Future or None): optional eighth element of + args, a future pointing to the model parameters already + scattered to the Dask workers Returns: errors (list): errors from differences between guessed and implied outer loop variables """ - bssmat, nssmat, TR_ss, Ig_baseline, factor_ss, p, client = args + if len(args) == 8: + ( + bssmat, + nssmat, + TR_ss, + Ig_baseline, + factor_ss, + p, + client, + scattered_p, + ) = args + else: + bssmat, nssmat, TR_ss, Ig_baseline, factor_ss, p, client = args + scattered_p = None # Rename the inputs r_p = guesses[0] @@ -1326,7 +1380,7 @@ def SS_fsolve(guesses, *args): new_factor, new_BQ, average_income_model, - ) = inner_loop(outer_loop_vars, p, client) + ) = inner_loop(outer_loop_vars, p, client, scattered_p) # Create list of errors in general equilibrium variables error_r_p = float(new_r_p - r_p) @@ -1438,6 +1492,7 @@ def run_SS(p, client=None): # Use the baseline solution to get starting values for the reform use_new_guesses = False # initialize this flag, switches to true # if baseline solution not work for reform + scattered_p = None # future for parameters scattered to Dask workers if p.baseline is False: baseline_ss_path = os.path.join(p.baseline_dir, "SS", "SS_vars.pkl") ss_solutions = utils.safe_read_pickle(baseline_ss_path) @@ -1489,6 +1544,10 @@ def run_SS(p, client=None): + BQ_items + [TRguess] ) + # Scatter the parameters to the workers once, rather + # than once per residual evaluation + scattered_p = scatter_params(p, client) + # Now solve for the steady state of the reform ss_params = ( b_guess, @@ -1498,6 +1557,7 @@ def run_SS(p, client=None): factor_ss, p, client, + scattered_p, ) # Solve for steady state using root finder @@ -1554,6 +1614,9 @@ def run_SS(p, client=None): else: TR_baseline = None Ig_baseline = None + # Scatter the parameters to the workers once, rather + # than once per residual evaluation + scattered_p = scatter_params(p, client) ss_params = ( b_guess, n_guess, @@ -1562,6 +1625,7 @@ def run_SS(p, client=None): factor_ss, p, client, + scattered_p, ) # Solve for steady state using root finder sol = opt.root( @@ -1616,6 +1680,7 @@ def run_SS(p, client=None): p, client, fsolve_flag, + scattered_p, ) if output["G"] < 0.0: warnings.warn( diff --git a/tests/test_SS.py b/tests/test_SS.py index f664426be..c2cb4969b 100644 --- a/tests/test_SS.py +++ b/tests/test_SS.py @@ -1410,3 +1410,136 @@ def test_initial_guesses(tmpdir, use_zeta): assert len(guesses) == 7 + p.J assert n_guess.shape == (p.S, p.J) assert b_guess.shape == (p.S, p.J) + + +class ScatterCountingClient: + """ + Minimal fake Dask client that counts calls to `scatter` and runs + submitted tasks synchronously. Used to check that model parameters + are scattered once per SS solve rather than once per residual + evaluation (see OG-Core issue #1211). + """ + + def __init__(self): + self.scatter_calls = 0 + + def scatter(self, obj, broadcast=False): + self.scatter_calls += 1 + return obj + + def submit(self, func, *args, **kwargs): + return func(*args) + + def gather(self, futures): + return list(futures) + + def __bool__(self): + return True + + +@pytest.fixture(scope="module") +def scatter_test_params(): + """Small (J=1) parameters object plus inner loop inputs.""" + _p = Specifications() + j1_updates = { + "J": 1, + "lambdas": np.array([1.0]), + "e": np.ones((80, 1)), + "beta_annual": [0.96], + "chi_b": [80], + "labor_income_tax_noncompliance_rate": [[0.0]], + "capital_income_tax_noncompliance_rate": [[0.0]], + "income_tax_filer": [[1]], + "wealth_tax_filer": [[1]], + "eta": np.ones((80, 1)) * (1 / 80), + "eta_RM": np.ones((80, 1)) * (1 / 80), + "replacement_rate_adjust": [[1.0]], + "omega": _p.omega.sum(axis=2, keepdims=True), + "omega_SS": _p.omega_SS.sum(axis=1, keepdims=True), + "omega_S_preTP": _p.omega_S_preTP.sum(axis=1, keepdims=True), + "rho": _p.rho[:, :, :1], + "rho_preTP": _p.rho_preTP[:, :1], + "imm_rates": _p.imm_rates[:, :, :1], + "imm_rates_preTP": _p.imm_rates_preTP[:, :1], + } + p = Specifications(baseline=True, num_workers=1) + p.update_specifications(j1_updates) + bssmat = np.ones((p.S, p.J)) * 0.07 + nssmat = np.ones((p.S, p.J)) * 0.4 * p.ltilde + r = 0.04 + w = firm.get_w_from_r(r, p, "SS") + outer_loop_vars = ( + bssmat, + nssmat, + r, + r, + w, + np.ones(p.M), + 1.3, + np.ones(p.J) * 0.00019646295986015257, + 0.12, + None, + 100000, + ) + return p, outer_loop_vars + + +def test_inner_loop_reuses_scattered_p(scatter_test_params): + """ + inner_loop should not re-scatter the parameters when the caller has + already scattered them, but must still scatter when called with the + legacy (outer_loop_vars, p, client) signature. + """ + p, outer_loop_vars = scatter_test_params + + # Legacy signature: scatters once per call + client = ScatterCountingClient() + SS.inner_loop(outer_loop_vars, p, client) + SS.inner_loop(outer_loop_vars, p, client) + assert client.scatter_calls == 2 + + # New signature: parameters scattered by the caller, never again + client = ScatterCountingClient() + scattered_p = SS.scatter_params(p, client) + assert client.scatter_calls == 1 + for _ in range(3): + SS.inner_loop(outer_loop_vars, p, client, scattered_p) + assert client.scatter_calls == 1 + + # No client => no scatter and no future + assert SS.scatter_params(p, None) is None + + +def test_SS_fsolve_scatters_once(scatter_test_params): + """ + Repeated residual evaluations (as done by the root finder) must not + re-scatter the parameters when a scattered future is passed through. + """ + p, outer_loop_vars = scatter_test_params + bssmat, nssmat = outer_loop_vars[0], outer_loop_vars[1] + r = outer_loop_vars[3] + w = outer_loop_vars[4] + guesses = ( + [r, r, w] + + list(np.ones(p.M)) + + [1.3] + + list(np.ones(p.J) * 0.00019646295986015257) + + [0.12, 100000] + ) + + # With a scattered future threaded through: one scatter total, + # regardless of the number of residual evaluations + client = ScatterCountingClient() + scattered_p = SS.scatter_params(p, client) + for _ in range(3): + SS.SS_fsolve( + guesses, bssmat, nssmat, None, None, None, p, client, scattered_p + ) + assert client.scatter_calls == 1 + + # Backwards compatible 7-element args tuple still works (and scatters + # once per evaluation, the old behavior) + client = ScatterCountingClient() + for _ in range(2): + SS.SS_fsolve(guesses, bssmat, nssmat, None, None, None, p, client) + assert client.scatter_calls == 2