# Based on https://github.com/astral-sh/uv-docker-example/blob/main/Dockerfile

# Use a PyTorch base image instead of a uv base image - easier to install uv than CUDA ;)
FROM pytorch/pytorch:2.4.1-cuda11.8-cudnn9-runtime

# Set up build-time environment variables
ARG PYTORCH_BUILD_TYPE
ENV PYTORCH_BUILD_TYPE=${PYTORCH_BUILD_TYPE}

# Install system dependencies for uv installer
RUN apt-get update && apt-get install -y --no-install-recommends curl ca-certificates && \
    rm -rf /var/lib/apt/lists/*

# Install uv using the installer
ADD https://astral.sh/uv/0.5.7/install.sh /uv-installer.sh
RUN sh /uv-installer.sh && rm /uv-installer.sh
ENV PATH="/root/.local/bin:$PATH"

# Enable bytecode compilation and copy from the cache instead of linking since it's a mounted volume.
# These are docker optimizations listed in uv docs.
ENV UV_COMPILE_BYTECODE=1
ENV UV_LINK_MODE=copy

# Install the project into the root directory
WORKDIR /

# Creating a venv before syncing lets us use the system site packages, namely pytorch
RUN bash -c "\
    uv venv --system-site-packages && \
    . .venv/bin/activate \
"

# Install the project's dependencies using the lockfile and settings.
# Pytorch build type is passed as an environment variable.
# To run pytorch on the GPU with CUDA, we can use the version installed in the base image.
# To run it on the CPU, we need a different version of pytorch for that device,
# which is defined as an extra in pyproject.toml.
RUN --mount=type=cache,target=/root/.cache/uv \
    --mount=type=bind,source=uv.lock,target=uv.lock \
    --mount=type=bind,source=pyproject.toml,target=pyproject.toml \
    if [ "$PYTORCH_BUILD_TYPE" = "cpu" ]; then \
        uv sync --frozen --no-install-project --no-dev --verbose --extra cpu; \
    else \
        uv sync --frozen --no-install-project --no-dev --verbose; \
    fi

# Then, add the rest of the project source code and install it.
# Installing separately from its dependencies allows optimal layer caching.
ADD . /
RUN --mount=type=cache,target=/root/.cache/uv \
    uv sync --frozen --no-dev --all-packages

# Make the entrypoint script executable.
# Can be used with: docker run <image> <experiment_name> <config_file>
RUN chmod +x /bls_autocoders_experiments/entrypoint.sh

# Place executables in the environment at the front of the path
ENV PATH="/.venv/bin:$PATH"