Pin dask and distributed for release (#4859)

This PR pins `dask` & `distributed` to `2022.9.2` for `22.10` release.

~~xref: rapidsai/cudf#11492
xref: rapidsai/cudf#11822

Authors:
  - GALI PREM SAGAR (https://github.com/galipremsagar)
  - Dante Gama Dessavre (https://github.com/dantegd)
  - Corey J. Nolet (https://github.com/cjnolet)

Approvers:
  - Ray Douglass (https://github.com/raydouglass)
  - Dante Gama Dessavre (https://github.com/dantegd)

URL: #4859

ci/gpu/build.sh

@@ -125,8 +125,8 @@ if [[ -z "$PROJECT_FLASH" || "$PROJECT_FLASH" == "0" ]]; then
 
     gpuci_logger "Install the main version of dask and distributed"
     set -x
-    pip install "git+https://github.com/dask/distributed.git@2022.7.1" --upgrade --no-deps
-    pip install "git+https://github.com/dask/dask.git@2022.7.1" --upgrade --no-deps
+    pip install "git+https://github.com/dask/distributed.git@2022.9.2" --upgrade --no-deps
+    pip install "git+https://github.com/dask/dask.git@2022.9.2" --upgrade --no-deps
     pip install "git+https://github.com/hdbscan/hdbscan.git@master" --force-reinstall --upgrade --no-deps
     set +x
 
@@ -199,8 +199,8 @@ else
     gpuci_logger "Install the main version of dask, distributed, and dask-glm"
     set -x
 
-    pip install "git+https://github.com/dask/distributed.git@2022.7.1" --upgrade --no-deps
-    pip install "git+https://github.com/dask/dask.git@2022.7.1" --upgrade --no-deps
+    pip install "git+https://github.com/dask/distributed.git@2022.9.2" --upgrade --no-deps
+    pip install "git+https://github.com/dask/dask.git@2022.9.2" --upgrade --no-deps
     pip install "git+https://github.com/dask/dask-glm@main" --force-reinstall --no-deps
     pip install "git+https://github.com/scikit-learn-contrib/hdbscan.git@master" --force-reinstall --upgrade --no-deps
     pip install sparse

conda/environments/cuml_dev_cuda11.0.yml

@@ -37,8 +37,8 @@ dependencies:
 - nltk
 - pip
 - pip:
-    - git+https://github.com/dask/dask.git@2022.7.1
-    - git+https://github.com/dask/distributed.git@2022.7.1
+    - git+https://github.com/dask/dask.git@2022.9.2
+    - git+https://github.com/dask/distributed.git@2022.9.2
     - git+https://github.com/dask/dask-glm@main
     - git+https://github.com/scikit-learn-contrib/hdbscan.git@master
     - sparse

conda/environments/cuml_dev_cuda11.2.yml

@@ -37,8 +37,8 @@ dependencies:
 - nltk
 - pip
 - pip:
-    - git+https://github.com/dask/dask.git@2022.7.1
-    - git+https://github.com/dask/distributed.git@2022.7.1
+    - git+https://github.com/dask/dask.git@2022.9.2
+    - git+https://github.com/dask/distributed.git@2022.9.2
     - git+https://github.com/dask/dask-glm@main
     - git+https://github.com/scikit-learn-contrib/hdbscan.git@master
     - sparse

conda/environments/cuml_dev_cuda11.4.yml

@@ -37,8 +37,8 @@ dependencies:
 - nltk
 - pip
 - pip:
-    - git+https://github.com/dask/dask.git@2022.7.1
-    - git+https://github.com/dask/distributed.git@2022.7.1
+    - git+https://github.com/dask/dask.git@2022.9.2
+    - git+https://github.com/dask/distributed.git@2022.9.2
     - git+https://github.com/dask/dask-glm.git@main
     - git+https://github.com/scikit-learn-contrib/hdbscan.git@master
     - sparse

conda/environments/cuml_dev_cuda11.5.yml

@@ -37,8 +37,8 @@ dependencies:
 - nltk
 - pip
 - pip:
-    - git+https://github.com/dask/dask.git@2022.7.1
-    - git+https://github.com/dask/distributed.git@2022.7.1
+    - git+https://github.com/dask/dask.git@2022.9.2
+    - git+https://github.com/dask/distributed.git@2022.9.2
     - git+https://github.com/dask/dask-glm@main
     - git+https://github.com/scikit-learn-contrib/hdbscan.git@master
     - sparse

conda/recipes/cuml/meta.yaml

@@ -59,8 +59,8 @@ requirements:
     - nccl>=2.9.9
     - ucx-py {{ ucx_py_version }}
     - ucx-proc=*=gpu
-    - dask==2022.7.1
-    - distributed==2022.7.1
+    - dask==2022.9.2
+    - distributed==2022.9.2
     - joblib >=0.11
     - {{ pin_compatible('cudatoolkit', max_pin='x', min_pin='x') }}
     - cuda-python >=11.5,<11.7.1

cpp/bench/sg/linkage.cu

@@ -73,7 +73,7 @@ class Linkage : public BlobsFixture<D> {
  private:
   int* labels;
   int* out_children;
-  raft::hierarchy::linkage_output<int, D> out_arrs;
+  raft::hierarchy::linkage_output<int> out_arrs;
 };
 
 std::vector<Params> getInputs()

cpp/include/cuml/cluster/linkage.hpp

@@ -46,7 +46,7 @@ void single_linkage_pairwise(const raft::handle_t& handle,
                              const float* X,
                              size_t m,
                              size_t n,
-                             raft::hierarchy::linkage_output<int, float>* out,
+                             raft::hierarchy::linkage_output<int>* out,
                              raft::distance::DistanceType metric,
                              int n_clusters = 5);
 
@@ -74,7 +74,7 @@ void single_linkage_neighbors(
   const float* X,
   size_t m,
   size_t n,
-  raft::hierarchy::linkage_output<int, float>* out,
+  raft::hierarchy::linkage_output<int>* out,
   raft::distance::DistanceType metric = raft::distance::DistanceType::L2Unexpanded,
   int c                               = 15,
   int n_clusters                      = 5);
@@ -83,7 +83,7 @@ void single_linkage_pairwise(const raft::handle_t& handle,
                              const float* X,
                              size_t m,
                              size_t n,
-                             raft::hierarchy::linkage_output<int64_t, float>* out,
+                             raft::hierarchy::linkage_output<int64_t>* out,
                              raft::distance::DistanceType metric,
                              int n_clusters = 5);
 

cpp/src/glm/ols.cuh

@@ -100,7 +100,7 @@ void olsFit(const raft::handle_t& handle,
     raft::linalg::sqrt(sample_weight, sample_weight, n_rows, stream);
     raft::matrix::matrixVectorBinaryMult(
       input, sample_weight, n_rows, n_cols, false, false, stream);
-    raft::linalg::map(
+    raft::linalg::map_k(
       labels,
       n_rows,
       [] __device__(math_t a, math_t b) { return a * b; },
@@ -129,7 +129,7 @@ void olsFit(const raft::handle_t& handle,
   if (sample_weight != nullptr) {
     raft::matrix::matrixVectorBinaryDivSkipZero(
       input, sample_weight, n_rows, n_cols, false, false, stream);
-    raft::linalg::map(
+    raft::linalg::map_k(
       labels,
       n_rows,
       [] __device__(math_t a, math_t b) { return a / b; },

cpp/src/glm/qn/glm_base.cuh

@@ -158,7 +158,7 @@ struct GLMBase : GLMDims {
         y.data,
         Z.data,
         sample_weights);
-      raft::linalg::map(
+      raft::linalg::map_k(
         Z.data,
         y.len,
         [dlz_copy] __device__(const T y, const T z, const T weight) {

cpp/src/glm/ridge.cuh

@@ -19,6 +19,7 @@
 #include <raft/core/cudart_utils.hpp>
 #include <raft/linalg/add.cuh>
 #include <raft/linalg/gemm.cuh>
+#include <raft/linalg/map.cuh>
 #include <raft/linalg/norm.cuh>
 #include <raft/linalg/subtract.cuh>
 #include <raft/linalg/svd.cuh>
@@ -196,7 +197,7 @@ void ridgeFit(const raft::handle_t& handle,
     raft::linalg::sqrt(sample_weight, sample_weight, n_rows, stream);
     raft::matrix::matrixVectorBinaryMult(
       input, sample_weight, n_rows, n_cols, false, false, stream);
-    raft::linalg::map(
+    raft::linalg::map_k(
       labels,
       n_rows,
       [] __device__(math_t a, math_t b) { return a * b; },
@@ -218,7 +219,7 @@ void ridgeFit(const raft::handle_t& handle,
   if (sample_weight != nullptr) {
     raft::matrix::matrixVectorBinaryDivSkipZero(
       input, sample_weight, n_rows, n_cols, false, false, stream);
-    raft::linalg::map(
+    raft::linalg::map_k(
       labels,
       n_rows,
       [] __device__(math_t a, math_t b) { return a / b; },

cpp/src/hierarchy/linkage.cu

@@ -28,7 +28,7 @@ void single_linkage_pairwise(const raft::handle_t& handle,
                              const float* X,
                              size_t m,
                              size_t n,
-                             raft::cluster::linkage_output<int, float>* out,
+                             raft::cluster::linkage_output<int>* out,
                              raft::distance::DistanceType metric,
                              int n_clusters)
 {
@@ -40,7 +40,7 @@ void single_linkage_neighbors(const raft::handle_t& handle,
                               const float* X,
                               size_t m,
                               size_t n,
-                              raft::cluster::linkage_output<int, float>* out,
+                              raft::cluster::linkage_output<int>* out,
                               raft::distance::DistanceType metric,
                               int c,
                               int n_clusters)

cpp/src/solver/cd.cuh

@@ -194,7 +194,7 @@ void cdFit(const raft::handle_t& handle,
     raft::linalg::sqrt(sample_weight, sample_weight, n_rows, stream);
     raft::matrix::matrixVectorBinaryMult(
       input, sample_weight, n_rows, n_cols, false, false, stream);
-    raft::linalg::map(
+    raft::linalg::map_k(
       labels,
       n_rows,
       [] __device__(math_t a, math_t b) { return a * b; },
@@ -286,7 +286,7 @@ void cdFit(const raft::handle_t& handle,
   if (sample_weight != nullptr) {
     raft::matrix::matrixVectorBinaryDivSkipZero(
       input, sample_weight, n_rows, n_cols, false, false, stream);
-    raft::linalg::map(
+    raft::linalg::map_k(
       labels,
       n_rows,
       [] __device__(math_t a, math_t b) { return a / b; },

cpp/src_prims/metrics/kl_divergence.cuh

@@ -67,7 +67,7 @@ DataT kl_divergence(const DataT* modelPDF, const DataT* candidatePDF, int size,
   rmm::device_scalar<DataT> d_KLDVal(stream);
   RAFT_CUDA_TRY(cudaMemsetAsync(d_KLDVal.data(), 0, sizeof(DataT), stream));
 
-  raft::linalg::mapThenSumReduce<DataT, KLDOp<DataT>, 256, const DataT*>(
+  raft::linalg::mapThenSumReduce<DataT, KLDOp<DataT>, size_t, 256, const DataT*>(
     d_KLDVal.data(), (size_t)size, KLDOp<DataT>(), stream, modelPDF, candidatePDF);
 
   DataT h_KLDVal;

cpp/test/sg/linkage_test.cu

@@ -81,7 +81,7 @@ class LinkageTest : public ::testing::TestWithParam<LinkageInputs<T, IdxT>> {
 
     handle.sync_stream(handle.get_stream());
 
-    raft::hierarchy::linkage_output<IdxT, T> out_arrs;
+    raft::hierarchy::linkage_output<IdxT> out_arrs;
     out_arrs.labels = labels.data();
 
     rmm::device_uvector<IdxT> out_children((params.n_row - 1) * 2, handle.get_stream());

cpp/test/sg/tsne_test.cu

@@ -67,7 +67,7 @@ float get_kl_div(TSNEParams& params,
   auto get_emb_dist = [=] __device__(const int64_t i, const int64_t j) {
     return emb_dists[i * n + j];
   };
-  raft::linalg::map(Qs, total_nn, get_emb_dist, stream, input_matrix.rows(), input_matrix.cols());
+  raft::linalg::map_k(Qs, total_nn, get_emb_dist, stream, input_matrix.rows(), input_matrix.cols());
 
   const float dof      = fmaxf(params.dim - 1, 1);  // degree of freedom
   const float exponent = (dof + 1.0) / 2.0;

python/cuml/cluster/agglomerative.pyx

@@ -34,7 +34,7 @@ from cuml.metrics.distance_type cimport DistanceType
 
 cdef extern from "raft/sparse/hierarchy/common.h" namespace "raft::hierarchy":
 
-    cdef cppclass linkage_output_int_float:
+    cdef cppclass linkage_output_int:
         int m
         int n_clusters
         int n_leaves
@@ -49,7 +49,7 @@ cdef extern from "cuml/cluster/linkage.hpp" namespace "ML":
         const float *X,
         size_t m,
         size_t n,
-        linkage_output_int_float *out,
+        linkage_output_int *out,
         DistanceType metric,
         int n_clusters
     ) except +
@@ -59,7 +59,7 @@ cdef extern from "cuml/cluster/linkage.hpp" namespace "ML":
         const float *X,
         size_t m,
         size_t n,
-        linkage_output_int_float *out,
+        linkage_output_int *out,
         DistanceType metric,
         int c,
         int n_clusters
@@ -202,7 +202,7 @@ class AgglomerativeClustering(Base, ClusterMixin, CMajorInputTagMixin):
         cdef uintptr_t labels_ptr = self.labels_.ptr
         cdef uintptr_t children_ptr = self.children_.ptr
 
-        cdef linkage_output_int_float linkage_output
+        cdef linkage_output_int linkage_output
         linkage_output.children = <int*>children_ptr
         linkage_output.labels = <int*>labels_ptr
 
@@ -215,13 +215,13 @@ class AgglomerativeClustering(Base, ClusterMixin, CMajorInputTagMixin):
         if self.connectivity == 'knn':
             single_linkage_neighbors(
                 handle_[0], <float*>input_ptr, <int> n_rows,
-                <int> n_cols, <linkage_output_int_float*> &linkage_output,
+                <int> n_cols, <linkage_output_int*> &linkage_output,
                 <DistanceType> metric, <int>self.n_neighbors,
                 <int> self.n_clusters)
         elif self.connectivity == 'pairwise':
             single_linkage_pairwise(
                 handle_[0], <float*>input_ptr, <int> n_rows,
-                <int> n_cols, <linkage_output_int_float*> &linkage_output,
+                <int> n_cols, <linkage_output_int*> &linkage_output,
                 <DistanceType> metric, <int> self.n_clusters)
         else:
             raise ValueError("'connectivity' can only be one of "

python/cuml/tests/dask/test_dask_one_hot_encoder.py

@@ -52,7 +52,8 @@ def test_onehot_inverse_transform(client, drop):
     enc = OneHotEncoder(drop=drop)
     ohe = enc.fit_transform(X)
     inv = enc.inverse_transform(ohe)
-    assert_frame_equal(inv.compute().to_pandas(), df.to_pandas())
+    assert_frame_equal(inv.compute().to_pandas().reset_index(drop=True),
+                       df.to_pandas())
 
 
 @pytest.mark.mg
@@ -158,7 +159,8 @@ def test_onehot_drop_idx_first(client):
     ref = sk_enc.fit_transform(X_ary)
     cp.testing.assert_array_equal(ohe.compute(), ref)
     inv = enc.inverse_transform(ohe)
-    assert_frame_equal(inv.compute().to_pandas(), X.to_pandas())
+    assert_frame_equal(inv.compute().to_pandas().reset_index(drop=True),
+                       X.to_pandas())
 
 
 @pytest.mark.mg
@@ -175,7 +177,8 @@ def test_onehot_drop_one_of_each(client):
     ref = sk_enc.fit_transform(X_ary)
     cp.testing.assert_array_equal(ohe.compute(), ref)
     inv = enc.inverse_transform(ohe)
-    assert_frame_equal(inv.compute().to_pandas(), X.to_pandas())
+    assert_frame_equal(inv.compute().to_pandas().reset_index(drop=True),
+                       X.to_pandas())
 
 
 @pytest.mark.mg