Skip to content

Commit 36d04de

Browse files
nitins17nitins17
committed
update scripts
1 parent 0bfc868 commit 36d04de

File tree

5 files changed

+13
-167
lines changed

5 files changed

+13
-167
lines changed

‎ci/.bazelrc‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -437,4 +437,4 @@ build:debug --config debug_symbols -c fastbuild
437437
try-import %workspace%/.jax_configure.bazelrc
438438

439439
# Load rc file with user-specific options.
440-
try-import %workspace%/.bazelrc.user
440+
try-import %workspace%/.bazelrc.user

‎ci/README.md‎

Lines changed: 1 addition & 122 deletions
Original file line numberDiff line numberDiff line change
@@ -6,125 +6,4 @@
66
> directory are stable and appropriate documentation around its usage is in
77
> place.
88
9-
********************************************************************************
10-
11-
The CI folder contains the configuration files and scripts used to build, test,
12-
and upload JAX artifacts.
13-
14-
## JAX's Official CI and Build/Test Scripts
15-
16-
JAX's official CI jobs run the scripts in this folder. The CI scripts require
17-
an env file to be set in `JAXCI_ENV_FILE` that sets various configuration settings.
18-
These "env" files are structured by their build type. For e.g.,
19-
`ci/envs/build_artifacts/jaxlib` contains the configs for building the `jaxlib`
20-
package. The scripts are intended to be used across different platforms and
21-
architectures and currently supports the following systems: Linux x86,
22-
Linux Arm64, Mac x86, Mac Arm64, Windows x86.
23-
24-
25-
If you would like to test these scripts, follow the instructions below.
26-
27-
### Choose how you would like to build:
28-
<details>
29-
<summary> Build CLI </summary>
30-
31-
Follow the instructions below to invoke the build CLI to build a JAX artifact
32-
of your choice. The CLI can build the `jaxlib`, `jax-cuda-plugin`, and the
33-
`jax-cuda-pjrt` artifacts. Note that all commands are meant to be run from the
34-
root of this repository.
35-
36-
By default, the CLI runs in local mode and will pick the "local_" configs in
37-
the `ci/.bazelrc` file. On Linux systems, Bazel defaults to using GCC
38-
as the default compiler. To change this, add `--use_clang` to your command. This
39-
requires Clang to be present on the system and in the path. If your Clang binary
40-
is not on the path, set its path using `--clang_path`.
41-
42-
**Build Modes**
43-
44-
If you want to run with the configs that the CI builds use, switch the mode by
45-
setting `--mode=ci`. Please note CI mode has a dependency on a custom toolchain
46-
that JAX uses. The build expects this toolchain to be present on the system. As
47-
such, CI mode is usually run from within a Docker container. See `JAXCI_DOCKER_IMAGE`
48-
inside `ci/build_artfacts` to know which image we use for each platform.
49-
50-
**RBE support**
51-
52-
If you are running this on a Linux x86 or a Windows machine, you have the option
53-
to use RBE to speed up the build. Please note this requires permissions to JAX's
54-
remote worker pool and RBE configs. To enable RBE, set `--use_rbe` to you command.
55-
56-
**Changing Python version**
57-
58-
If you would like to change the Python version of the artifact, add
59-
`--python_version=<python_version>` to your command. E.g. `--python_version=3.11`.
60-
By default, the CLI uses Python 3.12.
61-
62-
**Local XLA dependency**
63-
64-
JAX artifacts built by the CLI depend on XLA version pinned in JAX's
65-
`workspace.bzl`. If would like to depend on the XLA from your local system,
66-
set `--local_xla_path` to its path.
67-
68-
**Dry Run**
69-
70-
If you would like to just invoke a dry run, add `--dry_run` to your command.
71-
This will print the `bazel` command that the CLI would have ended up invoking.
72-
73-
**Some example invocations**
74-
75-
1. For building `jaxlib`, run `python ci/cli/build.py jaxlib`
76-
2. For building `jax-cuda-plugin` for Python 3.11, run `python ci/cli/build.py jax-cuda-pjrt --python_version=3.11`
77-
3. For building `jax-cuda-pjrt` for Python 3.10 with RBE, run `python ci/cli/build.py jax-cuda-pjrt --use_rbe --python_version=3.10`
78-
79-
</details>
80-
81-
<details>
82-
<summary> Shell Script </summary>
83-
84-
The artifact building script (`ci/build_artifacts.sh`) invokes the build CLI,
85-
`ci/cli/build.py` which in turn invokes the bazel command that builds the
86-
requested JAX artifact. Follow the instructions below to invoke the CI script
87-
to build a JAX artifact of your choice. These scripts can build the `jax`,
88-
`jaxlib`, `jax-cuda-plugin`, and the `jax-cuda-pjrt` artifacts. Note that all
89-
commands are meant to be run from the root of this repository.
90-
91-
**Docker (soft prerequisite)**
92-
93-
The CI scripts are recommended to be run in Docker where possible. This ensures
94-
the right build environment is set up before we can build the artifact. If you
95-
would like to disable Docker, run:
96-
97-
```
98-
export JAXCI_SETUP_DOCKER=0
99-
export JAXCI_CLI_BUILD_MODE=local
100-
```
101-
102-
**Changing Python version**
103-
104-
By default, the build will use Python 3.12. If you would like to change this,
105-
set `JAXCI_HERMETIC_PYTHON_VERSION`. E.g.`export JAXCI_HERMETIC_PYTHON_VERSION=3.11`
106-
107-
**RBE support**
108-
109-
If you are running this on a Linux x86 or a Windows machine, you have the option
110-
to use RBE to speed up the build. Please note this requires permissions to JAX's
111-
remote worker pool and RBE configs. To enable RBE, run `export JAXCI_BUILD_ARTIFACT_WITH_RBE=1`.
112-
113-
**How to run the script**
114-
115-
```
116-
1. Set JAXCI_ENV_FILE to one of the envs inside ci/build_artifacts based the
117-
artifact you want to build.
118-
E.g. export JAXCI_ENV_FILE=ci/envs/build_artifacts/jaxlib
119-
2. Run: bash ci/build_artifacts.sh
120-
```
121-
122-
**Known Bugs**
123-
124-
1. Building `jax` fails due to Python missing the `build` dependency.
125-
2. Auditwheel script fails on Linux Arm64's Docker image due to Python missing
126-
the `auditwheel` dependency
127-
3. If RBE is used to build the target for Windows, building the wheel fails
128-
due to a permission denied error.
129-
130-
</details>
9+
********************************************************************************

