From 7013286e2dc93aa90da96445b888da2912f139ed Mon Sep 17 00:00:00 2001
From: marcrasi <marcrasi@google.com>
Date: Tue, 15 Dec 2020 11:48:43 -0800
Subject: [PATCH] make Dockerfile able to use standard toolchain (#1149)

---
 Dockerfile                   |  2 ++
 Utilities/presubmit_build.sh | 11 ++++++++++-
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/Dockerfile b/Dockerfile
index 4a98d95c4..212e700ec 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -3,6 +3,7 @@ FROM gcr.io/swift-tensorflow/base-deps-cuda10.2-cudnn7-ubuntu18.04
 # Allows the caller to specify the toolchain to use.
 ARG swift_tf_url=https://storage.googleapis.com/swift-tensorflow-artifacts/nightlies/latest/swift-tensorflow-DEVELOPMENT-notf-ubuntu18.04.tar.gz
 ARG bazel_version=3.1.0
+ARG TENSORFLOW_USE_STANDARD_TOOLCHAIN=NO
 
 RUN if test -d /swift-apis/google-cloud-sdk; then \
   mv /swift-apis/google-cloud-sdk /opt/google-cloud-sdk; \
@@ -61,6 +62,7 @@ RUN cmake
       -D CMAKE_BUILD_TYPE=Release                                               \
       -D CMAKE_INSTALL_PREFIX=/swift-tensorflow-toolchain/usr                   \
       -D CMAKE_Swift_COMPILER=/swift-tensorflow-toolchain/usr/bin/swiftc        \
+      -D TENSORFLOW_USE_STANDARD_TOOLCHAIN="$TENSORFLOW_USE_STANDARD_TOOLCHAIN" \
       -G Ninja                                                                  \
       -S /swift-apis
 RUN cmake --build /BinaryCache/tensorflow-swift-apis --verbose
diff --git a/Utilities/presubmit_build.sh b/Utilities/presubmit_build.sh
index ac1fdb34e..59ed8b228 100644
--- a/Utilities/presubmit_build.sh
+++ b/Utilities/presubmit_build.sh
@@ -10,6 +10,12 @@ gcloud beta auth configure-docker
 # 'swift_tf_bigstore_gfile', if it exists.
 if [[ ! -z ${swift_tf_bigstore_gfile+x} ]]; then
   export swift_tf_url="${swift_tf_bigstore_gfile/\/bigstore/https://storage.googleapis.com}"
+  case "$swift_tf_url" in
+    *stock*) export TENSORFLOW_USE_STANDARD_TOOLCHAIN=YES ;;
+    *) export TENSORFLOW_USE_STANDARD_TOOLCHAIN=NO ;;
+  esac
+else
+  export TENSORFLOW_USE_STANDARD_TOOLCHAIN=NO
 fi
 
 # Help debug the job's disk space.
@@ -26,7 +32,10 @@ df -h
 
 cd github/swift-apis
 cp -R /opt/google-cloud-sdk .
-sudo -E docker build -t build-img -f Dockerfile --build-arg swift_tf_url .
+sudo -E docker build -t build-img -f Dockerfile \
+  --build-arg swift_tf_url \
+  --build-arg TENSORFLOW_USE_STANDARD_TOOLCHAIN \
+  .
 
 sudo docker create --name build-container build-img
 mkdir -p "$KOKORO_ARTIFACTS_DIR/swift_apis_benchmarks"