build.axs.sh

#/bin/bash

#
# Copyright (c) 2021-2024 Krai Ltd.
#
# SPDX-License-Identifier: MIT.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of
# this software and associated documentation files (the "Software"), to deal in
# the Software without restriction, including without limitation the rights to
# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
# of the Software, and to permit persons to whom the Software is furnished to do
# so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#

_SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
cd ${_SCRIPT_DIR}

source ./functions.sh

function exit_if_error() {
  if [ "${?}" != "0" ]; then exit 1; fi
}

if [[ $# < 1 ]]; then
  echo "Please enter the model name to build the Docker image for (one of: bert, resnet50, retinanet, sdxl).";

  exit 1;
fi
MODEL=$1
echo "Building AXS (QAIC-independent) image for '${MODEL}' ..."

if [[ ${MODEL} == "resnet50" ]]; then
  _IMAGENET=${IMAGENET:-"full"}
  if [[ "${_IMAGENET}" == "full" || "${_IMAGENET}" == "preprocessed" ]]
  then
    _IMAGENET_SUFFIX="full"
  else
    _IMAGENET_SUFFIX="min"
  fi
  _IMAGENET_BUILD_ARG="--build-arg IMAGENET=${_IMAGENET_SUFFIX}"
  _DOCKER_IMAGE_NAME="krai/axs.${MODEL}.${_IMAGENET_SUFFIX}"
  docker tag krai/axs.imagenet.${_IMAGENET_SUFFIX}:latest imagenet:latest
else
  _DOCKER_IMAGE_NAME="krai/axs.${MODEL}"
fi

# Control AXS version and use of private repositories.
_AXS_VERSION_CHECKOUT=${AXS_VERSION_CHECKOUT:-"mlperf_4.0"}
_DEV=${DEV:-"no"}
if [[ "${_DEV}" == "no" ]]; then
  _AXS_PRIVATE_PAT=""
elif [[ "${_DEV}" == "yes" ]]; then
  _AXS_PRIVATE_PAT=${AXS_PRIVATE_PAT}
fi

_DOCKER_OS=${DOCKER_OS:-deb}
_DOCKER_MODEL_IMAGE="${_DOCKER_IMAGE_NAME}:${_DOCKER_OS}_latest"
_DOCKER_COMMON_IMAGE="krai/axs.common:${_DOCKER_OS}_latest"

_NO_CACHE=${NO_CACHE:-"no"}
if [[ "${_NO_CACHE}" == "yes" ]]; then
  __NO_CACHE="--no-cache"
fi

echo "Running '$0' ..."
print_variables "${!_@}"
echo
echo "Press Ctrl-C to break ..."
sleep 10

if [[ "$(docker images -q ${_DOCKER_COMMON_IMAGE} 2> /dev/null)" == "" ]]; then
  cd ./base/ && AXS_VERSION_CHECKOUT=${_AXS_VERSION_CHECKOUT} DEV=${_DEV} ./build.axs.common.sh && cd ../
  exit_if_error
fi

if [[ ${MODEL} == "resnet50" ]]; then
  if [[ "$(docker images -q krai/axs.imagenet.${_IMAGENET_SUFFIX}:latest 2> /dev/null)" == "" ]]; then
    echo "Building ImageNet image for '${MODEL}' ...";
    cd ./imagenet/ && IMAGENET=${_IMAGENET} ./build.imagenet.sh && cd ../
    exit_if_error
  fi
fi

echo
echo "Building image: '${_DOCKER_MODEL_IMAGE}'"
echo
read -d '' CMD <<END_OF_CMD
cd ./${MODEL} && \
docker build ${__NO_CACHE} \
--build-arg DOCKER_OS=${_DOCKER_OS} \
--build-arg AXS_PRIVATE_PAT=${_AXS_PRIVATE_PAT} \
--build-arg AXS_VERSION_CHECKOUT=${_AXS_VERSION_CHECKOUT} \
${_IMAGENET_BUILD_ARG} \
-t ${_DOCKER_MODEL_IMAGE} \
-f Dockerfile.axs .
END_OF_CMD
echo "Command: ${CMD}"

if [ -z "${DRY_RUN}" ]; then
  eval ${CMD}
fi
exit_if_error

echo
echo "DONE (building '${_DOCKER_MODEL_IMAGE}')"
echo