Add new CI scripts for running Pytests

nitins17 · Google-ML-Automation · commit 031c0acf5053 · 2024-11-29T12:08:17.000-08:00
This commit adds the new CI scripts for running Pytests. It makes use of the pytest envs inside the "ci/envs/run_tests" folder to control the build behavior. For e.g: for running the GPU tests with Pytest, we will need to run `./ci/run_pytest.sh ./ci/envs/run_tests/pytest_gpu.env`. Note that Pytests need JAX wheels to be installed on the system to work. The `install_wheels_locally.sh` script installs these wheels in CI builds.

PiperOrigin-RevId: 701331411
diff --git a/ci/envs/default.env b/ci/envs/default.env
@@ -57,3 +57,13 @@ export JAXCI_MAX_TESTS_PER_GPU=${JAXCI_MAX_TESTS_PER_GPU:-12}
 # Sets the value of `JAX_ENABLE_X64` in the test scripts. CI builds override
 # this value in the Github action workflow files.
 export JAXCI_ENABLE_X64=${JAXCI_ENABLE_X64:-0}
+
+# 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:-}
+
+# 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}}
diff --git a/ci/run_pytest_cpu.sh b/ci/run_pytest_cpu.sh
@@ -0,0 +1,45 @@
+#!/bin/bash
+# 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.
+# ==============================================================================
+# Runs Pyest CPU tests. Requires a jaxlib wheel to be present
+# inside the $JAXCI_OUTPUT_DIR (../dist)
+#
+# -e: abort script if one command fails
+# -u: error if undefined variable used
+# -x: log all commands
+# -o history: record shell history
+# -o allexport: export all functions and variables to be available to subscripts
+set -exu -o history -o allexport
+
+# Source default JAXCI environment variables.
+source ci/envs/default.env
+
+# Install jaxlib wheel inside the $JAXCI_OUTPUT_DIR directory on the system.
+echo "Installing wheels locally..."
+source ./ci/utilities/install_wheels_locally.sh
+
+# Set up the build environment.
+source "ci/utilities/setup_build_environment.sh"
+
+"$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
+
+echo "Running CPU tests..."
+"$JAXCI_PYTHON" -m pytest -n auto --tb=short --maxfail=20 tests examples
diff --git a/ci/run_pytest_gpu.sh b/ci/run_pytest_gpu.sh
@@ -0,0 +1,61 @@
+#!/bin/bash
+# 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.
+# ==============================================================================
+# Runs Pyest CPU tests. Requires the jaxlib, jax-cuda-plugin, and jax-cuda-pjrt
+# wheels to be present inside $JAXCI_OUTPUT_DIR (../dist)
+#
+# -e: abort script if one command fails
+# -u: error if undefined variable used
+# -x: log all commands
+# -o history: record shell history
+# -o allexport: export all functions and variables to be available to subscripts
+set -exu -o history -o allexport
+
+# 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.
+echo "Installing wheels locally..."
+source ./ci/utilities/install_wheels_locally.sh
+
+# Set up the build environment.
+source "ci/utilities/setup_build_environment.sh"
+
+"$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 \
+tests examples \
+--deselect=tests/multi_device_test.py::MultiDeviceTest::test_computation_follows_data \
+--deselect=tests/xmap_test.py::XMapTest::testCollectivePermute2D \
+--deselect=tests/multiprocess_gpu_test.py::MultiProcessGpuTest::test_distributed_jax_visible_devices \
+--deselect=tests/compilation_cache_test.py::CompilationCacheTest::test_task_using_cache_metric
diff --git a/ci/run_pytest_tpu.sh b/ci/run_pytest_tpu.sh
@@ -0,0 +1,61 @@
+#!/bin/bash
+# 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.
+# ==============================================================================
+# Runs Pyest CPU tests. Requires a jaxlib wheel to be present
+# inside $JAXCI_OUTPUT_DIR (../dist)
+#
+# -e: abort script if one command fails
+# -u: error if undefined variable used
+# -x: log all commands
+# -o history: record shell history
+# -o allexport: export all functions and variables to be available to subscripts
+set -exu -o history -o allexport
+
+# Source default JAXCI environment variables.
+source ci/envs/default.env
+
+# Install jaxlib wheel inside the $JAXCI_OUTPUT_DIR directory on the system.
+echo "Installing wheels locally..."
+source ./ci/utilities/install_wheels_locally.sh
+
+# Set up the build environment.
+source "ci/utilities/setup_build_environment.sh"
+
+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()))"
+
+"$JAXCI_PYTHON" -c 'import sys; print("python version:", sys.version)'
+"$JAXCI_PYTHON" -c 'import jax; print("jax version:", jax.__version__)'
+"$JAXCI_PYTHON" -c 'import jaxlib; print("jaxlib version:", jaxlib.__version__)'
+strings /usr/local/lib/"$JAXCI_PYTHON"/site-packages/libtpu/libtpu.so | grep 'Built on'
+"$JAXCI_PYTHON" -c 'import jax; print("libtpu version:",jax.lib.xla_bridge.get_backend().platform_version)'
+
+echo "Running TPU tests..."
+export JAX_PLATFORMS=tpu,cpu
+# Run single-accelerator tests in parallel
+export JAX_ENABLE_TPU_XDIST=true
+
+"$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
+"$JAXCI_PYTHON" -m pytest -s tests/pallas/tpu_pallas_test.py::PallasCallPrintTest
+
+# Run multi-accelerator across all chips
+"$JAXCI_PYTHON" -m pytest --tb=short --maxfail=20 -m "multiaccelerator" tests
diff --git a/ci/utilities/install_wheels_locally.sh b/ci/utilities/install_wheels_locally.sh
@@ -0,0 +1,33 @@
+#!/bin/bash
+# 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.
+# ==============================================================================
+#
+# Install wheels stored in `JAXCI_OUTPUT_DIR` on the system using the Python
+# binary set in JAXCI_PYTHON. Use the absolute path to the `find` utility to
+# avoid using the Windows version of `find` on Msys.
+WHEELS=( $(/usr/bin/find "$JAXCI_OUTPUT_DIR/" -type f \( -name "*jaxlib*" -o -name "*jax*cuda*pjrt*" -o -name "*jax*cuda*plugin*" \)) )
+
+if [[ -z "$WHEELS" ]]; then
+  echo "ERROR: No wheels found under $JAXCI_OUTPUT_DIR"
+  exit 1
+fi
+
+echo "Installing the following wheels:"
+echo "${WHEELS[@]}"
+"$JAXCI_PYTHON" -m pip install "${WHEELS[@]}"
+
+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 .
