Identify lack of scalability in gwas_linear_regression

It appears that this function does not scale well when run on a cluster.

Notes from my most recent attempt:

The code I ran is here: https://github.com/related-sciences/ukb-gwas-pipeline-nealelab/blob/4f862e31b8093d25fdaa8da7f841b9be8583cda4/scripts/gwas.py#L268
This works on a dataset with fewer variants (chr XY, which has ~8k variants compared to ~141k for chr 21)
The operation utilizes only one or two workers in a cluster of 20 n1-highmem-8 instances:

CPU utilization across worker VMs

Status Page

Drilling in on one of the workers that is running all the tasks, I see that the only not obviously parallelizable task it seems to be running is "solve-triangular":

The job ultimately failed with the error "ValueError: Could not find dependent ('transpose-e1c6cc7244771a105b73686cc88c4e43', 42, 21). Check worker logs".

Several of the workers show log messages like this:

distributed.worker - INFO - Dependent not found: ('rechunk-merge-66cac011d34e1c66cde96678a9e011b5', 0, 21) 0 . Asking scheduler

Perhaps this is what happens when one node unexpectedly becomes unreachable? I'm not sure.

I will run this again on a smaller dataset that didn't fail to get a performance report and task graph screenshot (which doesn't work on this data because the UI won't render so many nodes).

@mrocklin

Notes from a more detailed performance report resulting from running this for a smaller dataset (that succeeds):

(renamed to .txt to avoid github attachment filter)
gwas-height-performance-report.html.txt

This zoomed-out view of the task stream in the report doesn't strike me as very healthy:

Task graph:

@mrocklin (cc: @ravwojdyla) do you have any suggestions on how to identify why the work isn't being distributed well on larger datasets for this workflow?

Thank you for producing the performance report. If you want to publish these in the future then you may also want to look into gist.github.com and https://raw.githack.com/ .

I've only looked very briefly at it, but the thing that stands out the most is the 760s transfer times leading up to sum tasks. At expected bandwidths, these would be 150GB payloads, which I'm assuming is higher than you're expecting. It's also odd in how synchronized these transfers are, they all end within a few seconds of each other.

@quasiben

ValueError: Could not find dependent ('transpose-e1c6cc7244771a105b73686cc88c4e43', 42, 21). Check worker logs

I'm curious, do worker logs report anything strange?

Also I'm curious, which version of distributed are you running? (adding this to the performance report here: dask/distributed#4249)

cc'ing @quasiben , who cares a bit about this space and has a lot of experience tracking down similar performance problems.

Also it looks like some of the workers are under memory pressure. I wonder if maybe they're struggling to accept a large-ish payload.

…

On Tue, Nov 17, 2020 at 11:13 AM Eric Czech ***@***.***> wrote: Notes from a more detailed performance report resulting from running this for a smaller dataset (that succeeds): (renamed to .txt to avoid github attachment filter) gwas-height-performance-report.html.txt <https://github.com/pystatgen/sgkit/files/5555640/gwas-height-performance-report.html.txt> This zoomed-out view of the task stream in the report doesn't strike me as very healthy: [image: Screen Shot 2020-11-17 at 2 07 57 PM] <https://user-images.githubusercontent.com/6130352/99435626-62c72080-28de-11eb-9627-560ead8131e8.png> Task graph: [image: Screen Shot 2020-11-17 at 1 54 52 PM] <https://user-images.githubusercontent.com/6130352/99435949-d406d380-28de-11eb-8097-2d3ba8e9bdaf.png> @mrocklin <https://github.com/mrocklin> (cc: @ravwojdyla <https://github.com/ravwojdyla>) do you have any suggestions on how to identify why the work isn't being distributed well on larger datasets for this workflow? — You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub <https://github.com/pystatgen/sgkit/issues/390#issuecomment-729142413>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AACKZTEPOAG7V7AMMB4M3GTSQLDPHANCNFSM4TY6UJFA> .

When I looked I saw some fairly lengthy disk-read/writes which I assume to be dask spilling. Which might correspond to the workers under memory pressure

Two minutes is a long time for a disk read/write. My initial guess is that there is some other external factor causing contention. That's just a wild guess though. Looking at the worker/thread ratio, I wonder if it would make more sense to have far more workers with fewer threads each. Perhaps try four threads per worker?

…

