VenusFactory2/docker-compose.yaml at main · ai4protein/VenusFactory2 · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
# Two service profiles share the same Dockerfile, only the base image and the
# torch/pyg wheel URLs differ.
#
# Usage:
#   cp .env.example .env                            # create local config (once)
#   docker compose --profile gpu up -d --build      # GPU (needs nvidia-container-toolkit)
#   docker compose --profile cpu up -d --build      # CPU-only
#
# Override port:
#   VENUS_PORT=8000 docker compose --profile gpu up -d

services:
  venus-gpu:
    profiles: ["gpu"]
    image: venusfactory/ce:gpu-${VENUS_FACTORY_VERSION:-0.1.0}
    build:
      context: .
      dockerfile: docker/Dockerfile
      args:
        BASE_IMAGE: nvidia/cuda:12.8.0-runtime-ubuntu24.04
        TORCH_INDEX_URL: https://download.pytorch.org/whl/cu128
        PYG_INDEX_URL: https://data.pyg.org/whl/torch-2.8.0+cu128.html
    ports:
      - "${VENUS_PORT:-7861}:7861"
    env_file:
      - path: .env
        required: false
    volumes:
      - ./temp_outputs:/app/temp_outputs
      - ./ckpt:/app/ckpt
    restart: unless-stopped
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]

  venus-cpu:
    profiles: ["cpu"]
    image: venusfactory/ce:cpu-${VENUS_FACTORY_VERSION:-0.1.0}
    build:
      context: .
      dockerfile: docker/Dockerfile
      args:
        BASE_IMAGE: ubuntu:24.04
        TORCH_INDEX_URL: https://download.pytorch.org/whl/cpu
        PYG_INDEX_URL: https://data.pyg.org/whl/torch-2.8.0+cpu.html
    ports:
      - "${VENUS_PORT:-7861}:7861"
    env_file:
      - path: .env
        required: false
    volumes:
      - ./temp_outputs:/app/temp_outputs
      - ./ckpt:/app/ckpt
    restart: unless-stopped