ci: 添加 Docker 镜像自动构建工作流

.dockerignore

@@ -3,6 +3,7 @@
 
 # Allow necessary files
 !astrai/
-!scripts/tools/
+!scripts/
+!assets/
 !pyproject.toml
 !README.md

.github/workflows/docker.yml

@@ -0,0 +1,50 @@
+name: Build and Push Docker Image
+
+on:
+  push:
+    tags:
+      - 'v*'
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      packages: write
+
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - name: Set up QEMU
+        uses: docker/setup-qemu-action@v3
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+
+      - name: Login to GitHub Container Registry
+        uses: docker/login-action@v3
+        with:
+          registry: ghcr.io
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Extract metadata
+        id: meta
+        uses: docker/metadata-action@v5
+        with:
+          images: ghcr.io/${{ github.repository }}
+          tags: |
+            type=ref,event=tag
+            type=raw,value=latest
+
+      - name: Build and push
+        uses: docker/build-push-action@v5
+        with:
+          context: .
+          platforms: linux/amd64
+          push: true
+          tags: ${{ steps.meta.outputs.tags }}
+          labels: ${{ steps.meta.outputs.labels }}
+          cache-from: type=gha
+          cache-to: type=gha,mode=max

Dockerfile

@@ -1,49 +1,54 @@
-# AstrAI Dockerfile
-# Multi-stage build for optimized image size
+# AstrAI Dockerfile - Multi-stage Build (Optimized)
 
-# Build stage
-FROM python:3.12-slim AS builder
+# Build stage - use base image with minimal build tools
+FROM nvidia/cuda:12.6.0-base-ubuntu24.04 AS builder
 
 WORKDIR /app
 