On Tue, Nov 17, 2020 at 12:39 PM Benjamin Zaitlen ***@***.***> wrote: When I looked I saw some fairly lengthy disk-read/writes which I assume to be dask spilling. Which might correspond to the workers under memory pressure [image: Screen Shot 2020-11-17 at 3 32 36 PM] <https://user-images.githubusercontent.com/1403768/99444378-20a3dc00-28ea-11eb-8890-fe2cf70240fd.png> — You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub <https://github.com/pystatgen/sgkit/issues/390#issuecomment-729193609>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AACKZTGNBX6YOZINCQP2AZDSQLNO3ANCNFSM4TY6UJFA> .

Thank you for producing the performance report. If you want to publish these in the future then you may also want to look into gist.github.com and https://raw.githack.com/ .

👍

I'm curious, do worker logs report anything strange?

I see messages like "Worker stream died during communication" so I'm sure a couple workers in the cluster had died. At the end I was down to 18 instead of 20. Full log from one worker is here jic.

When I looked I saw some fairly lengthy disk-read/writes which I assume to be dask spilling. Which might correspond to the workers under memory pressure

Thanks Ben, what do you make of the "transfer-sub" tasks (the long red bars)? Do you have any intuition for what's happening in those?

Memory pressure does seem to be part of the problem -- I rechunked my input to 1/16th of the original chunk size and the whole job has progressed further. Parallel utilization is still pretty disappointing across the whole cluster:

CPU utilization across workers after rechunking input to 1/16th of original

@quasiben

Looking at the worker/thread ratio, I wonder if it would make more sense to
have far more workers with fewer threads each. Perhaps try four threads
per worker?

Alright, I can try that. Hey @quasiben, how do you set the number of workers per VM in Cloud Provider?

Maybe you can choose smaller VMs?

…

On Tue, Nov 17, 2020, 1:05 PM Eric Czech ***@***.***> wrote: Looking at the worker/thread ratio, I wonder if it would make more sense to have far more workers with fewer threads each. Perhaps try four threads per worker? Alright, I can try that. Hey @quasiben <https://github.com/quasiben>, how do you set the number of workers per VM in Cloud Provider? — You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub <https://github.com/pystatgen/sgkit/issues/390#issuecomment-729213000>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AACKZTDB4HWUPL3D2RSO7C3SQLQQJANCNFSM4TY6UJFA> .

@quasiben

Alright, I can try that. Hey @quasiben, how do you set the number of workers per VM in Cloud Provider?

That's a good question. I don't think this is supported but shouldn't be too hard. You can control the number of threads with worker_options={"nthreads": 2 }. As @mrocklin suggest, this often is tuned with smaller VMs. Can I ask you to file an issue on dask-cloudprovider to support multiple workers per VM ?

Can I ask you to file an issue on dask-cloudprovider to support multiple workers per VM ?

You bet! dask/dask-cloudprovider#173

The larger dataset with smaller chunks did ultimately finish with no errors. Here are a couple readouts:

Performance report (26M): https://drive.google.com/file/d/1feWLKNrjQkslKDIZ7T39fPCNQQDBrQFs/view?usp=sharing

It doesn't seem like any of network, disk, or cpu are even close to being saturated so I assume there is some room for improvement.

As a very rough estimate, this task takes about 5 hrs on a single 64 vCPU VM and 3 hours on a cluster of 20 8 vCPU VMs (160 vCPUs). That would imply ~2 hours with perfect scaling so this should be an approximate ceiling for improvements.

Log: dask_gwas_chr21_log.txt

I will try again with smaller VMs and see if there are any major differences.

@ravwojdyla

FYI @ravwojdyla and I have been talking a bit about some similar observations on a much simpler workflow in https://github.com/related-sciences/data-team/issues/38 (private). One conclusion there was that the individual objects on GCS are so small that API requests aren't efficient yet bigger chunk sizes would start to make the workloads fail. The GCS objects for a chunk are roughly 2MiB on disk but >100MB in memory. Another potential explanation for this behavior is that the GCS objects aren't being loaded asynchronously. We're still investigating both.

@martindurant

I don't know if you're aware of the work done by @martindurant on nicer async IO for remote storage and zarr, but he might be good to talk to.

It looks like you're still primarily blocked by a few oddly long transfers. They're less prominent now, which is good, but probably worth investigating further.

(sorry, wrong thread)

@eric-czech

That's right, I was only increasing or decreasing the number of workers for Dask CP. I never explicitly set the number of threads per worker.

Thanks for the confirmation @eric-czech. https://github.com/pystatgen/sgkit/issues/437#issuecomment-762395287 shows how to adjust that for vanilla Dask distributed Client (which is trivial). In the dask-cloudprovider realm it's a bit more complicated, dask/dask-cloudprovider#173 is about support for multiple workers per VM, but that's actually not the case we need (since most computation is numpy based, it should be fine (and even more optimal) to have a single worker per VM, unless we believe GIL is a problem), what we need is a way to adjust how many worker-threads a single worker has, here's how you can do that for dask-cloudprovider's GCP cluster:

