ARG PYTHON_VERSION=310
ARG CUDA_VERSION=cu121
ARG BASE_IMAGE=quay.io/ado/ado:latest-py${PYTHON_VERSION}-${CUDA_VERSION}
FROM $BASE_IMAGE

# Install prerequisites for other packages
# AP: we need to install torch before flash-attn
RUN uv pip install --system "torch>=2.4,<2.5.0" && \
    uv pip install --system --no-build-isolation "flash-attn==2.6.3"

## Install ado-sfttrainer
RUN cd ado && \
    uv pip install --system \
    'plugins/actuators/sfttrainer'