#! /bin/bash

# run with `bazel run --config=mlinfra_v1 //scripts/build/mlinfra_platform:build_on_pipedream -- tensorflow24`

set -euo pipefail
set -x

# sudo apt install python3-dev python3-pip python3-venv
# sudo apt-get install build-essential openjdk-11-jdk python zip unzip
# sudo aptitude install --assume-yes build-essential openjdk-11-jdk zip unzip

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'`

BASE_DIR=/mnt/sun-tcs01/infra/tools/mlinfra/build_tensorflow24
WORK_DIR=$BASE_DIR/`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-3.1.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/3.1.0/bazel-3.1.0-dist.zip
    unzip -d $BAZEL_DIR bazel-3.1.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
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 --no-deps
    $VENV_DIR/bin/pip install -U keras_preprocessing --no-deps
fi

# get TF
# cd $WORK_DIR

# This should probably be $WORK_DIR, but it's slow, so for testing i'm reusing it
TF_DIR=$BASE_DIR/tensorflow
if [ -d $TF_DIR ]
then
    echo "Skipping tensorflow clone; $TF_DIR already exists."
else
    cd `dirname $TF_DIR`
    git clone git@git.zooxlabs.com:zooxco/tensorflow.git
fi

cd $TF_DIR
git fetch origin mlinfra-v2.4.1
git checkout mlinfra-v2.4.1

CUDA_DIR=$EXTERNAL_DIR/cuda_repo/cuda/cuda_cudart/targets/x86_64-linux

# 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 CUDA_TOOLKIT_PATH="$CUDA_DIR"
build --action_env TF_CUDA_COMPUTE_CAPABILITIES="7.5,8.0,8.6"
build --action_env GCC_HOST_COMPILER_PATH="/usr/bin/gcc-7"
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=-march=native
build:opt --copt=-Wno-sign-compare
build:opt --host_copt=-march=native
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

# LD_LIBRARY_PATH=$CUDA_DIR/lib
# For some reason, these two commands must be separate or else Bazel says "Could not find bazel-bin"
TF_CUDA_PATHS=$EXTERNAL_DIR/cuda_repo/cuda/cuda_cudart/targets/x86_64-linux:$EXTERNAL_DIR/cuda_repo/cuda/cuda_nvcc \
    $BAZEL_DIR/output/bazel --nohome_rc build --config=opt -c opt //tensorflow/tools/pip_package:build_pip_package
PYTHON_BIN_PATH=$VENV_DIR/bin/python ./bazel-bin/tensorflow/tools/pip_package/build_pip_package $WORK_DIR/pip_tensorflow_pkg

# install the package in a fresh venv and test it
cd $ORIG_DIR # Avoid importing from the actual TensorFlow directory
virtualenv $WORK_DIR/venv
$WORK_DIR/venv/bin/pip install $WORK_DIR/pip_tensorflow_pkg/tensorflow-*.whl
$WORK_DIR/venv/bin/python -c "import tensorflow as tf;print(tf.reduce_sum(tf.random.normal([1000, 1000])))"

# Upload the wheel
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
