Early scalability demonstrations

[alimanfoo]

It would be very useful to have some simple demonstrations that using sgkit and the underlying stack we can run some computations over large input datasets, showing we are making real steps towards the goal of interactive analysis of biobank-scale datasets, given a sufficiently powerful HPC/cloud/GPU cluster.

The main purpose of this is to be able to give some highlights in a short report on what has been achieved from the sgkit development work to date.

By "demonstration" I mean just a simple benchmark result.

E.g., being able to say that we can run a PCA on UK biobank-sized data using a N node cluster in Google Cloud and compute the result in M minutes.

E.g., being able to say we can run a pairwise distance computation on the MalariaGEN Ag1000G phase 2 data and compute the result in M minutes on a N node CPU cluster in Google Cloud. Even better would be being able to say we can also run the same computation in M' minutes on a N' node GPU cluster.

I.e., I just need a couple of data points here showing proof of concept, I don't need any detailed performance or scalability benchmarking.

Raising this issue to share suggestions for computations/datasets to use, and post back any results.

[alimanfoo]

Following up the suggestion to run a pairwise distance computation on the Ag1000G phase 2 data, with some more details on the input data.

The data are available in scikit-allel-style zarrs in GCS. The full set of SNP calls and genotypes is at this path: gs://ag1000g-release/phase2.AR1/variation/main/zarr/all/ag1000g.phase2.ar1. It can be opened e.g.:

import fsspec
import zarr

store = fsspec.get_mapper('gs://ag1000g-release/phase2.AR1/variation/main/zarr/all/ag1000g.phase2.ar1')
callset_snps = zarr.open_consolidated(store=store)

These data are grouped by chromosome arm. A demonstration for just a single chromosome arm would be sufficient. 2R is the largest, and so has the most SNPs. Alternatively whole genome is also interesting if equally convenient to compute. In case it helps, here are the genotype calls for 2R:

gt = callset_snps['2R/calldata/GT']

One route to bring these data into sgkit is via the open_vcfzarr function. Another route may be to use the data that @tomwhite has previously converted to the sgkit zarr representation. Either would be fine.

@aktech I'll follow up by email to get you access to our jupyterhub service on Google Cloud so you can try this on a dask kubernetes cluster if useful.

cc @jeromekelleher

[eric-czech]

E.g., being able to say that we can run a PCA on UK biobank-sized data using a N node cluster in Google Cloud and compute the result in M minutes.

FYI @alimanfoo this notebook runs PCA on the UKB PLINK data using a 50 node 8 vCPU cluster (n1-standard-8). It takes ~30 seconds to run on a 108,621 variant x 141,070 sample (~60GB) alt allele count matrix. I chose that size based on what UKB did for their own PCA calculations. For comparison, running an axis-wise reduce function takes about 3 seconds (e.g. get the call rate for all variants).

I also tried this on a 20 node cluster and saw the same PCA step take 50 seconds, so the operations don't scale perfectly -- the 50 node cluster should have only taken 20 seconds otherwise.

I originally tried this with tsqr and non-random PCA but that kept blowing past memory limits, even on "highmem" instances. I suspect something in tsqr produces n_samples x n_samples results for chunks that include all samples but relatively tiny numbers of variants. I think it's an awkward fit for arrays this shape regardless so I plan on getting comfortable with using randomized PCA (as in the notebook).

[aktech]

One route to bring these data into sgkit is via the open_vcfzarr function. Another route may be to use the data that @tomwhite has previously converted to the sgkit zarr representation. Either would be fine.

@tomwhite Can you share, where can I get this? or would you suggest I use open_vcfzarr function?

[tomwhite]

@aktech I think the simplest thing for the moment is to use read_vcfzarr on a single contig (e.g. 2R) - that way you'll avoid all the rechunking issues that we're still trying to get resolved. (Rechunking is only needed when we are combining contigs - for a single contig no rechunking is necessary.)

Note that read_vcfzarr doesn't support cloud stores yet, so you'll need to copy the contents of gs://ag1000g-release/phase2.AR1/variation/main/zarr/all/ag1000g.phase2.ar1/2R to the local filesystem (or change it to read using fsspec).

[aktech]

Oh, alright. Thanks @tomwhite I'll give it a go.

[tomwhite]

BTW the sgkit notebooks I'm working on are here: https://github.com/tomwhite/shiny-train/tree/sgkit/notebooks/gwss. (Very much WIP.)

[aktech]

Note that read_vcfzarr doesn't support cloud stores yet, so you'll need to copy the contents of gs://ag1000g-release/phase2.AR1/variation/main/zarr/all/ag1000g.phase2.ar1/2R to the local filesystem (or change it to read using fsspec).