-# Install build dependencies
-RUN apt-get update && apt-get install -y --no-install-recommends \
-    build-essential \
+# Install Python 3.12 and minimal build dependencies
+RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
+    python3.12 \
+    python3.12-dev \
+    python3.12-venv \
+    gcc \
+    g++ \
     && rm -rf /var/lib/apt/lists/*
 
-# Copy project files first for version extraction
+# Create isolated virtual environment
+RUN python3.12 -m venv --copies /opt/venv
+ENV PATH="/opt/venv/bin:$PATH"
+
+# Copy source code and install dependencies
 COPY astrai/ ./astrai/
 COPY pyproject.toml .
-
-# Install dependencies
 RUN pip install --no-cache-dir --upgrade pip \
-    && pip install --no-cache-dir .[dev]
+    && pip install --no-cache-dir . \
+    --extra-index-url https://download.pytorch.org/whl/cu126
 
 # Production stage
-FROM python:3.12-slim AS production
+FROM nvidia/cuda:12.6.0-base-ubuntu24.04 AS production
 
 WORKDIR /app
 
-# Install runtime dependencies
-RUN apt-get update && apt-get install -y --no-install-recommends \
+# Install Python 3.12 runtime
+RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
+    python3.12 \
     && rm -rf /var/lib/apt/lists/*
 
-# Copy installed packages from builder
-COPY --from=builder /usr/local/lib/python3.12/site-packages /usr/local/lib/python3.12/site-packages
-COPY --from=builder /usr/local/bin /usr/local/bin
+# Copy virtual environment from builder
+COPY --from=builder /opt/venv /opt/venv
+ENV PATH="/opt/venv/bin:$PATH"
 
 # Copy application code
 COPY astrai/ ./astrai/
-COPY scripts/tools/ ./scripts/tools/
+COPY scripts/ ./scripts/
+COPY assets/ ./assets/
 COPY pyproject.toml .
+COPY README.md .
 
 # Create non-root user
-RUN useradd -m -u 1000 astrai && chown -R astrai:astrai /app
+RUN useradd -m astrai && chown -R astrai:astrai /app
 USER astrai
 
-# Set environment variables
-ENV PYTHONUNBUFFERED=1
-ENV PYTHONDONTWRITEBYTECODE=1
-
-# Default command
-CMD ["python", "-m", "astrai.inference.server"]
+ENV PYTHONUNBUFFERED=1 \
+    PYTHONDONTWRITEBYTECODE=1

README.md

@@ -84,6 +84,30 @@ python scripts/tools/train.py \
 python scripts/tools/generate.py --param_path=/path/to/param_path
 ```
 
+#### Docker
+
+Build and run with Docker (recommended for GPU environments):
+
+```bash
+# Build image
+docker build -t astrai:latest .
+
+# Run with GPU support
+docker run --gpus all -it astrai:latest
+
+# Run with specific GPUs
+docker run --gpus '"device=0,1"' -it astrai:latest
+
+# Run inference server
+docker run --gpus all -p 8000:8000 astrai:latest \
+  python -m scripts.tools.server --port 8000 --device cuda
+
+# Run with volume mount for data
+docker run --gpus all -v /path/to/data:/data -it astrai:latest
+```
+
+> **Note**: `--gpus all` is required for CUDA support. Without it, `torch.cuda.is_available()` will return `False`.
+
 #### Start HTTP Server
 
 Start the inference server with OpenAI-compatible HTTP API:

assets/docs/README-zh-CN.md

@@ -85,6 +85,30 @@ python scripts/tools/train.py \
 python scripts/tools/generate.py --param_path=/path/to/param_path
 ```
 
+#### Docker
+
+使用 Docker 构建和运行（推荐用于 GPU 环境）：
+
+```bash
+# 构建镜像
+docker build -t astrai:latest .
+
+# 启用 GPU 运行
+docker run --gpus all -it astrai:latest
+
+# 指定特定 GPU
+docker run --gpus '"device=0,1"' -it astrai:latest
+
+# 运行推理服务
+docker run --gpus all -p 8000:8000 astrai:latest \
+  python -m scripts.tools.server --port 8000 --device cuda
+
+# 挂载数据卷
+docker run --gpus all -v /path/to/data:/data -it astrai:latest
+```
+
+> **注意**: 必须使用 `--gpus all` 才能启用 CUDA 支持，否则 `torch.cuda.is_available()` 将返回 `False`。
+
 #### 启动 HTTP 服务
 
 启动推理服务器，支持 OpenAI 兼容的 HTTP API：

scripts/demo/download.py

@@ -1,13 +1,41 @@
+import argparse
 from pathlib import Path
 
 from huggingface_hub import snapshot_download
 
 PROJECT_ROOT = Path(__file__).resolve().parents[2]
-PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
+DEFAULT_LOCAL_DIR = Path(PROJECT_ROOT, "params")
+DEFAULT_REPO_ID = "ViperEk/KHAOSZ"
 
 if __name__ == "__main__":
-    snapshot_download(
-        repo_id="ViperEk/KHAOSZ",
-        local_dir=PARAMETER_ROOT,
-        force_download=True,
+    parser = argparse.ArgumentParser(
+        description="Download model parameters from HuggingFace"
     )
+    parser.add_argument(
+        "--repo-id",
+        type=str,
+        default=DEFAULT_REPO_ID,
+        help=f"HuggingFace repo ID (default: {DEFAULT_REPO_ID})",
+    )
+    parser.add_argument(
+        "--local-dir",
+        type=Path,
+        default=DEFAULT_LOCAL_DIR,
+        help=f"Local directory to save model (default: {DEFAULT_LOCAL_DIR})",
+    )
+    parser.add_argument(
+        "--force",
+        action="store_true",
+        help="Force download even if files exist",
+    )
+    args = parser.parse_args()
+
+    print(f"Downloading model from {args.repo_id} to {args.local_dir}")
+
+    snapshot_download(
+        repo_id=args.repo_id,
+        local_dir=args.local_dir,
+        force_download=args.force,
+    )
+
+    print("Download complete!")