Run on Dask

[tomwhite]

For very large datasets it might be beneficial to run on Dask (although there is a lot of mileage in machines with lots of cores and RAM).

Joblib supports Dask as a backend, however the parallel implementation of pynndescent keeps all heap updates in memory, which may be limiting. A more distributed approach is to use Dask bags to shuffle updates across machines - this is slower but should support larger datasets. I made an initial (incomplete) attempt at this here: https://github.com/tomwhite/pynndescent/tree/dask

[lmcinnes]

I'll try to take a look when I get some time. I have also had a chat with Matt Rocklin about a distributed pynndescent on dask, and he had some good ideas for minimizing shuffles (at the cost of more iterations). Depending on performance costs and how distributed the arrays are that may also be worth pursuing.

[jamestwebber]

Hello again both of you! 👋

Is this still an issue? I've recently been implementing some workflows in Dasks and distributed kNN would be really great because it's one of the slower steps.

[lmcinnes]

The parallelism has since been reworked and runs via numba now. With that said a dask distributed version of pynndescent would be an awesome thing to have -- allowing users to potentially work with datasets larger than memory. I would be happy to try and work with someone with dask expertise who wanted to implement such a thing.

[jamestwebber]

I started checking out @tomwhite's branch but I see that it's pretty outdated at this point–it uses some functions from the pynndescent.threaded module which is no longer around due to the numba refactor that you mention. I see there's a threaded_rp_trees module with a note about dask in the comment.

I'm far from a dask expert but I have been talking to some of them recently, I could try to get a PR going with some help from @lmcinnes if you can provide guidance?

[lmcinnes]

Hi @jamestwebber, and thanks for being willing to look into this. The threaded rp-trees remain around because it was a template of a plan for how to handle rp-trees if the input data was (potentially distributed) dask array. I think the basic idea there is amenable to a dask style approach. The rest would need some discussion to work out the best way to handle distributed data so that we can try to retain some locality in computation and not just have giant data shuffles going on the whole time. I would happy to try and discuss some ideas for that in more detail soon.

[jamestwebber]

Sounds great. If I could figure out how to contact you, I would try to set up a time (unless you prefer to discuss via Github). But you're very elusive. 😂

[lmcinnes]

you can reach me at leland DOT mcinnes AT gmail DOT com

[murilocn]

Hi. I'm really interested in collaborating on a Dask version of PyNNDescent. Did you develop the initial project?

[jamestwebber]

I wish! Unfortunately I've done basically nothing on this project because of too many other priorities.

But it's still a bottleneck for my workflow so I'd like to solve it. Maybe having a collaborator would provide the spark! I'll send you an email. 🙂