A Simple Setup

[gauenk]

Thank you for the excellent tutorial. You have reduced my dev time by x1000 by providing this excellent minimum working example.

I wanted to offer a simpler setup.py. I am writing a module, and I want the module to talk with both Pytorch and Jax. To this end, I didn't want to use a CMake file since my Pytorch module currently doesn't need one. Instead, I am using the torch setup tools, described here.

I think your CMake file logic is good still since it shows how to do it (I had no idea how). This version just requires less code.

[dfm]

Thanks!! I thought about this approach too, but I didn't know how to feel about having torch as a required build dependency. It would be cool if the CUDAExtension stuff was forked :D

But either way, you'll need to add torch as a build dependency in pyproject.toml to get this to work!

[gauenk]

That's a good point -- pytorch just got added. Right now I have to install pytorch in a weird way (requesting CUDA 11 binaries) so I am omitting it originally.Anyway, I thought sharing this is worth while. And someone can checkout this PR for a shorter setup.

Btw: I agree with you -- I would love it they forked the CUDAExtension stuff too. Maybe I will submit a feature request once my project's pace lets up a bit. Happy hacking!

[dfm]

Interesting! Is the cuda-enabled version of pytorch required for the CUDAExtension stuff to work? I guess it makes sense that it would be...

[gauenk]

Good question. I have actually not tested this. For my life in computer vision, if the code is not CUDA enabled then I have no business using it ;P

[dfm]

Right - but this question does matter in terms of the build process, since this is setup for an isolated build that separately installs all the build dependencies. Even if you're using the CUDA-enabled runtime libraries.

[gauenk]

You are correct -- the question for sure matters. I'll add some detail. I went through the process documenting the procedure properly.

[Detail 1]: The install won't work for me (when GPU is enabled) with the current pyproject.toml because of jax.

1. The standard installation of jax does not recognize my GPU (even when pytorch does). So I install jax with the following lines:

python -m pip install https://storage.googleapis.com/jax-releases/cuda11/jaxlib-0.3.10+cuda11.cudnn805-cp38-none-manylinux2014_x86_64.whl
python -m pip install jax==0.3.10

[Detail 2]: We do need a pytorch install with a matching CUDA version to our system to use the simpler setup.py.

1. I install pytorch like so: pip install torch==1.12.1+cu113 torchvision==0.13.0+cu113 torchaudio==0.11.0 --extra-index-url https://download.pytorch.org/whl/cu113 before running the setup.py. via python -m pip install .

2a. I must link the pytorch lib so we have access to libc10.so. This is completed with the following line:
export LD_LIBRARY_PATH=`python -c "import torch; print(torch.__file__[:-11] + 'lib')"

2b. Alternatively, running import torch before importing any custom module fixes the issue.

[Detail 3]: I actually can't get the install to work with the pyproject.toml at all. I think there are known issues with torch (see here, and here). So I delete pyproject.toml and then run the install: python -m pip install ..

[dfm]

Thanks for these details! In this case, I'm inclined to stick with the currently implemented procedure, but it would be nice to add a few words about this to the README (perhaps including a link to your code when the time comes?) if you have any interested in opening a PR adding some discussion? Either way - thanks again!