JAXCI Environment Variables

This docpage describes the various JAXCI environment variables that are used in the CI scripts and their behaviors. These variables are used to control the behavior of the CI scripts such as the Python version used, path to JAX/XLA repo, if to clone XLA repo, etc.

Name	Default Value	Behavior	Usage
JAXCI_JAX_GIT_DIR	Present working directory: $(pwd)	Path to the JAX's Git directory.	Usage
JAXCI_HERMETIC_PYTHON_VERSION	System default	Controls the version of hermetic Python to use. This affects the Bazel commands only such as when building artifacts or when running the Bazel test scripts.	Usage
JAXCI_CUDA_VERSION	12	Controls the CUDA version to use when building the JAX artifacts or running the tests.	Usage
JAXCI_ROCM_VERSION	7	Controls the ROCm version to use when building the JAX artifacts or running the tests.	Usage
JAXCI_XLA_GIT_DIR	Unset	When using a local copy of XLA, this points to the root of the XLA git repository.	Usage
JAXCI_CLONE_MAIN_XLA	0	If set to 1, the XLA repository is cloned at HEAD and its path is set in JAXCI_XLA_GIT_DIR	Usage
JAXCI_XLA_COMMIT	Unset	Allows overriding the XLA commit that is used when using a local copy of XLA.	Usage
JAXCI_OUTPUT_DIR	$(pwd)/dist	Controls the location where the artifacts are written to. The directory will be automatically created if it does not exist.	Usage
JAXCI_BUILD_ARTIFACT_WITH_RBE	0	When set to 1, Bazel will use RBE to build the artifacts. Requires gcloud authentication and only certain platforms support RBE so this typically only set in CI builds	Usage
JAXCI_WRITE_TO_BAZEL_REMOTE_CACHE	0	When set to 1, Bazel will also try to push new cache entries to the cache bucket. Since writes to the bucket require authentication, this flag is enabled only for CI builds. Note that the builds using RBE use the RBE cache and not Bazel's remote cache, therefore this variable is a no-op if JAXCI_BUILD_ARTIFACT_WITH_RBE is set to 1. When JAXCI_BUILD_ARTIFACT_WITH_RBE and JAXCI_WRITE_TO_BAZEL_REMOTE_CACHE are both not set, Bazel will still read from the public cache bucket to try to speed up the build.	Usage
JAXCI_ARTIFACT_TYPE	"default"	Controls the type of artifacts to build. Valid values are "default", "release", "nightly". This affects the wheel tag and metadata, see ci/build_artifacts.sh to understand how.	Usage
JAXCI_WHEEL_RC_VERSION	Unset	During the release process, we build a Release Candidate (RC) wheel in addition to the release wheel. This environment variable sets the version of the RC wheel to build. Values are set internally.	Usage
JAXCI_PYTHON	python${JAXCI_HERMETIC_PYTHON_VERSION}	Points to the system Python binary to use. It used by scripts that make use of the system Python such as the Pytest scripts.	Usage
JAXCI_ENABLE_X64	0	By default, JAX enforces single-precision numbers to mitigate the Numpy API’s tendency to aggressively promote operands to double. When set to 1, the tests will use double-precision numbers.	Usage
JAXCI_TPU_CORES	Unset	Sets the number of TPU cores for the TPU machine type. Values are set in the workflow files.	Usage
JAXCI_RUN_FULL_TPU_TEST_SUITE	0	When set to 1, the full TPU test suite is run. Otherwise, a subset of tests is run.	Usage
JAXCI_JAX_PYPI_EXTRAS	Unset	Used to control the installation of JAX extras from PyPI. See JAX's setup.py for the list of valid values.	Usage
JAXCI_BUILD_JAXLIB	true	Used to control the value of build_jaxlib flag.	Usage
JAXCI_BUILD_JAX	true	Used to control the value of build_jax flag.	Usage
JAXCI_BAZEL_OUTPUT_BASE	Unset	Used to control the output base for Bazel builds.	Usage
JAXCI_BAZEL_CPU_RBE_MODE	test	Used to control whether to run or build the CPU test targets.	Usage

Docker Specific Environment Variables

[!NOTE] The following environment variables only affect the build if the run_docker_container.sh script was invoked to start a Docker container and the build is running inside that container. Typically, this would be the internal CI builds and local builds. Note that while GitHub actions use the same Docker images, they do not invoke "run_docker_container.sh" as they leverage built-in containerization features to run jobs within a container.

Name	Default Value	Behavior	Usage
JAXCI_DOCKER_WORK_DIR	"/jax"	The path on the container where the JAX Git repository is mounted to.	Usage
JAXCI_DOCKER_ARGS	Empty String	Space separated string of additional arguments that will be passed when starting the Docker container	Usage
JAXCI_DOCKER_IMAGE	Depends on the system (see ci/envs/docker.env)	Docker image to pull	Usage