Numba: no cuda support detected?

[arunoruto]

Describe the bug

I am currently trying to work with Numba on a system with an Nvidia GPU. For this, I am trying to use the numbaWithNumba package found in nixpkgs, but is_available() gives me False all the time and numba -s returns an incomplete setup status for CUDA.

Steps To Reproduce

1. Install numbaWithNumba in a shell via nix-shell -p python3Packages.numbaWithNumba
2. Run python -c "import numba.cuda; print(numba.cuda.is_available())" -> results in False
3. Run numba -s ->

System info:
/nix/store/c1rgbgz8d826xbyqcz6k68vvwkxbkh0m-python3-3.12.7-env/bin/python3.12: No module named pip
--------------------------------------------------------------------------------
__Time Stamp__
Report started (local time)                   : 2024-12-14 23:58:05.694113
UTC start time                                : 2024-12-14 22:58:05.694121
Running time (s)                              : 0.03438

__Hardware Information__
Machine                                       : x86_64
CPU Name                                      : skylake
CPU Count                                     : 16
Number of accessible CPUs                     : 16
List of accessible CPUs cores                 : 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
CFS Restrictions (CPUs worth of runtime)      : None

CPU Features                                  : 64bit adx aes avx avx2 bmi bmi2
                                                clflushopt cmov crc32 cx16 cx8
                                                f16c fma fsgsbase fxsr invpcid
                                                lzcnt mmx movbe pclmul popcnt
                                                prfchw rdrnd rdseed rtm sahf sgx
                                                sse sse2 sse3 sse4.1 sse4.2 ssse3
                                                xsave xsavec xsaveopt xsaves

Memory Total (MB)                             : 64133
Memory Available (MB)                         : 62099

__OS Information__
Platform Name                                 : Linux-6.6.63-x86_64-with-glibc2.40
Platform Release                              : 6.6.63
OS Name                                       : Linux
OS Version                                    : #1-NixOS SMP PREEMPT_DYNAMIC Fri Nov 22 14:38:37 UTC 2024
OS Specific Version                           : ?
Libc Version                                  : glibc 2.40

__Python Information__
Python Compiler                               : GCC 13.3.0
Python Implementation                         : CPython
Python Version                                : 3.12.7
Python Locale                                 : en_GB.UTF-8

__Numba Toolchain Versions__
Numba Version                                 : 0.61.0
llvmlite Version                              : 0.43.0

__LLVM Information__
LLVM Version                                  : 14.0.6

__CUDA Information__
CUDA Device Initialized                       : True
CUDA Driver Version                           : ?
CUDA Runtime Version                          : ?
CUDA NVIDIA Bindings Available                : ?
CUDA NVIDIA Bindings In Use                   : ?
CUDA Minor Version Compatibility Available    : ?
CUDA Minor Version Compatibility Needed       : ?
CUDA Minor Version Compatibility In Use       : ?
CUDA Detect Output:
Found 1 CUDA devices
id 0    b'NVIDIA TITAN X (Pascal)'                              [SUPPORTED]
                      Compute Capability: 6.1
                           PCI Device ID: 0
                              PCI Bus ID: 1
                                    UUID: GPU-b08e26d8-65ac-1443-3a8c-1fc9d3bf98a7
                                Watchdog: Disabled
             FP32/FP64 Performance Ratio: 32
Summary:
        1/1 devices are supported

CUDA Libraries Test Output:
None

__NumPy Information__
NumPy Version                                 : 1.26.4
NumPy Supported SIMD features                 : ('MMX', 'SSE', 'SSE2', 'SSE3', 'SSSE3', 'SSE41', 'POPCNT', 'SSE42', 'AVX', 'F16C', 'FMA3', 'AVX2')
NumPy Supported SIMD dispatch                 : ('SSSE3', 'SSE41', 'POPCNT', 'SSE42', 'AVX', 'F16C', 'FMA3', 'AVX2', 'AVX512F', 'AVX512CD', 'AVX512_KNL', 'AVX512_KNM', 'AVX512_SKX', 'AVX512_CLX', 'AVX512_CNL', 'AVX512_ICL', 'AVX512_SPR')
NumPy Supported SIMD baseline                 : ('SSE', 'SSE2', 'SSE3')
NumPy AVX512_SKX support detected             : False

