Fix XLA tools build

iree-org · Jan 6, 2024 · 6bd82b5 · 6bd82b5
1 parent ca9244f
commit 6bd82b5
Show file tree

Hide file tree

Showing 20 changed files with 140 additions and 93 deletions.
diff --git a/.github/workflows/db_import.yml b/.github/workflows/db_import.yml
@@ -25,7 +25,7 @@ jobs:
       - cpu
       - os-family=Linux
     env:
-      DOCKER_IMAGE: gcr.io/iree-oss/openxla-benchmark/db_import@sha256:2c427286d467cdae7a206fdfe038bf32dbf9fbe68281866b571abeefe82907cd
+      DOCKER_IMAGE: gcr.io/iree-oss/openxla-benchmark/db_import@sha256:6a1b89844a2866b5d9aa8bef4aa56cf4788ba499bc8e23e76c3626819189db19
     steps:
       - name: "Checking out PR repository"
         uses: actions/checkout@e2f20e631ae6d7dd3b768f56a5d2af784dd54791  # v2.5.0

diff --git a/.github/workflows/run_comparative_benchmark.yml b/.github/workflows/run_comparative_benchmark.yml
@@ -11,6 +11,7 @@ name: Comparative Benchmarks
 on:
   # Will only run when manually triggered.
   workflow_dispatch:
+  pull_request:
 
 concurrency:
   # A PR number if a pull request and otherwise the commit hash. This cancels
@@ -78,7 +79,7 @@ jobs:
         run: |
           mkdir -p "${XLA_TOOLS_DIR}"
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
-            "gcr.io/iree-oss/openxla-benchmark/cuda11.8-cudnn8.9@sha256:02a54fbe596d6770a7e3bf7eac243c5bd71a3008cdc4da8cf8b8b79f95417710" \
+            "gcr.io/iree-oss/openxla-benchmark/cuda11.8-cudnn8.9@sha256:a0d6339be0828ab859a7d11282840e147884ad151c586c2983d63b698a1c27c9" \
             ./comparative_benchmark/xla_hlo/build_xla_tools.sh \
               "${XLA_TOOLS_DIR}" \
               "${CUDA_VERSION}"
@@ -133,7 +134,7 @@ jobs:
           RESULTS_PATH="${RESULTS_DIR}/${XLA_HLO_RESULTS_JSON}"
           docker run --gpus all --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
             --env "OOBI_XLA_TOOLS_DIR=${XLA_TOOLS_DIR}" \
-            "gcr.io/iree-oss/openxla-benchmark/cuda11.8-cudnn8.9@sha256:02a54fbe596d6770a7e3bf7eac243c5bd71a3008cdc4da8cf8b8b79f95417710" \
+            "gcr.io/iree-oss/openxla-benchmark/cuda11.8-cudnn8.9@sha256:a0d6339be0828ab859a7d11282840e147884ad151c586c2983d63b698a1c27c9" \
             ./comparative_benchmark/xla_hlo/benchmark_all.sh \
               "${TARGET_DEVICE}"\
               "${RESULTS_PATH}"
@@ -145,7 +146,7 @@ jobs:
         run: |
           RESULTS_PATH="${RESULTS_DIR}/${JAX_XLA_RESULTS_JSON}"
           docker run --gpus all --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
-            "gcr.io/iree-oss/openxla-benchmark/cuda11.8-cudnn8.9@sha256:02a54fbe596d6770a7e3bf7eac243c5bd71a3008cdc4da8cf8b8b79f95417710" \
+            "gcr.io/iree-oss/openxla-benchmark/cuda11.8-cudnn8.9@sha256:a0d6339be0828ab859a7d11282840e147884ad151c586c2983d63b698a1c27c9" \
             ./comparative_benchmark/jax/benchmark_xla.sh \
               "${TARGET_DEVICE}"\
               "${RESULTS_PATH}"
@@ -158,7 +159,7 @@ jobs:
 #        run: |
 #          RESULTS_PATH="${RESULTS_DIR}/${JAX_IREE_RESULTS_JSON}"
 #          docker run --gpus all --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
