• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/env bash
2# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#     http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15# ==============================================================================
16# Build the Python PIP installation package for TensorFlow and install
17# the package.
18#
19# Usage:
20#   pip_new.sh
21#
22# Required step(s):
23#   Run configure.py prior to running this script.
24#
25# Required environment variable(s):
26#   CONTAINER_TYPE:      (CPU | GPU)
27#   OS_TYPE:             (UBUNTU | MACOS)
28#   TF_PYTHON_VERSION:   ( python3.6 | python3.7 | python3.8 )
29#   TF_BUILD_FLAGS:      Bazel build flags.
30#                          e.g. TF_BUILD_FLAGS="--config=opt"
31#   TF_TEST_FLAGS:       Bazel test flags.
32#                          e.g. TF_TEST_FLAGS="--verbose_failures=true \
33#                               --build_tests_only --test_output=errors"
34#   TF_TEST_FILTER_TAGS: Filtering tags for bazel tests. More specifically,
35#                        input tags for `--test_filter_tags` flag.
36#                          e.g. TF_TEST_FILTER_TAGS="no_pip,-nomac,no_oss"
37#   TF_TEST_TARGETS:     Bazel test targets.
38#                          e.g. TF_TEST_TARGETS="//tensorflow/... \
39#                               -//tensorflow/contrib/... \
40#                               -//tensorflow/python/..."
41#   IS_NIGHTLY:          Nightly run flag.
42#                          e.g. IS_NIGHTLY=1  # nightly runs
43#                          e.g. IS_NIGHTLY=0  # non-nightly runs
44#
45# Optional environment variables. If provided, overwrites any default values.
46#   TF_PIP_TESTS:        PIP tests to run. If NOT specified, skips all tests.
47#                          e.g. TF_PIP_TESTS="test_pip_virtualenv_clean \
48#                               test_pip_virtualenv_clean \
49#                               test_pip_virtualenv_oss_serial"
50#   TF_PROJECT_NAME:     Name of the project. This string will be pass onto
51#                        the wheel file name. For nightly builds, it will be
52#                        overwritten to 'tf_nightly'. For gpu builds, '_gpu'
53#                        will be appended.
54#                          e.g. TF_PROJECT_NAME="tensorflow"
55#                          e.g. TF_PROJECT_NAME="tf_nightly_gpu"
56#   TF_PIP_TEST_ROOT:    Root directory for building and testing pip pkgs.
57#                          e.g. TF_PIP_TEST_ROOT="pip_test"
58#   TF_BUILD_BOTH_GPU_PACKAGES:    (1 | 0)
59#                                  1 will build both tensorflow (w/ gpu support)
60#                                  and tensorflow-gpu pip package. Will
61#                                  automatically handle adding/removing of _gpu
62#                                  suffix depending on what project name was
63#                                  passed. Only work for Ubuntu.
64#   TF_BUILD_BOTH_CPU_PACKAGES:    (1 | 0)
65#                                  1 will build both tensorflow (no gpu support)
66#                                  and tensorflow-cpu pip package. Will
67#                                  automatically handle adding/removing of _cpu
68#                                  suffix depending on what project name was
69#                                  passed. Only work for MacOS
70#
71# To-be-deprecated variable(s).
72#   GIT_TAG_OVERRIDE:    Values for `--git_tag_override`. This flag gets passed
73#                        in as `--action_env` for bazel build and tests.
74#   TF_BUILD_INSTALL_EXTRA_PIP_PACKAGES:
75#                        Additional pip packages to be installed.
76#                        Caveat: pip version needs to be checked prior.
77#
78# ==============================================================================
79
80# set bash options
81set -e
82set -x
83
84###########################################################################
85# General helper function(s)
86###########################################################################
87
88# Strip leading and trailing whitespaces
89str_strip () {
90  echo -e "$1" | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//'
91}
92
93# Convert string to all lower case
94lowercase() {
95  if [[ -z "${1}" ]]; then
96    die "Nothing to convert to lowercase. No argument given."
97  fi
98  echo "${1}" | tr '[:upper:]' '[:lower:]'
99}
100
101check_global_vars() {
102  # Check container type
103  if ! [[ ${CONTAINER_TYPE} == "cpu" ]] && \
104     ! [[ ${CONTAINER_TYPE} == "rocm" ]] && \
105     ! [[ ${CONTAINER_TYPE} == "gpu" ]]; then
106    die "Error: Provided CONTAINER_TYPE \"${CONTAINER_TYPE}\" "\
107        "is not supported."
108  fi
109  # Check OS type
110  if ! [[ ${OS_TYPE} == "ubuntu" ]] && \
111     ! [[ ${OS_TYPE} == "macos" ]]; then
112    die "Error: Provided OS_TYPE \"${OS_TYPE}\" is not supported."
113  fi
114  # Check build flags
115  if [[ -z ${TF_BUILD_FLAGS} ]]; then
116    die "Error: TF_BUILD_FLAGS is not specified."
117  fi
118  # Check test flags
119  if [[ -z ${TF_TEST_FLAGS} ]]; then
120    die "Error: TF_TEST_FLAGS is not specified."
121  fi
122  # Check test filter tags
123  if [[ -z ${TF_TEST_FILTER_TAGS} ]]; then
124    die "Error: TF_TEST_FILTER_TAGS is not specified."
125  fi
126  # Check test targets
127  if [[ -z ${TF_TEST_TARGETS} ]]; then
128    die "Error: TF_TEST_TARGETS is not specified."
129  fi
130  # Check nightly status
131  if [[ -z ${IS_NIGHTLY} ]]; then
132    die "Error: IS_NIGHTLY is not specified."
133  fi
134}
135
136add_test_filter_tag() {
137  EMPTY=""
138  while true; do
139    FILTER="${1:$EMPTY}"
140    if ! [[ $TF_TEST_FILTER_TAGS == *"${FILTER}"* ]]; then
141      TF_TEST_FILTER_TAGS="${FILTER},${TF_TEST_FILTER_TAGS}"
142    fi
143    shift
144    if [[ -z "${1}" ]]; then
145      break
146    fi
147  done
148}
149
150remove_test_filter_tag() {
151  EMPTY=""
152  while true; do
153    FILTER="${1:$EMPTY}"
154    TF_TEST_FILTER_TAGS="$(echo ${TF_TEST_FILTER_TAGS} | sed -e 's/^'${FILTER}',//g' -e 's/,'${FILTER}'//g')"
155    shift
156    if [[ -z "${1}" ]]; then
157      break
158    fi
159  done
160}
161
162# Clean up bazel build & test flags with proper configuration.
163update_bazel_flags() {
164  # Add git tag override flag if necessary.
165  GIT_TAG_STR=" --action_env=GIT_TAG_OVERRIDE"
166  if [[ -z "${GIT_TAG_OVERRIDE}" ]] && \
167    ! [[ ${TF_BUILD_FLAGS} = *${GIT_TAG_STR}* ]]; then
168    TF_BUILD_FLAGS+="${GIT_TAG_STR}"
169  fi
170  # Clean up whitespaces
171  TF_BUILD_FLAGS=$(str_strip "${TF_BUILD_FLAGS}")
172  TF_TEST_FLAGS=$(str_strip "${TF_TEST_FLAGS}")
173  # Cleaned bazel flags
174  echo "Bazel build flags (cleaned):\n" "${TF_BUILD_FLAGS}"
175  echo "Bazel test flags (cleaned):\n" "${TF_TEST_FLAGS}"
176}
177
178update_test_filter_tags() {
179  # Add test filter tags
180  # This script is for validating built PIP packages. Add pip tags.
181  add_test_filter_tag -no_pip -nopip
182  # MacOS filter tags
183  if [[ ${OS_TYPE} == "macos" ]]; then
184    remove_test_filter_tag nomac no_mac
185    add_test_filter_tag -nomac -no_mac
186  fi
187  echo "Final test filter tags: ${TF_TEST_FILTER_TAGS}"
188}
189
190# Check currently running python and pip version
191check_python_pip_version() {
192  # Check if only the major version of python is provided by the user.
193  MAJOR_VER_ONLY=0
194  if [[ ${#PYTHON_VER} -lt 9 ]]; then
195    # User only provided major version (e.g. 'python3' instead of 'python3.7')
196    MAJOR_VER_ONLY=1
197  fi
198
199  # Retrieve only the version number of the user requested python.
200  PYTHON_VER_REQUESTED=${PYTHON_VER:6:3}
201  echo "PYTHON_VER_REQUESTED: ${PYTHON_VER_REQUESTED}"
202
203  # Retrieve only the version numbers of the python & pip in use currently.
204  PYTHON_VER_IN_USE=$(python --version 2>&1)
205  PYTHON_VER_IN_USE=${PYTHON_VER_IN_USE:7:3}
206  PIP_VER_IN_USE=$(${PIP_BIN_PATH} --version)
207  PIP_VER_IN_USE=${PIP_VER_IN_USE:${#PIP_VER_IN_USE}-4:3}
208
209  # If only major versions are applied, drop minor versions.
210  if [[ $MAJOR_VER_ONLY == 1 ]]; then
211    PYTHON_VER_IN_USE=${PYTHON_VER_IN_USE:0:1}
212    PIP_VER_IN_USE=${PIP_VER_IN_USE:0:1}
213  fi
214
215  # Check if all versions match.
216  echo -e "User requested python version: '${PYTHON_VER_REQUESTED}'\n" \
217    "Detected python version in use: '${PYTHON_VER_IN_USE}'\n"\
218    "Detected pip version in use: '${PIP_VER_IN_USE}'"
219  if ! [[ $PYTHON_VER_REQUESTED == $PYTHON_VER_IN_USE ]]; then
220    die "Error: Mismatch in python versions detected."
221  else:
222    echo "Python and PIP versions in use match the requested."
223  fi
224}
225
226###########################################################################
227# Setup: directories, local/global variables
228###########################################################################
229
230# Script directory and source necessary files.
231SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
232source "${SCRIPT_DIR}/builds_common.sh"
233
234# Required global variables
235# Checks on values for these vars are done in "Build TF PIP Package" section.
236CONTAINER_TYPE=$(lowercase "${CONTAINER_TYPE}")
237OS_TYPE=$(lowercase "${OS_TYPE}")
238PYTHON_VER=$(lowercase "${TF_PYTHON_VERSION}")
239
240# Python bin path
241if [[ -z "$PYTHON_BIN_PATH" ]]; then
242  die "Error: PYTHON_BIN_PATH was not provided. Did you run configure?"
243fi
244
245# Set optional environment variables; set to default in case not user defined.
246DEFAULT_PIP_TESTS="" # Do not run any tests by default
247DEFAULT_PROJECT_NAME="tensorflow"
248DEFAULT_PIP_TEST_ROOT="pip_test"
249DEFAULT_BUILD_BOTH_GPU_PACKAGES=0
250DEFAULT_BUILD_BOTH_CPU_PACKAGES=0
251# Take in optional global variables
252PIP_TESTS=${TF_PIP_TESTS:-$DEFAULT_PIP_TESTS}
253PROJECT_NAME=${TF_PROJECT_NAME:-$DEFAULT_PROJECT_NAME}
254PIP_TEST_ROOT=${TF_PIP_TEST_ROOT:-$DEFAULT_PIP_TEST_ROOT}
255BUILD_BOTH_GPU_PACKAGES=${TF_BUILD_BOTH_GPU_PACKAGES:-$DEFAULT_BUILD_BOTH_GPU_PACKAGES}
256BUILD_BOTH_CPU_PACKAGES=${TF_BUILD_BOTH_CPU_PACKAGES:-$DEFAULT_BUILD_BOTH_CPU_PACKAGES}
257
258# Local variables
259PIP_WHL_DIR="${KOKORO_ARTIFACTS_DIR}/tensorflow/${PIP_TEST_ROOT}/whl"
260mkdir -p "${PIP_WHL_DIR}"
261PIP_WHL_DIR=$(realpath "${PIP_WHL_DIR}") # Get absolute path
262WHL_PATH=""
263# Determine the major.minor versions of python being used (e.g., 3.7).
264# Useful for determining the directory of the local pip installation.
265PY_MAJOR_MINOR_VER=$(${PYTHON_BIN_PATH} -c "print(__import__('sys').version)" 2>&1 | awk '{ print $1 }' | head -n 1 | cut -c1-3)
266
267if [[ -z "${PY_MAJOR_MINOR_VER}" ]]; then
268  die "ERROR: Unable to determine the major.minor version of Python."
269fi
270echo "Python binary path to be used in PIP install: ${PYTHON_BIN_PATH} "\
271"(Major.Minor version: ${PY_MAJOR_MINOR_VER})"
272PYTHON_BIN_PATH_INIT=${PYTHON_BIN_PATH}
273PIP_BIN_PATH="$(which pip${PY_MAJOR_MINOR_VER})"
274
275# PIP packages
276INSTALL_EXTRA_PIP_PACKAGES="h5py portpicker scipy scikit-learn ${TF_BUILD_INSTALL_EXTRA_PIP_PACKAGES}"
277
278###########################################################################
279# Build TF PIP Package
280###########################################################################
281
282# First remove any already existing binaries for a clean start and test.
283if [[ -d ${PIP_TEST_ROOT} ]]; then
284  echo "Test root directory ${PIP_TEST_ROOT} already exists. Deleting it."
285  sudo rm -rf ${PIP_TEST_ROOT}
286fi
287
288# Check that global variables are properly set.
289check_global_vars
290
291# Check if in a virtualenv and exit if yes.
292IN_VENV=$(python -c 'import sys; print("1" if sys.version_info.major == 3 and sys.prefix != sys.base_prefix else "0")')
293if [[ "$IN_VENV" == "1" ]]; then
294  echo "It appears that we are already in a virtualenv. Deactivating..."
295  deactivate || source deactivate || die "FAILED: Unable to deactivate from existing virtualenv."
296fi
297
298# Obtain the path to python binary as written by ./configure if it was run.
299if [[ -e tools/python_bin_path.sh ]]; then
300  source tools/python_bin_path.sh
301fi
302# Assume PYTHON_BIN_PATH is exported by the script above or the caller.
303if [[ -z "$PYTHON_BIN_PATH" ]]; then
304  die "PYTHON_BIN_PATH was not provided. Did you run configure?"
305fi
306
307# Bazel build the file.
308PIP_BUILD_TARGET="//tensorflow/tools/pip_package:build_pip_package"
309# Clean bazel cache.
310bazel clean
311# Clean up and update bazel flags
312update_bazel_flags
313# Build. This outputs the file `build_pip_package`.
314bazel build \
315  --action_env=PYTHON_BIN_PATH=${PYTHON_BIN_PATH} \
316  ${TF_BUILD_FLAGS} \
317  ${PIP_BUILD_TARGET} \
318  || die "Error: Bazel build failed for target: '${PIP_BUILD_TARGET}'"
319
320###########################################################################
321# Test function(s)
322###########################################################################
323
324test_pip_virtualenv() {
325  # Get args
326  VENV_DIR_NAME=$1
327  shift
328  TEST_TYPE_FLAG=$1
329
330  # Check test type args
331  if ! [[ ${TEST_TYPE_FLAG} == "--oss_serial" ]] && \
332     ! [[ ${TEST_TYPE_FLAG} == "--clean" ]] && \
333     ! [[ ${TEST_TYPE_FLAG} == "" ]]; then
334     die "Error: Wrong test type given. TEST_TYPE_FLAG=${TEST_TYPE_FLAG}"
335  fi
336
337  # Create virtualenv directory for test
338  VENV_DIR="${PIP_TEST_ROOT}/${VENV_DIR_NAME}"
339
340  # Activate virtualenv
341  create_activate_virtualenv ${TEST_TYPE_FLAG} ${VENV_DIR}
342  # Install TF with pip
343  TIME_START=$SECONDS
344  install_tensorflow_pip "${WHL_PATH}"
345  TIME_ELAPSED=$(($SECONDS - $TIME_START))
346  echo "Time elapsed installing tensorflow = ${TIME_ELAPSED} seconds"
347
348  # cd to a temporary directory to avoid picking up Python files in the source
349  # tree.
350  TMP_DIR=$(mktemp -d)
351  pushd "${TMP_DIR}"
352
353  # Run a quick check on tensorflow installation.
354  RET_VAL=$(python -c "import tensorflow as tf; t1=tf.constant([1,2,3,4]); t2=tf.constant([5,6,7,8]); print(tf.add(t1,t2).shape)")
355
356  # Return to original directory. Remove temp dirs.
357  popd
358  sudo rm -rf "${TMP_DIR}"
359
360  # Check result to see if tensorflow is properly installed.
361  if ! [[ ${RET_VAL} == *'(4,)'* ]]; then
362    echo "PIP test on virtualenv (non-clean) FAILED"
363    return 1
364  fi
365
366  # Install extra pip packages, if specified.
367  for PACKAGE in ${INSTALL_EXTRA_PIP_PACKAGES}; do
368    echo "Installing extra pip package required by test-on-install: ${PACKAGE}"
369
370    ${PIP_BIN_PATH} install ${PACKAGE}
371    if [[ $? != 0 ]]; then
372      echo "${PIP_BIN_PATH} install ${PACKAGE} FAILED."
373      deactivate || source deactivate || die "FAILED: Unable to deactivate from existing virtualenv."
374      return 1
375    fi
376  done
377
378  # Run bazel test.
379  run_test_with_bazel ${TEST_TYPE_FLAG}
380  RESULT=$?
381
382  # Deactivate from virtualenv.
383  deactivate || source deactivate || die "FAILED: Unable to deactivate from existing virtualenv."
384  sudo rm -rf "${VENV_DIR}"
385
386  return $RESULT
387}
388
389###########################################################################
390# Test helper function(s)
391###########################################################################
392
393create_activate_virtualenv() {
394  VIRTUALENV_FLAGS="--system-site-packages"
395  if [[ "${1}" == "--clean" ]]; then
396    VIRTUALENV_FLAGS=""
397    shift
398  elif [[ "${1}" == "--oss_serial" ]]; then
399    shift
400  fi
401
402  VIRTUALENV_DIR="${1}"
403  if [[ -d "${VIRTUALENV_DIR}" ]]; then
404    if sudo rm -rf "${VIRTUALENV_DIR}"
405    then
406      echo "Removed existing virtualenv directory: ${VIRTUALENV_DIR}"
407    else
408      die "Failed to remove existing virtualenv directory: ${VIRTUALENV_DIR}"
409    fi
410  fi
411
412  if mkdir -p "${VIRTUALENV_DIR}"
413  then
414    echo "Created virtualenv directory: ${VIRTUALENV_DIR}"
415  else
416    die "FAILED to create virtualenv directory: ${VIRTUALENV_DIR}"
417  fi
418
419  # Use the virtualenv from the default python version (i.e., python-virtualenv)
420  # to create the virtualenv directory for testing. Use the -p flag to specify
421  # the python version inside the to-be-created virtualenv directory.
422  ${PYTHON_BIN_PATH_INIT} -m virtualenv -p ${PYTHON_BIN_PATH_INIT} ${VIRTUALENV_FLAGS} ${VIRTUALENV_DIR} || \
423    ${PYTHON_BIN_PATH_INIT} -m venv ${VIRTUALENV_DIR} || \
424    die "FAILED: Unable to create virtualenv"
425
426  source "${VIRTUALENV_DIR}/bin/activate" || \
427    die "FAILED: Unable to activate virtualenv in ${VIRTUALENV_DIR}"
428
429  # Update .tf_configure.bazelrc with venv python path for bazel test.
430  PYTHON_BIN_PATH="$(which python)"
431  yes "" | ./configure
432}
433
434install_tensorflow_pip() {
435  if [[ -z "${1}" ]]; then
436    die "Please provide a proper wheel file path."
437  fi
438
439  # Set path to pip.
440  PIP_BIN_PATH="${PYTHON_BIN_PATH} -m pip"
441
442  # Print python and pip bin paths
443  echo "PYTHON_BIN_PATH to be used to install the .whl: ${PYTHON_BIN_PATH}"
444  echo "PIP_BIN_PATH to be used to install the .whl: ${PIP_BIN_PATH}"
445
446  # Upgrade pip so it supports tags such as cp27mu, manylinux2010 etc.
447  echo "Upgrade pip in virtualenv"
448
449  # NOTE: pip install --upgrade pip leads to a documented TLS issue for
450  # some versions in python
451  curl https://bootstrap.pypa.io/get-pip.py | ${PYTHON_BIN_PATH} || \
452    die "Error: pip install (get-pip.py) FAILED"
453
454  # Check that requested python version matches configured one.
455  check_python_pip_version
456
457  # Force upgrade of setuptools. We need it to install pips using
458  # `install_requires` notation introduced in setuptools >=20.5. The default
459  # version of setuptools is 5.5.1.
460  ${PIP_BIN_PATH} install --upgrade setuptools || \
461    die "Error: setuptools install, upgrade FAILED"
462
463  # Force tensorflow reinstallation. Otherwise it may not get installed from
464  # last build if it had the same version number as previous build.
465  PIP_FLAGS="--upgrade --force-reinstall"
466  ${PIP_BIN_PATH} install ${PIP_FLAGS} ${WHL_PATH} || \
467    die "pip install (forcing to reinstall tensorflow) FAILED"
468  echo "Successfully installed pip package ${WHL_PATH}"
469
470  # Install the future package in the virtualenv. Installing it in user system
471  # packages does not appear to port it over when creating a virtualenv.
472  #   ImportError: No module named builtins
473  ${PIP_BIN_PATH} install --upgrade "future>=0.17.1" || \
474    die "Error: future install, upgrade FAILED"
475
476  # Install the gast package in the virtualenv. Installing it in user system
477  # packages does not appear to port it over when creating a virtualenv.
478  ${PIP_BIN_PATH} install --upgrade "gast==0.4.0" || \
479    die "Error: gast install, upgrade FAILED"
480
481}
482
483run_test_with_bazel() {
484  IS_OSS_SERIAL=0
485  if [[ "${1}" == "--oss_serial" ]]; then
486    IS_OSS_SERIAL=1
487  fi
488  TF_GPU_COUNT=${TF_GPU_COUNT:-4}
489
490  # PIP tests should have a "different" path. Different than the one we place
491  # virtualenv, because we are deleting and recreating it here.
492  PIP_TEST_PREFIX=bazel_pip
493  TEST_ROOT=$(pwd)/${PIP_TEST_PREFIX}
494  sudo rm -rf $TEST_ROOT
495  mkdir -p $TEST_ROOT
496  ln -s $(pwd)/tensorflow $TEST_ROOT/tensorflow
497
498  if [[ "${IS_OSS_SERIAL}" == "1" ]]; then
499    remove_test_filter_tag -no_oss
500    add_test_filter_tag oss_serial
501  else
502    add_test_filter_tag -oss_serial
503  fi
504
505  # Clean the bazel cache
506  bazel clean
507  # Clean up flags before running bazel commands
508  update_bazel_flags
509  # Clean up and update test filter tags
510  update_test_filter_tags
511
512  # Figure out how many concurrent tests we can run and do run the tests.
513  BAZEL_PARALLEL_TEST_FLAGS=""
514  if [[ $CONTAINER_TYPE == "gpu" ]] || [[ $CONTAINER_TYPE == "rocm" ]]; then
515    # Number of test threads is the number of GPU cards available.
516    if [[ $OS_TYPE == "macos" ]]; then
517      BAZEL_PARALLEL_TEST_FLAGS="--local_test_jobs=1"
518    else
519      BAZEL_PARALLEL_TEST_FLAGS="--local_test_jobs=${TF_GPU_COUNT} \
520        --run_under=//tensorflow/tools/ci_build/gpu_build:parallel_gpu_execute"
521    fi
522  else
523    # Number of test threads is the number of physical CPUs.
524    if [[ $OS_TYPE == "macos" ]]; then
525      BAZEL_PARALLEL_TEST_FLAGS="--local_test_jobs=$(sysctl -n hw.ncpu)"
526    else
527      BAZEL_PARALLEL_TEST_FLAGS="--local_test_jobs=$(grep -c ^processor /proc/cpuinfo)"
528    fi
529  fi
530
531  if [[ ${IS_OSS_SERIAL} == 1 ]]; then
532    BAZEL_PARALLEL_TEST_FLAGS="--local_test_jobs=1"
533  fi
534
535  TEST_TARGETS_SYMLINK=""
536  for TARGET in ${TF_TEST_TARGETS[@]}; do
537    TARGET_NEW=$(echo ${TARGET} | sed -e "s/\/\//\/\/${PIP_TEST_PREFIX}\//g")
538    TEST_TARGETS_SYMLINK+="${TARGET_NEW} "
539  done
540  echo "Test targets (symlink): ${TEST_TARGETS_SYMLINK}"
541
542  # Run the test.
543  bazel test --build_tests_only ${TF_TEST_FLAGS} ${BAZEL_PARALLEL_TEST_FLAGS} --test_tag_filters=${TF_TEST_FILTER_TAGS} -k -- ${TEST_TARGETS_SYMLINK}
544
545  unlink ${TEST_ROOT}/tensorflow
546}
547
548run_all_tests() {
549  if [[ -z "${PIP_TESTS}" ]]; then
550    echo "No test was specified to run. Skipping all tests."
551    return 0
552  fi
553  FAIL_COUNTER=0
554  PASS_COUNTER=0
555  for TEST in ${PIP_TESTS[@]}; do
556
557    # Run tests.
558    case "${TEST}" in
559    "test_pip_virtualenv_clean")
560      test_pip_virtualenv venv_clean --clean
561      ;;
562    "test_pip_virtualenv_non_clean")
563      test_pip_virtualenv venv
564      ;;
565    "test_pip_virtualenv_oss_serial")
566      test_pip_virtualenv venv_oss --oss_serial
567      ;;
568    *)
569      die "No matching test ${TEST} was found. Stopping test."
570      ;;
571    esac
572
573    # Check and update the results.
574    RETVAL=$?
575
576    # Update results counter
577    if [ ${RETVAL} -eq 0 ]; then
578      echo "Test (${TEST}) PASSED. (PASS COUNTER: ${PASS_COUNTER})"
579      PASS_COUNTER=$(($PASS_COUNTER+1))
580    else
581      echo "Test (${TEST}) FAILED. (FAIL COUNTER: ${FAIL_COUNTER})"
582      FAIL_COUNTER=$(($FAIL_COUNTER+1))
583    fi
584  done
585  printf "${PASS_COUNTER} PASSED | ${FAIL_COUNTER} FAILED"
586  if [[ "${FAIL_COUNTER}" == "0" ]]; then
587    printf "PIP tests ${COLOR_GREEN}PASSED${COLOR_NC}\n"
588    return 0
589  else:
590    printf "PIP tests ${COLOR_RED}FAILED${COLOR_NC}\n"
591    return 1
592  fi
593}
594
595###########################################################################
596# Build TF PIP Wheel file
597###########################################################################
598
599# Update the build flags for building whl.
600# Flags: GPU, OS, tf_nightly, project name
601GPU_FLAG=""
602NIGHTLY_FLAG=""
603
604# TF Nightly flag
605if [[ "$IS_NIGHTLY" == 1 ]]; then
606  # If 'nightly' is not specified in the project name already, then add.
607  if ! [[ $PROJECT_NAME == *"nightly"* ]]; then
608    echo "WARNING: IS_NIGHTLY=${IS_NIGHTLY} but requested project name \
609    (PROJECT_NAME=${PROJECT_NAME}) does not include 'nightly' string. \
610    Renaming it to 'tf_nightly'."
611    PROJECT_NAME="tf_nightly"
612  fi
613  NIGHTLY_FLAG="--nightly_flag"
614fi
615
616# CPU / GPU flag
617if [[ ${CONTAINER_TYPE} == "gpu" ]]; then
618  GPU_FLAG="--gpu"
619  if ! [[ $PROJECT_NAME == *"gpu"* ]]; then
620    # Only update PROJECT_NAME if TF_PROJECT_NAME is not set
621    if [[ -z "${TF_PROJECT_NAME}" ]]; then
622      echo "WARNING: GPU is specified but requested project name (PROJECT_NAME=${PROJECT_NAME}) \
623      does not include 'gpu'. Appending '_gpu' to the project name."
624      PROJECT_NAME="${PROJECT_NAME}_gpu"
625    fi
626  fi
627fi
628
629if [[ ${CONTAINER_TYPE} == "rocm" ]]; then
630  GPU_FLAG="--rocm"
631  if ! [[ $PROJECT_NAME == *"rocm"* ]]; then
632    # Only update PROJECT_NAME if TF_PROJECT_NAME is not set
633    if [[ -z "${TF_PROJECT_NAME}" ]]; then
634      echo "WARNING: ROCM is specified but requested project name (PROJECT_NAME=${PROJECT_NAME}) \
635      does not include 'rocm'. Appending '_rocm' to the project name."
636      PROJECT_NAME="${PROJECT_NAME}_rocm"
637    fi
638  fi
639fi
640
641./bazel-bin/tensorflow/tools/pip_package/build_pip_package ${PIP_WHL_DIR} ${GPU_FLAG} ${NIGHTLY_FLAG} "--project_name" ${PROJECT_NAME} || die "build_pip_package FAILED"
642
643PY_DOTLESS_MAJOR_MINOR_VER=$(echo $PY_MAJOR_MINOR_VER | tr -d '.')
644if [[ $PY_DOTLESS_MAJOR_MINOR_VER == "2" ]]; then
645  PY_DOTLESS_MAJOR_MINOR_VER="27"
646fi
647
648# Set wheel path and verify that there is only one .whl file in the path.
649WHL_PATH=$(ls "${PIP_WHL_DIR}"/"${PROJECT_NAME}"-*"${PY_DOTLESS_MAJOR_MINOR_VER}"*"${PY_DOTLESS_MAJOR_MINOR_VER}"*.whl)
650if [[ $(echo "${WHL_PATH}" | wc -w) -ne 1 ]]; then
651  echo "ERROR: Failed to find exactly one built TensorFlow .whl file in "\
652  "directory: ${PIP_WHL_DIR}"
653fi
654
655WHL_DIR=$(dirname "${WHL_PATH}")
656
657# Print the size of the wheel file.
658echo "Size of the PIP wheel file built: $(ls -l ${WHL_PATH} | awk '{print $5}')"
659
660# Build the other GPU package.
661if [[ "$BUILD_BOTH_GPU_PACKAGES" -eq "1" ]] || [[ "$BUILD_BOTH_CPU_PACKAGES" -eq "1" ]]; then
662
663  if [[ "$BUILD_BOTH_GPU_PACKAGES" -eq "1" ]] && [[ "$BUILD_BOTH_CPU_PACKAGES" -eq "1" ]]; then
664    die "ERROR: TF_BUILD_BOTH_GPU_PACKAGES and TF_BUILD_BOTH_GPU_PACKAGES cannot both be set. No additional package will be built."
665  fi
666
667  echo "====================================="
668  if [[ "$BUILD_BOTH_GPU_PACKAGES" -eq "1" ]]; then
669    if ! [[ ${OS_TYPE} == "ubuntu" ]]; then
670      die "ERROR: pip_new.sh only support building both GPU wheels on ubuntu."
671    fi
672    echo "Building the other GPU pip package."
673    PROJECT_SUFFIX="gpu"
674  else
675    if ! [[ ${OS_TYPE} == "macos" ]]; then
676      die "ERROR: pip_new.sh only support building both CPU wheels on macos."
677    fi
678    echo "Building the other CPU pip package."
679    PROJECT_SUFFIX="cpu"
680  fi
681
682  # Check container type
683  if ! [[ ${CONTAINER_TYPE} == ${PROJECT_SUFFIX} ]]; then
684    die "Error: CONTAINER_TYPE needs to be \"${PROJECT_SUFFIX}\" to build ${PROJECT_SUFFIX} packages. Got"\
685        "\"${CONTAINER_TYPE}\" instead."
686  fi
687  if [[ "$PROJECT_NAME" == *_${PROJECT_SUFFIX} ]]; then
688    NEW_PROJECT_NAME=${PROJECT_NAME}"_${PROJECT_SUFFIX}"
689  else
690    NEW_PROJECT_NAME="${PROJECT_NAME}_${PROJECT_SUFFIX}"
691  fi
692  echo "The given ${PROJECT_SUFFIX} \$PROJECT_NAME is ${PROJECT_NAME}. The additional ${PROJECT_SUFFIX}"\
693  "pip package will have project name ${NEW_PROJECT_NAME}."
694
695  ./bazel-bin/tensorflow/tools/pip_package/build_pip_package ${PIP_WHL_DIR} ${GPU_FLAG} ${NIGHTLY_FLAG} "--project_name" ${NEW_PROJECT_NAME} || die "build_pip_package FAILED"
696fi
697
698# Run tests (if any is specified).
699run_all_tests
700
701
702if [[ ${OS_TYPE} == "ubuntu" ]] && \
703   ! [[ ${CONTAINER_TYPE} == "rocm" ]] ; then
704  # Avoid Python3.6 abnormality by installing auditwheel here.
705  set +e
706  pip3 show auditwheel || "pip${PY_MAJOR_MINOR_VER}" show auditwheel
707  pip3 install auditwheel==2.0.0 || "pip${PY_MAJOR_MINOR_VER}" install auditwheel==2.0.0
708  sudo pip3 install auditwheel==2.0.0 || \
709    sudo "pip${PY_MAJOR_MINOR_VER}" install auditwheel==2.0.0
710  set -e
711  auditwheel --version
712
713  for WHL_PATH in $(ls ${PIP_WHL_DIR}/*.whl); do
714    # Repair the wheels for cpu manylinux2010
715    echo "auditwheel repairing ${WHL_PATH}"
716    auditwheel repair --plat manylinux2010_x86_64 -w "${WHL_DIR}" "${WHL_PATH}"
717
718    WHL_BASE_NAME=$(basename "${WHL_PATH}")
719    AUDITED_WHL_NAME="${WHL_DIR}"/$(echo "${WHL_BASE_NAME//linux/manylinux2010}")
720    if [[ -f ${AUDITED_WHL_NAME} ]]; then
721      WHL_PATH=${AUDITED_WHL_NAME}
722      echo "Repaired manylinux2010 wheel file at: ${WHL_PATH}"
723    else
724      die "WARNING: Cannot find repaired wheel."
725    fi
726  done
727fi
728echo "EOF: Successfully ran pip_new.sh"
729