510 lines
14 KiB
Bash
510 lines
14 KiB
Bash
#!/usr/bin/env bash
|
|
|
|
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
|
|
#
|
|
# Start a bash, mount REPO_MOUNT_POINT to be current directory.
|
|
#
|
|
# Usage: docker/bash.sh [-i|--interactive] [--net=host] [-t|--tty]
|
|
# [--mount MOUNT_DIR] [--repo-mount-point REPO_MOUNT_POINT]
|
|
# [--dry-run] [--name NAME] [--privileged]
|
|
# <DOCKER_IMAGE_NAME> [--] [COMMAND]
|
|
#
|
|
# Usage: docker/bash.sh <CONTAINER_NAME>
|
|
# Starts an interactive session
|
|
#
|
|
# Usage2: docker/bash.sh [-i] <CONTAINER_NAME> [COMMAND]
|
|
# Execute command in the docker image, default non-interactive
|
|
# With -i, execute interactively.
|
|
#
|
|
|
|
set -euo pipefail
|
|
|
|
function show_usage() {
|
|
cat <<EOF
|
|
Usage: docker/bash.sh [-i|--interactive] [--net=host] [-t|--tty]
|
|
[--cpus NUM_CPUS] [--mount MOUNT_DIR]
|
|
[--repo-mount-point REPO_MOUNT_POINT]
|
|
[--dry-run] [--name NAME]
|
|
<DOCKER_IMAGE_NAME> [--] [COMMAND]
|
|
|
|
|
|
-h, --help
|
|
|
|
Display this help message.
|
|
|
|
-i, --interactive
|
|
|
|
Start the docker session in interactive mode.
|
|
|
|
-t, --tty
|
|
|
|
Start the docker session with a pseudo terminal (tty).
|
|
|
|
--cpus NUM_CPUS
|
|
|
|
Limit the number of CPU cores to be used.
|
|
|
|
--net=host
|
|
|
|
Expose servers run into the container to the host, passing the
|
|
"--net=host" argument through to docker. On MacOS, this is
|
|
instead passed as "-p 8888:8888" since the host networking driver
|
|
isn't supported.
|
|
|
|
--mount MOUNT_DIR
|
|
|
|
Expose MOUNT_DIR as an additional mount point inside the docker
|
|
container. The mount point inside the container is the same as
|
|
the folder location outside the container. This option can be
|
|
specified multiple times.
|
|
|
|
--repo-mount-point REPO_MOUNT_POINT
|
|
|
|
The directory inside the docker container at which the TVM
|
|
repository should be mounted, and is used as the workspace inside
|
|
the docker container.
|
|
|
|
If unspecified, the mount location depends on the environment. If
|
|
running inside Jenkins, the mount location will be /workspace.
|
|
Otherwise, the mount location of the repository will be the same
|
|
as the external location of the repository, to maintain
|
|
compatibility with git-worktree.
|
|
|
|
--no-gpu
|
|
|
|
Do not use GPU device drivers even if using an CUDA Docker image
|
|
|
|
--dry-run
|
|
|
|
Print the docker command to be run, but do not execute it.
|
|
|
|
--build-dir BUILD_DIR
|
|
|
|
The build directory of TVM. This is appended to LD_LIBRARY_PATH
|
|
|
|
--env
|
|
|
|
Pass an environment variable through to the container.
|
|
|
|
--name
|
|
|
|
Set the name of the docker container, and the hostname that will
|
|
appear inside the container.
|
|
|
|
--privileged
|
|
|
|
Give extended privileges to this container.
|
|
|
|
DOCKER_IMAGE_NAME
|
|
|
|
The name of the docker container to be run. This can be an
|
|
explicit name of a docker image (e.g. "tlcpack/ci-gpu:v0.76") or
|
|
can be a shortcut as defined in the TVM Jenkinsfile
|
|
(e.g. "ci_gpu").
|
|
|
|
COMMAND
|
|
|
|
The command to be run inside the docker container. If this is set
|
|
to "bash", the --interactive, --tty and --net=host flags are set.
|
|
If no command is specified, defaults to "bash". If the command
|
|
contains dash-prefixed arguments, the command should be preceded
|
|
by -- to indicate arguments that are not intended for bash.sh.
|
|
|
|
EOF
|
|
}
|
|
|
|
|
|
#################################
|
|
### Start of argument parsing ###
|
|
#################################
|
|
|
|
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd -P)"
|
|
REPO_DIR="$(dirname "${SCRIPT_DIR}")"
|
|
|
|
DRY_RUN=false
|
|
INTERACTIVE=false
|
|
TTY=false
|
|
USE_NET_HOST=false
|
|
USE_GPU=true
|
|
DOCKER_IMAGE_NAME=
|
|
COMMAND=bash
|
|
MOUNT_DIRS=( )
|
|
CONTAINER_NAME=
|
|
|
|
# TODO(Lunderberg): Remove this if statement and always set to
|
|
# "${REPO_DIR}". The consistent directory for Jenkins is currently
|
|
# necessary to allow cmake build commands to run in CI after the build
|
|
# steps.
|
|
# TODO(https://github.com/apache/tvm/issues/11952):
|
|
# Figure out a better way to keep the same path
|
|
# between build and testing stages.
|
|
if [[ -n "${JENKINS_HOME:-}" ]]; then
|
|
REPO_MOUNT_POINT=/workspace
|
|
else
|
|
REPO_MOUNT_POINT="${REPO_DIR}"
|
|
fi
|
|
|
|
|
|
function parse_error() {
|
|
echo "$@" >&2
|
|
show_usage >&2
|
|
exit 1
|
|
}
|
|
|
|
# Handle joined flags, such as interpreting -ih as -i -h. Either rewrites
|
|
# the current argument if it is a joined argument, or shifts all arguments
|
|
# otherwise. Should be called as "eval $break_joined_flag" where joined
|
|
# flags are possible. Can't use a function definition, because it needs
|
|
# to overwrite the parent scope's behavior.
|
|
break_joined_flag='if (( ${#1} == 2 )); then shift; else set -- -"${1#-i}" "${@:2}"; fi'
|
|
|
|
DOCKER_ENV=( )
|
|
DOCKER_FLAGS=( )
|
|
|
|
while (( $# )); do
|
|
case "$1" in
|
|
-h|--help)
|
|
show_usage
|
|
exit 0
|
|
;;
|
|
|
|
-i*|--interactive)
|
|
INTERACTIVE=true
|
|
eval $break_joined_flag
|
|
;;
|
|
|
|
-t*|--tty)
|
|
TTY=true
|
|
eval $break_joined_flag
|
|
;;
|
|
|
|
--cpus)
|
|
DOCKER_FLAGS+=(--cpus "$2")
|
|
shift 2
|
|
;;
|
|
|
|
--net=host)
|
|
USE_NET_HOST=true
|
|
shift
|
|
;;
|
|
|
|
--net)
|
|
DOCKER_FLAGS+=( --net "$2" )
|
|
shift 2
|
|
;;
|
|
|
|
--mount)
|
|
if [[ -n "$2" ]]; then
|
|
MOUNT_DIRS+=("$2")
|
|
shift 2
|
|
else
|
|
parse_error 'ERROR: --mount requires a non-empty argument'
|
|
fi
|
|
;;
|
|
|
|
--mount=?*)
|
|
MOUNT_DIRS+=("${1#*=}")
|
|
shift
|
|
;;
|
|
|
|
--name)
|
|
if [[ -n "$2" ]]; then
|
|
CONTAINER_NAME="$2"
|
|
shift 2
|
|
else
|
|
parse_error 'ERROR: --name requires a non empty argument'
|
|
fi
|
|
;;
|
|
|
|
--privileged)
|
|
DOCKER_FLAGS+=( "--privileged" )
|
|
shift 1
|
|
;;
|
|
|
|
--env)
|
|
DOCKER_ENV+=( --env "$2" )
|
|
shift 2
|
|
;;
|
|
|
|
--volume)
|
|
DOCKER_FLAGS+=( --volume "$2" )
|
|
shift 2
|
|
;;
|
|
|
|
-e)
|
|
DOCKER_ENV+=( --env "$2" )
|
|
shift 2
|
|
;;
|
|
|
|
-v)
|
|
DOCKER_FLAGS+=( --volume "$2" )
|
|
shift 2
|
|
;;
|
|
|
|
--dry-run)
|
|
DRY_RUN=true
|
|
shift
|
|
;;
|
|
|
|
--no-gpu)
|
|
USE_GPU=false
|
|
shift
|
|
;;
|
|
|
|
--repo-mount-point)
|
|
if [[ -n "$2" ]]; then
|
|
REPO_MOUNT_POINT="$2"
|
|
shift 2
|
|
else
|
|
parse_error 'ERROR: --repo-mount-point requires a non-empty argument'
|
|
fi
|
|
;;
|
|
|
|
--repo-mount-point=?*)
|
|
REPO_MOUNT_POINT="${1#*=}"
|
|
shift
|
|
;;
|
|
|
|
--build-dir)
|
|
DOCKER_ENV+=( --env LD_LIBRARY_PATH=${REPO_MOUNT_POINT}/${2})
|
|
shift 2
|
|
;;
|
|
|
|
--)
|
|
shift
|
|
COMMAND=( "$@" )
|
|
break
|
|
;;
|
|
|
|
-*|--*)
|
|
echo "Error: Unknown flag: $1" >&2
|
|
echo " If this flag is intended to be passed to the" >&2
|
|
echo " docker command, please add -- before the docker" >&2
|
|
echo " command (e.g. docker/bash.sh ci_gpu -- build -j2)" >&2
|
|
show_usage >&2
|
|
exit 1
|
|
;;
|
|
|
|
*)
|
|
# First positional argument is the image name, all
|
|
# remaining below to the COMMAND.
|
|
if [[ -z "${DOCKER_IMAGE_NAME}" ]]; then
|
|
DOCKER_IMAGE_NAME=$1
|
|
shift
|
|
else
|
|
COMMAND=( "$@" )
|
|
break
|
|
fi
|
|
;;
|
|
esac
|
|
done
|
|
|
|
|
|
|
|
if [[ -z "${DOCKER_IMAGE_NAME}" ]]; then
|
|
echo "Error: Missing DOCKER_IMAGE_NAME" >&2
|
|
show_usage >&2
|
|
fi
|
|
|
|
if [[ ${COMMAND[@]+"${COMMAND[@]}"} = bash ]]; then
|
|
INTERACTIVE=true
|
|
TTY=true
|
|
USE_NET_HOST=true
|
|
fi
|
|
|
|
###############################
|
|
### End of argument parsing ###
|
|
###############################
|
|
|
|
source "$(dirname $0)/dev_common.sh" || exit 2
|
|
|
|
DOCKER_MOUNT=( )
|
|
DOCKER_DEVICES=( )
|
|
# If the user gave a shortcut defined in the Jenkinsfile, use it.
|
|
EXPANDED_SHORTCUT=$(lookup_image_spec "${DOCKER_IMAGE_NAME}")
|
|
if [ -n "${EXPANDED_SHORTCUT}" ]; then
|
|
if [ "${CI+x}" == "x" ]; then
|
|
DOCKER_IMAGE_NAME="${EXPANDED_SHORTCUT}"
|
|
else
|
|
python3 ci/scripts/jenkins/determine_docker_images.py "$DOCKER_IMAGE_NAME" 2> /dev/null
|
|
DOCKER_IMAGE_NAME=$(cat ".docker-image-names/$DOCKER_IMAGE_NAME")
|
|
if [[ "$DOCKER_IMAGE_NAME" == *"tlcpackstaging"* ]]; then
|
|
echo "WARNING: resolved docker image to fallback tag in tlcpackstaging" >&2
|
|
fi
|
|
fi
|
|
fi
|
|
|
|
# Set up working directories
|
|
|
|
DOCKER_FLAGS+=( --workdir "${REPO_MOUNT_POINT}" )
|
|
DOCKER_MOUNT+=( --volume "${REPO_DIR}":"${REPO_MOUNT_POINT}"
|
|
--volume "${SCRIPT_DIR}":/docker
|
|
)
|
|
|
|
# Set up CI-specific environment variables
|
|
DOCKER_ENV+=( --env CI_BUILD_HOME="${REPO_MOUNT_POINT}"
|
|
--env CI_BUILD_USER="$(id -u -n)"
|
|
--env CI_BUILD_UID="$(id -u)"
|
|
--env CI_BUILD_GROUP="$(id -g -n)"
|
|
--env CI_BUILD_GID="$(id -g)"
|
|
--env CI_PYTEST_ADD_OPTIONS="${CI_PYTEST_ADD_OPTIONS:-}"
|
|
--env CI_IMAGE_NAME="${DOCKER_IMAGE_NAME}"
|
|
)
|
|
|
|
# Remove the container once it finishes running (--rm).
|
|
DOCKER_FLAGS+=(--rm)
|
|
|
|
# Share the PID namespace (--pid=host). The process inside does not
|
|
# have pid 1 and SIGKILL is propagated to the process inside, allowing
|
|
# jenkins to kill it if needed. This is only necessary for docker
|
|
# daemons running as root.
|
|
if [ -z "${DOCKER_IS_ROOTLESS}" ]; then
|
|
DOCKER_FLAGS+=(--pid=host)
|
|
fi
|
|
|
|
|
|
# Expose services running in container to the host.
|
|
if $USE_NET_HOST; then
|
|
if [[ $(uname) == "Darwin" ]]; then
|
|
# Docker's host networking driver isn't supported on macOS.
|
|
# Use default bridge network and expose port for jupyter notebook.
|
|
DOCKER_FLAGS+=( -p 8888:8888 )
|
|
else
|
|
DOCKER_FLAGS+=(--net=host)
|
|
fi
|
|
fi
|
|
|
|
# Set up interactive sessions
|
|
if ${INTERACTIVE}; then
|
|
DOCKER_FLAGS+=( --interactive )
|
|
fi
|
|
|
|
if ${TTY}; then
|
|
DOCKER_FLAGS+=( --tty )
|
|
fi
|
|
|
|
# Setup the docker name and the hostname inside the container
|
|
if [[ ! -z "${CONTAINER_NAME}" ]]; then
|
|
DOCKER_FLAGS+=( --name ${CONTAINER_NAME} --hostname ${CONTAINER_NAME})
|
|
fi
|
|
|
|
# Expose external directories to the docker container
|
|
for MOUNT_DIR in ${MOUNT_DIRS[@]+"${MOUNT_DIRS[@]}"}; do
|
|
DOCKER_MOUNT+=( --volume "${MOUNT_DIR}:${MOUNT_DIR}" )
|
|
done
|
|
|
|
# Use nvidia-docker for GPU container. If nvidia-docker is not
|
|
# available, fall back to using "--gpus all" flag, requires docker
|
|
# version 19.03 or higher.
|
|
|
|
if [[ "$USE_GPU" == "true" ]] && [[ "${DOCKER_IMAGE_NAME}" == *"gpu"* || "${DOCKER_IMAGE_NAME}" == *"cuda"* ]]; then
|
|
if type nvidia-docker 1> /dev/null 2> /dev/null; then
|
|
DOCKER_BINARY=nvidia-docker
|
|
else
|
|
DOCKER_BINARY=docker
|
|
DOCKER_FLAGS+=( --gpus all )
|
|
fi
|
|
|
|
# nvidia-docker treats Vulkan as a graphics API, so we need to
|
|
# request passthrough of graphics APIs. This could also be set in
|
|
# the Dockerfile.
|
|
DOCKER_ENV+=( --env NVIDIA_DRIVER_CAPABILITIES=compute,graphics,utility )
|
|
|
|
# But as of nvidia-docker version 2.6.0-1, we still need to pass
|
|
# through the nvidia icd files ourselves.
|
|
ICD_SEARCH_LOCATIONS=(
|
|
# https://github.com/KhronosGroup/Vulkan-Loader/blob/master/loader/LoaderAndLayerInterface.md#icd-discovery-on-linux
|
|
/usr/local/etc/vulkan/icd.d
|
|
/usr/local/share/vulkan/icd.d
|
|
/etc/vulkan/icd.d
|
|
/usr/share/vulkan/icd.d
|
|
# https://github.com/NVIDIA/libglvnd/blob/master/src/EGL/icd_enumeration.md#icd-installation
|
|
/etc/glvnd/egl_vendor.d
|
|
/usr/share/glvnd/egl_vendor.d
|
|
)
|
|
for filename in $(find "${ICD_SEARCH_LOCATIONS[@]}" -name "*nvidia*.json" 2> /dev/null); do
|
|
DOCKER_MOUNT+=( --volume "${filename}":"${filename}":ro )
|
|
done
|
|
|
|
else
|
|
DOCKER_BINARY=docker
|
|
fi
|
|
|
|
|
|
|
|
# Pass any restrictions of allowed CUDA devices from the host to the
|
|
# docker container.
|
|
if [[ -n ${CUDA_VISIBLE_DEVICES:-} ]]; then
|
|
DOCKER_ENV+=( --env CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES}" )
|
|
fi
|
|
|
|
|
|
|
|
# Set TVM import path inside the docker image
|
|
if [[ "${DOCKER_IMAGE_NAME}" == *"ci"* ]]; then
|
|
DOCKER_ENV+=( --env PYTHONPATH="${REPO_MOUNT_POINT}"/python )
|
|
fi
|
|
|
|
# Add ROCm devices and set ROCM_ENABLED=1 which is used in the with_the_same_user script
|
|
# to add the user to the video group
|
|
if [[ "${DOCKER_IMAGE_NAME}" == *"rocm"* && -d "/dev/dri" ]]; then
|
|
DOCKER_DEVICES+=( --device=/dev/kfd --device=/dev/dri )
|
|
DOCKER_ENV+=( --env ROCM_ENABLED=1 )
|
|
fi
|
|
|
|
# When running from a git worktree, also mount the original git dir.
|
|
if [ -f "${REPO_DIR}/.git" ]; then
|
|
git_dir=$(cd ${REPO_DIR} && git rev-parse --git-common-dir)
|
|
if [ "${git_dir}" != "${REPO_DIR}/.git" ]; then
|
|
DOCKER_MOUNT+=( --volume "${git_dir}:${git_dir}" )
|
|
fi
|
|
fi
|
|
|
|
# If the docker daemon is running as root, use the TVM-provided
|
|
# "with_the_same_user" script to update the PID. When using rootless
|
|
# docker, this step is unnecessary.
|
|
if [ -z "${DOCKER_IS_ROOTLESS}" ]; then
|
|
COMMAND=(
|
|
bash --login /docker/with_the_same_user
|
|
${COMMAND[@]+"${COMMAND[@]}"}
|
|
)
|
|
fi
|
|
|
|
# Print arguments.
|
|
echo "REPO_DIR: ${REPO_DIR}"
|
|
echo "DOCKER CONTAINER NAME: ${DOCKER_IMAGE_NAME}"
|
|
echo ""
|
|
|
|
echo Running \'${COMMAND[@]+"${COMMAND[@]}"}\' inside ${DOCKER_IMAGE_NAME}...
|
|
|
|
DOCKER_CMD=(${DOCKER_BINARY} run
|
|
${DOCKER_FLAGS[@]+"${DOCKER_FLAGS[@]}"}
|
|
${DOCKER_ENV[@]+"${DOCKER_ENV[@]}"}
|
|
${DOCKER_MOUNT[@]+"${DOCKER_MOUNT[@]}"}
|
|
${DOCKER_DEVICES[@]+"${DOCKER_DEVICES[@]}"}
|
|
"${DOCKER_IMAGE_NAME}"
|
|
${COMMAND[@]+"${COMMAND[@]}"}
|
|
)
|
|
|
|
if ${DRY_RUN}; then
|
|
echo ${DOCKER_CMD[@]+"${DOCKER_CMD[@]}"}
|
|
else
|
|
${DOCKER_CMD[@]+"${DOCKER_CMD[@]}"}
|
|
fi
|