matrix-pygmalion-bot/runpod/runpod-worker-transformers/Dockerfile

#ARG BASE_IMAGE=nvidia/cuda:12.0.1-cudnn8-runtime-ubuntu22.04
#ARG BASE_IMAGE=nvidia/cuda:11.6.2-cudnn8-devel-ubuntu20.04
ARG BASE_IMAGE=runpod/pytorch:3.10-2.0.0-117
#ARG BASE_IMAGE=nvcr.io/nvidia/pytorch:23.03-py3
FROM ${BASE_IMAGE} as dev-base

ARG MODEL_NAME
ENV MODEL_NAME=${MODEL_NAME}

WORKDIR /
SHELL ["/bin/bash", "-o", "pipefail", "-c"]
ENV DEBIAN_FRONTEND noninteractive\
    SHELL=/bin/bash

RUN apt-get update --yes && \
    # - apt-get upgrade is run to patch known vulnerabilities in apt-get packages as
    #   the ubuntu base image is rebuilt too seldom sometimes (less than once a month)
    #apt-get upgrade --yes && \
    apt install --yes --no-install-recommends \
#      build-essential \
      ca-certificates \
      git \
      git-lfs \
      wget \
      curl \
      bash \
      libgl1 \
      software-properties-common \
      openssh-server && \
      apt-get clean && rm -rf /var/lib/apt/lists/* && \
      echo "en_US.UTF-8 UTF-8" > /etc/locale.gen

#RUN apt-key del 7fa2af80 && \
#    apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/3bf863cc.pub
#
#RUN add-apt-repository ppa:deadsnakes/ppa && \
#    apt-get install python3.10 python3.10-dev python3.10-venv python3-pip -y --no-install-recommends && \
#    update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.10 1 && \
#    update-alternatives --install /usr/bin/python python /usr/bin/python3 1 && \
##    update-alternatives --install /usr/bin/pip pip /usr/bin/pip3 1 && \
#    update-alternatives --set python3 /usr/bin/python3.10 && \
#    update-alternatives --set python /usr/bin/python3 && \
#    apt-get clean && rm -rf /var/lib/apt/lists/*

RUN apt-get update --yes && \
    apt install --yes --no-install-recommends \
      python3 python3-dev python3-venv python3-pip && \
    apt-get clean && rm -rf /var/lib/apt/lists/*

RUN pip3 install --upgrade pip && \
#    pip3 install torch torchvision torchaudio --extra-index-url=https://download.pytorch.org/whl/cu118 && \
    pip3 install bitsandbytes && \
    pip3 install safetensors && \
    pip3 install diffusers && \
    pip3 install transformers accelerate xformers triton && \
    pip3 install huggingface-hub && \
    pip3 install runpod

RUN mkdir -p /workspace
WORKDIR /workspace

COPY model_fetcher.py /workspace/
RUN python3 model_fetcher.py --model_name=${MODEL_NAME}
#RUN git lfs install && \
#    git clone --depth 1 https://huggingface.co/${MODEL_NAME}

COPY runpod_infer.py /workspace/
COPY test_input.json /workspace/

CMD python3 -u runpod_infer.py --model_name=${MODEL_NAME}