devenv + cuda + uv + numba

[arunoruto]

This isn't much of an issue, but some results from research on how to enable numbas CUDA support by using devenv and uv under the hood. I used to use uv2nix and a flake config, but this configuration approach made it so much simpler! If there isn't much to add, feel free to close it. Maybe it will help someone in the future or it can be documented if you see the need :)

I am currenlty using numba for a lot of my python projects to speed things up and the CPU JIT compilation works fine under devenv. The problem arises when trying to use the GPU to accelerate a workflow. It is a similar issue like with numpy not finding zlib by default and being needed included in languages.python.libraries. For numba, I found a small fix:

  env = {
    CUDA_HOME = lib.getDev (
      pkgs.symlinkJoin {
        name = "cudatoolkit";
        paths = with pkgs.cudaPackages; [
          cudatoolkit
          (lib.getStatic cuda_cudart)
        ];
        postBuild = ''
          ln -s $out/lib $out/lib64
        '';
      }
    );
    NUMBA_CUDA_DRIVER = "/run/opengl-driver/lib/libcuda.so";
  };

The CUDA_HOME is needed to find the libraries, since it usually tries to find the lib files in lib64, while nix places them in lib. NUMBA_CUDA_DRIVER is also needed and can be enabled on the system with hardware.opengl.enable = true and configured nvidia settings.

[domenkozar]

Hey! Reminds me of https://github.com/cachix/devenv/pull/422, maybe we do need a module?

[arunoruto]

Hey! Reminds me of #422, maybe we do need a module?

I skimed over the PR and I guess a CUDA module would be handy, but it will be complicated to tailor it to everyone's needs. But a small module which would allow setting the CUDA_ variables and install the respective package would be nice. Make it modular, so people can replace individual components with what they need.

I will see if I can work up a PoC, something more simple than the PR, and work from there.

[domenkozar]

Hey! Reminds me of #422, maybe we do need a module?

I skimed over the PR and I guess a CUDA module would be handy, but it will be complicated to tailor it to everyone's needs. But a small module which would allow setting the CUDA_ variables and install the respective package would be nice. Make it modular, so people can replace individual components with what they need.

I will see if I can work up a PoC, something more simple than the PR, and work from there.

That would be great :)