[BUG] CUDA ERROR: initialization error in experiment grid

[Gaiejj]

Required prerequisites

• I have read the documentation https://omnisafe.readthedocs.io.
• I have searched the Issue Tracker and Discussions that this hasn't already been reported. (+1 or comment there if it has.)
• Consider asking first in a Discussion.

What version of OmniSafe are you using?

0.5.0

System information

3.8.19 (default, Mar 20 2024, 19:58:24)
[GCC 11.2.0] linux
0.5.0

Problem description

We have found bugs in using experiment_grid for parallel experiments on certain torch versions or CUDA driver versions, specifically: CUDA ERROR: initialization error. We believe this is a bug in CUDA during the parallel training process. Through this error message prompt and related community issues, we have tentatively set the solution as:

1. look at omnisafe/common/experiment_grid.py
2. import multiprocessing as mp
3. modify pool = Pool(max_workers=num_pool) to pool = Pool(max_workers=num_pool, mp_context=mp.get_context('spawn')) in line 445.

If you encounter similar problems during the use of the OmniSafe experiment grid, you can refer to this solution. If you have a better solution, you are welcome to leave a message under this issue!

Reference:

• https://stackoverflow.com/questions/33748750/cuda-error-initialization-error-when-using-parallel-in-python
• https://pytorch.org/docs/stable/notes/multiprocessing.html#cuda-in-multiprocessing

Reproducible example code

Command lines:

python run_experiment_grid.py

Traceback

No response

Expected behavior

No response

Additional context

No response