Dockerfile.megamolbart

# Copyright 2021 NVIDIA Corporation
FROM nvcr.io/nvidia/pytorch:20.11-py3

RUN apt-get update \
    && apt-get upgrade -y \
    && DEBIAN_FRONTEND=noninteractive apt-get install -y wget git unzip tmux libxrender1 \
    && rm -rf /var/lib/apt/lists/*

ENV UCX_LOG_LEVEL error
ENV PATH /opt/conda/bin:$PATH
ENV CUDA_HOME /usr/local/cuda
SHELL ["/bin/bash", "-c"]
RUN echo "source activate base" > /etc/bash.bashrc

## Environment setup
COPY megamolbart/conda/env.yml /tmp/.
RUN conda env update --name base -f /tmp/env.yml && conda clean -afy

## PySMILES -- requirements handled by conda environment
RUN git clone https://github.com/MolecularAI/pysmilesutils.git --branch master /opt/pysmilesutils \
    && cd /opt/pysmilesutils; pip install .

## MegaMolBART and Megatron code
RUN git clone https://github.com/MolecularAI/MolBART.git --branch megatron-molbart-with-zinc /opt/MolBART \
    && cd /opt/MolBART/megatron_molbart/Megatron-LM-v1.1.5-3D_parallelism && pip install . \
    && cd /opt/MolBART; pip install .
ENV PYTHONPATH /opt/MolBART/megatron_molbart:/opt/MolBART

# Copy and add generated code to PYTHONPATH
COPY common/ /tmp/common
RUN cd /tmp/common; pip install .

RUN mkdir -p /opt/nvidia/cuchem/grpc
COPY common/generated /opt/nvidia/cuchem/grpc
ENV PYTHONPATH /opt/nvidia/cuchem/grpc:$PYTHONPATH

COPY megamolbart/ /opt/nvidia/megamolbart/
CMD cd /opt/nvidia/megamolbart && python3 launch.py