ViperEkura
71b6e3aaaf
- rewrite docker.sh with gpu default and --no-gpu override - inject host uid/gid via ASTRAI_UID/GID in train.sh compose() - filter readonly UID/GID lines when sourcing .env.train - build image user via USER_UID/USER_GID args matching host uid/gid - pass all GPUs (count: all) and filter by CUDA_VISIBLE_DEVICES inside the container - forward NCCL vars through compose environment
99 lines
2.0 KiB
Bash
Executable File
99 lines
2.0 KiB
Bash
Executable File
#!/bin/bash
|
|
set -e
|
|
|
|
IMAGE_NAME="${ASTRAI_IMAGE:-astrai}"
|
|
IMAGE_TAG="${ASTRAI_TAG:-latest}"
|
|
PORT="8000"
|
|
GPU=true
|
|
RUN_ARGS=()
|
|
|
|
usage() {
|
|
cat <<EOF
|
|
Usage: $0 <command> [options]
|
|
|
|
Commands:
|
|
build Build the image
|
|
run [--] [ARGS] Run a container; ARGS after -- are passed to the container
|
|
|
|
Options:
|
|
--gpu Enable GPU support (default)
|
|
--no-gpu Disable GPU support
|
|
--port PORT Host port for run (default: 8000)
|
|
-h, --help Show this help
|
|
|
|
Environment:
|
|
ASTRAI_IMAGE Image name (default: astrai)
|
|
ASTRAI_TAG Image tag (default: latest)
|
|
|
|
Examples:
|
|
$0 build
|
|
$0 run
|
|
$0 run --port 8080 -- python -m scripts.tools.server --port 8000 --device cuda
|
|
EOF
|
|
}
|
|
|
|
build_image() {
|
|
docker build -t "${IMAGE_NAME}:${IMAGE_TAG}" .
|
|
}
|
|
|
|
run_container() {
|
|
local gpu_args=()
|
|
[ "$GPU" = true ] && gpu_args=(--gpus all)
|
|
docker run "${gpu_args[@]}" -p "${PORT}:8000" "${IMAGE_NAME}:${IMAGE_TAG}" "$@"
|
|
}
|
|
|
|
main() {
|
|
local command=""
|
|
|
|
while [[ $# -gt 0 ]]; do
|
|
case "$1" in
|
|
build|run)
|
|
command="$1"
|
|
shift
|
|
;;
|
|
--gpu)
|
|
GPU=true
|
|
shift
|
|
;;
|
|
--no-gpu)
|
|
GPU=false
|
|
shift
|
|
;;
|
|
--port)
|
|
PORT="$2"
|
|
shift 2
|
|
;;
|
|
--)
|
|
shift
|
|
RUN_ARGS=("$@")
|
|
break
|
|
;;
|
|
-h|--help)
|
|
usage
|
|
exit 0
|
|
;;
|
|
*)
|
|
echo "Unknown option: $1" >&2
|
|
usage >&2
|
|
exit 1
|
|
;;
|
|
esac
|
|
done
|
|
|
|
case "$command" in
|
|
build)
|
|
build_image
|
|
;;
|
|
run)
|
|
run_container "${RUN_ARGS[@]}"
|
|
;;
|
|
*)
|
|
echo "No command specified. Use --help for usage" >&2
|
|
usage >&2
|
|
exit 1
|
|
;;
|
|
esac
|
|
}
|
|
|
|
main "$@"
|