Compile TensorFlow Serving with CUDA support (July 2017)

compile_tensorflow_serving.sh

#!/bin/bash


TENSORFLOW_COMMIT=bef67104a0da34b6cac0d801168f2e42a1ab7921 # July 7, 2017
TENSORFLOW_SERVING_COMMIT=15cb2ace25493161d47b0d76047b3997d371e59c # July 14, 2017
MODELS_COMMIT=3d97b007cdb2c73f23daf9141437796095436c22 # July 7, 2017

if [ -z $TENSORFLOW_SERVING_REPO_PATH ]; then
	TENSORFLOW_SERVING_REPO_PATH="serving"
fi
INITIAL_PATH=$(pwd)
export TF_NEED_CUDA=1
export TF_NEED_GCP=1
export TF_NEED_JEMALLOC=1
export TF_NEED_HDFS=0
export TF_NEED_OPENCL=0
export TF_NEED_MKL=0
export TF_NEED_VERBS=0
export TF_NEED_MPI=0
export TF_ENABLE_XLA=0
export TF_CUDA_VERSION=8.0
export TF_CUDNN_VERSION=5
export TF_CUDA_CLANG=0
export TF_CUDA_COMPUTE_CAPABILITIES="3.5,5.2,6.1"
CUDA_PATH="/usr/local/cuda"
if [ ! -d $CUDA_PATH ]; then
	CUDA_PATH="/opt/cuda"
fi
export CUDA_TOOLKIT_PATH=$CUDA_PATH
export CUDNN_INSTALL_PATH=$CUDA_PATH
export GCC_HOST_COMPILER_PATH=$(which gcc-5 || which gcc)
export PYTHON_BIN_PATH=$(which python)
export CC_OPT_FLAGS="-march=native"

function python_path {
  "$PYTHON_BIN_PATH" - <<END
from __future__ import print_function
import site
import os

try:
  input = raw_input
except NameError:
  pass

python_paths = []
if os.getenv('PYTHONPATH') is not None:
  python_paths = os.getenv('PYTHONPATH').split(':')
try:
  library_paths = site.getsitepackages()
except AttributeError:
 from distutils.sysconfig import get_python_lib
 library_paths = [get_python_lib()]
all_paths = set(python_paths + library_paths)

paths = []
for path in all_paths:
  if os.path.isdir(path):
    paths.append(path)

if len(paths) == 1:
  print(paths[0])
else:
  ret_paths = ",".join(paths)
  print(ret_paths)
END
}

export PYTHON_LIB_PATH=$(python_path)

cd $TENSORFLOW_SERVING_REPO_PATH
cd tensorflow
git reset --hard
git fetch
cd ../tf_models
git reset --hard
git fetch
cd ..
git reset --hard
git fetch
git submodule update --init --recursive
git checkout $TENSORFLOW_SERVING_COMMIT
git submodule update --init --recursive
cd tf_models
git checkout $MODELS_COMMIT
cd ../tensorflow
git submodule update --init --recursive
git checkout $TENSORFLOW_COMMIT
git submodule update --init --recursive
./configure
cd ..

# force to use gcc-5 to compile CUDA
if [ -e $(which gcc-5) ]; then
	sed -i.bak 's/"gcc"/"gcc-5"/g' tensorflow/third_party/gpus/cuda_configure.bzl
fi

# Error fix.  Ref: https://github.com/tensorflow/serving/issues/327
sed -i.bak 's/external\/nccl_archive\///g' tensorflow/tensorflow/contrib/nccl/kernels/nccl_manager.h
sed -i.bak 's/external\/nccl_archive\///g' tensorflow/tensorflow/contrib/nccl/kernels/nccl_ops.cc

cd ..
bazel build -c opt --config=cuda --verbose_failures @tf_serving//tensorflow_serving/model_servers:tensorflow_model_server

cd $INITIAL_PATH