Error related with "with_cuda“

[Ticotico410]

I followed the instruction in README.md to download the project and created the corresponding virtual environment. However, when I try to run the demo, the error occurs. I use Ubuntu 20.04 + Python 3.10.16 + CUDA 12.8.93 + jax-cuda12-plugin==0.5.3 + jax==0.5.3, but cannot find extra feature 'with_cuda'.

The error is shown as follow:
(dial-mpc) ycb@ycb-Dell-G15-5511:~/ERL-Research/projects/dial-mpc$ dial-mpc --example unitree_h1_jog
Traceback (most recent call last):
File "/home/ycb/anaconda3/envs/dial-mpc/lib/python3.10/site-packages/pkg_resources/init.py", line 3103, in requires
deps.extend(dm[safe_extra(ext)])
KeyError: 'with_cuda'

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

Traceback (most recent call last):
File "/home/ycb/anaconda3/envs/dial-mpc/bin/dial-mpc", line 33, in
sys.exit(load_entry_point('dial-mpc', 'console_scripts', 'dial-mpc')())
File "/home/ycb/anaconda3/envs/dial-mpc/bin/dial-mpc", line 25, in importlib_load_entry_point
return next(matches).load()
File "/home/ycb/anaconda3/envs/dial-mpc/lib/python3.10/importlib/metadata/init.py", line 171, in load
module = import_module(match.group('module'))
File "/home/ycb/anaconda3/envs/dial-mpc/lib/python3.10/importlib/init.py", line 126, in import_module
return _bootstrap._gcd_import(name[level:], package, level)
File "", line 1050, in _gcd_import
File "", line 1027, in _find_and_load
File "", line 1006, in _find_and_load_unlocked
File "", line 688, in _load_unlocked
File "", line 883, in exec_module
File "", line 241, in _call_with_frames_removed
File "/home/ycb/ERL-Research/projects/dial-mpc/dial_mpc/core/dial_core.py", line 17, in
from jax_cosmo.scipy.interpolate import InterpolatedUnivariateSpline
File "/home/ycb/anaconda3/envs/dial-mpc/lib/python3.10/site-packages/jax_cosmo/init.py", line 2, in
from pkg_resources import DistributionNotFound
File "/home/ycb/anaconda3/envs/dial-mpc/lib/python3.10/site-packages/pkg_resources/init.py", line 3657, in
def _initialize_master_working_set() -> None:
File "/home/ycb/anaconda3/envs/dial-mpc/lib/python3.10/site-packages/pkg_resources/init.py", line 3640, in _call_aside
f(*args, **kwargs)
File "/home/ycb/anaconda3/envs/dial-mpc/lib/python3.10/site-packages/pkg_resources/init.py", line 3669, in _initialize_master_working_set
working_set = _declare_state('object', 'working_set', WorkingSet._build_master())
File "/home/ycb/anaconda3/envs/dial-mpc/lib/python3.10/site-packages/pkg_resources/init.py", line 654, in _build_master
ws.require(requires)
File "/home/ycb/anaconda3/envs/dial-mpc/lib/python3.10/site-packages/pkg_resources/init.py", line 1061, in require
needed = self.resolve(parse_requirements(requirements))
File "/home/ycb/anaconda3/envs/dial-mpc/lib/python3.10/site-packages/pkg_resources/init.py", line 893, in resolve
new_requirements = dist.requires(req.extras)[::-1]
File "/home/ycb/anaconda3/envs/dial-mpc/lib/python3.10/site-packages/pkg_resources/init.py", line 3105, in requires
raise UnknownExtra(f"{self} has no such extra feature {ext!r}") from e
pkg_resources.UnknownExtra: jax-cuda12-plugin 0.5.3 has no such extra feature 'with_cuda'

[josephrocca]

Can confirm this is an issue. Replication is simple:

docker run --rm -it --gpus all pytorch/pytorch:2.6.0-cuda12.6-cudnn9-devel bash
git clone https://github.com/LeCar-Lab/dial-mpc.git --depth 1
cd dial-mpc
pip3 install -e .
dial-mpc --list-examples

I solved this by downgrading jax to 0.5.2 - i.e. changed jax[cuda12] to jax[cuda12]==0.5.2 in install_requires in setup.py, and then re-execute pip3 install -e .

[HaoruXue]

Will test ASAP and see if we need to freeze dependency versions

[Ticotico410]

it works, appreciate it!