feat: add Docker support and reverse-proxy compatibility

Add Docker infrastructure for running the training WebUI in containers: - Dockerfile based on PyTorch CUDA base image with layer-cached deps - docker-compose.yml with GPU support and nginx reverse proxy - nginx.conf with WebSocket support for Gradio Code fixes for container environments: - Stream training subprocess stdout/stderr to Docker logs - Support GRADIO_ROOT_PATH env var for reverse proxy (nginx/Traefik) - Echo startup URL to stdout for container log discovery All changes are backward-compatible: without Docker or env vars, behavior is identical to before.
2026-09-20 10:53:32 +08:00 · 2026-08-11 19:15:10 +09:30
parent 616d3d3e63
commit 85df305079
5 changed files with 217 additions and 1 deletions
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -0,0 +1,51 @@
+# ─────────────────────────────────────────────────────────────────────
+# VoxCPM Training WebUI — Docker image
+# ─────────────────────────────────────────────────────────────────────
+# Base: PyTorch with CUDA for GPU-accelerated LoRA fine-tuning.
+# Build context should be the project root:
+#
+#   docker build -f docker/Dockerfile -t voxcpm-training .
+#
+# ─────────────────────────────────────────────────────────────────────
+FROM pytorch/pytorch:2.5.1-cuda12.4-cudnn9-devel
+
+LABEL maintainer="OpenBMB <openbmb@gmail.com>"
+LABEL description="VoxCPM LoRA Training WebUI with GPU support"
+
+# Avoid interactive prompts during package installation
+ENV DEBIAN_FRONTEND=noninteractive
+
+# System deps required by Python packages:
+#   git        — setuptools_scm needs it to resolve version in pyproject.toml
+#   libsndfile1 — C library backing the 'soundfile' Python package
+#   ffmpeg     — audio codec support for torchaudio/librosa
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        git \
+        libsndfile1 \
+        ffmpeg \
+    && rm -rf /var/lib/apt/lists/*
+
+WORKDIR /app
+
+# Layer 1: Install dependencies only (cached unless pyproject.toml changes)
+# Create a minimal package stub so pip can resolve deps without real source.
+COPY pyproject.toml /app/
+RUN mkdir -p /app/src/voxcpm && echo '__version__ = "0.0.0"' > /app/src/voxcpm/__init__.py
+ENV SETUPTOOLS_SCM_PRETEND_VERSION=0.0.0
+RUN pip install --no-cache-dir -e .
+
+# Layer 2: Copy full project source (cheap rebuild on code changes)
+COPY . /app/
+
+# Create default directories and declare volumes
+RUN mkdir -p /app/lora /app/models /app/output
+VOLUME ["/app/models", "/app/lora", "/app/output"]
+
+EXPOSE 7860
+
+# Environment variables for configuration
+ENV GRADIO_SERVER_PORT=7860
+ENV GRADIO_ROOT_PATH=""
+
+# Default: launch training WebUI
+CMD ["python", "lora_ft_webui.py"]
--- a/docker/README.md
+++ b/docker/README.md
@@ -0,0 +1,104 @@
+# Docker Support for VoxCPM Training WebUI
+
+Run the VoxCPM LoRA fine-tuning WebUI in a Docker container with full GPU support and nginx reverse proxy.
+
+## Prerequisites
+
+- Docker Engine 19.03+ with [NVIDIA Container Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html)
+- NVIDIA GPU with CUDA 12.4+ compatible drivers
+- At least 16 GB GPU VRAM (24 GB+ recommended for larger models)
+
+## Quick Start
+
+```bash
+# From the project root directory:
+docker compose -f docker/docker-compose.yml up --build
+```
+
+This starts:
+- **training-webui** — the Gradio-based training interface on port 7860
+- **nginx** — reverse proxy serving the WebUI at `http://localhost/webui/`
+
+Access the WebUI at **http://localhost/webui/**.
+
+## Direct Access (no proxy)
+
+If you want to bypass nginx and access Gradio directly:
+
+```bash
+docker compose -f docker/docker-compose.yml up --build training-webui
+```
+
+Set `GRADIO_ROOT_PATH=` (empty) in the compose file when running without the proxy, then access at `http://localhost:7860`.
+
+## Building Manually
+
+```bash
+# Build the image
+docker build -f docker/Dockerfile -t voxcpm-training .
+
+# Run with GPU access (no reverse proxy)
+docker run --gpus all -p 7860:7860 \
+    -v ./models:/app/models \
+    -v ./lora:/app/lora \
+    -v ./output:/app/output \
+    voxcpm-training
+```
+
+## Model Weights
+
+Models are **auto-downloaded** from HuggingFace Hub on first use. The `/app/models` volume persists them across container restarts so they don't need to be re-downloaded.
+
+To pre-populate (avoids download at startup):
+
+```
+models/
+├── openbmb__VoxCPM2/       # VoxCPM2 (preferred)
+└── openbmb__VoxCPM1.5/     # VoxCPM1.5 (fallback)
+```
+
+## Environment Variables
+
+| Variable | Default | Description |
+|----------|---------|-------------|
+| `GRADIO_SERVER_PORT` | `7860` | Port for the WebUI server |
+| `GRADIO_ROOT_PATH` | `""` | URL prefix when behind a reverse proxy (e.g., `/webui`) |
+
+## Reverse Proxy
+
+The included `docker-compose.yml` ships with an nginx reverse proxy that serves the WebUI at `/webui/`. The `GRADIO_ROOT_PATH=/webui` env var ensures Gradio generates correct URLs for assets and WebSocket connections.
+
+### Custom nginx config
+
+Edit `docker/nginx.conf` to change the location prefix or add TLS.
+
+### Traefik Example (labels)
+
+```yaml
+labels:
+  - "traefik.http.routers.voxcpm.rule=PathPrefix(`/webui`)"
+  - "traefik.http.services.voxcpm.loadbalancer.server.port=7860"
+```
+
+## Viewing Training Logs
+
+Training subprocess output is streamed to stdout, visible via:
+
+```bash
+docker compose -f docker/docker-compose.yml logs -f training-webui
+```
+
+## Volumes
+
+| Mount Point | Purpose |
+|-------------|---------|
+| `/app/models` | Pre-trained model weights (read-only OK) |
+| `/app/lora` | LoRA checkpoints — training output is saved here |
+| `/app/output` | Additional training artifacts |
+
+## Troubleshooting
+
+- **"no NVIDIA GPU detected"**: Ensure the NVIDIA Container Toolkit is installed and `docker run --gpus all nvidia-smi` works.
+- **OOM errors**: Reduce batch size in the WebUI or use a GPU with more VRAM.
+- **WebUI not accessible**: Check that port 80 (nginx) or 7860 (direct) isn't blocked by a firewall.
+- **WebSocket errors behind proxy**: Ensure your proxy forwards `Upgrade` and `Connection` headers (the included nginx.conf handles this).
--- a/docker/docker-compose.yml
+++ b/docker/docker-compose.yml
@@ -0,0 +1,34 @@
+version: "3.8"
+
+services:
+  training-webui:
+    build:
+      context: ..
+      dockerfile: docker/Dockerfile
+    ports:
+      - "7860:7860"
+    volumes:
+      - ../models:/app/models        # Pre-downloaded model weights
+      - ../lora:/app/lora            # LoRA checkpoints (input/output)
+      - ../output:/app/output        # Training output artifacts
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: 1
+              capabilities: [gpu]
+    environment:
+      - GRADIO_SERVER_PORT=7860
+      - GRADIO_ROOT_PATH=/webui      # Matches nginx location block
+    restart: unless-stopped
+
+  nginx:
+    image: nginx:alpine
+    ports:
+      - "80:80"
+    volumes:
+      - ./nginx.conf:/etc/nginx/conf.d/default.conf:ro
+    depends_on:
+      - training-webui
+    restart: unless-stopped
--- a/docker/nginx.conf
+++ b/docker/nginx.conf
@@ -0,0 +1,21 @@
+server {
+    listen 80;
+    server_name _;
+
+    location /webui/ {
+        proxy_pass http://training-webui:7860/;
+        proxy_set_header Host $host;
+        proxy_set_header X-Real-IP $remote_addr;
+        proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
+        proxy_set_header X-Forwarded-Proto $scheme;
+
+        # WebSocket support (required for Gradio)
+        proxy_http_version 1.1;
+        proxy_set_header Upgrade $http_upgrade;
+        proxy_set_header Connection "upgrade";
+
+        # Increase timeouts for long-running training operations
+        proxy_read_timeout 300s;
+        proxy_send_timeout 300s;
+    }
+}
--- a/lora_ft_webui.py
+++ b/lora_ft_webui.py
@@ -500,6 +500,7 @@ def start_training(

        assert training_process.stdout is not None
        for line in training_process.stdout:
+            print(line, end="", flush=True)  # Stream to stdout (Docker logs)
            training_log += line
            # Keep log size manageable
            if len(training_log) > 100000:
@@ -1324,4 +1325,9 @@ with gr.Blocks(title="VoxCPM LoRA WebUI", theme=gr.themes.Soft(), css=custom_css
 if __name__ == "__main__":
    # Ensure lora directory exists
    os.makedirs("lora", exist_ok=True)
-    app.queue().launch(server_name="0.0.0.0", server_port=7860)
+
+    port = int(os.environ.get("GRADIO_SERVER_PORT", "7860"))
+    root_path = os.environ.get("GRADIO_ROOT_PATH", "")
+
+    print(f"\U0001f399\ufe0f  VoxCPM Training WebUI: http://0.0.0.0:{port}{root_path}", flush=True)
+    app.queue().launch(server_name="0.0.0.0", server_port=port, root_path=root_path)