source install wheels script from the build scripts

Author: nitins17

ci/run_pytest_cpu.sh

@@ -26,8 +26,8 @@ set -exu -o history -o allexport
 # Inherit default JAXCI environment variables.
 source ci/envs/default.env
 
-# Install jaxlib wheel on the system. 
-export JAXCI_INSTALL_WHEELS_LOCALLY=1
+echo "Installing wheels locally..."
+source ./ci/utilities/install_wheels_locally.sh
 
 # Set up the build environment.
 source "ci/utilities/setup_build_environment.sh"

ci/run_pytest_gpu.sh

@@ -27,7 +27,8 @@ set -exu -o history -o allexport
 source ci/envs/default.env
 
 # Install jaxlib, jax-cuda-plugin, and jax-cuda-pjrt wheels on the system.
-export JAXCI_INSTALL_WHEELS_LOCALLY=1
+echo "Installing wheels locally..."
+source ./ci/utilities/install_wheels_locally.sh
 
 # Set up the build environment.
 source "ci/utilities/setup_build_environment.sh"

ci/run_pytest_tpu.sh

@@ -28,7 +28,8 @@ source ci/envs/default.env
 
 # Install jaxlib wheel on the system. Requires a jaxlib wheel to be present
 # inside $JAXCI_OUTPUT_DIR (../dist)
-export JAXCI_INSTALL_WHEELS_LOCALLY=1
+echo "Installing wheels locally..."
+source ./ci/utilities/install_wheels_locally.sh
 
 # Set up the build environment.
 source "ci/utilities/setup_build_environment.sh"

ci/utilities/setup_build_environment.sh

@@ -72,10 +72,4 @@ if [[ $(uname -s) =~ "MSYS_NT" ]]; then
   echo 'Converting MSYS Linux-like paths to Windows paths (for Docker, Python, etc.)'
   # Convert all "_DIR" variables to Windows paths.
   source <(python3 ./ci/utilities/convert_msys_paths_to_win_paths.py)
-fi
-
-# When running Pytests, we need to install the wheels locally.
-if [[ "$JAXCI_INSTALL_WHEELS_LOCALLY" == 1 ]]; then
-   echo "Installing wheels locally..."
-   source ./ci/utilities/install_wheels_locally.sh
 fi