Multi-gpu KNN build for UMAP using all-neighbors API

[jinsolp]

Description

Allows multi-gpu knn graph building in UMAP using the all-neighbors API.

PRs that need to be merged before this one

• Wrapper for all-neighbors knn graph building cuvs#785
• Python wrapper for device_resources_snmg raft#2666

Changes in cuML UMAP usage

from pylibraft.common import DeviceResourcesSNMG

# if want to use multi GPU
multigpu_handle = DeviceResourcesSNMG()
umap_nnd = UMAP(handle = multigpu_handle, 
                 build_algo="nn_descent", 
                 build_kwds={"nnd_n_nearest_clusters": 2, 
                               "nnd_n_clusters": 8, 
                               "nnd_graph_degree": 32, 
                               "nnd_max_iterations": 20
                 })

Closes #6729

[divyegala]

The hints in Python docs are awesome. But again, there is no choice for a user of UMAP to select multi-GPU build. We must not do this automatically.

[jinsolp]

@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 : )

[divyegala]

@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.

[divyegala]

Single-GPU C++ and Python LGTM, but I'd like @csadorf to assign a Python reviewer as well. For multi-GPU, I'd like to see tests added perhaps in the dask module to ensure the environment has enough GPUs.

[viclafargue]

Thanks @jinsolp! Just two small comments

[csadorf]

I would like to review this PR prior to merge.

[divyegala]

Cpp approval only. LGTM!

[csadorf]

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

[jinsolp]

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

[divyegala]

This PR is blocked on some changes from rapidsai/cuvs#944

[csadorf]

A combination of nnd_n_clusters>1 and data_on_host=False (the default) will currently break user code, because batching is not supported on device.

Agreed mitigation approach is to auto-set data_on_host=True with a deprecation warning in 25.06 and then require data_on_host=True in combination with nnd_n_clusters>1 as of 25.08.

No longer blocked.

[csadorf]

Just a tiny suggestion for language, otherwise LGTM.

[divyegala]

/merge