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

Dockerfile

@@ -57,8 +57,13 @@ COPY docs/ ./docs/
 COPY pyproject.toml .
 COPY README.md .
 
-# Create non-root user
+# Create non-root user matching the host uid/gid (passed via build args)
-RUN useradd -m astrai && chown -R astrai:astrai /app
+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 \

docker-compose.yml

@@ -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]

scripts/docker.sh

@@ -1,257 +1,95 @@
 #!/bin/bash
-
-# AstrAI Docker Script
-# Build and manage Docker images
-
 set -e
 
-# Colors
+IMAGE_NAME="${ASTRAI_IMAGE:-astrai}"
-RED='\033[0;31m'
+IMAGE_TAG="${ASTRAI_TAG:-latest}"
-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="."
 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
+        case "$1" in
-            build)
+            build|run)
-                COMMAND="build"
+                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)
+            --no-gpu)
-                echo "Usage: $0 <command> [options]"
+                GPU=false
-                echo ""
+                shift
-                echo "Commands:"
+                ;;
-                echo "  build     Build Docker image"
+            --port)
-                echo "  run       Run container"
+                PORT="$2"
-                echo "  push      Push image to registry"
+                shift 2
-                echo "  remove    Remove image"
+                ;;
-                echo "  info      Show image information"
+            --)
-                echo "  logs      Show container logs"
+                shift
-                echo ""
+                RUN_ARGS=("$@")
-                echo "Options:"
+                break
-                echo "  --image NAME       Image name (default: astrai)"
+                ;;
-                echo "  --tag TAG          Image tag (default: latest)"
+            -h|--help)
-                echo "  --registry URL     Registry URL for push"
+                usage
-                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"
                 exit 0
                 ;;
             *)
-                if [ -z "$COMMAND" ]; then
+                echo "Unknown option: $1" >&2
-                    print_error "Unknown command: $1"
+                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"
+            run_container "${RUN_ARGS[@]}"
-            ;;
-        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
             ;;
         *)
-            print_error "Unknown command: $COMMAND"
+            echo "No command specified. Use --help for usage" >&2
+            usage >&2
             exit 1
             ;;
     esac

scripts/train.sh

@@ -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() {