Analysis of Dask distributed workloads

[talumbau]

To get started, I found that I wasn't able to easily capture the output of the dask distributed job. I found that I needed to use PYTHONUNBUFFERED in order to get output to the console and to tee:

export PYTHONUNBUFFERED=1
python run_og_usa.py 2>&1 | tee out2.log 

While the job is running, it looks like you can get the status by pointing a browser to http://localhost:8787/status. That will be my next place to look.

I have found that when I run the example reform policy starting here the job will not actually run to completion, even after 24+ hours. So the goal will be to get profiling information just from the baseline policy in line 62.

[talumbau]

I confirm that I can see a Dask status page at http://localhost:8787/status when running the example with just the baseline policy:

But the graphs for the actual workers don't change much in the baseline computation. I think I'll need to dig deeper in ogcore.execute to understand more.

I noticed that the calibration computation takes about 10 minutes before getting to the actual policy computation, so I make a pickle file of the initially created Calibration object so I can just skip to the runner section and continue to profile.

[jdebacker]

@talumbau Great- thanks for looking into this. Please let us know if there are specific questions we can help you with.

[talumbau]

Even the baseline policy has been running for about 24 hours. Just looking at top it seems like the jobs don't run at a very high CPU usage:

That's unexpected, at least for me. The first significant part of the calculation appears to be here in SS:

https://github.com/PSLmodels/OG-Core/blob/master/ogcore/SS.py#L229-L259

where some root finding is done. The loop has no interloop dependencies so the idea is just to create a series of lazy computations and give them to the dask client to execute in parallel. Seems reasonable but I'm not sure it's making good progress. Some print statements that come later in that file don't appear in my hung output so I think something is going wrong in this inner_loop function.

[jdebacker]

@talumbau This is not my experience. I've had some hanging like this in the reform policy run, but not baseline. I will try to recreate this today or over the weekend.

Can I ask - what hardware and version of Python are you using?

[talumbau]

Yes please do. It might be that something changed recently in a completely clean environment setup.

The machine is a pretty beefy Linux machine. 256GB of RAM, AMD Ryzen Threadripper 64-Cores

(ogusa-dev) talumbau@talumbau:~/src/OG-Core/ogcore$ python
Python 3.11.7 | packaged by conda-forge | (main, Dec 23 2023, 14:43:09) [GCC 12.3.0] on linux
Type "help", "copyright", "credits" or "license" for more information.

(ogusa-dev) talumbau@talumbau:~/src/OG-Core/ogcore$ head /proc/meminfo 
MemTotal:       263723852 kB
MemFree:         4166336 kB
MemAvailable:   239645460 kB
Buffers:         1308396 kB
Cached:         222421328 kB
SwapCached:         1772 kB
Active:         18479096 kB
Inactive:       221188056 kB
Active(anon):   14797508 kB
Inactive(anon):  1324464 kB

[talumbau]

Strange situation continues.... I have discovered that the computation is not proceeding past the scatter of the Parameters here:

 scattered_p = client.scatter(p, broadcast=True)

So for some reason we never return from the scatter. I wonder if there is some kind of circular reference that is causing Dask to just send Python objects to the workers forever. I took a quick look and it seems like this code hasn't changed for a while, so I don't have a good explanation for how this can be the behavior for me but be working for you.

[jdebacker]

@talumbau Today I:

1. Synced my local OG-USA master branch with the upstream repo
2. Rebuilt the ogusa-dev environment
3. Installed ogusa from source
4. Successfully ran the run_og_usa.py script. Last few lines of output:

