Utilize SIG Build Docker Images (#2515)

* Utilize SIG Build build image Co-authored-by: bhack <[email protected]>
tensorflow · Dec 28, 2021 · d19c3ef · d19c3ef
1 parent 33727be
commit d19c3ef
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 25 deletions.
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
@@ -124,11 +124,13 @@ 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
       - run: |
           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
@@ -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
@@ -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
@@ -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