Merge pull request #240 from ROCm/ci-upstream-sync-127_1

CI: 02/25/25 upstream sync

Author: charleshofer

.bazelrc

@@ -2,7 +2,7 @@
 # All default build options below. These apply to all build commands.
 # #############################################################################
 # Make Bazel print out all options from rc files.
-build --announce_rc
+common --announce_rc
 
 # By default, execute all actions locally.
 build --spawn_strategy=local
@@ -11,7 +11,7 @@ build --spawn_strategy=local
 # automatically when building on Linux.
 build --enable_platform_specific_config
 
-build --experimental_cc_shared_library
+common --experimental_cc_shared_library
 
 # Do not use C-Ares when building gRPC.
 build --define=grpc_no_ares=true

.github/workflows/asan.yaml

@@ -62,7 +62,8 @@ jobs:
         run: |
           source ${GITHUB_WORKSPACE}/venv/bin/activate
           cd jax
-          pip install -r build/test-requirements.txt
+          pip install uv~=0.5.30
+          uv pip install -r build/test-requirements.txt
       - name: Build and install JAX
         env:
           ASAN_OPTIONS: detect_leaks=0
@@ -73,8 +74,8 @@ jobs:
             --bazel_options=--color=yes \
             --bazel_options=--copt=-fsanitize=address \
             --clang_path=/usr/bin/clang-18
-          pip install dist/jaxlib-*.whl
-          pip install -e .
+          uv pip install dist/jaxlib-*.whl \
+            -e .
       - name: Run tests
         env:
           ASAN_OPTIONS: detect_leaks=0

.github/workflows/bazel_cuda_non_rbe.yml

@@ -42,8 +42,12 @@ on:
 
 jobs:
   run-tests:
+    defaults:
+      run:
+        # Explicitly set the shell to bash
+        shell: bash
     runs-on: ${{ inputs.runner }}
-    container: "us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/nosla-cuda12.3-cudnn9.1-ubuntu20.04-manylinux2014-multipython:latest"
+    container: "us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build-cuda12.3-cudnn9.1:latest"
 
     env:
       JAXCI_HERMETIC_PYTHON_VERSION: ${{ inputs.python }}
@@ -68,11 +72,23 @@ jobs:
           echo "ARCH=${arch}" >> $GITHUB_ENV
           echo "PYTHON_MAJOR_MINOR=${python_major_minor}" >> $GITHUB_ENV
       - name: Download the wheel artifacts from GCS
+        id: download-wheel-artifacts
+        # Set continue-on-error to true to prevent actions from failing the workflow if this step
+        # fails. Instead, we verify the outcome in the next step so that we can print a more
+        # informative error message.
+        continue-on-error: true
         run: >-
           mkdir -p $(pwd)/dist &&
           gsutil -m cp -r "${{ inputs.gcs_download_uri }}/jaxlib*${PYTHON_MAJOR_MINOR}*${OS}*${ARCH}*.whl" $(pwd)/dist/ &&
           gsutil -m cp -r "${{ inputs.gcs_download_uri }}/jax*cuda*plugin*${PYTHON_MAJOR_MINOR}*${OS}*${ARCH}*.whl" $(pwd)/dist/ &&
           gsutil -m cp -r "${{ inputs.gcs_download_uri }}/jax*cuda*pjrt*${OS}*${ARCH}*.whl" $(pwd)/dist/
+      - name: Skip the test run if the wheel artifacts were not downloaded successfully
+        if: steps.download-wheel-artifacts.outcome == 'failure'
+        run: |
+          echo "Failed to download wheel artifacts from GCS. Please check if the wheels were"
+          echo "built successfully by the artifact build jobs and are available in the GCS bucket."
+          echo "Skipping the test run."
+          exit 1
       # Halt for testing
       - name: Wait For Connection
         uses: google-ml-infra/actions/ci_connection@main

.github/workflows/build_artifacts.yml

@@ -131,7 +131,7 @@ jobs:
         with:
           halt-dispatch-input: ${{ inputs.halt-for-connection }}
       - name: Build ${{ inputs.artifact }}
-        timeout-minutes: 30
+        timeout-minutes: 60
         run: ./ci/build_artifacts.sh "${{ inputs.artifact }}"
       - name: Upload artifacts to a GCS bucket (non-Windows runs)
         if: >-

.github/workflows/ci-build.yaml

@@ -67,7 +67,7 @@ jobs:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies
       run: |
-        pip install uv
+        pip install uv~=0.5.30
         uv pip install --system .[minimum-jaxlib] -r build/test-requirements.txt
 
     - name: Run tests
@@ -105,7 +105,7 @@ jobs:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies
       run: |
-        pip install uv
+        pip install uv~=0.5.30
         uv pip install --system -r docs/requirements.txt
     - name: Test documentation
       env:
@@ -133,7 +133,7 @@ jobs:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies
       run: |
