FROM quay.io/pawsey/rocm-mpich-base:rocm6.1-mpich3.4.3-ubuntu22

## Add System Dependencies
# needs libxrender1 ImportError: libXrender.so.1: cannot open shared object file: No such file or directory (from pdbeccdutils.core import ccd_reader) throws the error
RUN apt-get update \
    && DEBIAN_FRONTEND=noninteractive \
    && apt-get install --no-install-recommends -y \
        build-essential \
        git \
        wget \
        libxrender1 \
    && rm -rf /var/lib/apt/lists/* \
    && apt-get autoremove -y \
    && apt-get clean

## ROCm environment variables

ENV ROCM_RELEASE=6.1.0
ENV ROCM_PATH=opt/rocm-$ROCM_RELEASE
ENV PATH=$ROCM_PATH/bin:$ROCM_PATH/llvm/bin:$PATH
ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$ROCM_PATH/lib

## Install miniforge
# install into
RUN set -eux ; \
  curl -LO https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-x86_64.sh ; \
  bash ./Miniforge3-* -b -p /opt/miniforge3 -s ; \
  rm -rf ./Miniforge3-*
ENV PATH=/opt/miniforge3/bin:$PATH

## install pip
RUN set -eux ; \
  mamba install -y -c conda-forge -c bioconda -c defaults pip

## Install some Python dependencies
RUN pip install pytest \
    requests \
    biopandas --no-cache-dir


WORKDIR /app/alphafold

## Clone and install the package + requirements
ARG GIT_TAG=main
RUN git clone https://github.com/amorehead/alphafold3-pytorch-lightning-hydra . --branch ${GIT_TAG} \
    # && git checkout main \
    && python -m pip install .

## install PyTorch compatible with ROCM6.1
RUN pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.1 --no-cache-dir
RUN pip install pytorch-triton-rocm --no-cache-dir
