From 839a95cb628ad89ebf1b5ea1a06447b2bc977e46 Mon Sep 17 00:00:00 2001
From: Abhishek Thakur <abhishek4@gmail.com>
Date: Wed, 22 Nov 2023 12:18:22 +0100
Subject: [PATCH] smaller docker image

---
 .dockerignore |  3 ++-
 Dockerfile    | 29 ++++++++++++++++-------------
 2 files changed, 18 insertions(+), 14 deletions(-)

diff --git a/.dockerignore b/.dockerignore
index dc9169a8fc..3e7d946c42 100644
--- a/.dockerignore
+++ b/.dockerignore
@@ -8,4 +8,5 @@ test.py
 .DS_Store
 .vscode/
 op*
-op_*
\ No newline at end of file
+op_*
+.git
\ No newline at end of file
diff --git a/Dockerfile b/Dockerfile
index becac787ed..ba7bf78e59 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,4 +1,4 @@
-FROM nvidia/cuda:12.1.0-cudnn8-devel-ubuntu22.04
+FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04
 
 ENV DEBIAN_FRONTEND=noninteractive \
     TZ=UTC
@@ -6,10 +6,10 @@ ENV DEBIAN_FRONTEND=noninteractive \
 ENV PATH="${HOME}/miniconda3/bin:${PATH}"
 ARG PATH="${HOME}/miniconda3/bin:${PATH}"
 
-RUN mkdir -p /tmp/model
-RUN chown -R 1000:1000 /tmp/model
-RUN mkdir -p /tmp/data
-RUN chown -R 1000:1000 /tmp/data
+RUN mkdir -p /tmp/model && \
+    chown -R 1000:1000 /tmp/model && \
+    mkdir -p /tmp/data && \
+    chown -R 1000:1000 /tmp/data
 
 RUN apt-get update &&  \
     apt-get upgrade -y &&  \
@@ -28,7 +28,8 @@ RUN apt-get update &&  \
     git \
     git-lfs \
     libgl1 \
-    && rm -rf /var/lib/apt/lists/*
+    && rm -rf /var/lib/apt/lists/* && \
+    apt-get clean
 
 
 RUN curl -s https://packagecloud.io/install/repositories/github/git-lfs/script.deb.sh | bash && \
@@ -58,11 +59,13 @@ RUN conda create -p /app/env -y python=3.10
 
 SHELL ["conda", "run","--no-capture-output", "-p","/app/env", "/bin/bash", "-c"]
 
-RUN conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia && conda clean -ya
-COPY --chown=1000:1000 . /app/
-
-RUN pip install -e .
+RUN conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia && \
+    conda clean -ya && \
+    conda install -c "nvidia/label/cuda-12.1.0" cuda-nvcc && conda clean -ya
 
-RUN python -m nltk.downloader punkt
-RUN autotrain setup
-RUN pip install flash-attn
\ No newline at end of file
+COPY --chown=1000:1000 . /app/
+RUN pip install -e . && \
+    python -m nltk.downloader punkt && \
+    autotrain setup && \
+    pip install flash-attn && \
+    pip cache purge
\ No newline at end of file