Add cuda-python as explicit dependency.

rapidsai · Jan 12, 2022 · 206535d · 206535d
1 parent 2454435
commit 206535d
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 4 deletions.
diff --git 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
@@ -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
@@ -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"