FROM nvidia/cuda:12.3.2-cudnn9-devel-ubuntu22.04


# Set environment variables
ENV PYTHONUNBUFFERED=1 \
    DEBIAN_FRONTEND=noninteractive \
    CUDA_HOME=/usr/local/cuda \
    PATH=/usr/local/cuda/bin:$PATH \
    LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH \
    NVIDIA_VISIBLE_DEVICES=all \
    NVIDIA_DRIVER_CAPABILITIES=compute,utility

# Install system dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
    python3 \
    python3-pip \
    python3-dev \
    build-essential \
    ffmpeg \
    libsndfile1 \
    curl \
    && rm -rf /var/lib/apt/lists/*

# Create a non-root user
# In your Dockerfile:
RUN useradd -m -s /bin/bash appuser && \
    mkdir -p /home/appuser/.cache/torch && \
    mkdir -p /home/appuser/.triton && \
    chown -R appuser:appuser /home/appuser && \
    chmod 755 /home/appuser/.triton

# Add these ENV variables 
ENV TRITON_CACHE_DIR=/home/appuser/.triton \
    TORCH_INDUCTOR_CACHE_DIR=/home/appuser/.triton

# Set cache directories environment variables
ENV TORCH_HOME=/home/appuser/.cache/torch \
    TRITON_CACHE_DIR=/home/appuser/.triton

# Upgrade pip and install build tools
RUN python3 -m pip install --upgrade pip setuptools wheel

WORKDIR /app
RUN chown appuser:appuser /app

COPY --chown=appuser:appuser requirements.txt .

# Install PyTorch with CUDA support
RUN pip3 install --no-cache-dir torch==2.1.2 torchvision==0.16.2 torchaudio==2.1.2 --index-url https://download.pytorch.org/whl/cu121

# Install other requirements
RUN pip3 install --no-cache-dir -r requirements.txt

COPY --chown=appuser:appuser . .

# Switch to non-root user
USER appuser

# Simple CUDA test
RUN python3 -c "import torch; print('CUDA available:', torch.cuda.is_available()); print('CUDA version:', torch.version.cuda); print('Device count:', torch.cuda.device_count())"

EXPOSE 8000

CMD ["python3", "server.py"]