Segmentation fault on next(scenarios)

[ixanezis]

Just trying to run the simplest thing:

from waymax import config
from waymax import dataloader

scenarios = dataloader.simulator_state_generator(config.WOD_1_1_0_TRAINING)
scenario = next(scenarios)
print("hello")

It does not reach the latest line, crashes with segmentation fault on scenario = next(scenarios). Is there a way to debug this out?

[gaoyinfeng]

Is your fault similar to this?
#17

[ixanezis]

I guess it is.
However, simply using higher cuda driver version did not help me.
After trying many various configurations, I have come up with a (probably) working one:

FROM us-docker.pkg.dev/colab-images/public/runtime

RUN pip install --upgrade pip
RUN pip install git+https://github.com/waymo-research/waymax.git@main#egg=waymo-waymax

I did not succeed starting from tensorflow/tensorflow, nvidia/cuda:11.8.0-devel-ubuntu22.04, and even pyhf/cuda:0.7.2-jax-cuda-11.8.0-cudnn8 which has passed the next(scenarios) step, but actually broke down on visualization with some "cannot load cublas...so" error.

[gaoyinfeng]

I also found it quite misleading to update CUDA version. I managed to make an env and successfully run demos without error printed after several trials. Here are my steps:

1. conda create -n waymax python=3.11
2. conda activate waymax
3. conda install cudatoolkit=11.8
4. conda install cudnn=8.2
5. install jaxlib=0.4.7 for python3.11 + cuda11.8 + cudnn8.2 from wheel
6. pip install jax=0.4.9 and tensorflow2.13.0
7. manually install all other dependencies from the setup.py file with the lowest version that shows in the file
8. manually reinstall jax=0.4.9 since it may update in step 7
9. pip install jupyter, since you need it to run demos and visualize driving scenarios

I am not sure the CUDA and cudnn are correctly linked because my machine installed cuda 11.8 in the base env.
Hope this can help you!

[iansu81]

I tried the dependencies mentioned above, but it seems like I'm stuck in an infinite loop of changing dependency versions. Has anyone successfully run Waymax in a Docker image?