-        pip install uv
+        pip install uv~=0.5.30
         uv pip install --system -r docs/requirements.txt
     - name: Render documentation
       run: |
@@ -159,8 +159,9 @@ jobs:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies
       run: |
-        pip install uv
-        uv pip install --system .[minimum-jaxlib] tensorflow -r build/test-requirements.txt
+        pip install uv~=0.5.30
+        uv pip install --system .[minimum-jaxlib] -r build/test-requirements.txt
+        uv pip install --system --pre tensorflow==2.19.0rc0
 
     - name: Run tests
       env:
@@ -189,6 +190,7 @@ jobs:
         python-version: 3.12
     - name: Install JAX
       run: |
+        pip install uv~=0.5.30
         pip install uv
         uv pip install --system .
     - name: Build and install example project

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

@@ -59,42 +59,38 @@ jobs:
           git config --global --add safe.directory "$GITHUB_WORKSPACE"
       - name: Install JAX test requirements
         run: |
-          $PYTHON -m pip install -U -r build/test-requirements.txt
-          $PYTHON -m pip install -U -r build/collect-profile-requirements.txt
+          $PYTHON -m uv pip install -U -r build/test-requirements.txt -r build/collect-profile-requirements.txt
       - name: Install JAX
         run: |
-          $PYTHON -m pip uninstall -y jax jaxlib libtpu
+          $PYTHON -m uv pip uninstall jax jaxlib libtpu
           if [ "${{ matrix.jaxlib-version }}" == "head" ]; then
             # Build and install jaxlib at head
             $PYTHON build/build.py build --wheels=jaxlib \
                     --bazel_options=--config=rbe_linux_x86_64 \
                     --local_xla_path="$(pwd)/xla" \
                     --verbose
 
