install_utils

#!/bin/bash
set -euo pipefail

clean_slate()
{
	echo "Cleaning up state"
	# uninstall any global installations
	pip uninstall --yes torchserve || true
	pip uninstall --yes torch-model-archiver || true
	sleep 5
}

install_java_deps()
{
  set +e
  JAVA_VERSION=$(java -version 2>&1 | head -1 | cut -d'"' -f2 | sed '/^1\./s///' | cut -d'.' -f1)
	if [ "$JAVA_VERSION" != "17" ]
  then
    if [[ "$OSTYPE" == "darwin"* ]]; then
      brew install openjdk@17
    else
      sudo apt-get install -y openjdk-17-jdk
    fi
  fi
  set -e
}

install_torch_deps()
{
  if is_gpu_instance && [ ! -z "$1" ];
  then
	  pip install -U -r requirements_$1.txt -f https://download.pytorch.org/whl/torch_stable.html
	else
	  pip install -U -r requirements.txt
	fi
}

install_pytest_suite_deps()
{
  pip install -U -r requirements/developer.txt
}

install_bert_dependencies()
{
  pip install transformers==3.5.1
}

build_frontend() {
  if frontend/gradlew -p frontend clean build;
    then
    echo "Frontend build suite execution successful"
  else
    echo "Frontend build suite execution failed!!! Check logs for more details"
    exit 1
  fi
}

run_backend_pytest()
{
  if python -m pytest --cov-report html:htmlcov --cov=ts/ ts/tests/unit_tests/;
  then
    echo "Backend test suite execution successful"
  else
    echo "Backend test suite execution failed!!! Check logs for more details"
    exit 1
  fi
}

run_backend_python_linting()
{
  if pylint -rn --rcfile=./ts/tests/pylintrc ts/.;
  then
    echo "Backend python linting suite execution successful"
  else
    echo "Backend python linting execution failed!!! Check logs for more details"
    exit 1
  fi
}

run_model_archiver_python_linting()
{
  cd model-archiver
  if pylint -rn --rcfile=./model_archiver/tests/pylintrc model_archiver/.;
  then
    echo "Model archiver python linting suite execution successful"
  else
    echo "Model archiver python linting execution failed!!! Check logs for more details"
    exit 1
  fi
  cd ..
}

run_model_archiver_UT_suite()
{
  cd model-archiver
  if python -m pytest --cov-report html:htmlcov_ut --cov=model_archiver/ model_archiver/tests/unit_tests/;
  then
    echo "Model-archiver UT test suite execution successfully"
  else
    echo "Model-archiver UT test suite execution failed!!! Check logs for more details"
    exit 1
  fi
  cd ..
}

run_model_archiver_IT_suite()
{
  cd model-archiver
  if python -m pytest --cov-report html:htmlcov_it --cov=model_archiver/ model_archiver/tests/integ_tests/;
  then
    echo "Model-archiver IT test suite execution successful"
  else
    echo "Model-archiver IT test suite execution failed!!! Check logs for more details"
    exit 1
  fi
  cd ..
}

run_workflow_archiver_python_linting()
{
  cd workflow-archiver
  if pylint -rn --rcfile=./workflow_archiver/tests/pylintrc workflow_archiver/.;
  then
    echo "Workflow archiver python linting suite execution successful"
  else
    echo "Workflow archiver python linting execution failed!!! Check logs for more details"
    exit 1
  fi
  cd ..
}

run_workflow_archiver_UT_suite()
{
  cd workflow-archiver
  if python -m pytest --cov-report html:htmlcov_ut --cov=workflow_archiver/ workflow_archiver/tests/unit_tests/;
  then
    echo "Workflow-archiver UT test suite execution successfully"
  else
    echo "Workflow-archiver UT test suite execution failed!!! Check logs for more details"
    exit 1
  fi
  cd ..
}

run_workflow_archiver_IT_suite()
{
  cd workflow-archiver
  if python -m pytest --cov-report html:htmlcov_it --cov=workflow_archiver/ workflow_archiver/tests/integ_tests/;
  then
    echo "Workflow-archiver IT test suite execution successful"
  else
    echo "Workflow-archiver IT test suite execution failed!!! Check logs for more details"
    exit 1
  fi
  cd ..
}

build_install_server()
{
	pip install .
}

build_install_model_archiver()
{
	echo $PWD
	cd model-archiver
	if pip install .;
  then
    echo "Successfully installed torch-model-archiver"
  else
    echo "torch-model-archiver installation failed"
    exit 1
  fi
  cd ..
}

build_install_workflow_archiver()
{
	echo $PWD
	cd workflow-archiver
	if pip install .;
  then
    echo "Successfully installed torch-workflow-archiver"
  else
    echo "torch-workflow-archiver installation failed"
    exit 1
  fi
  cd ..
}

start_torchserve()
{
  echo "Starting TorchServe"
  torchserve --start --model-store model_store &
  pid=$!
  count=$(ps -A| grep $pid |wc -l)
  if [[ $count -eq 1 ]]
  then
          if wait $pid; then
                  echo "Successfully started TorchServe"
          else
                  echo "TorchServe start failed (returned $?)"
                  exit 1
          fi
  else
          echo "Successfully started TorchServe"
  fi

  sleep 10
}

stop_torchserve()
{
  torchserve --stop
  sleep 10
}


# Takes model name and mar name from model zoo as input
register_model()
{
  echo "Registering $1 model"
  response=$(curl --write-out %{http_code} --silent --output /dev/null --retry 5 -X POST "http://localhost:8081/models?url=https://torchserve.pytorch.org/mar_files/$1.mar&initial_workers=1&synchronous=true&model_name=$1")

  if [ ! "$response" == 200 ]
  then
      echo "Failed to register model with torchserve"
      cleanup
      exit 1
  else
      echo "Successfully registered $1 model with torchserve"
  fi
}

# Takes model URL and payload path as input
run_inference()
{
  for i in {1..4}
  do
    echo "Running inference on $1 model"
    response=$(curl --write-out %{http_code} --silent --output /dev/null --retry 5 -X POST http://localhost:8080/predictions/$1 -T $2)

    if [ ! "$response" == 200 ]
    then
        echo "Failed to run inference on $1 model"
        cleanup
        exit 1
    else
        echo "Successfully ran infernece on $1 model."
    fi
  done
}


unregister_model()
{
  echo "Unregistering $1 model"
  response=$(curl --write-out %{http_code} --silent --output /dev/null --retry 5 -X DELETE "http://localhost:8081/models/$1")

  if [ ! "$response" == 200 ]
  then
      echo "Failed to register $1 model with torchserve"
      cleanup
      exit 1
  else
      echo "Successfully registered $1 model with torchserve"
  fi
}

clean_up_build_residuals()
{
  rm -rf ts/__pycache__/
  rm -rf ts/metrics/__pycache__/
  rm -rf ts/protocol/__pycache__/
  rm -rf ts/utils/__pycache__/
}

is_gpu_instance(){
  if command -v nvidia-smi;
  then
    nvidia-smi | grep 'NVIDIA-SMI has failed'
    if [ $? == 0 ];
    then
      return 1
    else
      return 0
    fi

  else
    return 1
  fi
}

run_markdown_link_checker(){
  STATUS=0
  set +e
  for i in $(find . -name \*.md); do # Whitespace-safe and recursive
        markdown-link-check $i --config link_check_config.json -q
        if [ $?  -ne  0 ] ; then
          echo Broken links in $i
          STATUS=1
        fi
  done
  set -e
  exit $STATUS
}