fix(python): Fix detection of cuda library on hosted runner (#554)

Closes #551.
apache · Jul 9, 2024 · e19c17f · e19c17f
1 parent fa8c3da
commit e19c17f
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 14 deletions.
diff --git a/.github/workflows/build-and-test-device.yaml b/.github/workflows/build-and-test-device.yaml
@@ -28,6 +28,7 @@ on:
       - 'CMakeLists.txt'
       - '.github/workflows/build-and-test-device.yaml'
       - 'src/nanoarrow/**'
+      - 'python/setup.py'
 
 permissions:
   contents: read
@@ -132,18 +133,30 @@ jobs:
           cd build
           ctest -T memcheck .
 
-      - name: Check Python bindings with CUDA
+      - name: Upload memcheck results
+        if: failure() && matrix.config.label == 'default-build'
+        uses: actions/upload-artifact@main
+        with:
+          name: nanoarrow-device-memcheck
+          path: build/Testing/Temporary/MemoryChecker.*.log
+
+      - name: Set up Python
+        if: matrix.config.label == 'with-cuda'
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+          cache: 'pip'
+
+      - name: Build Python bindings with CUDA
         if: matrix.config.label == 'with-cuda'
         env:
           NANOARROW_PYTHON_CUDA_HOME: "/usr/local/cuda"
         run: |
           cd python
           python3 -m pip install ".[test]" -vv
-          python3 -m pytest -vv
 
-      - name: Upload memcheck results
-        if: failure() && matrix.config.label == 'default-build'
-        uses: actions/upload-artifact@main
-        with:
-          name: nanoarrow-device-memcheck
-          path: build/Testing/Temporary/MemoryChecker.*.log
+      - name: Test Python bindings with CUDA
+        if: matrix.config.label == 'with-cuda'
+        run: |
+          cd python
+          python3 -m pytest -vv
diff --git a/python/setup.py b/python/setup.py
@@ -83,16 +83,16 @@ def get_version(pkg_path):
     if not include_dir.is_dir():
         raise ValueError(f"CUDA include directory does not exist: '{include_dir}'")
 
-    lib_dirs = [d for d in possible_libs if d.exists()]
-    if not lib_dirs:
-        lib_dirs_err = ", ".join(f"'{d}" for d in possible_libs)
-        raise ValueError(f"Can't find CUDA library directory. Checked {lib_dirs_err}")
-
     device_include_dirs.append(str(include_dir))
-    device_library_dirs.append(str(lib_dirs[0].parent))
     device_libraries.append("cuda")
     device_define_macros.append(("NANOARROW_DEVICE_WITH_CUDA", 1))
 
+    # Library might be already in a system library directory such that no -L flag
+    # is needed
+    lib_dirs = [d for d in possible_libs if d.exists()]
+    if lib_dirs:
+        device_library_dirs.append(str(lib_dirs[0].parent))
+
 
 def nanoarrow_extension(
     name, *, nanoarrow_c=False, nanoarrow_device=False, nanoarrow_ipc=False