Pairwise distance scalability

[alimanfoo]

Raising this issue to revisit the scalability of our pairwise distance calculation and whether it's worth returning to a map-reduce style implementation that would allow chunking along both dimensions.

In the work that @aktech is doing on early scalability demonstrations (#345) there are some memory usage difficulties emerging. @aktech is I believe trying to run a pairwise distance computation over data from Ag1000G phase 2, using all SNPs and samples from a single chromosome arm. This is of the order ~20 million variants and ~1000 samples. With the current implementation, it is hard to get this to run on systems with average memory/CPU ratios, below 12G/CPU. My understanding is that, essentially, this is because the pairwise distance implementation currently does not support chunking along the variants dimension, and so to reduce memory footprint you need short chunks along the samples dimension. Depending on how the input data have been chunked natively, this may be suboptimal, i.e., you may need to run the computation with sample chunks that are (much) shorter than the native storage.

If this is correct, then it raises two questions for discussion.

First, should we revisit the map-reduce implementation of pairwise distance? This would allow chunking on both samples and variants dimensions, and so could naturally make use of whatever the underlying chunking of the data in storage, without large memory footprint.

Secondly, do we ever really need to run pairwise distance on arrays that are large in the variants dimension? I.e., do we care about scaling this up to large numbers of variants? xref https://github.com/pystatgen/sgkit/pull/306#issuecomment-714654217

[jeromekelleher]

Secondly, do we ever really need to run pairwise distance on arrays that are large in the variants dimension? I.e., do we care about scaling this up to large numbers of variants? xref #306 (comment)

I think you mean in the samples dimension? (If so, I agree - ~10K is as many samples as we could hope to support for O(n^2) like this)

[alimanfoo]

Secondly, do we ever really need to run pairwise distance on arrays that are large in the variants dimension? I.e., do we care about scaling this up to large numbers of variants? xref #306 (comment)

On this, it is true that for many analyses it's entirely reasonable to work around memory limitations by downsampling data along the variants dimensions. E.g., often when doing a PCA we randomly downsample variants, and this gives a very reasonable approximation. I.e., typically information about population structure is very clear from just a fraction of the variants.

In other words, the scalability test that @aktech is trying, running pairwise distance over all 20 million SNPs from a chromosome arm, is perhaps somewhat artificial, as you could reasonably downsample to 100,000 SNPs and get the information you need.

On the other side, we do have use cases where we want to compute pairwise distance using all SNPs. E.g., during QC, we check for samples that are technically identical, and check these match our expectations for replicate samples. To do this we compute pairwise distance at all SNPs. (To avoid unnecessary computation, we do an initial pass with a downsampled set of SNPs to exclude pairs of individuals that are obviously not identical, then rerun with all SNPs on remaining samples. But that's a detail.)

[alimanfoo]

Secondly, do we ever really need to run pairwise distance on arrays that are large in the variants dimension? I.e., do we care about scaling this up to large numbers of variants? xref #306 (comment)

I think you mean in the samples dimension? (If so, I agree - ~10K is as many samples as we could hope to support for O(n^2) like this)

No, I did mean the variants dimension.

Re the samples dimension, yes agreed, probably ~10k samples is the most we could hope to support due to computational complexity (although if you have access to GPUs then perhaps that allows you to reach a little further).

My main point was about the variants dimension. Computation scales linearly on the variants dimension, there's no computational reason to expect any limit there. We currently have a practical limit due to memory requirements, which could be removed if we supported chunking along the variants dimension via the original map/reduce implementation. Should we revisit that?

[jeromekelleher]

I see, thanks. Well, from my perspective I'd find it surprising as a user to be limited on the number of variants we could process by memory restrictions - since it's a sample x sample comparison.

How have you used this in the past? Do you subset down to a representative set of sites, or just throw compute at it and let it run over the whole lot?

[alimanfoo]

I see, thanks. Well, from my perspective I'd find it surprising as a user to be limited on the number of variants we could process by memory restrictions - since it's a sample x sample comparison.

Yes, exactly.

How have you used this in the past? Do you subset down to a representative set of sites, or just throw compute at it and let it run over the whole lot?

We've done both, more details above.

[alimanfoo]

Just to add, memory restrictions will be even tighter when moving this to GPU.

[ravwojdyla]

Re the samples dimension, yes agreed, probably ~10k samples is the most we could hope to support due to computational complexity (although if you have access to GPUs then perhaps that allows you to reach a little further).

@alimanfoo I'm curious, I understand that the O(n^2) complexity is vicious, but I'm curious how exactly have you arrived at ~10k (including potential horizontal scaling via dask)?

Also a quick point regarding our current implementation, our current blockwise:

    x_distance = da.blockwise(
        # Lambda wraps reshape for broadcast
        lambda _x, _y: metric_ufunc(_x[:, None, :], _y),
        "jk",
        x,
        "ji",
        x,
        "ki",
        dtype="float64",
        concatenate=True,
    )

we specify concatenate=True, from dask:

Any index, like i missing from the output index is interpreted as a contraction (note that this differs from Einstein convention; repeated indices do not imply contraction.) In the case of a contraction the passed function should expect an iterable of blocks on any array that holds that index. To receive arrays concatenated along contracted dimensions instead pass concatenate=True.

So afaiu in #345 we are concatenating on the variants blocks, and thus potentially the memory issue(?). @aktech what's the reasoning behind concatenate=True? I might have missed it, but have we evaluated an implementation that takes concatenate=None (default)?

[jeromekelleher]

@alimanfoo I'm curious, I understand that the O(n^2) complexity is vicious, but I'm curious how exactly have you arrived at ~10k (including potential horizontal scaling via dask)?

I pulled this one out of thin air first. It was just an arbitrary number from experience - in practise, somewhere between 1K and 100K.

[aktech]

So afaiu in #345 we are concatenating on the variants blocks, and thus potentially the memory issue(?). @aktech what's the reasoning behind concatenate=True? I might have missed it, but have we evaluated an implementation that takes concatenate=None (default)?

@ravwojdyla The concatenate=True as you noted is concatenation along variant blocks is made to pass full vectors to the distance metrics functions for distance calculation. If it is None the x and y will be passed to that lamba function as a list of chunks and you cannot calculate distance metric with non-full vectors i.e. chunks, unless you do it via Map-Reduce approach, which we did evaluated, here is the implementation in this commit. This was slower than the current implementation, but perhaps seems like more scalable.

To summarise what I think of the current state is it's worth evaluating the map reduce implementation for scalability.

[ravwojdyla]

Thanks @aktech, I was curious about the use of blockwise with concatenate=False, and specifically the memory footprint. The commit you reference doesn't use blockwise but instead works directly on blocks. I understand that both might look similar. I was also curious if the iterable of blocks on any array that holds that index was lazy or not(?) I can't find a documentation on this, and when I try it on a dummy data, I actually get a np.ndarray (which would suggest it isn't lazy), in which case the memory consumption is likely the same for concatenate equal True/False, which would be rather disappointing (it would still be interesting to double check that tho, plus whether it duplicates the memory required for chunks), I've also noticed @eric-czech comment:

I'm also thinking that there is a way to express the loop we have now that references .blocks as a blockwise call like this, though I haven't been able to get it to work yet.

which might suggest we have not evaluated it.

+1 to trying the "map-reduce" implementation on the #345. Going forward iff the performance is important right now, we could also have both implementations that switch based on the shape of the data (and/or parameter), unfortunately that would have a real maintenance cost (so that's something we need to be cautious about).

[aktech]

Thanks @aktech, I was curious about the use of blockwise with concatenate=False, and specifically the memory footprint. The commit you reference doesn't use blockwise but instead works directly on blocks. I understand that both might look similar. I was also curious if the iterable of blocks on any array that holds that index was lazy or not(?) I can't find a documentation on this, and when I try it on a dummy data, I actually get a np.ndarray (which would suggest it isn't lazy), in which case the memory consumption is likely the same for concatenate equal True/False, which would be rather disappointing (it would still be interesting to double check that tho, plus whether it duplicates the memory required for chunks), I've also noticed @eric-czech comment:

@ravwojdyla Yes, that's very true, I was also expecting the iterable of blocks to be lazy, which would make the memory consumption very low compared to the current implementation.

I only realised this after I wrote a map reduce implementation with blockwise yesterday. Here is the map reduce implementation with blockwise: https://github.com/aktech/sgkit/blob/pairwise-mr/sgkit/distance/api.py#L53. It might not be the most efficient though (could probably be improved), but gives a rough idea of how it might look like. I tried it yesterday and the memory consumption was not any better (if not worse). So far the the previous map reduce implementation I shared in the commit in my previous comment was better than these two.

Also, one of the reasons that the iterable of blocks are not lazy is probably due to the fact, that the func (blockwise's first argument) is not called unless you call .compute on it. You would rather expect it would call it before you do .compute on it and make an efficient graph. I don't know if that's a feature or a bug, but in most of dask's internal codebase blockwise is called with concatenate=True, which gives me the impression that they are not scalable in both dimenstions.

[ravwojdyla]

@aktech so in general it "seems" like for blockwise iff you have a contraction for specific index, all blocks along that index must fit into memory regardless of the concatenate option (which might actually double the memory required for those blocks, first to keep the chunks, then to build up the input for the blockwise function (?), tho this part likely depends on dask). I wish this was in the documentation, plus the iterable of blocks is a bit misleading. I wonder if we should open an issue to add that information, wdyt?

[aktech]

@aktech so in general it "seems" like for blockwise iff you have a contraction for specific index, all block along that index must fit into memory regardless of the concatenate option

Yes, true.

(which might actually double the memory required for those blocks, first to keep the chunks, then to build up the input for the blockwise function (?), tho this part likely depends on dask).

What do you mean by "keep the chunks", those chunks are computed once, from the original array to make the input for blockwise's func I believe. The memory consumption would be twice atleast for sure from the chunks from _x and _y, which are passed to the func.

I wish this was in the documentation, plus the iterable of blocks is a bit misleading. I wonder if we should open an issue to add that information, wdyt?

Yes, I think so. That would be worth mentioning.

[ravwojdyla]

What do you mean by "keep the chunks", those chunks are computed once, from the original array to make the input for blockwise's func I believe. The memory consumption would be twice atleast for sure from the chunks from _x and _y, which are passed to the func.

@aktech yea, per process (threads within process will reuse memory if needed), I just wonder at which stage exactly will Dask free up the memory required for chunks on each process and how exactly will it build up the input for blockwise. On a dummy data it looks like the chunks are coming in a ndarray (not py list). Also if the cluster has multiple processes and they require the same chunks they will likely fetch that from a process that computed them. So essentially it might be that we would need memory for chunks, _and_ memory for concatenated array or an "array of chunks" (depending on the concatenate parameter). The documentation says that blockwise for concatenate=None will pass list or iterable (depending where you look) of blocks, but locally I'm actually seeing a numpy array of blocks being pass into the function given to blockwise, which afaiu might requires more memory to build up the array than to pass a list/iterable of materialized/numpy chunks (depending on how soon can dask free up memory of the chunk). I wonder if this is some kind of local optimization or a bug. (clarification https://github.com/pystatgen/sgkit/issues/375#issuecomment-724204043)

Edit: but just to be clear, I agree that the blockwise implementations do not look at all promising for any computation that has a large number of elements in the contraction index.

[aktech]

On a dummy data it looks like the chunks are coming in a ndarray (not py list).

Ah that's interesting, for me that is only in the first iteration on the func you'll see an empty numpy array, if you check on the next iteration it will be a Python list. That execution of the function happens even before you call .compute, with empty arrays. Do you see arrays in all iterations?