[Docker] Use Cuda 11 (#13691)

build-docker.sh

@@ -16,7 +16,7 @@ key="$1"
 case $key in
     --gpu)
     GPU="-gpu"
-    BASE_IMAGE="nvidia/cuda:10.1-cudnn7-runtime-ubuntu18.04"
+    BASE_IMAGE="nvidia/cuda:11.0-cudnn8-runtime-ubuntu18.04"
     ;;
     --no-cache-build)
     NO_CACHE="--no-cache"

ci/travis/build-docker-images.py

@@ -84,7 +84,7 @@ def _build_cpu_gpu_images(image_name, no_cache=True) -> List[str]:
         build_args = {}
         if image_name == "base-deps":
             build_args["BASE_IMAGE"] = (
-                "nvidia/cuda:10.1-cudnn7-runtime-ubuntu18.04"
+                "nvidia/cuda:11.0-cudnn8-runtime-ubuntu18.04"
                 if gpu == "-gpu" else "ubuntu:focal")
         else:
             build_args["GPU"] = gpu

ci/travis/build-docker-images.sh

@@ -22,7 +22,7 @@ build_and_push_tags() {
     # $2 tag for image (e.g. hash of commit)
     for GPU in "" "-gpu" 
     do 
-        BASE_IMAGE=$(if [ "$GPU" ]; then echo "nvidia/cuda:10.1-cudnn7-runtime-ubuntu18.04"; else echo "ubuntu:focal"; fi;)
+        BASE_IMAGE=$(if [ "$GPU" ]; then echo "nvidia/cuda:11.0-cudnn8-runtime-ubuntu18.04"; else echo "ubuntu:focal"; fi;)
         FULL_NAME_WITH_TAG="rayproject/$1:$2$GPU"
         NIGHTLY_FULL_NAME_WITH_TAG="rayproject/$1:nightly$GPU"
         docker build --no-cache --build-arg GPU="$GPU" --build-arg BASE_IMAGE="$BASE_IMAGE" --build-arg WHEEL_PATH=".whl/$WHEEL" --label "SHA=$2" -t "$FULL_NAME_WITH_TAG" /"$ROOT_DIR"/docker/"$1"

docker/base-deps/Dockerfile

@@ -1,6 +1,6 @@
 # The base-deps Docker image installs main libraries needed to run Ray
 
-# The GPU option is nvidia/cuda:10.1-cudnn7-runtime-ubuntu18.04
+# The GPU option is nvidia/cuda:11.0-cudnn8-runtime-ubuntu18.04
 ARG BASE_IMAGE="ubuntu:focal"
 FROM ${BASE_IMAGE}
 # If this arg is not "autoscaler" then no autoscaler requirements will be included

python/requirements_ml_docker.txt

@@ -1,3 +1,6 @@
 ipython
-tensorflow-gpu
-torch
+tensorflow-gpu>=2.4.0
+-f https://download.pytorch.org/whl/torch_stable.html
+torch==1.7.1+cu110 
+-f https://download.pytorch.org/whl/torch_stable.html
+torchvision==0.8.2+cu110 

release/rllib_tests/unit_gpu_tests/requirements.txt

@@ -1,7 +1,9 @@
 ray[rllib]
 ray
-torch==1.6+cu101
-torchvision==0.7.0+cu101
+-f https://download.pytorch.org/whl/torch_stable.html
+torch==1.7.1+cu110 
+-f https://download.pytorch.org/whl/torch_stable.html
+torchvision==0.8.2+cu110 
 boto3==1.4.8
 cython==0.29.0
 pytest