__SVML Information__
SVML State, config.USING_SVML                 : False
SVML Library Loaded                           : False
llvmlite Using SVML Patched LLVM              : False
SVML Operational                              : False

__Threading Layer Information__
TBB Threading Layer Available                 : False
+--> Disabled due to Unknown import problem.
OpenMP Threading Layer Available              : True
+-->Vendor: GNU
Workqueue Threading Layer Available           : True
+-->Workqueue imported successfully.

__Numba Environment Variable Information__
None found.

__Conda Information__
Conda not available.

__Installed Packages__
Couldn't retrieve packages info.

__Error log__
Error (pip): Command '('/nix/store/c1rgbgz8d826xbyqcz6k68vvwkxbkh0m-python3-3.12.7-env/bin/python3.12', '-m', 'pip', 'list')' returned non-zero exit status 1.

__Warning log__
Warning (cuda): Probing CUDA failed (device and driver present, runtime problem?)
(cuda) <class 'OSError'>: libcudart.so: cannot open shared object file: No such file or directory
Warning: Conda not available.
 Error was [Errno 2] No such file or directory: 'conda'

Warning (no file): /sys/fs/cgroup/cpuacct/cpu.cfs_quota_us
Warning (no file): /sys/fs/cgroup/cpuacct/cpu.cfs_period_us
--------------------------------------------------------------------------------
If requested, please copy and paste the information between
the dashed (----) lines, or from a given specific section as
appropriate.

=============================================================
IMPORTANT: Please ensure that you are happy with sharing the
contents of the information present, any information that you
wish to keep private you should remove before sharing.
=============================================================

On the contrary, using jaxWithCuda works fine. I can import the lib, execute code and get results back.

Expected behavior

1. The python command should return True
2. Running numba -s should result in
   • no questions marks
   • libcudart.so should be found

Additional context

This happens on both unstable and 24.11. I also recently started using the cache from the cuda-team, but I had the problem before I added their substituter.

Metadata

• system: "x86_64-linux"
• host os: Linux 6.12.3, NixOS, 24.11 (Vicuna), 24.11.20241210.a0f3e10
• multi-user?: yes
• sandbox: yes
• version: nix-env (Nix) 2.24.10
• nixpkgs: /nix/store/va0p2i72cm2ljwm084a0g6ji41s5qnyz-source

Notify maintainers

@mweinelt I know you are not listed as a maintainer, but you added a few cuda related things to numba, so maybe you can help me/clarify things a bit.

---

Note for maintainers: Please tag this issue in your PR.

---

Add a 👍 reaction to issues you find important.

[mweinelt]

Defering to the @NixOS/cuda-maintainers.

[arunoruto]

I did some further troubleshooting and I am maybe onto something...
There was a similar issue on numbas issue tracker, and it was recommended to run python -c "from numba import cuda; cuda.cudadrv.libs.test()". The output on my NixOS system is:

Finding driver from candidates:
        libcuda.so
        libcuda.so.1
        /run/opengl-driver/lib/libcuda.so
        /run/opengl-driver/lib/libcuda.so.1
        /usr/lib/libcuda.so
        /usr/lib/libcuda.so.1
        /usr/lib64/libcuda.so
        /usr/lib64/libcuda.so.1
Using loader <class 'ctypes.CDLL'>
        Trying to load driver...        ok
                Loaded from libcuda.so.1
        Mapped libcuda.so paths:
                /nix/store/9xyylgpsldqfbx7abjyvz42amharnm3d-nvidia-x11-565.77-6.6.64/lib/libcuda.so.565.77
