Merge branch 'main' into dching/event-timing-message

Author: leofang

.bandit

@@ -0,0 +1,2 @@
+[bandit]
+skips = B101,B311

.github/actions/fetch_ctk/action.yml

@@ -9,15 +9,22 @@ inputs:
   cuda-version:
     required: true
     type: string
+  cuda-components:
+    description: "A list of the CTK components to install as a comma-separated list. e.g. 'cuda_nvcc,cuda_nvrtc,cuda_cudart'"
+    required: false
+    type: string
+    default: "cuda_nvcc,cuda_cudart,cuda_nvrtc,cuda_profiler_api,cuda_cccl,cuda_sanitizer_api,libnvjitlink"
 
 runs:
   using: composite
   steps:
     - name: Set up CTK cache variable
       shell: bash --noprofile --norc -xeuo pipefail {0}
       run: |
-        echo "CTK_CACHE_KEY=mini-ctk-${{ inputs.cuda-version }}-${{ inputs.host-platform }}" >> $GITHUB_ENV
-        echo "CTK_CACHE_FILENAME=mini-ctk-${{ inputs.cuda-version }}-${{ inputs.host-platform }}.tar.gz" >> $GITHUB_ENV
+        HASH=$(echo -n "${{ inputs.cuda-components }}" | sha256sum | awk '{print $1}')
+        echo "CTK_CACHE_KEY=mini-ctk-${{ inputs.cuda-version }}-${{ inputs.host-platform }}-$HASH" >> $GITHUB_ENV
+        echo "CTK_CACHE_FILENAME=mini-ctk-${{ inputs.cuda-version }}-${{ inputs.host-platform }}-$HASH.tar.gz" >> $GITHUB_ENV
+        echo "CTK_CACHE_COMPONENTS=${{ inputs.cuda-components }}" >> $GITHUB_ENV
 
     - name: Install dependencies
       uses: ./.github/actions/install_unix_deps
@@ -78,18 +85,16 @@ runs:
           rm $CTK_COMPONENT_COMPONENT_FILENAME
         }
 
-        # Get headers and shared libraries in place
-        # Note: the existing artifact would need to be manually deleted (ex: through web UI)
-        # if this list is changed, as the artifact actions do not offer any option for us to
-        # invalidate the artifact.
-        populate_cuda_path cuda_nvcc
-        populate_cuda_path cuda_cudart
-        populate_cuda_path cuda_nvrtc
-        populate_cuda_path cuda_profiler_api
-        populate_cuda_path cuda_cccl
-        if [[ "$(cut -d '.' -f 1 <<< ${{ inputs.cuda-version }})" -ge 12 ]]; then
-          populate_cuda_path libnvjitlink
+        # Conditionally strip out libnvjitlink for CUDA versions < 12
+        if [[ "$(cut -d '.' -f 1 <<< ${{ inputs.cuda-version }})" -lt 12 ]]; then
+          CTK_CACHE_COMPONENTS="${CTK_CACHE_COMPONENTS//libnvjitlink/}"
         fi
+        # Cleanup stray commas after removing components
+        CTK_CACHE_COMPONENTS="${CTK_CACHE_COMPONENTS//,,/,}"
+        # Get headers and shared libraries in place
+        for item in $(echo $CTK_CACHE_COMPONENTS | tr ',' ' '); do
+            populate_cuda_path "$item"
+        done
         ls -l $CUDA_PATH
 
         # Prepare the cache

.github/workflows/bandit.yml

@@ -0,0 +1,19 @@
+name: "Static Analysis: Bandit Scan"
+
+on:
+  push:
+    branches:
+      - "pull-request/[0-9]+"
+      - "main"
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}-${{ github.event_name }}
+  cancel-in-progress: true
+
+jobs:
+  analyze:
+    runs-on: ubuntu-latest
+    permissions:
+      security-events: write
+    steps:
+      - name: Perform Bandit Analysis
+        uses: PyCQA/bandit-action@8a1b30610f61f3f792fe7556e888c9d7dffa52de  # v1.0.0

.github/workflows/codeql.yml

@@ -0,0 +1,39 @@
+name: "Static Analysis: CodeQL Scan"
+
+on:
+  push:
+    branches:
+      - "pull-request/[0-9]+"
+      - "main"
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}-${{ github.event_name }}
+  cancel-in-progress: true
+
+jobs:
+  analyze:
+    name: Analyze (${{ matrix.language }})
+    runs-on: ubuntu-latest
+    permissions:
+      security-events: write
+
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+        - language: python
+          build-mode: none
+    steps:
+    - name: Checkout repository
+      uses: actions/checkout@v4
+
+    - name: Initialize CodeQL
+      uses: github/codeql-action/init@v3
+      with:
+        languages: ${{ matrix.language }}
+        build-mode: ${{ matrix.build-mode }}
+        queries: security-extended
+
+    - name: Perform CodeQL Analysis
+      uses: github/codeql-action/analyze@v3
+      with:
+        category: "/language:${{matrix.language}}"