-#            "gcr.io/iree-oss/openxla-benchmark/cuda11.8-cudnn8.9@sha256:02a54fbe596d6770a7e3bf7eac243c5bd71a3008cdc4da8cf8b8b79f95417710" \
+#            "gcr.io/iree-oss/openxla-benchmark/cuda11.8-cudnn8.9@sha256:a0d6339be0828ab859a7d11282840e147884ad151c586c2983d63b698a1c27c9" \
 #            ./comparative_benchmark/jax/benchmark_iree.sh \
 #              "${TARGET_DEVICE}"\
 #              "${RESULTS_PATH}"
@@ -170,7 +171,7 @@ jobs:
         run: |
           RESULTS_PATH="${RESULTS_DIR}/${TF_XLA_RESULTS_JSON}"
           docker run --gpus all --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
-            "gcr.io/iree-oss/openxla-benchmark/cuda11.8-cudnn8.9@sha256:02a54fbe596d6770a7e3bf7eac243c5bd71a3008cdc4da8cf8b8b79f95417710" \
+            "gcr.io/iree-oss/openxla-benchmark/cuda11.8-cudnn8.9@sha256:a0d6339be0828ab859a7d11282840e147884ad151c586c2983d63b698a1c27c9" \
             ./comparative_benchmark/tf_xla/benchmark_all.sh \
               "${TARGET_DEVICE}"\
               "${RESULTS_PATH}"
@@ -182,24 +183,24 @@ jobs:
         run: |
           RESULTS_PATH="${RESULTS_DIR}/${PT_INDUCTOR_RESULTS_JSON}"
           docker run --gpus all --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
-            "gcr.io/iree-oss/openxla-benchmark/cuda11.8-cudnn8.9@sha256:02a54fbe596d6770a7e3bf7eac243c5bd71a3008cdc4da8cf8b8b79f95417710" \
+            "gcr.io/iree-oss/openxla-benchmark/cuda11.8-cudnn8.9@sha256:a0d6339be0828ab859a7d11282840e147884ad151c586c2983d63b698a1c27c9" \
             ./comparative_benchmark/pt_inductor/benchmark_all.sh \
               "${TARGET_DEVICE}"\
               "${RESULTS_PATH}"
           gcloud storage cp "${RESULTS_PATH}" "${RESULTS_GCS_DIR}/"
 
-  benchmark_on_c2-standard-16:
+  benchmark_on_c2-standard-60:
     needs: [setup, build_xla_tools]
     timeout-minutes: 1440
     runs-on:
       - self-hosted  # must come first
       - runner-group=${{ needs.setup.outputs.runner-group }}
       - environment=prod
-      - machine-type=c2-standard-16
+      - machine-type=c2-standard-60
     env:
       BENCHMARK_GCS_DIR: ${{ needs.setup.outputs.benchmark-gcs-dir }}
       RESULTS_DIR: results-dir
-      TARGET_DEVICE: c2-standard-16
+      TARGET_DEVICE: c2-standard-60
       XLA_TOOLS_DIR: ${{ needs.build_xla_tools.outputs.xla-tools-dir }}
       XLA_TOOLS_DIR_ARCHIVE: ${{ needs.build_xla_tools.outputs.xla-tools-dir-archive }}
       XLA_TOOLS_DIR_GCS_ARTIFACT: ${{ needs.build_xla_tools.outputs.xla-tools-dir-gcs-artifact }}
@@ -223,7 +224,7 @@ jobs:
           RESULTS_PATH="${RESULTS_DIR}/${XLA_HLO_RESULTS_JSON}"
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
             --env "OOBI_XLA_TOOLS_DIR=${XLA_TOOLS_DIR}" \
-            "gcr.io/iree-oss/openxla-benchmark/base@sha256:2dbee52eaa63e62137682f0eda701ac4cf59b8e16395daa757f6e1906b52dd82" \
+            "gcr.io/iree-oss/openxla-benchmark/base-python3.10@sha256:1c9751631c5fd3b532b11af307af4b9436cdc0a1246d79c065b8cc79945b8386" \
             ./comparative_benchmark/xla_hlo/benchmark_all.sh \
               "${TARGET_DEVICE}"\
               "${RESULTS_PATH}"