Finding nvvm from Nix store
        Located at libnvvm.so
        Trying to open library...       ERROR: failed to open nvvm:
libnvvm.so: cannot open shared object file: No such file or directory
Finding nvrtc from Nix store
        Located at libnvrtc.so
        Trying to open library...       ERROR: failed to open nvrtc:
libnvrtc.so: cannot open shared object file: No such file or directory
Finding cudart from Nix store
        Located at libcudart.so
        Trying to open library...       ERROR: failed to open cudart:
libcudart.so: cannot open shared object file: No such file or directory
Finding cudadevrt from <unknown>
        Located at libcudadevrt.a
        Checking library...     ERROR: failed to find cudadevrt:
libcudadevrt.a not found
Finding libdevice from Nix store
        Located at /nix/store/3v3r8yyfrpjkwb3mzpwvz77ybpsp7wgp-cuda_nvcc-12.4.99/nvvm/libdevice/libdevice.10.bc
        Checking library...     ok

If I run the same command on Ubuntu, where everything is detected correctly, I get the following output:

Finding driver from candidates:
        libcuda.so
        libcuda.so.1
        /usr/lib/libcuda.so
        /usr/lib/libcuda.so.1
        /usr/lib64/libcuda.so
        /usr/lib64/libcuda.so.1
Using loader <class 'ctypes.CDLL'>
        Trying to load driver...        ok
                Loaded from libcuda.so
        Mapped libcuda.so paths:
                /usr/lib/x86_64-linux-gnu/libcuda.so.550.127.05
Finding nvvm from System
        Located at /usr/local/cuda/nvvm/lib64/libnvvm.so.4.0.0
        Trying to open library...       ok
Finding nvrtc from System
        Located at /usr/local/cuda/lib64/libnvrtc.so.12.6.85
        Trying to open library...       ok
Finding cudart from System
        Located at /usr/local/cuda/lib64/libcudart.so.12.6.77
        Trying to open library...       ok
Finding cudadevrt from System
        Located at /usr/local/cuda/lib64/libcudadevrt.a
        Checking library...     ok
Finding libdevice from System
        Located at /usr/local/cuda/nvvm/libdevice/libdevice.10.bc
        Checking library...     ok

It seems like libraries found in /usr/local/cuda/lib64 on ubuntu are not found on NixOS, but they are present on my system, but not in the lib64 folder, but in lib. I guess numba is searching in the wrong folder now?
Maybe I could patch cudatoolkit to make a symlink of lib to lib64.

EDIT: I just did a small test, and it seems like I am right. If I add the $CUDA_HOME/lib folder to LD_LIBRARY_PATH it detects all libraries just fine! (except the static one, since it is in its own folder)

[ConnorBaker]

From the patch file we have vendored, it looks like Numba expects everything to exist in a single directory. If that's the case, that won't work with our packages which use splayed outputs.

To work around that, I suggest using symlinkJoin to create a derivation with all of the relevant CUDA libraries in a single output. From what I can see from a quick look at the runtime failure, that would be something like

symlinkJoin {
  name = "cudatoolkit";
  paths = with cudaPackages; [ cuda_nvcc cuda_nvrtc cuda_cudart ];
}

I may be forgetting a library :l

Give that a try and let me know how that works.

EDIT: Depending on how symlinkJoin is implemented, you may need to include the result of both getDev and getLib on each of those (and probably getBin on cuda_nvcc as well).

[arunoruto]

@ConnorBaker Thanks for the tip! I guess I got it running now to a point I can work with!

I experimented with getDev and getLib and got this:

nix-repl> "${lib.getDev pkgs.cudaPackages.cuda_nvcc}"
"/nix/store/3v3r8yyfrpjkwb3mzpwvz77ybpsp7wgp-cuda_nvcc-12.4.99"

nix-repl> "${lib.getLib pkgs.cudaPackages.cuda_nvcc}"
"/nix/store/3v3r8yyfrpjkwb3mzpwvz77ybpsp7wgp-cuda_nvcc-12.4.99"

