| #!/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] |
| # [--shell SHELL] [--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 |
| |
| ############################################## |
| ### Section 1: Constants and defaults ### |
| ############################################## |
| |
| 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= |
| USER_SHELL= |
| |
| # 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 |
| |
| ############################################## |
| ### Section 2: Helper functions ### |
| ############################################## |
| |
| function show_usage() { |
| cat <<EOF |
| Usage: docker/bash.sh [-i|--interactive] [--net=host] [-t|--tty] |
| [--shell SHELL] [--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). |
| |
| --shell SHELL |
| |
| Specify a shell (e.g. --shell zsh) to use instead of bash |
| when entering auto-interactive mode (no command or command |
| is "bash"). |
| |
| --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. |
| |
| Environment Variables: |
| |
| TVM_DEV_DOCKER_MOUNTS |
| |
| Space-separated list of additional mount specifications. |
| Each entry is either: |
| /path/to/dir - bind mount at same path inside container |
| /src:/dst - bind mount src to dst inside container |
| Source paths are validated; missing sources are skipped with a warning. |
| |
| EOF |
| } |
| |
| function parse_error() { |
| echo "$@" >&2 |
| show_usage >&2 |
| exit 1 |
| } |
| |
| # Detect if Docker daemon is running in rootless mode. |
| # Uses DOCKER_IS_ROOTLESS env var as override (for CI), otherwise auto-detects. |
| # Result is cached after first call. |
| _ROOTLESS_CACHED= |
| _ROOTLESS_RESULT=1 |
| function detect_rootless() { |
| if [ -n "${_ROOTLESS_CACHED}" ]; then |
| return $_ROOTLESS_RESULT |
| fi |
| _ROOTLESS_CACHED=1 |
| if [ -n "${DOCKER_IS_ROOTLESS:-}" ]; then |
| # env var override for CI |
| _ROOTLESS_RESULT=0 |
| return 0 |
| fi |
| if docker info -f '{{.SecurityOptions}}' 2>/dev/null | grep -q "rootless"; then |
| _ROOTLESS_RESULT=0 |
| return 0 |
| fi |
| _ROOTLESS_RESULT=1 |
| return 1 |
| } |
| |
| # Validate a mount source path exists, warn if missing. |
| # Usage: validate_mount /path/to/source |
| # Returns 0 if source exists, 1 if missing. |
| function validate_mount() { |
| local src="$1" |
| if [ -e "$src" ]; then |
| return 0 |
| else |
| echo "Warning: Mount source '$src' not found on host, skipping." >&2 |
| return 1 |
| fi |
| } |
| |
| |
| ############################################## |
| ### Section 3: Argument parsing ### |
| ############################################## |
| |
| # 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 |
| if [[ "$1" == "--tty" ]]; then |
| shift |
| else |
| # Short form: -t or combined like -th |
| eval $break_joined_flag |
| fi |
| ;; |
| |
| --shell) |
| if (( $# >= 2 )) && [[ -n "$2" ]]; then |
| USER_SHELL="$2" |
| shift 2 |
| else |
| parse_error 'ERROR: --shell requires a non-empty argument' |
| fi |
| ;; |
| |
| --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 |
| |
| # Smart defaults: "bash" command triggers interactive mode |
| if [[ ${COMMAND[@]+"${COMMAND[@]}"} = bash ]]; then |
| INTERACTIVE=true |
| TTY=true |
| fi |
| |
| |
| ############################################## |
| ### Section 4: Image resolution ### |
| ############################################## |
| |
| # Save user-set DOCKER_IS_ROOTLESS before dev_common.sh overwrites it. |
| # detect_rootless() should only see the user's explicit env var, not |
| # the auto-detection done by dev_common.sh. |
| _USER_ROOTLESS="${DOCKER_IS_ROOTLESS+__set__}" |
| _USER_ROOTLESS_VAL="${DOCKER_IS_ROOTLESS:-}" |
| |
| source "$(dirname $0)/dev_common.sh" || exit 2 |
| |
| # Restore: if user had set DOCKER_IS_ROOTLESS, keep their value; |
| # otherwise unset the one dev_common.sh created. |
| if [ "$_USER_ROOTLESS" = "__set__" ]; then |
| DOCKER_IS_ROOTLESS="$_USER_ROOTLESS_VAL" |
| else |
| unset DOCKER_IS_ROOTLESS |
| fi |
| unset _USER_ROOTLESS _USER_ROOTLESS_VAL |
| |
| 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 |
| |
| |
| ############################################## |
| ### Section 5: Mount construction ### |
| ############################################## |
| |
| # Working directories |
| DOCKER_FLAGS+=( --workdir "${REPO_MOUNT_POINT}" ) |
| DOCKER_MOUNT+=( --volume "${REPO_DIR}":"${REPO_MOUNT_POINT}" |
| --volume "${SCRIPT_DIR}":/docker |
| ) |
| |
| # Expose external directories from --mount flags |
| for MOUNT_DIR in ${MOUNT_DIRS[@]+"${MOUNT_DIRS[@]}"}; do |
| DOCKER_MOUNT+=( --volume "${MOUNT_DIR}:${MOUNT_DIR}" ) |
| done |
| |
| # Parse TVM_DEV_DOCKER_MOUNTS env var: space-separated mount specs |
| # Each entry is either /path (same path) or /src:/dst |
| if [[ -n "${TVM_DEV_DOCKER_MOUNTS:-}" ]]; then |
| for mount_spec in ${TVM_DEV_DOCKER_MOUNTS}; do |
| if [[ "$mount_spec" == *:* ]]; then |
| # src:dst format |
| mount_src="${mount_spec%%:*}" |
| if validate_mount "$mount_src"; then |
| DOCKER_MOUNT+=( --volume "$mount_spec" ) |
| fi |
| else |
| # Same path format |
| if validate_mount "$mount_spec"; then |
| DOCKER_MOUNT+=( --volume "$mount_spec:$mount_spec" ) |
| fi |
| fi |
| done |
| 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 |
| |
| |
| ############################################## |
| ### Section 6: Environment construction ### |
| ############################################## |
| |
| # 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 ! detect_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 |
| |
| # 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 |
| |
| |
| ############################################## |
| ### Section 7: GPU/device setup ### |
| ############################################## |
| |
| DOCKER_BINARY=docker |
| |
| # Use --gpus all for GPU/CUDA containers |
| if [[ "$USE_GPU" == "true" ]] && [[ "${DOCKER_IMAGE_NAME}" == *"gpu"* || "${DOCKER_IMAGE_NAME}" == *"cuda"* ]]; then |
| DOCKER_FLAGS+=( --gpus all ) |
| DOCKER_ENV+=( --env NVIDIA_DRIVER_CAPABILITIES=compute,graphics,utility ) |
| 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 |
| |
| |
| ############################################## |
| ### Section 8: Entry command ### |
| ############################################## |
| |
| # Apply --shell override: when command is "bash" (auto-interactive) |
| # and a custom shell was requested via --shell <name>, use that shell |
| if [[ ${COMMAND[@]+"${COMMAND[@]}"} = bash ]] && [[ -n "$USER_SHELL" ]]; then |
| COMMAND=( "$USER_SHELL" ) |
| fi |
| |
| # Build the entry command |
| if ! detect_rootless; then |
| # 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. |
| COMMAND=( |
| bash --login /docker/with_the_same_user |
| ${COMMAND[@]+"${COMMAND[@]}"} |
| ) |
| fi |
| |
| |
| ############################################## |
| ### Section 9: Final docker run ### |
| ############################################## |
| |
| # 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 |
