FROM pytorch/pytorch:2.1.2-cuda12.1-cudnn8-devel

ARG PIP_SOURCE=https://pypi.tuna.tsinghua.edu.cn/simple
ARG BUILD_TMP=/build_tmp

RUN mkdir $BUILD_TMP

# install basic dependencies
COPY sources.list /etc/apt/sources.list
RUN apt-get clean all && apt-get update && apt-get install -y --no-install-recommends \
    python3.8 \
    git \
    curl \
    wget

# install python requirements
RUN pip install flash-attn==2.4.2 --no-build-isolation

COPY requirements.txt $BUILD_TMP/.
RUN pip install -i $PIP_SOURCE -r $BUILD_TMP/requirements.txt

# set local
ENV LANG C.UTF-8 LC_ALL=C.UTF-8

# clean up
RUN rm -rf $BUILD_TMP

# initialize workspace
WORKDIR /workspace

CMD ["/bin/bash", "/workspace/start.sh"]