Skip to content

Commit

Permalink
Reduce CI Docker image size (#660)
Browse files Browse the repository at this point in the history
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
  • Loading branch information
vmilosevic authored Sep 12, 2024
1 parent 665fc4c commit 235d441
Show file tree
Hide file tree
Showing 5 changed files with 22 additions and 13 deletions.
25 changes: 17 additions & 8 deletions .github/Dockerfile.ci
Original file line number Diff line number Diff line change
@@ -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 && \
Expand All @@ -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
Expand All @@ -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
7 changes: 3 additions & 4 deletions env/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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)
1 change: 1 addition & 0 deletions runtime/tools/python/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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}
Expand Down
1 change: 1 addition & 0 deletions runtime/tools/python/requirements.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
torch==2.3.0 --index-url https://download.pytorch.org/whl/cpu
1 change: 0 additions & 1 deletion runtime/tools/python/setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down

0 comments on commit 235d441

Please sign in to comment.