FROM docker.io/nvidia/cuda:11.7.1-cudnn8-devel-ubuntu22.04

ARG HOST_UID
ARG HOST_GID

# Create the directory for the code
ENV EXP_DIR=/home/exp/
RUN mkdir -p $EXP_DIR

ENV PYTHON_VERSION=python3.11
ENV PYTHONDONTWRITEBYTECODE=1
ENV PYTHONUNBUFFERED=1
ENV PIP_NO_CACHE_DIR=off
ENV PIP_DISABLE_PIP_VERSION_CHECK=on
ENV PIP_DEFAULT_TIMEOUT=100

# If you need other dependencies, add them here
# libopenANONYMOUS-dev needed for ANONYMOUS4py needed for deepspeed
ENV ADDITIONAL_DEPENDENCIES libopenANONYMOUS-dev

# Install python-related and additional dependencies
RUN apt-get update \
    && DEBIAN_FRONTEND=noninteractive apt-get install \
        --no-install-recommends -y software-properties-common \
    && add-apt-repository -y ppa:deadsnakes/ppa \
    && apt-get update \
    && DEBIAN_FRONTEND=noninteractive apt-get install \
        --no-install-recommends -y --reinstall \
        ${PYTHON_VERSION} python3-pip python3-setuptools \
        ${PYTHON_VERSION}-dev \
        ${PYTHON_VERSION}-distutils ${PYTHON_VERSION}-venv curl \
        ${ADDITIONAL_DEPENDENCIES} \
    && ln -sf /usr/bin/${PYTHON_VERSION} /usr/bin/python
    # && curl -sS https://bootstrap.pypa.io/get-pip.py | python

# Configure and install Poetry
# Based on https://stackoverflow.com/questions/72465421/how-to-use-poetry-with-docker
ENV POETRY_VERSION=1.3.2
ENV POETRY_HOME=/opt/poetry
ENV POETRY_VENV=/opt/poetry-venv
ENV POETRY_CACHE_DIR=/opt/.cache

# Install poetry separated from system interpreter
RUN python -m venv $POETRY_VENV \
    && $POETRY_VENV/bin/pip install -U pip setuptools \
    && $POETRY_VENV/bin/pip install poetry==${POETRY_VERSION}

# Add `poetry` to PATH
ENV PATH="${PATH}:${POETRY_VENV}/bin"

WORKDIR $EXP_DIR
# ---- Same until here with the base Dockerfile ----

# Copy and install dependencies seperately first to avoid triggering
# a dependency reinstallation on every change in the code
# Copy using poetry.lock* in case it doesn't exist yet
COPY ./pyproject.toml ./poetry.lock* $EXP_DIR
# Create some content required files to make poetry happy
RUN touch ./README.md \
    && poetry install --no-root --no-cache --with dev --with docker