GCPCluster(projectid="foobar",
           n_workers=1,
           # here you could also disable spilling etc if you are in the "tuning mode"
           worker_options={"nthreads": 2, "memory_limit": 1.0},
           machine_type="n1-standard-4")

and now instead of 1 worker with 4 threads (avg 3.75GB per thread), you get 1 worker with avg 7.5GB per thread. "memory_limit": 1.0 is paramount (without that option you would get 2 threads with avg 3.75GB per thread). And thus you get more memory per thread, and potentially can use larger chunks (less tasks overhead) or avoid spilling. Here I would also mentioned that you might want to adjust the numpy threads at this point: dask/dask-cloudprovider#230

Here I would also reiterate that https://github.com/pystatgen/sgkit/issues/437#issuecomment-758849588 has other options to adjust "memory limits": worker resources and more recently layer annotations dask/dask#6701 (but I haven't tried them with GCPCluster).

Some summary and takeaways from the meeting today:

"square" chunks {variant: 5216, sample: 5792} result in spilling and poor performance
"short and fat" chunks {variant: 652, sample: 5792} mitigate the spilling but we still see significant transfer time overhead
in this workflow matmul does contraction in the sample axis
I start to develop this heuristic: in matmul chunking of the contraction axis (sample in our case) will have major impact on the amount of communication required, and in our case variant size/chunking will influence memory. I run a couple of matmul tests to validate that, you can see reports here:
- https://ghcdn.rawgit.org/ravwojdyla/dask/rav/matmul_stats/report_w4_v2E%2B02_s1E%2B03.html
- https://ghcdn.rawgit.org/ravwojdyla/dask/rav/matmul_stats/report_w4_v2E%2B02_s1E%2B05.html
  So above by just adjusting chunking in contraction axis for matmul we cut the walltime from 30s to 8s (and cumullative transfer time from 30s to 1.5s) So it's not surprising that we see difference (slowdowns) between chromosomes (for example XY has barely 8K variants, while chr 21 ~141K and chr 11 ~530K). And there is space for improvement in the communication time.

A hypothesis I would explore: tune the size of chunks in a way that reduces the transfer needed in the "contraction"/reduction axis.

To be more concrete, we saw benefit of going from 5216 -> 652 in the variant axis, specifically no spilling (less memory overhead), and increasing the size of samples should reduce the communication/transfer time. So overall to be more precise, I believe we should try to:

increase the chunk size in the sample axis as much as practical
and if we start spilling reduce the chunk size in the variant axis
if we need to increase WT/Memory ratio previous comments describe that
so a "degenerated" case would be chunking: {variant: 1, sample: MAX}.

Looking at the https://github.com/pystatgen/sgkit/issues/390#issuecomment-748205731, I annotate the moment we switch from chr. 21 to chr. 11 and observe slow down:

notice that it leads to spilling
chr. 11 has has 530k variants compared to 140k for chr21
again size/chunking in the variant axis had influence on the memory required, and led to spilling and slow downs
so for a "constant" cluster, maybe there exists a chunk size that is optimal for all chromosomes, but we might end up with chunking that depends on the chromosome, fortunately we have all the information to make that decision and it would be just a matter of coming up with a function that goes from chromosome size to chunking

@ravwojdyla

Thanks @ravwojdyla - that's very interesting and useful.

I have now reproduced the slowdown using @eric-czech's simulated data from #438. Running on a cluster of 8 n1-standard-8 workers I get the following execution times for multiples of the XY dataset:

Multiple	Time (s)	Perf report
1	30.91	pr_8444_365941_25_dist_1_gs.html.txt
2	40.35	pr_16888_365941_25_dist_1_gs.html.txt
4	80.84	pr_33776_365941_25_dist_1_gs.html.txt
8	696.48	pr_67552_365941_25_dist_1_gs.html.txt

Notice that going from 4x to 8x produces a disproportionate increase in time taken. And this correlates with spilling, as shown on the GCP monitoring snapshot (spilling starts at 12:32). It also matches the whitespace in the taskstream.

This is using @ravwojdyla's new matmul implementation. The code I used to run the benchmark is here: https://github.com/tomwhite/gwas-benchmark

I think the next thing to try is different chunk sizes as suggested in Rafal's last comment.