Max absolute value resource constraint error: 2.7881147422054164e-07
Checking time path for violations of constraints.
Max Euler error, savings:  2.8039792709932954e-12
Max Euler error labor supply:  1.7783552408445757e-12
Time path iteration complete.
It took 3013.398489713669 seconds to get that part done.
run time =  3013.3985810279846
Percentage changes in aggregates: Year                    Variable  2023  2024  2025  ...  2031  2032  2023-2032    SS
0                    GDP ($Y_t$) -0.08 -0.07 -0.13  ... -0.02  0.01      -0.07 -0.36
1            Consumption ($C_t$) -0.01  0.69 -0.02  ... -0.26 -0.27      -0.08 -0.72
2          Capital Stock ($K_t$) -0.33 -0.35 -0.50  ... -0.40 -0.35      -0.44 -1.49
3                  Labor ($L_t$)  0.07  0.10  0.11  ...  0.21  0.23       0.15  0.34
4     Real interest rate ($r_t$) -3.11 -3.02 -2.83  ... -2.79 -2.82      -2.83 -1.39
5                      Wage rate -0.15 -0.17 -0.23  ... -0.23 -0.22      -0.23 -0.70

[talumbau]

Got it. OK that's a good sanity check. So the question is why would this occur on my machine. The issue I'm seeing sounds exactly like this (somewhat old) Github issue on the dask repo. Will dig a bit more. Btw, do you have these versions of dask and distributed?

Type "help", "copyright", "credits" or "license" for more information.
>>> import dask
>>> dask.__version__
'2024.1.1'
>>> import distributed
>>> distributed.__version__
'2024.1.1'
>>> 

[jdebacker]

Yep, same versions of dask and distributed here:

(base) jason.debacker@JDEBACKER-8 ~ % conda activate ogusa-dev
python
(ogusa-dev) jason.debacker@JDEBACKER-8 ~ % python
iPython 3.11.7 | packaged by conda-forge | (main, Dec 23 2023, 14:38:52) [Clang 16.0.6 ] on darwin
Type "help", "copyright", "credits" or "license" for more information.
>>> import dask
dask.>>> dask.__version__
'2024.1.1'
>>> import distributed
>>> distributed.__version__
'2024.1.1'
>>>

[talumbau]

The only difference I can see between our setups would be the Mac/Linux difference (and I guess that your Python was built with clang and mine with gcc). If that makes a difference in the operation of scatter, that would be bad. I'll try some of the suggested workarounds I'm seeing on StackOverflow. It's encouraging at least that I'm seeing that others have hit this error over the last couple of years, but I don't see anything very recent.

[talumbau]

Continuing to debug. Dask distributed loops forever in the while loop here:

https://github.com/dask/distributed/blob/main/distributed/scheduler.py#L6823

coming from the client.scatter... Trying to figure out why.

[talumbau]

OK. I figured out the basics of the problem but I don't have a deep understanding of dask distributed defaults. The fix for this is to manually set the n_workers of the Dask Client to be the same as num_workers in the example, so just change these lines to:

    num_workers = min(multiprocessing.cpu_count(), 7)
    client = Client(n_workers=num_workers)

This allows the scatter to proceed and I can see that each worker is doing useful work now:

There is apparently some distinction between the num_workers argument that one can give to a dask client for an actual computing task and the number of work processes spun up for a LocalCluster when creating a client. It looks like if the number of workers for a LocalClient is larger than some expected value when using things like scatter, the scatter operation will hang because it's not getting the data to all of the workers. The difference between our two machines is that my machine has 128 cores. I noticed that my default LocalCluster from the Client creation was set to 16. I'm guessing Dask has a default value for the n_workers parameter that is something like min(local cores, 16). I think if the computation is set up to usefully use 7 cores at once, I should put a PR with the small change above. @jdebacker does that sound good?

[talumbau]

Confirming that this code in distributed returns 16 processes when you pass in n=128.

[talumbau]

For the final output of run_ogusa_example.py, I got this:

Max absolute value resource constraint error: 2.878503971640023e-07
Checking time path for violations of constraints.
Max Euler error, savings:  1.3273826482418372e-12 
Max Euler error labor supply:  1.715960706860642e-12
Time path iteration complete.
It took 10943.8991522789 seconds to get that part done.
run time =  10943.904386043549

I noticed that this is longer than yours by about 3x. Is that unexpected?

[jdebacker]

@talumbau, a PR with that small change would be great. Nice work figuring this out!

