From 5e5a465fb3426c78665b37f94dca4e8d08f5714f Mon Sep 17 00:00:00 2001
From: Alex Shraer <shralex@gmail.com>
Date: Sat, 13 Jun 2026 13:23:14 +0000
Subject: [PATCH] Configure gemini-investigate on build failure for
 UploadDockerImages.yml

---
 .github/workflows/UploadDockerImages.yml | 20 ++++++++++++++++++--
 1 file changed, 18 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/UploadDockerImages.yml b/.github/workflows/UploadDockerImages.yml
index 93f0b081bc..212bf7d7c4 100644
--- a/.github/workflows/UploadDockerImages.yml
+++ b/.github/workflows/UploadDockerImages.yml
@@ -95,15 +95,31 @@ jobs:
 
   notify_failure:
     name: Notify failed build
-    needs: [build-and-push]
-    if: ${{ failure() && inputs.image_suffix == '' }}
+    needs: [build_and_upload_maxtext_package, build-and-push]
+    if: ${{ always() }}
     runs-on: ubuntu-latest
     permissions:
       issues: write
     steps:
       - name: Create issue on failure
+        if: ${{ contains(needs.*.result, 'failure') && inputs.image_suffix == '' }}
         uses: jayqi/failed-build-issue-action@1a893bbf43ef1c2a8705e2b115cd4f0fe3c5649b
         with:
           github-token: ${{ secrets.GITHUB_TOKEN }}
           title-template: "MaxText Docker Image Build Failure"
           label-name: "docker-image-build-failure"
+
+  investigate_failure:
+    name: Investigate failed build
+    needs: [build_and_upload_maxtext_package, build-and-push, notify_failure]
+    if: ${{ always() && contains(needs.*.result, 'failure') && inputs.image_suffix == '' }}
+    uses: ./.github/workflows/gemini-investigate.yml
+    permissions:
+      contents: 'read'
+      id-token: 'write'
+      issues: 'write'
+      pull-requests: 'write'
+      actions: 'read'
+    with:
+      failed_run_id: '${{ github.run_id }}'
+    secrets: inherit