@@ -235,7 +236,7 @@ jobs:
         run: |
           RESULTS_PATH="${RESULTS_DIR}/${JAX_XLA_RESULTS_JSON}"
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
-            "gcr.io/iree-oss/openxla-benchmark/base@sha256:2dbee52eaa63e62137682f0eda701ac4cf59b8e16395daa757f6e1906b52dd82" \
+            "gcr.io/iree-oss/openxla-benchmark/base-python3.10@sha256:1c9751631c5fd3b532b11af307af4b9436cdc0a1246d79c065b8cc79945b8386" \
             ./comparative_benchmark/jax/benchmark_xla.sh \
               "${TARGET_DEVICE}"\
               "${RESULTS_PATH}"
@@ -247,7 +248,7 @@ jobs:
         run: |
           RESULTS_PATH="${RESULTS_DIR}/${JAX_IREE_RESULTS_JSON}"
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
-            "gcr.io/iree-oss/openxla-benchmark/base@sha256:2dbee52eaa63e62137682f0eda701ac4cf59b8e16395daa757f6e1906b52dd82" \
+            "gcr.io/iree-oss/openxla-benchmark/base-python3.10@sha256:1c9751631c5fd3b532b11af307af4b9436cdc0a1246d79c065b8cc79945b8386" \
             ./comparative_benchmark/jax/benchmark_iree.sh \
               "${TARGET_DEVICE}"\
               "${RESULTS_PATH}"
@@ -259,7 +260,7 @@ jobs:
         run: |
           RESULTS_PATH="${RESULTS_DIR}/${TF_XLA_RESULTS_JSON}"
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
-            "gcr.io/iree-oss/openxla-benchmark/base@sha256:2dbee52eaa63e62137682f0eda701ac4cf59b8e16395daa757f6e1906b52dd82" \
+            "gcr.io/iree-oss/openxla-benchmark/base-python3.10@sha256:1c9751631c5fd3b532b11af307af4b9436cdc0a1246d79c065b8cc79945b8386" \
             ./comparative_benchmark/tf_xla/benchmark_all.sh \
               "${TARGET_DEVICE}"\
               "${RESULTS_PATH}"
@@ -271,7 +272,7 @@ jobs:
         run: |
           RESULTS_PATH="${RESULTS_DIR}/${PT_INDUCTOR_RESULTS_JSON}"
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
-            "gcr.io/iree-oss/openxla-benchmark/base@sha256:2dbee52eaa63e62137682f0eda701ac4cf59b8e16395daa757f6e1906b52dd82" \
+            "gcr.io/iree-oss/openxla-benchmark/base-python3.10@sha256:1c9751631c5fd3b532b11af307af4b9436cdc0a1246d79c065b8cc79945b8386" \
             ./comparative_benchmark/pt_inductor/benchmark_all.sh \
               "${TARGET_DEVICE}"\
               "${RESULTS_PATH}"

diff --git a/.github/workflows/run_convperf.yml b/.github/workflows/run_convperf.yml
@@ -11,6 +11,7 @@ name: convperf
 on:
   # Will only run when manually triggered.
   workflow_dispatch:
+  pull_request:
 
 concurrency:
   # A PR number if a pull request and otherwise the commit hash. This cancels
@@ -71,7 +72,7 @@ jobs:
           mkdir ${RESULTS_DIR}
           # Run under privileged mode because system-level scheduling is used.
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" --privileged \
-            gcr.io/iree-oss/openxla-benchmark/convperf@sha256:fda453f1414085558a5c4a2250db0010841c73a082379997f507a58b19cb8aaf \
+            gcr.io/iree-oss/openxla-benchmark/convperf@sha256:ee13c0bf4a03d7671b8764bd1631c5afd0bba57eda789ce4923eb18d5790c84c \
           ./experimental/microbenchmarks/convperf/build_and_run_convperf.sh "${BUILD_DIR}" "${RESULTS_DIR}"
       - name: "Uploading results"
         run: |

diff --git a/.github/workflows/run_ggml_benchmark.yml b/.github/workflows/run_ggml_benchmark.yml
@@ -10,6 +10,7 @@ name: GGML Benchmarks
 
 on:
   workflow_dispatch:
+  pull_request:
 
 concurrency:
   # A PR number if a pull request and otherwise the commit hash. This cancels
