Re-factor build CLI to a subcommand based approach

This commit reworks the JAX build CLI to a subcommand based approach where artifacts are now defined as CLI subcommands. The new structure offers a clear and organized CLI that enables users to execute specific build tasks without having to navigate through a monolithic script.

Each subcommand has specific options and arguments that apply to its respective build process. This allows users to execute targeted build commands with relevant options (e.g CUDA arguments only apply to CUDA subcommands, ROCM arguments only apply to ROCM subcommands, etc.). This would also reduce the complexity and the potential for errors during the build process.  Segregating functionalities into distinct subcommands also simplifies the code which should help with the maintenance and future extensions.

A new flag, `--use_ci_bazelrc_flags`, is introduced and is intended to be used in CI builds. When set, the build CLI will apply the "ci_"/"rbe_" config from the .bazelrc.

There is also a transition from using `subprocess.check_output` to `asyncio.create_subprocess_shell` for executing the build commands which allows for streaming logs and helps in showing the build progress in real time.

Usage:
* Building `jaxlib`:
```
python build/build.py jaxlib --python_version=3.10
```
* Building `jax-cuda-plugin`:
```
python build/build.py jax-cuda-plugin --cuda_version=12.3.2 --cudnn_version=9.1.1 --python_version=3.10
```
* Building `jax-rocm-pjrt`:
```
python build/build.py jax-rocm-pjrt --rocm_version=60 --rocm_path=/path/to/rocm
```
* Using a local XLA path:
```
python build/build.py jaxlib --local_xla_path=/path/to/xla
```
* Updating requirements_lock.txt files:
```
python build/build.py requirements_update --python_version=3.10
```
PiperOrigin-RevId: 691466647

.github/workflows/asan.yaml

@@ -65,9 +65,9 @@ jobs:
         run: |
           source ${GITHUB_WORKSPACE}/venv/bin/activate
           cd jax
-          python build/build.py \
-            --bazel_options=--color=yes \
-            --bazel_options=--copt=-fsanitize=address \
+          python build/build.py jaxlib --verbose \
+            --bazel_build_options=--verbose_failures=true \
+            --bazel_build_options=--copt=-fsanitize=address \
             --clang_path=/usr/bin/clang-18
           pip install dist/jaxlib-*.whl
           pip install -e .

.github/workflows/wheel_win_x64.yml

@@ -40,10 +40,9 @@ jobs:
           python -m pip install -r build/test-requirements.txt
           python -m pip install --upgrade numpy==2.0.0 scipy==1.13.1
           "C:\\msys64\\;C:\\msys64\\usr\\bin\\;" >> $env:GITHUB_PATH
-          python.exe build\build.py `
-            --bazel_options=--color=yes `
-            --bazel_options=--config=win_clang `
-            --verbose
+          python.exe build\build.py jaxlib --verbose `
+            --bazel_build_options=--verbose_failures=true `
+            --bazel_build_options=--config=win_clang
 
       - uses: actions/upload-artifact@b4b15b8c7c6ac21ea08fcf65892d2ee8f75cf882  # v4.4.3
         with:

.github/workflows/windows_ci.yml

@@ -49,9 +49,9 @@ jobs:
           python -m pip install -r build/test-requirements.txt
           python -m pip install --upgrade numpy==2.0.0 scipy==1.13.1
           "C:\\msys64\\;C:\\msys64\\usr\\bin\\;" >> $env:GITHUB_PATH
-          python.exe build\build.py `
-            --bazel_options=--color=yes `
-            --bazel_options=--config=win_clang
+          python.exe build\build.py jaxlib --verbose `
+            --bazel_build_options=--verbose_failures=true `
+            --bazel_build_options=--config=win_clang
 
       - uses: actions/upload-artifact@b4b15b8c7c6ac21ea08fcf65892d2ee8f75cf882  # v4.4.3
         with: