From a21f9b58cf784f98e04a2dd11fb128ca3327da6b Mon Sep 17 00:00:00 2001
From: Torsten Scholak <torsten.scholak@googlemail.com>
Date: Mon, 21 Oct 2024 08:31:35 -0400
Subject: [PATCH] Optimize Docker Build Layers and Add Sudo Privileges for
 Fast-LLM Container (#2)

---
 Dockerfile | 27 +++++++++++++++++----------
 1 file changed, 17 insertions(+), 10 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 2302c8bb..b12f966b 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,30 +1,37 @@
+# syntax=docker/dockerfile:1.7-labs
 FROM nvcr.io/nvidia/pytorch:24.07-py3
 
-# git-lfs is needed to interact with the huggingface hub
+# Install git-lfs for Huggingface hub interaction and sudo for system adjustments
 RUN apt-get update \
-    && apt-get install git-lfs \
+    && apt-get install --no-install-recommends -y git-lfs sudo util-linux \
     && rm -rf /var/lib/apt/lists/* \
     && git lfs install
 
+# Add a user for Fast-LLM with sudo privileges for runtime adjustments
 ARG FAST_LLM_USER_ID=1000
+RUN useradd -m -u $FAST_LLM_USER_ID -s /bin/bash fast_llm \
+    && echo 'fast_llm ALL=(ALL) NOPASSWD: ALL' >> /etc/sudoers
 
-RUN useradd -m -u $FAST_LLM_USER_ID -s /bin/bash -d /home/fast_llm fast_llm
 USER fast_llm
 WORKDIR /app
-ENV PYTHONPATH=/app:/app:/app/Megatron-LM
-ENV PATH=$PATH:/app:/home/fast_llm/.local/bin/
 
-COPY --chown=fast_llm setup.py setup.cfg ./
-COPY --chown=fast_llm fast_llm/__init__.py ./fast_llm/
+# Environment settings for Python and PATH
+ENV PYTHONPATH=/app:/app/Megatron-LM \
+    PATH=$PATH:/home/fast_llm/.local/bin/
 
+# Copy the dependency files and install dependencies
+COPY --chown=fast_llm setup.py setup.cfg pyproject.toml ./
 RUN PIP_NO_INPUT=1 pip3 install --no-cache-dir -e ".[CORE,OPTIONAL,DEV]"
 
-COPY --chown=fast_llm fast_llm/csrc/ ./fast_llm/csrc/
+# Compile the C++ extensions (fast_llm/csrc)
+COPY --chown=fast_llm ./fast_llm/csrc/ fast_llm/csrc/
 RUN make -C ./fast_llm/csrc/
 
+# Copy the rest of the code
 COPY --chown=fast_llm ./Megatron-LM Megatron-LM
 COPY --chown=fast_llm ./examples examples
 COPY --chown=fast_llm ./tests tests
 COPY --chown=fast_llm ./tools tools
-COPY --chown=fast_llm ./fast_llm fast_llm
-COPY --chown=fast_llm fast_llm/tools/train.py pyproject.toml ./
+
+# Copy the main source code for Fast-LLM and install in editable mode
+COPY --exclude=./fast_llm/csrc/ --chown=fast_llm ./fast_llm/ fast_llm/