File size: 2,233 Bytes
499bc4b | 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 | FROM nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04
# Build arg for GPU architectures
ARG CUDA_ARCHITECTURES="8.0;8.6"
ENV DEBIAN_FRONTEND=noninteractive
# Install system dependencies
RUN apt update && \
apt install -y \
python3 python3-pip git wget curl cmake ninja-build \
libgl1 libglib2.0-0 ffmpeg && \
apt clean
WORKDIR /workspace
COPY requirements.txt .
# Upgrade pip first
RUN pip install --upgrade pip setuptools wheel
# Install requirements if exists
RUN pip install -r requirements.txt
# Install PyTorch with CUDA support
RUN pip install --extra-index-url https://download.pytorch.org/whl/cu124 \
torch==2.6.0+cu124 torchvision==0.21.0+cu124
# Install SageAttention from git (patch GPU detection)
ENV TORCH_CUDA_ARCH_LIST="${CUDA_ARCHITECTURES}"
ENV FORCE_CUDA="1"
ENV MAX_JOBS="1"
COPY <<EOF /tmp/patch_setup.py
import os
with open('setup.py', 'r') as f:
content = f.read()
# Get architectures from environment variable
arch_list = os.environ.get('TORCH_CUDA_ARCH_LIST')
arch_set = '{' + ', '.join([f'"{arch}"' for arch in arch_list.split(';')]) + '}'
# Replace the GPU detection section
old_section = '''compute_capabilities = set()
device_count = torch.cuda.device_count()
for i in range(device_count):
major, minor = torch.cuda.get_device_capability(i)
if major < 8:
warnings.warn(f"skipping GPU {i} with compute capability {major}.{minor}")
continue
compute_capabilities.add(f"{major}.{minor}")'''
new_section = 'compute_capabilities = ' + arch_set + '''
print(f"Manually set compute capabilities: {compute_capabilities}")'''
content = content.replace(old_section, new_section)
with open('setup.py', 'w') as f:
f.write(content)
EOF
RUN git clone https://github.com/thu-ml/SageAttention.git /tmp/sageattention && \
cd /tmp/sageattention && \
python3 /tmp/patch_setup.py && \
pip install --no-build-isolation .
# --- FIX: Copy application code ---
COPY . /workspace
# ----------------------------------
RUN useradd -u 1000 -ms /bin/bash user
RUN chown -R user:user /workspace
RUN mkdir /home/user/.cache && \
chown -R user:user /home/user/.cache
COPY entrypoint.sh /workspace/entrypoint.sh
ENTRYPOINT ["/workspace/entrypoint.sh"]
|