‎ci/build_artifacts.sh‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
# limitations under the License.
1515
# ==============================================================================
1616
# Source JAXCI environment variables.
17-
source "ci/utilities/setup_envs.sh"
17+
source "ci/utilities/setup_envs.sh" "$1"
1818
# Set up the build environment.
1919
source "ci/utilities/setup_build_environment.sh"
2020

‎ci/envs/default‎

Lines changed: 1 addition & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -64,40 +64,6 @@ export JAXCI_DOCKER_WORK_DIR="/jax"
6464
export JAXCI_DOCKER_IMAGE=""
6565
export JAXCI_DOCKER_ARGS=""
6666

67-
# #############################################################################
68-
# Test specific environment variables.
69-
# #############################################################################
70-
71-
# Used by envs inside ci/build_artifacts. When set to 1, we disable x64 mode
72-
# and clone XLA at HEAD.
73-
export JAXCI_SETUP_TEST_ENVIRONMENT=${JAXCI_SETUP_TEST_ENVIRONMENT:-0}
74-
75-
# Set when running tests locally where we need the wheels to be installed on
76-
# the system.
77-
export JAXCI_INSTALL_WHEELS_LOCALLY=0
78-
79-
# JAXCI_PYTHON is used to install the wheels locally. It needs to match the
80-
# version of the hermetic Python used by Bazel.
81-
export JAXCI_PYTHON=python${JAXCI_HERMETIC_PYTHON_VERSION}
82-
83-
# Bazel test environment variables.
84-
export JAXCI_RUN_BAZEL_TEST_CPU=0
85-
export JAXCI_RUN_BAZEL_TEST_GPU_LOCAL=0
86-
export JAXCI_RUN_BAZEL_TEST_GPU_RBE=0
87-
88-
# Pytest environment variables.
89-
export JAXCI_RUN_PYTEST_CPU=0
90-
export JAXCI_RUN_PYTEST_GPU=0
91-
export JAXCI_RUN_PYTEST_TPU=0
92-
export JAXCI_TPU_CORES=""
93-
94-
# If set to 1, the script will clone the main XLA repository at HEAD, set its
95-
# path in JAXCI_XLA_GIT_DIR and use it to build the artifacts or run the tests.
96-
export JAXCI_CLONE_MAIN_XLA=${JAXCI_CLONE_MAIN_XLA:-0}
97-
98-
# Enable this globally across all builds.
99-
export JAX_SKIP_SLOW_TESTS=true
100-
10167
# #############################################################################
10268
# Variables that can be overridden by the user.
10369
# #############################################################################
@@ -112,4 +78,4 @@ export JAXCI_XLA_COMMIT=${JAXCI_XLA_COMMIT:-}
11278
# When running tests locally, JAX artifacts are built with CUDA 12.3 and then
11379
# tested with a Docker image with CUDA 12.3 and CUDA 12.1. By default, we set
11480
# the CUDA version of the Docker image to 12.3.
115-
export JAXCI_DOCKER_CUDA_VERSION=${JAXCI_DOCKER_CUDA_VERSION:-12.3}
81+
export JAXCI_DOCKER_CUDA_VERSION=${JAXCI_DOCKER_CUDA_VERSION:-12.3}

‎ci/utilities/setup_envs.sh‎

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -17,13 +17,14 @@
1717
# Source JAXCI environment variables.
1818

1919
# If the user has not passed in an JAXCI_ENV_FILE, exit.
20-
if [[ -z "${JAXCI_ENV_FILE}" ]]; then
21-
echo "JAXCI_ENV_FILE is not set."
22-
echo "setup_envs.sh requires that `JAXCI_ENV_FILE` be set."
23-
echo "If you are looking to build JAX artifacts, please set JAXCI_ENV_FILE"
24-
echo "to an env file in the ci/envs/build_artifacts directory."
25-
echo "If you are looking to run JAX tests, please set JAXCI_ENV_FILE to an"
26-
echo "env file in the ci/envs/run_tests directory."
20+
if [[ -z "$1" ]]; then
21+
echo "ERROR: No argument passed."
22+
echo "setup_envs.sh requires that a path to a JAX CI env file be passed as"
23+
echo "an argument when invoking the build scripts."
24+
echo "If you are looking to build JAX artifacts, please pass in a"
25+
echo "corresponding env file from the ci/envs/build_artifacts directory."
26+
echo "If you are looking to run JAX tests, please pass in a"
27+
echo "corresponding env file from the ci/envs/run_tests directory."
2728
exit 1
2829
fi
2930

@@ -45,4 +46,4 @@ env | grep ^JAXCI_ > "$user_set_jaxci_envs"
4546
# -o history: record shell history
4647
# -o allexport: export all functions and variables to be available to subscripts
4748
set -exuo pipefail -o history -o allexport
48-
source "$JAXCI_ENV_FILE"
49+
source "$1"

0 commit comments

Comments
 (0)