ARG BASE_IMAGE="nvidia/cuda:12.2.2-devel-ubuntu22.04"

FROM ${BASE_IMAGE}
ARG HOMEDIRECTORY="/myhome"
ENV HOMEDIRECTORY=$HOMEDIRECTORY

USER root
RUN apt-get update && \
   apt-get install -y  --no-install-recommends  \
    curl \
    python3 \
    python3-pip \
    python3-dev \
    poppler-utils \
    gcc \
    git \
    git-lfs \
    htop \
    libgl1 \
    libglib2.0-0 \
    ncdu \
    openssh-client \
    openssh-server \
    psmisc \
    rsync \
    screen \
    sudo \
    tmux \
    unzip \
    vim \
    wget && \
    wget -q  https://github.com/justjanne/powerline-go/releases/download/v1.24/powerline-go-linux-"$(dpkg --print-architecture)" -O /usr/local/bin/powerline-shell && \
    chmod a+x /usr/local/bin/powerline-shell

RUN ln -s /usr/bin/python3 /usr/bin/python
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cu117


# setup ssh
RUN ssh-keygen -A
RUN sed -i 's/#*PermitRootLogin prohibit-password/PermitRootLogin yes/' /etc/ssh/sshd_config
EXPOSE 22

# Make the root user's home directory /myhome (the default for run.ai),
# and allow to login with password 'root'.
RUN echo 'root:root' | chpasswd
RUN sed -i 's|:root:/root:|:root:/myhome:|' /etc/passwd

ENTRYPOINT sudo service ssh start && /bin/bash