.pre-commit-config.yaml

@@ -6,16 +6,23 @@ ci:
     autoupdate_branch: ''
     autoupdate_commit_msg: '[pre-commit.ci] pre-commit autoupdate'
     autoupdate_schedule: quarterly
-    skip: []
+    skip: [bandit]
     submodules: false
 
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.11.4
+    rev: 971923581912ef60a6b70dbf0c3e9a39563c9d47  #v0.11.4
     hooks:
       - id: ruff
         args: [--fix, --show-fixes]
       - id: ruff-format
+  - repo: https://github.com/PyCQA/bandit
+    rev: 8ff25e07e487f143571cc305e56dd0253c60bc7b  #v1.8.3
+    hooks:
+      - id: bandit
+        args:
+          - --ini
+          - .bandit
 
 default_language_version:
       python: python3

cuda_bindings/cuda/__init__.pxd

@@ -60,36 +60,37 @@ cdef int cuPythonInit() except -1 nogil:
         except:
             handle = None
 
-        # Else try default search
-        if not handle:
-            LOAD_LIBRARY_SAFE_CURRENT_DIRS = 0x00002000
-            try:
-                handle = win32api.LoadLibraryEx("nvrtc64_120_0.dll", 0, LOAD_LIBRARY_SAFE_CURRENT_DIRS)
-            except:
-                pass
-
-        # Final check if DLLs can be found within pip installations
+        # Check if DLLs can be found within pip installations
         if not handle:
+            LOAD_LIBRARY_SEARCH_DEFAULT_DIRS = 0x00001000
+            LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR = 0x00000100
             site_packages = [site.getusersitepackages()] + site.getsitepackages()
             for sp in site_packages:
                 mod_path = os.path.join(sp, "nvidia", "cuda_nvrtc", "bin")
-                if not os.path.isdir(mod_path):
-                    continue
-                os.add_dll_directory(mod_path)
-            LOAD_LIBRARY_SEARCH_DEFAULT_DIRS = 0x00001000
-            LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR = 0x00000100
-            try:
-                handle = win32api.LoadLibraryEx(
-                    # Note: LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR needs an abs path...
-                    os.path.join(mod_path, "nvrtc64_120_0.dll"),
-                    0, LOAD_LIBRARY_SEARCH_DEFAULT_DIRS | LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR)
-
-                # Note: nvrtc64_120_0.dll calls into nvrtc-builtins64_*.dll which is
-                # located in the same mod_path.
-                # Update PATH environ so that the two dlls can find each other
-                os.environ["PATH"] = os.pathsep.join((os.environ.get("PATH", ""), mod_path))
-            except:
-                pass
+                if os.path.isdir(mod_path):
+                    os.add_dll_directory(mod_path)
+                    try:
+                        handle = win32api.LoadLibraryEx(
+                            # Note: LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR needs an abs path...
+                            os.path.join(mod_path, "nvrtc64_120_0.dll"),
+                            0, LOAD_LIBRARY_SEARCH_DEFAULT_DIRS | LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR)
+
+                        # Note: nvrtc64_120_0.dll calls into nvrtc-builtins64_*.dll which is
+                        # located in the same mod_path.
+                        # Update PATH environ so that the two dlls can find each other
+                        os.environ["PATH"] = os.pathsep.join((os.environ.get("PATH", ""), mod_path))
+                    except:
+                        pass
+                    else:
+                        break
+            else:
+                # Else try default search
+                # Only reached if DLL wasn't found in any site-package path
+                LOAD_LIBRARY_SAFE_CURRENT_DIRS = 0x00002000
+                try:
+                    handle = win32api.LoadLibraryEx("nvrtc64_120_0.dll", 0, LOAD_LIBRARY_SAFE_CURRENT_DIRS)
+                except:
+                    pass
 
         if not handle:
             raise RuntimeError('Failed to LoadLibraryEx nvrtc64_120_0.dll')

cuda_bindings/cuda/__init__.py

@@ -56,40 +56,30 @@ cdef load_library(const int driver_ver):
 
         # First check if the DLL has been loaded by 3rd parties
         try:
