Merge branch 'main' of https://github.com/jax-ml/jax

Author: nitins17

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

@@ -54,6 +54,10 @@ jobs:
         with:
           repository: openxla/xla
           path: xla
+      # We need to mark the GitHub workspace as safe as otherwise git commands will fail.
+      - name: Mark GitHub workspace as safe
+        run: |
+          git config --global --add safe.directory "$GITHUB_WORKSPACE"
       - name: Install JAX test requirements
         run: |
           $PYTHON -m pip install -U -r build/test-requirements.txt
@@ -63,9 +67,11 @@ jobs:
           $PYTHON -m pip uninstall -y jax jaxlib libtpu
           if [ "${{ matrix.jaxlib-version }}" == "head" ]; then
             # Build and install jaxlib at head
-            $PYTHON build/build.py --bazel_options=--config=rbe_linux_x86_64 \
-                    --bazel_options="--override_repository=xla=$(pwd)/xla" \
-                    --bazel_options=--color=yes
+            $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

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

@@ -28,7 +28,7 @@ jobs:
       with:
         repository: data-apis/array-api-tests
         # TODO(jakevdp) update this to a stable release/tag when available.
-        ref: 'ad81cf6c3721d9dbeb168bdab49c962b6b38c0d5'  # Latest commit as of 2024-11-20
+        ref: '1572b129c6682211abfe139e112592226c361a6c'  # Latest commit as of 2024-12-04
         submodules: 'true'
         path: 'array-api-tests'
     - name: Set up Python ${{ matrix.python-version }}

build/rocm/Dockerfile.ms

@@ -1,5 +1,6 @@
 ################################################################################
-FROM ubuntu:20.04 AS rocm_base
+ARG BASE_DOCKER=ubuntu:22.04
+FROM $BASE_DOCKER AS rocm_base
 ################################################################################
 
 RUN --mount=type=cache,target=/var/cache/apt \

build/rocm/ci_build

