Fix blacklist-related bug in pip test

The bug was introduced during earlier refactoring. It muted the GPU-specific test blacklist.
Change: 118793075

2 files changed, 33 insertions, 13 deletions

diff --git a/tensorflow/tools/ci_build/builds/pip.sh b/tensorflow/tools/ci_build/builds/pip.sh
index 96b32dd1f4..6522975f28 100755
--- a/tensorflow/tools/ci_build/builds/pip.sh
+++ b/tensorflow/tools/ci_build/builds/pip.sh
@@ -77,23 +77,16 @@ for ARG in $@; do
 done
 
 PIP_BUILD_TARGET="//tensorflow/tools/pip_package:build_pip_package"
+GPU_FLAG=""
 if [[ ${CONTAINER_TYPE} == "cpu" ]]; then
   bazel build -c opt ${PIP_BUILD_TARGET} || die "Build failed."
 elif [[ ${CONTAINER_TYPE} == "gpu" ]]; then
   bazel build -c opt --config=cuda ${PIP_BUILD_TARGET} || die "Build failed."
+  GPU_FLAG="--gpu"
 else
   die "Unrecognized container type: \"${CONTAINER_TYPE}\""
 fi
 
-echo "PY_TEST_WHITELIST: ${PY_TEST_WHITELIST}"
-echo "PY_TEST_BLACKLIST: ${PY_TEST_BLACKLIST}"
-echo "PY_TEST_GPU_BLACKLIST: ${PY_TEST_GPU_BLACKLIST}"
-
-# Append GPU-only test blacklist
-if [[ ${CONTAINER_TYPE} == "gpu" ]]; then
-  PY_TEST_BLACKLIST="${PY_TEST_BLACKLIST}:${PY_TEST_GPU_BLACKLIST}"
-fi
-
 # If still in a virtualenv, deactivate it first
 if [[ ! -z "$(which deactivate)" ]]; then
   echo "It appears that we are already in a virtualenv. Deactivating..."
@@ -191,7 +184,7 @@ fi
 # Call test_installation.sh to perform test-on-install
 DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 
-"${DIR}/test_installation.sh" --virtualenv || \
+"${DIR}/test_installation.sh" --virtualenv ${GPU_FLAG} || \
     die "PIP tests-on-install FAILED"
 
 # Optional: Run the tutorial tests
diff --git a/tensorflow/tools/ci_build/builds/test_installation.sh b/tensorflow/tools/ci_build/builds/test_installation.sh
index f658357ceb..b6385ac47f 100755
--- a/tensorflow/tools/ci_build/builds/test_installation.sh
+++ b/tensorflow/tools/ci_build/builds/test_installation.sh
@@ -18,12 +18,15 @@
 # and run the Python unit tests from the source code on the installation
 #
 # Usage:
-#   test_installation.sh [--virtualenv]
+#   test_installation.sh [--virtualenv] [--gpu]
 #
 # If the flag --virtualenv is set, the script will use "python" as the Python
 # binary path. Otherwise, it will use tools/python_bin_path.sh to determine
 # the Python binary path.
 #
+# The --gpu flag informs the script that this is a GPU build, so that the
+# appropriate test blacklists can be applied accordingly.
+#
 # When executing the Python unit tests, the script obeys the shell
 # variables: PY_TEST_WHITELIST, PY_TEST_BLACKLIST, PY_TEST_GPU_BLACKLIST,
 #
@@ -88,6 +91,10 @@ fi
 
 # =============================================================================
 
+echo "PY_TEST_WHITELIST: ${PY_TEST_WHITELIST}"
+echo "PY_TEST_BLACKLIST: ${PY_TEST_BLACKLIST}"
+echo "PY_TEST_GPU_BLACKLIST: ${PY_TEST_GPU_BLACKLIST}"
+
 
 # Helper functions
 # Get the absolute path from a path
@@ -104,9 +111,24 @@ die() {
 # Script directory
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 
+# Process input arguments
+IS_VIRTUALENV=0
+IS_GPU=0
+while true; do
+  if [[ "$1" == "--virtualenv" ]]; then
+    IS_VIRTUALENV=1
+  elif [[ "$1" == "--gpu" ]]; then
+    IS_GPU=1
+  fi
+  shift
+
+  if [[ -z "$1" ]]; then
+    break
+  fi
+done
+
 # Obtain the path to Python binary
-# source tools/python_bin_path.sh
-if [[ "$1" == "--virtualenv" ]]; then
+if [[ ${IS_VIRTUALENV} == "1" ]]; then
   PYTHON_BIN_PATH="$(which python)"
 else
   source tools/python_bin_path.sh
@@ -118,6 +140,11 @@ if [[ -z "${PYTHON_BIN_PATH}" ]]; then
 "did you run configure?"
 fi
 
+# Append GPU-only test blacklist
+if [[ ${IS_GPU} == "1" ]]; then
+  PY_TEST_BLACKLIST="${PY_TEST_BLACKLIST}:${PY_TEST_GPU_BLACKLIST}"
+fi
+
 # Determine the major and minor versions of Python being used (e.g., 2.7)
 # This info will be useful for determining the directory of the local pip
 # installation of Python