-            handle = win32api.GetModuleHandle(dll_name)
+            return win32api.GetModuleHandle(dll_name)
         except:
             pass
-        else:
-            break
 
         # Next, check if DLLs are installed via pip
         for sp in get_site_packages():
             mod_path = os.path.join(sp, "nvidia", "nvJitLink", "bin")
-            if not os.path.isdir(mod_path):
-                continue
-            os.add_dll_directory(mod_path)
-        try:
-            handle = win32api.LoadLibraryEx(
-                # Note: LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR needs an abs path...
-                os.path.join(mod_path, dll_name),
-                0, LOAD_LIBRARY_SEARCH_DEFAULT_DIRS | LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR)
-        except:
-            pass
-        else:
-            break
-
+            if os.path.isdir(mod_path):
+                os.add_dll_directory(mod_path)
+                try:
+                    return win32api.LoadLibraryEx(
+                        # Note: LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR needs an abs path...
+                        os.path.join(mod_path, dll_name),
+                        0, LOAD_LIBRARY_SEARCH_DEFAULT_DIRS | LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR)
+                except:
+                    pass
         # Finally, try default search
+        # Only reached if DLL wasn't found in any site-package path
         try:
-            handle = win32api.LoadLibrary(dll_name)
+            return win32api.LoadLibrary(dll_name)
         except:
             pass
-        else:
-            break
-    else:
-        raise RuntimeError('Failed to load nvJitLink')
 
-    assert handle != 0
-    return handle
+    raise RuntimeError('Failed to load nvJitLink')
 
 
 cdef int _check_or_init_nvjitlink() except -1 nogil:

cuda_bindings/cuda/bindings/_bindings/cynvrtc.pyx.in

@@ -41,7 +41,7 @@ cdef void* __nvvmGetProgramLog = NULL
 
 
 cdef inline list get_site_packages():
-    return [site.getusersitepackages()] + site.getsitepackages()
+    return [site.getusersitepackages()] + site.getsitepackages() + ["conda"]
 
 
 cdef load_library(const int driver_ver):
@@ -50,44 +50,42 @@ cdef load_library(const int driver_ver):
     for suffix in get_nvvm_dso_version_suffix(driver_ver):
         if len(suffix) == 0:
             continue
-        dll_name = "nvvm64_40_0"
+        dll_name = "nvvm64_40_0.dll"
 
         # First check if the DLL has been loaded by 3rd parties
         try:
-            handle = win32api.GetModuleHandle(dll_name)
+            return win32api.GetModuleHandle(dll_name)
         except:
             pass
-        else:
-            break
 
-        # Next, check if DLLs are installed via pip
+        # Next, check if DLLs are installed via pip or conda
         for sp in get_site_packages():
-            mod_path = os.path.join(sp, "nvidia", "cuda_nvcc", "nvvm", "bin")
-            if not os.path.isdir(mod_path):
-                continue
-            os.add_dll_directory(mod_path)
-        try:
-            handle = win32api.LoadLibraryEx(
-                # Note: LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR needs an abs path...
-                os.path.join(mod_path, dll_name),
-                0, LOAD_LIBRARY_SEARCH_DEFAULT_DIRS | LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR)
-        except:
-            pass
-        else:
-            break
+            if sp == "conda":
+                # nvvm is not under $CONDA_PREFIX/lib, so it's not in the default search path
+                conda_prefix = os.environ.get("CONDA_PREFIX")
+                if conda_prefix is None:
+                    continue
+                mod_path = os.path.join(conda_prefix, "Library", "nvvm", "bin")
+            else:
+                mod_path = os.path.join(sp, "nvidia", "cuda_nvcc", "nvvm", "bin")
+            if os.path.isdir(mod_path):
+                os.add_dll_directory(mod_path)
+                try:
+                    return win32api.LoadLibraryEx(
+                        # Note: LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR needs an abs path...
+                        os.path.join(mod_path, dll_name),
+                        0, LOAD_LIBRARY_SEARCH_DEFAULT_DIRS | LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR)
+                except:
+                    pass
 
         # Finally, try default search
+        # Only reached if DLL wasn't found in any site-package path
         try:
-            handle = win32api.LoadLibrary(dll_name)
+            return win32api.LoadLibrary(dll_name)
         except:
             pass
-        else:
-            break
-    else:
-        raise RuntimeError('Failed to load nvvm')
 
-    assert handle != 0
-    return handle
+    raise RuntimeError('Failed to load nvvm')
 
 
 cdef int _check_or_init_nvvm() except -1 nogil: