Dispatch thrust versions and upgrade rmm.

dmlc · Sep 24, 2021 · a2a72d2 · a2a72d2
1 parent 0ee11da
commit a2a72d2
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 5 deletions.
diff --git a/src/common/device_helpers.cuh b/src/common/device_helpers.cuh
@@ -708,6 +708,10 @@ constexpr std::pair<int, int> CUDAVersion() {
 #endif  // defined(__CUDACC_VER_MAJOR__)
 }
 
+constexpr std::pair<int32_t, int32_t> ThrustVersion() {
+  return std::make_pair(THRUST_MAJOR_VERSION, THRUST_MAJOR_VERSION);
+}
+
 namespace detail {
 template <typename T>
 using TypedDiscardCTK114 = thrust::discard_iterator<T>;
@@ -721,9 +725,9 @@ class TypedDiscard : public thrust::discard_iterator<T> {
 
 template <typename T>
 using TypedDiscard =
-    std::conditional_t<((CUDAVersion().first == 11 &&
-                         CUDAVersion().second >= 4) ||
-                        CUDAVersion().first > 11),
+    std::conditional_t<((ThrustVersion().first == 1 &&
+                         ThrustVersion().second >= 12) ||
+                        ThrustVersion().first > 1),
                        detail::TypedDiscardCTK114<T>, detail::TypedDiscard<T>>;
 
 /**

diff --git a/tests/ci_build/Dockerfile.gpu b/tests/ci_build/Dockerfile.gpu
@@ -19,7 +19,7 @@ ENV PATH=/opt/python/bin:$PATH
 # Create new Conda environment with cuDF, Dask, and cuPy
 RUN \
     conda create -n gpu_test -c rapidsai-nightly -c rapidsai -c nvidia -c conda-forge -c defaults \
-        python=3.7 cudf=21.08* rmm=21.08* cudatoolkit=$CUDA_VERSION_ARG dask dask-cuda dask-cudf cupy=9.1* \
+        python=3.8 cudf=21.10* rmm=21.10* cudatoolkit=$CUDA_VERSION_ARG dask dask-cuda=21.10* dask-cudf=21.10* cupy=9.1* \
         numpy pytest scipy scikit-learn pandas matplotlib wheel python-kubernetes urllib3 graphviz hypothesis
 
 ENV GOSU_VERSION 1.10

diff --git a/tests/ci_build/Dockerfile.rmm b/tests/ci_build/Dockerfile.rmm
@@ -29,7 +29,7 @@ ENV PATH=/opt/python/bin:$PATH
 # Create new Conda environment with RMM
 RUN \
     conda create -n gpu_test -c rapidsai-nightly -c rapidsai -c nvidia -c conda-forge -c defaults \
-        python=3.7 rmm=21.08* cudatoolkit=$CUDA_VERSION_ARG
+        python=3.8 rmm=21.10* cudatoolkit=$CUDA_VERSION_ARG
 
 ENV GOSU_VERSION 1.10