-            $PYTHON -m pip install dist/*.whl
-
-            # Install "jax" at head
-            $PYTHON  -m pip install -U -e .
-
-            # Install libtpu
-            $PYTHON -m pip install --pre libtpu \
-              -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+            # Install jaxlib, "jax" at head, and libtpu
+            $PYTHON -m uv pip install dist/*.whl \
+               -U -e . \
+               --pre libtpu -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
           elif [ "${{ matrix.jaxlib-version }}" == "pypi_latest" ]; then
-            $PYTHON -m pip install .[tpu] \
+            $PYTHON -m uv pip install .[tpu] \
               -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
 
           elif [ "${{ matrix.jaxlib-version }}" == "nightly" ]; then
-            $PYTHON -m pip install --pre . -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html
-            $PYTHON -m pip install --pre libtpu \
-              -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
-            $PYTHON -m pip install requests
+            $PYTHON -m uv pip install \
+              --pre . -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html \
+              libtpu -f https://storage.googleapis.com/jax-releases/libtpu_releases.html \
+              requests
 
           elif [ "${{ matrix.jaxlib-version }}" == "nightly+oldest_supported_libtpu" ]; then
-            $PYTHON -m pip install --pre . -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html
             # TODO(phawkins): switch to libtpu, when the oldest release we support is a libtpu release.
-            $PYTHON -m pip install --pre libtpu-nightly==0.1.dev${{ env.LIBTPU_OLDEST_VERSION_DATE }} \
-              -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
-            $PYTHON -m pip install requests
+            $PYTHON -m uv pip install \
+              --pre . -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html \
+              libtpu-nightly==0.1.dev${{ env.LIBTPU_OLDEST_VERSION_DATE }} \
+              -f https://storage.googleapis.com/jax-releases/libtpu_releases.html \
+              requests
           else
             echo "Unknown jaxlib-version: ${{ matrix.jaxlib-version }}"
             exit 1

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

@@ -74,8 +74,7 @@ jobs:
           git config --global --add safe.directory "$GITHUB_WORKSPACE"
       - name: Install JAX test requirements
         run: |
-          $JAXCI_PYTHON -m pip install -U -r build/test-requirements.txt
-          $JAXCI_PYTHON -m pip install -U -r build/collect-profile-requirements.txt
+          $JAXCI_PYTHON -m uv pip install -U -r build/test-requirements.txt -r build/collect-profile-requirements.txt
       - name: Build jaxlib at head with latest XLA
         run: |
           # Build and install jaxlib at head
@@ -86,7 +85,7 @@ jobs:
                   --verbose
 
           # Install libtpu
-          $JAXCI_PYTHON -m pip install --pre libtpu \
+          $JAXCI_PYTHON -m uv pip install --pre libtpu \
             -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
       # Halt for testing
       - name: Wait For Connection

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

@@ -37,9 +37,8 @@ jobs:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies
       run: |
-        pip install uv
-        uv pip install --system .[ci]
-        uv pip install --system pytest-xdist -r array-api-tests/requirements.txt
+        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

.github/workflows/metal_plugin_ci.yml

@@ -35,15 +35,14 @@ jobs:
           rm -rf ${GITHUB_WORKSPACE}/jax-metal-venv
           python3 -m venv ${GITHUB_WORKSPACE}/jax-metal-venv
           source  ${GITHUB_WORKSPACE}/jax-metal-venv/bin/activate
-          pip install -U pip numpy wheel
-          pip install absl-py pytest
+          pip install uv~=0.5.30
+          uv pip install -U pip numpy wheel absl-py pytest
           if [[ "${{ matrix.jaxlib-version }}" == "nightly" ]]; then
-            pip install --pre jaxlib \
+            uv pip install --pre jaxlib \
               -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html
           fi;
           cd jax
-          pip install .
-          pip install jax-metal
+          uv pip install . jax-metal
       - name: Run test
         run: |
           source  ${GITHUB_WORKSPACE}/jax-metal-venv/bin/activate

.github/workflows/pytest_cpu.yml

@@ -29,6 +29,11 @@ on:
         type: string
         required: true
         default: "0"
+      install-jax-current-commit:
+        description: "Should the 'jax' package be installed from the current commit?"
+        type: string
+        required: true
+        default: "1"
       gcs_download_uri:
         description: "GCS location prefix from where the artifacts should be downloaded"
         required: true
@@ -57,6 +62,7 @@ jobs:
       JAXCI_HERMETIC_PYTHON_VERSION: "${{ inputs.python }}"
       JAXCI_PYTHON: "python${{ inputs.python }}"
       JAXCI_ENABLE_X64: "${{ inputs.enable-x64 }}"
+      JAXCI_INSTALL_JAX_CURRENT_COMMIT: "${{ inputs.install-jax-current-commit }}"
 
     steps:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
@@ -79,18 +85,47 @@ jobs:
           echo "ARCH=${arch}" >> $GITHUB_ENV
           echo "PYTHON_MAJOR_MINOR=${python_major_minor}" >> $GITHUB_ENV
       - name: Download jaxlib wheel from GCS (non-Windows runs)
+        id: download-wheel-artifacts-nw
+        # Set continue-on-error to true to prevent actions from failing the workflow if this step
+        # fails. Instead, we verify the outcome in the step below so that we can print a more
+        # informative error message.
+        continue-on-error: true
         if: ${{ !contains(inputs.runner, 'windows-x86') }}
-        run: >-
+        run: |
           mkdir -p $(pwd)/dist &&
           gsutil -m cp -r "${{ inputs.gcs_download_uri }}/jaxlib*${PYTHON_MAJOR_MINOR}*${OS}*${ARCH}*.whl" $(pwd)/dist/
+
+          # Download the "jax" wheel from GCS if inputs.install-jax-current-commit is not set to 1
+          if [[ "${{ inputs.install-jax-current-commit }}" != 1 ]]; then
+            gsutil -m cp -r "${{ inputs.gcs_download_uri }}"/jax*py3*none*any.whl $(pwd)/dist/
+          fi
       - name: Download jaxlib wheel from GCS (Windows runs)
+        id: download-wheel-artifacts-w
+        # Set continue-on-error to true to prevent actions from failing the workflow if this step
+        # fails. Instead, we verify the outcome in step below so that we can print a more
+        # informative error message.
+        continue-on-error: true
         if: ${{ contains(inputs.runner, 'windows-x86') }}
         shell: cmd
-        run: >-
-          mkdir dist &&
-          gsutil -m cp -r "${{ inputs.gcs_download_uri }}/jaxlib*%PYTHON_MAJOR_MINOR%*%OS%*%ARCH%*.whl" dist/
+        run: |
+          mkdir dist
+          @REM Use `call` so that we can run sequential gsutil commands on Windows
+          @REM See https://github.com/GoogleCloudPlatform/gsutil/issues/233#issuecomment-196150652
+          call gsutil -m cp -r "${{ inputs.gcs_download_uri }}/jaxlib*%PYTHON_MAJOR_MINOR%*%OS%*%ARCH%*.whl" dist/
+
+          @REM Download the "jax" wheel from GCS if inputs.install-jax-current-commit is not set to 1
+          if not "${{ inputs.install-jax-current-commit }}"=="1" (
+            call gsutil -m cp -r "${{ inputs.gcs_download_uri }}"/jax*py3*none*any.whl dist/
+          )
+      - name: Skip the test run if the wheel artifacts were not downloaded successfully
+        if: steps.download-wheel-artifacts-nw.outcome == 'failure' || steps.download-wheel-artifacts-w.outcome == 'failure'
+        run: |
+          echo "Failed to download wheel artifacts from GCS. Please check if the wheels were"
+          echo "built successfully by the artifact build jobs and are available in the GCS bucket."
+          echo "Skipping the test run."
+          exit 1
       - name: Install Python dependencies
-        run: $JAXCI_PYTHON -m pip install -r build/requirements.in
+        run: $JAXCI_PYTHON -m uv pip install -r build/requirements.in
       # Halt for testing
       - name: Wait For Connection
         uses: google-ml-infra/actions/ci_connection@main