fail install package with GPU

[Borda]

hello, not sure what I am doing wrong but I was not able to install this package with Linux environment... https://dev.azure.com/Lightning-AI/Tutorials/_build/results?buildId=209848&view=logs&j=39b3c594-f55d-5bf7-f035-2947d37d58a7&t=76638fff-2d3c-56ee-85c7-984ea34044a4

pip install torchvision rl-warp-drive==2.5 matplotlib 'pytorch-lightning>=2.0, <2.1.0' torchtext 'torchmetrics>=1.0, <1.3' ffmpeg-python 'numpy <2.0' 'torch>=1.8.1, <2.1.0' --quiet --extra-index-url=https://download.pytorch.org/whl/cu118
  error: subprocess-exited-with-error
  
  × Building wheel for pycuda (pyproject.toml) did not run successfully.
  │ exit code: 1
  ╰─> [12328 lines of output]
      ***************************************************************
      *** WARNING: nvcc not in path.
      *** May need to set CUDA_INC_DIR for installation to succeed.
      ***************************************************************
      *************************************************************
      *** I have detected that you have not run configure.py.
      *************************************************************
      *** Additionally, no global config files were found.
      *** I will go ahead with the default configuration.
      *** In all likelihood, this will not work out.
      ***
      *** See README_SETUP.txt for more information.
      ***
      *** If the build does fail, just re-run configure.py with the
      *** correct arguments, and then retry. Good luck!
      *************************************************************
      *** HIT Ctrl-C NOW IF THIS IS NOT WHAT YOU WANT
      *************************************************************
      Continuing in 10 seconds...
      Continuing in 9 seconds...
...

and many logs ending with:

      In file included from src/cpp/cuda.cpp:4:
      src/cpp/cuda.hpp:14:10: fatal error: cuda.h: No such file or directory
         14 | #include <cuda.h>
            |          ^~~~~~~~
      compilation terminated.
      error: command '/usr/bin/x86_64-linux-gnu-gcc' failed with exit code 1
      [end of output]
  
  note: This error originates from a subprocess, and is likely not a problem with pip.
  ERROR: Failed building wheel for pycuda
ERROR: ERROR: Failed to build installable wheels for some pyproject.toml based projects (pycuda)

[Emerald01]

You can configure your CUDA environment by using the Docker similar to the following

FROM nvcr.io/nvidia/pytorch:23.04-py3
LABEL description="warpdrive-env"

WORKDIR /home/
RUN chmod a+rwx /home/

RUN pip install --upgrade pip
RUN pip install pycuda==2022.2.2
RUN pip install numba==0.56.4

Then you can pip install warpdrive, it should work

[Borda]

Do I need these specific pycuda and numba versions?

[Emerald01]

Not necessary. I can also use higher versions but the base image may require higher version too, depending the CUDA drivers etc. The example is the image I usually use. And there is not much difference in functionalities or speed for them.