use XLA base image for TPU testing (#2536)

* drop py3.6

* use base image

* typo

* skip extra

* drop cache

.github/workflows/docker-builds.yml

@@ -62,7 +62,7 @@ jobs:
     runs-on: ubuntu-20.04
     strategy:
       matrix:
-        python_version: [3.6, 3.7]
+        python_version: [3.7]  # TODO: add 3.6, for now the Google docker image has missing 3.6
     steps:
       - name: Checkout
         uses: actions/checkout@v2

dockers/tpu-extras/Dockerfile

@@ -26,6 +26,10 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     update-alternatives --install /usr/bin/python${PYTHON_VERSION%%.*} python${PYTHON_VERSION%%.*} /usr/bin/python${PYTHON_VERSION} 1 && \
     update-alternatives --install /usr/bin/python python /usr/bin/python${PYTHON_VERSION} 1 && \
 
+# Disable cache
+    pip config set global.cache-dir false && \
+    pip install "pip>20.1" -U  && \
+
 # Install Pytorch XLA
     py_version=${PYTHON_VERSION/./} && \
     # Python 3.7 wheels are available. Replace cp36-cp36m with cp37-cp37m
@@ -42,7 +46,8 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
 
 # Cleaning
     apt-get autoremove -y && \
-    apt-get clean
+    apt-get clean && \
+    rm -rf /root/.cache
 
 # Install pytorch-lightning dependencies.
 RUN \
@@ -58,7 +63,8 @@ RUN \
     python -c "fname = 'requirements/extra.txt' ; lines = [line for line in open(fname).readlines() if not line.startswith('horovod')] ; open(fname, 'w').writelines(lines)" && \
     pip install --requirement ./requirements/extra.txt && \
     cd .. && \
-    rm -rf pytorch-lightning-master
+    rm -rf pytorch-lightning-master && \
+    rm -rf /root/.cache
 
 ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/lib/"
 

dockers/tpu-tests/Dockerfile

@@ -1,69 +1,30 @@
-FROM google/cloud-sdk:slim
+ARG PYTHON_VERSION=3.7
+
+FROM pytorchlightning/pytorch_lightning:XLA-extras-py${PYTHON_VERSION}
 
 # Build args.
 ARG GITHUB_REF=refs/heads/master
 ARG TEST_IMAGE=0
 
 # This Dockerfile installs pytorch/xla 3.7 wheels. There are also 3.6 wheels available; see below.
-ENV PYTHON_VERSION=3.7
 
-SHELL ["/bin/bash", "-c"]
-
-RUN apt-get update && apt-get install -y --no-install-recommends \
-        python${PYTHON_VERSION} \
-        python${PYTHON_VERSION}-dev \
-        $( [ ${PYTHON_VERSION%%.*} -ge 3 ] && echo "python${PYTHON_VERSION%%.*}-distutils" ) \
-        build-essential \
-        cmake \
-        git \
-        curl \
-        wget \
-        ca-certificates \
-        libomp5 \
-    && \
-
-# install python dependencies
-    wget https://bootstrap.pypa.io/get-pip.py --progress=bar:force:noscroll --no-check-certificate && \
-    python${PYTHON_VERSION} get-pip.py && \
-    rm get-pip.py && \
-
-# Set the default python and install PIP packages
-    update-alternatives --install /usr/bin/python${PYTHON_VERSION%%.*} python${PYTHON_VERSION%%.*} /usr/bin/python${PYTHON_VERSION} 1 && \
-    update-alternatives --install /usr/bin/python python /usr/bin/python${PYTHON_VERSION} 1
-
-RUN py_version=${PYTHON_VERSION/./} && \
-    # Python 3.7 wheels are available. Replace cp36-cp36m with cp37-cp37m
-    gsutil cp "gs://tpu-pytorch/wheels/torch-nightly-cp${py_version}-cp${py_version}m-linux_x86_64.whl" . && \
-    gsutil cp "gs://tpu-pytorch/wheels/torch_xla-nightly-cp${py_version}-cp${py_version}m-linux_x86_64.whl" . && \
-    gsutil cp "gs://tpu-pytorch/wheels/torchvision-nightly-cp${py_version}-cp${py_version}m-linux_x86_64.whl" . && \
-    pip install "torch-nightly-cp${py_version}-cp${py_version}m-linux_x86_64.whl" && \
-    pip install "torch_xla-nightly-cp${py_version}-cp${py_version}m-linux_x86_64.whl" && \
-    pip install "torchvision-nightly-cp${py_version}-cp${py_version}m-linux_x86_64.whl" && \
-    rm "torch-nightly-cp${py_version}-cp${py_version}m-linux_x86_64.whl" && \
-    rm "torch_xla-nightly-cp${py_version}-cp${py_version}m-linux_x86_64.whl" && \
-    rm "torchvision-nightly-cp${py_version}-cp${py_version}m-linux_x86_64.whl" && \
-    pip install mkl
-
-ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/lib/"
-
-RUN python -c "import torch; print(torch.__version__)"
+#SHELL ["/bin/bash", "-c"]
 
 # Install pytorch-lightning at the current PR, plus dependencies.
 RUN git clone https://github.com/PyTorchLightning/pytorch-lightning.git && \
     cd pytorch-lightning && \
     git fetch origin $GITHUB_REF:CI && \
     git checkout CI && \
-    pip install --requirement ./requirements/base.txt
+    pip install --requirement ./requirements/base.txt --no-cache-dir
 
 # If using this image for tests, intall more dependencies and don"t delete
 # the source code where the tests live.
 RUN \
-    # TODO: use conda sources if possible
     # drop Horovod
-    # python -c "fname = './pytorch-lightning/requirements/extra.txt' ; lines = [ln for ln in open(fname).readlines() if not ln.startswith('horovod')] ; open(fname, 'w').writelines(lines)" && \
-    # pip install -r pytorch-lightning/requirements/extra.txt ; && \
+    #python -c "fname = 'pytorch-lightning/requirements/extra.txt' ; lines = [line for line in open(fname).readlines() if not line.startswith('horovod')] ; open(fname, 'w').writelines(lines)" && \
+    #pip install --requirement pytorch-lightning/requirements/extra.txt --no-cache-dir && \
     if [ $TEST_IMAGE -eq 1 ] ; then \
-        pip install -r pytorch-lightning/requirements/test.txt ; \
+        pip install --requirement pytorch-lightning/requirements/test.txt --no-cache-dir ; \
     else \
         rm -rf pytorch-lightning ; \
     fi