mirror of
https://github.com/NVIDIA/nvbench.git
synced 2026-03-14 20:27:24 +00:00
307 lines
10 KiB
Bash
Executable File
307 lines
10 KiB
Bash
Executable File
#!/usr/bin/env bash
|
|
|
|
set -euo pipefail
|
|
|
|
# Ensure the script is being executed in the nvbench/ root
|
|
cd "$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )/..";
|
|
|
|
print_help() {
|
|
echo "Usage: $0 [-c|--cuda <CUDA version>] [-H|--host <Host compiler>] [-d|--docker]"
|
|
echo "Launch a development container. If no CUDA version or Host compiler are specified,"
|
|
echo "the top-level devcontainer in .devcontainer/devcontainer.json will be used."
|
|
echo ""
|
|
echo "Options:"
|
|
echo " -c, --cuda Specify the CUDA version. E.g., 12.2"
|
|
echo " -H, --host Specify the host compiler. E.g., gcc12"
|
|
echo " -d, --docker Launch the development environment in Docker directly without using VSCode."
|
|
echo " --gpus gpu-request GPU devices to add to the container ('all' to pass all GPUs)."
|
|
echo " -e, --env list Set additional container environment variables."
|
|
echo " -v, --volume list Bind mount a volume."
|
|
echo " -h, --help Display this help message and exit."
|
|
}
|
|
|
|
# Assign variable one scope above the caller
|
|
# Usage: local "$1" && _upvar $1 "value(s)"
|
|
# Param: $1 Variable name to assign value to
|
|
# Param: $* Value(s) to assign. If multiple values, an array is
|
|
# assigned, otherwise a single value is assigned.
|
|
# See: http://fvue.nl/wiki/Bash:_Passing_variables_by_reference
|
|
_upvar() {
|
|
if unset -v "$1"; then
|
|
if (( $# == 2 )); then
|
|
eval $1=\"\$2\";
|
|
else
|
|
eval $1=\(\"\${@:2}\"\);
|
|
fi;
|
|
fi
|
|
}
|
|
|
|
parse_options() {
|
|
local -;
|
|
set -euo pipefail;
|
|
|
|
# Read the name of the variable in which to return unparsed arguments
|
|
local UNPARSED="${!#}";
|
|
# Splice the unparsed arguments variable name from the arguments list
|
|
set -- "${@:1:$#-1}";
|
|
|
|
local OPTIONS=c:e:H:dhv
|
|
local LONG_OPTIONS=cuda:,env:,host:,gpus:,volume:,docker,help
|
|
# shellcheck disable=SC2155
|
|
local PARSED_OPTIONS=$(getopt -n "$0" -o "${OPTIONS}" --long "${LONG_OPTIONS}" -- "$@")
|
|
|
|
# shellcheck disable=SC2181
|
|
if [[ $? -ne 0 ]]; then
|
|
exit 1
|
|
fi
|
|
|
|
eval set -- "${PARSED_OPTIONS}"
|
|
|
|
while true; do
|
|
case "$1" in
|
|
-c|--cuda)
|
|
cuda_version="$2"
|
|
shift 2
|
|
;;
|
|
-e|--env)
|
|
env_vars+=("$1" "$2")
|
|
shift 2
|
|
;;
|
|
-H|--host)
|
|
host_compiler="$2"
|
|
shift 2
|
|
;;
|
|
--gpus)
|
|
gpu_request="$2"
|
|
shift 2
|
|
;;
|
|
-d|--docker)
|
|
docker_mode=true
|
|
shift
|
|
;;
|
|
-h|--help)
|
|
print_help
|
|
exit 0
|
|
;;
|
|
-v|--volume)
|
|
volumes+=("$1" "$2")
|
|
shift 2
|
|
;;
|
|
--)
|
|
shift
|
|
_upvar "${UNPARSED}" "${@}"
|
|
break
|
|
;;
|
|
*)
|
|
echo "Invalid option: $1"
|
|
print_help
|
|
exit 1
|
|
;;
|
|
esac
|
|
done
|
|
}
|
|
|
|
# shellcheck disable=SC2155
|
|
launch_docker() {
|
|
local -;
|
|
set -euo pipefail
|
|
|
|
inline_vars() {
|
|
cat - \
|
|
`# inline local workspace folder` \
|
|
| sed "s@\${localWorkspaceFolder}@$(pwd)@g" \
|
|
`# inline local workspace folder basename` \
|
|
| sed "s@\${localWorkspaceFolderBasename}@$(basename "$(pwd)")@g" \
|
|
`# inline container workspace folder` \
|
|
| sed "s@\${containerWorkspaceFolder}@${WORKSPACE_FOLDER:-}@g" \
|
|
`# inline container workspace folder basename` \
|
|
| sed "s@\${containerWorkspaceFolderBasename}@$(basename "${WORKSPACE_FOLDER:-}")@g" \
|
|
`# translate local envvars to shell syntax` \
|
|
| sed -r 's/\$\{localEnv:([^\:]*):?(.*)\}/${\1:-\2}/g'
|
|
}
|
|
|
|
args_to_path() {
|
|
local -a keys=("${@}")
|
|
keys=("${keys[@]/#/[}")
|
|
keys=("${keys[@]/%/]}")
|
|
echo "$(IFS=; echo "${keys[*]}")"
|
|
}
|
|
|
|
json_string() {
|
|
python3 -c "import json,sys; print(json.load(sys.stdin)$(args_to_path "${@}"))" 2>/dev/null | inline_vars
|
|
}
|
|
|
|
json_array() {
|
|
python3 -c "import json,sys; [print(f'\"{x}\"') for x in json.load(sys.stdin)$(args_to_path "${@}")]" 2>/dev/null | inline_vars
|
|
}
|
|
|
|
json_map() {
|
|
python3 -c "import json,sys; [print(f'{k}=\"{v}\"') for k,v in json.load(sys.stdin)$(args_to_path "${@}").items()]" 2>/dev/null | inline_vars
|
|
}
|
|
|
|
devcontainer_metadata_json() {
|
|
docker inspect --type image --format '{{json .Config.Labels}}' "$DOCKER_IMAGE" \
|
|
| json_string '"devcontainer.metadata"'
|
|
}
|
|
|
|
###
|
|
# Read relevant values from devcontainer.json
|
|
###
|
|
|
|
local devcontainer_json="${path}/devcontainer.json";
|
|
|
|
# Read image
|
|
local DOCKER_IMAGE="$(json_string '"image"' < "${devcontainer_json}")"
|
|
# Always pull the latest copy of the image
|
|
docker pull "$DOCKER_IMAGE"
|
|
|
|
# Read workspaceFolder
|
|
local WORKSPACE_FOLDER="$(json_string '"workspaceFolder"' < "${devcontainer_json}")"
|
|
# Read remoteUser
|
|
local REMOTE_USER="$(json_string '"remoteUser"' < "${devcontainer_json}")"
|
|
# If remoteUser isn't in our devcontainer.json, read it from the image's "devcontainer.metadata" label
|
|
if test -z "${REMOTE_USER:-}"; then
|
|
REMOTE_USER="$(devcontainer_metadata_json | json_string "-1" '"remoteUser"')"
|
|
fi
|
|
# Read runArgs
|
|
local -a RUN_ARGS="($(json_array '"runArgs"' < "${devcontainer_json}"))"
|
|
# Read initializeCommand
|
|
local -a INITIALIZE_COMMAND="($(json_array '"initializeCommand"' < "${devcontainer_json}"))"
|
|
# Read containerEnv
|
|
local -a ENV_VARS="($(json_map '"containerEnv"' < "${devcontainer_json}" | sed -r 's/(.*)=(.*)/--env \1=\2/'))"
|
|
# Read mounts
|
|
local -a MOUNTS="($(
|
|
tee < "${devcontainer_json}" \
|
|
1>/dev/null \
|
|
>(json_array '"mounts"') \
|
|
>(json_string '"workspaceMount"') \
|
|
| xargs -r -I% echo --mount '%'
|
|
))"
|
|
|
|
###
|
|
# Update run arguments and container environment variables
|
|
###
|
|
|
|
# Only pass `-it` if the shell is a tty
|
|
if ! ${CI:-'false'} && tty >/dev/null 2>&1 && (exec </dev/tty); then
|
|
RUN_ARGS+=("-it")
|
|
fi
|
|
|
|
for flag in rm init; do
|
|
if [[ " ${RUN_ARGS[*]} " != *" --${flag} "* ]]; then
|
|
RUN_ARGS+=("--${flag}")
|
|
fi
|
|
done
|
|
|
|
# Prefer the user-provided --gpus argument
|
|
if test -n "${gpu_request:-}"; then
|
|
RUN_ARGS+=(--gpus "${gpu_request}")
|
|
else
|
|
# Otherwise read and infer from hostRequirements.gpu
|
|
local GPU_REQUEST="$(json_string '"hostRequirements"' '"gpu"' < "${devcontainer_json}")"
|
|
if test "${GPU_REQUEST:-false}" = true; then
|
|
RUN_ARGS+=(--gpus all)
|
|
elif test "${GPU_REQUEST:-false}" = optional && \
|
|
command -v nvidia-container-runtime >/dev/null 2>&1; then
|
|
RUN_ARGS+=(--gpus all)
|
|
fi
|
|
fi
|
|
|
|
RUN_ARGS+=(--workdir "${WORKSPACE_FOLDER:-/home/coder/nvbench}")
|
|
|
|
if test -n "${REMOTE_USER:-}"; then
|
|
ENV_VARS+=(--env NEW_UID="$(id -u)")
|
|
ENV_VARS+=(--env NEW_GID="$(id -g)")
|
|
ENV_VARS+=(--env REMOTE_USER="$REMOTE_USER")
|
|
RUN_ARGS+=(-u root:root)
|
|
RUN_ARGS+=(--entrypoint "${WORKSPACE_FOLDER:-/home/coder/nvbench}/.devcontainer/docker-entrypoint.sh")
|
|
fi
|
|
|
|
if test -n "${SSH_AUTH_SOCK:-}"; then
|
|
ENV_VARS+=(--env "SSH_AUTH_SOCK=/tmp/ssh-auth-sock")
|
|
MOUNTS+=(--mount "source=${SSH_AUTH_SOCK},target=/tmp/ssh-auth-sock,type=bind")
|
|
fi
|
|
|
|
# Append user-provided volumes
|
|
if test -v volumes && test ${#volumes[@]} -gt 0; then
|
|
MOUNTS+=("${volumes[@]}")
|
|
fi
|
|
|
|
# Append user-provided envvars
|
|
if test -v env_vars && test ${#env_vars[@]} -gt 0; then
|
|
ENV_VARS+=("${env_vars[@]}")
|
|
fi
|
|
|
|
# Run the initialize command before starting the container
|
|
if test "${#INITIALIZE_COMMAND[@]}" -gt 0; then
|
|
eval "${INITIALIZE_COMMAND[*]@Q}"
|
|
fi
|
|
|
|
exec docker run \
|
|
"${RUN_ARGS[@]}" \
|
|
"${ENV_VARS[@]}" \
|
|
"${MOUNTS[@]}" \
|
|
"${DOCKER_IMAGE}" \
|
|
"$@"
|
|
}
|
|
|
|
launch_vscode() {
|
|
local -;
|
|
set -euo pipefail;
|
|
# Since Visual Studio Code allows only one instance per `devcontainer.json`,
|
|
# this code prepares a unique temporary directory structure for each launch of a devcontainer.
|
|
# By doing so, it ensures that multiple instances of the same environment can be run
|
|
# simultaneously. The script replicates the `devcontainer.json` from the desired CUDA
|
|
# and compiler environment into this temporary directory, adjusting paths to ensure the
|
|
# correct workspace is loaded. A special URL is then generated to instruct VSCode to
|
|
# launch the development container using this temporary configuration.
|
|
local workspace="$(basename "$(pwd)")"
|
|
local tmpdir="$(mktemp -d)/${workspace}"
|
|
mkdir -p "${tmpdir}"
|
|
mkdir -p "${tmpdir}/.devcontainer"
|
|
cp -arL "${path}/devcontainer.json" "${tmpdir}/.devcontainer"
|
|
sed -i "s@\\${localWorkspaceFolder}@$(pwd)@g" "${tmpdir}/.devcontainer/devcontainer.json"
|
|
local path="${tmpdir}"
|
|
local hash="$(echo -n "${path}" | xxd -pu - | tr -d '[:space:]')"
|
|
local url="vscode://vscode-remote/dev-container+${hash}/home/coder/nvbench"
|
|
|
|
local launch=""
|
|
if type open >/dev/null 2>&1; then
|
|
launch="open"
|
|
elif type xdg-open >/dev/null 2>&1; then
|
|
launch="xdg-open"
|
|
fi
|
|
|
|
if [ -n "${launch}" ]; then
|
|
echo "Launching VSCode Dev Container URL: ${url}"
|
|
code --new-window "${tmpdir}"
|
|
exec "${launch}" "${url}" >/dev/null 2>&1
|
|
fi
|
|
}
|
|
|
|
main() {
|
|
local -a unparsed;
|
|
parse_options "$@" unparsed;
|
|
set -- "${unparsed[@]}";
|
|
|
|
# If no CTK/Host compiler are provided, just use the default environment
|
|
if [[ -z ${cuda_version:-} ]] && [[ -z ${host_compiler:-} ]]; then
|
|
path=".devcontainer"
|
|
else
|
|
path=".devcontainer/cuda${cuda_version}-${host_compiler}"
|
|
if [[ ! -f "${path}/devcontainer.json" ]]; then
|
|
echo "Unknown CUDA [${cuda_version}] compiler [${host_compiler}] combination"
|
|
echo "Requested devcontainer ${path}/devcontainer.json does not exist"
|
|
exit 1
|
|
fi
|
|
fi
|
|
|
|
if ${docker_mode:-'false'}; then
|
|
launch_docker "$@"
|
|
else
|
|
launch_vscode
|
|
fi
|
|
}
|
|
|
|
main "$@"
|