@@ -54,17 +55,17 @@ jobs:
           fi
           echo "benchmark-gcs-dir=${BENCHMARK_GCS_DIR}" >> "${GITHUB_OUTPUT}"
 
-  benchmark_on_c2-standard-16:
+  benchmark_on_c2-standard-60:
     needs: [setup]
     runs-on:
       - self-hosted  # must come first
       - runner-group=${{ needs.setup.outputs.runner-group }}
       - environment=prod
-      - machine-type=c2-standard-16
+      - machine-type=c2-standard-60
     env:
       BENCHMARK_GCS_DIR: ${{ needs.setup.outputs.benchmark-gcs-dir }}
       RESULTS_DIR: results-dir
-      TARGET_DEVICE: c2-standard-16
+      TARGET_DEVICE: c2-standard-60
       GGML_BUILD_DIR: build-dir
     steps:
       - name: "Checking out PR repository"
@@ -77,7 +78,7 @@ jobs:
       - name: "Building GGML CPU"
         run: |
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
-            "gcr.io/iree-oss/openxla-benchmark/base@sha256:2dbee52eaa63e62137682f0eda701ac4cf59b8e16395daa757f6e1906b52dd82" \
+            "gcr.io/iree-oss/openxla-benchmark/base-python3.11@sha256:03f729350ad2fda2cc572c202b95562064031194ceb6bb146695ae1f0c4f180f" \
             ./experimental/ggml/build_ggml.sh \
               "${TARGET_DEVICE}" \
               "${GGML_BUILD_DIR}"
@@ -88,7 +89,7 @@ jobs:
         run: |
           RESULTS_PATH="${RESULTS_DIR}/${GGML_RESULTS_JSON}"
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
-            "gcr.io/iree-oss/openxla-benchmark/base@sha256:2dbee52eaa63e62137682f0eda701ac4cf59b8e16395daa757f6e1906b52dd82" \
+            "gcr.io/iree-oss/openxla-benchmark/base-python3.11@sha256:03f729350ad2fda2cc572c202b95562064031194ceb6bb146695ae1f0c4f180f" \
             ./experimental/ggml/benchmark_ggml.sh \
               "${TARGET_DEVICE}" \
               "${GGML_BUILD_DIR}" \
@@ -118,7 +119,7 @@ jobs:
         run: |
           mkdir -p "${GGML_BUILD_DIR}"
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
-            "gcr.io/iree-oss/openxla-benchmark/android@sha256:4c219b20213c7f26edf848b6924a95991c0d359463709d4ee55b499eeaa3349b" \
+            "gcr.io/iree-oss/openxla-benchmark/android@sha256:5f4f50580fd055c0c3537d50228b447a207a21a43c352422970d715298d8279b" \
             ./experimental/ggml/build_ggml.sh \
               "${TARGET_DEVICE}" \
               "${GGML_BUILD_DIR}"

diff --git a/.github/workflows/run_mmperf.yml b/.github/workflows/run_mmperf.yml
@@ -70,7 +70,7 @@ jobs:
       - name: "Building mmperf for CPU"
         run: |
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
-            gcr.io/iree-oss/openxla-benchmark/mmperf@sha256:c972ce5b2144de0786f103611fecbd88d93dd45ecd068f8c97d98c08677cee57 \
+            gcr.io/iree-oss/openxla-benchmark/mmperf@sha256:b6d95efb9429ae05a4400b3ce003ea4600595f8bb995fd202f80d122b749f3b1 \
           ./experimental/microbenchmarks/mmperf/build_mmperf.sh "cpu" "${BUILD_DIR}"
       - name: "Removing unused files"
         run: |
@@ -114,7 +114,7 @@ jobs:
       - name: "Building mmperf for CUDA"
         run: |
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
-            gcr.io/iree-oss/openxla-benchmark/mmperf@sha256:c972ce5b2144de0786f103611fecbd88d93dd45ecd068f8c97d98c08677cee57 \
+            gcr.io/iree-oss/openxla-benchmark/mmperf@sha256:b6d95efb9429ae05a4400b3ce003ea4600595f8bb995fd202f80d122b749f3b1 \
           ./experimental/microbenchmarks/mmperf/build_mmperf.sh "cuda" "${BUILD_DIR}"
       - name: "Removing unused files"
         run: |
