Improve setup by removing pycuda dependency and adding cuda runtime and cublas to RPATH

lmdeploy/__init__.py

@@ -21,7 +21,11 @@ def bootstrap():
     if os.name == 'nt' and has_turbomind:
         if sys.version_info[:2] >= (3, 8):
             CUDA_PATH = os.getenv('CUDA_PATH')
-            os.add_dll_directory(os.path.join(CUDA_PATH, 'bin'))
+            assert CUDA_PATH is not None, 'Can not find $env:CUDA_PATH'
+            dll_path = os.path.join(CUDA_PATH, 'bin')
+            print('Add dll path {dll_path}, please note cuda version '
+                  'should >= 11.3 when compiled with cuda 11')
+            os.add_dll_directory(dll_path)
 
 
 bootstrap()

lmdeploy/pytorch/utils.py

@@ -5,16 +5,17 @@
 from typing import Dict, Sequence
 
 import psutil
-import pycuda.driver as drv
+from pynvml import (nvmlDeviceGetHandleByIndex, nvmlDeviceGetMemoryInfo,
+                    nvmlInit)
 
 
 def get_gpu_memory(id: int = 0) -> int:
     """Returns the free and total physical memory of the GPU in bytes."""
-    drv.init()
-    dev = drv.Device(id)
-    cxt = dev.make_context()
-    free, total = drv.mem_get_info()
-    cxt.pop()
+    nvmlInit()
+    handle = nvmlDeviceGetHandleByIndex(id)
+    mem_info = nvmlDeviceGetMemoryInfo(handle)
+    free = mem_info.free
+    total = mem_info.total
     return free, total
 
 

requirements/runtime.txt

@@ -2,8 +2,8 @@ fastapi
 fire
 mmengine-lite
 numpy
-pycuda
 pydantic>2.0.0
+pynvml
 safetensors
 sentencepiece
 shortuuid

setup.py

@@ -26,6 +26,32 @@ def check_ext_modules():
     return False
 
 
+def get_cuda_pkgs():
+    arg_name = '--cuda='
+    arg_value = None
+    for arg in sys.argv[1:]:
+        if arg.startswith(arg_name):
+            arg_value = arg[len(arg_name):]
+            sys.argv.remove(arg)
+            break
+
+    cuda_pkgs = []
+    if arg_value == '11':
+        cuda_pkgs = [
+            'nvidia-nccl-cu11', 'nvidia-cuda-runtime-cu11',
+            'nvidia-cublas-cu11'
+        ]
+    elif arg_value == '12':
+        cuda_pkgs = [
+            'nvidia-nccl-cu12', 'nvidia-cuda-runtime-cu12',
+            'nvidia-cublas-cu12'
+        ]
+    return cuda_pkgs
+
+
+cuda_pkgs = get_cuda_pkgs()
+
+
 def parse_requirements(fname='requirements.txt', with_version=True):
     """Parse the package dependencies listed in a file but strips specific
     versioning information.
@@ -42,21 +68,6 @@ def parse_requirements(fname='requirements.txt', with_version=True):
     """
     require_fpath = fname
 
-    def get_nccl_pkg():
-        arg_name = '--cuda='
-        arg_value = None
-        for arg in sys.argv[1:]:
-            if arg.startswith(arg_name):
-                arg_value = arg[len(arg_name):]
-                sys.argv.remove(arg)
-                break
-
-        if arg_value == '11':
-            return 'nvidia-nccl-cu11'
-        elif arg_value == '12':
-            return 'nvidia-nccl-cu12'
-        return None
-
     def parse_line(line):
         """Parse information from a line in a requirements text file."""
         if line.startswith('-r '):
@@ -113,9 +124,7 @@ def gen_packages_items():
                 yield item
 
     packages = list(gen_packages_items())
-    nccl_pkg = get_nccl_pkg()
-    if nccl_pkg is not None:
-        packages += [nccl_pkg]
+    packages += cuda_pkgs
     return packages
 
 

src/turbomind/python/CMakeLists.txt

@@ -17,6 +17,13 @@ target_link_libraries(${PROJECT_NAME} PRIVATE TransformerTritonBackend
     LlamaTritonBackend)
 target_compile_features(${PROJECT_NAME} PRIVATE cxx_std_14)
 
+set(_INSTALL_CUDA_RPATH
+    "\$ORIGIN"
+    "\$ORIGIN/../../nvidia/nccl/lib/"
+    "\$ORIGIN/../../nvidia/cuda_runtime/lib/"
+    "\$ORIGIN/../../nvidia/cublas/lib/"
+)
 set_target_properties(${PROJECT_NAME} PROPERTIES
-        BUILD_RPATH "\$ORIGIN"
-	INSTALL_RPATH "\$ORIGIN;\$ORIGIN/../../nvidia/nccl/lib/")
+    BUILD_RPATH "\$ORIGIN"
+    INSTALL_RPATH "${_INSTALL_CUDA_RPATH}"
+)