Test GitHub app

.bazelrc

@@ -390,8 +390,6 @@ common:rbe --spawn_strategy=remote,worker,standalone,local
 # Attempt to minimize the amount of data transfer between bazel and the remote
 # workers:
 common:rbe --remote_download_toplevel
-# Speed up repos creation by downloading ".tar" dists from the mirror.
-common:rbe --config=use_tar_archive_files
 test:rbe --test_env=USER=anon
 
 # RBE configs for Linux x86
@@ -415,8 +413,6 @@ common:rbe_linux_x86_64 --config=ci_linux_x86_64
 
 common:rbe_linux_x86_64_cuda_common --config=rbe_linux_x86_64_base
 common:rbe_linux_x86_64_cuda_common --repo_env=REMOTE_GPU_TESTING=1
-# Speed up CUDA repos creation by downloading ".tar" dists from the mirror.
-common:rbe_linux_x86_64_cuda_common --config=use_tar_archive_files
 # Enable forward compatibility for CUDA builds because RBE docker image doesn't
 # have latest CUDA drivers installed.
 common:rbe_linux_x86_64_cuda_common --@cuda_driver//:enable_forward_compatibility=true

.bazelversion

@@ -1 +1 @@
-7.4.1
+7.7.0

.github/workflows/asan.yaml

@@ -14,6 +14,11 @@ on:
     paths:
       - '**/workflows/asan.yaml'
 permissions: {}
+
+env:
+  UV_DEFAULT_INDEX: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple"
+  PIP_INDEX_URL: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple"
+
 jobs:
   asan:
     # Don't execute in fork due to runner type

.github/workflows/bazel_cuda_h100_b200.yml

@@ -55,8 +55,8 @@ jobs:
           done
     outputs:
       any_changed: ${{ steps.changed-files.outputs.any_changed || 'false' }}
-  run_tests: 
-    needs: changed_files 
+  run_tests:
+    needs: changed_files
     if: ${{ github.event.repository.fork == false && (github.event_name == 'schedule' || needs.changed_files.outputs.any_changed == 'true' || github.event_name == 'workflow_dispatch' || contains(github.event.pull_request.labels.*.name, 'CI Optional GPU Presubmit')) }}
     runs-on: linux-x86-a4-224-b200-1gpu
     container: 'us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest'
@@ -142,4 +142,6 @@ jobs:
             //tests:fused_attention_stablehlo_test_gpu \
             //tests:gpu_tests \
             //tests:python_callback_test_gpu \
-            //tests:ragged_collective_test_gpu
+            //tests:ragged_collective_test_gpu \
+            //tests/multiprocess:gpu_tests \
+            //jax/experimental/jax2tf/tests/multiprocess:gpu_tests

.github/workflows/bazel_test_tpu.yml

@@ -7,7 +7,7 @@
 # run-tests:
 #    - Downloads the jaxlib wheel from a GCS bucket if build_jaxlib is false.
 #    - Downloads the libtpu wheels.
-#    - Executes the `run_bazel_test_cpu.sh` script, which performs the following actions:
+#    - Executes the `run_bazel_test_tpu.sh` script, which performs the following actions:
 #      - Runs the TPU tests with Bazel.
 name: CI - Bazel test TPU
 on:
@@ -79,6 +79,10 @@ on:
         default: 'false'
         type: string
 permissions: {}
+
+env:
+  PIP_INDEX_URL: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple"
+
 jobs:
   run-bazel-tests:
     defaults:

.github/workflows/ci-build.yaml

@@ -15,6 +15,10 @@ concurrency:
   group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
   cancel-in-progress: true
 
+env:
+  UV_DEFAULT_INDEX: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple"
+  PIP_INDEX_URL: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple"
+
 jobs:
   lint_and_typecheck:
     runs-on: ubuntu-latest

.github/workflows/cloud-tpu-ci-nightly.yml

@@ -22,6 +22,10 @@ on:
 # document and enforce the permissions here.
 permissions:
   contents: read
+
+env:
+  UV_DEFAULT_INDEX: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple"
+
 jobs:
   cloud-tpu-test:
     strategy:

.github/workflows/culprit_finder.yml

@@ -0,0 +1,32 @@
+name: Culprit Finder Workflow
+
+on:
+  pull_request:
+    branches:
+      - main
+  workflow_dispatch:
+    inputs:
+      workflow-to-debug:
+        description: "Which workflow should be run?"
+        type: string
+        default: ""
+        required: true
+      json_vars:
+        type: string
+        description: 'A JSON string of key-value pairs that will be passed as environment variables to the workflow. (e.g., {"NAME": "test", "VERSION": "1.2.3"})'
+        required: false
+        default: '{"TEST": "my-cli-test"}'
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
+
+permissions: {}
+
+jobs:
+  bisect-hello-world:
+    if: ${{ inputs.workflow-to-debug == 'hello_world.yml' }}
+    uses: ./.github/workflows/hello_world.yml
+
+    name: Bisect Hello World
+    with:
+      json_vars: ${{ inputs.json_vars }}

.github/workflows/hello_world.yml