I noticed that this is longer than yours by about 3x. Is that unexpected?

Yes. But you go many warnings re memory usage throughout, I have seen that happen. But if it ran straight through, I'm surprised at the difference in time given that we both ran the simulation using the max of 7 cores.

[talumbau]

OK, I re-ran on my machine with the above fix and just used the basic cProfile tool:

python -m cProfile -o output.pstat run_og_usa.py

Then I used snakeviz to take a look:

The claim from this profiling is that TPI takes the majority of the time. The section for SS (the long rectangle) is basically dominated by root finding, which makes sense. The section for TPI shows time spent in either the distributed function gather (the majority) or the scatter function (smaller rectangle). I need to break this down more and figure out how to capture profiling info across all of the workers. It does look like there's some inefficiency in how Python data is distributed to the cluster, but I really need to see the per-worker breakdowns to understand what's going on. @jdebacker @rickecon is the breakdown of TPI vs. SS total work fairly accurate as far as you know?

[jdebacker]

@talumbau Yes, I would expect the TPI to take about 20-60x the time of the SS solution.

@rickecon and I are noticing serious memory issues when using distributed since we have increase the size of the parameters object (which is passed, in full, to each worker in the distributed calls).

[talumbau]

Yes, I have a suspicion that there's overhead in serlializing and sending the parameters object in full but I haven't measured it yet. I did find that I could get a reasonable profile of the Dask workers in TPI by using the built-in Dask profiler context manager. That resulted in a nice HTML profile but I can't upload it to this issue because apparently HTML is not supported for uploads.

My goal is just to measure how much time each worker is spending in scipy.opt.root. If the serialize/deserialize of the parameters object is taking too much time, we can fix that first. But once it is mostly opt.root in SS and TPI, I have identified a replacement routine that runs on GPUs or TPUs. So, the goal then would be to just isolate some representative function calls and see if there is a speedup from using similar root finding functions that are optimized for accelerator hardware. Does that sound like a reasonable path?

[jdebacker]

@talumbau Sounds like a great plan. Definitely interested in the overhead with the parameters object.

[talumbau]

I figured out how to use Dask pre-loading to get cProfile data from every worker. See this PR if you want to try this:

#90

However, I accidentally did a test run while still using the built-in Dask profiling capability (so I had double profiling going on).
So, basically I was running TPI with the performance_report(filename='dask-profile.html') context manager and the above profiling.This made the job take way too long, so I'm skeptical of the cProfile results I got. I'm going to re-run with just the profiling in the above PR and nothing else going on in TPI. I have verified that the above dumps out a pstats file for each worker so then we will be able to see time spent in root finding vs. the other stuff.

[talumbau]

Somewhat disappointing results, in that I didn't get a trace that really indicated "time spent on root finding" vs "time spent on communication overhead". The profiling data mostly looks like this:

