# Expecting base image to be the image built by ./Dockerfile.ec2.cpu
ARG BASE_IMAGE=""

FROM $BASE_IMAGE

LABEL maintainer="Amazon AI"
LABEL dlc_major_version="1"

ARG PYTHON=python3
# Add arguments to achieve the version, python and url
ARG PYTHON_VERSION=3.8.13

# The smdebug pipeline relies for following format to perform string replace and trigger DLC pipeline for validating
# the nightly builds. Therefore, while updating the smdebug version, please ensure that the format is not disturbed.
ARG SMDEBUG_VERSION=1.0.22

ENV SAGEMAKER_TRAINING_MODULE=sagemaker_pytorch_container.training:main

# swap the pytorch training wheel with additional smdebug and smmdp features
ARG PT_TRAINING_URL=https://aws-pytorch-unified-cicd-binaries.s3.us-west-2.amazonaws.com/r1.10.2_sm/20220911-190136/1be40bd9b7251ad77e72b81863bcf59bc933e3ca/torch-1.10.2%2Bcpu-cp38-cp38-linux_x86_64.whl

WORKDIR /

# Copy workaround script for incorrect hostname
COPY changehostname.c /
COPY start_with_right_hostname.sh /usr/local/bin/start_with_right_hostname.sh

# Install scikit-learn and pandas
RUN conda install -y -c conda-forge \
    scikit-learn \
    pandas

RUN pip install --upgrade pip --trusted-host pypi.org --trusted-host files.pythonhosted.org \
 && pip install --no-cache-dir -U \
    # smdebug==${SMDEBUG_VERSION} \
    smclarify \
    "sagemaker>=2,<3" \
    sagemaker-experiments==0.* \
    "sagemaker-pytorch-training>=2.5.1,<3" \
 && pip uninstall -y torch \
 && pip install --no-cache-dir -U ${PT_TRAINING_URL}

# Install smdebug from souce
RUN cd /tmp \
  && git clone -b ${SMDEBUG_VERSION} https://github.com/awslabs/sagemaker-debugger \
  && cd sagemaker-debugger \
  && python setup.py install \
  && rm -rf /tmp/*

# Install extra packages
# numba 0.54 only works with numpy>=1.20. See https://github.com/numba/numba/issues/7339
RUN pip install --no-cache-dir -U \
    "bokeh>=2.3,<3" \
    "imageio>=2.9,<3" \
    "opencv-python>=4.6,<5" \
    "plotly>=5.1,<6" \
    "seaborn>=0.11,<1" \
    "numba<0.54" \
    "shap>=0.39,<1"

RUN chmod +x /usr/local/bin/start_with_right_hostname.sh

RUN HOME_DIR=/root \
 && curl -o ${HOME_DIR}/oss_compliance.zip https://aws-dlinfra-utilities.s3.amazonaws.com/oss_compliance.zip \
 && unzip ${HOME_DIR}/oss_compliance.zip -d ${HOME_DIR}/ \
 && cp ${HOME_DIR}/oss_compliance/test/testOSSCompliance /usr/local/bin/testOSSCompliance \
 && chmod +x /usr/local/bin/testOSSCompliance \
 && chmod +x ${HOME_DIR}/oss_compliance/generate_oss_compliance.sh \
 && ${HOME_DIR}/oss_compliance/generate_oss_compliance.sh ${HOME_DIR} ${PYTHON} \
 && rm -rf ${HOME_DIR}/oss_compliance* \
 && rm -rf /tmp/tmp*

ENTRYPOINT ["bash", "-m", "start_with_right_hostname.sh"]
CMD ["/bin/bash"]