@divyegala

But again, there is no choice for a user of UMAP to select multi-GPU build. We must not do this automatically.

I am working on multi-gpu resource on raft side. Currently the plan is to wrap this in pylibraft and advise users to explicitly pass this as the handle if they want to use multi-gpu. Like this;

from pylibraft.common import DeviceResourcesSNMG

# if want to use multi GPU
multigpu_handle = DeviceResourcesSNMG()
# can also do this indicating which GPUs they want to use
multigpu_handle = DeviceResourcesSNMG([2,4,6,7])

# pass the multigpu_handle as the handle to UMAP
umap_nnd = UMAP(handle = multigpu_handle, 
                 build_algo="nn_descent", 
                 build_kwds={"n_nearest_clusters": 2, 
                            "n_clusters": 8, 
                            "nn_descent": {'graph_degree': 32, 'max_iterations': 20}
                 })

If nothing is passed to the handle, then the default handle will be used, and this results in a single-gpu run.

Do you think there is a better way to let users do a multi-gpu build? Any advice would be greatly appreciated : )

@jinsolp thanks for the explanation, that API looks good to me. Opt-in behavior is perfect.

We need to be very explicit about the fact that this is not multi-GPU UMAP but rather multi-GPU KNN step in UMAP. That's because users of cuml-dask expect that their data can be distributed across GPUs and we would still be able to create a learned ML model, but in this use-case data will need to be on a single-GPU for other parts of UMAP.

I would like to review this PR prior to merge.

[After discussing offline]
For now, decided to stick to current nnd_ prefixed build_kwds without any nested kwds since we just expose nn descent build for UMAP. This PR is no longer a breaking change, and will just have an additional nnd_n_nearest_clusters in build_kwds.

@jcrist Can you approve this PR assuming that your concerns were addressed?

Changed docs (content staying almost same, change in style) in .pyx file because of docs build failure in CI
@csadorf

/merge