Necessary minimum to make the dev image build again.

[CI] switch bazel to 5.1.1 and update dev docker container

Author: bashimao

.github/workflows/ci_test.yml

@@ -10,7 +10,7 @@ on:
       - master
       - r*
 env:
-  USE_BAZEL_VERSION: "3.7.2"
+  USE_BAZEL_VERSION: "5.1.1"
 
 jobs:
   yapf-test:

.github/workflows/release.yml

@@ -17,7 +17,7 @@ on:
 env:
   MIN_PY_VERSION: '3.7'
   MAX_PY_VERSION: '3.9'
-  USE_BAZEL_VERSION: "3.7.2"
+  USE_BAZEL_VERSION: "5.1.1"
   HOROVOD_VERSION: '0.23.0'
 
 jobs:
@@ -167,11 +167,15 @@ jobs:
     name: Upload dev container to DockerHub
     needs: [release-wheel, test-with-bazel]
     runs-on: ubuntu-18.04
+    strategy:
+      matrix:
+        py-version: ['3.7', '3.8', '3.9']
     if: (github.event_name == 'push' && github.ref == 'refs/heads/master')
     steps:
       - uses: actions/checkout@v2
       - run: |
           set -e -x
           echo ${{ secrets.DOCKER_PW }} | docker login --username ${{ secrets.DOCKER_USER }} --password-stdin
+          export PY_VERSION=${{ matrix.py-version }}
           bash .github/workflows/github_build_dev_container.sh
-          docker push tfra/dev_container:latest-cpu
+          docker push tfra/dev_container:latest-python${{ matrix.py-version }}

README.md

