Clean up ray-ml requirements

docker/ray-ml/Dockerfile

@@ -19,20 +19,26 @@ RUN sudo apt-get update \
         libgl1-mesa-dev \
         unzip \
         unrar \
-    && $HOME/anaconda3/bin/pip install -U pip \
-    && $HOME/anaconda3/bin/pip --use-deprecated=legacy-resolver --no-cache-dir install -r requirements.txt \
-    && $HOME/anaconda3/bin/pip --no-cache-dir install -U -r requirements_rllib.txt \
-    && $HOME/anaconda3/bin/pip --no-cache-dir install -U -r requirements_train.txt \
-    && $HOME/anaconda3/bin/pip --no-cache-dir install -U -r requirements_tune.txt \
-    && $HOME/anaconda3/bin/pip --no-cache-dir install -U -r requirements_ml_docker.txt \
-    && $HOME/anaconda3/bin/pip --no-cache-dir install -U -r requirements_upstream.txt \
+    && $HOME/anaconda3/bin/pip --no-cache-dir install -U pip \
+    && $HOME/anaconda3/bin/pip --no-cache-dir install -U \
+           -r requirements.txt \
+           -r requirements_rllib.txt \
+           -r requirements_train.txt \
+           -r requirements_tune.txt \
+           -r requirements_upstream.txt \
+    # explicitly install (overwrite) pytorch with CUDA support
+    && $HOME/anaconda3/bin/pip --no-cache-dir install -U \
+           -r requirements_ml_docker.txt \
     # Remove dataclasses & typing because they are included in Python > 3.6
     && if [ $(python -c 'import sys; print(sys.version_info.minor)') != "6" ]; then \
         $HOME/anaconda3/bin/pip uninstall dataclasses typing -y; fi  \
-    && sudo rm requirements.txt && sudo rm requirements_ml_docker.txt \
-    && sudo rm requirements_tune.txt && sudo rm requirements_rllib.txt \
-    && sudo rm requirements_train.txt \
-    && sudo rm requirements_upstream.txt \
+    && sudo rm \
+        requirements.txt \
+        requirements_ml_docker.txt \
+        requirements_tune.txt \
+        requirements_rllib.txt \
+        requirements_train.txt \
+        requirements_upstream.txt \
     && sudo apt-get clean
 
 # Make sure tfp is installed correctly and matches tf version.

python/requirements.txt

@@ -85,7 +85,6 @@ pytest-lazy-fixture
 pytest-timeout
 pytest-virtualenv
 scikit-learn==0.22.2
-tensorflow==2.5.1
 testfixtures
 werkzeug
 xlrd

python/requirements/ml/requirements_dl.txt

@@ -1,20 +1,15 @@
 # These requirements are used for the CI and CPU-only Docker images so we install CPU only versions of torch.
-# For GPU Docker images, you should install requirements_ml_docker.txt instead.
+# For GPU Docker images, you should install requirements_ml_docker.txt afterwards.
 
+tensorflow==2.6.3
+tensorflow-probability==0.14.1
 
-# If you make changes to this file, please also change `requirements_ml_docker.txt`!
-
-
-# TODO(amogkam): Remove after https://github.com/tensorflow/tensorflow/issues/52922 is fixed.
-keras==2.6.0
-tensorflow==2.6.0
-tensorflow-probability==0.14.0
+# If you make changes below this line, please also change `requirements_ml_docker.txt`!
 
 torch==1.9.0;sys_platform=="darwin"
 torchvision==0.10.0;sys_platform=="darwin"
 
 # On non-OSX machines only install CPU version of torch and torchvision
 -f https://download.pytorch.org/whl/torch_stable.html
 torch==1.9.0+cpu;sys_platform!="darwin"
--f https://download.pytorch.org/whl/torch_stable.html
-torchvision==0.10.0+cpu;sys_platform!="darwin"
+torchvision==0.10.0+cpu;sys_platform!="darwin"

python/requirements/ml/requirements_rllib.txt

@@ -32,7 +32,6 @@ higher==0.2.1
 pyglet==1.5.15
 imageio-ffmpeg==0.4.5
 # For JSON reader/writer.
-smart_open==5.0.0
 # Ray Serve example
 starlette==0.16.0
 # ONNX

python/requirements/ml/requirements_tune.txt

@@ -33,7 +33,6 @@ shortuuid==1.0.1
 scikit-learn==0.24.2
 scikit-optimize==0.8.1
 sigopt==7.5.0
-smart_open==5.1.0
 timm==0.4.5
 transformers==4.10.0
 wandb==0.12.5

python/requirements_ml_docker.txt

@@ -5,13 +5,6 @@ tblib
 
 # If you make changes to anything below this line, please also change `requirements_dl.txt`!
 
-# TODO(amogkam): Remove after https://github.com/tensorflow/tensorflow/issues/52922 is fixed.
-keras==2.6.0
-# In TF >v2, GPU support is included in the base package.
-tensorflow==2.6.0
-tensorflow-probability==0.14.0
-
 -f https://download.pytorch.org/whl/torch_stable.html
 torch==1.9.0+cu111
--f https://download.pytorch.org/whl/torch_stable.html
 torchvision==0.10.0+cu111