welsh-whisperx/docker-compose.gpu.yml at main · techiaith/welsh-whisperx · 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
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
services:

  application:
     image: techiaith/whisperx-api
     # Removed container_name to allow potential scaling
     restart: always
     environment:
       - REDIS_HOST=${REDIS_HOST}
       - REDIS_PORT=${REDIS_PORT}
       - CLEANUP_ENABLED=${CLEANUP_ENABLED}
       - FILE_RETENTION_DAYS=${FILE_RETENTION_DAYS}
       - CLEANUP_SCHEDULE=${CLEANUP_SCHEDULE}
       - RECORDINGS_DIR=${RECORDINGS_DIR}
     build:
       context: ${PWD}/api
       dockerfile: Dockerfile
     ports:
       - "5511:8000"
     volumes:
       - ${PWD}/recordings:/recordings
       - short_form_audio:/recordings_short
     depends_on:
       - redis


  # High-priority worker: /transcribe/ and /keyboard/ requests
  # Runs on GPU 1 — shares GPU with worker-default
  worker-high:
    image: techiaith/whisperx-worker
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids: ['${CUDA_DEVICE_INDEX_HIGH:-1}']
              capabilities: [gpu]
    restart: always
    environment:
      - WORKER_QUEUES=high_priority
      - WHISPER_MODEL_NAME=${WHISPER_MODEL_NAME}
      - WHISPER_MODEL_DEVICE=gpu
      - WHISPER_MODEL_LANGUAGE=${WHISPER_MODEL_LANGUAGE}
      - WAV2VEC2_MODEL=${WAV2VEC2_MODEL}
      - HF_AUTH_TOKEN=${HF_AUTH_TOKEN}
      - REDIS_HOST=${REDIS_HOST}
      - REDIS_PORT=${REDIS_PORT}
      - TORCH_FORCE_WEIGHTS_ONLY_LOAD=0
      - DIARIZE_MIN_DURATION_SECONDS=${DIARIZE_MIN_DURATION_SECONDS:-30}
      - SAVE_FILES_MIN_DURATION_SECONDS=${SAVE_FILES_MIN_DURATION_SECONDS:-30}
      - ALIGN_MIN_DURATION_SECONDS=${ALIGN_MIN_DURATION_SECONDS:-30}
    build:
      context: ${PWD}/worker
      dockerfile: Dockerfile
      args:
        DEVICE: gpu
    volumes:
       - ${PWD}/recordings:/recordings
       - short_form_audio:/recordings_short
    depends_on:
       - redis

  # Second high-priority worker on GPU 0 — dedicated GPU.
  # Handles overflow high-priority requests while GPU 1 is busy with long-form.
  worker-high-2:
    image: techiaith/whisperx-worker
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids: ['${CUDA_DEVICE_INDEX_HIGH_2:-0}']
              capabilities: [gpu]
    restart: always
    environment:
      - WORKER_QUEUES=high_priority
      - WHISPER_MODEL_NAME=${WHISPER_MODEL_NAME}
      - WHISPER_MODEL_DEVICE=gpu
      - WHISPER_MODEL_LANGUAGE=${WHISPER_MODEL_LANGUAGE}
      - WAV2VEC2_MODEL=${WAV2VEC2_MODEL}
      - HF_AUTH_TOKEN=${HF_AUTH_TOKEN}
      - REDIS_HOST=${REDIS_HOST}
      - REDIS_PORT=${REDIS_PORT}
      - TORCH_FORCE_WEIGHTS_ONLY_LOAD=0
      - DIARIZE_MIN_DURATION_SECONDS=${DIARIZE_MIN_DURATION_SECONDS:-30}
      - SAVE_FILES_MIN_DURATION_SECONDS=${SAVE_FILES_MIN_DURATION_SECONDS:-30}
      - ALIGN_MIN_DURATION_SECONDS=${ALIGN_MIN_DURATION_SECONDS:-30}
    volumes:
       - ${PWD}/recordings:/recordings
       - short_form_audio:/recordings_short
    depends_on:
       - redis

  # Default worker: /transcribe_long_form/ requests (normal and low priority)
  # Runs on GPU 1 — shares GPU with worker-high
  # No build: section — uses the image built by worker-high
  worker-default:
    image: techiaith/whisperx-worker
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids: ['${CUDA_DEVICE_INDEX_DEFAULT:-1}']
              capabilities: [gpu]
    restart: always
    environment:
      - WORKER_QUEUES=default
      - WHISPER_MODEL_NAME=${WHISPER_MODEL_NAME}
      - WHISPER_MODEL_DEVICE=gpu
      - WHISPER_MODEL_LANGUAGE=${WHISPER_MODEL_LANGUAGE}
      - WAV2VEC2_MODEL=${WAV2VEC2_MODEL}
      - HF_AUTH_TOKEN=${HF_AUTH_TOKEN}
      - REDIS_HOST=${REDIS_HOST}
      - REDIS_PORT=${REDIS_PORT}
      - TORCH_FORCE_WEIGHTS_ONLY_LOAD=0
      - DIARIZE_MIN_DURATION_SECONDS=${DIARIZE_MIN_DURATION_SECONDS:-30}
      - SAVE_FILES_MIN_DURATION_SECONDS=${SAVE_FILES_MIN_DURATION_SECONDS:-30}
      - ALIGN_MIN_DURATION_SECONDS=${ALIGN_MIN_DURATION_SECONDS:-30}
    volumes:
       - ${PWD}/recordings:/recordings
    depends_on:
       - redis


  # Alignment workers: CPU-only, handle /align/ and /align_long_form/ requests.
  # Runs wav2vec2 on CPU to avoid CUDA OOM on long audio files.
  # No GPU reservation needed — wav2vec2 is small and CPU is fast enough.
  # Two replicas so a long alignment doesn't block short /align/ requests.
  # Scale further with: docker compose -f docker-compose.gpu.yml up -d --scale worker-alignment=N
  worker-alignment:
    image: techiaith/whisperx-worker
    deploy:
      replicas: 2
    restart: always
    environment:
      - WORKER_QUEUES=alignment
      - WHISPER_MODEL_NAME=${WHISPER_MODEL_NAME}
      - WHISPER_MODEL_DEVICE=cpu
      - WHISPER_MODEL_LANGUAGE=${WHISPER_MODEL_LANGUAGE}
      - WAV2VEC2_MODEL=${WAV2VEC2_MODEL}
      - HF_AUTH_TOKEN=${HF_AUTH_TOKEN}
      - REDIS_HOST=${REDIS_HOST}
      - REDIS_PORT=${REDIS_PORT}
      - TORCH_FORCE_WEIGHTS_ONLY_LOAD=0
    volumes:
       - ${PWD}/recordings:/recordings
       - short_form_audio:/recordings_short
    depends_on:
       - redis

  redis:
    image: redis:7.2.1
    # Keep container_name for Redis as it's a singleton service
    container_name: whisperx-redis
    hostname: redis
    restart: always
    volumes:
       - redis_data:/data

volumes:
    redis_data: {}
    short_form_audio:
      driver: local
      driver_opts:
        type: tmpfs
        device: tmpfs
        o: size=256m