@@ -0,0 +1,53 @@
+name: Hello World
+on:
+  pull_request:
+    branches:
+      - main
+  workflow_dispatch: # allows triggering the workflow run manually
+  workflow_call:
+    inputs:
+      json_vars:
+        type: string
+        description: 'A JSON string of key-value pairs that will be passed as environment variables to the workflow. (e.g., {"NAME": "test", "VERSION": "1.2.3"})'
+        required: false
+        default: ''
+
+permissions: {}
+
+jobs:
+  build:
+
+    runs-on: linux-x86-n2-16
+    container: 'us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest'
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+      - name: Parse and set dynamic env vars
+        run: |
+          echo "Setting environment variables from input..."
+          echo '${{ inputs.json_vars }}' | jq -r 'to_entries[] | "\(.key)=\(.value)"' >> $GITHUB_ENV
+      - name: Simulate build
+        run: |
+          echo $GITHUB_ENV
+          echo $TEST
+          echo "Building the project..."
+          sleep 10
+          echo "Build complete...."
+          exit 1
+
+  test:
+    runs-on: linux-x86-n2-16
+    container: 'us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest'
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+      - name: Intermediate step
+        run: |
+          echo "Running tests."
+          sleep 15
+          echo "Tests complete........"
+      - name: Simulate testing
+        run: |
+          echo "Running tests."
+          sleep 15
+          echo "Tests complete."

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

@@ -12,6 +12,10 @@ concurrency:
   group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
   cancel-in-progress: true
 permissions: {}
+
+env:
+  UV_DEFAULT_INDEX: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple"
+
 jobs:
   build:
     runs-on: linux-x86-n4-16

.github/workflows/numpy_nightly.yml

@@ -24,6 +24,9 @@ concurrency:
   # Don't cancel in-progress jobs for main/release branches.
   cancel-in-progress: ${{ !contains(github.ref, 'release/') && github.ref != 'main' }}
 
+env:
+  UV_DEFAULT_INDEX: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple"
+
 jobs:
   test-nightly-numpy:
     defaults:

.github/workflows/oldest_supported_numpy.yml

@@ -16,6 +16,9 @@ concurrency:
   # Don't cancel in-progress jobs for main/release branches.
   cancel-in-progress: ${{ !contains(github.ref, 'release/') && github.ref != 'main' }}
 
+env:
+  UV_DEFAULT_INDEX: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple"
+
 jobs:
   test-oldest-supported-numpy:
     if: ${{ github.event.repository.fork == false && !startsWith(github.head_ref, 'release/') }}

.github/workflows/pytest_cpu.yml

@@ -36,9 +36,19 @@ on:
         type: string
       halt-for-connection:
         description: 'Should this workflow run wait for a remote connection?'
-        type: boolean
-        required: false
-        default: false
+        type: string
+        default: 'no'
+      max-processes:
+        description: 'The value to set for `--maxprocesses` when running pytest. Leave empty to use max processes.'
+        type: string
+        default: ''
+
+permissions: {}
+
+env:
+  UV_DEFAULT_INDEX: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple"
+  PIP_INDEX_URL: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple"
+
 jobs:
   run-tests:
     defaults:
@@ -58,6 +68,7 @@ jobs:
       JAXCI_HERMETIC_PYTHON_VERSION: "${{ inputs.python }}"
       JAXCI_PYTHON: "python${{ inputs.python }}"
       JAXCI_ENABLE_X64: "${{ inputs.enable-x64 }}"
+      MAX_PROCESSES: ${{ inputs.max-processes }}
 
     steps:
       - uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0

.github/workflows/pytest_cuda.yml

@@ -45,9 +45,13 @@ on:
         type: string
       halt-for-connection:
         description: 'Should this workflow run wait for a remote connection?'
-        type: boolean
-        required: false
-        default: false
+        type: string
+        default: 'no'
+permissions: {}
+
+env:
+  UV_DEFAULT_INDEX: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple"
+
 jobs:
   run-tests:
     defaults:

.github/workflows/pytest_tpu.yml

@@ -60,6 +60,10 @@ on:
         type: string
         default: 'no'
 permissions: {}
+
+env:
+  UV_DEFAULT_INDEX: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple"
+
 jobs:
   run-tests:
     defaults:

.github/workflows/rocm-ci.yml

@@ -46,7 +46,7 @@ jobs:
             dist_docker \
             --image-tag $TEST_IMAGE
       - name: Archive jax wheels
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@330a01c490aca151604b8cf639adc76d48f6c5d4 # v5.0.0
         with:
           name: rocm_jax_r${{ env.ROCM_VERSION }}_py${{ env.PYTHON_VERSION  }}_id${{ github.run_id }}
           path: ${{ env.WORKSPACE_DIR }}/dist/*.whl

.github/workflows/tsan.yaml

@@ -14,6 +14,11 @@ on:
       - '**/workflows/tsan.yaml'
       - '**/workflows/tsan-suppressions*.txt'
 permissions: {}
+
+env:
+  UV_DEFAULT_INDEX: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple"
+  PIP_INDEX_URL: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple"
+
 jobs:
   tsan:
     runs-on: linux-x86-n4-64

.github/workflows/upstream-nightly.yml

@@ -20,6 +20,11 @@ on:
     paths:
       - '**workflows/upstream-nightly.yml'
 permissions: {}
+
+env:
+  UV_DEFAULT_INDEX: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple"
+  PIP_INDEX_URL: "https://us-python.pkg.dev/ml-oss-artifacts-published/pypi-mirror/simple"
+
 jobs:
   upstream-dev:
     runs-on: linux-x86-n4-64