Re-factor build CLI to a subcommand based approach

[copybara-service]

Re-factor build CLI to a subcommand based approach

This commit reworks the JAX build CLI to a subcommand based approach where CLI use cases are now defined as subcommands. Two subcommands are defined: build and requirements_update. "build" is to be used when wanting to build a JAX wheel package. "requirements_update" is to be used when wanting to update the requirements_lock.txt files. 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 arguments that apply to its respective build process. In addition, arguments are separated into groups to achieve a cleaner separation and improves the readability when the CLI subcommands are run with --help. It also makes it clear as to which parts of the build they affect. E.g: CUDA arguments only apply to CUDA builds, ROCM arguments only apply to ROCM builds, etc. This reduces 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.

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 build --wheels=jaxlib --python_version=3.10

• Building jax-cuda-plugin:

python build/build.py build --wheels=jax-cuda-plugin --cuda_version=12.3.2 --cudnn_version=9.1.1 --python_version=3.10

• Building multiple packages:

python build/build.py build --wheels=jaxlib,jax-cuda-plugin,jax-cuda-pjrt --cuda_version=12.3.2 --cudnn_version=9.1.1 --python_version=3.10

• Building jax-rocm-pjrt:

python build/build.py build --wheels=jax-rocm-pjrt --rocm_version=60 --rocm_path=/path/to/rocm

• Using a local XLA path:

python build/build.py build --wheels=jaxlib --local_xla_path=/path/to/xla

• Updating requirements_lock.txt files:

python build/build.py requirements_update --python_version=3.10

For more details on each argument and to see available options, run:

python build/build.py build --help

or

python build/build.py requirements_update --help