Update for Dockerfile to change to cu128 to enable Blackwell GPU

This commit is contained in:
The Guy 2025-07-16 09:29:42 -04:00
parent 7bd4cfceef
commit 3204232cd4

View file

@ -29,7 +29,8 @@ RUN python3 -m venv /app/venv
ENV PATH="/app/venv/bin:$PATH"
# Install PyTorch with CUDA support and other dependencies
RUN pip3 install --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124 && \
# Using CUDA 12.8 for Blackwell sm_120 support
RUN pip3 install --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128 && \
pip3 install --no-cache-dir -r requirements.txt
# Copy project files
@ -38,10 +39,12 @@ COPY --chown=appuser:appuser . .
# Set environment variables
ENV PYTHONUNBUFFERED=1 \
PYTHONPATH=/app \
USE_GPU=true
USE_GPU=true \
CUDA_LAUNCH_BLOCKING=1 \
TORCH_USE_CUDA_DSA=1
# Expose the port
EXPOSE 5005
# Run FastAPI server with uvicorn
CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "5005", "--workers", "1"]
CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "5005", "--workers", "1"]