Understand Hail GWAS regression implementation

• sgkit cost estimate on UKB data: https://github.com/related-sciences/ukb-gwas-pipeline-nealelab/issues/32
• hail cost estimate on UKB data: https://github.com/Nealelab/UK_Biobank_GWAS/issues/37
• Work to improve our costs: https://github.com/pystatgen/sgkit/issues/390

We’d like to get our costs closer to Hail’s costs. To do so, it would be helpful to understand the Hail implementation and see if there are any ideas in their implementation that we might reuse in ours.

[tomwhite]

I tried with chunks: {variant: 64, sample: -1} and it was a bit faster than map_blocks on this dataset:

• original code without chunking in the samples dimension: 68 s * 128 cores / 67552 variants / phenotype = 0.13 core seconds/variant/phenotype

Here’s the notebook and performance report.

[tomwhite]

I tried another experiment, where I used Dask map_blocks to independently process each block of variants. This is akin to what Hail does (except I’m still doing covariate projection as discussed above). It’s important that the array is not chunked in the samples dimension, which means that the chunk size in the variants dimension has to be quite small. I used 64 to give ~100MB chunks.

On 8x data the processing time on a 16 node cluster was 77s, compared to 110s from the equivalent run in https://github.com/pystatgen/sgkit/issues/390#issuecomment-768332568. This is a 1.4x speedup.

Translating this into normalized numbers (using https://github.com/pystatgen/sgkit/issues/390#issuecomment-768380382):

• Original implementation: 150 s * 960 cores / 141910 variants / phenotype = 1.01 core seconds/variant/phenotype
• Improved chunking: 185 s * 192 cores / 141910 variants / phenotype = 0.25 core seconds/variant/phenotype
• map_blocks: 77 s * 128 cores / 67552 variants / phenotype = 0.15 core seconds/variant/phenotype

This is a ~6x speedup from the original, and if we could use preemptible instances to get a ~5x cost saving, I think that would put us in the same cost ballpark as Hail.

Ideally Dask would do this kind of optimization for us so we didn’t have to resort to map_blocks, but it’s good to know that this is a technique we can fall back to if needed.

Here’s the notebook I used, and the performance report.