I had a look at the shapes and chunk sizes of the intermediate variables in the gwas function. To do this I broke out each computation into a separate cell in this "chunk report" notebook, so you can get an idea of what each operation is doing.

One thing that stuck out was the simple fact that the total memory needed for the computation (for 8x data) exceeds the cluster memory (see the bottom of the notebook) - but it doesn't for 4x data. So this explains why disk spilling is inevitable, and why we see a big slowdown from 4x to 8x. When I doubled the cluster size to 16 nodes, there was barely any spilling, and the 8x computation took ~200s as opposed to ~700s (on the 8 node cluster). A good improvement, but still lots of whitespace in the task stream. (See the performance report)

Next, I tried matching the output of the first, "outer" matmul (XC @ LS) to have same chunk sizes as the XL array, which it is later combined with. If XC is not chunked in the first dimension, then XC @ LS will not be chunked in the first dimension, and when XL - XC @ LS is computed there is a lot of data transfer since XL has roughly square chunks, while XC @ LS has tall skinny chunks. (It's easier to see this graphically in the chunk report notebook.) I think this is why there are a lot of "transfer-sum" bars in the task stream that has lots of whitespace.

This helped a lot - the time is now down to 110s, so only 2.75x slower than the 4x data case from before. The task graph has very little whitespace now, which is a significant improvement. (See the performance report)

(Side note: I get a Dask warning saying 'PerformanceWarning: Increasing number of chunks by factor of 64' which actually is a good thing in this case!)

The notebook I used for the computation is here: https://nbviewer.jupyter.org/github/tomwhite/gwas-benchmark/blob/1a73ff865d100da7ea51c73456aa1a8526ff29b7/gwas_simulation.ipynb

Summary:

Use a cluster large enough to ensure all of the intermediate arrays can fit in cluster memory to avoid (very expensive) disk spilling.
Be careful about "outer" matmul operations since the output chunking probably isn't what you expect.

I ran the benchmark on 16x data (135,104 variants) using the XC rechunking change from the previous comment. It took 185s on a 24 node n1-standard-8 cluster.

This data is comparable in size to chr21 (141,910 variants), and according to https://github.com/related-sciences/ukb-gwas-pipeline-nealelab/issues/32, it took 150s on a 60 node n1-highmem-16 cluster.

150 s * 960 cores / phenotype = 144,000 core seconds/phenotype
185 s * 192 cores / phenotype = 35,520 core seconds/phenotype

So this is a ~4x speedup.

What is the cost improvement?

What would it cost if we used preemptible instances? I know Eric had problems with our workloads running reliably on preemptible instances, so I'm curious if that's even an option for us.

@ravwojdyla

Cost improvement is probably a bit better than 4x since I wasn't using highmem instances (more like 5x maybe) - but these are all estimates, so there's likely to be some variation. (Also we don't know if the scaling is linear all the way up to chr1 size data.)

I'm not sure about how well preemptible instances work with Dask in general, or for our workload.

There is still potentially a lot of scope for reducing transfer times (e.g. for the second perf report in https://github.com/pystatgen/sgkit/issues/390#issuecomment-768332568, the compute time is ~5000s and the transfer time is ~3000s), by improving the chunking further. My change looks like it fixed one of the more egregious cases, but there are likely other things we can do along the lines that @ravwojdyla has suggested. This would involve going through the chunk report in more detail.

Adjust chunksize in gwas_linear_regression to reduce data transfer be…

To summarize the state of this issue, I think the main ways to improve performance have been identified (and fixed in a couple of cases):

Improve scalability of Dask matmul (fixed in Rewrite matmul as blockwise without concatenated contraction dask/dask#7000)
Rechunking the covariates array (fixed in Adjust chunksize in gwas_linear_regression to reduce data transfer between workers #454), ~4x speedup
Persisting the input dataset in cluster memory (Investigate persisting input dataset in cluster memory on GWAS performance #449), ~1.5x speedup
Using preemptible instances (Investigate use of preemptible GCP instances for GWAS #453), up to ~5x cost saving
~~Using standard linear regression, not a mixed model~~ Using map_blocks to do independent linear regressions (Understand Hail GWAS regression implementation #448), ~1.4x speedup
Further rechunking improvements (Investigate further chunking improvements for better GWAS performance #461), unknown speedup

Thanks for the update Tom. This has been a very effective case study to motivate general performance enhancements. Should you ever find some free time (ha!) I encourage you to write things up . I think that there are some folks out there that would find it interesting.

…

On Tue, Feb 9, 2021 at 4:10 AM Tom White ***@***.***> wrote: To summarize the state of this issue, I think the main ways to improve performance have been identified (and fixed in a couple of cases): 1. Improve scalability of Dask matmul (fixed in dask/dask#7000 <dask/dask#7000>) 2. Rechunking the covariates array (fixed in #454 <https://github.com/pystatgen/sgkit/pull/454>), ~4x speedup 3. Persisting the input dataset in cluster memory (#449 <https://github.com/pystatgen/sgkit/issues/449>), ~1.5x speedup 4. Using preemptible instances (#453 <https://github.com/pystatgen/sgkit/issues/453>), up to ~5x cost saving 5. Using standard linear regression, not a mixed model (#448 <https://github.com/pystatgen/sgkit/issues/448>), ~2x speedup 6. Further rechunking improvements (#461 <https://github.com/pystatgen/sgkit/issues/461>), unknown speedup — You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub <https://github.com/pystatgen/sgkit/issues/390#issuecomment-775822748>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AACKZTCU5DAM3HRIJLQIW73S6EC2BANCNFSM4TY6UJFA> .

Another summary of the latest state. I think these three changes should solve the problem for the time being:

Improve scalability of Dask matmul (fixed in dask/dask#7000)
Rechunking so there is no chunking in the samples dimension (https://github.com/pystatgen/sgkit/issues/448#issuecomment-780655217), ~7x speedup
Using preemptible instances (Investigate use of preemptible GCP instances for GWAS #453), up to ~5x cost saving

The next step would be to run on a suitable subset of the original UKB data to see if they work well enough in practice. (Note that we also have map_blocks as a fallback if 2 above doesn't scale well enough for our purposes.)

@tomwhite

@tomwhite thanks for the summary and hard work! If I may, I would also like to document some potential experiments to further improve the performance (when we get back to this issue):

the https://github.com/pystatgen/sgkit/issues/448#issuecomment-780655217 uses chunking {variants: 16, samples: -1}. -1/MAX reduces the transfer time, that said cumulative transfer time is still roughly 25% of the computation time. As mentioned in https://github.com/pystatgen/sgkit/issues/390#issuecomment-764950336, at this time we could explore maximising variant size (without triggering spilling!). In a sense if we freeze some parameters like cluster type/size etc, we could frame this problem as an optimisation of a function of variant and sample chunk sizes with some standard techniques and an educated starting point, but at this point it might be enough to just continue manual optimisation (amusingly, it also reminded me of Herodotou, Herodotos, et al. "Starfish: A Self-tuning System for Big Data Analytics." Cidr. Vol. 11. No. 2011. 2011). So this is one area to explore.
zarr/IO still takes significant amount of time, which is another area to experiment with, see knowledge dump in Genetics data IO performance stats/doc #437. First step here should be to capture more detailed metrics (including profiles of this step, for example I would be curious to see how much time we spend decompressing data) and making an educated decision. This could be as simple as experimenting with the native zarr chunk size, changing compression parameters, all the way to exploring an alternative layouts or formats that could better serve the specific properties of genetic data or computation (like high compressibility, limited data types and values etc), or using cache/hot storage layer (which btw could also be useful with preemptive VM if used in a creative way).

Navigation Menu

Repository navigation

Identify lack of scalability in gwas_linear_regression #390

34 remaining items

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Identify lack of scalability in gwas_linear_regression #390

Description

Activity

eric-czech commented on Nov 17, 2020

mrocklin commented on Nov 17, 2020

mrocklin commented on Nov 17, 2020

mrocklin commented on Nov 17, 2020

quasiben commented on Nov 17, 2020

mrocklin commented on Nov 17, 2020

eric-czech commented on Nov 17, 2020

eric-czech commented on Nov 17, 2020

mrocklin commented on Nov 17, 2020

quasiben commented on Nov 17, 2020

eric-czech commented on Nov 17, 2020

eric-czech commented on Nov 18, 2020

eric-czech commented on Nov 18, 2020

mrocklin commented on Nov 18, 2020

martindurant commented on Nov 18, 2020

34 remaining items

ravwojdyla commented on Jan 20, 2021

ravwojdyla commented on Jan 21, 2021

tomwhite commented on Jan 25, 2021

tomwhite commented on Jan 27, 2021

tomwhite commented on Jan 27, 2021

hammer commented on Jan 27, 2021

tomwhite commented on Jan 27, 2021

tomwhite commented on Feb 9, 2021

mrocklin commented on Feb 10, 2021

tomwhite commented on Feb 18, 2021

ravwojdyla commented on Feb 19, 2021

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions