Skip to content

Commit

Permalink
Scatter Parameters to Workers once in SS
Browse files Browse the repository at this point in the history
 - Instead of scattering the Parameters object prior to each call to the
   dask client, scatter it once and put the resulting future in module
   level scope. Retrieve it using `global` and then pass to the client
   in the inner loop. This significantly reduces the run time of SS.
  • Loading branch information
talumbau committed Mar 16, 2024
1 parent 0f9c5a0 commit 04e3212
Showing 1 changed file with 15 additions and 4 deletions.
19 changes: 15 additions & 4 deletions ogcore/SS.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,13 @@
"""
VERBOSE = True

"""
A global future for the Parameters object for client workers.
This is scattered once and place at module scope, then used
by the client in the inner loop.
"""
scattered_p = None

"""
------------------------------------------------------------------------
Define Functions
Expand Down Expand Up @@ -207,6 +214,8 @@ def inner_loop(outer_loop_vars, p, client):
units
"""
# Retrieve the "scattered" Parameters object.
global scattered_p
# unpack variables to pass to function
bssmat, nssmat, r_p, r, w, p_m, Y, BQ, TR, factor = outer_loop_vars

Expand All @@ -222,10 +231,6 @@ def inner_loop(outer_loop_vars, p, client):
ubi = p.ubi_nom_array[-1, :, :] / factor

lazy_values = []
if client:
scattered_p = client.scatter(p, broadcast=True)
else:
scattered_p = p
for j in range(p.J):
guesses = np.append(bssmat[:, j], nssmat[:, j])
euler_params = (
Expand Down Expand Up @@ -1125,6 +1130,12 @@ def run_SS(p, client=None):
results
"""
global scattered_p
if client:
scattered_p = client.scatter(p, broadcast=True)
else:
scattered_p = p

# Create list of deviation factors for initial guesses of r and TR
dev_factor_list = [
[1.00, 1.0],
Expand Down

0 comments on commit 04e3212

Please sign in to comment.