FROM nvidia/cuda:12.1.1-devel-ubuntu22.04

RUN apt-get update \
  && DEBIAN_FRONTEND=noninteractive apt-get install \
       -y --no-install-recommends \
       postgresql \
       vim \
       python3 \
       python3-dev \
       python3-pip \
       git \
       ninja-build \
       openssh-server \
       tini \
  && rm -rf /var/lib/apt/lists/*

WORKDIR /transformer

ADD ./requirements.txt ./
ADD ./install_requirements.sh ./

RUN ./install_requirements.sh && pip cache remove '*'

RUN mkdir /root/.ssh \
    && chmod 700 /root/.ssh \
    && echo 'ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQC9iKDoH5tGGwR5KfFDr6FpukVGm5GwUE8FUlQkiklZ91puqY/lvhqMyuDEQufhJiPmGT0IJPe22bPoTLL2PC6OFxlQg2BwAzxs+iWLbjoo83M+hShjDSc+p2MdupwMVqrRnO+J+BgTf3zP+a6Ug+1kdmwie56y9bToSqVxxZN+rtAPuQ9aqCesRD7W3NYJ+9MtmyNJS4zK4OynoebWDDw6K7F0TF9X5GqZeAhgW10Y1z/kYnwWJNHkHXPOvT2kE/UpqGWQMLS61tGOU4fd9e9XFrWJmqjvKWqzPzLfxxdtyDg/Ou/QJfg3LrXqTJ8+pRF7o30hxLQG8FgbgakJkK35 m4burns@Marcs-MacBook-Air-2.local' > /root/.ssh/authorized_keys

RUN mkdir /workspace

ENV WANDB_MODE=disabled

RUN mkdir checkpoints
ADD ./t5-small ./t5-small
ADD ./*.py ./
ADD ./*.json ./
ADD ./run.sh ./

CMD [ "/usr/bin/tini", "/bin/sh", "--", "-c", "ulimit -n 524288 && service ssh start && sleep infinity" ]