# This is based loosely on :
# https://github.com/aws/amazon-sagemaker-examples/blob/master/advanced_functionality/custom-training-containers/script-mode-container/docker/Dockerfile

FROM 763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:2.0.0-gpu-py310-cu118-ubuntu20.04-ec2

LABEL maintainer="Amazon Braket"
LABEL major_version="1"

# prevent stopping by user interaction
ENV DEBIAN_FRONTEND noninteractive
ENV DEBCONF_NONINTERACTIVE_SEEN true

ENV PYTHONDONTWRITEBYTECODE=1
ENV PYTHONUNBUFFERED=1
ENV PYTHONIOENCODING=UTF-8
ENV LANG=C.UTF-8
ENV LC_ALL=C.UTF-8


ARG PYTHON=python3.10
ARG PYTHON_PIP=python3-pip
ARG PIP=pip3
ARG PYTHON_VERSION=3.10.11

RUN apt-get update \
 && apt-get install -y --no-install-recommends \
    build-essential \
    openssh-client \
    openssh-server \
    ca-certificates \
    curl \
    git \
    libtemplate-perl \
    libssl1.1 \
    openssl \
    unzip \
    vim \
    wget \
    zlib1g-dev \
    # Install dependent library for OpenCV
    libgtk2.0-dev \
 && rm -rf /var/lib/apt/lists/*

RUN apt-get update \
 && apt-get install -y --no-install-recommends \
    libbz2-dev \
    libc6-dev \
    libffi-dev \
    libgdbm-dev \
    liblzma-dev \
    libncursesw5-dev \
    libreadline-gplv2-dev \
    libsqlite3-dev \
    libssl-dev \
    tk-dev \
 && rm -rf /var/lib/apt/lists/* \
 && apt-get clean

# Installing Python3
RUN wget https://www.python.org/ftp/python/$PYTHON_VERSION/Python-$PYTHON_VERSION.tgz \
 && tar -xvf Python-$PYTHON_VERSION.tgz \
 && cd Python-$PYTHON_VERSION \
 && ./configure && make && make install \
 && rm -rf ../Python-$PYTHON_VERSION*

# Upgrading pip and creating symbolic link for python3
RUN ${PIP} --no-cache-dir install --upgrade pip setuptools
RUN ln -s $(which ${PYTHON}) /usr/local/bin/python \
 && ln -s $(which ${PIP}) /usr/bin/pip

RUN apt-get update && apt-get -y install cmake protobuf-compiler


# Installing our custom python libraries
RUN ${PIP} install --no-cache --upgrade \
        amazon-braket-default-simulator==1.18.1 \
        amazon-braket-schemas==1.18.0 \
        amazon-braket-pennylane-plugin==1.18.0 \
        amazon-braket-sdk==1.50.0 \
        awscli==1.29.4 \
        botocore==1.31.4 \
        boto3==1.28.3 \
        dask==2023.1.1 \
        decorator==4.4.0 \
        h5py==3.8.0 \
        ipykernel==5.3.4 \
        jinja2==2.11.3 \
        keras==2.12.0 \
        markupsafe==2.1.2 \
        matplotlib==3.6.3 \
        numpy==1.23.5 \
        openfermion==1.5.1 \
        pandas==2.0.0 \
        PennyLane==0.30.0 \
        PennyLane-Lightning[gpu]==0.30.0 \
        pydantic==1.10.6 \
        requests==2.26.0 \
        scikit-learn==1.2.2 \
        six==1.16.0 \
        scipy==1.9.3 \
        typing_extensions==4.3.0

RUN ${PIP} install --no-cache --upgrade sagemaker-training==4.4.10

# install cuQuantum
RUN wget https://developer.download.nvidia.com/compute/cuquantum/redist/cuquantum/linux-x86_64/cuquantum-linux-x86_64-23.03.0.20-archive.tar.xz && tar xvf ./cuquantum-linux-x86_64-23.03.0.20-archive.tar.xz
RUN mv /cuquantum-linux-x86_64-23.03.0.20-archive /opt/cuquantum
 
ENV LD_LIBRARY_PATH="/opt/cuquantum/lib:${LD_LIBRARY_PATH}"


RUN apt-get clean


# Add EFA and SMDDP to LD library path
ENV LD_LIBRARY_PATH="/opt/conda/lib/python${PYTHON_SHORT_VERSION}/site-packages/smdistributed/dataparallel/lib:$LD_LIBRARY_PATH"
ENV LD_LIBRARY_PATH=/opt/amazon/efa/lib/:$LD_LIBRARY_PATH


# Setup our entry point

COPY braket_container.py /opt/ml/code/braket_container.py

ENV SAGEMAKER_PROGRAM braket_container.py