where the biggest chunk is listed as "epoll". My instinct is that this is because the opt.root is called as a future and this might mess up the profiling (if, for example, the actual function execution is in another thread and not the main thread of the worker process). If that theory is correct, then the actual work is about 88% and the comm overhead is about 12% (that's the overall contribution of the other "chunk" in the profile).

I think I'm going to find it worth my time to continue to tweak the profiling until I'm convinced the above theory is correct. It's typically not a good idea to try performance optimization based on a theory of what is happening. Best to keep plugging away until I know for sure.

[talumbau]

Also, I updated #90 with a second commit that does two things:

• I get rid of the reform policy just to makes things simpler
• I noticed that the Client is built, used, del'd, then re-built after Calibration. So I create the custom LocalCluster and dump the profiling twice (once for each client).

It's surprising that Dask built-in profiling here doesn't provide the basic needed info here. Could be I'm doing something wrong....

[talumbau]

I put some wrappers around the root finding calls to get some data on how much they are called, average time per call, etc. in this PR: https://github.com/PSLmodels/OG-Core/pull/904/files

I dumped that to a file and processed the log with pandas in this gist: https://gist.github.com/talumbau/5ae4f134bbda11ccf231b2b06fa1d63c

That gives this data:

SS totals
pid
1319897    5417.061035
1320389       3.586065
1320391      16.631678
1320396      18.278505
1320401       4.214920
1320405       3.994411
1320409      20.373022
1320412      20.316383
Name: exec_times, dtype: float32
SS mean
pid
1319897    5417.061035
1320389       0.128074
1320391       0.296994
1320396       0.326402
1320401       0.150533
1320405       0.142658
1320409       0.363804
1320412       0.534642
Name: exec_times, dtype: float32
TPI totals
pid
1320389    24934.210938
1320391    27634.402344
1320396     9755.297852
1320401    25821.816406
1320405    47013.824219
1320409     8526.523438
1320412     6819.452637
Name: exec_times, dtype: float32
TPI mean
pid
1320389    2.154888
1320391    1.731479
1320396    0.788689
1320401    1.797676
1320405    2.507003
1320409    0.689346
1320412    0.657360
Name: exec_times, dtype: float32

So one observation is that there is one very long call to opt.root in the SS portion of the computation. Also, I wonder if it might be best to completely serialize the SS portion of the computation because the workers don't spend a big portion of time in that function.

TPI definitely has more time spent in root finding, although it's interesting that a given call is not very time consuming. At the end of this run, the output was:

It took 26263.215485334396 seconds to get that part done.
run time =  26263.21924138069

so I would not say I totally understand how the sum of the time in root finding exceeds the wallclock time of the computation for some of the work processes. Each process is set up to have multiple threads but I don't believe this should help for this kind of computation. A bad situation would be if somehow we are starting multiple calls to opt.root for one process and using something like async to have multiple threads in one Python process await-ing on the call to opt.root. That would not provide a benefit. We would be better off in Python spawning a separate process for each call. Generally, threading can work for Python loads that might be IO bound or network bound, but since root finding is compute bound, in Python we need more processes.

[jdebacker]

We determined that the long call to opt.root in the SS portion is due to the call to SS.SS_fsolve within an opt.root.

The other calls to opt.root are spread across processes via Dask inside the opt.root call that calls SS.SS_fsolve.

[talumbau]

I made some progress here. Since you mentioned that Parameters doesn't change between iterations, I am using client.scatter just once when we start SS and then not scattering in the inner loop. This saves a lot of time for SS. @jdebacker @rickecon can you try running with this PR and see if it improves your run times?

PSLmodels/OG-Core#911

Next step would be to try something similar for TPI.

[talumbau]

I tried a similar trick with TPI (so both using the PR above and a similar change in TPI.py) and the output definitely is cleaner. Very few log messages like this:

distributed.utils_perf - WARNING - full garbage collections took 30% CPU time recently (threshold: 10%)

However, I have noticed that the run time for run_og_usa.py. has increased recently (I'm trying to stay current with main on OGUSA and OGCORE. Does this overall runtime look reasonable to you right now?

Max absolute value resource constraint error: 2.878503971640023e-07
Checking time path for violations of constraints.
Max Euler error, savings:  1.3273826482418372e-12
Max Euler error labor supply:  1.715960706860642e-12
Time path iteration complete.
It took 34532.09340286255 seconds to get that part done.
run time =  34532.09715175629

[jdebacker]

@talumbau This single scatter works really well for me -- solving SS in just 19 seconds!

But the TPI run time of 34.5k seconds does seem long to me.

[talumbau]

ok, let's consider merging the SS change. I'll put up the TPI change as well and then if you're willing to do TPI with that we can do an A/B comparison on TPI on your machine. Still not sure why my TPI runs seems to be really long sometimes.

[jdebacker]

Sounds like a plan. I see some CI tests failing for the SS change, so let me work those out and make a PR to your branch to get the tests passing before we merge.

[talumbau]

Oh yes, I see. I think there are some paths into that inner_loop function that don't first go through a path to set scattered_p from the initial value of None. Good tests! I'll take a look or if you have a fix that would be great too.