feat: rework docker workflow for gpu-first training
- 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
This commit is contained in:
+7
-2
@@ -57,8 +57,13 @@ COPY docs/ ./docs/
|
||||
COPY pyproject.toml .
|
||||
COPY README.md .
|
||||
|
||||
# Create non-root user
|
||||
RUN useradd -m astrai && chown -R astrai:astrai /app
|
||||
# Create non-root user matching the host uid/gid (passed via build args)
|
||||
ARG USER_UID=1000
|
||||
ARG USER_GID=1000
|
||||
RUN groupadd -g "${USER_GID}" astrai \
|
||||
&& useradd -m -u "${USER_UID}" -g astrai astrai \
|
||||
&& chown -R astrai:astrai /app
|
||||
ENV HOME=/home/astrai
|
||||
USER astrai
|
||||
|
||||
ENV PYTHONUNBUFFERED=1 \
|
||||
|
||||
+12
-4
@@ -5,7 +5,9 @@ services:
|
||||
dockerfile: Dockerfile
|
||||
args:
|
||||
CUDA_TAG: ${CUDA_TAG:-cu128}
|
||||
user: "${UID:-1000}:${GID:-1000}"
|
||||
USER_UID: ${ASTRAI_UID:-1000}
|
||||
USER_GID: ${ASTRAI_GID:-1000}
|
||||
user: "${ASTRAI_UID:-1000}:${ASTRAI_GID:-1000}"
|
||||
ports:
|
||||
- "8000:8000"
|
||||
volumes:
|
||||
@@ -33,7 +35,9 @@ services:
|
||||
dockerfile: Dockerfile
|
||||
args:
|
||||
CUDA_TAG: ${CUDA_TAG:-cu128}
|
||||
user: "${UID:-1000}:${GID:-1000}"
|
||||
USER_UID: ${ASTRAI_UID:-1000}
|
||||
USER_GID: ${ASTRAI_GID:-1000}
|
||||
user: "${ASTRAI_UID:-1000}:${ASTRAI_GID:-1000}"
|
||||
ports:
|
||||
- "8000:8000"
|
||||
volumes:
|
||||
@@ -54,8 +58,10 @@ services:
|
||||
dockerfile: Dockerfile
|
||||
args:
|
||||
CUDA_TAG: ${CUDA_TAG:-cu128}
|
||||
USER_UID: ${ASTRAI_UID:-1000}
|
||||
USER_GID: ${ASTRAI_GID:-1000}
|
||||
init: true
|
||||
user: "${UID:-1000}:${GID:-1000}"
|
||||
user: "${ASTRAI_UID:-1000}:${ASTRAI_GID:-1000}"
|
||||
volumes:
|
||||
- ${TRAIN_DATA_DIR:-./data}:/data:ro
|
||||
- ${TRAIN_MODEL_DIR:-./params}:/models/base:ro
|
||||
@@ -67,6 +73,8 @@ services:
|
||||
- CHECKPOINT_ROOT=/checkpoints
|
||||
- TRAIN_GPU_COUNT=${TRAIN_GPU_COUNT:-all}
|
||||
- CUDA_VISIBLE_DEVICES
|
||||
- NCCL_P2P_DISABLE
|
||||
- NCCL_NET_GDR_LEVEL
|
||||
entrypoint: ["bash", "/app/scripts/docker/train-entrypoint.sh"]
|
||||
ipc: ${TRAIN_IPC_MODE:-host}
|
||||
stop_grace_period: ${TRAIN_STOP_GRACE_PERIOD:-10m}
|
||||
@@ -81,5 +89,5 @@ services:
|
||||
reservations:
|
||||
devices:
|
||||
- driver: nvidia
|
||||
count: ${TRAIN_GPU_COUNT:-all}
|
||||
count: all
|
||||
capabilities: [gpu]
|
||||
|
||||
+67
-229
@@ -1,257 +1,95 @@
|
||||
#!/bin/bash
|
||||
|
||||
# AstrAI Docker Script
|
||||
# Build and manage Docker images
|
||||
|
||||
set -e
|
||||
|
||||
# Colors
|
||||
RED='\033[0;31m'
|
||||
GREEN='\033[0;32m'
|
||||
YELLOW='\033[1;33m'
|
||||
BLUE='\033[0;34m'
|
||||
NC='\033[0m' # No Color
|
||||
|
||||
# Default values
|
||||
IMAGE_NAME="astrai"
|
||||
IMAGE_TAG="latest"
|
||||
REGISTRY=""
|
||||
CONTAINER_ID=""
|
||||
|
||||
# Print colored messages
|
||||
print_info() {
|
||||
echo -e "${BLUE}[INFO]${NC} $1"
|
||||
}
|
||||
|
||||
print_success() {
|
||||
echo -e "${GREEN}[SUCCESS]${NC} $1"
|
||||
}
|
||||
|
||||
print_error() {
|
||||
echo -e "${RED}[ERROR]${NC} $1"
|
||||
}
|
||||
|
||||
print_warning() {
|
||||
echo -e "${YELLOW}[WARNING]${NC} $1"
|
||||
}
|
||||
|
||||
# Check if Docker is installed
|
||||
check_docker() {
|
||||
if ! command -v docker &> /dev/null; then
|
||||
print_error "Docker is not installed"
|
||||
exit 1
|
||||
fi
|
||||
print_success "Docker version: $(docker --version)"
|
||||
}
|
||||
|
||||
# Build Docker image
|
||||
build_image() {
|
||||
local dockerfile="${1:-Dockerfile}"
|
||||
local context="${2:-.}"
|
||||
|
||||
if [ ! -f "$dockerfile" ]; then
|
||||
print_error "Dockerfile not found: $dockerfile"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
print_info "Building Docker image: ${IMAGE_NAME}:${IMAGE_TAG}"
|
||||
docker build -t "${IMAGE_NAME}:${IMAGE_TAG}" -f "$dockerfile" "$context"
|
||||
print_success "Image built successfully"
|
||||
}
|
||||
|
||||
# Run container
|
||||
run_container() {
|
||||
local port="${1:-8000}"
|
||||
local gpu="${2:-false}"
|
||||
|
||||
print_info "Running container on port $port..."
|
||||
|
||||
if [ "$gpu" = true ]; then
|
||||
docker run --gpus all -p "${port}:8000" "${IMAGE_NAME}:${IMAGE_TAG}"
|
||||
else
|
||||
docker run -p "${port}:8000" "${IMAGE_NAME}:${IMAGE_TAG}"
|
||||
fi
|
||||
}
|
||||
|
||||
# Push image to registry
|
||||
push_image() {
|
||||
if [ -z "$REGISTRY" ]; then
|
||||
print_error "Registry not set. Use --registry option"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
local full_tag="${REGISTRY}/${IMAGE_NAME}:${IMAGE_TAG}"
|
||||
print_info "Tagging image: ${full_tag}"
|
||||
docker tag "${IMAGE_NAME}:${IMAGE_TAG}" "$full_tag"
|
||||
|
||||
print_info "Pushing image to registry..."
|
||||
docker push "$full_tag"
|
||||
print_success "Image pushed successfully"
|
||||
}
|
||||
|
||||
# Remove image
|
||||
remove_image() {
|
||||
print_info "Removing image: ${IMAGE_NAME}:${IMAGE_TAG}"
|
||||
docker rmi "${IMAGE_NAME}:${IMAGE_TAG}" 2>/dev/null || print_warning "Image not found"
|
||||
print_success "Image removed"
|
||||
}
|
||||
|
||||
# Show image info
|
||||
show_info() {
|
||||
print_info "Image information:"
|
||||
docker images "${IMAGE_NAME}"
|
||||
}
|
||||
|
||||
# Show logs
|
||||
show_logs() {
|
||||
local container_id="$1"
|
||||
if [ -z "$container_id" ]; then
|
||||
print_error "Container ID required"
|
||||
exit 1
|
||||
fi
|
||||
docker logs "$container_id"
|
||||
}
|
||||
|
||||
# Main function
|
||||
main() {
|
||||
echo "========================================"
|
||||
echo " AstrAI Docker Management"
|
||||
echo "========================================"
|
||||
echo ""
|
||||
|
||||
COMMAND=""
|
||||
DOCKERFILE="Dockerfile"
|
||||
CONTEXT="."
|
||||
IMAGE_NAME="${ASTRAI_IMAGE:-astrai}"
|
||||
IMAGE_TAG="${ASTRAI_TAG:-latest}"
|
||||
PORT="8000"
|
||||
GPU=false
|
||||
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=""
|
||||
|
||||
# Parse arguments
|
||||
while [[ $# -gt 0 ]]; do
|
||||
case $1 in
|
||||
build)
|
||||
COMMAND="build"
|
||||
case "$1" in
|
||||
build|run)
|
||||
command="$1"
|
||||
shift
|
||||
;;
|
||||
run)
|
||||
COMMAND="run"
|
||||
shift
|
||||
;;
|
||||
push)
|
||||
COMMAND="push"
|
||||
shift
|
||||
;;
|
||||
remove|rm)
|
||||
COMMAND="remove"
|
||||
shift
|
||||
;;
|
||||
info)
|
||||
COMMAND="info"
|
||||
shift
|
||||
;;
|
||||
logs)
|
||||
COMMAND="logs"
|
||||
shift
|
||||
;;
|
||||
--image)
|
||||
IMAGE_NAME="$2"
|
||||
shift 2
|
||||
;;
|
||||
--tag)
|
||||
IMAGE_TAG="$2"
|
||||
shift 2
|
||||
;;
|
||||
--registry)
|
||||
REGISTRY="$2"
|
||||
shift 2
|
||||
;;
|
||||
--dockerfile)
|
||||
DOCKERFILE="$2"
|
||||
shift 2
|
||||
;;
|
||||
--context)
|
||||
CONTEXT="$2"
|
||||
shift 2
|
||||
;;
|
||||
--port)
|
||||
PORT="$2"
|
||||
shift 2
|
||||
;;
|
||||
--container)
|
||||
CONTAINER_ID="$2"
|
||||
shift 2
|
||||
;;
|
||||
--gpu)
|
||||
GPU=true
|
||||
shift
|
||||
;;
|
||||
--help)
|
||||
echo "Usage: $0 <command> [options]"
|
||||
echo ""
|
||||
echo "Commands:"
|
||||
echo " build Build Docker image"
|
||||
echo " run Run container"
|
||||
echo " push Push image to registry"
|
||||
echo " remove Remove image"
|
||||
echo " info Show image information"
|
||||
echo " logs Show container logs"
|
||||
echo ""
|
||||
echo "Options:"
|
||||
echo " --image NAME Image name (default: astrai)"
|
||||
echo " --tag TAG Image tag (default: latest)"
|
||||
echo " --registry URL Registry URL for push"
|
||||
echo " --dockerfile FILE Dockerfile path (default: Dockerfile)"
|
||||
echo " --context PATH Build context (default: .)"
|
||||
echo " --port PORT Port for run (default: 8000)"
|
||||
echo " --container ID Container ID for logs"
|
||||
echo " --gpu Enable GPU support"
|
||||
echo " --help Show this help message"
|
||||
echo ""
|
||||
echo "Examples:"
|
||||
echo " $0 build"
|
||||
echo " $0 build --tag v1.0.0"
|
||||
echo " $0 run --port 8080"
|
||||
echo " $0 run --gpu"
|
||||
echo " $0 logs --container abc123"
|
||||
echo " $0 push --registry ghcr.io/username"
|
||||
--no-gpu)
|
||||
GPU=false
|
||||
shift
|
||||
;;
|
||||
--port)
|
||||
PORT="$2"
|
||||
shift 2
|
||||
;;
|
||||
--)
|
||||
shift
|
||||
RUN_ARGS=("$@")
|
||||
break
|
||||
;;
|
||||
-h|--help)
|
||||
usage
|
||||
exit 0
|
||||
;;
|
||||
*)
|
||||
if [ -z "$COMMAND" ]; then
|
||||
print_error "Unknown command: $1"
|
||||
echo "Unknown option: $1" >&2
|
||||
usage >&2
|
||||
exit 1
|
||||
fi
|
||||
shift
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
check_docker
|
||||
|
||||
case "$COMMAND" in
|
||||
case "$command" in
|
||||
build)
|
||||
build_image "$DOCKERFILE" "$CONTEXT"
|
||||
build_image
|
||||
;;
|
||||
run)
|
||||
run_container "$PORT" "$GPU"
|
||||
;;
|
||||
push)
|
||||
push_image
|
||||
;;
|
||||
remove)
|
||||
remove_image
|
||||
;;
|
||||
info)
|
||||
show_info
|
||||
;;
|
||||
logs)
|
||||
show_logs "$CONTAINER_ID"
|
||||
;;
|
||||
"")
|
||||
print_error "No command specified. Use --help for usage"
|
||||
exit 1
|
||||
run_container "${RUN_ARGS[@]}"
|
||||
;;
|
||||
*)
|
||||
print_error "Unknown command: $COMMAND"
|
||||
echo "No command specified. Use --help for usage" >&2
|
||||
usage >&2
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
|
||||
+6
-2
@@ -42,8 +42,9 @@ EOF
|
||||
load_env() {
|
||||
if [[ -f "${ENV_FILE}" ]]; then
|
||||
set -a
|
||||
# UID/GID are readonly in bash; compose gets them via ASTRAI_UID/GID in compose()
|
||||
# shellcheck disable=SC1090
|
||||
source "${ENV_FILE}"
|
||||
source <(grep -v -E '^[[:space:]]*(UID|GID)=' "${ENV_FILE}")
|
||||
set +a
|
||||
fi
|
||||
|
||||
@@ -75,7 +76,10 @@ compose() {
|
||||
if [[ -f "${ENV_FILE}" ]]; then
|
||||
command+=(--env-file "${ENV_FILE}")
|
||||
fi
|
||||
"${command[@]}" "$@"
|
||||
|
||||
# Inject the host user into compose so container processes share the
|
||||
# checkpoint directory ownership (bash UID/GID are readonly).
|
||||
ASTRAI_UID="$(id -u)" ASTRAI_GID="$(id -g)" "${command[@]}" "$@"
|
||||
}
|
||||
|
||||
init_environment() {
|
||||
|
||||
Reference in New Issue
Block a user