Developer setup error for `gpu-tests`

[kshitijgoel007]

➜  workspaces git clone --recurse-submodules git@github.com:learnsyslab/crazyflow.git

Cloning into 'crazyflow'...
remote: Enumerating objects: 3020, done.
remote: Counting objects: 100% (596/596), done.
remote: Compressing objects: 100% (176/176), done.
remote: Total 3020 (delta 481), reused 436 (delta 418), pack-reused 2424 (from 2)
Receiving objects: 100% (3020/3020), 14.37 MiB | 17.62 MiB/s, done.
Resolving deltas: 100% (1918/1918), done.
Submodule 'submodules/drone-controllers' (https://github.com/learnsyslab/drone-controllers.git) registered for path 'submodules/drone-controllers'
Submodule 'submodules/drone-models' (https://github.com/learnsyslab/drone-models.git) registered for path 'submodules/drone-models'
Cloning into '/home/kshitijgoel/workspaces/crazyflow/submodules/drone-controllers'...
remote: Enumerating objects: 292, done.        
remote: Counting objects: 100% (62/62), done.        
remote: Compressing objects: 100% (51/51), done.        
remote: Total 292 (delta 10), reused 43 (delta 9), pack-reused 230 (from 1)        
Receiving objects: 100% (292/292), 257.19 KiB | 8.57 MiB/s, done.
Resolving deltas: 100% (133/133), done.
Cloning into '/home/kshitijgoel/workspaces/crazyflow/submodules/drone-models'...
remote: Enumerating objects: 1804, done.        
remote: Counting objects: 100% (450/450), done.        
remote: Compressing objects: 100% (196/196), done.        
remote: Total 1804 (delta 307), reused 361 (delta 247), pack-reused 1354 (from 2)        
Receiving objects: 100% (1804/1804), 6.21 MiB | 23.56 MiB/s, done.
Resolving deltas: 100% (1051/1051), done.
Submodule path 'submodules/drone-controllers': checked out 'a5f39c9a6b3efb75420f146cadf270d415cc1ae2'
Submodule path 'submodules/drone-models': checked out '51076d84485e4935013cbfcb475058f8f3ec461b'
➜  workspaces cd crazyflow 
➜  crazyflow git:(main) ls
benchmark  CITATION.cff  crazyflow  docs  examples  LICENSE  pixi.lock  properdocs.yml  pyproject.toml  README.md  submodules  tests
➜  crazyflow git:(main) pixi shell -e gpu-tests

(crazyflow:gpu-tests) ➜  crazyflow git:(main) python -c "from crazyflow.sim import Sim; sim = Sim(); sim.reset(); print('OK')"
Failed to import warp: No module named 'warp'
Failed to import mujoco_warp: No module named 'warp'
Jax plugin configuration error: Exception when calling jax_plugins.xla_cuda12.initialize()
Traceback (most recent call last):
  File "/home/kshitijgoel/workspaces/crazyflow/.pixi/envs/gpu-tests/lib/python3.12/site-packages/jax_plugins/xla_cuda12/__init__.py", line 201, in _version_check
    version = get_version()
              ^^^^^^^^^^^^^
RuntimeError: jaxlib/cuda/versions_helpers.cc:85: operation cusparseGetProperty(MAJOR_VERSION, &major) failed: The cuSPARSE library was not found.

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/kshitijgoel/workspaces/crazyflow/.pixi/envs/gpu-tests/lib/python3.12/site-packages/jax/_src/xla_bridge.py", line 495, in discover_pjrt_plugins
    plugin_module.initialize()
  File "/home/kshitijgoel/workspaces/crazyflow/.pixi/envs/gpu-tests/lib/python3.12/site-packages/jax_plugins/xla_cuda12/__init__.py", line 370, in initialize
    _check_cuda_versions(raise_on_first_error = True)
  File "/home/kshitijgoel/workspaces/crazyflow/.pixi/envs/gpu-tests/lib/python3.12/site-packages/jax_plugins/xla_cuda12/__init__.py", line 266, in _check_cuda_versions
    _version_check("cuSPARSE", cuda_versions.cusparse_get_version,
  File "/home/kshitijgoel/workspaces/crazyflow/.pixi/envs/gpu-tests/lib/python3.12/site-packages/jax_plugins/xla_cuda12/__init__.py", line 205, in _version_check
    raise RuntimeError(err_msg) from e
RuntimeError: Unable to load cuSPARSE. Is it installed?
An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.
OK
(crazyflow:gpu-tests) ➜  crazyflow git:(main)  

[amacati]

Hi, can you tell me more about your machine? I assume you do have an NVIDIA GPU? Is regular JAX working with GPU arrays for you?

[kshitijgoel007]

I am on an Ubuntu 22.04 machine.

I assume you do have an NVIDIA GPU?

Yes, I have a NVIDIA RTX 5090 currently.

Is regular JAX working with GPU arrays for you?

It was not working, thanks for pointing it out. It turns out I had a discrepancy with my global LD_LIBRARY_PATH setting (which was pointing to cuSPARSE from CUDA 12.8). I have since fixed this and brought everything to CUDA 13.

After the change above, I obtained the following in a virtual environment (just using pip install -U "jax[cuda13]")

(.venv) ➜  crazyflow python -c "import jax; print('Backend:', jax.default_backend()); print('Devices:', jax.devices())"
Backend: gpu
Devices: [CudaDevice(id=0)]

After that, I removed the crazyflow directory and started fresh with the instructions

git clone --recurse-submodules git@github.com:learnsyslab/crazyflow.git
cd crazyflow
pixi shell -e gpu-tests

Then I ran the test command and got OK output

(crazyflow:gpu-tests) ➜  crazyflow git:(main) ✗ python -c "from crazyflow.sim import Sim; sim = Sim(device='gpu'); sim.reset(); print('OK')"
Failed to import warp: No module named 'warp'
Failed to import mujoco_warp: No module named 'warp'
OK

However, there is a failure due to warp. This is not noted in the instructions. Is this expected behavior?

[amacati]

This is expected, yes. We use MJX for collision checking, which by default checks if warp is installed and prints a warning otherwise. So this is not an error, just an info that you cannot switch to the warp implementation, which we don't use anyways, and can be disregarded.

[kshitijgoel007]

Thanks for the prompt response. Nicely written simulator! Closing this issue as completed.