# Use PyTorch base image with CUDA support (much smaller than building from scratch)
FROM pytorch/pytorch:2.6.0-cuda12.4-cudnn9-runtime

# Set environment variables
ENV DEBIAN_FRONTEND=noninteractive \
    PYTHONUNBUFFERED=TRUE \
    PYTHONFAULTHANDLER=1 \
    MPLCONFIGDIR=/tmp/matplotlib \
    CHAI_DOWNLOADS_DIR=/tmp/chai_downloads \
    HF_HOME=/tmp/hf_home

# Set working directory
WORKDIR /workspace

# Install chai_lab and compatible transformers in a single layer
RUN pip install --no-cache-dir \
    chai_lab==0.5.2 \
    "transformers>=4.30.0,<4.45.0"

# Make package dirs writable for any user (K8s may run as non-root)
RUN chmod -R a+rw /opt/conda/lib/python3.11/site-packages/chai_lab/ && \
    mkdir -p /opt/conda/lib/python3.11/site-packages/downloads && \
    chmod -R a+rw /opt/conda/lib/python3.11/site-packages/downloads && \
    mkdir -p /tmp/chai_downloads /tmp/hf_home /tmp/matplotlib && \
    chmod -R 777 /tmp/chai_downloads /tmp/hf_home /tmp/matplotlib

# Verify installations
RUN python -c "import torch; print(f'PyTorch: {torch.__version__}')" && \
    python -c "from transformers import EsmModel; print('EsmModel: OK')" && \
    python -c "import chai_lab; print('chai_lab: OK')" && \
    chai --help

# Add entry point script
COPY entrypoint.sh /workspace/
RUN chmod +x /workspace/entrypoint.sh

# Set entry point
ENTRYPOINT ["/workspace/entrypoint.sh"]