nix-repl> "${lib.getBin pkgs.cudaPackages.cuda_nvcc}"
"/nix/store/3v3r8yyfrpjkwb3mzpwvz77ybpsp7wgp-cuda_nvcc-12.4.99"

nix-repl> "${lib.getExe pkgs.cudaPackages.cuda_nvcc}"
"/nix/store/3v3r8yyfrpjkwb3mzpwvz77ybpsp7wgp-cuda_nvcc-12.4.99/bin/nvcc"

nix-repl> "${lib.getDev pkgs.cudaPackages.cuda_nvrtc}"
"/nix/store/yyag2ikqh0nvbjw193lar506q41vhknl-cuda_nvrtc-12.4.99-dev"

nix-repl> "${lib.getLib pkgs.cudaPackages.cuda_nvrtc}"
"/nix/store/4q7rxvd6dwnkrsyiif4f75xh1l6xcjz5-cuda_nvrtc-12.4.99-lib"

nix-repl> "${lib.getDev pkgs.cudaPackages.cuda_cudart}"
"/nix/store/c9ssqifwhvxnrk2s4j1mxkmsw1jwnd5f-cuda_cudart-12.4.99-dev"

nix-repl> "${lib.getLib pkgs.cudaPackages.cuda_cudart}"
"/nix/store/x9lnlkmvfvvph8r723z3vrkv0r2fh1fa-cuda_cudart-12.4.99-lib"

cuda_nvcc can be just included like that, no need to differentiate, while nvrtc and cudart could need both dev and lib. Nevertheless, I didn't really need all of them.
What really made the trick work, is the postBuild hook of symlinkJoin:

cuda-symlink = pkgs.symlinkJoin {
  name = "cudatoolkit";
  paths = with pkgs.cudaPackages; [
    cudatoolkit
    # cuda_nvcc
    # (lib.getLib cuda_nvrtc)
    # (lib.getDev cuda_nvrtc)
    # (lib.getLib cuda_cudart)
    # (lib.getDev cuda_cudart)
    (lib.getStatic cuda_cudart)
  ];
  postBuild = ''
    ln -s $out/lib $out/lib64
  '';
};

Seems like only cudatoolkit and the static library of cuda_cudart are needed and in the hook I made a symlink of the lib64 folder. Since I am using a flake, I just exported a CUDA_HOME variable pointing to my cuda-symlink derivation and now numba -s displays everything nicely!

How should be proceed on this topic? Maybe "fix" this upstream too?

[34j]

Minimal flake for Numba with CUDA:

./flake.nix:

{
  description = "A flake providing a dev shell for Numba with CUDA without installing Numba via nix. Also supports PyTorch yet being minimal for Numba with CUDA.";

  inputs = {
    nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable";
  };

  outputs =
    { self, nixpkgs }:
    let
      system = "x86_64-linux"; # Adjust if needed
      pkgs = import nixpkgs {
        system = system;
        config.allowUnfree = true;
      };
      cudatookit-with-cudart-to-lib64 = pkgs.symlinkJoin {
        name = "cudatoolkit";
        paths = with pkgs.cudaPackages; [
          cudatoolkit
          (pkgs.lib.getStatic cuda_cudart)
        ];
        postBuild = ''
          ln -s $out/lib $out/lib64
        '';
      };
    in
    {
      devShells.${system}.default = pkgs.mkShell {
        shellHook = ''
          # Required for both PyTorch and Numba to find CUDA
          export CUDA_PATH=${cudatookit-with-cudart-to-lib64}

          # Required for both PyTorch and Numba, adds necessary paths for dynamic linking
          export LD_LIBRARY_PATH=${
            pkgs.lib.makeLibraryPath [
              "/run/opengl-driver" # Needed to find libGL.so, required by both PyTorch and Numba
            ]
          }:$LD_LIBRARY_PATH
        '';
      };
    };
}