@@ -138,13 +138,13 @@ jobs:
           gcloud storage cp "${BUILD_DIR_ARCHIVE}" "${BUILD_DIR_GCS_ARTIFACT}"
           echo "build-dir-gcs-artifact=${BUILD_DIR_GCS_ARTIFACT}" >> "${GITHUB_OUTPUT}"
 
-  benchmark_on_c2-standard-16:
+  benchmark_on_c2-standard-60:
     needs: [setup, build_cpu]
     runs-on:
       - self-hosted  # must come first
       - runner-group=${{ needs.setup.outputs.runner-group }}
       - environment=prod
-      - machine-type=c2-standard-16
+      - machine-type=c2-standard-60
     env:
       RESULTS_DIR: mmperf-results-cpu
       GCS_UPLOAD_PARENT_DIR: "gs://mmperf-benchmark-artifacts/cpu"
@@ -163,7 +163,7 @@ jobs:
         run: |
           mkdir ${RESULTS_DIR}
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
-            gcr.io/iree-oss/openxla-benchmark/mmperf@sha256:c972ce5b2144de0786f103611fecbd88d93dd45ecd068f8c97d98c08677cee57 \
+            gcr.io/iree-oss/openxla-benchmark/mmperf@sha256:b6d95efb9429ae05a4400b3ce003ea4600595f8bb995fd202f80d122b749f3b1 \
           ./experimental/microbenchmarks/mmperf/run_mmperf.sh "${BUILD_DIR}" "${RESULTS_DIR}"
       - name: "Uploading results"
         run: |
@@ -195,7 +195,7 @@ jobs:
         run: |
           mkdir ${RESULTS_DIR}
           docker run --gpus all --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
-            gcr.io/iree-oss/openxla-benchmark/mmperf@sha256:c972ce5b2144de0786f103611fecbd88d93dd45ecd068f8c97d98c08677cee57 \
+            gcr.io/iree-oss/openxla-benchmark/mmperf@sha256:b6d95efb9429ae05a4400b3ce003ea4600595f8bb995fd202f80d122b749f3b1 \
           ./experimental/microbenchmarks/mmperf/run_mmperf.sh "${BUILD_DIR}" "${RESULTS_DIR}"
       - name: "Uploading results"
         run: |

diff --git a/.github/workflows/run_mobile_comparative_benchmark.yml b/.github/workflows/run_mobile_comparative_benchmark.yml
@@ -13,6 +13,7 @@ on:
   schedule:
     # Scheduled to run at 09:00 UTC and 21:00 UTC.
     - cron: '0 09,21 * * *'
+  pull_request:
 
 concurrency:
   # A PR number if a pull request and otherwise the commit hash. This cancels
@@ -98,7 +99,7 @@ jobs:
               --env IREE_SOURCE_DIR="${IREE_SOURCE_DIR}" \
               --env IREE_INSTALL_DIR="${IREE_INSTALL_DIR}" \
               --env ANDROID_PLATFORM_VERSION="${ANDROID_PLATFORM_VERSION}" \
-              "gcr.io/iree-oss/openxla-benchmark/android@sha256:4c219b20213c7f26edf848b6924a95991c0d359463709d4ee55b499eeaa3349b" \
+              "gcr.io/iree-oss/openxla-benchmark/android@sha256:5f4f50580fd055c0c3537d50228b447a207a21a43c352422970d715298d8279b" \
                 ./experimental/iree/build_iree.sh \
                     "${X86_BUILD_DIR}" \
                     "${ANDROID_BUILD_DIR}"
@@ -159,7 +160,7 @@ jobs:
 
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
             --env OOBI_TEMP_DIR="${TF_SOURCE_DIR}" \
-            "gcr.io/iree-oss/openxla-benchmark/tflite-android@sha256:e02858dd5998dbae6798e11c8646bcff59135ffeb35e0f08a52977106979f1d2" \
+            "gcr.io/iree-oss/openxla-benchmark/tflite-android@sha256:818b99bc21ae85e18bd5a8c78c1dc7c4ec7c7bca9cfaf3b6164e7b0b0528202f" \
               ./experimental/tflite/build_tflite.sh \
                   "${X86_BUILD_DIR}" \
                   "${ANDROID_BUILD_DIR}"
