From 0f2c03d17d507bd091be4a7dc560c9690ac8dece Mon Sep 17 00:00:00 2001
From: Mahesh Doijade <mdoijade@nvidia.com>
Date: Mon, 7 Jun 2021 18:42:30 +0530
Subject: [PATCH 1/3] fix build failure of mst and knn test by adding cuda
 stream arg to rmm::device_buffer

---
 cpp/test/mst.cu         | 18 +++++++++---------
 cpp/test/spatial/knn.cu |  6 ++++--
 2 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/cpp/test/mst.cu b/cpp/test/mst.cu
index 215c6f6548..b0007c73fd 100644
--- a/cpp/test/mst.cu
+++ b/cpp/test/mst.cu
@@ -198,15 +198,15 @@ class MSTTest
       MSTTestInput<vertex_t, edge_t, weight_t>>::GetParam();
     iterations = mst_input.iterations;
 
-    csr_d.offsets =
-      rmm::device_buffer(mst_input.csr_h.offsets.data(),
-                         mst_input.csr_h.offsets.size() * sizeof(edge_t));
-    csr_d.indices =
-      rmm::device_buffer(mst_input.csr_h.indices.data(),
-                         mst_input.csr_h.indices.size() * sizeof(vertex_t));
-    csr_d.weights =
-      rmm::device_buffer(mst_input.csr_h.weights.data(),
-                         mst_input.csr_h.weights.size() * sizeof(weight_t));
+    csr_d.offsets = rmm::device_buffer(
+      mst_input.csr_h.offsets.data(),
+      mst_input.csr_h.offsets.size() * sizeof(edge_t), handle.get_stream());
+    csr_d.indices = rmm::device_buffer(
+      mst_input.csr_h.indices.data(),
+      mst_input.csr_h.indices.size() * sizeof(vertex_t), handle.get_stream());
+    csr_d.weights = rmm::device_buffer(
+      mst_input.csr_h.weights.data(),
+      mst_input.csr_h.weights.size() * sizeof(weight_t), handle.get_stream());
   }
 
   void TearDown() override {}
diff --git a/cpp/test/spatial/knn.cu b/cpp/test/spatial/knn.cu
index cfd4ecc9d1..2b1ef89f7a 100644
--- a/cpp/test/spatial/knn.cu
+++ b/cpp/test/spatial/knn.cu
@@ -107,11 +107,13 @@ class KNNTest : public ::testing::TestWithParam<KNNInputs> {
       }
     }
     rmm::device_buffer input_d = rmm::device_buffer(
-      row_major_input.data(), row_major_input.size() * sizeof(float));
+      row_major_input.data(), row_major_input.size() * sizeof(float),
+      handle_.get_stream());
     float *input_ptr = static_cast<float *>(input_d.data());
 
     rmm::device_buffer labels_d = rmm::device_buffer(
-      params_.labels.data(), params_.labels.size() * sizeof(int));
+      params_.labels.data(), params_.labels.size() * sizeof(int),
+      handle_.get_stream());
     int *labels_ptr = static_cast<int *>(labels_d.data());
 
     raft::allocate(input_, rows_ * cols_, true);

From 484b0820f0cc3ce3c33b1a0ebb90a505513aa7ea Mon Sep 17 00:00:00 2001
From: Mahesh Doijade <mdoijade@nvidia.com>
Date: Mon, 7 Jun 2021 18:59:34 +0530
Subject: [PATCH 2/3] temp commit for test rerun

---
 cpp/test/mst.cu | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cpp/test/mst.cu b/cpp/test/mst.cu
index b0007c73fd..d7aa76500b 100644
--- a/cpp/test/mst.cu
+++ b/cpp/test/mst.cu
@@ -58,7 +58,7 @@ weight_t prims(CSRHost<vertex_t, edge_t, weight_t> &csr_h) {
   auto n_vertices = csr_h.offsets.size() - 1;
 
   bool active_vertex[n_vertices];
-  // bool mst_set[csr_h.n_edges];
+  //  bool mst_set[csr_h.n_edges];
   weight_t curr_edge[n_vertices];
 
   for (auto i = 0; i < n_vertices; i++) {

From 04f656f97ee82b75046501615fa2927781880a6c Mon Sep 17 00:00:00 2001
From: Mahesh Doijade <mdoijade@nvidia.com>
Date: Mon, 7 Jun 2021 19:14:01 +0530
Subject: [PATCH 3/3] use ucx-py version 0.21 to temp resolve ci build failures

---
 ci/gpu/build.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ci/gpu/build.sh b/ci/gpu/build.sh
index b4eef5ebd0..1707b72f44 100644
--- a/ci/gpu/build.sh
+++ b/ci/gpu/build.sh
@@ -51,7 +51,7 @@ gpuci_conda_retry install -c conda-forge -c rapidsai -c rapidsai-nightly -c nvid
       "rmm=${MINOR_VERSION}" \
       "dask-cudf=${MINOR_VERSION}" \
       "dask-cuda=${MINOR_VERSION}" \
-      "ucx-py=${MINOR_VERSION}" \
+      "ucx-py=0.21" \
       "rapids-build-env=${MINOR_VERSION}.*" \
       "rapids-notebook-env=${MINOR_VERSION}.*" \
       "rapids-doc-env=${MINOR_VERSION}.*"