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