#!/bin/bash

# run with `bazel run --config=mlinfra_v1 //scripts/build/mlinfra_platform:build_on_pipedream -- tensorflow22 tensorflow22_cpu` etc.
set -euo pipefail
set -x

ORIG_DIR=`pwd`
SCRIPTS_DIR=`pwd`/scripts/build
# This is a gross hack, but it works.
EXTERNAL_DIR=`pwd | sed -e 's#execroot.*#external#g'`

BRANCH=$1
shift

PATCH_DIR=$1
shift

CPU_OR_GPU=$1
shift

UPLOAD=$1
shift

# We're going to do some awful things to /usr/, so we need to make sure we're on a PD node inside a docker container.
set +x
if [[ `hostname | fgrep desktop` != "" ]]
then
    echo "Skipping doing some awful things to /.  hostname:" `hostname`
else
    if [[ `cat /proc/1/cpuset` == "/" ]]
    then
        echo "Skipping doing some awful things to /.  cat /proc/1/cpuset:" `cat /proc/1/cpuset`
    else
        if [[ "${PIPEDREAM_USER:-}" == "" ]]
        then
            echo "Skipping doing some awful things to /.  PIPEDREAM_USER not set: ${PIPEDREAM_USER:-}"
        else
            set -x
            # This is where we destroy /usr /etc etc.
            rm $EXTERNAL_DIR/gcc_7_4/BUILD.bazel
            sudo cp -r $EXTERNAL_DIR/gcc_7_4/ /
            sudo rm /usr/bin/gcc
            sudo ln -s /usr/bin/gcc-7 /usr/bin/gcc
        fi
    fi
fi

set -x
echo $PATH
which gcc
gcc --version

BASE_DIR=/mnt/sun-tcs01/infra/tools/mlinfra/build_tensorflow22
WORK_DIR=/tmp/build_tensorflow22/`date "+%FT%H.%M.%S-$$"`
mkdir -p $WORK_DIR

# get and compile bazel for the version we need if it hasn't already been
BAZEL_DIR=$BASE_DIR/bazel-2.0.0
# Note: if you need to change anything about how bazel is compiled, you must
# uncomment the following so that your changes don't pollute builds on other
# branches that don't have your changes yet.
# BAZEL_DIR=$WORK_DIR/bazel-2.0.0
if [ -d $BAZEL_DIR ]
then
    echo "Skipping bazel compile; $BAZEL_DIR already exists."
else
    mkdir -p $BAZEL_DIR
    cd $BASE_DIR
    wget https://github.com/bazelbuild/bazel/releases/download/2.0.0/bazel-2.0.0-dist.zip
    unzip -d $BAZEL_DIR bazel-2.0.0-dist.zip

    cd $BAZEL_DIR
    env EXTRA_BAZEL_ARGS="--host_javabase=@local_jdk//:jdk" bash ./compile.sh
    # The bazel binary we need to compile TF with should be located in output/ relative to the unzipped Bazel source.
fi

VENV_DIR=$BASE_DIR/build_venv
# Note: if you need to change anything about how the venv is made, you must
# uncomment the following so that your changes don't pollute builds on other
# branches that don't have your changes yet.
# VENV_DIR=$WORK_DIR/build_venv
if [ -d $VENV_DIR ]
then
    echo "Skipping venv creation; $VENV_DIR already exists."
else
    # set up the venv we're going to use, install the same versions of packages we use in driving
    # $WORK_DIR/venv/bin/pip install -U pip six 'numpy<1.19.0' wheel setuptools mock 'future>=0.17.1'
    virtualenv $VENV_DIR

    # Ran this to generate the below:
    # for p in six numpy wheel setuptools mock future
    # do
    #     echo `fgrep $p== build/requirements3.generated.txt` '\'
    # done

    # ERROR: Requested numpy==1.16.4.post1 from https://pypi.zooxlabs.com/numpy/1.16.4/numpy-1.16.4.post1-cp36-cp36m-linux_x86_64.whl#md5=76bdd6176e3b60e80b971274b0640fe1 has different version in metadata: '1.16.4'
    # Maybe can revert to 1.16.4.post1 now, haven't checked
    $VENV_DIR/bin/pip install -U pip==9.0.1

    $VENV_DIR/bin/pip install -U six==1.12.0 \
    numpy==1.16.4 \
    wheel==0.34.2 \
    setuptools==41.0.0 \
    mock==2.0.0 \
    pytest-mock==1.10.0 \
    future==0.18.2 \

    $VENV_DIR/bin/pip install -U keras_applications==1.0.8 --no-deps
    $VENV_DIR/bin/pip install -U keras_preprocessing==1.1.2 --no-deps

    # For the TF patching process
    $VENV_DIR/bin/pip install astor==0.8.1
