NNDescent Dask performance depends on import order
@csadorf is already working on this.
Since Nov 17, 2025.
- Dominant language
- Cuda
- Stars
- 854
- Forks
- 236
- Avg merge
- 3d 3h
- Merged PRs (30d)
- 62
Description
Summary
The cuvs.neighbors.nn_descent module exhibits significant performance degradation when imported after a Dask cluster is created, compared to importing before cluster creation.
We should consider programmatically mitigating this issue, issuing warnings when the import order may cause problems, and/or clarifying the correct usage in the documentation.
Originally reported in https://github.com/rapidsai/cuvs/issues/1384 .
Observed Behavior
Testing with N=100,000 samples, K=100 features:
Pre-import (cuVS imported before Dask cluster):
- Run 1: 0.58s
- Run 2: 0.55s
- Run 3: 0.56s
- Average: ~0.56s
Post-import (cuVS imported after Dask cluster):
- Run 1: 1.49s
- Run 2: 1.62s
- Run 3: 1.53s
- Average: ~1.55s
Performance impact: ~2.8x slowdown when importing after cluster creation.
Key Code Pattern
The issue occurs when cuVS is imported after Dask cluster creation:
from dask_cuda import LocalCUDACluster
from dask.distributed import Client
# Creating cluster FIRST
cluster = LocalCUDACluster()
client = Client(cluster)
# Then importing cuVS - causes slowdown
from cuvs.neighbors import nn_descent
import cupy as cp
Test Environment
- RAPIDS version: 25.12
- Python: 3.13
- CUDA_VISIBLE_DEVICES: 0,1
- LocalCUDACluster with 2 GPUs
Workaround
Import cuVS and CuPy before creating the Dask cluster:
from cuvs.neighbors import nn_descent
import cupy as cp
from dask_cuda import LocalCUDACluster
from dask.distributed import Client
cluster = LocalCUDACluster()
client = Client(cluster)
Reproduction Scripts
test_nn_descent_import_order.py
import argparse
from dask_cuda import LocalCUDACluster
from dask.distributed import Client
parser = argparse.ArgumentParser()
parser.add_argument("--import-order", choices=["pre", "post"], default="post",
help="Import cuvs before (pre) or after (post) creating Dask cluster")
args = parser.parse_args()
if args.import_order == "pre":
from sklearn.datasets import make_blobs
from cuvs.neighbors import nn_descent
import cupy as cp
if __name__ == "__main__":
cluster = LocalCUDACluster()
client = Client(cluster)
if args.import_order == "post":
from sklearn.datasets import make_blobs
from cuvs.neighbors import nn_descent
import cupy as cp
N = 100000
K = 100
X, y = make_blobs(n_samples=N, n_features=K, random_state=12)
X = cp.asarray(X, dtype="float32")
idxparams = nn_descent.IndexParams(
graph_degree=15,
intermediate_graph_degree=None,
metric="sqeuclidean",
)
idx = nn_descent.build(idxparams, dataset=X)
client.close()
cluster.close()
test_import_order.sh
#!/bin/bash
echo "Testing cuVS nn_descent import order performance issue"
echo "======================================================="
echo
echo "Running with PRE-import (cuVS imported before Dask cluster):"
for i in {1..3}; do
echo -n " Run $i: "
python test_nn_descent_import_order.py --import-order pre 2>/dev/null | grep "Time:"
done
echo
echo "Running with POST-import (cuVS imported after Dask cluster):"
for i in {1..3}; do
echo -n " Run $i: "
python test_nn_descent_import_order.py --import-order post 2>/dev/null | grep "Time:"
done
echo
echo "Expected: POST should be significantly slower than PRE"
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.