FROM pytorch/pytorch:1.11.0-cuda11.3-cudnn8-runtime

RUN apt-get update && apt-get install -y --no-install-recommends \
    jq \
    build-essential \
    cmake \
    gcc
RUN rm -rf /var/lib/apt/lists/*

COPY lib/changehostname.c /
COPY lib/start_with_right_hostname.sh /usr/local/bin/start_with_right_hostname.sh
RUN chmod +x /usr/local/bin/start_with_right_hostname.sh

COPY dist/sagemaker_pytorch_training-*.tar.gz /sagemaker_pytorch_training.tar.gz
RUN pip install --no-cache-dir /sagemaker_pytorch_training.tar.gz && \
    rm /sagemaker_pytorch_training.tar.gz

ENV SAGEMAKER_TRAINING_MODULE=sagemaker_pytorch_container.training:main

WORKDIR /

# Starts framework
ENTRYPOINT ["bash", "-m", "start_with_right_hostname.sh"]