Dealing with Dask memory issues in the run script, #2

[jdebacker]

Following up on PR #85, I noticed that despite this changes I was again seeing memory warnings:

2024-02-08 09:49:00,004 - distributed.utils_perf - WARNING - full garbage collections took 14% CPU time recently (threshold: 10%)

And no, or slow, progress in when running run_og_usa.py when running with OG-Core 0.11.0 (which increases the size of several objects in the Specifications class object that contains the model parameters.

After some profiling, I noticed that the "bytes stored" in the dask processes became quite large with the tax function estimation, but then remained quite high after that task completed:

This PR helps deal with that issue by closing and deleting the client after tax function estimation and then instantiating a new client before the model solution algorithm is called.

This has seemed to avoid these issues.

[rickecon]

@jdebacker. This looks great. Thanks for doing this. Will merge as soon as tests pass.

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (7a6757b) 70.82% compared to head (b726392) 70.82%.

Additional details and impacted files

@@           Coverage Diff           @@
##           master      #87   +/-   ##
=======================================
  Coverage   70.82%   70.82%           
=======================================
  Files          22       22           
  Lines        1340     1340           
=======================================
  Hits          949      949           
  Misses        391      391           

Flag	Coverage Δ
unittests	70.82% <ø> (ø)

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

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[rickecon]

Thanks @jdebacker. Merging.