ARG BASE_IMAGE=""
FROM $BASE_IMAGE

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

ENV SAGEMAKER_TRAINING_MODULE sagemaker_tensorflow_container.training:main

ARG SMDEBUG_VERSION=1.0.12
ARG PYTHON=python3.8
# Add arguments to achieve the version, python and url
ARG PYTHON_VERSION=3.8.2

# dependencies for kebros 
RUN apt-get update && apt-get install -y --no-install-recommends \
    libgtk2.0-dev \
    && apt-get install -y -qq libkrb5-dev \
    && apt-get install -y -qq libsasl2-dev libsasl2-modules \
    && apt-get install -y -qq krb5-user \
&& rm -rf /var/lib/apt/lists/*

RUN $PYTHON -m 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"

# Starting sagemaker v2.92.0, the requirement of importlib-metadata was restricted from >=1.4.0 to >=1.4.0,<2.0
# However, the markdown package (dependency of tensorboard package) requires importlib-metadata to be >=4.4.
# Issue: https://github.com/aws/sagemaker-python-sdk/pull/3138
RUN $PYTHON -m pip install --no-cache-dir -U \
    "sagemaker>=2,<3" \
    sagemaker-experiments==0.* \
    "sagemaker-tensorflow>=2.7,<2.8" \
    "sagemaker-tensorflow-training>=20" \ 
    smclarify 

RUN $PYTHON -m pip install --no-cache-dir -U \
    "sagemaker-studio-analytics-extension==0.0.2" \ 
    "sparkmagic<1" \
    "sagemaker-studio-sparkmagic-lib<1" 

# install smdebug directly the specific branch
RUN git clone -b $SMDEBUG_VERSION https://github.com/awslabs/sagemaker-debugger.git \
    && cd sagemaker-debugger && $PYTHON setup.py install && cd .. && rm -rf sagemaker-debugger 

# remove tmp files
RUN rm -rf /tmp/git-secrets

ADD https://raw.githubusercontent.com/aws/deep-learning-containers/master/src/deep_learning_container.py /usr/local/bin/deep_learning_container.py

RUN chmod +x /usr/local/bin/deep_learning_container.py

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*

RUN curl https://aws-dlc-licenses.s3.amazonaws.com/tensorflow-2.7/license.txt -o /license.txt

CMD ["/bin/bash"]