mirror of
https://github.com/jamiepine/voicebox.git
synced 2026-09-15 04:40:40 -07:00
fix(docker): add ROCm GPU support via compose overlay (#630)
* fix(docker): add ROCm GPU support via compose overlay Fixes #618. The Docker image installs CPU-only PyTorch from PyPI by default, so even when users correctly pass /dev/kfd and /dev/dri device nodes into the container, torch.cuda.is_available() returns False and the GPU is reported as "None (CPU only)". Changes: - Dockerfile: add PYTORCH_VARIANT build arg (default: cpu). When set to "rocm", the ROCm-enabled PyTorch wheels are installed from the pytorch.org/whl/rocm6.3 index before requirements.txt runs, so pip sees the ROCm build as already satisfying the torch>=2.2.0 constraint and does not overwrite it with the CPU wheel. The render and video groups are created with parameterised GIDs (RENDER_GID / VIDEO_GID, defaulting to Ubuntu 22.04 values) and the voicebox user is added to both groups so it can open /dev/kfd and /dev/dri. - docker-compose.rocm.yml: new compose overlay that wires everything together — PYTORCH_VARIANT=rocm build arg, /dev/kfd + /dev/dri device passthrough, group_add for render/video, HSA_OVERRIDE_GFX_VERSION (defaults to 11.0.0 for RDNA3/Strix Halo with a comment listing values for RDNA2/RDNA1/Vega), and PYTORCH_HIP_ALLOC_CONF for the memory allocator. Usage: docker compose -f docker-compose.yml -f docker-compose.rocm.yml up --build - docker-compose.yml: add a comment pointing to the ROCm overlay. The CPU default path is unchanged — no extra build time, no size increase. Co-authored-by: Cursor <[email protected]> * fix(docker): address review comments on ROCm overlay Two issues raised in PR review: 1. CodeRabbit: `docker compose up --build-arg` is not supported by the `up` subcommand. Replaced the GID override instructions with the correct env-var export pattern. Added RENDER_GID and VIDEO_GID to `build.args` using ${VAR:-default} interpolation so a single export covers both the Dockerfile group creation and the runtime group_add. Changed group_add entries from hardcoded strings to the same interpolated vars so host GIDs stay in sync end-to-end. 2. @Xarianne: ROCm 6.3 does not support RDNA 4 (RX 9000 series) cards. Added a ROCM_VERSION build arg (default 6.3) to both the Dockerfile and docker-compose.rocm.yml so users can set ROCM_VERSION=7.2 for RDNA 4 support without editing any files. Added RDNA 4 / 12.0.0 to the HSA_OVERRIDE_GFX_VERSION comment table. Co-authored-by: Cursor <[email protected]> --------- Co-authored-by: Cursor <[email protected]>
This commit is contained in:
Ankush
+47
-1
@@ -1,8 +1,15 @@
|
||||
# ============================================================
|
||||
# Voicebox — Local TTS Server with Web UI (CPU)
|
||||
# Voicebox — Local TTS Server with Web UI
|
||||
# 3-stage build: Frontend → Python deps → Runtime
|
||||
#
|
||||
# Build variants:
|
||||
# CPU (default): docker compose up --build
|
||||
# ROCm (AMD GPU): docker compose -f docker-compose.yml -f docker-compose.rocm.yml up --build
|
||||
# ============================================================
|
||||
|
||||
# Top-level ARG so it is visible to all stages.
|
||||
ARG PYTORCH_VARIANT=cpu
|
||||
|
||||
# === Stage 1: Build frontend ===
|
||||
FROM oven/bun:1 AS frontend
|
||||
|
||||
@@ -24,6 +31,9 @@ RUN cd web && bunx --bun vite build
|
||||
# === Stage 2: Build Python dependencies ===
|
||||
FROM python:3.11-slim AS backend-builder
|
||||
|
||||
# Re-declare ARG inside the stage (Docker scoping requirement).
|
||||
ARG PYTORCH_VARIANT=cpu
|
||||
|
||||
WORKDIR /build
|
||||
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
@@ -34,6 +44,21 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
RUN pip install --no-cache-dir --upgrade pip
|
||||
|
||||
COPY backend/requirements.txt .
|
||||
|
||||
# ROCm version to pull PyTorch wheels for. Default is 6.3 (supports RDNA1/2/3).
|
||||
# Set ROCM_VERSION=7.2 for RDNA 4 (RX 9000 series) support.
|
||||
ARG ROCM_VERSION=6.3
|
||||
|
||||
# When building the ROCm variant, install the ROCm-enabled PyTorch wheels
|
||||
# first so that the subsequent requirements.txt install sees them as already
|
||||
# satisfying the torch/torchaudio constraints and leaves them in place.
|
||||
# The CPU path skips this step and installs torch from PyPI as before.
|
||||
RUN if [ "$PYTORCH_VARIANT" = "rocm" ]; then \
|
||||
pip install --no-cache-dir --prefix=/install \
|
||||
torch torchaudio \
|
||||
--index-url "https://download.pytorch.org/whl/rocm${ROCM_VERSION}"; \
|
||||
fi
|
||||
|
||||
RUN pip install --no-cache-dir --prefix=/install -r requirements.txt
|
||||
RUN pip install --no-cache-dir --prefix=/install --no-deps chatterbox-tts
|
||||
RUN pip install --no-cache-dir --prefix=/install --no-deps hume-tada
|
||||
@@ -44,10 +69,31 @@ RUN pip install --no-cache-dir --prefix=/install \
|
||||
# === Stage 3: Runtime ===
|
||||
FROM python:3.11-slim
|
||||
|
||||
# Re-declare ARG inside the stage (Docker scoping requirement).
|
||||
ARG PYTORCH_VARIANT=cpu
|
||||
|
||||
# ROCm device access requires the container user to belong to the render
|
||||
# and video groups. GIDs are parameterised to match the host; Ubuntu 22.04+
|
||||
# defaults are used here. Override via env vars (docker-compose.rocm.yml
|
||||
# passes them through automatically):
|
||||
# export RENDER_GID=$(getent group render | cut -d: -f3)
|
||||
# export VIDEO_GID=$(getent group video | cut -d: -f3)
|
||||
ARG RENDER_GID=992
|
||||
ARG VIDEO_GID=44
|
||||
RUN if [ "$PYTORCH_VARIANT" = "rocm" ]; then \
|
||||
groupadd -f -g ${RENDER_GID} render && \
|
||||
groupadd -f -g ${VIDEO_GID} video; \
|
||||
fi
|
||||
|
||||
# Create non-root user for security
|
||||
RUN groupadd -r voicebox && \
|
||||
useradd -r -g voicebox -m -s /bin/bash voicebox
|
||||
|
||||
# ROCm: add voicebox user to render+video so it can open /dev/kfd and /dev/dri.
|
||||
RUN if [ "$PYTORCH_VARIANT" = "rocm" ]; then \
|
||||
usermod -aG render,video voicebox; \
|
||||
fi
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
# Install only runtime system dependencies
|
||||
|
||||
@@ -0,0 +1,64 @@
|
||||
# ROCm (AMD GPU) overlay for Voicebox
|
||||
#
|
||||
# Usage:
|
||||
# docker compose -f docker-compose.yml -f docker-compose.rocm.yml up --build
|
||||
#
|
||||
# Prerequisites on the host:
|
||||
# 1. ROCm drivers installed — https://rocm.docs.amd.com/projects/install-on-linux
|
||||
# 2. Current user in the 'render' and 'video' groups:
|
||||
# sudo usermod -aG render,video $USER (then log out/in)
|
||||
#
|
||||
# If the render/video GIDs on your host differ from the Ubuntu 22.04 defaults
|
||||
# (render=992, video=44), export them before running compose so the build arg
|
||||
# and group_add values both stay in sync automatically:
|
||||
# export RENDER_GID=$(getent group render | cut -d: -f3)
|
||||
# export VIDEO_GID=$(getent group video | cut -d: -f3)
|
||||
# docker compose -f docker-compose.yml -f docker-compose.rocm.yml up --build
|
||||
#
|
||||
# ROCm version (ROCM_VERSION):
|
||||
# Default is 6.3 (supports RDNA1/2/3). For RDNA 4 (RX 9000 series) use 7.2:
|
||||
# export ROCM_VERSION=7.2
|
||||
# docker compose -f docker-compose.yml -f docker-compose.rocm.yml up --build
|
||||
|
||||
services:
|
||||
voicebox:
|
||||
build:
|
||||
context: .
|
||||
args:
|
||||
PYTORCH_VARIANT: rocm
|
||||
# These build args read from env vars so a single export covers both the
|
||||
# Dockerfile group creation and the runtime group_add below.
|
||||
RENDER_GID: ${RENDER_GID:-992}
|
||||
VIDEO_GID: ${VIDEO_GID:-44}
|
||||
ROCM_VERSION: ${ROCM_VERSION:-6.3}
|
||||
|
||||
# Pass the AMD GPU device nodes into the container.
|
||||
# /dev/kfd — ROCm compute interface (required for GPU inference)
|
||||
# /dev/dri — DRM render nodes (required for display/memory access)
|
||||
devices:
|
||||
- /dev/kfd
|
||||
- /dev/dri
|
||||
|
||||
# Grant access to the render and video groups so the non-root user
|
||||
# inside the container can open the GPU device nodes.
|
||||
# These reference the same env vars used in build.args above so a
|
||||
# single export keeps Dockerfile groups and runtime group_add in sync.
|
||||
group_add:
|
||||
- "${RENDER_GID:-992}" # render
|
||||
- "${VIDEO_GID:-44}" # video
|
||||
|
||||
environment:
|
||||
# HSA_OVERRIDE_GFX_VERSION forces the ROCm runtime to treat the GPU as a
|
||||
# specific GFX version when auto-detection fails or the GPU is newer than
|
||||
# the ROCm release. app.py sets 10.3.0 (RDNA2) by default; override here
|
||||
# for your GPU family:
|
||||
# RDNA4 / RX 9000 series: 12.0.0 (requires ROCM_VERSION=7.2)
|
||||
# RDNA3 / RX 7000 series / Strix Halo: 11.0.0
|
||||
# RDNA2 / RX 6000 series: 10.3.0
|
||||
# RDNA1 / RX 5000 series: 10.1.0
|
||||
# Vega / GCN5: 9.0.0
|
||||
- HSA_OVERRIDE_GFX_VERSION=11.0.0
|
||||
|
||||
# Tune the ROCm memory allocator to reduce fragmentation during
|
||||
# multi-engine inference (TTS + STT + LLM running concurrently).
|
||||
- PYTORCH_HIP_ALLOC_CONF=garbage_collection_threshold:0.8,max_split_size_mb:512
|
||||
@@ -1,3 +1,7 @@
|
||||
# Voicebox — CPU build (default)
|
||||
# For AMD ROCm GPU acceleration use the overlay:
|
||||
# docker compose -f docker-compose.yml -f docker-compose.rocm.yml up --build
|
||||
|
||||
services:
|
||||
voicebox:
|
||||
build: .
|
||||
|
||||
Reference in New Issue
Block a user