From 235d441fd14c40566c6afd3d54ef16ae8dbe141f Mon Sep 17 00:00:00 2001
From: Vladimir Milosevic <157983820+vmilosevic@users.noreply.github.com>
Date: Thu, 12 Sep 2024 17:05:32 +0200
Subject: [PATCH] Reduce CI Docker image size (#660)

Reduce CI Docker image size

Our CI Docker image is a lot bigger than the recommended image size because we use it to cache and preinstall dependencies, the biggest of which is LLVM. This causes issues like long image pull time and fills up the docker cache. Also, having a big image prevents us from using Github built-in runners.

The current image size is about 27GB, with this fix, it should go down to 11GB.

Update Cmake env build to build dependencies locally in tt-mlir/env/build folder, and only install to the toolchain folder.
Use multi-stage build so only the toolchain folder is transferred to the final image
---
 .github/Dockerfile.ci                 | 25 +++++++++++++++++--------
 env/CMakeLists.txt                    |  7 +++----
 runtime/tools/python/CMakeLists.txt   |  1 +
 runtime/tools/python/requirements.txt |  1 +
 runtime/tools/python/setup.py         |  1 -
 5 files changed, 22 insertions(+), 13 deletions(-)
 create mode 100644 runtime/tools/python/requirements.txt

diff --git a/.github/Dockerfile.ci b/.github/Dockerfile.ci
index d96206b29..c5eed0727 100644
--- a/.github/Dockerfile.ci
+++ b/.github/Dockerfile.ci
@@ -1,18 +1,19 @@
 ARG GIT_SHA
 ARG FROM_TAG=${GIT_SHA:-latest}
 
-FROM ghcr.io/tenstorrent/tt-mlir/tt-mlir-base-ubuntu-22-04:${FROM_TAG}
+FROM ghcr.io/tenstorrent/tt-mlir/tt-mlir-base-ubuntu-22-04:${FROM_TAG} AS ci-build
 SHELL ["/bin/bash", "-c"]
 
 # Create a directory for the build and toolchain
+ARG GIT_SHA
 ENV PROJECT_NAME=tt-mlir
 ENV BUILD_DIR=/home/build
 ENV TTMLIR_TOOLCHAIN_DIR=/opt/ttmlir-toolchain
-RUN mkdir -p $BUILD_DIR && \
-    mkdir -p $TTMLIR_TOOLCHAIN_DIR
 
+RUN echo "Building $PROJECT_NAME at $GIT_SHA"
 
-ARG GIT_SHA
+RUN mkdir -p $BUILD_DIR && \
+    mkdir -p $TTMLIR_TOOLCHAIN_DIR
 
 # Clone the project and update submodules
 RUN git clone https://github.com/tenstorrent/$PROJECT_NAME.git $BUILD_DIR/$PROJECT_NAME && \
@@ -21,7 +22,8 @@ RUN git clone https://github.com/tenstorrent/$PROJECT_NAME.git $BUILD_DIR/$PROJE
 
 # Build the toolchain
 WORKDIR $BUILD_DIR/$PROJECT_NAME
-RUN cmake -B env/build env && \
+RUN source env/activate && \
+    cmake -B env/build env && \
     cmake --build env/build
 
 # Self-test
@@ -46,11 +48,18 @@ RUN source env/activate && \
 
 # Run clang-tidy
 RUN source env/activate && \
-    cmake --build build -- clang-tidy || true
+    cmake --build build -- clang-tidy
 
 # Run the tests
 RUN source env/activate && \
     cmake --build build -- check-ttmlir
 
-# Clean up the build directory
-RUN rm -rf $BUILD_DIR/$PROJECT_NAME
+# Final stage
+FROM ghcr.io/tenstorrent/tt-mlir/tt-mlir-base-ubuntu-22-04:${FROM_TAG} AS ci
+
+# Copy the TTMLIR_TOOLCHAIN_DIR from the previous stage
+ENV TTMLIR_TOOLCHAIN_DIR=/opt/ttmlir-toolchain
+RUN echo "Copying from ci-build stage $TTMLIR_TOOLCHAIN_DIR"
+COPY --from=ci-build $TTMLIR_TOOLCHAIN_DIR $TTMLIR_TOOLCHAIN_DIR
+
+RUN du -h --max-depth=2 $TTMLIR_TOOLCHAIN_DIR
diff --git a/env/CMakeLists.txt b/env/CMakeLists.txt
index 72dc18ddd..883d0bc30 100644
--- a/env/CMakeLists.txt
+++ b/env/CMakeLists.txt
@@ -21,7 +21,6 @@ add_custom_target(python-venv ALL COMMAND TTMLIR_TOOLCHAIN=${TTMLIR_TOOLCHAIN_DI
 
 ExternalProject_Add(
     flatbuffers
-    PREFIX ${TTMLIR_TOOLCHAIN_DIR}
     CMAKE_GENERATOR Ninja
     CMAKE_ARGS
     -DCMAKE_BUILD_TYPE=Release
@@ -36,9 +35,8 @@ ExternalProject_Add(
 # -DMLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS=ON set if debug
 ExternalProject_Add(
     llvm-project
-    PREFIX ${TTMLIR_TOOLCHAIN_DIR}
     # Super hacky way to install the python dependencies before the build
-    PATCH_COMMAND bash -c "source ${CMAKE_CURRENT_SOURCE_DIR}/activate && pip install -r ${TTMLIR_TOOLCHAIN_DIR}/src/llvm-project/mlir/python/requirements.txt"
+    PATCH_COMMAND bash -c "source ${CMAKE_CURRENT_SOURCE_DIR}/activate && pip install -r mlir/python/requirements.txt"
     CMAKE_GENERATOR Ninja
     CMAKE_ARGS
     -DPython3_FIND_VIRTUALENV=ONLY
@@ -77,4 +75,5 @@ ExternalProject_Add(stablehlo
    INSTALL_COMMAND ""
 )
 
-add_custom_target(llvm-lit ALL COMMAND cp ${TTMLIR_TOOLCHAIN_DIR}/src/llvm-project-build/bin/llvm-lit ${TTMLIR_TOOLCHAIN_DIR}/bin/llvm-lit DEPENDS llvm-project)
+add_custom_target(llvm-lit ALL COMMAND cp llvm-project-prefix/src/llvm-project-build/bin/llvm-lit ${TTMLIR_TOOLCHAIN_DIR}/bin/llvm-lit DEPENDS llvm-project)
+add_custom_target(run-clang-tidy-install ALL COMMAND cp llvm-project-prefix/src/llvm-project/clang-tools-extra/clang-tidy/tool/run-clang-tidy.py ${TTMLIR_TOOLCHAIN_DIR}/bin/run-clang-tidy.py DEPENDS llvm-project)
diff --git a/runtime/tools/python/CMakeLists.txt b/runtime/tools/python/CMakeLists.txt
index 84810a4cf..e58bec029 100644
--- a/runtime/tools/python/CMakeLists.txt
+++ b/runtime/tools/python/CMakeLists.txt
@@ -5,6 +5,7 @@ add_custom_target(ttrt-copy-files
 
 add_custom_target(ttrt
   COMMAND rm -f build/*.whl
+  COMMAND python -m pip install -r requirements.txt
   COMMAND TTMLIR_ENABLE_RUNTIME=${TTMLIR_ENABLE_RUNTIME}
           TT_RUNTIME_ENABLE_TTNN=${TT_RUNTIME_ENABLE_TTNN}
           TT_RUNTIME_ENABLE_TTMETAL=${TT_RUNTIME_ENABLE_TTMETAL}
diff --git a/runtime/tools/python/requirements.txt b/runtime/tools/python/requirements.txt
new file mode 100644
index 000000000..8bfab8347
--- /dev/null
+++ b/runtime/tools/python/requirements.txt
@@ -0,0 +1 @@
+torch==2.3.0 --index-url https://download.pytorch.org/whl/cpu
diff --git a/runtime/tools/python/setup.py b/runtime/tools/python/setup.py
index 19e693a75..6f3f701a3 100644
--- a/runtime/tools/python/setup.py
+++ b/runtime/tools/python/setup.py
@@ -73,7 +73,6 @@
 
 if enable_runtime:
     assert enable_ttmetal or enable_ttnn, "At least one runtime must be enabled"
-    install_requires += ["torch"]
 
     for dylib in runlibs:
         shutil.copy(