@@ -134,8 +134,15 @@ but maybe this doc can help you : [Extract headers from TensorFlow compiling dir
 At the same time, we find some OPs used by TRFA have better performance, so we highly recommend you update TensorFlow to 2.x.
 
 #### Installing from Source
+
+For all developers, we recommend you use the development docker containers which are all GPU enabled:
+```shell
+docker pull tfra/dev_container:latest-python3.8  # "3.7", "3.9" are all avaliable.
+docker run --privileged --gpus all -it --rm -v $(pwd):$(pwd) tfra/dev_container:latest-3.8
+```
+
 ##### CPU Only
-You can also install from source. This requires the [Bazel](https://bazel.build/) build system (version == 3.7.2).
+You can also install from source. This requires the [Bazel](https://bazel.build/) build system (version == 5.1.1).
 Please install a TensorFlow on your compiling machine, The compiler needs to know the version of Tensorflow and 
 its headers according to the installed TensorFlow. 
 
@@ -158,18 +165,20 @@ pip install artifacts/tensorflow_recommenders_addons-*.whl
 ##### GPU Support
 Only `TF_NEED_CUDA=1` is required and other environment variables are optional:
 ```shell
-PY_VERSION="3.7" \
-TF_NEED_CUDA=1 \
-TF_CUDA_VERSION=11.2 \
-TF_CUDNN_VERSION=8.1 \
-CUDA_TOOLKIT_PATH="/usr/local/cuda" \
-CUDNN_INSTALL_PATH="/usr/lib/x86_64-linux-gnu" \
+export TF_VERSION="2.5.1"  # "2.7.0", "2.5.1" are well tested.
+export PY_VERSION="3.8" 
+export TF_NEED_CUDA=1
+export TF_CUDA_VERSION=11.2
+export TF_CUDNN_VERSION=8.1
+export CUDA_TOOLKIT_PATH="/usr/local/cuda"
+export CUDNN_INSTALL_PATH="/usr/lib/x86_64-linux-gnu"
+
 python configure.py
 ```
 And then build the pip package and install:
 ```shell
 bazel build --enable_runfiles build_pip_pkg
-TF_NEED_CUDA=1 bazel-bin/build_pip_pkg artifacts
+bazel-bin/build_pip_pkg artifacts
 pip install artifacts/tensorflow_recommenders_addons_gpu-*.whl
 ```
 ##### Apple Silicon Support (Beta Release)

configure.py

@@ -40,10 +40,10 @@ def _VALID_BAZEL_VERSION(tf_version):
         'refering to the previous COMMIT to compile properly by themselves.')
     return target_bazel
   elif tf_version >= "2.0.0":
-    target_bazel = "3.7.2"
+    target_bazel = "5.1.1"
     logging.info(
         'To ensure code compatibility with Bazel rules_foreign_cc component, '
-        'we specify Bazel version greater than 3.7.2 '
+        'we specify Bazel version greater than 5.1.1 '
         'for Tensorflow versions greater than 2.0.0.')
     return target_bazel
   else:

tools/build_dev_container.sh

@@ -1,10 +1,11 @@
 #!/usr/bin/env bash
 
 set -x -e
-
-DOCKER_BUILDKIT=1 docker build \
+docker build \
     -f tools/docker/dev_container.Dockerfile \
     --build-arg TF_VERSION=2.5.1 \
-    --build-arg TF_PACKAGE=tensorflow-cpu \
-    --target dev_container_cpu \
-    -t tfaddons/dev_container:latest-cpu ./
+    --build-arg TF_PACKAGE=tensorflow-gpu \
+    --build-arg PY_VERSION=$PY_VERSION \
+    --no-cache \
+    --target dev_container \
+    -t tfra/dev_container:latest-python$PY_VERSION ./

tools/docker/cpu_tests.Dockerfile

@@ -2,7 +2,7 @@
 FROM python:3.7 as build_wheel
 
 ARG TF_VERSION=2.5.1
-ARG USE_BAZEL_VERSION=3.7.2
+ARG USE_BAZEL_VERSION=5.1.1
 ARG MPI_VERSION="4.1.1"
 ARG HOROVOD_VERSION="0.23.0"
 

tools/docker/cuda11.0-cudnn8-ubuntu18.04-manylinux2010-multipython.Dockerfile

@@ -79,7 +79,7 @@ RUN apt-get update && apt-get install -y \
     rm -rf /var/lib/apt/lists/*
 
 COPY install/install_bazel.sh /install/
-RUN /install/install_bazel.sh "3.7.2"
+RUN /install/install_bazel.sh "5.1.1"
 
 COPY install/build_and_install_python.sh /install/
 RUN /install/build_and_install_python.sh "3.5.9"

tools/docker/cuda11.2-cudnn8.1-ubuntu18.04-manylinux2010-multipython.Dockerfile

@@ -77,17 +77,14 @@ RUN apt-get update && apt-get install -y \
 RUN chmod 777 /tmp/
 WORKDIR /tmp/
 
-COPY install/install_trt.sh /install/
-RUN /install/install_trt.sh "7.2.3-1+cuda11.1"
-
 COPY install/install_nccl.sh /install/
 RUN /install/install_nccl.sh "2.8.4-1+cuda11.2"
 
 COPY install/install_rocksdb.sh /install/
 RUN /install/install_rocksdb.sh "6.22.1"
 
 COPY install/install_bazel.sh /install/
-RUN /install/install_bazel.sh "3.7.2"
+RUN /install/install_bazel.sh "5.1.1"
 
 COPY install/build_and_install_python.sh /install/
 RUN /install/build_and_install_python.sh "3.6.9"

tools/docker/dev_container.Dockerfile

@@ -1,27 +1,29 @@
 #syntax=docker/dockerfile:1.1.5-experimental
-FROM tensorflow/tensorflow:2.1.0-custom-op-ubuntu16 as dev_container_cpu
+ARG PY_VERSION
+ARG IMAGE_TYPE
+
+# Currenly all of our dev images are GPU capable but at a cost of being quite large.
+# See https://github.com/tensorflow/build/pull/47
+FROM tensorflow/build:latest-python$PY_VERSION as dev_container
 ARG TF_PACKAGE
 ARG TF_VERSION
 
-# Temporary until custom-op container is updated
-RUN ln -sf /usr/bin/python3 /usr/bin/python
-RUN ln -sf /usr/local/bin/pip3 /usr/local/bin/pip
+RUN pip uninstall $TF_PACKAGE -y
 RUN pip install --default-timeout=1000 $TF_PACKAGE==$TF_VERSION
 
 COPY tools/install_deps /install_deps
 COPY requirements.txt /tmp/requirements.txt
-RUN pip install -r /install_deps/black.txt \
-    -r /install_deps/flake8.txt \
+RUN pip install -r /install_deps/yapf.txt \
     -r /install_deps/pytest.txt \
     -r /install_deps/typedapi.txt \
     -r /tmp/requirements.txt
 
 RUN bash /install_deps/buildifier.sh
 RUN bash /install_deps/clang-format.sh
 
-ENV TFRA_DEV_CONTAINER="1"
+ENV ADDONS_DEV_CONTAINER="1"
 
 # Clean up
 RUN apt-get autoremove -y \
     && apt-get clean -y \
-    && rm -rf /var/lib/apt/lists/*
+    && rm -rf /var/lib/apt/lists/*

tools/run_build.sh

@@ -6,5 +6,5 @@ set -e
 
 DOCKER_BUILDKIT=1 docker build \
     -f tools/docker/sanity_check.Dockerfile \
-    --build-arg USE_BAZEL_VERSION=${USE_BAZEL_VERSION:-"3.7.2"} \
+    --build-arg USE_BAZEL_VERSION=${USE_BAZEL_VERSION:-"5.1.1"} \
     --target=${1} ./