Add GitHub Action for XLA CPU CI tests to test new self-hosted runners

PiperOrigin-RevId: 720332744

.github/workflows/ci.yml

@@ -0,0 +1,51 @@
+# Copyright 2025 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+name: CI
+permissions:
+  contents: read
+on:
+  pull_request:
+  push:
+    branches:
+      - main
+
+
+jobs:
+  Tests:
+    strategy:
+      matrix:
+        job_info: [
+          {
+            pool: "linux-x86-n2-16",
+            container: "us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest",
+            kokoro_job_name: "xla-linux-x86-cpu",
+            pretty_name: "XLA Linux x86 CPU",
+          }
+        ]
+    name: ${{ matrix.job_info.pretty_name }}
+    runs-on: ${{ matrix.job_info.pool }}
+    container: ${{ matrix.job_info.container }}
+    defaults:
+      run:
+        shell: bash
+    timeout-minutes: 30
+    steps:
+      - name: "Checking out repository"
+        uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # v4.1.1
+      - name: "Run build.py"
+        env:
+          # TODO(ddunleavy): refactor build.py to not depend on this env var
+          KOKORO_JOB_NAME: ${{ matrix.job_info.kokoro_job_name }}
+        run: build_tools/ci/build.py

build_tools/ci/build.py

@@ -88,6 +88,7 @@ def _write_to_sponge_config(key, value) -> None:
 class BuildType(enum.Enum):
   """Enum representing all types of builds."""
   CPU_X86 = enum.auto()
+  CPU_X86_SELF_HOSTED = enum.auto()
   CPU_ARM64 = enum.auto()
   GPU = enum.auto()
   GPU_CONTINUOUS = enum.auto()
@@ -157,10 +158,12 @@ def docker_run_command(self, *, command: str, **kwargs: Any) -> List[str]:
   def commands(self) -> List[List[str]]:
     """Returns list of commands for a build."""
     cmds = []
-    cmds.append([
-        f"{_KOKORO_ARTIFACTS_DIR}/github/xla/.kokoro/generate_index_html.sh",
-        "index.html",
-    ])
+
+    if "self_hosted" not in self.type_.name.lower():
+      cmds.append([
+          f"{_KOKORO_ARTIFACTS_DIR}/github/xla/.kokoro/generate_index_html.sh",
+          "index.html",
+      ])
     if self.repo != "openxla/xla":
       _, repo_name = self.repo.split("/")
 
@@ -292,6 +295,16 @@ def nvidia_gpu_build_with_compute_capability(
     test_tag_filters=cpu_x86_tag_filter,
     options=_DEFAULT_BAZEL_OPTIONS,
 )
+_CPU_X86_SELF_HOSTED_BUILD = Build(
+    type_=BuildType.CPU_X86_SELF_HOSTED,
+    repo="openxla/xla",
+    image_url=None,
+    configs=("warnings", "nonccl", "rbe_linux_cpu"),
+    target_patterns=_XLA_DEFAULT_TARGET_PATTERNS,
+    build_tag_filters=cpu_x86_tag_filter,
+    test_tag_filters=cpu_x86_tag_filter,
+    options=_DEFAULT_BAZEL_OPTIONS,
+)
 
 cpu_arm_tag_filter = (
     "-no_oss",
@@ -463,6 +476,7 @@ def nvidia_gpu_build_with_compute_capability(
     "tensorflow/xla/jax/gpu/build_gpu": _JAX_GPU_BUILD,
     "tensorflow/xla/tensorflow/cpu/build_cpu": _TENSORFLOW_CPU_BUILD,
     "tensorflow/xla/tensorflow/gpu/build_gpu": _TENSORFLOW_GPU_BUILD,
+    "xla-linux-x86-cpu": _CPU_X86_SELF_HOSTED_BUILD,
 }
 
 
@@ -488,7 +502,8 @@ def main():
     return
 
   build = _KOKORO_JOB_NAME_TO_BUILD_MAP[kokoro_job_name]
-
+  logging.info("build.type_: %s", build.type_)
+  logging.info("build.commands(): %s", build.commands())
   for cmd in build.commands():
     sh(cmd)
 

build_tools/ci/golden_commands.txt

@@ -16,6 +16,11 @@ docker exec xla_ci bazel test --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvi
 docker exec xla_ci bazel analyze-profile profile.json.gz
 docker stop xla_ci
 # END BuildType.CPU_X86
+# BEGIN BuildType.CPU_X86_SELF_HOSTED
+parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd --config=warnings --config=nonccl --config=rbe_linux_cpu --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --nobuild -- //xla/... //build_tools/... @tsl//tsl/...
+bazel test --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd --config=warnings --config=nonccl --config=rbe_linux_cpu --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async -- //xla/... //build_tools/... @tsl//tsl/...
+bazel analyze-profile profile.json.gz
+# END BuildType.CPU_X86_SELF_HOSTED
 # BEGIN BuildType.GPU
 $KOKORO_ARTIFACTS_DIR/github/xla/.kokoro/generate_index_html.sh index.html
 nvidia-smi