diff --git a/.github/workflows/docker-tt-smi.yml b/.github/workflows/docker-tt-smi.yml
index d60a4cf487..444266ccdb 100644
--- a/.github/workflows/docker-tt-smi.yml
+++ b/.github/workflows/docker-tt-smi.yml
@@ -11,6 +11,11 @@ on:
         description: "TT SMI version"
         required: true
         default: "3.0.25"
+      tag_latest:
+        description: "Update 'latest'"
+        type: boolean
+        required: true
+        default: false
 
 jobs:
   build-tt-smi:
@@ -38,8 +43,12 @@ jobs:
             --build-arg IMAGE_NAME=${IMAGE_NAME} \
             --build-arg TT_SMI_VERSION=${{ inputs.tt_smi_version }} \
             --build-arg BUILD_DATE=$(date --utc --iso-8601=seconds)Z \
-            --tag ${IMAGE_NAME}:latest
-          VERSION=$(docker inspect --format '{{ index .Config.Labels "org.opencontainers.image.version" }}' ${IMAGE_NAME})
-          docker tag ${IMAGE_NAME}:latest ${IMAGE_NAME}:${VERSION}
+            --tag ${IMAGE_NAME}:${{ inputs.tt_smi_version }}
+          VERSION=$(docker inspect --format '{{ index .Config.Labels "org.opencontainers.image.version" }}' ${IMAGE_NAME}:${{ inputs.tt_smi_version }})
+          docker tag ${IMAGE_NAME}:${{ inputs.tt_smi_version }} ${IMAGE_NAME}:${VERSION}
           docker push ${IMAGE_NAME}:${VERSION}
-          docker push ${IMAGE_NAME}:latest
+      - name: Tag and push latest
+        if: ${{ inputs.tag_latest }}
+        run: |
+          docker tag ${{ inputs.image_name }}:${{ inputs.tt_smi_version }} ${{ inputs.image_name }}:latest
+          docker push ${{ inputs.image_name }}:latest
diff --git a/docker/tt-smi/Dockerfile b/docker/tt-smi/Dockerfile
index b8417e4259..da190db313 100644
--- a/docker/tt-smi/Dockerfile
+++ b/docker/tt-smi/Dockerfile
@@ -1,4 +1,4 @@
-FROM ubuntu:20.04
+FROM ubuntu:22.04
 
 ARG IMAGE_NAME
 ARG TT_SMI_VERSION
@@ -10,7 +10,8 @@ RUN \
     apt-get update && \
     apt-get install -y curl git python3 python3-pip && \
     curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y && \
-    pip3 install --no-cache-dir git+https://github.com/tenstorrent/tt-smi@v${TT_SMI_VERSION} && \
+    pip3 install --no-cache-dir --upgrade pip setuptools wheel tomli && \
+    pip3 install --no-cache-dir --no-build-isolation git+https://github.com/tenstorrent/tt-smi@v${TT_SMI_VERSION} && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*