Distributed tskit

[jeromekelleher]

It turns out that with a few straightforward changes to the stats API (and perhaps other places) we can support not-only thread based parellelism, but full distributed computation using dask. Here's an example based on the divergence matrix computation (which is compute intensive) in #2736:

def divergence_matrix(ts, chunks=16):

    @dask.delayed
    def worker(ts, interval):
        return ts.divergence_matrix(windows=interval)

    work = ts._chunk_sequence_by_tree(chunks)

    ts_delayed = dask.delayed(ts)
    tasks = [worker(ts_delayed, interval) for interval in work]
    results = dask.compute(*tasks)

    return sum(results)


ts = msprime.sim_ancestry(5000, recombination_rate=1, sequence_length=1000, random_seed=2)
print(ts)
D = divergence_matrix(ts, 4)
print(D)

This is pretty cool, but I don't think we should go all in and add Dask as a hard dependency in tskit. It's just too heavy and problematic a dependency, and would be really annoying for things like msprime that would never use it.

We could make a "tskit.distributed" (or something) package which provides wrappers like the one above for key functions that would benefit from this kind of distribution.

Any thoughts?

[benjeffery]

Seeing as tsinfer is heading this way too it makes sense to offer this method of parallelism in tskit.

One option would be to make dask be an optional dependency that is only imported as needed with a friendly error message? This would add a lot of cruft to the TreeSequence class though, so I'm leaning towards a separate module as you suggest.