From d19c3efa5e4b9924e5f1e0b603fe0a22b8efcbfe Mon Sep 17 00:00:00 2001
From: Sean Morgan <seanmorgan@outlook.com>
Date: Mon, 27 Dec 2021 19:51:40 -0800
Subject: [PATCH] Utilize SIG Build Docker Images (#2515)

* Utilize SIG Build build image
Co-authored-by: bhack <bhack@users.noreply.github.com>
---
 .github/workflows/release.yml         |  4 +++-
 tools/build_dev_container.sh          |  7 ++++---
 tools/docker/build_wheel.Dockerfile   | 19 ++-----------------
 tools/docker/dev_container.Dockerfile | 10 ++++++----
 4 files changed, 15 insertions(+), 25 deletions(-)

diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
index e8cb3da6f6..e4ddcee49f 100644
--- a/.github/workflows/release.yml
+++ b/.github/workflows/release.yml
@@ -124,6 +124,8 @@ jobs:
     name: Upload dev container to DockerHub
     needs: [release-wheel, test-with-bazel]
     runs-on: ubuntu-18.04
+    env:
+      PY_VERSION: '3.9'
     if: (github.event_name == 'push' && github.ref == 'refs/heads/master')
     steps:
       - uses: actions/checkout@v2
@@ -131,4 +133,4 @@ jobs:
           set -e -x
           echo ${{ secrets.DOCKER_PW }} | docker login --username ${{ secrets.DOCKER_USER }} --password-stdin
           bash .github/workflows/github_build_dev_container.sh
-          docker push tfaddons/dev_container:latest-cpu
+          docker push tfaddons/dev_container:latest-gpu
diff --git a/tools/build_dev_container.sh b/tools/build_dev_container.sh
index 59cf84e7f7..4d57eeab0d 100755
--- a/tools/build_dev_container.sh
+++ b/tools/build_dev_container.sh
@@ -5,7 +5,8 @@ set -x -e
 docker build \
     -f tools/docker/dev_container.Dockerfile \
     --build-arg TF_VERSION=2.7.0 \
-    --build-arg TF_PACKAGE=tensorflow-cpu \
+    --build-arg TF_PACKAGE=tensorflow \
+    --build-arg PY_VERSION=$PY_VERSION \
     --no-cache \
-    --target dev_container_cpu \
-    -t tfaddons/dev_container:latest-cpu ./
+    --target dev_container \
+    -t tfaddons/dev_container:latest-gpu ./
diff --git a/tools/docker/build_wheel.Dockerfile b/tools/docker/build_wheel.Dockerfile
index db247b84c8..4ad7e610fc 100644
--- a/tools/docker/build_wheel.Dockerfile
+++ b/tools/docker/build_wheel.Dockerfile
@@ -1,23 +1,9 @@
 #syntax=docker/dockerfile:1.1.5-experimental
-ARG TF_VERSION
 ARG PY_VERSION
-FROM gcr.io/tensorflow-testing/nosla-cuda11.2-cudnn8.1-ubuntu18.04-manylinux2010-multipython as base_install
-ENV TF_NEED_CUDA="1"
-
-# Required for setuptools v50.0.0
-# https://setuptools.readthedocs.io/en/latest/history.html#v50-0-0
-# https://github.com/pypa/setuptools/issues/2352
-ENV SETUPTOOLS_USE_DISTUTILS=stdlib
-
-# Fix presented in
-# https://stackoverflow.com/questions/44967202/pip-is-showing-error-lsb-release-a-returned-non-zero-exit-status-1/44967506
-RUN echo "#! /usr/bin/python2.7" >> /usr/bin/lsb_release2
-RUN cat /usr/bin/lsb_release >> /usr/bin/lsb_release2
-RUN mv /usr/bin/lsb_release2 /usr/bin/lsb_release
+FROM tensorflow/build:latest-python$PY_VERSION as base_install
 
+ENV TF_NEED_CUDA="1"
 ARG PY_VERSION
-RUN ln -sf /usr/local/bin/python$PY_VERSION /usr/bin/python
-
 ARG TF_VERSION
 RUN python -m pip install --default-timeout=1000 tensorflow==$TF_VERSION
 
@@ -28,7 +14,6 @@ COPY requirements.txt .
 RUN python -m pip install -r requirements.txt
 
 COPY ./ /addons
-RUN rm /addons/.bazeliskrc
 WORKDIR /addons
 
 # -------------------------------------------------------------------
diff --git a/tools/docker/dev_container.Dockerfile b/tools/docker/dev_container.Dockerfile
index 2477b682e7..68ba49ee86 100644
--- a/tools/docker/dev_container.Dockerfile
+++ b/tools/docker/dev_container.Dockerfile
@@ -1,11 +1,13 @@
 #syntax=docker/dockerfile:1.1.5-experimental
-FROM gcr.io/tensorflow-testing/nosla-cuda11.2-cudnn8.1-ubuntu18.04-manylinux2010-multipython 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
 
-RUN ln -sf /usr/local/bin/python3.9 /usr/bin/python
-RUN ln -sf /usr/local/bin/pip3.9 /usr/local/bin/pip
-
 RUN pip install --default-timeout=1000 $TF_PACKAGE==$TF_VERSION
 
 COPY tools/install_deps /install_deps