Skip to content

Commit

Permalink
Simplify GPU detection in docker launch script (#2064)
Browse files Browse the repository at this point in the history
Co-authored-by: Antonin RAFFIN <[email protected]>
  • Loading branch information
sanowl and araffin authored Jan 7, 2025
1 parent 2b529e5 commit b7c64a1
Showing 1 changed file with 4 additions and 11 deletions.
15 changes: 4 additions & 11 deletions scripts/run_docker_gpu.sh
Original file line number Diff line number Diff line change
@@ -1,19 +1,12 @@
#!/bin/bash
# Launch an experiment using the docker gpu image

cmd_line="$@"

echo "Executing in the docker (gpu image):"
echo $cmd_line

# TODO: always use new-style once sufficiently widely used (probably 2021 onwards)
if [ -x "$(which nvidia-docker)" ]; then
# old-style nvidia-docker2
NVIDIA_ARG="--runtime=nvidia"
else
NVIDIA_ARG="--gpus all"
fi
# Using new-style GPU argument
NVIDIA_ARG="--gpus all"

docker run -it ${NVIDIA_ARG} --rm --network host --ipc=host \
--mount src=$(pwd),target=/home/mamba/stable-baselines3,type=bind stablebaselines/stable-baselines3:latest \
bash -c "cd /home/mamba/stable-baselines3/ && $cmd_line"
--mount src=$(pwd),target=/home/mamba/stable-baselines3,type=bind stablebaselines/stable-baselines3:latest \
bash -c "cd /home/mamba/stable-baselines3/ && $cmd_line"

0 comments on commit b7c64a1

Please sign in to comment.