fix invalid merge

Author: nitins17

ci/envs/default.env

@@ -66,4 +66,4 @@ 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}}
+export JAXCI_PYTHON=${JAXCI_PYTHON:-python${JAXCI_HERMETIC_PYTHON_VERSION}}

ci/run_pytest_gpu.sh

@@ -13,11 +13,7 @@
 # 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
@@ -27,39 +23,17 @@
 # -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
-
-"$JAXCI_PYTHON" -c "import jax; print(jax.default_backend()); print(jax.devices()); print(len(jax.devices()))"
-
-nvidia-smi
-export NCCL_DEBUG=WARN
-export TF_CPP_MIN_LOG_LEVEL=0
-
-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
@@ -80,7 +54,6 @@ export XLA_FLAGS=--xla_gpu_force_compilation_parallelism=1
 
 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 \

ci/run_pytest_tpu.sh

@@ -33,29 +33,27 @@ 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'
+strings /usr/local/lib/"$JAXCI_PYTHON"/dist-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..."
+# Set up common test environment variables
+export PY_COLORS=1
+export JAX_SKIP_SLOW_TESTS=true
 export JAX_PLATFORMS=tpu,cpu
-# Run single-accelerator tests in parallel
-export JAX_ENABLE_TPU_XDIST=true
+# End of common test environment variable setup
 
-"$JAXCI_PYTHON" -m pytest -n="$JAXCI_TPU_CORES" --tb=short \
+echo "Running TPU tests..."
+# Run single-accelerator tests in parallel
+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
+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

ci/utilities/convert_msys_paths_to_win_paths.py

@@ -16,11 +16,7 @@
 Converts MSYS Linux-like paths stored in env variables to Windows paths.
 
 This is necessary on Windows, because some applications do not understand/handle
-<<<<<<< HEAD
-Linux-like paths MSYS uses, for example, Docker.
-=======
 Linux-like paths MSYS uses, for example, Bazel.
->>>>>>> 5ade371c88a1f879556ec29867b173da49ae57f0
 """
 import argparse
 import os
@@ -47,8 +43,6 @@ def msys_to_windows_path(msys_path):
     print(f"Error converting path: {e}")
     return None
 
-<<<<<<< HEAD
-=======
 def should_convert(var: str,
                    convert: list[str] | None):
   """Check the variable name against convert list"""
@@ -57,24 +51,15 @@ def should_convert(var: str,
   else:
     return False
 
->>>>>>> 5ade371c88a1f879556ec29867b173da49ae57f0
 def main(parsed_args: argparse.Namespace):
   converted_paths = {}
 
   for var, value in os.environ.items():
-<<<<<<< HEAD
-    if (parsed_args.blacklist and var in parsed_args.blacklist) or not value:
-      continue
-    if "_DIR" in var or (args.whitelist and var in parsed_args.whitelist):
-      converted_path = msys_to_windows_path(value)
-      converted_paths[var] = converted_path
-=======
     if not value or not should_convert(var,
                                        parsed_args.convert):
       continue
     converted_path = msys_to_windows_path(value)
     converted_paths[var] = converted_path
->>>>>>> 5ade371c88a1f879556ec29867b173da49ae57f0
 
   var_str = '\n'.join(f'export {k}="{v}"'
                       for k, v in converted_paths.items())
@@ -86,19 +71,10 @@ def main(parsed_args: argparse.Namespace):
 if __name__ == '__main__':
   parser = argparse.ArgumentParser(description=(
       'Convert MSYS paths in environment variables to Windows paths.'))
-<<<<<<< HEAD
-  parser.add_argument('--blacklist',
-                      nargs='*',
-                      help='List of variables to ignore')
-  parser.add_argument('--whitelist',
-                      nargs='*',
-                      help='List of variables to include')
-=======
   parser.add_argument('--convert',
                       nargs='+',
                       required=True,
                       help='Space separated list of environment variables to convert. E.g: --convert env_var1 env_var2')
->>>>>>> 5ade371c88a1f879556ec29867b173da49ae57f0
   args = parser.parse_args()
 
-  main(args)
+  main(args)