[BUG] KNeighborsClassifier predict crashes for medium/large datasets

[Banus]

Describe the bug When using 20k or more samples to fit a KNeighborsClassifier, training goes on smoothly but the method predict crashes with the following message:

  File "cuml/neighbors/kneighbors_classifier.pyx", line 228, in cuml.neighbors.kneighbors_classifier.KNeighborsClassifier.predict
RuntimeError: Exception occured! file=/conda/conda-bld/libcuml_1591208841859/work/cpp/src_prims/selection/knn.cuh line=519: FAIL: call='cudaPeekAtLastError()'. Reason:invalid argument

Steps/Code to reproduce bug

from cuml.neighbors import KNeighborsClassifier
from sklearn.datasets import make_blobs
from sklearn.model_selection import train_test_split

X, y = make_blobs(n_samples=int(2e4), centers=5, n_features=248)
knn = KNeighborsClassifier(n_neighbors=3)

X_train, X_test, y_train, y_test =\
  train_test_split(X, y, train_size=0.80)
knn.fit(X_train, y_train)
knn.predict(X_test[:1, :])

Expected behavior predict should return the predictions for the given samples.

Environment details (please complete the following information):

• Environment location: [Bare-metal]
• Linux Distro/Architecture: [Ubuntu 18.04.4 amd64]
• GPU Model/Driver: [V100 and driver 440.64.00]
• CUDA: [10.2]
• Method of cuDF & cuML install: [conda]

  _libgcc_mutex             0.1                 conda_forge    conda-forge
  _openmp_mutex             4.5                      1_llvm    conda-forge
  aiohttp                   3.6.2            py37h516909a_0    conda-forge
  appdirs                   1.4.3                      py_1    conda-forge
  arrow-cpp                 0.15.0           py37h090bef1_2    conda-forge
  async-timeout             3.0.1                   py_1000    conda-forge
  attrs                     19.3.0                     py_0    conda-forge
  backcall                  0.2.0              pyh9f0ad1d_0    conda-forge
  bleach                    3.1.5              pyh9f0ad1d_0    conda-forge
  bokeh                     1.4.0            py37hc8dfbb8_1    conda-forge
  boost                     1.70.0           py37h9de70de_1    conda-forge
  boost-cpp                 1.70.0               h8e57a91_2    conda-forge
  brotli                    1.0.7             he1b5a44_1002    conda-forge
  brotlipy                  0.7.0           py37h8f50634_1000    conda-forge
  bzip2                     1.0.8                h516909a_2    conda-forge
  c-ares                    1.15.0            h516909a_1001    conda-forge
  ca-certificates           2020.6.20            hecda079_0    conda-forge
  cairo                     1.16.0            hcf35c78_1003    conda-forge
  certifi                   2020.6.20        py37hc8dfbb8_0    conda-forge
  cffi                      1.14.0           py37he30daa8_1  
  cfitsio                   3.470                hb60a0a2_2    conda-forge
  chardet                   3.0.4           py37hc8dfbb8_1006    conda-forge
  click                     7.1.2              pyh9f0ad1d_0    conda-forge
  click-plugins             1.1.1                      py_0    conda-forge
  cligj                     0.5.0                      py_0    conda-forge
  cloudpickle               1.4.1                      py_0    conda-forge
  colorcet                  2.0.1                      py_0    conda-forge
  cryptography              2.9.2            py37hb09aad4_0    conda-forge
  cudatoolkit               10.2.89              h6bb024c_0    nvidia
  cudf                      0.14.0                   py37_0    rapidsai
  cudnn                     7.6.5                cuda10.2_0  
  cugraph                   0.14.0                   py37_0    rapidsai
  cuml                      0.14.0          cuda10.2_py37_0    rapidsai
  cupy                      7.6.0            py37h940342b_0    conda-forge
  curl                      7.68.0               hf8cf82a_0    conda-forge
  cusignal                  0.14.0                   py37_0    rapidsai
  cuspatial                 0.14.0                   py37_0    rapidsai
  cuxfilter                 0.14.0                   py37_0    rapidsai
  cycler                    0.10.0                     py_2    conda-forge
  cytoolz                   0.10.1           py37h516909a_0    conda-forge
  dask                      2.19.0                     py_0    conda-forge
  dask-core                 2.19.0                     py_0    conda-forge
  dask-cuda                 0.14.1                   py37_0    rapidsai
  dask-cudf                 0.14.0                   py37_0    rapidsai
  dask-xgboost              0.2.0.dev28      cuda10.2py37_0    rapidsai
  datashader                0.10.0                     py_0    conda-forge
  datashape                 0.5.4                      py_1    conda-forge
  decorator                 4.4.2                      py_0    conda-forge
  defusedxml                0.6.0                      py_0    conda-forge
  distributed               2.19.0           py37hc8dfbb8_0    conda-forge
  dlpack                    0.2                  he1b5a44_1    conda-forge
  double-conversion         3.1.5                he1b5a44_2    conda-forge
  entrypoints               0.3             py37hc8dfbb8_1001    conda-forge
  expat                     2.2.9                he1b5a44_2    conda-forge
  fastavro                  0.23.5           py37h8f50634_0    conda-forge
  fastrlock                 0.5              py37h3340039_0    conda-forge
  fiona                     1.8.13           py37h900e953_0    conda-forge
  fontconfig                2.13.1            h86ecdb6_1001    conda-forge
  freetype                  2.10.2               he06d7ca_0    conda-forge
  freexl                    1.0.5             h14c3975_1002    conda-forge
  fsspec                    0.7.4                      py_0    conda-forge
  gdal                      3.0.2            py37hbb6b9fb_2    conda-forge
  geopandas                 0.8.0                      py_1    conda-forge
  geos                      3.7.2                he1b5a44_2    conda-forge
  geotiff                   1.5.1                hbd99317_7    conda-forge
  gflags                    2.2.2             he1b5a44_1002    conda-forge
  giflib                    5.1.7                h516909a_1    conda-forge
  glib                      2.65.0               h3eb4bd4_0  
  glog                      0.4.0                h49b9bf7_3    conda-forge
  grpc-cpp                  1.23.0               h18db393_0    conda-forge
  hdf4                      4.2.13            hf30be14_1003    conda-forge
  hdf5                      1.10.5          nompi_h3c11f04_1104    conda-forge
  heapdict                  1.0.1                      py_0    conda-forge
  icu                       64.2                 he1b5a44_1    conda-forge
  idna                      2.9                        py_1    conda-forge
  imagecodecs-lite          2019.12.3        py37h03ebfcd_1    conda-forge
  imageio                   2.8.0                      py_0    conda-forge
  importlib-metadata        1.6.1            py37hc8dfbb8_0    conda-forge
  importlib_metadata        1.6.1                         0    conda-forge
  ipykernel                 5.3.0            py37h43977f1_0    conda-forge
  ipython                   7.15.0           py37hc8dfbb8_0    conda-forge
  ipython_genutils          0.2.0                      py_1    conda-forge
  jedi                      0.17.1           py37hc8dfbb8_0    conda-forge
  jinja2                    2.11.2             pyh9f0ad1d_0    conda-forge
  joblib                    0.15.1                     py_0    conda-forge
  jpeg                      9d                   h516909a_0    conda-forge
  json-c                    0.13.1            hbfbb72e_1002    conda-forge
  jsonschema                3.2.0            py37hc8dfbb8_1    conda-forge
  jupyter-server-proxy      1.5.0                      py_0    conda-forge
  jupyter_client            6.1.3                      py_0    conda-forge
  jupyter_core              4.6.3            py37hc8dfbb8_1    conda-forge
  kealib                    1.4.13               hec59c27_0    conda-forge
  kiwisolver                1.2.0            py37h99015e2_0    conda-forge
  krb5                      1.16.4               h2fd8d38_0    conda-forge
  ld_impl_linux-64          2.33.1               h53a641e_7  
  libblas                   3.8.0               16_openblas    conda-forge
  libcblas                  3.8.0               16_openblas    conda-forge
  libcudf                   0.14.0               cuda10.2_0    rapidsai
  libcugraph                0.14.0               cuda10.2_0    rapidsai
  libcuml                   0.14.0               cuda10.2_0    rapidsai
  libcumlprims              0.14.1               cuda10.2_0    nvidia
  libcurl                   7.68.0               hda55be3_0    conda-forge
  libcuspatial              0.14.0               cuda10.2_0    rapidsai
  libdap4                   3.20.4               hd3bb157_0    conda-forge
  libedit                   3.1.20191231         h7b6447c_0  
  libevent                  2.1.10               hcdb4288_1    conda-forge
  libffi                    3.3                  he6710b0_1  
  libgcc-ng                 9.2.0                h24d8f2e_2    conda-forge
  libgdal                   3.0.2                hc7cfd23_2    conda-forge
  libgfortran-ng            7.5.0                hdf63c60_6    conda-forge
  libhwloc                  2.1.0                h3c4fd83_0    conda-forge
  libiconv                  1.15              h516909a_1006    conda-forge
  libkml                    1.3.0             h4fcabce_1010    conda-forge
  liblapack                 3.8.0               16_openblas    conda-forge
  libllvm8                  8.0.1                hc9558a2_0    conda-forge
  libnetcdf                 4.7.1           nompi_h94020b1_102    conda-forge
  libnvstrings              0.14.0               cuda10.2_0    rapidsai
  libopenblas               0.3.9                h5ec1e0e_0    conda-forge
  libpng                    1.6.37               hed695b0_1    conda-forge
  libpq                     11.5                 hd9ab2ff_2    conda-forge
  libprotobuf               3.8.0                h8b12597_0    conda-forge
  librmm                    0.14.0               cuda10.2_0    rapidsai
  libsodium                 1.0.17               h516909a_0    conda-forge
  libspatialindex           1.9.3                he1b5a44_3    conda-forge
  libspatialite             4.3.0a            h4f6d029_1032    conda-forge
  libssh2                   1.9.0                hab1572f_2    conda-forge
  libstdcxx-ng              9.1.0                hdf63c60_0  
  libtiff                   4.1.0                hfc65ed5_0    conda-forge
  libuuid                   2.32.1            h14c3975_1000    conda-forge
  libuv                     1.34.0               h516909a_0    conda-forge
  libwebp                   1.0.2                hf4e8a37_4    conda-forge
  libxcb                    1.13              h14c3975_1002    conda-forge
  libxgboost                1.1.0dev.rapidsai0.14      cuda10.2_0    rapidsai
  libxml2                   2.9.10               hee79883_0    conda-forge
  llvm-openmp               10.0.0               hc9558a2_0    conda-forge
  llvmlite                  0.32.1           py37h5202443_0    conda-forge
  locket                    0.2.0                      py_2    conda-forge
  lz4-c                     1.8.3             he1b5a44_1001    conda-forge
  markdown                  3.2.2                      py_0    conda-forge
  markupsafe                1.1.1            py37h8f50634_1    conda-forge
  matplotlib-base           3.2.2            py37h30547a4_0    conda-forge
  mistune                   0.8.4           py37h8f50634_1001    conda-forge
  msgpack-python            1.0.0            py37h99015e2_1    conda-forge
  multidict                 4.7.5            py37h8f50634_1    conda-forge
  multipledispatch          0.6.0                      py_0    conda-forge
  munch                     2.5.0                      py_0    conda-forge
  nbconvert                 5.6.1            py37hc8dfbb8_1    conda-forge
  nbformat                  5.0.6                      py_0    conda-forge
  nccl                      2.5.7.1              hc6a2c23_0    conda-forge
  ncurses                   6.2                  he6710b0_1  
  networkx                  2.4                        py_1    conda-forge
  nodejs                    13.13.0              hf5d1a2b_0    conda-forge
  notebook                  6.0.3            py37hc8dfbb8_0    conda-forge
  numba                     0.49.1           py37h0573a6f_0  
  numpy                     1.18.5           py37h8960a57_0    conda-forge
  nvstrings                 0.14.0                   py37_0    rapidsai
  olefile                   0.46                       py_0    conda-forge
  openjpeg                  2.3.1                h981e76c_3    conda-forge
  openssl                   1.1.1g               h516909a_0    conda-forge
  packaging                 20.4               pyh9f0ad1d_0    conda-forge
  pandas                    0.25.3           py37hb3f55d8_0    conda-forge
  pandoc                    2.9.2.1                       0    conda-forge
  pandocfilters             1.4.2                      py_1    conda-forge
  panel                     0.6.4                         0    conda-forge
  param                     1.9.3                      py_0    conda-forge
  parquet-cpp               1.5.1                         2    conda-forge
  parso                     0.7.0              pyh9f0ad1d_0    conda-forge
  partd                     1.1.0                      py_0    conda-forge
  pcre                      8.44                 he1b5a44_0    conda-forge
  pexpect                   4.8.0            py37hc8dfbb8_1    conda-forge
  pickleshare               0.7.5           py37hc8dfbb8_1001    conda-forge
  pillow                    7.1.2            py37h718be6c_0    conda-forge
  pip                       20.1.1                   py37_1  
  pixman                    0.38.0            h516909a_1003    conda-forge
  poppler                   0.67.0               h14e79db_8    conda-forge
  poppler-data              0.4.9                         1    conda-forge
  postgresql                11.5                 hc63931a_2    conda-forge
  proj                      6.2.1                hc80f0dc_0    conda-forge
  prometheus_client         0.8.0              pyh9f0ad1d_0    conda-forge
  prompt-toolkit            3.0.5                      py_1    conda-forge
  psutil                    5.7.0            py37h8f50634_1    conda-forge
  pthread-stubs             0.4               h14c3975_1001    conda-forge
  ptyprocess                0.6.0                   py_1001    conda-forge
  py-xgboost                1.1.0dev.rapidsai0.14  cuda10.2py37_0    rapidsai
  pyarrow                   0.15.0           py37h8b68381_1    conda-forge
  pycparser                 2.20                       py_0    conda-forge
  pyct                      0.4.6                      py_0    conda-forge
  pyct-core                 0.4.6                      py_0    conda-forge
  pyee                      7.0.2              pyh9f0ad1d_0    conda-forge
  pygments                  2.6.1                      py_0    conda-forge
  pynvml                    8.0.4                      py_0    conda-forge
  pyopenssl                 19.1.0                     py_1    conda-forge
  pyparsing                 2.4.7              pyh9f0ad1d_0    conda-forge
  pyppeteer                 0.0.25                     py_1    conda-forge
  pyproj                    2.4.2.post1      py37h12732c1_0    conda-forge
  pyrsistent                0.16.0           py37h8f50634_0    conda-forge
  pysocks                   1.7.1            py37hc8dfbb8_1    conda-forge
  python                    3.7.7                hcff3b4d_5  
  python-dateutil           2.8.1                      py_0    conda-forge
  python_abi                3.7                     1_cp37m    conda-forge
  pytz                      2020.1             pyh9f0ad1d_0    conda-forge
  pyviz_comms               0.7.5              pyh9f0ad1d_0    conda-forge
  pywavelets                1.1.1            py37h03ebfcd_1    conda-forge
  pyyaml                    5.3.1            py37h8f50634_0    conda-forge
  pyzmq                     19.0.1           py37hac76be4_0    conda-forge
  rapids                    0.14.0          cuda10.2_py37_2    rapidsai
  rapids-xgboost            0.14.0          cuda10.2_py37_2    rapidsai
  re2                       2020.04.01           he1b5a44_0    conda-forge
  readline                  8.0                  h7b6447c_0  
  requests                  2.24.0             pyh9f0ad1d_0    conda-forge
  rmm                       0.14.0                   py37_0    rapidsai
  rtree                     0.9.4            py37h8526d28_1    conda-forge
  scikit-image              0.17.2           py37h0da4684_1    conda-forge
  scikit-learn              0.23.1           py37h8a51577_0    conda-forge
  scipy                     1.4.1            py37ha3d9a3c_3    conda-forge
  send2trash                1.5.0                      py_0    conda-forge
  setuptools                47.3.1                   py37_0  
  shapely                   1.6.4           py37hec07ddf_1006    conda-forge
  simpervisor               0.3                        py_1    conda-forge
  six                       1.15.0             pyh9f0ad1d_0    conda-forge
  snappy                    1.1.8                he1b5a44_2    conda-forge
  sortedcontainers          2.2.2              pyh9f0ad1d_0    conda-forge
  spdlog                    1.6.1                hc9558a2_0    conda-forge
  sqlite                    3.32.3               h62c20be_0  
  tbb                       2018.0.5             h2d50403_0    conda-forge
  tblib                     1.6.0                      py_0    conda-forge
  terminado                 0.8.3            py37hc8dfbb8_1    conda-forge
  testpath                  0.4.4                      py_0    conda-forge
  threadpoolctl             2.1.0              pyh5ca1d4c_0    conda-forge
  thrift-cpp                0.12.0            hf3afdfd_1004    conda-forge
  tifffile                  2020.6.3                   py_0    conda-forge
  tiledb                    1.6.2                h7d710e0_2    conda-forge
  tk                        8.6.10               hbc83047_0  
  toolz                     0.10.0                     py_0    conda-forge
  tornado                   6.0.4            py37h8f50634_1    conda-forge
  tqdm                      4.46.1             pyh9f0ad1d_0    conda-forge
  traitlets                 4.3.3            py37hc8dfbb8_1    conda-forge
  tzcode                    2020a                h516909a_0    conda-forge
  ucx                       1.8.0+gf6ec8d4      cuda10.2_20    rapidsai
  ucx-py                    0.14.0+gf6ec8d4          py37_0    rapidsai
  uriparser                 0.9.3                he1b5a44_1    conda-forge
  urllib3                   1.25.9                     py_0    conda-forge
  wcwidth                   0.2.5              pyh9f0ad1d_0    conda-forge
  webencodings              0.5.1                      py_1    conda-forge
  websockets                8.1              py37h8f50634_1    conda-forge
  wheel                     0.34.2                   py37_0  
  xarray                    0.15.1                     py_0    conda-forge
  xerces-c                  3.2.2             h8412b87_1004    conda-forge
  xgboost                   1.1.0dev.rapidsai0.14  cuda10.2py37_0    rapidsai
  xorg-kbproto              1.0.7             h14c3975_1002    conda-forge
  xorg-libice               1.0.10               h516909a_0    conda-forge
  xorg-libsm                1.2.3             h84519dc_1000    conda-forge
  xorg-libx11               1.6.9                h516909a_0    conda-forge
  xorg-libxau               1.0.9                h14c3975_0    conda-forge
  xorg-libxdmcp             1.1.3                h516909a_0    conda-forge
  xorg-libxext              1.3.4                h516909a_0    conda-forge
  xorg-libxrender           0.9.10            h516909a_1002    conda-forge
  xorg-renderproto          0.11.1            h14c3975_1002    conda-forge
  xorg-xextproto            7.3.0             h14c3975_1002    conda-forge
  xorg-xproto               7.0.31            h14c3975_1007    conda-forge
  xz                        5.2.5                h7b6447c_0  
  yaml                      0.2.5                h516909a_0    conda-forge
  yarl                      1.3.0           py37h516909a_1000    conda-forge
  zeromq                    4.3.2                he1b5a44_2    conda-forge
  zict                      2.0.0                      py_0    conda-forge
  zipp                      3.1.0                      py_0    conda-forge
  zlib                      1.2.11               h7b6447c_3  
  zstd                      1.4.3                h3b9ef0a_0    conda-forge

Additional context Tried on a clean conda environment; also in the latest RAPIDS 0.15 nightly build.

[cjnolet]

@Banus

Thank you for opening this issue. I have verified this crash also occurs in our pytests when the data sizes are set to 20k rows and 248 cols.