Enable CI test for testing

.github/workflows/bazel_cpu_rbe.yml

@@ -26,7 +26,6 @@ concurrency:
 
 jobs:
   run_tests:
-    if: github.event.repository.fork == false
     runs-on: ${{ matrix.runner }}
     container: ${{ (contains(matrix.runner, 'linux-x86') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest') ||
                    (contains(matrix.runner, 'linux-arm64') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build-arm64:latest') }}
@@ -37,7 +36,7 @@ jobs:
     strategy:
       matrix:
         python: ["3.10", "3.13"]
-        runner: ["linux-x86-n2-16", "linux-arm64-c4a-16"]
+        runner: ["linux-x86-ct5lp-224-8tpu"]
         enable-x_64: [1, 0]
         exclude:
             # Exclude x64=1 on the oldest Python and x64=0 on the newest Python. As long as we have
@@ -59,5 +58,7 @@ jobs:
           halt-dispatch-input: ${{ inputs.halt-for-connection }}
       # Since we do not have a Linux Arm64 RBE pool, we do not run the tests on Arm64. Instead, we
       # cross-compile the tests on the Linux x86 RBE pool.
+      - name: sleep
+        run: sleep 10000
       - name: ${{ (contains(matrix.runner, 'linux-arm64') && 'Build' || 'Run') }} Bazel CPU Tests with RBE
-        run: ./ci/run_bazel_test_cpu_rbe.sh
+        run: ./ci/run_bazel_test_cpu_rbe.sh

.github/workflows/bazel_cuda_non_rbe.yml

@@ -18,7 +18,7 @@ on:
         description: "Which runner should the workflow run on?"
         type: string
         required: true
-        default: "linux-x86-n2-16"
+        default: "linux-x86-n2-32"
       python:
         description: "Which python version to test?"
         type: string
@@ -47,7 +47,7 @@ jobs:
         # Explicitly set the shell to bash
         shell: bash
     runs-on: ${{ inputs.runner }}
-    container: "us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build-cuda12.8-cudnn9.8:latest"
+    container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-cuda12.1-cudnn9.8:latest"
 
     env:
       JAXCI_HERMETIC_PYTHON_VERSION: ${{ inputs.python }}

.github/workflows/bazel_cuda_rbe.yml

@@ -27,7 +27,7 @@ concurrency:
 jobs:
   run_tests:
     runs-on: ${{ matrix.runner }}
-    container: 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest'
+    container: 'us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest'
     env:
       JAXCI_HERMETIC_PYTHON_VERSION: ${{ matrix.python }}
       JAXCI_ENABLE_X64: ${{ matrix.enable-x_64 }}
@@ -53,4 +53,4 @@ jobs:
         with:
           halt-dispatch-input: ${{ inputs.halt-for-connection }}
       - name: Run Bazel CUDA Tests with RBE
-        run: ./ci/run_bazel_test_cuda_rbe.sh
+        run: for i in $(seq 1 180); do echo "This is statement number $i at $(date)"; sleep 60; done

.github/workflows/ci-build.yaml

@@ -43,7 +43,7 @@ jobs:
 
   build:
     # Don't execute in fork due to runner type
-    if: ${{ github.repository == 'jax-ml/jax' }}
+    if: ${{ github.repository == 'jax-ml/jax' || github.repository == 'google-ml-infra/jax-fork' }}
     name: "build ${{ matrix.name-prefix }} (py ${{ matrix.python-version }} on ubuntu-20.04, x64=${{ matrix.enable-x64}})"
     runs-on: linux-x86-n2-32
     container:

.github/workflows/docker.yaml

@@ -0,0 +1,23 @@
+name: Docker Test Build
+
+on:
+  pull_request:
+    branches:
+      - main
+  push:
+    branches:
+      - main
+      - 'release/**'
+
+jobs:
+  run_tests:
+    runs-on: "linux-x86-n2-32"
+    container: "google/cloud-sdk:524.0.0"
+    name: "Test build container"
+# End Presubmit Naming Check github-cpu-presubmits
+    steps:
+      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
+      - name: Wait For Connection
+        uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c
+      - name: Sleep
+        run: sleep 1000000

.github/workflows/jax-array-api.yml

@@ -13,36 +13,21 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
-  build:
-
-    runs-on: ubuntu-latest
-    strategy:
-      matrix:
-        python-version: [3.11]
+  run:
+    # Containers must run in Linux based operating systems
+    runs-on: linux-x86-n2-32
+    # Docker Hub image that `container-job` executes in
+    container: us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/maxtext_jax_stable:latest
 
     steps:
-    - name: Checkout jax
-      uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
-    - name: Checkout array-api-tests
-      uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
-      with:
-        repository: data-apis/array-api-tests
-        # TODO(jakevdp) update this to a stable release/tag when available.
-        ref: 'c48410f96fc58e02eea844e6b7f6cc01680f77ce'  # Latest commit as of 2025-04-02
-        submodules: 'true'
-        path: 'array-api-tests'
-    - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065  # v5.6.0
-      with:
-        python-version: ${{ matrix.python-version }}
-    - name: Install dependencies
-      run: |
-        pip install uv~=0.5.30
-        uv pip install --system .[ci] pytest-xdist -r array-api-tests/requirements.txt
-    - name: Run the test suite
-      env:
-        ARRAY_API_TESTS_MODULE: jax.numpy
-        JAX_ENABLE_X64: 'true'
-      run: |
-        cd ${GITHUB_WORKSPACE}/array-api-tests
-        pytest -n auto array_api_tests --derandomize --disable-deadline --skips-file ${GITHUB_WORKSPACE}/tests/array_api_skips.txt
+      # Downloads a copy of the code in your repository before running CI tests
+      - uses: actions/checkout@v5
+
+      - name: Connect to Redis
+        # Runs a script that creates a Redis client, populates
+        # the client with data, and retrieves data
+        run: |
+          /usr/pathways/run/cloud_pathways_server_sanitized --server_port=29001 --node_type=resource_manager --instance_count=1 --instance_type=tpuv4:2x2x1, --gcs_scratch_location=gs://cloud-pathways-staging/tmp
+
+
+

test_docker/Dockerfile

@@ -0,0 +1,2 @@
+FROM scratch
+CMD ["/bin/echo", "Hello, World!"]