Investigate persisting input dataset in cluster memory on GWAS performance

[tomwhite]

In #390 it might help performance to persist the dosage data in cluster memory, and avoid loading it from Zarr for each phenotype. This could be tested using the benchmark using simulated data.

[tomwhite]

I did a quick test for this (notebook), by calling the following before calling the gwas function:

from dask.distributed import wait
XL = XL.persist()
wait(XL)

This has the effect of pinning the biggest matrix in the cluster memory.

The time went from 103s (load from Zarr + gwas) to 66s (gwas only), so since we run a computation for each phenotype on the same input, there are potential savings here of ~30%.

Performance reports:

• pr_20210201_67552_365941_25_dist_1_gs_baseline.html.txt
• pr_20210201_67552_365941_25_dist_1_gs_persist.html.txt

(BTW I noticed that the actual computation times shown in the task stream were ~75s and ~40s respectively, so the saving may actually be more like 50%. I wonder if the missing time goes to serializing the task graph - if so then hopefully this will be improved by some of the work going on in Dask and Dask Distributed to move to HighLevelGraph. There's some discussion of the problem in dask/distributed#3872. We might also look at increasing chunk sizes to alleviate the problem.)