Update Dockerfile

Dockerfile

@@ -80,6 +80,11 @@ RUN --mount=target=/root/packages.txt,source=packages.txt \
 RUN --mount=target=/root/on_startup.sh,source=on_startup.sh,readwrite \
 	bash /root/on_startup.sh
 
+# Install ninja, which might help speed up flash attn build time
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        ninja-build \
+        && rm -rf /var/lib/apt/lists/*
+
 #######################################
 # End root user section
 #######################################
@@ -91,11 +96,6 @@ RUN curl https://sh.rustup.rs -sSf | bash -s -- -y
 ENV PATH="$HOME/.cargo/bin:${PATH}"
 RUN chmod -R 777 $HOME/.cargo/
 
-# Install ninja, which might help speed up flash attn build time
-RUN apt-get update && apt-get install -y --no-install-recommends \
-        ninja-build \
-        && rm -rf /var/lib/apt/lists/*
-
 # Python packages
 RUN --mount=target=requirements.txt,source=requirements.txt \
     pip install --no-cache-dir --upgrade -r requirements.txt