From 206535d5886e282655eddeacd79e61992ec1aa9f Mon Sep 17 00:00:00 2001
From: Ashwin Srinath <shwina@users.noreply.github.com>
Date: Wed, 12 Jan 2022 10:50:47 -0500
Subject: [PATCH] Add cuda-python as explicit dependency.

---
 conda/environments/cudf_dev_cuda11.5.yml |  1 +
 conda/recipes/cudf/meta.yaml             |  2 +-
 python/cudf/cudf/utils/gpu_utils.py      | 10 +++++++---
 3 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/conda/environments/cudf_dev_cuda11.5.yml b/conda/environments/cudf_dev_cuda11.5.yml
index cc8d50a1717..841aaaa53f1 100644
--- a/conda/environments/cudf_dev_cuda11.5.yml
+++ b/conda/environments/cudf_dev_cuda11.5.yml
@@ -33,6 +33,7 @@ dependencies:
   - ipython
   - pandoc=<2.0.0
   - cudatoolkit=11.5
+  - cuda-python >=11.5,<12.0
   - pip
   - flake8=3.8.3
   - black=19.10
diff --git a/conda/recipes/cudf/meta.yaml b/conda/recipes/cudf/meta.yaml
index 2600ab358cc..1035655a190 100644
--- a/conda/recipes/cudf/meta.yaml
+++ b/conda/recipes/cudf/meta.yaml
@@ -52,7 +52,7 @@ requirements:
     - packaging
     - cachetools
     - ptxcompiler  # [linux64]  # CUDA enhanced compatibility. See https://github.com/rapidsai/ptxcompiler
-
+    - cuda-python >=11.5,<12.0
 test:                                   # [linux64]
   requires:                             # [linux64]
     - cudatoolkit {{ cuda_version }}.*  # [linux64]
diff --git a/python/cudf/cudf/utils/gpu_utils.py b/python/cudf/cudf/utils/gpu_utils.py
index afbd625cc2b..8947760e052 100644
--- a/python/cudf/cudf/utils/gpu_utils.py
+++ b/python/cudf/cudf/utils/gpu_utils.py
@@ -15,7 +15,7 @@ def validate_setup():
 
     import warnings
 
-    from cuda.cudart import cudaError_t
+    from cuda.cudart import cudaDeviceAttr, cudaError_t
 
     from rmm._cuda.gpu import (
         CUDARuntimeError,
@@ -68,7 +68,9 @@ def _try_get_old_or_new_symbols():
         # Cupy throws RunTimeException to get GPU count,
         # hence obtaining GPU count by in-house cpp api above
 
-        major_version = getDeviceAttribute("ComputeCapabilityMajor", 0)
+        major_version = getDeviceAttribute(
+            cudaDeviceAttr.cudaDevAttrComputeCapabilityMajor, 0
+        )
 
         if major_version >= 6:
             # You have a GPU with NVIDIA Pascal™ architecture or better
@@ -82,7 +84,9 @@ def _try_get_old_or_new_symbols():
             pass
         else:
             device_name = deviceGetName(0)
-            minor_version = getDeviceAttribute("ComputeCapabilityMinor", 0)
+            minor_version = getDeviceAttribute(
+                cudaDeviceAttr.cudaDevAttrComputeCapabilityMinor, 0
+            )
             warnings.warn(
                 f"You will need a GPU with NVIDIA Pascal™ or "
                 f"newer architecture"