The contents of gs://ag1000g-release/phase2.AR1/variation/main/zarr/all/ag1000g.phase2.ar1/2R is 57.58 GiB, which I cannot copy in the datalab as it only has less than 10 GiB of storage.

or change it to read using fsspec

By this did you mean modifying the function read_vcfzarr to read using fsspec?

[tomwhite]

it only has less than 10 GiB of storage

Is it possible to increase that?

By this did you mean modifying the function read_vcfzarr to read using fsspec?

Yes, since the input is zarr you should be able to call zarr. open_group with a MutableMapping obtained from fsspec. We already do something similar here, for example: https://github.com/pystatgen/sgkit/blob/master/sgkit/io/vcf/vcf_reader.py#L264

[aktech]

Is it possible to increase that?

I am not sure, its malariagen's datalab.

I am now doing this on my cluster, where I have enough disk space. I have pulled the contents of gs://ag1000g-release/phase2.AR1/variation/main/zarr/all/ag1000g.phase2.ar1/2R on the server now. I am getting the following error, while reading the data through read_vcfzarr

read_vcfzarr(path)
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-23-b318af759c7a> in <module>
----> 1 read_vcfzarr(path)

~/sgkit/sgkit/io/vcfzarr_reader.py in read_vcfzarr(path)
     44 
     45     # Index the contig names
---> 46     variants_chrom = da.from_zarr(vcfzarr["variants/CHROM"]).astype(str)
     47     variant_contig, variant_contig_names = encode_array(variants_chrom.compute())
     48     variant_contig = variant_contig.astype("int16")

~/.local/lib/python3.7/site-packages/zarr/hierarchy.py in __getitem__(self, item)
    347                          synchronizer=self._synchronizer)
    348         else:
--> 349             raise KeyError(item)
    350 
    351     def __setitem__(self, item, value):

KeyError: 'variants/CHROM'

Also, after I get through the error above, I am not sure which part of the xarray dataset I need to use for pairwise distance calculation, as in the pairwise distance needs a 2D matrix, is there a documentation somewhere which I can refer to for the same?

[jeromekelleher]

Also, after I get through the error above, I am not sure which part of the xarray dataset I need to use for pairwise distance calculation, as in the pairwise distance needs a 2D matrix, is there a documentation somewhere which I can refer to for the same?

Was there an example in one of @alimanfoo's notebooks for this? Maybe in the pairwise distance blogpost?

[aktech]

Was there an example in one of @alimanfoo's notebooks for this? Maybe in the pairwise distance blogpost?

I didn't find any notebooks neither in blog post nor in the prototype.

[jeromekelleher]

OK. Well, we'll be comparing the genotypes anyway - I'm not sure exactly how Alistair uses pairwise distance between genotypes, but I don't think it'll matter too much for a simple benchmark.

So, basically we're getting the distance between pairs of samples by comparing their genotype arrays.

[tomwhite]

@aktech try using #324 to read the data.

[aktech]

So, basically we're getting the distance between pairs of samples by comparing their genotype arrays.

Alright, let me check that.

@aktech try using #324 to read the data.

Thanks, let me have a look.

[aktech]

@tomwhite I am getting the same error, while trying to read the data: read_vcfzarr("shared/sgkit_data/2R"), am I missing something here?

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-5-7e129b306df9> in <module>
----> 1 read_vcfzarr("shared/sgkit_data/2R")

~/sgkit/sgkit/io/vcfzarr_reader.py in read_vcfzarr(path)
     42 
     43     # don't fix strings since it requires a pass over the whole dataset
---> 44     return _vcfzarr_to_dataset(vcfzarr, fix_strings=False)
     45 
     46 

~/sgkit/sgkit/io/vcfzarr_reader.py in _vcfzarr_to_dataset(vcfzarr, contig, variant_contig_names, fix_strings)
    168     if contig is None:
    169         # Get the contigs from variants/CHROM
--> 170         variants_chrom = da.from_zarr(vcfzarr["variants/CHROM"]).astype(str)
    171         variant_contig, variant_contig_names = encode_array(variants_chrom.compute())
    172         variant_contig = variant_contig.astype("int16")

~/.local/lib/python3.7/site-packages/zarr/hierarchy.py in __getitem__(self, item)
    347                          synchronizer=self._synchronizer)
    348         else:
--> 349             raise KeyError(item)
    350 
    351     def __setitem__(self, item, value):

KeyError: 'variants/CHROM'