Use self hosted runners for CI

.github/workflows/buildAndTestBazel.yml

@@ -23,29 +23,43 @@ on:
     # Run once a day
     - cron:  '0 12 * * *'
   workflow_dispatch:
+    inputs:
+      halt-for-connection:
+        description: 'Should this workflow run wait for a remote connection?'
+        type: choice
+        required: true
+        default: 'no'
+        options:
+          - 'yes'
+          - 'no'
 
 # Ensure that only a single job or workflow using the same
 # concurrency group will run at a time. This would cancel
 # any in-progress jobs in the same github workflow and github
 # ref (e.g. refs/heads/main or refs/pull/<pr_number>/merge).
 concurrency:
-  group: ${{ github.workflow }}-${{ github.ref }}
-  cancel-in-progress: true
+  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
+  cancel-in-progress: ${{ github.ref != 'main' }}
 
 jobs:
   bazel-build:
     strategy:
       fail-fast: false
-    runs-on: ${{ github.repository == 'openxla/stablehlo' && 'ubuntu-22.04-64core' ||  'ubuntu-22.04'  }}
-    # TODO: Use ML OSS Artifacts docker to manage deps needs to package with cmake tooling
-    # container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest"
+    runs-on: ${{ github.repository == 'openxla/stablehlo' && 'linux-x86-n2-16' ||  'ubuntu-22.04'  }}
+    container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest"
 
     steps:
     # Checkout StableHLO with all commits in the chain
     - name: Checkout StableHLO
-      uses: actions/checkout@v4
+      uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
       with:
         fetch-depth: $(( ${{ github.event_name == 'pull_request' && github.event.pull_request.commits || 0 }} + 1 ))
+        set-safe-directory: true
+
+    - name: "Wait For Connection"
+      uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c
+      with:
+        halt-dispatch-input: ${{ inputs.halt-for-connection }}
 
     - name: Validate LLVM Commit
       run: |
@@ -55,28 +69,22 @@ jobs:
       id: llvm-version
       shell: bash
       run: |
-        echo "version=$(cat ${{ github.workspace }}/build_tools/llvm_version.txt)" >> $GITHUB_OUTPUT
-
-    - name: Setup Bazelisk
-      uses: bazelbuild/setup-bazelisk@v3
-
-    - name: Mount Bazel Cache
-      uses: actions/cache@v4
-      with:
-        path: "~/.cache/bazel"
-        key: ${{ runner.os }}-stablehlo_bazelbuild-${{ hashFiles('.bazelversion', '.bazelrc', 'WORKSPACE', 'WORKSPACE.bazel', 'MODULE.bazel') }}
-        restore-keys: |
-          ${{ runner.os }}-stablehlo_bazelbuild-
+        echo "version=$(cat ./build_tools/llvm_version.txt)" >> $GITHUB_OUTPUT
 
     - name: Build and check docs
       shell: bash
       run: |
+        git config --global --add safe.directory "$(pwd)"
+        echo "$PWD"
+        git status
         ./build_tools/github_actions/ci_build_docs.sh -c
 
     - name: Build and Test StableHLO (Diff Only)
       if: ${{ github.event_name == 'pull_request' }}
       shell: bash
       run: |
+        echo "$PWD"
+        git status
         START_HASH=${{ github.event.pull_request.base.sha }}
         END_HASH=${{ github.event.pull_request.head.sha }}
         BAZEL_DIFF=/tmp/bazel-diff.jar
@@ -89,3 +97,9 @@ jobs:
       run: |
         ./build_tools/github_actions/ci_build_bazel.sh
 
+    - name: "Wait For Connection - On Error"
+      if: failure()
+      uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c
+      with:
+        halt-dispatch-input: ${{ inputs.halt-for-connection }}
+

WORKSPACE.bazel

@@ -17,38 +17,32 @@ workspace(name = "stablehlo")
 
 load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 
+# Initialize toolchains for ML projects.
+#
+# A hermetic build system is designed to produce completely reproducible builds for C++.
+# Details: https://github.com/google-ml-infra/rules_ml_toolchain
 http_archive(
-    name = "bazel_skylib",
-    sha256 = "74d544d96f4a5bb630d465ca8bbcfe231e3594e5aae57e1edbf17a6eb3ca2506",
-    urls = [
-        "https://storage.googleapis.com/mirror.tensorflow.org/github.com/bazelbuild/bazel-skylib/releases/download/1.3.0/bazel-skylib-1.3.0.tar.gz",
-        "https://github.com/bazelbuild/bazel-skylib/releases/download/1.3.0/bazel-skylib-1.3.0.tar.gz",
-    ],
+    name = "rules_ml_toolchain",
+    sha256 = "9bd46bc5e06a56a9335897be630d4b820c678281aa88302518c250921338ad22",
+    strip_prefix = "rules_ml_toolchain-40efd07eb8e6565e506562f36d7dc43cd83e5b32",
+    url = "https://github.com/google-ml-infra/rules_ml_toolchain/archive/40efd07eb8e6565e506562f36d7dc43cd83e5b32.tar.gz",
 )
 
-# This is needed since https://reviews.llvm.org/D143344.
-# Not sure if it's a bug or a feature, but it doesn't hurt to keep an additional
-# dependency here.
-http_archive(
-    name = "llvm_zstd",
-    build_file = "@llvm-raw//utils/bazel/third_party_build:zstd.BUILD",
-    sha256 = "7c42d56fac126929a6a85dbc73ff1db2411d04f104fae9bdea51305663a83fd0",
-    strip_prefix = "zstd-1.5.2",
-    urls = [
-        "https://github.com/facebook/zstd/releases/download/v1.5.2/zstd-1.5.2.tar.gz",
-    ],
+load(
+    "@rules_ml_toolchain//cc/deps:cc_toolchain_deps.bzl",
+    "cc_toolchain_deps",
 )
 
-# This is needed since https://reviews.llvm.org/D143320
-# Not sure if it's a bug or a feature, but it doesn't hurt to keep an additional
-# dependency here.
+cc_toolchain_deps()
+
+register_toolchains("@rules_ml_toolchain//cc:linux_x86_64_linux_x86_64")
+
 http_archive(
-    name = "llvm_zlib",
-    build_file = "@llvm-raw//utils/bazel/third_party_build:zlib-ng.BUILD",
-    sha256 = "e36bb346c00472a1f9ff2a0a4643e590a254be6379da7cddd9daeb9a7f296731",
-    strip_prefix = "zlib-ng-2.0.7",
+    name = "bazel_skylib",
+    sha256 = "74d544d96f4a5bb630d465ca8bbcfe231e3594e5aae57e1edbf17a6eb3ca2506",
     urls = [
-        "https://github.com/zlib-ng/zlib-ng/archive/refs/tags/2.0.7.zip",
+        "https://storage.googleapis.com/mirror.tensorflow.org/github.com/bazelbuild/bazel-skylib/releases/download/1.3.0/bazel-skylib-1.3.0.tar.gz",
+        "https://github.com/bazelbuild/bazel-skylib/releases/download/1.3.0/bazel-skylib-1.3.0.tar.gz",
     ],
 )