Update Dockerfile

Dockerfile

@@ -1,52 +1,34 @@
 FROM nvidia/cuda:12.5.1-cudnn-devel-ubuntu20.04
 
 ENV DEBIAN_FRONTEND=noninteractive \
-    TZ=Asia/Kolkata
+    TZ=Asia/Kolkata \
+    PYTHONUNBUFFERED=1
 
-# Remove any third-party apt sources to avoid issues with expiring keys.
-RUN rm -f /etc/apt/sources.list.d/*.list && \
+# Basic deps
+RUN RUN rm -f /etc/apt/sources.list.d/*.list && \
     apt-get update && apt-get install -y --no-install-recommends \
-    curl \
-    ca-certificates \
-    sudo \
-    git \
-    wget \
-    procps \
-    git-lfs \
-    zip \
-    unzip \
-    htop \
-    vim \
-    nano \
-    bzip2 \
-    libx11-6 \
-    build-essential \
-    libsndfile-dev \
-    software-properties-common \
- && rm -rf /var/lib/apt/lists/*
-
+    python3 python3-pip python3-venv git curl wget ca-certificates && \
+    rm -rf /var/lib/apt/lists/*
+    
 RUN add-apt-repository ppa:flexiondotorg/nvtop && \
     apt-get upgrade -y && \
     apt-get install -y --no-install-recommends nvtop
 
-RUN curl -sL https://deb.nodesource.com/setup_21.x  | bash - && \
-    apt-get install -y nodejs && \
-    npm install -g configurable-http-proxy
-
 # Create a working directory
 WORKDIR /app
 
-# Create non-root user
+# Create a non-root user and switch to it
 RUN adduser --disabled-password --gecos '' --shell /bin/bash user \
  && chown -R user:user /app
 RUN echo "user ALL=(ALL) NOPASSWD:ALL" > /etc/sudoers.d/90-user
 USER user
 
+# All users can use /home/user as their home directory
 ENV HOME=/home/user
 RUN mkdir $HOME/.cache $HOME/.config \
  && chmod -R 777 $HOME
 
-# Install Miniconda
+# Set up the Conda environment
 ENV CONDA_AUTO_UPDATE_CONDA=false \
     PATH=$HOME/miniconda/bin:$PATH
 RUN curl -sLo ~/miniconda.sh https://repo.continuum.io/miniconda/Miniconda3-py39_4.10.3-Linux-x86_64.sh \
@@ -58,41 +40,62 @@ RUN curl -sLo ~/miniconda.sh https://repo.continuum.io/miniconda/Miniconda3-py39
 WORKDIR $HOME/app
 
 #######################################
-# Root section
+# Start root user section
 #######################################
+
 USER root
 
-# Optional extra packages
-RUN apt-get update && apt-get install -y --no-install-recommends \
-    python3-dev python3-pip python3-venv \
- && rm -rf /var/lib/apt/lists/*
+# User Debian packages
+## Security warning : Potential user code executed as root (build time)
+RUN --mount=target=/root/packages.txt,source=packages.txt \
+    apt-get update && \
+    xargs -r -a /root/packages.txt apt-get install -y --no-install-recommends \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN --mount=target=/root/on_startup.sh,source=on_startup.sh,readwrite \
+	bash /root/on_startup.sh
 
 RUN mkdir /data && chown user:user /data
 
 #######################################
-# End root section
+# End root user section
 #######################################
+
 USER user
 
-# Install Python requirements for vLLM
-RUN pip install --no-cache-dir --upgrade pip && \
-    pip install --no-cache-dir vllm
+# # Python packages
+# RUN --mount=target=requirements.txt,source=requirements.txt \
+#     pip install --no-cache-dir --upgrade -r requirements.txt
+
+# Upgrade pip
+RUN python3 -m pip install --upgrade pip
 
-# Download model at build time (optional)
-RUN python3 -m vllm.entrypoints.api_server --model unsloth/llama-2-7b-bnb-4bit --download-only
+# Install CUDA-compatible torch first, then vLLM
+RUN pip install --no-cache-dir --index-url https://download.pytorch.org/whl/cu121 \
+    torch torchvision torchaudio && \
+    pip install --no-cache-dir vllm
+    pip install vllm --torch-backend=auto
 
-# Copy source code
+# Copy the current directory contents into the container at $HOME/app setting the owner to the user
 COPY --chown=user . $HOME/app
 
 RUN chmod +x start_server.sh
 
-# Set env for vLLM
+
+# Default env vars
+ENV MODEL_NAME="unsloth/llama-2-7b-bnb-4bit" \
+    VLLM_PORT=7860 \
+    HOST=0.0.0.0 \
+    TP_SIZE=1 \
+    GPU_MEMORY_UTILIZATION=0.90
+
 ENV PYTHONUNBUFFERED=1 \
-    VLLM_HOST=0.0.0.0 \
-    VLLM_PORT=7860
+	GRADIO_ALLOW_FLAGGING=never \
+	GRADIO_NUM_PORTS=1 \
+	GRADIO_SERVER_NAME=0.0.0.0 \
+	GRADIO_THEME=huggingface \
+	SYSTEM=spaces \
+	SHELL=/bin/bash
 
-# Expose port
-EXPOSE 7860
+CMD ["./start_server.sh"]
 
-# Start vLLM API server
-CMD ["python3", "-m", "vllm.entrypoints.api_server", "--model", "unsloth/llama-2-7b-bnb-4bit", "--host", "0.0.0.0", "--port", "7860"]