Reducing space of the CPU docker image (#11807)

* Install CPU only version of torch, as there is no need to install cuDNN and CUDA

* Install CPU only version of torch, as there is no need to install cuDNN and CUDA

* Install CPU only version of torch

There is no need to install cuDNN and CUDA.

* Update Dockerfile

* Updated torch installation to install CPU only

* Updated torch installation to install CPU only

* Install CPU-only torch

* Updated torch installation

Install CPU-only wheel

* Pinned torch to latest

* Unpinned versions of torch CPU

* Unpinned versions of torch CPU

* Updated Dockerfile

Co-authored-by: Deeptendu Santra <[email protected]>

* Remove nvidia-ml-py library

* Delete requirements_cpu.txt

* Added `optional_cpu.txt`

* Updated `optional_cpu.txt`'s comments

* Delete optional_cpu.txt

* Install only cpu speific variant libraries

* Minor change in Dockerfile

* Reverted back

* Defaulted to CPU-based libraries

* test with optional_gpu

* remove optional_gpu

---------

Co-authored-by: Abhimanyu <[email protected]>
Co-authored-by: Deeptendu Santra <[email protected]>
Co-authored-by: Deeptendu Santra <[email protected]>

docker/Dockerfile

@@ -23,28 +23,33 @@ RUN apt-get update && \
     pip3 install --upgrade pip && \
     pip3 install setuptools==58.5.3
 
-RUN pip3 install --no-cache-dir torch -f https://download.pytorch.org/whl/torch_stable.html
-RUN pip3 install --no-cache-dir --upgrade torch-scatter -f https://pytorch-geometric.com/whl
+
+RUN pip3 install --no-cache-dir torch --extra-index-url https://download.pytorch.org/whl/cpu
+RUN export ver=$(pip show torch | grep Version | cut -d ' ' -f2) && \
+    pip3 install --no-cache-dir --upgrade torch-scatter -f https://data.pyg.org/whl/torch-$ver.html
 
 # Install Ivy Upstream
 RUN git clone --progress --recurse-submodules https://github.com/unifyai/ivy --depth 1 && \
     cd ivy && \
+    sed -i '/nvidia/d' requirements/requirements.txt && \
     cat requirements/requirements.txt | grep -v "ivy-" | pip3 install --no-cache-dir -r /dev/stdin && \
     cat requirements/optional.txt | grep -v "ivy-" | pip3 install --no-cache-dir -r /dev/stdin && \
     python3 -m pip install --user -e . && \
     cd ivy_tests/test_array_api && \
     pip3 install --no-cache-dir -r requirements.txt
 
 # Install local requirements
-COPY requirements/requirements.txt .
+RUN cp ivy/requirements/requirements.txt .
 RUN pip3 install --no-cache-dir -r requirements.txt
 
 # Install local optional
 COPY requirements/optional.txt .
+COPY requirements/optional_gpu.txt .
 RUN pip3 install --no-cache-dir -r optional.txt
 
 COPY run_tests_CLI/test_dependencies.py .
-RUN python3 test_dependencies.py -fp requirements.txt,optional.txt && \
+RUN python3 test_dependencies.py -fp requirements.txt,optional_gpu.txt && \
     rm -rf requirements.txt && \
     rm -rf optional.txt && \
+    rm -rf optional_gpu.txt && \
     rm -rf test_dependencies.py

docker/DockerfileGPU

@@ -98,14 +98,14 @@ RUN apt-get install -y python-opengl
 # Ivy #
 # ----#
 
-# Make sure torch installed before torch-scatter (in optional.txt)
+# Make sure torch installed before torch-scatter (in optional_gpu.txt)
 RUN pip install --no-cache-dir torch==1.11.0+cu113 --extra-index-url https://download.pytorch.org/whl/cu113
 
 # Install Ivy Upstream
 RUN git clone --recurse-submodules https://github.com/unifyai/ivy --depth 1 && \
     cd ivy && \
     cat requirements/requirements.txt | grep -v "ivy-" | pip3 install --no-cache-dir -r /dev/stdin && \
-    cat requirements/optional.txt | grep -v "ivy-" | pip3 install --no-cache-dir -r /dev/stdin && \
+    cat requirements/optional_gpu.txt | grep -v "ivy-" | pip3 install --no-cache-dir -r /dev/stdin && \
     python3 -m pip install --user -e . && \
     cd ivy_tests/test_array_api && \
     pip3 install --no-cache-dir -r requirements.txt
@@ -115,14 +115,14 @@ COPY requirements/requirements.txt .
 RUN pip3 install --no-cache-dir -r requirements.txt
 
 # Install local optional
-COPY requirements/optional.txt .
-RUN pip3 install --no-cache-dir -r optional.txt
+COPY requirements/optional_gpu.txt .
+RUN pip3 install --no-cache-dir -r optional_gpu.txt
 
-# Install jax cuda after optional.txt, otherwise cpu version will override
+# Install jax cuda after optional_gpu.txt, otherwise cpu version will override
 RUN pip install --no-cache-dir jaxlib==0.3.14+cuda11.cudnn82 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 
 COPY run_tests_CLI/test_dependencies.py .
-RUN python3 test_dependencies.py -fp requirements.txt,optional.txt && \
+RUN python3 test_dependencies.py -fp requirements.txt,optional_gpu.txt && \
     rm -rf requirements.txt && \
-    rm -rf optional.txt && \
+    rm -rf optional_gpu.txt && \
     rm -rf test_dependencies.py

requirements/optional.txt

@@ -6,9 +6,9 @@ pymongo # unpinned
 redis # unpinned
 matplotlib # unpinned
 opencv-python # unpinned , mod_name=cv2
-jax # unpinned, we test for latest version now
+jax[cpu] # unpinned, we test for latest version now
 jaxlib # unpinned, we test for latest version now
-tensorflow  # unpinned, we test for latest version now
+tensorflow-cpu  # unpinned, we test for latest version now
 tensorflow-addons # unpinned, we test for latest version now, mod_name=tensorflow_addons
 tensorflow-probability # unpinned, we test for latest version now, mod_name=tensorflow_probability
 torch # unpinned, we test for latest version now

requirements/optional_gpu.txt

@@ -0,0 +1,21 @@
+h5py # unpinned
+pytest # unpinned
+networkx # unpinned
+hypothesis # unpinned
+pymongo # unpinned
+redis # unpinned
+matplotlib # unpinned
+opencv-python # unpinned , mod_name=cv2
+jax # unpinned, we test for latest version now
+jaxlib # unpinned, we test for latest version now
+tensorflow  # unpinned, we test for latest version now
+tensorflow-addons # unpinned, we test for latest version now, mod_name=tensorflow_addons
+tensorflow-probability # unpinned, we test for latest version now, mod_name=tensorflow_probability
+torch # unpinned, we test for latest version now
+torch-scatter # unpinned, mod_name=torch_scatter
+functorch # unpinned, we test for latest version now
+scipy # unpinned
+dm-haiku # unpinned mod_name=haiku
+pydriller
+tqdm
+coverage