# Start with AWS Deep Learning Container. See
# https://github.com/aws/deep-learning-containers/blob/master/available_images.md
FROM 763104351884.dkr.ecr.us-east-1.amazonaws.com/tensorflow-training:2.12.0-gpu-py310-cu118-ubuntu20.04-ec2

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

ARG PIP=pip3

RUN apt-get update \
 && apt-get install -y --no-install-recommends \
    vim\
 && rm -rf /var/lib/apt/lists/* \
 && apt-get clean

# 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 \        
        h5py==3.8.0 \
        ipykernel==5.3.4 \
        jinja2==3.1.2 \
        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==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}"

# Some packages can not be installed using pip (because they are not
# python managed packages) - install them using apt get

# Setup our entry point
COPY braket_container.py /opt/ml/code/braket_container.py

ENV SAGEMAKER_PROGRAM braket_container.py