Build docker image as part of a workflow

.github/Dockerfile.ci

@@ -1,11 +1,9 @@
-ARG GIT_SHA
-ARG FROM_TAG=${GIT_SHA:-latest}
+ARG FROM_TAG=latest
 
 FROM ghcr.io/tenstorrent/tt-xla/tt-xla-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-xla
 ENV BUILD_DIR=/home/build
 ENV TTMLIR_TOOLCHAIN_DIR=/opt/ttmlir-toolchain
@@ -18,8 +16,12 @@ RUN mkdir -p $BUILD_DIR && \
 # Copy the project from host, cloned in build-image.yml
 COPY . $BUILD_DIR/$PROJECT_NAME
 
-# Build the toolchain
 WORKDIR $BUILD_DIR/$PROJECT_NAME
+
+# Show last commit
+RUN git log -1
+
+# Build the toolchain
 RUN cmake -B toolchain -DTOOLCHAIN=ON third_party/ && \
     cd third_party/tt-mlir/src/tt-mlir/ && \
     source env/activate && \

.github/build-docker-images.sh

@@ -0,0 +1,52 @@
+#!/bin/bash
+# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+
+set -e
+
+REPO=tenstorrent/tt-xla
+BASE_IMAGE_NAME=ghcr.io/$REPO/tt-xla-base-ubuntu-22-04
+CI_IMAGE_NAME=ghcr.io/$REPO/tt-xla-ci-ubuntu-22-04
+
+# Compute the hash of the Dockerfile
+DOCKER_TAG=$(./.github/get-docker-tag.sh)
+echo "Docker tag: $DOCKER_TAG"
+
+# Are we on main branch
+ON_MAIN=$(git branch --show-current | grep -q main && echo "true" || echo "false")
+
+build_and_push() {
+    local image_name=$1
+    local dockerfile=$2
+    local on_main=$3
+
+
+    if docker manifest inspect $image_name:$DOCKER_TAG > /dev/null; then
+        echo "Image $image_name:$DOCKER_TAG already exists"
+    else
+        echo "Building image $image_name:$DOCKER_TAG"
+        docker build \
+            --progress=plain \
+            --build-arg FROM_TAG=$DOCKER_TAG \
+            -t $image_name:$DOCKER_TAG \
+            -t $image_name:latest \
+            -f $dockerfile .
+
+        echo "Pushing image $image_name:$DOCKER_TAG"
+        docker push $image_name:$DOCKER_TAG
+
+        # If we are on main branch also push the latest tag
+        if [ "$on_main" = "true" ]; then
+            echo "Pushing image $image_name:latest"
+            docker push $image_name:latest
+        fi
+    fi
+}
+
+build_and_push $BASE_IMAGE_NAME .github/Dockerfile.base $ON_MAIN
+build_and_push $CI_IMAGE_NAME .github/Dockerfile.ci $ON_MAIN
+
+echo "All images built and pushed successfully"
+echo "CI_IMAGE_NAME:"
+echo $CI_IMAGE_NAME:$DOCKER_TAG

.github/get-docker-tag.sh

@@ -0,0 +1,33 @@
+#!/bin/bash
+# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+
+# Calculate hash for docker image tag.
+# The hash is based on the MLIR docker tag  and the hash of the Dockerfile(s).
+
+# Exit immediately if a command exits with a non-zero status
+set -e
+
+# Execute this in a separate bash process
+(
+    # Read tt-mlir version from third_party/CMakeLists.txt and clone third_party/tt-mlir
+    # Get the MLIR docker tag
+    TT_MLIR_VERSION=$(grep -oP 'set\(TT_MLIR_VERSION "\K[^"]+' third_party/CMakeLists.txt)
+    if [ ! -d "third_party/tt-mlir" ]; then
+        git clone https://github.com/tenstorrent/tt-mlir.git third_party/tt-mlir --quiet
+    fi
+    cd third_party/tt-mlir
+    git fetch --quiet
+    git checkout $TT_MLIR_VERSION --quiet
+    if [ -f ".github/get-docker-tag.sh" ]; then
+        MLIR_DOCKER_TAG=$(.github/get-docker-tag.sh)
+    else
+        MLIR_DOCKER_TAG="default-tag"
+    fi
+    cd ../..
+)
+
+DOCKERFILE_HASH_FILES=".github/Dockerfile.base .github/Dockerfile.ci"
+DOCKERFILE_HASH=$( (echo $MLIR_DOCKER_TAG; sha256sum $DOCKERFILE_HASH_FILES) | sha256sum | cut -d ' ' -f 1)
+echo dt-$DOCKERFILE_HASH

.github/workflows/build-and-test.yml

@@ -81,11 +81,41 @@ jobs:
   #       echo "version=$version" >> $GITHUB_ENV
   #       echo $version
 
+  build-image:
+    runs-on: n300
+    outputs:
+      docker-image: ${{ steps.build.outputs.docker-image }}
+    steps:
+      - name: Fix permissions
+        shell: bash
+        run: sudo chown ubuntu:ubuntu -R $(pwd)
+
+      - name: Checkout repository
+        uses: actions/checkout@v4
+
+      - name: Log in to GitHub Container Registry
+        uses: docker/login-action@v3
+        with:
+          registry: ghcr.io
+          username: ${{ github.repository_owner }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Build Docker images and output the image name
+        id: build
+        shell: bash
+        run: |
+          # Output the image name
+          set pipefail
+          .github/build-docker-images.sh | tee docker.log
+          DOCKER_CI_IMAGE=$(tail -n 1 docker.log)
+          echo "DOCKER_CI_IMAGE $DOCKER_CI_IMAGE"
+          echo "docker-image=$DOCKER_CI_IMAGE" >> "$GITHUB_OUTPUT"
 
   # # Run tests on TT hardware
 
   build-and-run-tests:
     timeout-minutes: 120
+    needs: build-image
     strategy:
       fail-fast: false
       matrix:
@@ -99,8 +129,8 @@ jobs:
       - ${{ matrix.build.runs-on }}
 
     container:
-      image: ghcr.io/tenstorrent/tt-xla/tt-xla-ci-ubuntu-22-04:latest
-      options: --user root --device /dev/tenstorrent/0
+      image: ${{ needs.build-image.outputs.docker-image }}
+      options: --device /dev/tenstorrent/0
       volumes:
         - /dev/hugepages:/dev/hugepages
         - /dev/hugepages-1G:/dev/hugepages-1G