1#!/usr/bin/env bash 2# Copyright 2015 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 17 18set -e 19 20function is_absolute { 21 [[ "$1" = /* ]] || [[ "$1" =~ ^[a-zA-Z]:[/\\].* ]] 22} 23 24function real_path() { 25 is_absolute "$1" && echo "$1" || echo "$PWD/${1#./}" 26} 27 28function cp_external() { 29 local src_dir=$1 30 local dest_dir=$2 31 32 pushd . 33 cd "$src_dir" 34 for f in `find . ! -type d ! -name '*.py' ! -path '*local_config_cuda*' ! -path '*local_config_tensorrt*' ! -path '*local_config_syslibs*' ! -path '*org_tensorflow*' ! -path '*llvm-project/llvm/*'`; do 35 mkdir -p "${dest_dir}/$(dirname ${f})" 36 cp "${f}" "${dest_dir}/$(dirname ${f})/" 37 done 38 popd 39 40 mkdir -p "${dest_dir}/local_config_cuda/cuda/cuda/" 41 cp "${src_dir}/local_config_cuda/cuda/cuda/cuda_config.h" "${dest_dir}/local_config_cuda/cuda/cuda/" 42} 43 44function copy_xla_aot_runtime_sources() { 45 local src_dir=$1 46 local dst_dir=$2 47 48 local srcs_txt="tensorflow/tools/pip_package/xla_compiled_cpu_runtime_srcs.txt" 49 50 if [ ! -f "${src_dir}/${srcs_txt}" ]; then 51 echo Could not find source list file "${src_dir}/${srcs_txt}". 1>&2 52 return 0 53 fi 54 55 pushd $src_dir 56 for file in $(cat "${srcs_txt}") 57 do 58 # Sometimes $file has a prefix bazel-out/host/ we want to remove. 59 prefix=${file%%tensorflow/*} # Find the location of "tensorflow/*" 60 candidate_file=${file#$prefix} # Remove the prefix 61 if [ ! -z "$candidate_file" ]; then 62 file=$candidate_file 63 fi 64 dn=$(dirname $file) 65 if test -f "$file"; then 66 mkdir -p "${dst_dir}/${dn}" 67 cp $file "${dst_dir}/${file}" 68 else 69 echo "Missing xla source file: ${file}" 1>&2 70 fi 71 done 72 cp tensorflow/tools/pip_package/xla_build/CMakeLists.txt "${dst_dir}" 73 popd 74} 75 76function move_to_root_if_exists () { 77 arg_to_move="$1" 78 if [ -e "${arg_to_move}" ]; then 79 mv ${arg_to_move} ./ 80 fi 81} 82 83function reorganize_includes() { 84 TMPDIR="${1%/}" 85 pushd "${TMPDIR}/tensorflow/include/" 86 87 move_to_root_if_exists external/com_google_absl/absl 88 89 move_to_root_if_exists external/eigen_archive/Eigen 90 move_to_root_if_exists external/eigen_archive/unsupported 91 92 move_to_root_if_exists external/jsoncpp_git/include 93 rm -rf external/jsoncpp_git 94 95 move_to_root_if_exists external/com_google_protobuf/src/google 96 rm -rf external/com_google_protobuf/python 97 98 popd 99} 100 101PLATFORM="$(uname -s | tr 'A-Z' 'a-z')" 102function is_windows() { 103 if [[ "${PLATFORM}" =~ (cygwin|mingw32|mingw64|msys)_nt* ]]; then 104 true 105 else 106 false 107 fi 108} 109 110function prepare_src() { 111 if [ $# -lt 1 ] ; then 112 echo "No destination dir provided" 113 exit 1 114 fi 115 116 TMPDIR="${1%/}" 117 mkdir -p "$TMPDIR" 118 EXTERNAL_INCLUDES="${TMPDIR}/tensorflow/include/external" 119 XLA_AOT_RUNTIME_SOURCES="${TMPDIR}/tensorflow/xla_aot_runtime_src" 120 121 echo $(date) : "=== Preparing sources in dir: ${TMPDIR}" 122 123 if [ ! -d bazel-bin/tensorflow ]; then 124 echo "Could not find bazel-bin. Did you run from the root of the build tree?" 125 exit 1 126 fi 127 128 if is_windows; then 129 rm -rf ./bazel-bin/tensorflow/tools/pip_package/simple_console_for_window_unzip 130 mkdir -p ./bazel-bin/tensorflow/tools/pip_package/simple_console_for_window_unzip 131 echo "Unzipping simple_console_for_windows.zip to create runfiles tree..." 132 unzip -o -q ./bazel-bin/tensorflow/tools/pip_package/simple_console_for_windows.zip -d ./bazel-bin/tensorflow/tools/pip_package/simple_console_for_window_unzip 133 echo "Unzip finished." 134 # runfiles structure after unzip the python binary 135 cp -L \ 136 bazel-bin/tensorflow/tools/pip_package/simple_console_for_window_unzip/runfiles/org_tensorflow/LICENSE \ 137 "${TMPDIR}" 138 cp -LR \ 139 bazel-bin/tensorflow/tools/pip_package/simple_console_for_window_unzip/runfiles/org_tensorflow/tensorflow \ 140 "${TMPDIR}" 141 cp_external \ 142 bazel-bin/tensorflow/tools/pip_package/simple_console_for_window_unzip/runfiles \ 143 "${EXTERNAL_INCLUDES}/" 144 copy_xla_aot_runtime_sources \ 145 bazel-bin/tensorflow/tools/pip_package/simple_console_for_window_unzip/runfiles/org_tensorflow \ 146 "${XLA_AOT_RUNTIME_SOURCES}/" 147 RUNFILES=bazel-bin/tensorflow/tools/pip_package/simple_console_for_window_unzip/runfiles/org_tensorflow 148 # If oneDNN was built with openMP then copy the omp libs over 149 if [ -f "bazel-bin/external/llvm_openmp/libiomp5md.dll" ]; then 150 cp bazel-bin/external/llvm_openmp/libiomp5md.dll ${TMPDIR}/tensorflow/python 151 cp bazel-bin/external/llvm_openmp/libiomp5md.dll.if.lib ${TMPDIR}/tensorflow/python 152 fi 153 else 154 RUNFILES=bazel-bin/tensorflow/tools/pip_package/build_pip_package.runfiles/org_tensorflow 155 if [ -d bazel-bin/tensorflow/tools/pip_package/build_pip_package.runfiles/org_tensorflow/external ]; then 156 # Old-style runfiles structure (--legacy_external_runfiles). 157 cp -L \ 158 bazel-bin/tensorflow/tools/pip_package/build_pip_package.runfiles/org_tensorflow/LICENSE \ 159 "${TMPDIR}" 160 cp -LR \ 161 bazel-bin/tensorflow/tools/pip_package/build_pip_package.runfiles/org_tensorflow/tensorflow \ 162 "${TMPDIR}" 163 cp_external \ 164 bazel-bin/tensorflow/tools/pip_package/build_pip_package.runfiles/org_tensorflow/external \ 165 "${EXTERNAL_INCLUDES}" 166 copy_xla_aot_runtime_sources \ 167 bazel-bin/tensorflow/tools/pip_package/build_pip_package.runfiles/org_tensorflow \ 168 "${XLA_AOT_RUNTIME_SOURCES}" 169 # Copy MKL libs over so they can be loaded at runtime 170 so_lib_dir=$(ls $RUNFILES | grep solib) || true 171 if [ -n "${so_lib_dir}" ]; then 172 mkl_so_dir=$(ls ${RUNFILES}/${so_lib_dir} | grep mkl) || true 173 if [ -n "${mkl_so_dir}" ]; then 174 mkdir "${TMPDIR}/${so_lib_dir}" 175 cp -R ${RUNFILES}/${so_lib_dir}/${mkl_so_dir} "${TMPDIR}/${so_lib_dir}" 176 fi 177 fi 178 else 179 # New-style runfiles structure (--nolegacy_external_runfiles). 180 cp -L \ 181 bazel-bin/tensorflow/tools/pip_package/build_pip_package.runfiles/org_tensorflow/LICENSE \ 182 "${TMPDIR}" 183 cp -LR \ 184 bazel-bin/tensorflow/tools/pip_package/build_pip_package.runfiles/org_tensorflow/tensorflow \ 185 "${TMPDIR}" 186 cp_external \ 187 bazel-bin/tensorflow/tools/pip_package/build_pip_package.runfiles \ 188 "${EXTERNAL_INCLUDES}" 189 copy_xla_aot_runtime_sources \ 190 bazel-bin/tensorflow/tools/pip_package/build_pip_package.runfiles/org_tensorflow \ 191 "${XLA_AOT_RUNTIME_SOURCES}" 192 # Copy MKL libs over so they can be loaded at runtime 193 so_lib_dir=$(ls $RUNFILES | grep solib) || true 194 if [ -n "${so_lib_dir}" ]; then 195 mkl_so_dir=$(ls ${RUNFILES}/${so_lib_dir} | grep mkl) || true 196 if [ -n "${mkl_so_dir}" ]; then 197 mkdir "${TMPDIR}/${so_lib_dir}" 198 cp -R ${RUNFILES}/${so_lib_dir}/${mkl_so_dir} "${TMPDIR}/${so_lib_dir}" 199 fi 200 fi 201 fi 202 fi 203 204 mkdir -p ${TMPDIR}/third_party 205 cp -R $RUNFILES/third_party/eigen3 ${TMPDIR}/third_party 206 207 reorganize_includes "${TMPDIR}" 208 209 cp tensorflow/tools/pip_package/MANIFEST.in ${TMPDIR} 210 cp tensorflow/tools/pip_package/README ${TMPDIR}/README.md 211 cp tensorflow/tools/pip_package/setup.py ${TMPDIR} 212 213 rm -f ${TMPDIR}/tensorflow/libtensorflow_framework.so 214 rm -f ${TMPDIR}/tensorflow/libtensorflow_framework.so.[0-9].* 215 216 # TODO(annarev): copy over API files from tensorflow/api/_vN to tensorflow/ 217 # except tensorflow/api/_vN/lite/. 218 219 # TODO(b/150440817): support autocomplete for tf.keras 220 # Copy over keras API folder to the root directory 221 # so that autocomplete works as expected for all keras subimports. 222 # if [ -d "${TMPDIR}/tensorflow/_api/v1/" ] 223 # then 224 # cp -r ${TMPDIR}/tensorflow/python/keras/api/_v1/keras/ ${TMPDIR}/tensorflow/keras/ 225 # sed -i'.original' -e 's/.python.keras.api._v1/tensorflow/g' ${TMPDIR}/tensorflow/__init__.py 226 # else 227 # cp -r ${TMPDIR}/tensorflow/python/keras/api/_v2/keras/ ${TMPDIR}/tensorflow/keras/ 228 # sed -i'.original' -e 's/.python.keras.api._v2/tensorflow/g' ${TMPDIR}/tensorflow/__init__.py 229 # fi 230} 231 232function build_wheel() { 233 if [ $# -lt 2 ] ; then 234 echo "No src and dest dir provided" 235 exit 1 236 fi 237 238 TMPDIR="$1" 239 DEST="$2" 240 PKG_NAME_FLAG="$3" 241 242 # Before we leave the top-level directory, make sure we know how to 243 # call python. 244 if [[ -e tools/python_bin_path.sh ]]; then 245 source tools/python_bin_path.sh 246 fi 247 248 pushd ${TMPDIR} > /dev/null 249 250 rm -f MANIFEST 251 echo $(date) : "=== Building wheel" 252 "${PYTHON_BIN_PATH:-python}" setup.py bdist_wheel ${PKG_NAME_FLAG} >/dev/null 253 mkdir -p ${DEST} 254 cp dist/* ${DEST} 255 popd > /dev/null 256 echo $(date) : "=== Output wheel file is in: ${DEST}" 257} 258 259function usage() { 260 echo "Usage:" 261 echo "$0 [--src srcdir] [--dst dstdir] [options]" 262 echo "$0 dstdir [options]" 263 echo "" 264 echo " --src prepare sources in srcdir" 265 echo " will use temporary dir if not specified" 266 echo "" 267 echo " --dst build wheel in dstdir" 268 echo " if dstdir is not set do not build, only prepare sources" 269 echo "" 270 echo " Options:" 271 echo " --project_name <name> set project name to name" 272 echo " --cpu build tensorflow_cpu" 273 echo " --gpu build tensorflow_gpu" 274 echo " --gpudirect build tensorflow_gpudirect" 275 echo " --rocm build tensorflow_rocm" 276 echo " --nightly_flag build tensorflow nightly" 277 echo "" 278 exit 1 279} 280 281function main() { 282 PKG_NAME_FLAG="" 283 PROJECT_NAME="" 284 CPU_BUILD=0 285 GPU_BUILD=0 286 GPUDIRECT_BUILD=0 287 ROCM_BUILD=0 288 NIGHTLY_BUILD=0 289 SRCDIR="" 290 DSTDIR="" 291 CLEANSRC=1 292 while true; do 293 if [[ "$1" == "--help" ]]; then 294 usage 295 exit 1 296 elif [[ "$1" == "--nightly_flag" ]]; then 297 NIGHTLY_BUILD=1 298 elif [[ "$1" == "--gpu" ]]; then 299 GPU_BUILD=1 300 elif [[ "$1" == "--cpu" ]]; then 301 CPU_BUILD=1 302 elif [[ "$1" == "--gpudirect" ]]; then 303 GPUDIRECT_BUILD=1 304 elif [[ "$1" == "--rocm" ]]; then 305 ROCM_BUILD=1 306 elif [[ "$1" == "--project_name" ]]; then 307 shift 308 if [[ -z "$1" ]]; then 309 break 310 fi 311 PROJECT_NAME="$1" 312 elif [[ "$1" == "--src" ]]; then 313 shift 314 SRCDIR="$(real_path $1)" 315 CLEANSRC=0 316 elif [[ "$1" == "--dst" ]]; then 317 shift 318 DSTDIR="$(real_path $1)" 319 else 320 DSTDIR="$(real_path $1)" 321 fi 322 shift 323 324 if [[ -z "$1" ]]; then 325 break 326 fi 327 done 328 329 if [[ $(( GPU_BUILD + CPU_BUILD + GPUDIRECT_BUILD + ROCM_BUILD )) -gt "1" ]]; then 330 echo "Only one of [--gpu, --cpu, --gpudirect, --rocm] may be provided." 331 usage 332 exit 1 333 fi 334 335 if [[ -z "$DSTDIR" ]] && [[ -z "$SRCDIR" ]]; then 336 echo "No destination dir provided" 337 usage 338 exit 1 339 fi 340 341 if [[ -z "$SRCDIR" ]]; then 342 # make temp srcdir if none set 343 SRCDIR="$(mktemp -d -t tmp.XXXXXXXXXX)" 344 fi 345 346 prepare_src "$SRCDIR" 347 348 if [[ -z "$DSTDIR" ]]; then 349 # only want to prepare sources 350 exit 351 fi 352 353 if [[ -n ${PROJECT_NAME} ]]; then 354 PKG_NAME_FLAG="--project_name ${PROJECT_NAME}" 355 elif [[ ${NIGHTLY_BUILD} == "1" && ${GPU_BUILD} == "1" ]]; then 356 PKG_NAME_FLAG="--project_name tf_nightly_gpu" 357 elif [[ ${NIGHTLY_BUILD} == "1" && ${GPUDIRECT_BUILD} == "1" ]]; then 358 PKG_NAME_FLAG="--project_name tf_nightly_gpudirect" 359 elif [[ ${NIGHTLY_BUILD} == "1" && ${ROCM_BUILD} == "1" ]]; then 360 PKG_NAME_FLAG="--project_name tf_nightly_rocm" 361 elif [[ ${NIGHTLY_BUILD} == "1" && ${CPU_BUILD} == "1" ]]; then 362 PKG_NAME_FLAG="--project_name tf_nightly_cpu" 363 elif [[ ${NIGHTLY_BUILD} == "1" ]]; then 364 PKG_NAME_FLAG="--project_name tf_nightly" 365 elif [[ ${GPU_BUILD} == "1" ]]; then 366 PKG_NAME_FLAG="--project_name tensorflow_gpu" 367 elif [[ ${GPUDIRECT_BUILD} == "1" ]]; then 368 PKG_NAME_FLAG="--project_name tensorflow_gpudirect" 369 elif [[ ${ROCM_BUILD} == "1" ]]; then 370 PKG_NAME_FLAG="--project_name tensorflow_rocm" 371 elif [[ ${CPU_BUILD} == "1" ]]; then 372 PKG_NAME_FLAG="--project_name tensorflow_cpu" 373 fi 374 375 build_wheel "$SRCDIR" "$DSTDIR" "$PKG_NAME_FLAG" 376 377 if [[ $CLEANSRC -ne 0 ]]; then 378 rm -rf "${TMPDIR}" 379 fi 380} 381 382main "$@" 383