update scripts

nitins17 · nitins17 · commit 3b55033fd3b8 · 2024-09-19T17:57:41.000Z
diff --git a/ci/.bazelrc b/ci/.bazelrc
@@ -415,6 +415,38 @@ build:cross_compile_darwin_x86_64 --platform_mappings=platform_mappings
 build:rbe_cross_compile_darwin_x86_64 --config=cross_compile_darwin_x86_64
 build:rbe_cross_compile_darwin_x86_64 --config=rbe_cross_compile_base
 
+# #############################################################################
+# Test specific config options below. These are used when `bazel test` is run.
+# #############################################################################
+test --test_output=errors
+
+# Common configs for for running GPU tests.
+test:gpu --test_env=TF_CPP_MIN_LOG_LEVEL=0 --test_env=XLA_PYTHON_CLIENT_ALLOCATOR=platform
+
+# Non-multiaccelerator tests with one GPU apiece. These tests are run on RBE
+# and locally.
+test:non_multiaccelerator --config=gpu
+test:non_multiaccelerator --test_env=JAX_EXCLUDE_TEST_TARGETS=PmapTest.testSizeOverflow
+test:non_multiaccelerator --test_tag_filters=-multiaccelerator
+
+# Configs for running non-multiaccelerator tests locally
+test:non_multiaccelerator_local --config=non_multiaccelerator
+# Disable building jaxlib. Instead we depend on the local wheel.
+test:non_multiaccelerator_local --//jax:build_jaxlib=false
+
+# `JAX_ACCELERATOR_COUNT` needs to match the number of GPUs in the VM.
+test:non_multiaccelerator_local --test_env=JAX_TESTS_PER_ACCELERATOR=12 --test_env=JAX_ACCELERATOR_COUNT=4
+
+# The product of the `JAX_ACCELERATOR_COUNT`` and `JAX_TESTS_PER_ACCELERATOR`
+# should match the VM's CPU core count (set in `--local_test_jobs`).
+test:non_multiaccelerator_local --local_test_jobs=48
+
+# Multiaccelerator tests with all GPUs. These tests are only run locally
+# Disable building jaxlib. Instead we depend on the local wheel.
+test:multiaccelerator_local --config=gpu
+test:multiaccelerator_local --//jax:build_jaxlib=false
+test:multiaccelerator_local --jobs=8 --test_tag_filters=multiaccelerator
+
 #############################################################################
 # Some configs to make getting some forms of debug builds. In general, the
 # codebase is only regularly built with optimizations. Use 'debug_symbols' to
diff --git a/ci/envs/default b/ci/envs/default
@@ -64,6 +64,40 @@ export JAXCI_DOCKER_WORK_DIR="/jax"
 export JAXCI_DOCKER_IMAGE=""
 export JAXCI_DOCKER_ARGS=""
 
+# #############################################################################
+# Test specific environment variables.
+# #############################################################################
+
+# Used by envs inside ci/build_artifacts. When set to 1, we disable x64 mode
+# and clone XLA at HEAD.
+export JAXCI_SETUP_TEST_ENVIRONMENT=${JAXCI_SETUP_TEST_ENVIRONMENT:-0}
+
+# Set when running tests locally where we need the wheels to be installed on
+# the system.
+export JAXCI_INSTALL_WHEELS_LOCALLY=0
+
+# JAXCI_PYTHON is used to install the wheels locally. It needs to match the
+# version of the hermetic Python used by Bazel.
+export JAXCI_PYTHON=python${JAXCI_HERMETIC_PYTHON_VERSION}
+
+# Bazel test environment variables.
+export JAXCI_RUN_BAZEL_TEST_CPU=0
+export JAXCI_RUN_BAZEL_TEST_GPU_LOCAL=0
+export JAXCI_RUN_BAZEL_TEST_GPU_RBE=0
+
+# Pytest environment variables.
+export JAXCI_RUN_PYTEST_CPU=0
+export JAXCI_RUN_PYTEST_GPU=0
+export JAXCI_RUN_PYTEST_TPU=0
+export JAXCI_TPU_CORES=""
+
+# If set to 1, the script will clone the main XLA repository at HEAD, set its
+# path in JAXCI_XLA_GIT_DIR and use it to build the artifacts or run the tests.
+export JAXCI_CLONE_MAIN_XLA=${JAXCI_CLONE_MAIN_XLA:-0}
+
+# Enable this globally across all builds.
+export JAX_SKIP_SLOW_TESTS=true
+
 # #############################################################################
 # Variables that can be overridden by the user.
 # #############################################################################
diff --git a/ci/envs/run_tests/bazel_cpu b/ci/envs/run_tests/bazel_cpu
@@ -1,3 +1,17 @@
+# Copyright 2024 The JAX Authors.
+#
+# 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.
+# ==============================================================================
 # Inherit default JAXCI environment variables.
 source ci/envs/default
 
diff --git a/ci/envs/run_tests/bazel_gpu_local b/ci/envs/run_tests/bazel_gpu_local
@@ -1,3 +1,17 @@
+# Copyright 2024 The JAX Authors.
+#
+# 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.
+# ==============================================================================
 # Inherit default JAXCI environment variables.
 source ci/envs/default
 
diff --git a/ci/envs/run_tests/bazel_gpu_rbe b/ci/envs/run_tests/bazel_gpu_rbe
@@ -1,3 +1,17 @@
+# Copyright 2024 The JAX Authors.
+#
+# 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.
+# ==============================================================================
 # Inherit default JAXCI environment variables.
 source ci/envs/default
 
