From b4bcb9b60a0861753ebf78c90236807a2280fd70 Mon Sep 17 00:00:00 2001
From: Ian Rodney <ian.rodney@gmail.com>
Date: Wed, 27 Jan 2021 13:45:30 -0800
Subject: [PATCH] [Docker] Use Cuda 11  (#13691)

---
 build-docker.sh                                     | 2 +-
 ci/travis/build-docker-images.py                    | 2 +-
 ci/travis/build-docker-images.sh                    | 2 +-
 docker/base-deps/Dockerfile                         | 2 +-
 python/requirements_ml_docker.txt                   | 7 +++++--
 release/rllib_tests/unit_gpu_tests/requirements.txt | 6 ++++--
 6 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/build-docker.sh b/build-docker.sh
index 3a09b4896..b39336186 100755
--- a/build-docker.sh
+++ b/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"
diff --git a/ci/travis/build-docker-images.py b/ci/travis/build-docker-images.py
index a2ae7a18d..c549bc95e 100644
--- a/ci/travis/build-docker-images.py
+++ b/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
diff --git a/ci/travis/build-docker-images.sh b/ci/travis/build-docker-images.sh
index c894da23a..6463c880f 100755
--- a/ci/travis/build-docker-images.sh
+++ b/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"
diff --git a/docker/base-deps/Dockerfile b/docker/base-deps/Dockerfile
index a5bcfedbf..3aec50c99 100644
--- a/docker/base-deps/Dockerfile
+++ b/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
diff --git a/python/requirements_ml_docker.txt b/python/requirements_ml_docker.txt
index 6f610c468..c61ba0c05 100644
--- a/python/requirements_ml_docker.txt
+++ b/python/requirements_ml_docker.txt
@@ -1,3 +1,6 @@
 ipython
-tensorflow-gpu
-torch
\ No newline at end of file
+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 
\ No newline at end of file
diff --git a/release/rllib_tests/unit_gpu_tests/requirements.txt b/release/rllib_tests/unit_gpu_tests/requirements.txt
index 4f8897539..b8a991f74 100644
--- a/release/rllib_tests/unit_gpu_tests/requirements.txt
+++ b/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