fi

# Putting the TF checkout on TCS is slower than local disk
TF_DIR=$WORK_DIR/tensorflow
if [ -d $TF_DIR ]
then
    echo "Skipping tensorflow clone; $TF_DIR already exists."
    cd $TF_DIR
    git fetch origin $BRANCH
    git checkout $BRANCH
else
    cd `dirname $TF_DIR`
    git clone git@git.zooxlabs.com:zooxco/tensorflow.git --branch $BRANCH
    cd $TF_DIR
    git config user.email "mlinfra@zoox.com"
    git config user.name "ML Infra TF Updates"
fi

# Apply patches
# Note that eventually we shouldn't do this live any more, but for now it's okay for testing
$VENV_DIR/bin/python $SCRIPTS_DIR/mlinfra_platform/tensorflow/$PATCH_DIR/patch_tensorflow.py

CUDA_DIR=$WORK_DIR/cuda
cp -r $EXTERNAL_DIR/cuda_repo/cuda $CUDA_DIR

# Output from configure:
# Preconfigured Bazel build configs. You can use any of the below by adding "--config=<>" to your build command. See .bazelrc for more details.
#         --config=mkl            # Build with MKL support.
#         --config=monolithic     # Config for mostly static monolithic build.
#         --config=ngraph         # Build with Intel nGraph support.
#         --config=numa           # Build with NUMA support.
#         --config=dynamic_kernels        # (Experimental) Build kernels into separate shared objects.
#         --config=v2             # Build TensorFlow 2.x instead of 1.x.
# Preconfigured Bazel build configs to DISABLE default on features:
#         --config=noaws          # Disable AWS S3 filesystem support.
#         --config=nogcp          # Disable GCP support.
#         --config=nohdfs         # Disable HDFS support.
#         --config=nonccl         # Disable NVIDIA NCCL support.

# instead of running ./configure we should create the configure file we need
# This template is from running ./configure by hand once
# fix the compute capabilities (do we need prior to 7.5?)
cat > $TF_DIR/.tf_configure.bazelrc <<_CONFIG_EOF
build --action_env PYTHON_BIN_PATH="$VENV_DIR/bin/python"
build --action_env PYTHON_LIB_PATH="$VENV_DIR/lib/python3.6/site-packages"
build --python_path="$VENV_DIR/bin/python"
build --config=xla
build --action_env GCC_HOST_COMPILER_PATH="/usr/bin/gcc-7"
build --action_env CUDA_TOOLKIT_PATH="$CUDA_DIR"
build --action_env TF_CUDA_COMPUTE_CAPABILITIES="6.0,6.1,6.2,7.0,7.2,7.5"  # Add 8.0,8.6 for ampere
build --config=cuda
build --jobs=64
build --noshow_loading_progress
build --show_progress_rate_limit=30
build --curses=no
build --output_filter=DONT_MATCH_ANYTHING
build:opt --copt=-mavx2
build:opt --copt=-Wno-sign-compare
build:opt --host_copt=-mavx2
build:opt --define with_default_optimizations=true
test --flaky_test_attempts=3
test --test_size_filters=small,medium
test --test_env=LD_LIBRARY_PATH
test:v1 --test_tag_filters=-benchmark-test,-no_oss,-no_gpu,-oss_serial
test:v1 --build_tag_filters=-benchmark-test,-no_oss,-no_gpu
test:v2 --test_tag_filters=-benchmark-test,-no_oss,-no_gpu,-oss_serial,-v1only
test:v2 --build_tag_filters=-benchmark-test,-no_oss,-no_gpu,-v1only
build --action_env TF_CONFIGURE_IOS="0"
_CONFIG_EOF

