Take advantage of larger job runners

.github/workflows/ci-build-checks.yaml

@@ -45,22 +45,26 @@ jobs:
           cache: pip
       - name: Install Bazel on CI
         run: ./scripts/ci_install.sh
-      - name: Configure CI TF
-        run: echo "Y\n" | ./configure.sh
-      - name: Build Wheel Test
+      - name: Build & test wheel for TFQ
         run: ./scripts/build_pip_package_test.sh
       - name: Test Wheel
         run: ./scripts/run_example.sh
 
   bazel-tests:
     name: Library tests
-    runs-on: ubuntu-22.04
+    runs-on: linux-x86-n2-32
+    container:
+      image: us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest
     timeout-minutes: 60
 
     steps:
+      - name: Adjust the ml-build runner to our needs
+        run: |
+          # lsb-release (needed by setup-python) is missing in ml-build runners
+          sudo apt-get update && sudo apt-get install -y lsb-release
+          # This file, though it's empty, hangs the Bazel install step.
+          rm -f /etc/bazel.bazelrc
       - uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # v6.0.0
-      - name: Remove some unnecessary software to free up disk space
-        run: sudo rm -rf /usr/share/dotnet /usr/local/lib/android &
       - uses: actions/setup-python@e797f83bcb11b83ae66e0230d6156d7c80228e7c # v5
         with:
           python-version: '3.10'
@@ -69,7 +73,7 @@ jobs:
       - name: Install Bazel on CI
         run: ./scripts/ci_install.sh
       - name: Configure CI TF
-        run: echo "Y\n" | ./configure.sh
+        run: printf "Y\n" | ./configure.sh
       - name: Full Library Test
         run: ./scripts/test_all.sh
 
@@ -91,17 +95,24 @@ jobs:
   #     - name: Install Bazel on CI
   #       run: ./scripts/ci_install.sh
   #     - name: Configure CI TF
-  #       run: echo "Y\n" | ./configure.sh
+  #       run: printf "Y\n" | ./configure.sh
   #     - name: Leak Test qsim and src
   #       run: ./scripts/msan_test.sh
 
   tutorials-test:
     name: Tutorial tests
-    runs-on: ubuntu-22.04
-    needs: wheel-build
+    runs-on: linux-x86-n2-32
+    container:
+      image: us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest
     timeout-minutes: 60
 
     steps:
+      - name: Adjust the ml-build runner to our needs
+        run: |
+          # lsb-release (needed by setup-python) is missing in ml-build runners
+          sudo apt-get update && sudo apt-get install -y lsb-release
+          # This file, though it's empty, hangs the Bazel install step.
+          rm -f /etc/bazel.bazelrc
       - uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # v6.0.0
       - uses: actions/setup-python@e797f83bcb11b83ae66e0230d6156d7c80228e7c # v5
         with:
@@ -112,9 +123,7 @@ jobs:
         run: pip install --upgrade pip seaborn==0.10.0
       - name: Install Bazel on CI
         run: ./scripts/ci_install.sh
-      - name: Configure CI TF
-        run: echo "Y\n" | ./configure.sh
-      - name: Build Wheel
+      - name: Build TFQ Wheel
         run: ./scripts/build_pip_package_test.sh
       - name: Test Notebooks
         run: ./scripts/ci_validate_tutorials.sh