diff --git a/ci/envs/run_tests/pytest_cpu b/ci/envs/run_tests/pytest_cpu
@@ -1,3 +1,17 @@
+# Copyright 2024 The JAX Authors.
+#
+# 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.
+# ==============================================================================
 # Inherit default JAXCI environment variables.
 source ci/envs/default
 
diff --git a/ci/envs/run_tests/pytest_gpu b/ci/envs/run_tests/pytest_gpu
@@ -1,3 +1,17 @@
+# Copyright 2024 The JAX Authors.
+#
+# 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.
+# ==============================================================================
 # Inherit default JAXCI environment variables.
 source ci/envs/default
 
diff --git a/ci/envs/run_tests/pytest_tpu b/ci/envs/run_tests/pytest_tpu
@@ -1,3 +1,17 @@
+# Copyright 2024 The JAX Authors.
+#
+# 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.
+# ==============================================================================
 # Inherit default JAXCI environment variables.
 source ci/envs/default
 
diff --git a/ci/run_bazel_test.sh b/ci/run_bazel_test.sh
@@ -1,5 +1,5 @@
 #!/bin/bash
-# Copyright 2024 JAX Authors. All Rights Reserved.
+# Copyright 2024 The JAX Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -14,7 +14,7 @@
 # limitations under the License.
 # ==============================================================================
 # Source JAXCI environment variables.
-source "ci/utilities/setup_envs.sh"
+source "ci/utilities/setup_envs.sh" "$1"
 # Set up the build environment.
 source "ci/utilities/setup_build_environment.sh"
 
diff --git a/ci/run_pytest.sh b/ci/run_pytest.sh
@@ -1,5 +1,5 @@
 #!/bin/bash
-# Copyright 2024 JAX Authors. All Rights Reserved.
+# Copyright 2024 The JAX Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -14,7 +14,7 @@
 # limitations under the License.
 # ==============================================================================
 # Source JAXCI environment variables.
-source "ci/utilities/setup_envs.sh"
+source "ci/utilities/setup_envs.sh" "$1"
 # Set up the build environment.
 source "ci/utilities/setup_build_environment.sh"
 
@@ -40,13 +40,13 @@ fi
 if [[ $JAXCI_RUN_PYTEST_TPU == 1 ]]; then
   echo "Running TPU tests..."
   # Run single-accelerator tests in parallel
-  export JAX_ENABLE_TPU_XDIST=true 
+  export JAX_ENABLE_TPU_XDIST=true
   check_if_to_run_in_docker "$JAXCI_PYTHON" -m pytest -n="$JAXCI_TPU_CORES" --tb=short \
     --deselect=tests/pallas/tpu_pallas_test.py::PallasCallPrintTest \
     --maxfail=20 -m "not multiaccelerator" tests examples
- 
+
   # Run Pallas printing tests, which need to run with I/O capturing disabled.
-  export TPU_STDERR_LOG_LEVEL=0 
+  export TPU_STDERR_LOG_LEVEL=0
   check_if_to_run_in_docker "$JAXCI_PYTHON" -m pytest -s tests/pallas/tpu_pallas_test.py::PallasCallPrintTest
 
   # Run multi-accelerator across all chips
diff --git a/ci/utilities/install_wheels_locally.sh b/ci/utilities/install_wheels_locally.sh
@@ -1,5 +1,5 @@
 #!/bin/bash
-# Copyright 2024 JAX Authors. All Rights Reserved.
+# Copyright 2024 The JAX Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -15,19 +15,15 @@
 # ==============================================================================
 #
 # Install wheels stored in `JAXCI_OUTPUT_DIR` locally using the Python binary
-# set in JAXCI_PYTHON.
+# set in JAXCI_PYTHON. Use the absolute path to the `find` utility to avoid
+# using the Windows version of `find` on Windows.
+WHEELS=$(/usr/bin/find "$JAXCI_OUTPUT_DIR/" -type f \( -name "*jaxlib*" -o -name "*jax*cuda*pjrt*" -o -name "*jax*cuda*plugin*" \))
 
-# When running tests with Pytests, install wheels using the Python binary set
-# in JAXCI_PYTHON.
-if [[ $JAXCI_RUN_PYTEST_CPU == 1 ]] || [[ $JAXCI_RUN_PYTEST_GPU == 1 ]]; then
-  # Install the wheels inside the output directory. On Windows, the path needs
-  # to be in Linux format.
-  if [[ ! $(uname -s) =~ "MSYS_NT" ]]; then
-    bash -c "$JAXCI_PYTHON -m pip install $JAXCI_OUTPUT_DIR/*.whl"
-  else
-    bash -c "$JAXCI_PYTHON -m pip install $(cygpath $JAXCI_OUTPUT_DIR)/*.whl"
-  fi
+for wheel in "$WHEELS"; do
+  echo "Installing $(basename $wheel) ..."
+  "$JAXCI_PYTHON" -m pip install "$wheel"
+done
 
-  # Install JAX package at the current commit.
-  "$JAXCI_PYTHON" -m pip install -U -e .
-fi
+echo "Installing the JAX package in editable mode at the current commit..."
+# Install JAX package at the current commit.
+"$JAXCI_PYTHON" -m pip install -U -e .
