FROM nvidia/cuda:11.8.0-devel-ubuntu22.04 as build

RUN apt-get update && \
    apt-get install -y git python3-pip python3-venv wget unzip && \
    rm -rf /var/lib/apt/lists/*
RUN pip3 install --upgrade pip setuptools wheel

RUN git clone https://git.evulid.cc/cyberes/local-llm-server.git /local-llm-server

RUN python3 -m venv /jupyterlab
RUN /jupyterlab/bin/pip install jupyterlab
RUN /jupyterlab/bin/jupyter labextension disable "@jupyterlab/apputils-extension:announcements"

RUN mkdir -p /app
RUN wget https://github.com/rapiz1/rathole/releases/download/v0.4.8/rathole-x86_64-unknown-linux-gnu.zip -O /tmp/rathole.zip
RUN unzip -j /tmp/rathole.zip -d /tmp
RUN rm /tmp/rathole.zip
RUN cp /tmp/rathole /app

RUN python3 -m venv /venv

# Install PyTorch before installing VLLM in an attempt to ensure we use the right
# version for our CUDA install. (VLLM wants 2.0.1)
RUN /venv/bin/pip3 install torch==2.0.1 --index-url https://download.pytorch.org/whl/cu118

WORKDIR /local-llm-server

# We don't need to rebuild VLLM every time we build the container. But if we need
# to, uncomment the following line.
# ADD "https://www.random.org/cgi-bin/randbyte?nbytes=10&format=h" skipcache
RUN /venv/bin/pip install git+https://github.com/vllm-project/vllm

# Uninstall PyTorch since VLLM may have installed a version that is incompatible with
# our CUDA version.
RUN /venv/bin/pip3 uninstall -y torch
RUN /venv/bin/pip3 install torch==2.0.1 --index-url https://download.pytorch.org/whl/cu118


FROM nvidia/cuda:11.8.0-base-ubuntu22.04 as runtime

RUN apt-get update && apt-get install -y supervisor && rm -rf /var/lib/apt/lists/*

RUN useradd -ms /bin/bash apiserver
RUN usermod -s /bin/bash root

RUN apt-get update && \
    apt-get install -y python3 python3-pip wget aria2 git-lfs git openssh-server openssh-client nano tmux file && \
    rm -rf /var/lib/apt/lists/*

RUN pip3 install --upgrade pip setuptools wheel
RUN pip3 install glances

# Enable root SSH login
RUN sed -i 's/#PermitRootLogin prohibit-password/PermitRootLogin yes/' /etc/ssh/sshd_config

# Disable password SSH login
RUN sed -i 's/#PasswordAuthentication yes/PasswordAuthentication no/' /etc/ssh/sshd_config

# Create the necessary directory for sshd
RUN mkdir /var/run/sshd

COPY --from=build /local-llm-server /local-llm-server
COPY --from=build /venv /venv
COPY --from=build /app /app
COPY --from=build /jupyterlab /jupyterlab

COPY supervisord.conf /etc/supervisor/conf.d/supervisord.conf
COPY start-vllm.sh /app/start-vllm.sh
COPY start-container.sh /app/start.sh

RUN mkdir -p /var/log/app/

RUN chown -R apiserver:apiserver /local-llm-server && \
    chown -R apiserver:apiserver /app && \
    chown -R apiserver:apiserver /var/log/app/
RUN git config --global --add safe.directory /local-llm-server

RUN chmod +x /app/start.sh

ENV SHELL="/bin/bash"

# Expose Jupyter. We don't need to expose
# VLLM or SSH since rathole will tunnel those.
EXPOSE 8888

CMD /app/start.sh