# https://github.com/aws/deep-learning-containers/blob/master/available_images.md
# refer to the above page to pull latest Pytorch image

# docker image region us-west-2
FROM 763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.0.0-gpu-py310-cu118-ubuntu20.04-sagemaker

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

# version args
ARG TRANSFORMERS_VERSION
ARG DATASETS_VERSION
ARG DIFFUSERS_VERSION=0.16.1
ARG EVALUATE_VERSION=0.4.0
ARG ACCELERATE_VERSION=0.19.0
ARG PYTHON=python3

# install Hugging Face libraries and its dependencies
RUN pip install --no-cache-dir \
	transformers[sklearn,sentencepiece,audio,vision]==${TRANSFORMERS_VERSION} \
	datasets==${DATASETS_VERSION} \
	diffusers==${DIFFUSERS_VERSION} \
	Jinja2 \
	tensorboard \
	evaluate==${EVALUATE_VERSION} \
	accelerate==${ACCELERATE_VERSION}

RUN apt-get update \
 # TODO: Remove upgrade statements once packages are updated in base image
 && apt-get -y upgrade --only-upgrade systemd openssl cryptsetup \
 && apt install -y git-lfs \
 && apt-get clean \
 && rm -rf /var/lib/apt/lists/*

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*