FROM nvidia/cuda:12.1.0-cudnn8-runtime-ubuntu22.04

ENV DEBIAN_FRONTEND=noninteractive \
    PYTHONUNBUFFERED=1 \
    CUDA_HOME=/usr/local/cuda \
    PATH=/usr/local/cuda/bin:$PATH \
    LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH \
    TORCH_CUDA_ARCH_LIST="8.6" \
    NVIDIA_VISIBLE_DEVICES=all \
    NVIDIA_DRIVER_CAPABILITIES=all

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

# Python environment configuration
RUN ln -sf /usr/bin/python3 /usr/bin/python && \
    curl -sS https://bootstrap.pypa.io/get-pip.py | python3 && \
    python3 -m pip install --no-cache-dir pip==23.3.1 setuptools==69.0.3 wheel==0.42.0

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

WORKDIR /app
RUN mkdir -p /app/cache /app/src && \
    chmod 777 /app/cache

# Rest of your Dockerfile remains the same...
# # System-level configuration
# ENV DEBIAN_FRONTEND=noninteractive \
#     PYTHONUNBUFFERED=1 \
#     CUDA_HOME=/usr/local/cuda \
#     PATH=/usr/local/cuda/bin:$PATH \
#     LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH \
#     TORCH_CUDA_ARCH_LIST="8.6" \
#     NVIDIA_VISIBLE_DEVICES=all \
#     NVIDIA_DRIVER_CAPABILITIES=compute,utility

# # System dependencies - removed version pinning for base packages
# RUN apt-get update && apt-get install -y --no-install-recommends \
#     build-essential \
#     python3-dev \
#     python3-pip \
#     git \
#     ffmpeg \
#     libsndfile1 \
#     curl \
#     && apt-get clean \
#     && rm -rf /var/lib/apt/lists/*

# Python environment configuration
RUN ln -sf /usr/bin/python3 /usr/bin/python && \
    curl -sS https://bootstrap.pypa.io/get-pip.py | python3 && \
    python3 -m pip install --no-cache-dir pip==23.3.1 setuptools==69.0.3 wheel==0.42.0

# Application directory structure
WORKDIR /app
RUN mkdir -p /app/cache /app/src && \
    chmod 777 /app/cache

# Dependencies installation (two-stage for better caching)
COPY requirements.txt /app/src/
WORKDIR /app/src
RUN python3 -m pip install --no-cache-dir -r requirements.txt

# Application code
COPY . /app/src/

# Runtime environment
ENV MODEL_PATH=/app/cache \
    PYTHONPATH=/app/src:$PYTHONPATH \
    OMP_NUM_THREADS=1 \
    MKL_NUM_THREADS=1

# # GPU verification with comprehensive diagnostics
# RUN python3 -c 'import torch; \
#     assert torch.cuda.is_available(), "CUDA unavailable"; \
#     print(f"PyTorch: {torch.__version__}"); \
#     print(f"CUDA: {torch.version.cuda}"); \
#     print(f"GPU: {torch.cuda.get_device_name()}"); \
#     print(f"Arch: {torch.cuda.get_device_capability()}");'

EXPOSE 8000

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