@@ -89,9 +89,9 @@ def dist_wheels(
 
     mounts = [
         "-v",
-        "./:/jax",
+        os.path.abspath("./") + ":/jax",
         "-v",
-        "./wheelhouse:/wheelhouse",
+        os.path.abspath("./wheelhouse") + ":/wheelhouse",
     ]
 
     if xla_path:
@@ -143,6 +143,7 @@ def _fetch_jax_metadata(xla_path):
 
 def dist_docker(
     rocm_version,
+    base_docker,
     python_versions,
     xla_path,
     rocm_build_job="",
@@ -168,6 +169,7 @@ def dist_docker(
         "--build-arg=ROCM_VERSION=%s" % rocm_version,
         "--build-arg=ROCM_BUILD_JOB=%s" % rocm_build_job,
         "--build-arg=ROCM_BUILD_NUM=%s" % rocm_build_num,
+        "--build-arg=BASE_DOCKER=%s" % base_docker,
         "--build-arg=PYTHON_VERSION=%s" % python_version,
         "--build-arg=JAX_VERSION=%(jax_version)s" % md,
         "--build-arg=JAX_COMMIT=%(jax_commit)s" % md,
@@ -210,7 +212,7 @@ def test(image_name):
     # JAX and jaxlib are already installed from wheels
     mounts = [
         "-v",
-        "./:/jax",
+        os.path.abspath("./") + ":/jax",
     ]
 
     cmd.extend(mounts)
@@ -231,6 +233,12 @@ def test(image_name):
 
 def parse_args():
     p = argparse.ArgumentParser()
+    p.add_argument(
+        "--base-docker",
+         default="",
+        help="Argument to override base docker in dockerfile",
+    )
+
     p.add_argument(
         "--python-versions",
         type=lambda x: x.split(","),
@@ -308,6 +316,7 @@ def main():
         )
         dist_docker(
             args.rocm_version,
+            args.base_docker,
             args.python_versions,
             args.xla_source_dir,
             rocm_build_job=args.rocm_build_job,

build/rocm/ci_build.sh

@@ -48,7 +48,7 @@ PYTHON_VERSION="3.10"
 ROCM_VERSION="6.1.3"
 ROCM_BUILD_JOB=""
 ROCM_BUILD_NUM=""
-BASE_DOCKER="ubuntu:20.04"
+BASE_DOCKER="ubuntu:22.04"
 CUSTOM_INSTALL=""
 JAX_USE_CLANG=""
 POSITIONAL_ARGS=()
@@ -90,6 +90,10 @@ while [[ $# -gt 0 ]]; do
           ROCM_BUILD_NUM="$2"
           shift 2
           ;;
+        --base_docker)
+          BASE_DOCKER="$2"
+          shift 2
+          ;;
         --use_clang)
           JAX_USE_CLANG="$2"
           shift 2
@@ -154,6 +158,7 @@ fi
 # which is the ROCm image that is shipped for users to use (i.e. distributable).
 ./build/rocm/ci_build \
     --rocm-version $ROCM_VERSION \
+    --base-docker $BASE_DOCKER \
     --python-versions $PYTHON_VERSION \
     --xla-source-dir=$XLA_CLONE_DIR \
     --rocm-build-job=$ROCM_BUILD_JOB \

build/rocm/run_multi_gpu.sh

@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-set -eu
+set -xu
 
 # Function to run tests with specified GPUs
 run_tests() {

build/rocm/tools/get_rocm.py

@@ -229,7 +229,7 @@ def _build_installer_url(rocm_version, metadata):
 
     rv = parse_version(rocm_version)
 
-    base_url = "http://artifactory-cdn.amd.com/artifactory/list"
+    base_url = "https://artifactory-cdn.amd.com/artifactory/list"
 
     if md["ID"] == "ubuntu":
         fmt = "amdgpu-install-internal_%(rocm_major)s.%(rocm_minor)s-%(os_version)s-1_all.deb"

ci/envs/default.env

@@ -58,11 +58,12 @@ export JAXCI_MAX_TESTS_PER_GPU=${JAXCI_MAX_TESTS_PER_GPU:-12}
 # this value in the Github action workflow files.
 export JAXCI_ENABLE_X64=${JAXCI_ENABLE_X64:-0}
 
-# #############################################################################
-# Docker specific environment variables.
-# #############################################################################
+# Pytest specific environment variables below. Used in run_pytest_*.sh scripts.
+# Sets the number of TPU cores for the TPU machine type. These values are
+# defined in the TPU GitHub Actions workflow.
+export JAXCI_TPU_CORES=${JAXCI_TPU_CORES:-}
 
-# Docker specifc environment variables. Used by `run_docker_container.sh`
-export JAXCI_DOCKER_WORK_DIR="/jax"
-export JAXCI_DOCKER_IMAGE=""
-export JAXCI_DOCKER_ARGS=""
+# JAXCI_PYTHON points to the Python interpreter to use for installing JAX wheels
+# on the system. By default, it is set to match the version of the hermetic
+# Python used by Bazel for building the wheels.
+export JAXCI_PYTHON=${JAXCI_PYTHON:-python${JAXCI_HERMETIC_PYTHON_VERSION}}

ci/run_pytest_cpu.sh

@@ -14,7 +14,11 @@
 # limitations under the License.
 # ==============================================================================
 # Runs Pyest CPU tests. Requires a jaxlib wheel to be present
+<<<<<<< HEAD
 # inside $JAXCI_OUTPUT_DIR (../dist)
+=======
+# inside the $JAXCI_OUTPUT_DIR (../dist)
+>>>>>>> 5ade371c88a1f879556ec29867b173da49ae57f0
 #
 # -e: abort script if one command fails
 # -u: error if undefined variable used
@@ -23,20 +27,38 @@
 # -o allexport: export all functions and variables to be available to subscripts
 set -exu -o history -o allexport
 
+<<<<<<< HEAD
 # Inherit default JAXCI environment variables.
 source ci/envs/default.env
 
+=======
+# Source default JAXCI environment variables.
+source ci/envs/default.env
+
+# Install jaxlib wheel inside the $JAXCI_OUTPUT_DIR directory on the system.
+>>>>>>> 5ade371c88a1f879556ec29867b173da49ae57f0
 echo "Installing wheels locally..."
 source ./ci/utilities/install_wheels_locally.sh
 
 # Set up the build environment.
 source "ci/utilities/setup_build_environment.sh"
 
+<<<<<<< HEAD
 export PY_COLORS=1
 export JAX_SKIP_SLOW_TESTS=true
 
 "$JAXCI_PYTHON" -c "import jax; print(jax.default_backend()); print(jax.devices()); print(len(jax.devices()))"
 
 export TF_CPP_MIN_LOG_LEVEL=0
+=======
+"$JAXCI_PYTHON" -c "import jax; print(jax.default_backend()); print(jax.devices()); print(len(jax.devices()))"
+
+# Set up all test environment variables
+export PY_COLORS=1
+export JAX_SKIP_SLOW_TESTS=true
+export TF_CPP_MIN_LOG_LEVEL=0
+# End of test environment variable setup
+
+>>>>>>> 5ade371c88a1f879556ec29867b173da49ae57f0
 echo "Running CPU tests..."
 "$JAXCI_PYTHON" -m pytest -n auto --tb=short --maxfail=20 tests examples

ci/run_pytest_gpu.sh

@@ -13,7 +13,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+<<<<<<< HEAD
 # Runs Pyest CPU tests. Requires all jaxlib, jax-cuda-plugin, and jax-cuda-pjrt
+=======
+# Runs Pyest CPU tests. Requires the jaxlib, jax-cuda-plugin, and jax-cuda-pjrt
+>>>>>>> 5ade371c88a1f879556ec29867b173da49ae57f0
 # wheels to be present inside $JAXCI_OUTPUT_DIR (../dist)
 #
 # -e: abort script if one command fails
@@ -23,16 +27,25 @@
 # -o allexport: export all functions and variables to be available to subscripts
 set -exu -o history -o allexport
 
+<<<<<<< HEAD
 # Inherit default JAXCI environment variables.
 source ci/envs/default.env
 
 # Install jaxlib, jax-cuda-plugin, and jax-cuda-pjrt wheels on the system.
+=======
+# Source default JAXCI environment variables.
+source ci/envs/default.env
+
+# Install jaxlib, jax-cuda-plugin, and jax-cuda-pjrt wheels inside the
+# $JAXCI_OUTPUT_DIR directory on the system.
+>>>>>>> 5ade371c88a1f879556ec29867b173da49ae57f0
 echo "Installing wheels locally..."
 source ./ci/utilities/install_wheels_locally.sh
 
 # Set up the build environment.
 source "ci/utilities/setup_build_environment.sh"
 
+<<<<<<< HEAD
 export PY_COLORS=1
 export JAX_SKIP_SLOW_TESTS=true
 
@@ -46,6 +59,28 @@ echo "Running GPU tests..."
 export XLA_PYTHON_CLIENT_ALLOCATOR=platform
 export XLA_FLAGS=--xla_gpu_force_compilation_parallelism=1
 "$JAXCI_PYTHON" -m pytest -n 8 --tb=short --maxfail=20 \
+=======
+"$JAXCI_PYTHON" -c "import jax; print(jax.default_backend()); print(jax.devices()); print(len(jax.devices()))"
+
+nvidia-smi
+
+# Set up all test environment variables
+export PY_COLORS=1
+export JAX_SKIP_SLOW_TESTS=true
+export NCCL_DEBUG=WARN
+export TF_CPP_MIN_LOG_LEVEL=0
+
+# Set the number of processes to run to be 4x the number of GPUs.
+export gpu_count=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)
+export num_processes=`expr 4 \* $gpu_count`
+
+export XLA_PYTHON_CLIENT_ALLOCATOR=platform
+export XLA_FLAGS=--xla_gpu_force_compilation_parallelism=1
+# End of test environment variable setup
+
+echo "Running GPU tests..."
+"$JAXCI_PYTHON" -m pytest -n $num_processes --tb=short --maxfail=20 \
+>>>>>>> 5ade371c88a1f879556ec29867b173da49ae57f0
 tests examples \
 --deselect=tests/multi_device_test.py::MultiDeviceTest::test_computation_follows_data \
 --deselect=tests/xmap_test.py::XMapTest::testCollectivePermute2D \