if [ "$CPU_OR_GPU" == "cpu" ]
then

cat > $TF_DIR/.tf_configure.bazelrc <<_CONFIG_EOF
build --action_env PYTHON_BIN_PATH="$VENV_DIR/bin/python"
build --action_env PYTHON_LIB_PATH="$VENV_DIR/lib/python3.6/site-packages"
build --python_path="$VENV_DIR/bin/python"
build --config=xla
build --jobs=64
build --noshow_loading_progress
build --show_progress_rate_limit=30
build --curses=no
build --output_filter=DONT_MATCH_ANYTHING
build:opt --copt=-mavx2
build:opt --copt=-Wno-sign-compare
build:opt --host_copt=-mavx2
build:opt --define with_default_optimizations=true
test --flaky_test_attempts=3
test --test_size_filters=small,medium
test:v1 --test_tag_filters=-benchmark-test,-no_oss,-gpu,-oss_serial
test:v1 --build_tag_filters=-benchmark-test,-no_oss,-gpu
test:v2 --test_tag_filters=-benchmark-test,-no_oss,-gpu,-oss_serial,-v1only
test:v2 --build_tag_filters=-benchmark-test,-no_oss,-gpu,-v1only
build --action_env TF_CONFIGURE_IOS="0"
_CONFIG_EOF

fi

# Do we want a debug build?
DEBUG_OPTIONS=""
# DEBUG_OPTIONS="-c dbg --strip=never"

date

# Now we build the target that will make the wheel, then run it.
# For some reason, these two commands must be separate or else Bazel says "Could not find bazel-bin"
TF_CUDA_PATHS=$CUDA_DIR \
    $BAZEL_DIR/output/bazel --nohome_rc --output_base $WORK_DIR/bazel_output_base \
    build --verbose_failures $DEBUG_OPTIONS \
    --noshow_loading_progress --show_progress_rate_limit=30 --curses=no \
    --noincompatible_do_not_split_linking_cmdline \
    --config=opt -c opt //tensorflow/tools/pip_package:build_pip_package

cd bazel-bin/tensorflow/tools/pip_package/build_pip_package.runfiles/org_tensorflow/tensorflow
ln -s libtensorflow_cc.so.2.2.2 libtensorflow_cc.so.2.2
ln -s libtensorflow_cc.so.2.2 libtensorflow_cc.so.2
ln -s libtensorflow_cc.so.2 libtensorflow_cc.so

date

cd $TF_DIR
PYTHON_BIN_PATH=$VENV_DIR/bin/python ./bazel-bin/tensorflow/tools/pip_package/build_pip_package $WORK_DIR/pip_tensorflow_pkg

date

# Install the new TF package in a fresh venv and test it.
# The || true is because of glibc segfaults.
cd $ORIG_DIR # Avoid importing from the actual TensorFlow directory
virtualenv $WORK_DIR/test_venv
$WORK_DIR/test_venv/bin/pip install --quiet $WORK_DIR/pip_tensorflow_pkg/tensorflow-*.whl
$WORK_DIR/test_venv/bin/python -c "import tensorflow as tf;print(tf.reduce_sum(tf.random.normal([1000, 1000])))" || true
$WORK_DIR/test_venv/bin/python -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.get_compile_flags()))' || true
$WORK_DIR/test_venv/bin/python -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.get_link_flags()))' || true

unzip -l $WORK_DIR/pip_tensorflow_pkg/tensorflow-*.whl | fgrep .so

# Upload the wheel
if [[ "$UPLOAD" == "upload" ]]
then
    echo "Wheel to upload: $WORK_DIR/pip_tensorflow_pkg/tensorflow-*.whl"
    $SCRIPTS_DIR/init_z3_profile.py
    $SCRIPTS_DIR/upload-to-z3-and-s3.sh zoox-third-party $WORK_DIR/pip_tensorflow_pkg/tensorflow-*.whl | fgrep -v 'file(s) remaining'
fi
