feat: Update Docker configuration for GPU support and remove deprecated compose file

This commit is contained in:
JB
2025-10-05 16:05:16 -07:00
parent 338ff104e8
commit 66a0679e18
2 changed files with 20 additions and 19 deletions
+7 -11
View File
@@ -53,24 +53,20 @@ Browse to http://localhost:8000. Uploaded source files are stored in `/data/uplo
Set any of these with `-e VAR=value` when starting the container. Set any of these with `-e VAR=value` when starting the container.
### GPU-enabled build ### Docker Compose (GPU by default)
If you want CUDA acceleration inside the container, a GPU-aware Docker runtime (for example the NVIDIA Container Toolkit) is required. The repository ships an updated `abogen/Dockerfile` based on the CUDA runtime plus a helper Compose file. The repo includes `docker-compose.yaml`, which targets GPU hosts out of the box. Install the NVIDIA Container Toolkit and run:
```bash ```bash
# Build the GPU image (installs the matching CUDA PyTorch wheel) docker compose up -d --build
docker compose -f docker-compose.gpu.yml build
# Start the service with GPU access (--profile gpu in Compose v2 is optional)
docker compose -f docker-compose.gpu.yml up -d
``` ```
Useful overrides: Key build/runtime knobs:
- `TORCH_VERSION` pin a specific PyTorch release that matches your host driver (leave empty for latest). - `TORCH_VERSION` pin a specific PyTorch release that matches your driver (leave blank for the latest on the configured index).
- `TORCH_INDEX_URL` change the download index if you need a different CUDA build. - `TORCH_INDEX_URL` swap out the PyTorch download index when targeting a different CUDA build.
- `ABOGEN_DATA` host path that stores uploads/outputs (defaults to `./data`). - `ABOGEN_DATA` host path that stores uploads/outputs (defaults to `./data`).
The Compose file reserves a GPU via `device_requests`. Standard `docker run` works as well: CPU-only deployment: comment out the `deploy.resources.reservations.devices` block (and the optional `runtime: nvidia` line) inside the compose file. Compose will then run without requesting a GPU. If you prefer the classic CLI:
```bash ```bash
docker build -f abogen/Dockerfile -t abogen-gpu . docker build -f abogen/Dockerfile -t abogen-gpu .
+13 -8
View File
@@ -1,3 +1,5 @@
version: "3.9"
services: services:
abogen: abogen:
build: build:
@@ -6,7 +8,7 @@ services:
args: args:
TORCH_INDEX_URL: ${TORCH_INDEX_URL:-https://download.pytorch.org/whl/cu124} TORCH_INDEX_URL: ${TORCH_INDEX_URL:-https://download.pytorch.org/whl/cu124}
TORCH_VERSION: ${TORCH_VERSION:-} TORCH_VERSION: ${TORCH_VERSION:-}
image: abogen-gpu:latest image: abogen:latest
ports: ports:
- "${ABOGEN_PORT:-8000}:8000" - "${ABOGEN_PORT:-8000}:8000"
volumes: volumes:
@@ -16,15 +18,18 @@ services:
ABOGEN_PORT: 8000 ABOGEN_PORT: 8000
ABOGEN_UPLOAD_ROOT: /data/uploads ABOGEN_UPLOAD_ROOT: /data/uploads
ABOGEN_OUTPUT_ROOT: /data/outputs ABOGEN_OUTPUT_ROOT: /data/outputs
# --- GPU support -----------------------------------------------------
# These settings assume the NVIDIA Container Toolkit is installed.
# Leave them in place for GPU acceleration; comment out the entire block
# below if you are deploying to a CPU-only host.
deploy: deploy:
resources: resources:
reservations: reservations:
devices: devices:
- driver: nvidia - capabilities: [gpu]
count: all # driver: nvidia
capabilities: [gpu] # count: all
device_requests: # Runtime flag is only honored by legacy docker-compose (v1) CLI.
- driver: nvidia # Uncomment if you're still using it:
count: -1 # runtime: nvidia
capabilities: [gpu]
restart: unless-stopped restart: unless-stopped