@@ -226,7 +227,7 @@ jobs:
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
             --env IREE_COMPILE_PATH="${IREE_COMPILE_PATH}" \
             --env OOBI_TEMP_DIR="${IREE_TEMP_DIR}" \
-            "gcr.io/iree-oss/openxla-benchmark/base@sha256:2dbee52eaa63e62137682f0eda701ac4cf59b8e16395daa757f6e1906b52dd82" \
+            "gcr.io/iree-oss/openxla-benchmark/base-python3.11@sha256:03f729350ad2fda2cc572c202b95562064031194ceb6bb146695ae1f0c4f180f" \
               ./experimental/iree/compile_workloads.sh \
                   "${TARGET_DEVICE}" \
                   "${BENCHMARK_ARTIFACTS_DIR}"
@@ -239,7 +240,7 @@ jobs:
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
             --env OOBI_IREE_BENCHMARK_MODULE_PATH="${IREE_INSTALL_DIR}/bin/iree-benchmark-module" \
             --env OOBI_IREE_RUN_MODULE_PATH="${IREE_INSTALL_DIR}/bin/iree-run-module" \
-            "gcr.io/iree-oss/openxla-benchmark/base@sha256:2dbee52eaa63e62137682f0eda701ac4cf59b8e16395daa757f6e1906b52dd82" \
+            "gcr.io/iree-oss/openxla-benchmark/base-python3.11@sha256:03f729350ad2fda2cc572c202b95562064031194ceb6bb146695ae1f0c4f180f" \
             ./experimental/iree/benchmark_iree_x86.sh \
               "${TARGET_DEVICE}" \
               "${BENCHMARK_ARTIFACTS_DIR}" \
@@ -265,7 +266,7 @@ jobs:
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
             --env OOBI_TEMP_DIR="${TFLITE_TEMP_DIR}" \
             --env TFLITE_BENCHMARK_BINARY="${TFLITE_BUILD_DIR}/benchmark_model" \
-            "gcr.io/iree-oss/openxla-benchmark/base@sha256:2dbee52eaa63e62137682f0eda701ac4cf59b8e16395daa757f6e1906b52dd82" \
+            "gcr.io/iree-oss/openxla-benchmark/base-python3.11@sha256:03f729350ad2fda2cc572c202b95562064031194ceb6bb146695ae1f0c4f180f" \
             ./experimental/tflite/benchmark_tflite_x86.sh \
               "${TARGET_DEVICE}" \
               "${TFLITE_RESULTS_FILE}"
@@ -288,7 +289,7 @@ jobs:
           gcloud storage cp "${TFLITE_RESULTS_GCS_PATH}" "${TFLITE_RESULTS_FILE}"
 
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
-            "gcr.io/iree-oss/openxla-benchmark/base@sha256:2dbee52eaa63e62137682f0eda701ac4cf59b8e16395daa757f6e1906b52dd82" \
+            "gcr.io/iree-oss/openxla-benchmark/base-python3.11@sha256:03f729350ad2fda2cc572c202b95562064031194ceb6bb146695ae1f0c4f180f" \
             ./experimental/utils/generate_iree_vs_tflite_report.sh \
               "${IREE_RESULTS_FILE}" \
               "${TFLITE_RESULTS_FILE}" \
@@ -343,7 +344,7 @@ jobs:
           docker run --mount="type=bind,src="${PWD}",target=/work" --workdir="/work" \
             --env IREE_COMPILE_PATH="${IREE_COMPILE_PATH}" \
             --env OOBI_TEMP_DIR="${OOBI_TEMP_DIR}" \
-            "gcr.io/iree-oss/openxla-benchmark/android@sha256:4c219b20213c7f26edf848b6924a95991c0d359463709d4ee55b499eeaa3349b" \
+            "gcr.io/iree-oss/openxla-benchmark/android@sha256:5f4f50580fd055c0c3537d50228b447a207a21a43c352422970d715298d8279b" \
               ./experimental/iree/compile_workloads.sh \
                   "${TARGET_DEVICE}" \
                   "${BENCHMARK_ARTIFACTS_DIR}"